pax_global_header00006660000000000000000000000064151545161520014517gustar00rootroot0000000000000052 comment=b410fb26d01c7cefc4c6e3443a53562164001371 apache-arrow-java-19.0.0/000077500000000000000000000000001515451615200150765ustar00rootroot00000000000000apache-arrow-java-19.0.0/.asf.yaml000066400000000000000000000027211515451615200166130ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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://github.com/apache/infrastructure-asfyaml/blob/main/README.md github: description: "Official Java implementation of Apache Arrow" homepage: https://arrow.apache.org/java/ labels: - apache-arrow - java del_branch_on_merge: true enabled_merge_buttons: merge: false 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 apache-arrow-java-19.0.0/.cmake-format.py000066400000000000000000000056041515451615200201010ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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 apache-arrow-java-19.0.0/.editorconfig000066400000000000000000000014761515451615200175630ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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 apache-arrow-java-19.0.0/.env000066400000000000000000000041471515451615200156750ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # All of the following environment variables are required to set default values # for the parameters in docker-compose.yml. # empty prefix means that the docker-compose configuration will use named # volumes which potentially improves the performance on docker for macos and # docker for windows, it also prevents the contamination of the source # directory # a non-empty prefix means that directories from the host are bind-mounted # into the container, it should be set to ".docker/" on github actions to keep # the cache plugin functional DOCKER_VOLUME_PREFIX= # turn on inline build cache, this is a docker buildx feature documented # at https://github.com/docker/buildx#--cache-tonametypetypekeyvalue BUILDKIT_INLINE_CACHE=1 COMPOSE_DOCKER_CLI_BUILD=1 DOCKER_BUILDKIT=1 # different architecture notations ARCH=amd64 ARCH_ALIAS=x86_64 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.9 # Versions for various dependencies used to build artifacts # Keep in sync with apache/arrow ARROW_REPO_ROOT=./arrow VCPKG="4334d8b4c8916018600212ab4dd4bbdc343065d1" # 2025.09.17 Release apache-arrow-java-19.0.0/.github/000077500000000000000000000000001515451615200164365ustar00rootroot00000000000000apache-arrow-java-19.0.0/.github/CODEOWNERS000066400000000000000000000051461515451615200200370ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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 apache-arrow-java-19.0.0/.github/ISSUE_TEMPLATE/000077500000000000000000000000001515451615200206215ustar00rootroot00000000000000apache-arrow-java-19.0.0/.github/ISSUE_TEMPLATE/bug_report.yaml000066400000000000000000000021441515451615200236560ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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: Bug Report description: File a bug report labels: ["Type: bug"] assignees: [] body: - type: textarea id: description attributes: label: Describe the bug, including details regarding any error messages, version, and platform. description: Please include what you expected. validations: required: true apache-arrow-java-19.0.0/.github/ISSUE_TEMPLATE/config.yml000066400000000000000000000014561515451615200226170ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. blank_issues_enabled: false apache-arrow-java-19.0.0/.github/ISSUE_TEMPLATE/feature_request.yaml000066400000000000000000000022601515451615200247100ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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: Enhancement Request description: Request an enhancement to the project labels: ["Type: enhancement"] assignees: [] body: - type: markdown attributes: value: | Thanks for taking the time to share your feedback on ways Apache Arrow can be improved! - type: textarea id: description attributes: label: Describe the enhancement requested validations: required: true apache-arrow-java-19.0.0/.github/ISSUE_TEMPLATE/usage_question.yaml000066400000000000000000000036251515451615200245460ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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: Usage Question description: Ask a question labels: ["Type: usage"] assignees: [] body: - type: markdown attributes: value: > While we enable issues as a mechanism for new contributors and passers-by who are unfamiliar with Apache Software Foundation projects to ask questions and interact with the project, we encourage users to ask such questions on public mailing lists: * Development discussions: dev@arrow.apache.org (first subscribe by sending an e-mail to dev-subscribe@arrow.apache.org). * User discussions: user@arrow.apache.org (first subscribe by sending an e-mail to user-subscribe@arrow.apache.org). * Mailing list archives: https://arrow.apache.org/community/ Do not be surprised by responses to issues raised here directing you to those mailing lists, or to report a bug or feature request here. Thank you! - type: textarea id: description attributes: label: > Describe the usage question you have. Please include as many useful details as possible. validations: required: true apache-arrow-java-19.0.0/.github/dependabot.yml000066400000000000000000000022001515451615200212600ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. version: 2 updates: - package-ecosystem: "github-actions" directory: "/" schedule: interval: "weekly" commit-message: prefix: "MINOR: [CI] " open-pull-requests-limit: 10 - package-ecosystem: "maven" directory: "/" schedule: interval: "weekly" commit-message: prefix: "MINOR: " open-pull-requests-limit: 10 apache-arrow-java-19.0.0/.github/pull_request_template.md000066400000000000000000000002651515451615200234020ustar00rootroot00000000000000## What's Changed Please fill in a description of the changes here. **This contains breaking changes.** Closes #NNN. apache-arrow-java-19.0.0/.github/release.yml000066400000000000000000000020611515451615200206000ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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: - "*" apache-arrow-java-19.0.0/.github/workflows/000077500000000000000000000000001515451615200204735ustar00rootroot00000000000000apache-arrow-java-19.0.0/.github/workflows/comment_bot.yml000066400000000000000000000026501515451615200235270ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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: Comment Bot on: issue_comment: types: - created - edited jobs: issue_assign: name: "Assign issue" permissions: contents: read issues: write if: github.event.comment.body == 'take' runs-on: ubuntu-latest steps: - uses: actions/github-script@v8 with: github-token: ${{ secrets.GITHUB_TOKEN }} script: |- github.rest.issues.addAssignees({ owner: context.repo.owner, repo: context.repo.repo, issue_number: context.payload.issue.number, assignees: context.payload.comment.user.login }); apache-arrow-java-19.0.0/.github/workflows/dev.yml000066400000000000000000000031721515451615200217770ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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 on: pull_request: {} push: {} concurrency: group: ${{ github.repository }}-${{ github.ref }}-${{ github.workflow }} cancel-in-progress: true permissions: contents: read jobs: pre-commit: name: "pre-commit" runs-on: ubuntu-latest steps: - uses: actions/checkout@v6 with: fetch-depth: 0 persist-credentials: false - uses: actions/setup-python@v6 with: python-version: '3.x' - name: pre-commit (cache) uses: actions/cache@9255dc7a253b0ccc959486e2bca901246202afeb # v5.0.1 with: path: ~/.cache/pre-commit key: pre-commit-${{ hashFiles('.pre-commit-config.yaml') }} - name: pre-commit (--all-files) run: | python -m pip install pre-commit pre-commit run --show-diff-on-failure --color=always --all-files apache-arrow-java-19.0.0/.github/workflows/dev_pr.js000066400000000000000000000171131515451615200223130ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT 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"); }, }; apache-arrow-java-19.0.0/.github/workflows/dev_pr.yml000066400000000000000000000053731515451615200225050ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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@v6 with: fetch-depth: 0 persist-credentials: false - name: Ensure PR title format id: title-format uses: actions/github-script@v8 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@v8 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@v8 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@v8 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 }} apache-arrow-java-19.0.0/.github/workflows/dev_pr_milestone.sh000077500000000000000000000044101515451615200243670ustar00rootroot00000000000000#!/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 "$@" apache-arrow-java-19.0.0/.github/workflows/rc.yml000066400000000000000000000541711515451615200216320ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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@v6 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@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 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@70fc10c6e5e1ce46ad2ea6f2b72d43f7d47b13c3 # v8.0.0 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@v6 with: repository: apache/arrow path: arrow - name: Checkout apache/arrow-testing uses: actions/checkout@v6 with: repository: apache/arrow-testing path: arrow/testing - name: Checkout apache/parquet-testing uses: actions/checkout@v6 with: repository: apache/parquet-testing path: arrow/cpp/submodules/parquet-testing - uses: docker/login-action@b45d80f862d83dbcd57f89517bcf500b2ab88fb2 # v4.0.0 with: registry: ghcr.io username: ${{ github.actor }} password: ${{ secrets.GITHUB_TOKEN }} - name: Cache uses: actions/cache@9255dc7a253b0ccc959486e2bca901246202afeb # v5.0.1 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@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 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@70fc10c6e5e1ce46ad2ea6f2b72d43f7d47b13c3 # v8.0.0 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@v6 with: repository: apache/arrow path: arrow - name: Checkout apache/arrow-testing uses: actions/checkout@v6 with: repository: apache/arrow-testing path: arrow/testing - name: Checkout apache/parquet-testing uses: actions/checkout@v6 with: repository: apache/parquet-testing path: arrow/cpp/submodules/parquet-testing - name: Set up Python uses: actions/setup-python@v6 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@9255dc7a253b0ccc959486e2bca901246202afeb # v5.0.1 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@11)/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@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 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@70fc10c6e5e1ce46ad2ea6f2b72d43f7d47b13c3 # v8.0.0 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@v6 with: repository: apache/arrow path: arrow - name: Set up Java uses: actions/setup-java@v5 with: java-version: '11' 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@9255dc7a253b0ccc959486e2bca901246202afeb # v5.0.1 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@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 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@70fc10c6e5e1ce46ad2ea6f2b72d43f7d47b13c3 # v8.0.0 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@v6 with: repository: apache/arrow-testing path: testing - name: Cache ~/.m2 uses: actions/cache@9255dc7a253b0ccc959486e2bca901246202afeb # v5.0.1 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@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 with: name: release-binaries path: binaries/* - name: Upload docs uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 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@v6 with: cache: 'pip' - name: Download source archive uses: actions/download-artifact@70fc10c6e5e1ce46ad2ea6f2b72d43f7d47b13c3 # v8.0.0 with: name: release-source - name: Download Javadocs uses: actions/download-artifact@70fc10c6e5e1ce46ad2ea6f2b72d43f7d47b13c3 # v8.0.0 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@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 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@v6 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@70fc10c6e5e1ce46ad2ea6f2b72d43f7d47b13c3 # v8.0.0 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@70fc10c6e5e1ce46ad2ea6f2b72d43f7d47b13c3 # v8.0.0 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@70fc10c6e5e1ce46ad2ea6f2b72d43f7d47b13c3 # v8.0.0 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 }} apache-arrow-java-19.0.0/.github/workflows/release.yml000066400000000000000000000061551515451615200226450ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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: 5 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@v6 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}" apache-arrow-java-19.0.0/.github/workflows/test.yml000066400000000000000000000146541515451615200222070ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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: Test on: push: branches: - '**' - '!dependabot/**' tags: - '**' pull_request: concurrency: group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} cancel-in-progress: true permissions: contents: read env: DOCKER_VOLUME_PREFIX: ".docker/" jobs: ubuntu: 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, 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@v6 with: fetch-depth: 0 submodules: recursive - name: Cache Docker Volumes uses: actions/cache@v5 with: path: .docker key: maven-${{ matrix.jdk }}-${{ matrix.maven }}-${{ hashFiles('compose.yaml', '**/pom.xml', '**/*.java') }} restore-keys: maven-${{ matrix.jdk }}-${{ matrix.maven }}- - name: Execute Docker Build env: # Enables build caching, but not strictly required DEVELOCITY_ACCESS_KEY: ${{ secrets.DEVELOCITY_ACCESS_KEY }} run: | docker compose run \ -e CI=true \ -e "DEVELOCITY_ACCESS_KEY=$DEVELOCITY_ACCESS_KEY" \ ${{ matrix.image }} macos: name: ${{ matrix.arch }} macOS ${{ matrix.macos }} Java JDK ${{ matrix.jdk }} runs-on: macos-${{ matrix.macos }} if: ${{ !contains(github.event.pull_request.title, 'WIP') }} timeout-minutes: 30 strategy: fail-fast: false matrix: include: - arch: AMD64 jdk: 11 macos: 15-intel - arch: AArch64 jdk: 11 macos: latest steps: - name: Set up Java uses: actions/setup-java@v5 with: distribution: 'temurin' java-version: ${{ matrix.jdk }} - name: Checkout Arrow uses: actions/checkout@v6 with: fetch-depth: 0 submodules: recursive - name: Build shell: bash env: DEVELOCITY_ACCESS_KEY: ${{ secrets.DEVELOCITY_ACCESS_KEY }} run: ci/scripts/build.sh . build jni - name: Test shell: bash env: DEVELOCITY_ACCESS_KEY: ${{ secrets.DEVELOCITY_ACCESS_KEY }} run: ci/scripts/test.sh . build jni windows: name: AMD64 Windows Server 2022 Java JDK ${{ matrix.jdk }} runs-on: windows-latest if: ${{ !contains(github.event.pull_request.title, 'WIP') }} timeout-minutes: 30 strategy: fail-fast: false matrix: jdk: [11] steps: - name: Set up Java uses: actions/setup-java@v5 with: java-version: ${{ matrix.jdk }} distribution: 'temurin' - name: Checkout Arrow uses: actions/checkout@v6 with: fetch-depth: 0 submodules: recursive - name: Build shell: bash env: DEVELOCITY_ACCESS_KEY: ${{ secrets.DEVELOCITY_ACCESS_KEY }} run: ci/scripts/build.sh . build jni - name: Test shell: bash env: DEVELOCITY_ACCESS_KEY: ${{ secrets.DEVELOCITY_ACCESS_KEY }} run: ci/scripts/test.sh . build jni integration: name: AMD64 integration runs-on: ubuntu-latest timeout-minutes: 60 steps: - name: Checkout Arrow uses: actions/checkout@v6 with: fetch-depth: 0 repository: apache/arrow submodules: recursive - name: Checkout Arrow Rust uses: actions/checkout@v6 with: repository: apache/arrow-rs path: rust - name: Checkout Arrow nanoarrow uses: actions/checkout@v6 with: repository: apache/arrow-nanoarrow path: nanoarrow - name: Checkout Arrow .NET uses: actions/checkout@v6 with: repository: apache/arrow-dotnet path: dotnet - name: Checkout Arrow Go uses: actions/checkout@v6 with: repository: apache/arrow-go path: go - name: Checkout Arrow Java uses: actions/checkout@v6 with: path: java - name: Checkout Arrow JavaScript uses: actions/checkout@v6 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@v5 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 apache-arrow-java-19.0.0/.gitmodules000066400000000000000000000001311515451615200172460ustar00rootroot00000000000000[submodule "testing"] path = testing url = https://github.com/apache/arrow-testing.git apache-arrow-java-19.0.0/.mvn/000077500000000000000000000000001515451615200157545ustar00rootroot00000000000000apache-arrow-java-19.0.0/.mvn/develocity.xml000066400000000000000000000032771515451615200206560ustar00rootroot00000000000000 arrow https://develocity.apache.org false #{isFalse(env['CI'])} #{{'0.0.0.0'}} false apache-arrow-java-19.0.0/.mvn/extensions.xml000066400000000000000000000026501515451615200207000ustar00rootroot00000000000000 com.gradle develocity-maven-extension 2.3.4 com.gradle common-custom-user-data-maven-extension 2.1.0 apache-arrow-java-19.0.0/.pre-commit-config.yaml000066400000000000000000000041221515451615200213560ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # 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 name: Release Audit Tool language: system entry: | bash -c " \ git archive HEAD \ --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" apache-arrow-java-19.0.0/Brewfile000066400000000000000000000014631515451615200165640ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. brew "openjdk@11" brew "sccache" apache-arrow-java-19.0.0/CMakeLists.txt000066400000000000000000000073501515451615200176430ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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_minimum_required(VERSION 3.16) message(STATUS "Building using CMake version: ${CMAKE_VERSION}") # find_package() uses _ROOT variables. # https://cmake.org/cmake/help/latest/policy/CMP0074.html if(POLICY CMP0074) cmake_policy(SET CMP0074 NEW) endif() project(arrow-java-jni) if("${CMAKE_CXX_STANDARD}" STREQUAL "") set(CMAKE_CXX_STANDARD 17) endif() set(CMAKE_CXX_STANDARD_REQUIRED ON) # Components option(ARROW_JAVA_JNI_ENABLE_DEFAULT "Whether enable components by default or not" ON) option(ARROW_JAVA_JNI_ENABLE_C "Enable C data interface" ${ARROW_JAVA_JNI_ENABLE_DEFAULT}) option(ARROW_JAVA_JNI_ENABLE_DATASET "Enable dataset" ${ARROW_JAVA_JNI_ENABLE_DEFAULT}) option(ARROW_JAVA_JNI_ENABLE_GANDIVA "Enable Gandiva" ${ARROW_JAVA_JNI_ENABLE_DEFAULT}) option(ARROW_JAVA_JNI_ENABLE_ORC "Enable ORC" ${ARROW_JAVA_JNI_ENABLE_DEFAULT}) include(GNUInstallDirs) # ccache option(ARROW_JAVA_JNI_USE_CCACHE "Use ccache when compiling (if available)" ON) if(ARROW_USE_CCACHE AND NOT CMAKE_C_COMPILER_LAUNCHER AND NOT CMAKE_CXX_COMPILER_LAUNCHER) find_program(CCACHE ccache) if(CCACHE) message(STATUS "Using ccache: ${CCACHE}") set(CMAKE_C_COMPILER_LAUNCHER ${CCACHE}) set(CMAKE_CXX_COMPILER_LAUNCHER ${CCACHE}) # ARROW-3985: let ccache preserve C++ comments, because some of them may be # meaningful to the compiler set(ENV{CCACHE_COMMENTS} "1") endif() endif() # Build find_package(Java REQUIRED) find_package(JNI REQUIRED) include(UseJava) add_library(jni INTERFACE IMPORTED) set_target_properties(jni PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${JNI_INCLUDE_DIRS}") include(CTest) 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) target_link_libraries(arrow_java_test INTERFACE ArrowTesting::arrow_testing_static GTest::gtest_main) endif() # The ARROW_JAVA_JNI_ARCH_DIR will automatically be derived the normalized # operating system from system processor. The user can override this variable # if auto-detection fails. if("${ARROW_JAVA_JNI_ARCH_DIR}" STREQUAL "") if("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "aarch64") set(ARROW_JAVA_JNI_ARCH_DIR "aarch_64") elseif("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "i386") set(ARROW_JAVA_JNI_ARCH_DIR "x86_64") elseif("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "arm64") set(ARROW_JAVA_JNI_ARCH_DIR "aarch_64") elseif("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "AMD64") set(ARROW_JAVA_JNI_ARCH_DIR "x86_64") else() set(ARROW_JAVA_JNI_ARCH_DIR "${CMAKE_SYSTEM_PROCESSOR}") endif() endif() if(ARROW_JAVA_JNI_ENABLE_C) add_subdirectory(c) endif() if(ARROW_JAVA_JNI_ENABLE_DATASET) add_subdirectory(dataset) endif() if(ARROW_JAVA_JNI_ENABLE_GANDIVA) add_subdirectory(gandiva) endif() if(ARROW_JAVA_JNI_ENABLE_ORC) add_subdirectory(adapter/orc) endif() apache-arrow-java-19.0.0/CODE_OF_CONDUCT.md000066400000000000000000000016451515451615200177030ustar00rootroot00000000000000 # Code of Conduct * [Code of Conduct for The Apache Software Foundation][1] [1]: https://www.apache.org/foundation/policies/conduct.html apache-arrow-java-19.0.0/CONTRIBUTING.md000066400000000000000000000067711515451615200173420ustar00rootroot00000000000000 # How to contribute to Apache Arrow Java ## Did you find a bug? The Arrow Java project uses GitHub as a bug tracker. To report a bug, sign in to your GitHub account, navigate to [GitHub issues](https://github.com/apache/arrow-java/issues) and click on **New issue** . Before you create a new bug entry, we recommend you first search among 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. - [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 Any functionality change should have a GitHub issue opened. For minor changes that affect documentation, you do not need to open up a GitHub issue. Instead you can prefix the title of your PR with "MINOR: " if it meets one of the following: * Grammar, usage and spelling fixes that affect no more than 2 files * Documentation updates affecting no more than 2 files and not more than 500 words. ## Do you want to propose a significant new feature or an important refactoring? We ask that all discussions about major changes in the codebase happen publicly on the [arrow-dev mailing-list](https://mail-archives.apache.org/mod_mbox/arrow-dev/). ## Do you have questions about the source code, the build procedure or the development process? You can also ask on the mailing-list, see above. ## Further information Please read our [development documentation](https://arrow.apache.org/docs/developers/index.html) or look through the [New Contributor's Guide](https://arrow.apache.org/docs/developers/guide/index.html). apache-arrow-java-19.0.0/LICENSE.txt000066400000000000000000000276511515451615200167340ustar00rootroot00000000000000 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 product includes code from Netty 4.1.117.Final: * 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 -------------------------------------------------------------------------------- This product includes code from Google Guava 33.4.0-jre * memory/memory-core/src/main/java/org/apache/arrow/util/Preconditions.java Copyright: (C) 2011 The Guava Authors Home page: https://github.com/google/guava/ License: https://www.apache.org/licenses/LICENSE-2.0 apache-arrow-java-19.0.0/NOTICE.txt000066400000000000000000000231161515451615200166230ustar00rootroot00000000000000Apache 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 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 apache-arrow-java-19.0.0/README.md000066400000000000000000000137271515451615200163670ustar00rootroot00000000000000 # Arrow Java ## Getting Started The following guides explain the fundamental data structures used in the Java implementation of Apache Arrow. - [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/java/current/). ## Building from source Refer to [Building Apache Arrow](https://arrow.apache.org/docs/dev/developers/java/building.html) for documentation of environment setup and build instructions. ## Flatbuffers dependency Arrow uses Google's Flatbuffers to transport metadata. The java version of the library requires the generated flatbuffer classes can only be used with the same version that generated them. Arrow packages a version of the arrow-vector module that shades flatbuffers and arrow-format into a single JAR. Using the classifier "shade-format-flatbuffers" in your `pom.xml` will make use of this JAR, you can then exclude/resolve the original dependency to a version of your choosing. ### Updating the flatbuffers generated code 1. Verify that your version of flatc matches the declared dependency: ```bash $ flatc --version flatc version 25.1.24 $ grep "dep.fbs.version" pom.xml 25.1.24 ``` 2. Generate the flatbuffer java files by performing the following: ```bash cd $ARROW_HOME # remove the existing files rm -rf format/src # regenerate from the .fbs files flatc --java -o format/src/main/java arrow-format/*.fbs # prepend license header mvn spotless:apply -pl :arrow-format ``` ## Performance Tuning There are several system/environmental variables that users can configure. These trade off safety (they turn off checking) for speed. Typically they are only used in production settings after the code has been thoroughly tested without using them. * Bounds Checking for memory accesses: Bounds checking is on by default. You can disable it by setting either the system property(`arrow.enable_unsafe_memory_access`) or the environmental variable (`ARROW_ENABLE_UNSAFE_MEMORY_ACCESS`) to `true`. When both the system property and the environmental variable are set, the system property takes precedence. * null checking for gets: `ValueVector` get methods (not `getObject`) methods by default verify the slot is not null. You can disable it by setting either the system property(`arrow.enable_null_check_for_get`) or the environmental variable (`ARROW_ENABLE_NULL_CHECK_FOR_GET`) to `false`. When both the system property and the environmental variable are set, the system property takes precedence. ## Java Properties * `-Dio.netty.tryReflectionSetAccessible=true` should be set. This fixes `java.lang.UnsupportedOperationException: sun.misc.Unsafe or java.nio.DirectByteBuffer.(long, int) not available`. thrown by Netty. * To support duplicate fields in a `StructVector` enable `-Darrow.struct.conflict.policy=CONFLICT_APPEND`. Duplicate fields are ignored (`CONFLICT_REPLACE`) by default and overwritten. To support different policies for conflicting or duplicate fields set this JVM flag or use the correct static constructor methods for `StructVector`s. ## Java Code Style Guide 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, standard Java, org.\*, com.\* * Line length can be up to 120 characters * Operators for line wrapping are at end-of-line * Naming rules for methods, parameters, etc. have been relaxed * Disabled `NoFinalizer`, `OverloadMethodsDeclarationOrder`, and `VariableDeclarationUsageDistance` due to the existing code base. These rules should be followed when possible. Refer to [checkstyle.xml](dev/checkstyle/checkstyle.xml) for rule specifics. ## Test Logging Configuration When running tests, Arrow Java uses the Logback logger with SLF4J. By default, it uses the `logback.xml` present in the corresponding module's `src/test/resources` directory, which has the default log level set to `INFO`. Arrow Java can be built with an alternate logback configuration file using the following command run in the project root directory: ```bash mvn -Dlogback.configurationFile=file: ``` 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](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: ```bash mvn -Pintegration-tests ``` apache-arrow-java-19.0.0/adapter/000077500000000000000000000000001515451615200165165ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/000077500000000000000000000000001515451615200174655ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/pom.xml000066400000000000000000000047321515451615200210100ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-java-root 19.0.0 ../../pom.xml arrow-avro Arrow AVRO Adapter (Contrib/Experimental) A library for converting Avro data to Arrow data. http://maven.apache.org org.apache.arrow arrow-memory-core org.apache.arrow arrow-memory-netty runtime org.apache.arrow arrow-vector org.immutables value-annotations org.apache.avro avro ${dep.avro.version} org.apache.maven.plugins maven-compiler-plugin -Werror apache-arrow-java-19.0.0/adapter/avro/src/000077500000000000000000000000001515451615200202545ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/000077500000000000000000000000001515451615200212005ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/000077500000000000000000000000001515451615200221215ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/module-info.java000066400000000000000000000022721515451615200252050ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ 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; requires org.apache.arrow.vector; requires org.apache.avro; } apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/000077500000000000000000000000001515451615200227105ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/000077500000000000000000000000001515451615200241315ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200252635ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/000077500000000000000000000000001515451615200267035ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/000077500000000000000000000000001515451615200276525ustar00rootroot00000000000000ArrowToAvroUtils.java000066400000000000000000000755101515451615200337140ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } } apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrow.java000066400000000000000000000063411515451615200327460ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.io.IOException; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.avro.Schema; import org.apache.avro.io.Decoder; /** Utility class to convert Avro objects to columnar Arrow format objects. */ public class AvroToArrow { /** * Fetch the data from {@link Decoder} and convert it to Arrow objects. Only for testing purpose. * * @param schema avro schema. * @param decoder avro decoder * @param config configuration of the conversion. * @return Arrow Data Objects {@link VectorSchemaRoot} */ static VectorSchemaRoot avroToArrow(Schema schema, Decoder decoder, AvroToArrowConfig config) throws IOException { Preconditions.checkNotNull(schema, "Avro schema object cannot be null"); Preconditions.checkNotNull(decoder, "Avro decoder object cannot be null"); Preconditions.checkNotNull(config, "config cannot be null"); return AvroToArrowUtils.avroToArrowVectors(schema, decoder, config); } /** * Fetch the data from {@link Decoder} and iteratively convert it to Arrow objects. * * @param schema avro schema * @param decoder avro decoder * @param config configuration of the conversion. * @throws IOException on error */ public static AvroToArrowVectorIterator avroToArrowIterator( Schema schema, Decoder decoder, AvroToArrowConfig config) throws IOException { Preconditions.checkNotNull(schema, "Avro schema object cannot be null"); Preconditions.checkNotNull(decoder, "Avro decoder object cannot be null"); Preconditions.checkNotNull(config, "config cannot be null"); 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); } } AvroToArrowConfig.java000066400000000000000000000103361515451615200340140ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.Set; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.dictionary.DictionaryProvider; /** This class configures the Avro-to-Arrow conversion process. */ 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. */ private final int targetBatchSize; /** * The dictionary provider used for enum type. If avro schema has enum type, will create * dictionary and update this provider. */ private final DictionaryProvider.MapDictionaryProvider provider; /** 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. * * @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. */ AvroToArrowConfig( BufferAllocator allocator, int targetBatchSize, DictionaryProvider.MapDictionaryProvider provider, Set skipFieldNames) { 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; // 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() { return allocator; } public int getTargetBatchSize() { return targetBatchSize; } public DictionaryProvider.MapDictionaryProvider getProvider() { return provider; } public Set getSkipFieldNames() { return skipFieldNames; } public boolean isLegacyMode() { return legacyMode; } } AvroToArrowConfigBuilder.java000066400000000000000000000043331515451615200353230ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.HashSet; import java.util.Set; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.dictionary.DictionaryProvider; /** This class builds {@link AvroToArrowConfig}s. */ public class AvroToArrowConfigBuilder { private BufferAllocator allocator; private int targetBatchSize; private DictionaryProvider.MapDictionaryProvider provider; private Set skipFieldNames; /** Default constructor for the {@link AvroToArrowConfigBuilder}. */ public AvroToArrowConfigBuilder(BufferAllocator allocator) { this.allocator = allocator; this.targetBatchSize = AvroToArrowVectorIterator.DEFAULT_BATCH_SIZE; this.provider = new DictionaryProvider.MapDictionaryProvider(); this.skipFieldNames = new HashSet<>(); } public AvroToArrowConfigBuilder setTargetBatchSize(int targetBatchSize) { this.targetBatchSize = targetBatchSize; return this; } public AvroToArrowConfigBuilder setProvider(DictionaryProvider.MapDictionaryProvider provider) { this.provider = provider; return this; } public AvroToArrowConfigBuilder setSkipFieldNames(Set skipFieldNames) { this.skipFieldNames = skipFieldNames; return this; } /** This builds the {@link AvroToArrowConfig} from the provided params. */ public AvroToArrowConfig build() { return new AvroToArrowConfig(allocator, targetBatchSize, provider, skipFieldNames); } } AvroToArrowUtils.java000066400000000000000000001304441515451615200337120ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.vector.types.FloatingPointPrecision.DOUBLE; import static org.apache.arrow.vector.types.FloatingPointPrecision.SINGLE; import java.io.EOFException; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Locale; import java.util.Map; import java.util.Set; import java.util.stream.Collectors; import org.apache.arrow.adapter.avro.consumers.AvroArraysConsumer; import org.apache.arrow.adapter.avro.consumers.AvroBooleanConsumer; import org.apache.arrow.adapter.avro.consumers.AvroBytesConsumer; import org.apache.arrow.adapter.avro.consumers.AvroDoubleConsumer; import org.apache.arrow.adapter.avro.consumers.AvroEnumConsumer; import org.apache.arrow.adapter.avro.consumers.AvroFixedConsumer; import org.apache.arrow.adapter.avro.consumers.AvroFloatConsumer; import org.apache.arrow.adapter.avro.consumers.AvroIntConsumer; 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; import org.apache.arrow.adapter.avro.consumers.CompositeAvroConsumer; import org.apache.arrow.adapter.avro.consumers.Consumer; 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; 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.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.UnionVector; 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.TimeUnit; import org.apache.arrow.vector.types.Types; 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.arrow.vector.util.JsonStringArrayList; import org.apache.arrow.vector.util.ValueVectorUtility; import org.apache.avro.LogicalType; import org.apache.avro.LogicalTypes; import org.apache.avro.Schema; import org.apache.avro.io.Decoder; /** * Class that does most of the work to convert Avro data into Arrow columnar format Vector objects. */ public class AvroToArrowUtils { /** * Creates a {@link Consumer} from the {@link Schema} * *

This method currently performs following type mapping for Avro data types to corresponding * Arrow data types. * *

    *
  • STRING --> ArrowType.Utf8 *
  • INT --> ArrowType.Int(32, signed) *
  • LONG --> ArrowType.Int(64, signed) *
  • FLOAT --> ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE) *
  • DOUBLE --> ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE) *
  • BOOLEAN --> ArrowType.Bool *
  • BYTES --> ArrowType.Binary *
  • ARRAY --> ArrowType.List *
  • MAP --> ArrowType.Map *
  • FIXED --> ArrowType.FixedSizeBinary *
  • RECORD --> ArrowType.Struct *
  • UNION --> ArrowType.Union *
  • ENUM--> ArrowType.Int *
  • DECIMAL --> ArrowType.Decimal *
  • Date --> ArrowType.Date(DateUnit.DAY) *
  • TimeMillis --> ArrowType.Time(TimeUnit.MILLISECOND, 32) *
  • TimeMicros --> ArrowType.Time(TimeUnit.MICROSECOND, 64) *
  • TimestampMillis --> ArrowType.Timestamp(TimeUnit.MILLISECOND, null) *
  • TimestampMicros --> ArrowType.Timestamp(TimeUnit.MICROSECOND, null) *
*/ private static Consumer createConsumer(Schema schema, String name, AvroToArrowConfig config) { return createConsumer(schema, name, false, config, null); } private static Consumer createConsumer( Schema schema, String name, AvroToArrowConfig config, FieldVector vector) { return createConsumer(schema, name, false, config, vector); } /** * Create a consumer with the given Avro schema. * * @param schema avro schema * @param name arrow field name * @param consumerVector vector to keep in consumer, if v == null, will create a new vector via * field. * @return consumer */ private static Consumer createConsumer( Schema schema, String name, boolean nullable, AvroToArrowConfig config, FieldVector consumerVector) { Preconditions.checkNotNull(schema, "Avro schema object can't be null"); Preconditions.checkNotNull(config, "Config can't be null"); final BufferAllocator allocator = config.getAllocator(); final Schema.Type type = schema.getType(); final LogicalType logicalType = schema.getLogicalType(); final ArrowType arrowType; final FieldType fieldType; final FieldVector vector; final Consumer consumer; switch (type) { case UNION: 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, nullable, config, consumerVector); break; case MAP: consumer = createMapConsumer(schema, name, nullable, config, consumerVector); break; case RECORD: consumer = createStructConsumer(schema, name, nullable, config, consumerVector); break; case ENUM: consumer = createEnumConsumer(schema, name, nullable, config, consumerVector); break; case STRING: arrowType = new ArrowType.Utf8(); 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, config); if (logicalType instanceof LogicalTypes.Decimal) { arrowType = createDecimalArrowType((LogicalTypes.Decimal) logicalType, schema); fieldType = new FieldType( nullable, arrowType, /* dictionary= */ null, getMetaData(schema, extProps, config)); vector = createVector(consumerVector, fieldType, name, allocator); 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, config)); vector = createVector(consumerVector, fieldType, name, allocator); consumer = new AvroFixedConsumer((FixedSizeBinaryVector) vector, schema.getFixedSize()); } break; case INT: if (logicalType instanceof LogicalTypes.Date) { arrowType = new ArrowType.Date(DateUnit.DAY); 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, 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, 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, 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, config)); vector = createVector(consumerVector, fieldType, name, allocator); consumer = new AvroTimeMicroConsumer((TimeMicroVector) vector); } 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, config)); vector = createVector(consumerVector, fieldType, name, allocator); consumer = new AvroTimestampMillisConsumer((TimeStampMilliVector) vector); } 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, 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, 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, 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, config)); vector = createVector(consumerVector, fieldType, name, allocator); consumer = new AvroDoubleConsumer((Float8Vector) vector); break; case BYTES: if (logicalType instanceof LogicalTypes.Decimal) { 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); 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, 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, config)); vector = new NullVector(name, fieldType); // Respect nullability defined in fieldType consumer = new AvroNullConsumer((NullVector) vector); break; default: // no-op, shouldn't get here throw new UnsupportedOperationException( "Can't convert avro type %s to arrow type." + type.getName()); } return consumer; } private static ArrowType createDecimalArrowType(LogicalTypes.Decimal logicalType, Schema schema) { final int scale = logicalType.getScale(); final int precision = logicalType.getPrecision(); Preconditions.checkArgument( 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); 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) { SkipFunction skipFunction; Schema.Type type = schema.getType(); switch (type) { case UNION: List unionDelegates = schema.getTypes().stream().map(s -> createSkipConsumer(s)).collect(Collectors.toList()); skipFunction = decoder -> unionDelegates.get(decoder.readInt()).consume(decoder); break; case ARRAY: Consumer elementDelegate = createSkipConsumer(schema.getElementType()); skipFunction = decoder -> { for (long i = decoder.skipArray(); i != 0; i = decoder.skipArray()) { for (long j = 0; j < i; j++) { elementDelegate.consume(decoder); } } }; break; case MAP: Consumer valueDelegate = createSkipConsumer(schema.getValueType()); skipFunction = decoder -> { for (long i = decoder.skipMap(); i != 0; i = decoder.skipMap()) { for (long j = 0; j < i; j++) { decoder.skipString(); // Discard key valueDelegate.consume(decoder); } } }; break; case RECORD: List delegates = schema.getFields().stream() .map(field -> createSkipConsumer(field.schema())) .collect(Collectors.toList()); skipFunction = decoder -> { for (Consumer consumer : delegates) { consumer.consume(decoder); } }; break; case ENUM: skipFunction = decoder -> decoder.readEnum(); break; case STRING: skipFunction = decoder -> decoder.skipString(); break; case FIXED: skipFunction = decoder -> decoder.skipFixed(schema.getFixedSize()); break; case INT: skipFunction = decoder -> decoder.readInt(); break; case BOOLEAN: skipFunction = decoder -> decoder.skipFixed(1); break; case LONG: skipFunction = decoder -> decoder.readLong(); break; case FLOAT: skipFunction = decoder -> decoder.readFloat(); break; case DOUBLE: skipFunction = decoder -> decoder.readDouble(); break; case BYTES: skipFunction = decoder -> decoder.skipBytes(); break; case NULL: skipFunction = decoder -> {}; break; default: // no-op, shouldn't get here throw new UnsupportedOperationException("Invalid avro type: " + type.getName()); } 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<>(); final Set skipFieldNames = config.getSkipFieldNames(); Schema.Type type = schema.getType(); if (type == Schema.Type.RECORD) { for (Schema.Field field : schema.getFields()) { if (skipFieldNames.contains(field.name())) { consumers.add(createSkipConsumer(field.schema())); } else { Consumer consumer = createConsumer(field.schema(), field.name(), config); consumers.add(consumer); } } } else { Consumer consumer = createConsumer(schema, "", config); consumers.add(consumer); } return new CompositeAvroConsumer(consumers); } private static FieldVector createVector( FieldVector consumerVector, FieldType fieldType, String name, BufferAllocator allocator) { return consumerVector != null ? consumerVector : fieldType.createNewSingleVector(name, allocator, null); } private static String getDefaultFieldName(ArrowType type) { Types.MinorType minorType = Types.getMinorTypeForArrowType(type); return minorType.name().toLowerCase(Locale.ROOT); } private static Field avroSchemaToField(Schema schema, String name, AvroToArrowConfig config) { 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(); final List children = new ArrayList<>(); final FieldType fieldType; switch (type) { case UNION: 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); } break; case ARRAY: Schema elementSchema = schema.getElementType(); 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(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); Field structField = new Field( MapVector.DATA_VECTOR_NAME, structFieldType, Arrays.asList(keyField, valueField)); children.add(structField); fieldType = createFieldType( nullable, new ArrowType.Map(/* keysSorted= */ false), schema, externalProps, config); break; case RECORD: final Set skipFieldNames = config.getSkipFieldNames(); for (int i = 0; i < schema.getFields().size(); i++) { final Schema.Field field = schema.getFields().get(i); Schema childSchema = field.schema(); String fullChildName = String.format("%s.%s", name, field.name()); if (!skipFieldNames.contains(fullChildName)) { final Map extProps = new HashMap<>(); String doc = field.doc(); Set aliases = field.aliases(); if (doc != null) { extProps.put("doc", doc); } if (aliases != null && (!aliases.isEmpty() || config.isLegacyMode())) { extProps.put("aliases", convertAliases(aliases)); } children.add(avroSchemaToField(childSchema, fullChildName, config, extProps)); } } fieldType = createFieldType(nullable, new ArrowType.Struct(), schema, externalProps, config); break; case ENUM: DictionaryProvider.MapDictionaryProvider provider = config.getProvider(); int current = provider.getDictionaryIds().size(); int enumCount = schema.getEnumSymbols().size(); ArrowType.Int indexType = DictionaryEncoder.getIndexType(enumCount); fieldType = createFieldType( nullable, indexType, schema, externalProps, new DictionaryEncoding(current, /* ordered= */ false, /* indexType= */ indexType), config); break; case STRING: 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, schema); } else { fixedArrowType = new ArrowType.FixedSizeBinary(schema.getFixedSize()); } fieldType = createFieldType(nullable, fixedArrowType, schema, externalProps, config); break; case INT: final ArrowType intArrowType; if (logicalType instanceof LogicalTypes.Date) { intArrowType = new ArrowType.Date(DateUnit.DAY); } else if (logicalType instanceof LogicalTypes.TimeMillis) { intArrowType = new ArrowType.Time(TimeUnit.MILLISECOND, 32); } else { intArrowType = new ArrowType.Int(32, /* isSigned= */ true); } fieldType = createFieldType(nullable, intArrowType, schema, externalProps, config); break; case BOOLEAN: 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) { // 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); } fieldType = createFieldType(nullable, longArrowType, schema, externalProps, config); break; case FLOAT: fieldType = createFieldType( nullable, new ArrowType.FloatingPoint(SINGLE), schema, externalProps, config); break; case DOUBLE: 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, schema); } else { bytesArrowType = new ArrowType.Binary(); } fieldType = createFieldType(nullable, bytesArrowType, schema, externalProps, config); break; case NULL: fieldType = createFieldType(ArrowType.Null.INSTANCE, schema, externalProps, config); break; default: // no-op, shouldn't get here throw new UnsupportedOperationException(); } 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, boolean nullable, AvroToArrowConfig config, FieldVector consumerVector) { ListVector listVector; if (consumerVector == null) { final Field field = avroSchemaToField(schema, name, nullable, config, /* externalProps= */ null); listVector = (ListVector) field.createVector(config.getAllocator()); } else { listVector = (ListVector) consumerVector; } FieldVector dataVector = listVector.getDataVector(); // create delegate Schema childSchema = schema.getElementType(); Consumer delegate = createConsumer(childSchema, childSchema.getName(), config, dataVector); return new AvroArraysConsumer(listVector, delegate); } private static Consumer createStructConsumer( 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, nullable, config, createExternalProps(schema, config)); structVector = (StructVector) field.createVector(config.getAllocator()); } else { structVector = (StructVector) consumerVector; } Consumer[] delegates = new Consumer[schema.getFields().size()]; int vectorIndex = 0; for (int i = 0; i < schema.getFields().size(); i++) { Schema.Field childField = schema.getFields().get(i); Consumer delegate; // use full name to distinguish fields have same names between parent and child fields. final String fullChildName = String.format("%s.%s", name, childField.name()); if (skipFieldNames.contains(fullChildName)) { delegate = createSkipConsumer(childField.schema()); } else { delegate = createConsumer( childField.schema(), fullChildName, config, structVector.getChildrenFromFields().get(vectorIndex++)); } delegates[i] = delegate; } return new AvroStructConsumer(structVector, delegates); } private static Consumer createEnumConsumer( Schema schema, String name, boolean nullable, AvroToArrowConfig config, FieldVector consumerVector) { BaseIntVector indexVector; if (consumerVector == null) { final Field field = avroSchemaToField(schema, name, nullable, config, createExternalProps(schema, config)); indexVector = (BaseIntVector) field.createVector(config.getAllocator()); } else { indexVector = (BaseIntVector) consumerVector; } final int valueCount = schema.getEnumSymbols().size(); VarCharVector dictVector = new VarCharVector(name, config.getAllocator()); dictVector.allocateNewSafe(); dictVector.setValueCount(valueCount); for (int i = 0; i < valueCount; i++) { dictVector.set(i, schema.getEnumSymbols().get(i).getBytes(StandardCharsets.UTF_8)); } Dictionary dictionary = new Dictionary(dictVector, indexVector.getField().getDictionary()); config.getProvider().put(dictionary); return new AvroEnumConsumer(indexVector); } private static Consumer createMapConsumer( Schema schema, String name, boolean nullable, AvroToArrowConfig config, FieldVector consumerVector) { MapVector mapVector; if (consumerVector == null) { final Field field = avroSchemaToField(schema, name, nullable, config, /* externalProps= */ null); mapVector = (MapVector) field.createVector(config.getAllocator()); } else { mapVector = (MapVector) consumerVector; } // create delegate struct consumer StructVector structVector = (StructVector) mapVector.getDataVector(); // keys in avro map are always assumed to be strings. Consumer keyConsumer = new AvroStringConsumer((VarCharVector) structVector.getChildrenFromFields().get(0)); Consumer valueConsumer = createConsumer( schema.getValueType(), schema.getValueType().getName(), config, structVector.getChildrenFromFields().get(1)); AvroStructConsumer internalConsumer = new AvroStructConsumer(structVector, new Consumer[] {keyConsumer, valueConsumer}); return new AvroMapConsumer(mapVector, internalConsumer); } private static Consumer createUnionConsumer( Schema schema, String name, boolean nullableUnion, AvroToArrowConfig config, FieldVector consumerVector) { final int size = schema.getTypes().size(); UnionVector unionVector; if (consumerVector == null) { final Field field = avroSchemaToField(schema, name, config); unionVector = (UnionVector) field.createVector(config.getAllocator()); } else { unionVector = (UnionVector) consumerVector; } List childVectors = unionVector.getChildrenFromFields(); Consumer[] delegates = new Consumer[size]; Types.MinorType[] types = new Types.MinorType[size]; for (int i = 0; i < size; i++) { FieldVector child = childVectors.get(i); Schema subSchema = schema.getTypes().get(i); Consumer delegate = createConsumer(subSchema, subSchema.getName(), nullableUnion, config, child); delegates[i] = delegate; types[i] = child.getMinorType(); } return new AvroUnionsConsumer(unionVector, delegates, types); } /** * Read data from {@link Decoder} and generate a {@link VectorSchemaRoot}. * * @param schema avro schema * @param decoder avro decoder to read data from */ static VectorSchemaRoot avroToArrowVectors( Schema schema, Decoder decoder, AvroToArrowConfig config) throws IOException { List vectors = new ArrayList<>(); List consumers = new ArrayList<>(); final Set skipFieldNames = config.getSkipFieldNames(); Schema.Type type = schema.getType(); if (type == Schema.Type.RECORD) { for (Schema.Field field : schema.getFields()) { if (skipFieldNames.contains(field.name())) { consumers.add(createSkipConsumer(field.schema())); } else { Consumer consumer = createConsumer(field.schema(), field.name(), config); consumers.add(consumer); vectors.add(consumer.getVector()); } } } else { Consumer consumer = createConsumer(schema, "", config); consumers.add(consumer); vectors.add(consumer.getVector()); } long validConsumerCount = consumers.stream().filter(c -> !c.skippable()).count(); Preconditions.checkArgument( vectors.size() == validConsumerCount, "vectors size not equals consumers size."); List fields = vectors.stream().map(t -> t.getField()).collect(Collectors.toList()); VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, 0); CompositeAvroConsumer compositeConsumer = new CompositeAvroConsumer(consumers); int valueCount = 0; try { while (true) { ValueVectorUtility.ensureCapacity(root, valueCount + 1); compositeConsumer.consume(decoder); valueCount++; } } catch (EOFException eof) { // reach the end of encoder stream. root.setRowCount(valueCount); } catch (Exception e) { compositeConsumer.close(); throw new UnsupportedOperationException("Error occurs while consume process.", e); } return root; } // 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<>(); 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, AvroToArrowConfig config) { Map metadata = getMetaData(schema, config); if (externalProps != null) { metadata.putAll(externalProps); } return metadata; } /** Parse avro attributes and convert them to metadata. */ 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 && (!aliases.isEmpty() || config.isLegacyMode())) { extProps.put("aliases", convertAliases(aliases)); } return extProps; } private static FieldType createFieldType( 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, AvroToArrowConfig config) { return new FieldType( nullable, arrowType, dictionary, getMetaData(schema, externalProps, config)); } private static String convertAliases(Set aliases) { JsonStringArrayList jsonList = new JsonStringArrayList(aliases.size()); jsonList.addAll(aliases); return jsonList.toString(); } } AvroToArrowVectorIterator.java000066400000000000000000000131311515451615200355570ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.io.EOFException; 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; import org.apache.avro.Schema; import org.apache.avro.io.Decoder; /** VectorSchemaRoot iterator for partially converting avro data. */ public class AvroToArrowVectorIterator implements Iterator, AutoCloseable { public static final int NO_LIMIT_BATCH_SIZE = -1; public static final int DEFAULT_BATCH_SIZE = 1024; private final Decoder decoder; private final Schema schema; private final AvroToArrowConfig config; private CompositeAvroConsumer compositeConsumer; private org.apache.arrow.vector.types.pojo.Schema rootSchema; private VectorSchemaRoot nextBatch; private final int targetBatchSize; /** Construct an instance. */ private AvroToArrowVectorIterator(Decoder decoder, Schema schema, AvroToArrowConfig config) { this.decoder = decoder; this.schema = schema; this.config = config; this.targetBatchSize = config.getTargetBatchSize(); } /** Create a ArrowVectorIterator to partially convert data. */ public static AvroToArrowVectorIterator create( Decoder decoder, Schema schema, AvroToArrowConfig config) { AvroToArrowVectorIterator iterator = new AvroToArrowVectorIterator(decoder, schema, config); try { iterator.initialize(); return iterator; } catch (Exception e) { iterator.close(); throw new RuntimeException("Error occurs while creating iterator.", e); } } private void initialize() { // create consumers compositeConsumer = AvroToArrowUtils.createCompositeConsumer(schema, config); 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(); load(root); } private void consumeData(VectorSchemaRoot root) { int readRowCount = 0; try { while ((targetBatchSize == NO_LIMIT_BATCH_SIZE || readRowCount < targetBatchSize)) { compositeConsumer.consume(decoder); readRowCount++; } if (targetBatchSize == NO_LIMIT_BATCH_SIZE) { while (true) { ValueVectorUtility.ensureCapacity(root, readRowCount + 1); compositeConsumer.consume(decoder); readRowCount++; } } else { while (readRowCount < targetBatchSize) { compositeConsumer.consume(decoder); readRowCount++; } } root.setRowCount(readRowCount); } catch (EOFException eof) { // reach the end of encoder stream. root.setRowCount(readRowCount); } catch (Exception e) { compositeConsumer.close(); throw new RuntimeException("Error occurs while consuming data.", e); } } // Loads the next schema root or null if no more rows are available. private void load(VectorSchemaRoot root) { final int targetBatchSize = config.getTargetBatchSize(); if (targetBatchSize != NO_LIMIT_BATCH_SIZE) { ValueVectorUtility.preAllocate(root, targetBatchSize); } long validConsumerCount = compositeConsumer.getConsumers().stream().filter(c -> !c.skippable()).count(); Preconditions.checkArgument( root.getFieldVectors().size() == validConsumerCount, "Schema root vectors size not equals to consumers size."); compositeConsumer.resetConsumerVectors(root); // consume data consumeData(root); if (root.getRowCount() == 0) { root.close(); nextBatch = null; } else { nextBatch = root; } } @Override public boolean hasNext() { return nextBatch != null; } /** Gets the next vector. The user is responsible for freeing its resources. */ @Override public VectorSchemaRoot next() { Preconditions.checkArgument(hasNext()); VectorSchemaRoot returned = nextBatch; try { load(VectorSchemaRoot.create(rootSchema, config.getAllocator())); } catch (Exception e) { returned.close(); throw new RuntimeException("Error occurs while getting next schema root.", e); } return returned; } /** Clean up resources. */ @Override public void close() { if (nextBatch != null) { nextBatch.close(); } compositeConsumer.close(); } } apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/000077500000000000000000000000001515451615200316705ustar00rootroot00000000000000AvroArraysConsumer.java000066400000000000000000000043541515451615200362670ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex.ListVector; import org.apache.avro.io.Decoder; /** * Consumer which consume array type values from avro decoder. Write the data to {@link ListVector}. */ public class AvroArraysConsumer extends BaseAvroConsumer { private final Consumer delegate; /** Instantiate a ArrayConsumer. */ public AvroArraysConsumer(ListVector vector, Consumer delegate) { super(vector); this.delegate = delegate; } @Override public void consume(Decoder decoder) throws IOException { vector.startNewValue(currentIndex); long totalCount = 0; for (long count = decoder.readArrayStart(); count != 0; count = decoder.arrayNext()) { totalCount += count; ensureInnerVectorCapacity(totalCount); for (int element = 0; element < count; element++) { delegate.consume(decoder); } } vector.endValue(currentIndex, (int) totalCount); currentIndex++; } @Override public void close() throws Exception { super.close(); delegate.close(); } @Override public boolean resetValueVector(ListVector vector) { this.delegate.resetValueVector(vector.getDataVector()); return super.resetValueVector(vector); } void ensureInnerVectorCapacity(long targetCapacity) { while (vector.getDataVector().getValueCapacity() < targetCapacity) { vector.getDataVector().reAlloc(); } } } AvroBooleanConsumer.java000066400000000000000000000026121515451615200364000ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.BitVector; import org.apache.avro.io.Decoder; /** * Consumer which consume boolean type values from avro decoder. Write the data to {@link * BitVector}. */ public class AvroBooleanConsumer extends BaseAvroConsumer { /** Instantiate a AvroBooleanConsumer. */ public AvroBooleanConsumer(BitVector vector) { super(vector); } @Override public void consume(Decoder decoder) throws IOException { vector.set(currentIndex, decoder.readBoolean() ? 1 : 0); currentIndex++; } } AvroBytesConsumer.java000066400000000000000000000032651515451615200361140ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.nio.ByteBuffer; import org.apache.arrow.vector.VarBinaryVector; import org.apache.avro.io.Decoder; /** * Consumer which consume bytes type values from avro decoder. Write the data to {@link * VarBinaryVector}. */ public class AvroBytesConsumer extends BaseAvroConsumer { private ByteBuffer cacheBuffer; /** Instantiate a AvroBytesConsumer. */ public AvroBytesConsumer(VarBinaryVector vector) { super(vector); } @Override public void consume(Decoder decoder) throws IOException { // cacheBuffer is initialized null and create in the first consume, // if its capacity < size to read, decoder will create a new one with new capacity. cacheBuffer = decoder.readBytes(cacheBuffer); vector.setSafe(currentIndex, cacheBuffer, 0, cacheBuffer.limit()); currentIndex++; } } AvroDoubleConsumer.java000066400000000000000000000025671515451615200362440ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.Float8Vector; import org.apache.avro.io.Decoder; /** * Consumer which consume double type values from avro decoder. Write the data to {@link * Float8Vector}. */ public class AvroDoubleConsumer extends BaseAvroConsumer { /** Instantiate a AvroDoubleConsumer. */ public AvroDoubleConsumer(Float8Vector vector) { super(vector); } @Override public void consume(Decoder decoder) throws IOException { vector.set(currentIndex++, decoder.readDouble()); } } AvroEnumConsumer.java000066400000000000000000000026501515451615200357270ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.BaseIntVector; import org.apache.arrow.vector.IntVector; import org.apache.avro.io.Decoder; /** * Consumer which consume enum type values from avro decoder. Write the data to {@link IntVector}. */ public class AvroEnumConsumer extends BaseAvroConsumer { /** Instantiate a AvroEnumConsumer. */ public AvroEnumConsumer(BaseIntVector vector) { super(vector); } @Override public void consume(Decoder decoder) throws IOException { vector.setWithPossibleTruncate(currentIndex++, decoder.readEnum()); } } AvroFixedConsumer.java000066400000000000000000000030331515451615200360560ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.FixedSizeBinaryVector; import org.apache.avro.io.Decoder; /** * Consumer which consume fixed type values from avro decoder. Write the data to {@link * org.apache.arrow.vector.FixedSizeBinaryVector}. */ public class AvroFixedConsumer extends BaseAvroConsumer { private final byte[] reuseBytes; /** Instantiate a AvroFixedConsumer. */ public AvroFixedConsumer(FixedSizeBinaryVector vector, int size) { super(vector); reuseBytes = new byte[size]; } @Override public void consume(Decoder decoder) throws IOException { decoder.readFixed(reuseBytes); vector.setSafe(currentIndex++, reuseBytes); } } AvroFloatConsumer.java000066400000000000000000000025621515451615200360720ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.Float4Vector; import org.apache.avro.io.Decoder; /** * Consumer which consume float type values from avro decoder. Write the data to {@link * Float4Vector}. */ public class AvroFloatConsumer extends BaseAvroConsumer { /** Instantiate a AvroFloatConsumer. */ public AvroFloatConsumer(Float4Vector vector) { super(vector); } @Override public void consume(Decoder decoder) throws IOException { vector.set(currentIndex++, decoder.readFloat()); } } AvroIntConsumer.java000066400000000000000000000025311515451615200355530ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.IntVector; import org.apache.avro.io.Decoder; /** * Consumer which consume int type values from avro decoder. Write the data to {@link IntVector}. */ public class AvroIntConsumer extends BaseAvroConsumer { /** Instantiate a AvroIntConsumer. */ public AvroIntConsumer(IntVector vector) { super(vector); } @Override public void consume(Decoder decoder) throws IOException { vector.set(currentIndex++, decoder.readInt()); } } AvroLongConsumer.java000066400000000000000000000025551515451615200357260ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.BigIntVector; import org.apache.avro.io.Decoder; /** * Consumer which consume long type values from avro decoder. Write the data to {@link * BigIntVector}. */ public class AvroLongConsumer extends BaseAvroConsumer { /** Instantiate a AvroLongConsumer. */ public AvroLongConsumer(BigIntVector vector) { super(vector); } @Override public void consume(Decoder decoder) throws IOException { vector.set(currentIndex++, decoder.readLong()); } } AvroMapConsumer.java000066400000000000000000000046461515451615200355470ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.avro.io.Decoder; /** * Consumer which consume map type values from avro decoder. Write the data to {@link MapVector}. */ public class AvroMapConsumer extends BaseAvroConsumer { private final Consumer delegate; /** Instantiate a AvroMapConsumer. */ public AvroMapConsumer(MapVector vector, Consumer delegate) { super(vector); this.delegate = delegate; } @Override public void consume(Decoder decoder) throws IOException { vector.startNewValue(currentIndex); long totalCount = 0; for (long count = decoder.readMapStart(); count != 0; count = decoder.mapNext()) { totalCount += count; ensureInnerVectorCapacity(totalCount); for (int element = 0; element < count; element++) { delegate.consume(decoder); } } vector.endValue(currentIndex, (int) totalCount); currentIndex++; } @Override public void close() throws Exception { super.close(); delegate.close(); } @Override public boolean resetValueVector(MapVector vector) { this.delegate.resetValueVector(vector.getDataVector()); return super.resetValueVector(vector); } void ensureInnerVectorCapacity(long targetCapacity) { StructVector innerVector = (StructVector) vector.getDataVector(); for (FieldVector v : innerVector.getChildrenFromFields()) { while (v.getValueCapacity() < targetCapacity) { v.reAlloc(); } } } } AvroNullConsumer.java000066400000000000000000000024631515451615200357370ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullVector; import org.apache.avro.io.Decoder; /** * Consumer which consume null type values from avro decoder. Corresponding to {@link * org.apache.arrow.vector.NullVector}. */ public class AvroNullConsumer extends BaseAvroConsumer { public AvroNullConsumer(NullVector vector) { super(vector); } @Override public void consume(Decoder decoder) throws IOException { currentIndex++; } } AvroNullableConsumer.java000066400000000000000000000047511515451615200365650ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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(); } } AvroStringConsumer.java000066400000000000000000000032371515451615200362730ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.nio.ByteBuffer; import org.apache.arrow.vector.VarCharVector; import org.apache.avro.io.Decoder; /** * Consumer which consume string type values from avro decoder. Write the data to {@link * VarCharVector}. */ public class AvroStringConsumer extends BaseAvroConsumer { private ByteBuffer cacheBuffer; /** Instantiate a AvroStringConsumer. */ public AvroStringConsumer(VarCharVector vector) { super(vector); } @Override public void consume(Decoder decoder) throws IOException { // cacheBuffer is initialized null and create in the first consume, // if its capacity < size to read, decoder will create a new one with new capacity. cacheBuffer = decoder.readBytes(cacheBuffer); vector.setSafe(currentIndex++, cacheBuffer, 0, cacheBuffer.limit()); } } AvroStructConsumer.java000066400000000000000000000044551515451615200363140ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.AutoCloseables; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.avro.io.Decoder; /** * Consumer which consume nested record type values from avro decoder. Write the data to {@link * org.apache.arrow.vector.complex.StructVector}. */ public class AvroStructConsumer extends BaseAvroConsumer { private final Consumer[] delegates; /** Instantiate a AvroStructConsumer. */ public AvroStructConsumer(StructVector vector, Consumer[] delegates) { super(vector); this.delegates = delegates; } @Override public void consume(Decoder decoder) throws IOException { ensureInnerVectorCapacity(currentIndex + 1); for (int i = 0; i < delegates.length; i++) { delegates[i].consume(decoder); } vector.setIndexDefined(currentIndex); currentIndex++; } @Override public void close() throws Exception { super.close(); AutoCloseables.close(delegates); } @Override public boolean resetValueVector(StructVector vector) { for (int i = 0; i < delegates.length; i++) { delegates[i].resetValueVector(vector.getChildrenFromFields().get(i)); } return super.resetValueVector(vector); } void ensureInnerVectorCapacity(long targetCapacity) { for (FieldVector v : vector.getChildrenFromFields()) { while (v.getValueCapacity() < targetCapacity) { v.reAlloc(); } } } } AvroUnionsConsumer.java000066400000000000000000000054701515451615200363010ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.AutoCloseables; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.types.Types; import org.apache.avro.io.Decoder; /** * Consumer which consume unions type values from avro decoder. Write the data to {@link * org.apache.arrow.vector.complex.UnionVector}. */ public class AvroUnionsConsumer extends BaseAvroConsumer { private Consumer[] delegates; private Types.MinorType[] types; /** Instantiate an AvroUnionConsumer. */ public AvroUnionsConsumer(UnionVector vector, Consumer[] delegates, Types.MinorType[] types) { super(vector); this.delegates = delegates; this.types = types; } @Override public void consume(Decoder decoder) throws IOException { int fieldIndex = decoder.readInt(); ensureInnerVectorCapacity(currentIndex + 1, fieldIndex); Consumer delegate = delegates[fieldIndex]; vector.setType(currentIndex, types[fieldIndex]); // In UnionVector we need to set sub vector writer position before consume a value // because in the previous iterations we might not have written to the specific union sub // vector. delegate.setPosition(currentIndex); delegate.consume(decoder); currentIndex++; } @Override public void close() throws Exception { super.close(); AutoCloseables.close(delegates); } @Override public boolean resetValueVector(UnionVector vector) { for (int i = 0; i < delegates.length; i++) { delegates[i].resetValueVector(vector.getChildrenFromFields().get(i)); } return super.resetValueVector(vector); } void ensureInnerVectorCapacity(long targetCapacity, int fieldIndex) { ValueVector fieldVector = vector.getChildrenFromFields().get(fieldIndex); if (fieldVector.getMinorType() == Types.MinorType.NULL) { return; } while (fieldVector.getValueCapacity() < targetCapacity) { fieldVector.reAlloc(); } } } BaseAvroConsumer.java000066400000000000000000000032651515451615200357000ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.FieldVector; /** * Base class for non-skippable avro consumers. * * @param vector type. */ public abstract class BaseAvroConsumer implements Consumer { protected T vector; protected int currentIndex; /** * Constructs a base avro consumer. * * @param vector the vector to consume. */ public BaseAvroConsumer(T vector) { this.vector = vector; } @Override public void addNull() { currentIndex++; } @Override public void setPosition(int index) { currentIndex = index; } @Override public FieldVector getVector() { return vector; } @Override public void close() throws Exception { vector.close(); } @Override public boolean resetValueVector(T vector) { this.vector = vector; this.currentIndex = 0; return true; } } CompositeAvroConsumer.java000066400000000000000000000040461515451615200367660ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.List; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.avro.io.Decoder; /** Composite consumer which hold all consumers. It manages the consume and cleanup process. */ public class CompositeAvroConsumer implements AutoCloseable { private final List consumers; public List getConsumers() { return consumers; } public CompositeAvroConsumer(List consumers) { this.consumers = consumers; } /** Consume decoder data. */ public void consume(Decoder decoder) throws IOException { for (Consumer consumer : consumers) { consumer.consume(decoder); } } /** Reset vector of consumers with the given {@link VectorSchemaRoot}. */ public void resetConsumerVectors(VectorSchemaRoot root) { int index = 0; for (Consumer consumer : consumers) { if (consumer.resetValueVector(root.getFieldVectors().get(index))) { index++; } } } @Override public void close() { // clean up try { AutoCloseables.close(consumers); } catch (Exception e) { throw new RuntimeException("Error occurs in close.", e); } } } Consumer.java000066400000000000000000000041021515451615200342440ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * Interface that is used to consume values from avro decoder. * * @param The vector within consumer or its delegate, used for partially consume purpose. */ public interface Consumer extends AutoCloseable { /** * Consume a specific type value from avro decoder and write it to vector. * * @param decoder avro decoder to read data * @throws IOException on error */ void consume(Decoder decoder) throws IOException; /** Add null value to vector by making writer position + 1. */ void addNull(); /** Set the position to write value into vector. */ void setPosition(int index); /** Get the vector within the consumer. */ FieldVector getVector(); /** Close this consumer when occurs exception to avoid potential leak. */ @Override void close() throws Exception; /** * Reset the vector within consumer for partial read purpose. * * @return true if reset is successful, false if reset is not needed. */ boolean resetValueVector(T vector); /** Indicates whether the consumer is type of {@link SkipConsumer}. */ default boolean skippable() { return false; } } SkipConsumer.java000066400000000000000000000032351515451615200351010ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 which skip (throw away) data from the decoder. */ public class SkipConsumer implements Consumer { private final SkipFunction skipFunction; public SkipConsumer(SkipFunction skipFunction) { this.skipFunction = skipFunction; } @Override public void consume(Decoder decoder) throws IOException { skipFunction.apply(decoder); } @Override public void addNull() {} @Override public void setPosition(int index) {} @Override public FieldVector getVector() { return null; } @Override public void close() throws Exception {} @Override public boolean resetValueVector(FieldVector vector) { return false; } @Override public boolean skippable() { return true; } } SkipFunction.java000066400000000000000000000021001515451615200350610ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.avro.io.Decoder; /** Adapter function to skip (throw away) data from the decoder. */ @FunctionalInterface public interface SkipFunction { void apply(Decoder decoder) throws IOException; } apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/000077500000000000000000000000001515451615200333025ustar00rootroot00000000000000AvroDateConsumer.java000066400000000000000000000026711515451615200373150ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.DateDayVector; import org.apache.avro.io.Decoder; /** * Consumer which consume date type values from avro decoder. Write the data to {@link * DateDayVector}. */ public class AvroDateConsumer extends BaseAvroConsumer { /** Instantiate a AvroDateConsumer. */ public AvroDateConsumer(DateDayVector vector) { super(vector); } @Override public void consume(Decoder decoder) throws IOException { vector.set(currentIndex++, decoder.readInt()); } } AvroDecimal256Consumer.java000066400000000000000000000053311515451615200402270ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } } AvroDecimalConsumer.java000066400000000000000000000052401515451615200377710ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.DecimalVector; import org.apache.avro.io.Decoder; /** * Consumer which consume decimal type values from avro decoder. Write the data to {@link * DecimalVector}. */ public abstract class AvroDecimalConsumer extends BaseAvroConsumer { /** Instantiate a AvroDecimalConsumer. */ public AvroDecimalConsumer(DecimalVector vector) { super(vector); } /** Consumer for decimal logical type with original bytes type. */ public static class BytesDecimalConsumer extends AvroDecimalConsumer { private ByteBuffer cacheBuffer; /** Instantiate a BytesDecimalConsumer. */ public BytesDecimalConsumer(DecimalVector 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 <= 16, "Decimal bytes length should <= 16."); cacheBuffer.get(bytes); vector.setBigEndian(currentIndex++, bytes); } } /** Consumer for decimal logical type with original fixed type. */ public static class FixedDecimalConsumer extends AvroDecimalConsumer { private byte[] reuseBytes; /** Instantiate a FixedDecimalConsumer. */ public FixedDecimalConsumer(DecimalVector vector, int size) { super(vector); Preconditions.checkArgument(size <= 16, "Decimal bytes length should <= 16."); reuseBytes = new byte[size]; } @Override public void consume(Decoder decoder) throws IOException { decoder.readFixed(reuseBytes); vector.setBigEndian(currentIndex++, reuseBytes); } } } AvroTimeMicroConsumer.java000066400000000000000000000027271515451615200403320ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.TimeMicroVector; import org.apache.avro.io.Decoder; /** * Consumer which consume date time-micro values from avro decoder. Write the data to {@link * TimeMicroVector}. */ public class AvroTimeMicroConsumer extends BaseAvroConsumer { /** Instantiate a AvroTimeMicroConsumer. */ public AvroTimeMicroConsumer(TimeMicroVector vector) { super(vector); } @Override public void consume(Decoder decoder) throws IOException { vector.set(currentIndex++, decoder.readLong()); } } AvroTimeMillisConsumer.java000066400000000000000000000027311515451615200405050ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.TimeMilliVector; import org.apache.avro.io.Decoder; /** * Consumer which consume date time-millis values from avro decoder. Write the data to {@link * TimeMilliVector}. */ public class AvroTimeMillisConsumer extends BaseAvroConsumer { /** Instantiate a AvroTimeMilliConsumer. */ public AvroTimeMillisConsumer(TimeMilliVector vector) { super(vector); } @Override public void consume(Decoder decoder) throws IOException { vector.set(currentIndex++, decoder.readInt()); } } AvroTimestampMicrosConsumer.java000066400000000000000000000030041515451615200415470ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.TimeStampMicroVector; import org.apache.avro.io.Decoder; /** * Consumer which consumes local-timestamp-micros values from avro decoder. Write the data to {@link * TimeStampMicroVector}. */ public class AvroTimestampMicrosConsumer extends BaseAvroConsumer { /** Instantiate a AvroTimestampMicroConsumer. */ public AvroTimestampMicrosConsumer(TimeStampMicroVector vector) { super(vector); } @Override public void consume(Decoder decoder) throws IOException { vector.set(currentIndex++, decoder.readLong()); } } AvroTimestampMicrosTzConsumer.java000066400000000000000000000030151515451615200420670ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } AvroTimestampMillisConsumer.java000066400000000000000000000030041515451615200415440ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.TimeStampMilliVector; import org.apache.avro.io.Decoder; /** * Consumer which consume local-timestamp-millis values from avro decoder. Write the data to {@link * TimeStampMilliVector}. */ public class AvroTimestampMillisConsumer extends BaseAvroConsumer { /** Instantiate a AvroTimestampMillisConsumer. */ public AvroTimestampMillisConsumer(TimeStampMilliVector vector) { super(vector); } @Override public void consume(Decoder decoder) throws IOException { vector.set(currentIndex++, decoder.readLong()); } } AvroTimestampMillisTzConsumer.java000066400000000000000000000030141515451615200420630ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } AvroTimestampNanosConsumer.java000066400000000000000000000027741515451615200414060ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } AvroTimestampNanosTzConsumer.java000066400000000000000000000030021515451615200417050ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/000077500000000000000000000000001515451615200316605ustar00rootroot00000000000000AvroBigIntProducer.java000066400000000000000000000040411515451615200361530ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroBooleanProducer.java000066400000000000000000000026141515451615200363620ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } AvroBytesProducer.java000066400000000000000000000033361515451615200360730ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroEnumProducer.java000066400000000000000000000026221515451615200357060ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++)); } } AvroFixedSizeBinaryProducer.java000066400000000000000000000041371515451615200400440ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroFixedSizeListProducer.java000066400000000000000000000051131515451615200375260ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } AvroFloat2Producer.java000066400000000000000000000030241515451615200361260ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroFloat4Producer.java000066400000000000000000000027201515451615200361320ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroFloat8Producer.java000066400000000000000000000027241515451615200361420ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroIntProducer.java000066400000000000000000000040001515451615200355240ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroListProducer.java000066400000000000000000000050141515451615200357130ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } AvroMapProducer.java000066400000000000000000000050001515451615200355100ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } AvroNullProducer.java000066400000000000000000000025251515451615200357160ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroNullableProducer.java000066400000000000000000000044441515451615200365440ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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(); } } AvroSmallIntProducer.java000066400000000000000000000027351515451615200365320ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroStringProducer.java000066400000000000000000000032771515451615200362570ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroStructProducer.java000066400000000000000000000046111515451615200362660ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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)); } } } AvroTinyIntProducer.java000066400000000000000000000027231515451615200364020ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroUint1Producer.java000066400000000000000000000027731515451615200360110ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroUint2Producer.java000066400000000000000000000027761515451615200360150ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroUint4Producer.java000066400000000000000000000030061515451615200360020ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroUint8Producer.java000066400000000000000000000031031515451615200360040ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } BaseAvroProducer.java000066400000000000000000000042471515451615200356610ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; } } CompositeAvroProducer.java000066400000000000000000000037531515451615200367520ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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)); } } } DictionaryDecodingProducer.java000066400000000000000000000033641515451615200377200ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } Producer.java000066400000000000000000000033001515451615200342230ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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(); } apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/000077500000000000000000000000001515451615200332725ustar00rootroot00000000000000AvroDateDayProducer.java000066400000000000000000000024561515451615200377340ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } AvroDateMilliProducer.java000066400000000000000000000036321515451615200402620ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroDecimal256Producer.java000066400000000000000000000033211515451615200402040ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } AvroDecimalProducer.java000066400000000000000000000046661515451615200377640ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } AvroTimeMicroProducer.java000066400000000000000000000025321515451615200403040ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } AvroTimeMilliProducer.java000066400000000000000000000025271515451615200403050ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } AvroTimeNanoProducer.java000066400000000000000000000034611515451615200401300ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroTimeSecProducer.java000066400000000000000000000041001515451615200377360ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroTimestampMicroProducer.java000066400000000000000000000026371515451615200413570ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } AvroTimestampMicroTzProducer.java000066400000000000000000000027711515451615200416740ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } AvroTimestampMilliProducer.java000066400000000000000000000026371515451615200413540ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } AvroTimestampMilliTzProducer.java000066400000000000000000000027711515451615200416710ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } AvroTimestampNanoProducer.java000066400000000000000000000026261515451615200411770ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } AvroTimestampNanoTzProducer.java000066400000000000000000000027621515451615200415160ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } AvroTimestampSecProducer.java000066400000000000000000000041261515451615200410130ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } AvroTimestampSecTzProducer.java000066400000000000000000000043611515451615200413320ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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++; } } apache-arrow-java-19.0.0/adapter/avro/src/test/000077500000000000000000000000001515451615200212335ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/java/000077500000000000000000000000001515451615200221545ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/java/org/000077500000000000000000000000001515451615200227435ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/java/org/apache/000077500000000000000000000000001515451615200241645ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200253165ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/java/org/apache/arrow/adapter/000077500000000000000000000000001515451615200267365ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/000077500000000000000000000000001515451615200277055ustar00rootroot00000000000000ArrowToAvroDataTest.java000066400000000000000000003505521515451615200343620ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } } } } ArrowToAvroSchemaTest.java000066400000000000000000002224031515451615200347020ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } AvroLogicalTypesTest.java000066400000000000000000000164351515451615200345710ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.math.BigDecimal; import java.nio.ByteBuffer; import java.time.LocalDateTime; import java.util.ArrayList; import java.util.Arrays; import java.util.List; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.util.DateUtility; import org.apache.avro.Conversions; import org.apache.avro.Schema; import org.apache.avro.generic.GenericFixed; import org.junit.jupiter.api.Test; public class AvroLogicalTypesTest extends AvroTestBase { @Test public void testTimestampMicros() throws Exception { Schema schema = getSchema("logical/test_timestamp_micros.avsc"); List data = Arrays.asList(10000L, 20000L, 30000L, 40000L, 50000L); List expected = Arrays.asList( DateUtility.getLocalDateTimeFromEpochMicro(10000), DateUtility.getLocalDateTimeFromEpochMicro(20000), DateUtility.getLocalDateTimeFromEpochMicro(30000), DateUtility.getLocalDateTimeFromEpochMicro(40000), DateUtility.getLocalDateTimeFromEpochMicro(50000)); VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(expected, vector); } @Test public void testTimestampMillis() throws Exception { Schema schema = getSchema("logical/test_timestamp_millis.avsc"); List data = Arrays.asList(10000L, 20000L, 30000L, 40000L, 50000L); List expected = Arrays.asList( DateUtility.getLocalDateTimeFromEpochMilli(10000), DateUtility.getLocalDateTimeFromEpochMilli(20000), DateUtility.getLocalDateTimeFromEpochMilli(30000), DateUtility.getLocalDateTimeFromEpochMilli(40000), DateUtility.getLocalDateTimeFromEpochMilli(50000)); VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(expected, vector); } @Test public void testTimeMicros() throws Exception { Schema schema = getSchema("logical/test_time_micros.avsc"); List data = Arrays.asList(10000L, 20000L, 30000L, 40000L, 50000L); VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(data, vector); } @Test public void testTimeMillis() throws Exception { Schema schema = getSchema("logical/test_time_millis.avsc"); List data = Arrays.asList(100, 200, 300, 400, 500); List expected = Arrays.asList( DateUtility.getLocalDateTimeFromEpochMilli(100), DateUtility.getLocalDateTimeFromEpochMilli(200), DateUtility.getLocalDateTimeFromEpochMilli(300), DateUtility.getLocalDateTimeFromEpochMilli(400), DateUtility.getLocalDateTimeFromEpochMilli(500)); VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(expected, vector); } @Test public void testDate() throws Exception { Schema schema = getSchema("logical/test_date.avsc"); List data = Arrays.asList(100, 200, 300, 400, 500); VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(data, vector); } @Test public void testDecimalWithOriginalBytes() throws Exception { Schema schema = getSchema("logical/test_decimal_with_original_bytes.avsc"); List data = new ArrayList<>(); List expected = new ArrayList<>(); Conversions.DecimalConversion conversion = new Conversions.DecimalConversion(); for (int i = 0; i < 5; i++) { BigDecimal value = new BigDecimal(i * i).setScale(2); ByteBuffer buffer = conversion.toBytes(value, schema, schema.getLogicalType()); data.add(buffer); expected.add(value); } VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(expected, vector); } @Test public void testDecimalWithOriginalFixed() throws Exception { Schema schema = getSchema("logical/test_decimal_with_original_fixed.avsc"); List data = new ArrayList<>(); List expected = new ArrayList<>(); Conversions.DecimalConversion conversion = new Conversions.DecimalConversion(); for (int i = 0; i < 5; i++) { BigDecimal value = new BigDecimal(i * i).setScale(2); GenericFixed fixed = conversion.toFixed(value, schema, schema.getLogicalType()); data.add(fixed); expected.add(value); } VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(expected, vector); } @Test public void testInvalidDecimalPrecision() throws Exception { Schema schema = getSchema("logical/test_decimal_invalid1.avsc"); List data = new ArrayList<>(); Conversions.DecimalConversion conversion = new Conversions.DecimalConversion(); for (int i = 0; i < 5; i++) { BigDecimal value = new BigDecimal(i * i).setScale(2); ByteBuffer buffer = conversion.toBytes(value, schema, schema.getLogicalType()); data.add(buffer); } IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () -> writeAndRead(schema, data)); assertTrue(e.getMessage().contains("Precision must be in range of 1 to 76")); } @Test public void testFailedToCreateDecimalLogicalType() throws Exception { // For decimal logical type, if avro validate schema failed, it will not create logical type, // and the schema will be treated as its original type. // java.lang.IllegalArgumentException: Invalid decimal scale: -1 (must be positive) Schema schema1 = getSchema("logical/test_decimal_invalid2.avsc"); assertNull(schema1.getLogicalType()); // java.lang.IllegalArgumentException: Invalid decimal scale: 40 (greater than precision: 20) Schema schema2 = getSchema("logical/test_decimal_invalid3.avsc"); assertNull(schema2.getLogicalType()); // java.lang.IllegalArgumentException: fixed(1) cannot store 30 digits (max 2) Schema schema3 = getSchema("logical/test_decimal_invalid4.avsc"); assertNull(schema3.getLogicalType()); } } AvroSkipFieldTest.java000066400000000000000000000626171515451615200340470ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; import java.util.HashSet; import java.util.Set; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.types.Types; import org.apache.avro.Schema; import org.apache.avro.generic.GenericData; import org.apache.avro.generic.GenericRecord; import org.junit.jupiter.api.Test; public class AvroSkipFieldTest extends AvroTestBase { @Test public void testSkipUnionWithOneField() throws Exception { Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f0"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); Schema schema = getSchema("skip/test_skip_union_before.avsc"); Schema expectedSchema = getSchema("skip/test_skip_union_one_field_expected.avsc"); ArrayList data = new ArrayList<>(); ArrayList expectedData = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, "test" + i); record.put(1, i % 2 == 0 ? "test" + i : null); record.put(2, i % 2 == 0 ? "test" + i : i); record.put(3, i); data.add(record); GenericRecord expectedRecord = new GenericData.Record(expectedSchema); expectedRecord.put(0, record.get(1)); expectedRecord.put(1, record.get(2)); expectedRecord.put(2, record.get(3)); expectedData.add(expectedRecord); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(expectedSchema, expectedData, root); } @Test public void testSkipUnionWithNullableOneField() throws Exception { Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f1"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); Schema schema = getSchema("skip/test_skip_union_before.avsc"); Schema expectedSchema = getSchema("skip/test_skip_union_nullable_field_expected.avsc"); ArrayList data = new ArrayList<>(); ArrayList expectedData = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, "test" + i); record.put(1, i % 2 == 0 ? "test" + i : null); record.put(2, i % 2 == 0 ? "test" + i : i); record.put(3, i); data.add(record); GenericRecord expectedRecord = new GenericData.Record(expectedSchema); expectedRecord.put(0, record.get(0)); expectedRecord.put(1, record.get(2)); expectedRecord.put(2, record.get(3)); expectedData.add(expectedRecord); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(expectedSchema, expectedData, root); } @Test public void testSkipUnionWithMultiFields() throws Exception { Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f2"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); Schema schema = getSchema("skip/test_skip_union_before.avsc"); Schema expectedSchema = getSchema("skip/test_skip_union_multi_fields_expected.avsc"); ArrayList data = new ArrayList<>(); ArrayList expectedData = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, "test" + i); record.put(1, i % 2 == 0 ? "test" + i : null); record.put(2, i % 2 == 0 ? "test" + i : i); record.put(3, i); data.add(record); GenericRecord expectedRecord = new GenericData.Record(expectedSchema); expectedRecord.put(0, record.get(0)); expectedRecord.put(1, record.get(1)); expectedRecord.put(2, record.get(3)); expectedData.add(expectedRecord); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(expectedSchema, expectedData, root); } @Test public void testSkipMapField() throws Exception { Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f1"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); Schema schema = getSchema("skip/test_skip_map_before.avsc"); Schema expectedSchema = getSchema("skip/test_skip_map_expected.avsc"); HashMap map = new HashMap(); map.put("key1", "value1"); map.put("key2", "value3"); ArrayList data = new ArrayList<>(); ArrayList expectedData = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, "test" + i); record.put(1, map); record.put(2, i % 2 == 0); data.add(record); GenericRecord expectedRecord = new GenericData.Record(expectedSchema); expectedRecord.put(0, record.get(0)); expectedRecord.put(1, record.get(2)); expectedData.add(expectedRecord); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(expectedSchema, expectedData, root); } @Test public void testSkipArrayField() throws Exception { Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f1"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); Schema schema = getSchema("skip/test_skip_array_before.avsc"); Schema expectedSchema = getSchema("skip/test_skip_array_expected.avsc"); ArrayList data = new ArrayList<>(); ArrayList expectedData = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, "test" + i); record.put(1, Arrays.asList("test" + i, "test" + i)); record.put(2, i % 2 == 0); data.add(record); GenericRecord expectedRecord = new GenericData.Record(expectedSchema); expectedRecord.put(0, record.get(0)); expectedRecord.put(1, record.get(2)); expectedData.add(expectedRecord); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(expectedSchema, expectedData, root); } @Test public void testSkipMultiFields() throws Exception { Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f1"); skipFieldNames.add("f2"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); Schema schema = getSchema("test_record.avsc"); Schema expectedSchema = getSchema("skip/test_skip_multi_fields_expected.avsc"); ArrayList data = new ArrayList<>(); ArrayList expectedData = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, "test" + i); record.put(1, i); record.put(2, i % 2 == 0); data.add(record); GenericRecord expectedRecord = new GenericData.Record(expectedSchema); expectedRecord.put(0, record.get(0)); expectedData.add(expectedRecord); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(expectedSchema, expectedData, root); } @Test public void testSkipStringField() throws Exception { Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f2"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); Schema schema = getSchema("skip/test_skip_base1.avsc"); Schema expectedSchema = getSchema("skip/test_skip_string_expected.avsc"); ArrayList data = new ArrayList<>(); ArrayList expectedData = new ArrayList<>(); for (int i = 0; i < 5; i++) { final byte[] testBytes = ("test" + i).getBytes(StandardCharsets.UTF_8); GenericRecord record = new GenericData.Record(schema); GenericData.Fixed fixed = new GenericData.Fixed(schema.getField("f0").schema()); fixed.bytes(testBytes); record.put(0, fixed); GenericData.EnumSymbol symbol = new GenericData.EnumSymbol(schema.getField("f1").schema(), "TEST" + i % 2); record.put(1, symbol); record.put(2, "testtest" + i); record.put(3, ByteBuffer.wrap(testBytes)); data.add(record); GenericRecord expectedRecord = new GenericData.Record(expectedSchema); expectedRecord.put(0, testBytes); expectedRecord.put(1, (byte) i % 2); expectedRecord.put(2, testBytes); expectedData.add(expectedRecord); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(expectedSchema, expectedData, root); } @Test public void testSkipBytesField() throws Exception { Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f3"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); Schema schema = getSchema("skip/test_skip_base1.avsc"); Schema expectedSchema = getSchema("skip/test_skip_bytes_expected.avsc"); ArrayList data = new ArrayList<>(); ArrayList expectedData = new ArrayList<>(); for (int i = 0; i < 5; i++) { final byte[] testBytes = ("test" + i).getBytes(StandardCharsets.UTF_8); GenericRecord record = new GenericData.Record(schema); GenericData.Fixed fixed = new GenericData.Fixed(schema.getField("f0").schema()); fixed.bytes(testBytes); record.put(0, fixed); GenericData.EnumSymbol symbol = new GenericData.EnumSymbol(schema.getField("f1").schema(), "TEST" + i % 2); record.put(1, symbol); record.put(2, "testtest" + i); record.put(3, ByteBuffer.wrap(testBytes)); data.add(record); GenericRecord expectedRecord = new GenericData.Record(expectedSchema); expectedRecord.put(0, testBytes); expectedRecord.put(1, (byte) i % 2); expectedRecord.put(2, record.get(2)); expectedData.add(expectedRecord); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(expectedSchema, expectedData, root); } @Test public void testSkipFixedField() throws Exception { Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f0"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); Schema schema = getSchema("skip/test_skip_base1.avsc"); Schema expectedSchema = getSchema("skip/test_skip_fixed_expected.avsc"); ArrayList data = new ArrayList<>(); ArrayList expectedData = new ArrayList<>(); for (int i = 0; i < 5; i++) { final byte[] testBytes = ("test" + i).getBytes(StandardCharsets.UTF_8); GenericRecord record = new GenericData.Record(schema); GenericData.Fixed fixed = new GenericData.Fixed(schema.getField("f0").schema()); fixed.bytes(testBytes); record.put(0, fixed); GenericData.EnumSymbol symbol = new GenericData.EnumSymbol(schema.getField("f1").schema(), "TEST" + i % 2); record.put(1, symbol); record.put(2, "testtest" + i); record.put(3, ByteBuffer.wrap(testBytes)); data.add(record); GenericRecord expectedRecord = new GenericData.Record(expectedSchema); expectedRecord.put(0, (byte) i % 2); expectedRecord.put(1, record.get(2)); expectedRecord.put(2, record.get(3)); expectedData.add(expectedRecord); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(expectedSchema, expectedData, root); } @Test public void testSkipEnumField() throws Exception { Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f1"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); Schema schema = getSchema("skip/test_skip_base1.avsc"); Schema expectedSchema = getSchema("skip/test_skip_fixed_expected.avsc"); ArrayList data = new ArrayList<>(); ArrayList expectedData = new ArrayList<>(); for (int i = 0; i < 5; i++) { final byte[] testBytes = ("test" + i).getBytes(StandardCharsets.UTF_8); GenericRecord record = new GenericData.Record(schema); GenericData.Fixed fixed = new GenericData.Fixed(schema.getField("f0").schema()); fixed.bytes(testBytes); record.put(0, fixed); GenericData.EnumSymbol symbol = new GenericData.EnumSymbol(schema.getField("f1").schema(), "TEST" + i % 2); record.put(1, symbol); record.put(2, "testtest" + i); record.put(3, ByteBuffer.wrap(testBytes)); data.add(record); GenericRecord expectedRecord = new GenericData.Record(expectedSchema); expectedRecord.put(0, testBytes); expectedRecord.put(1, record.get(2)); expectedRecord.put(2, record.get(3)); expectedData.add(expectedRecord); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(expectedSchema, expectedData, root); } @Test public void testSkipBooleanField() throws Exception { Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f0"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); Schema schema = getSchema("skip/test_skip_base2.avsc"); Schema expectedSchema = getSchema("skip/test_skip_boolean_expected.avsc"); ArrayList data = new ArrayList<>(); ArrayList expectedData = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, i % 2 == 0); record.put(1, i); record.put(2, (long) i); record.put(3, (float) i); record.put(4, (double) i); data.add(record); GenericRecord expectedRecord = new GenericData.Record(expectedSchema); expectedRecord.put(0, record.get(1)); expectedRecord.put(1, record.get(2)); expectedRecord.put(2, record.get(3)); expectedRecord.put(3, record.get(4)); expectedData.add(expectedRecord); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(expectedSchema, expectedData, root); } @Test public void testSkipIntField() throws Exception { Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f1"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); Schema schema = getSchema("skip/test_skip_base2.avsc"); Schema expectedSchema = getSchema("skip/test_skip_int_expected.avsc"); ArrayList data = new ArrayList<>(); ArrayList expectedData = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, i % 2 == 0); record.put(1, i); record.put(2, (long) i); record.put(3, (float) i); record.put(4, (double) i); data.add(record); GenericRecord expectedRecord = new GenericData.Record(expectedSchema); expectedRecord.put(0, record.get(0)); expectedRecord.put(1, record.get(2)); expectedRecord.put(2, record.get(3)); expectedRecord.put(3, record.get(4)); expectedData.add(expectedRecord); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(expectedSchema, expectedData, root); } @Test public void testSkipLongField() throws Exception { Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f2"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); Schema schema = getSchema("skip/test_skip_base2.avsc"); Schema expectedSchema = getSchema("skip/test_skip_long_expected.avsc"); ArrayList data = new ArrayList<>(); ArrayList expectedData = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, i % 2 == 0); record.put(1, i); record.put(2, (long) i); record.put(3, (float) i); record.put(4, (double) i); data.add(record); GenericRecord expectedRecord = new GenericData.Record(expectedSchema); expectedRecord.put(0, record.get(0)); expectedRecord.put(1, record.get(1)); expectedRecord.put(2, record.get(3)); expectedRecord.put(3, record.get(4)); expectedData.add(expectedRecord); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(expectedSchema, expectedData, root); } @Test public void testSkipFloatField() throws Exception { Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f3"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); Schema schema = getSchema("skip/test_skip_base2.avsc"); Schema expectedSchema = getSchema("skip/test_skip_float_expected.avsc"); ArrayList data = new ArrayList<>(); ArrayList expectedData = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, i % 2 == 0); record.put(1, i); record.put(2, (long) i); record.put(3, (float) i); record.put(4, (double) i); data.add(record); GenericRecord expectedRecord = new GenericData.Record(expectedSchema); expectedRecord.put(0, record.get(0)); expectedRecord.put(1, record.get(1)); expectedRecord.put(2, record.get(2)); expectedRecord.put(3, record.get(4)); expectedData.add(expectedRecord); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(expectedSchema, expectedData, root); } @Test public void testSkipDoubleField() throws Exception { Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f4"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); Schema schema = getSchema("skip/test_skip_base2.avsc"); Schema expectedSchema = getSchema("skip/test_skip_double_expected.avsc"); ArrayList data = new ArrayList<>(); ArrayList expectedData = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, i % 2 == 0); record.put(1, i); record.put(2, (long) i); record.put(3, (float) i); record.put(4, (double) i); data.add(record); GenericRecord expectedRecord = new GenericData.Record(expectedSchema); expectedRecord.put(0, record.get(0)); expectedRecord.put(1, record.get(1)); expectedRecord.put(2, record.get(2)); expectedRecord.put(3, record.get(3)); expectedData.add(expectedRecord); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(expectedSchema, expectedData, root); } @Test public void testSkipRecordField() throws Exception { Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f0"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); Schema schema = getSchema("skip/test_skip_record_before.avsc"); Schema nestedSchema = schema.getFields().get(0).schema(); ArrayList data = new ArrayList<>(); Schema expectedSchema = getSchema("skip/test_skip_record_expected.avsc"); ArrayList expectedData = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); GenericRecord nestedRecord = new GenericData.Record(nestedSchema); nestedRecord.put(0, "test" + i); nestedRecord.put(1, i); record.put(0, nestedRecord); record.put(1, i); data.add(record); GenericRecord expectedRecord = new GenericData.Record(expectedSchema); expectedRecord.put(0, i); expectedData.add(expectedRecord); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(expectedSchema, expectedData, root); } @Test public void testSkipNestedFields() throws Exception { Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f0.f0"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); Schema schema = getSchema("test_nested_record.avsc"); Schema nestedSchema = schema.getFields().get(0).schema(); ArrayList data = new ArrayList<>(); Schema expectedSchema = getSchema("skip/test_skip_second_level_expected.avsc"); Schema expectedNestedSchema = expectedSchema.getFields().get(0).schema(); ArrayList expectedData = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); GenericRecord nestedRecord = new GenericData.Record(nestedSchema); nestedRecord.put(0, "test" + i); nestedRecord.put(1, i); record.put(0, nestedRecord); data.add(record); GenericRecord expectedRecord = new GenericData.Record(expectedSchema); GenericRecord expectedNestedRecord = new GenericData.Record(expectedNestedSchema); expectedNestedRecord.put(0, nestedRecord.get(1)); expectedRecord.put(0, expectedNestedRecord); expectedData.add(expectedRecord); } VectorSchemaRoot root = writeAndRead(schema, data); checkNestedRecordResult(expectedSchema, expectedData, root); } @Test public void testSkipThirdLevelField() throws Exception { Schema firstLevelSchema = getSchema("skip/test_skip_third_level_expected.avsc"); Schema secondLevelSchema = firstLevelSchema.getFields().get(0).schema(); Schema thirdLevelSchema = secondLevelSchema.getFields().get(0).schema(); ArrayList data = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord firstLevelRecord = new GenericData.Record(firstLevelSchema); GenericRecord secondLevelRecord = new GenericData.Record(secondLevelSchema); GenericRecord thirdLevelRecord = new GenericData.Record(thirdLevelSchema); thirdLevelRecord.put(0, i); thirdLevelRecord.put(1, "test" + i); thirdLevelRecord.put(2, i % 2 == 0); secondLevelRecord.put(0, thirdLevelRecord); firstLevelRecord.put(0, secondLevelRecord); data.add(firstLevelRecord); } // do not skip any fields first VectorSchemaRoot root1 = writeAndRead(firstLevelSchema, data); assertEquals(1, root1.getFieldVectors().size()); assertEquals(Types.MinorType.STRUCT, root1.getFieldVectors().get(0).getMinorType()); StructVector secondLevelVector = (StructVector) root1.getFieldVectors().get(0); assertEquals(1, secondLevelVector.getChildrenFromFields().size()); assertEquals( Types.MinorType.STRUCT, secondLevelVector.getChildrenFromFields().get(0).getMinorType()); StructVector thirdLevelVector = (StructVector) secondLevelVector.getChildrenFromFields().get(0); assertEquals(3, thirdLevelVector.getChildrenFromFields().size()); // skip third level field and validate Set skipFieldNames = new HashSet<>(); skipFieldNames.add("f0.f0.f0"); config = new AvroToArrowConfigBuilder(config.getAllocator()) .setSkipFieldNames(skipFieldNames) .build(); VectorSchemaRoot root2 = writeAndRead(firstLevelSchema, data); assertEquals(1, root2.getFieldVectors().size()); assertEquals(Types.MinorType.STRUCT, root2.getFieldVectors().get(0).getMinorType()); StructVector secondStruct = (StructVector) root2.getFieldVectors().get(0); assertEquals(1, secondStruct.getChildrenFromFields().size()); assertEquals( Types.MinorType.STRUCT, secondStruct.getChildrenFromFields().get(0).getMinorType()); StructVector thirdStruct = (StructVector) secondStruct.getChildrenFromFields().get(0); assertEquals(2, thirdStruct.getChildrenFromFields().size()); assertEquals(Types.MinorType.INT, thirdStruct.getChildrenFromFields().get(0).getMinorType()); assertEquals(Types.MinorType.BIT, thirdStruct.getChildrenFromFields().get(1).getMinorType()); } } apache-arrow-java-19.0.0/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/AvroTestBase.java000066400000000000000000000206521515451615200331170ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.File; import java.io.FileInputStream; import java.io.FileOutputStream; import java.io.InputStream; import java.lang.reflect.Method; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.List; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.util.Text; import org.apache.avro.Schema; 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.DatumWriter; import org.apache.avro.io.DecoderFactory; import org.apache.avro.io.EncoderFactory; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.io.TempDir; public class AvroTestBase { @TempDir public File TMP; protected AvroToArrowConfig config; @BeforeEach public void init() { BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); config = new AvroToArrowConfigBuilder(allocator).build(); } public static Schema getSchema(String schemaName) throws Exception { try { // Attempt to use JDK 9 behavior of getting the module then the resource stream from the // module. // Note that this code is caller-sensitive. Method getModuleMethod = Class.class.getMethod("getModule"); Object module = getModuleMethod.invoke(TestWriteReadAvroRecord.class); Method getResourceAsStreamFromModule = module.getClass().getMethod("getResourceAsStream", String.class); try (InputStream is = (InputStream) getResourceAsStreamFromModule.invoke(module, "/schema/" + schemaName)) { return new Schema.Parser().parse(is); } } catch (NoSuchMethodException ex) { // Use JDK8 behavior. try (InputStream is = TestWriteReadAvroRecord.class.getResourceAsStream("/schema/" + schemaName)) { return new Schema.Parser().parse(is); } } } protected VectorSchemaRoot writeAndRead(Schema schema, List data) throws Exception { File dataFile = new File(TMP, "test.avro"); try (FileOutputStream fos = new FileOutputStream(dataFile); FileInputStream fis = new FileInputStream(dataFile)) { BinaryEncoder encoder = new EncoderFactory().directBinaryEncoder(fos, null); DatumWriter writer = new GenericDatumWriter<>(schema); BinaryDecoder decoder = new DecoderFactory().directBinaryDecoder(fis, null); for (Object value : data) { writer.write(value, encoder); } return AvroToArrow.avroToArrow(schema, decoder, config); } } protected void checkArrayResult(List> expected, ListVector vector) { assertEquals(expected.size(), vector.getValueCount()); for (int i = 0; i < expected.size(); i++) { checkArrayElement(expected.get(i), vector.getObject(i)); } } protected void checkArrayElement(List expected, List actual) { assertEquals(expected.size(), actual.size()); for (int i = 0; i < expected.size(); i++) { Object value1 = expected.get(i); Object value2 = actual.get(i); if (value1 == null) { assertTrue(value2 == null); continue; } if (value2 instanceof byte[]) { value2 = ByteBuffer.wrap((byte[]) value2); } else if (value2 instanceof Text) { value2 = value2.toString(); } assertEquals(value1, value2); } } protected void checkPrimitiveResult(List data, FieldVector vector) { assertEquals(data.size(), vector.getValueCount()); for (int i = 0; i < data.size(); i++) { Object value1 = data.get(i); Object value2 = vector.getObject(i); if (value1 == null) { assertTrue(value2 == null); continue; } if (value2 instanceof byte[]) { value2 = ByteBuffer.wrap((byte[]) value2); if (value1 instanceof byte[]) { value1 = ByteBuffer.wrap((byte[]) value1); } } else if (value2 instanceof Text) { value2 = value2.toString(); } else if (value2 instanceof Byte) { value2 = ((Byte) value2).intValue(); } assertEquals(value1, value2); } } protected void checkRecordResult(Schema schema, List data, VectorSchemaRoot root) { assertEquals(data.size(), root.getRowCount()); assertEquals(schema.getFields().size(), root.getFieldVectors().size()); for (int i = 0; i < schema.getFields().size(); i++) { ArrayList fieldData = new ArrayList(); for (GenericRecord record : data) { fieldData.add(record.get(i)); } checkPrimitiveResult(fieldData, root.getFieldVectors().get(i)); } } protected void checkNestedRecordResult( Schema schema, List data, VectorSchemaRoot root) { assertEquals(data.size(), root.getRowCount()); assertTrue(schema.getFields().size() == 1); final Schema nestedSchema = schema.getFields().get(0).schema(); final StructVector structVector = (StructVector) root.getFieldVectors().get(0); for (int i = 0; i < nestedSchema.getFields().size(); i++) { ArrayList fieldData = new ArrayList(); for (GenericRecord record : data) { GenericRecord nestedRecord = (GenericRecord) record.get(0); fieldData.add(nestedRecord.get(i)); } checkPrimitiveResult(fieldData, structVector.getChildrenFromFields().get(i)); } } // belows are for iterator api protected void checkArrayResult(List> expected, List vectors) { int valueCount = vectors.stream().mapToInt(v -> v.getValueCount()).sum(); assertEquals(expected.size(), valueCount); int index = 0; for (ListVector vector : vectors) { for (int i = 0; i < vector.getValueCount(); i++) { checkArrayElement(expected.get(index++), vector.getObject(i)); } } } protected void checkRecordResult( Schema schema, List data, List roots) { roots.forEach( root -> { assertEquals(schema.getFields().size(), root.getFieldVectors().size()); }); for (int i = 0; i < schema.getFields().size(); i++) { List fieldData = new ArrayList(); List vectors = new ArrayList<>(); for (GenericRecord record : data) { fieldData.add(record.get(i)); } final int columnIndex = i; roots.forEach(root -> vectors.add(root.getFieldVectors().get(columnIndex))); checkPrimitiveResult(fieldData, vectors); } } protected void checkPrimitiveResult(List data, List vectors) { int valueCount = vectors.stream().mapToInt(v -> v.getValueCount()).sum(); assertEquals(data.size(), valueCount); int index = 0; for (FieldVector vector : vectors) { for (int i = 0; i < vector.getValueCount(); i++) { Object value1 = data.get(index++); Object value2 = vector.getObject(i); if (value1 == null) { assertNull(value2); continue; } if (value2 instanceof byte[]) { value2 = ByteBuffer.wrap((byte[]) value2); if (value1 instanceof byte[]) { value1 = ByteBuffer.wrap((byte[]) value1); } } else if (value2 instanceof Text) { value2 = value2.toString(); } assertEquals(value1, value2); } } } } AvroToArrowIteratorTest.java000066400000000000000000000227601515451615200352770ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.io.EOFException; import java.io.File; import java.io.FileInputStream; import java.io.FileOutputStream; import java.io.IOException; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Arrays; import java.util.List; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.complex.ListVector; import org.apache.avro.Schema; import org.apache.avro.generic.GenericData; 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.DatumWriter; import org.apache.avro.io.Decoder; import org.apache.avro.io.DecoderFactory; import org.apache.avro.io.EncoderFactory; import org.apache.avro.util.Utf8; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class AvroToArrowIteratorTest extends AvroTestBase { @BeforeEach @Override public void init() { final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); this.config = new AvroToArrowConfigBuilder(allocator).setTargetBatchSize(3).build(); } private void writeDataToFile(Schema schema, List data, File dataFile) throws Exception { try (FileOutputStream fos = new FileOutputStream(dataFile)) { BinaryEncoder encoder = EncoderFactory.get().directBinaryEncoder(fos, null); DatumWriter writer = new GenericDatumWriter<>(schema); for (Object value : data) { writer.write(value, encoder); } encoder.flush(); } } private AvroToArrowVectorIterator convert(Schema schema, FileInputStream fis) throws Exception { BinaryDecoder decoder = DecoderFactory.get().directBinaryDecoder(fis, null); return AvroToArrow.avroToArrowIterator(schema, decoder, config); } @Test public void testStringType() throws Exception { Schema schema = getSchema("test_primitive_string.avsc"); List data = Arrays.asList("v1", "v2", "v3", "v4", "v5"); File dataFile = new File(TMP, "test.avro"); writeDataToFile(schema, data, dataFile); List roots = new ArrayList<>(); List vectors = new ArrayList<>(); try (FileInputStream fis = new FileInputStream(dataFile); AvroToArrowVectorIterator iterator = convert(schema, fis)) { while (iterator.hasNext()) { VectorSchemaRoot root = iterator.next(); FieldVector vector = root.getFieldVectors().get(0); roots.add(root); vectors.add(vector); } } checkPrimitiveResult(data, vectors); AutoCloseables.close(roots); } @Test public void testNullableStringType() throws Exception { Schema schema = getSchema("test_nullable_string.avsc"); List data = new ArrayList<>(); List expected = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); String value = i % 2 == 0 ? "test" + i : null; record.put(0, value); expected.add(value); data.add(record); } File dataFile = new File(TMP, "test.avro"); writeDataToFile(schema, data, dataFile); List roots = new ArrayList<>(); List vectors = new ArrayList<>(); try (FileInputStream fis = new FileInputStream(dataFile); AvroToArrowVectorIterator iterator = convert(schema, fis)) { while (iterator.hasNext()) { VectorSchemaRoot root = iterator.next(); FieldVector vector = root.getFieldVectors().get(0); roots.add(root); vectors.add(vector); } } checkPrimitiveResult(expected, vectors); AutoCloseables.close(roots); } @Test public void testRecordType() throws Exception { Schema schema = getSchema("test_record.avsc"); ArrayList data = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, "test" + i); record.put(1, i); record.put(2, i % 2 == 0); data.add(record); } File dataFile = new File(TMP, "test.avro"); writeDataToFile(schema, data, dataFile); List roots = new ArrayList<>(); try (FileInputStream fis = new FileInputStream(dataFile); AvroToArrowVectorIterator iterator = convert(schema, fis)) { while (iterator.hasNext()) { roots.add(iterator.next()); } } checkRecordResult(schema, data, roots); AutoCloseables.close(roots); } @Test public void testArrayType() throws Exception { Schema schema = getSchema("test_array.avsc"); List> data = Arrays.asList( Arrays.asList("11", "222", "999"), Arrays.asList("12222", "2333", "1000"), Arrays.asList("1rrr", "2ggg"), Arrays.asList("1vvv", "2bbb"), Arrays.asList("1fff", "2")); File dataFile = new File(TMP, "test.avro"); writeDataToFile(schema, data, dataFile); List roots = new ArrayList<>(); List vectors = new ArrayList<>(); try (FileInputStream fis = new FileInputStream(dataFile); AvroToArrowVectorIterator iterator = convert(schema, fis)) { while (iterator.hasNext()) { VectorSchemaRoot root = iterator.next(); roots.add(root); vectors.add((ListVector) root.getFieldVectors().get(0)); } } checkArrayResult(data, vectors); AutoCloseables.close(roots); } @Test public void runLargeNumberOfRows() throws Exception { Schema schema = getSchema("test_large_data.avsc"); int x = 0; final int targetRows = 600000; Decoder fakeDecoder = new FakeDecoder(targetRows); try (AvroToArrowVectorIterator iter = AvroToArrow.avroToArrowIterator( schema, fakeDecoder, new AvroToArrowConfigBuilder(config.getAllocator()).build())) { while (iter.hasNext()) { VectorSchemaRoot root = iter.next(); x += root.getRowCount(); root.close(); } } assertEquals(targetRows, x); } /** Fake avro decoder to test large data. */ private static class FakeDecoder extends Decoder { private int numRows; FakeDecoder(int numRows) { this.numRows = numRows; } // note that Decoder has no hasNext() API, assume enum is the first type in schema // and fixed is the last type in schema and they are unique. private void validate() throws EOFException { if (numRows <= 0) { throw new EOFException(); } } @Override public void readNull() throws IOException {} @Override public boolean readBoolean() throws IOException { return false; } @Override public int readInt() throws IOException { return 0; } @Override public long readLong() throws IOException { return 0; } @Override public float readFloat() throws IOException { return 0; } @Override public double readDouble() throws IOException { return 0; } @Override public Utf8 readString(Utf8 old) throws IOException { return new Utf8("test123test123" + numRows); } @Override public String readString() throws IOException { return "test123test123" + numRows; } @Override public void skipString() throws IOException {} @Override public ByteBuffer readBytes(ByteBuffer old) throws IOException { return ByteBuffer.allocate(0); } @Override public void skipBytes() throws IOException {} @Override public void readFixed(byte[] bytes, int start, int length) throws IOException { // fixed type is last column, after read value, decrease numRows numRows--; } @Override public void skipFixed(int length) throws IOException {} @Override public int readEnum() throws IOException { // enum type is first column, validate numRows first. validate(); return 0; } @Override public long readArrayStart() throws IOException { return 5; } @Override public long arrayNext() throws IOException { return 0; } @Override public long skipArray() throws IOException { return 0; } @Override public long readMapStart() throws IOException { return 5; } @Override public long mapNext() throws IOException { return 0; } @Override public long skipMap() throws IOException { return 0; } @Override public int readIndex() throws IOException { return 0; } } } AvroToArrowTest.java000066400000000000000000000402661515451615200335660ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import org.apache.arrow.vector.FieldVector; 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.avro.Schema; import org.apache.avro.generic.GenericData; import org.apache.avro.generic.GenericRecord; import org.junit.jupiter.api.Test; public class AvroToArrowTest extends AvroTestBase { @Test public void testStringType() throws Exception { Schema schema = getSchema("test_primitive_string.avsc"); List data = Arrays.asList("v1", "v2", "v3", "v4", "v5"); VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(data, vector); } @Test public void testNullableStringType() throws Exception { Schema schema = getSchema("test_nullable_string.avsc"); ArrayList data = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, i % 2 == 0 ? "test" + i : null); data.add(record); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(schema, data, root); } @Test public void testRecordType() throws Exception { Schema schema = getSchema("test_record.avsc"); ArrayList data = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, "test" + i); record.put(1, i); record.put(2, i % 2 == 0); data.add(record); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(schema, data, root); } @Test public void testFixedAttributes() throws Exception { Schema schema = getSchema("attrs/test_fixed_attr.avsc"); List data = new ArrayList<>(); for (int i = 0; i < 5; i++) { byte[] value = ("value" + i).getBytes(StandardCharsets.UTF_8); GenericData.Fixed fixed = new GenericData.Fixed(schema); fixed.bytes(value); data.add(fixed); } VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); Map metadata = vector.getField().getMetadata(); assertEquals("fixed doc", metadata.get("doc")); assertEquals("[\"alias1\",\"alias2\"]", metadata.get("aliases")); } @Test public void testEnumAttributes() throws Exception { Schema schema = getSchema("attrs/test_enum_attrs.avsc"); List data = Arrays.asList( new GenericData.EnumSymbol(schema, "SPADES"), new GenericData.EnumSymbol(schema, "HEARTS"), new GenericData.EnumSymbol(schema, "DIAMONDS"), new GenericData.EnumSymbol(schema, "CLUBS"), new GenericData.EnumSymbol(schema, "SPADES")); VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); Map metadata = vector.getField().getMetadata(); assertEquals("enum doc", metadata.get("doc")); assertEquals("[\"alias1\",\"alias2\"]", metadata.get("aliases")); } @Test public void testRecordAttributes() throws Exception { Schema schema = getSchema("attrs/test_record_attrs.avsc"); Schema nestedSchema = schema.getFields().get(0).schema(); ArrayList data = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); GenericRecord nestedRecord = new GenericData.Record(nestedSchema); nestedRecord.put(0, "test" + i); nestedRecord.put(1, i); record.put(0, nestedRecord); data.add(record); } VectorSchemaRoot root = writeAndRead(schema, data); StructVector structVector = (StructVector) root.getFieldVectors().get(0); Map structMeta = structVector.getField().getMetadata(); Map childMeta1 = structVector.getChildByOrdinal(0).getField().getMetadata(); Map childMeta2 = structVector.getChildByOrdinal(1).getField().getMetadata(); assertEquals("f0 doc", structMeta.get("doc")); assertEquals("[\"f0.a1\"]", structMeta.get("aliases")); assertEquals("f1 doc", childMeta1.get("doc")); assertEquals("[\"f1.a1\",\"f1.a2\"]", childMeta1.get("aliases")); assertEquals("f2 doc", childMeta2.get("doc")); assertEquals("[\"f2.a1\",\"f2.a2\"]", childMeta2.get("aliases")); } @Test public void testNestedRecordType() throws Exception { Schema schema = getSchema("test_nested_record.avsc"); Schema nestedSchema = schema.getFields().get(0).schema(); ArrayList data = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); GenericRecord nestedRecord = new GenericData.Record(nestedSchema); nestedRecord.put(0, "test" + i); nestedRecord.put(1, i); record.put(0, nestedRecord); data.add(record); } VectorSchemaRoot root = writeAndRead(schema, data); checkNestedRecordResult(schema, data, root); } @Test public void testEnumType() throws Exception { Schema schema = getSchema("test_primitive_enum.avsc"); List data = Arrays.asList( new GenericData.EnumSymbol(schema, "SPADES"), new GenericData.EnumSymbol(schema, "HEARTS"), new GenericData.EnumSymbol(schema, "DIAMONDS"), new GenericData.EnumSymbol(schema, "CLUBS"), new GenericData.EnumSymbol(schema, "SPADES")); List expectedIndices = Arrays.asList(0, 1, 2, 3, 0); VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(expectedIndices, vector); VarCharVector dictVector = (VarCharVector) config.getProvider().lookup(0).getVector(); assertEquals(4, dictVector.getValueCount()); assertEquals("SPADES", dictVector.getObject(0).toString()); assertEquals("HEARTS", dictVector.getObject(1).toString()); assertEquals("DIAMONDS", dictVector.getObject(2).toString()); assertEquals("CLUBS", dictVector.getObject(3).toString()); } @Test public void testIntType() throws Exception { Schema schema = getSchema("test_primitive_int.avsc"); List data = Arrays.asList(1, 2, 3, 4, 5); VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(data, vector); } @Test public void testNullableIntType() throws Exception { Schema schema = getSchema("test_nullable_int.avsc"); ArrayList data = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, i % 2 == 0 ? i : null); data.add(record); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(schema, data, root); } @Test public void testLongType() throws Exception { Schema schema = getSchema("test_primitive_long.avsc"); List data = Arrays.asList(1L, 2L, 3L, 4L, 5L); VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(data, vector); } @Test public void testNullableLongType() throws Exception { Schema schema = getSchema("test_nullable_long.avsc"); ArrayList data = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, i % 2 == 0 ? (long) i : null); data.add(record); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(schema, data, root); } @Test public void testFloatType() throws Exception { Schema schema = getSchema("test_primitive_float.avsc"); List data = Arrays.asList(1.1f, 2.2f, 3.3f, 4.4f, 5.5f); VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(data, vector); } @Test public void testNullableFloatType() throws Exception { Schema schema = getSchema("test_nullable_float.avsc"); ArrayList data = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, i % 2 == 0 ? i + 0.1f : null); data.add(record); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(schema, data, root); } @Test public void testDoubleType() throws Exception { Schema schema = getSchema("test_primitive_double.avsc"); List data = Arrays.asList(1.1, 2.2, 3.3, 4.4, 5.5); VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(data, vector); } @Test public void testNullableDoubleType() throws Exception { Schema schema = getSchema("test_nullable_double.avsc"); ArrayList data = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, i % 2 == 0 ? i + 0.1 : null); data.add(record); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(schema, data, root); } @Test public void testBytesType() throws Exception { Schema schema = getSchema("test_primitive_bytes.avsc"); List data = Arrays.asList( ByteBuffer.wrap("value1".getBytes(StandardCharsets.UTF_8)), ByteBuffer.wrap("value2".getBytes(StandardCharsets.UTF_8)), ByteBuffer.wrap("value3".getBytes(StandardCharsets.UTF_8)), ByteBuffer.wrap("value4".getBytes(StandardCharsets.UTF_8)), ByteBuffer.wrap("value5".getBytes(StandardCharsets.UTF_8))); VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(data, vector); } @Test public void testNullableBytesType() throws Exception { Schema schema = getSchema("test_nullable_bytes.avsc"); ArrayList data = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put( 0, i % 2 == 0 ? ByteBuffer.wrap(("test" + i).getBytes(StandardCharsets.UTF_8)) : null); data.add(record); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(schema, data, root); } @Test public void testBooleanType() throws Exception { Schema schema = getSchema("test_primitive_boolean.avsc"); List data = Arrays.asList(true, false, true, false, true); VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(data, vector); } @Test public void testNullableBooleanType() throws Exception { Schema schema = getSchema("test_nullable_boolean.avsc"); ArrayList data = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, i % 2 == 0 ? true : null); data.add(record); } VectorSchemaRoot root = writeAndRead(schema, data); checkRecordResult(schema, data, root); } @Test public void testArrayType() throws Exception { Schema schema = getSchema("test_array.avsc"); List> data = Arrays.asList( Arrays.asList("11", "222", "999"), Arrays.asList("12222", "2333", "1000"), Arrays.asList("1rrr", "2ggg"), Arrays.asList("1vvv", "2bbb"), Arrays.asList("1fff", "2")); VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkArrayResult(data, (ListVector) vector); } @Test public void testMapType() throws Exception { Schema schema = getSchema("test_map.avsc"); List keys = Arrays.asList("key1", "key2", "key3", "key4", "key5", "key6"); List vals = Arrays.asList("val1", "val2", "val3", "val4", "val5", "val6"); List data = new ArrayList<>(); LinkedHashMap map1 = new LinkedHashMap(); map1.put(keys.get(0), vals.get(0)); map1.put(keys.get(1), vals.get(1)); data.add(map1); LinkedHashMap map2 = new LinkedHashMap(); map2.put(keys.get(2), vals.get(2)); map2.put(keys.get(3), vals.get(3)); data.add(map2); LinkedHashMap map3 = new LinkedHashMap(); map3.put(keys.get(4), vals.get(4)); map3.put(keys.get(5), vals.get(5)); data.add(map3); VectorSchemaRoot root = writeAndRead(schema, data); MapVector vector = (MapVector) root.getFieldVectors().get(0); checkPrimitiveResult(keys, vector.getDataVector().getChildrenFromFields().get(0)); checkPrimitiveResult(vals, vector.getDataVector().getChildrenFromFields().get(1)); assertEquals(0, vector.getOffsetBuffer().getInt(0)); assertEquals(2, vector.getOffsetBuffer().getInt(1 * 4)); assertEquals(4, vector.getOffsetBuffer().getInt(2 * 4)); assertEquals(6, vector.getOffsetBuffer().getInt(3 * 4)); } @Test public void testFixedType() throws Exception { Schema schema = getSchema("test_fixed.avsc"); List data = new ArrayList<>(); List expected = new ArrayList<>(); for (int i = 0; i < 5; i++) { byte[] value = ("value" + i).getBytes(StandardCharsets.UTF_8); expected.add(value); GenericData.Fixed fixed = new GenericData.Fixed(schema); fixed.bytes(value); data.add(fixed); } VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(expected, vector); } @Test public void testUnionType() throws Exception { Schema schema = getSchema("test_union.avsc"); ArrayList data = new ArrayList<>(); ArrayList expected = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, i % 2 == 0 ? "test" + i : i); expected.add(i % 2 == 0 ? "test" + i : i); data.add(record); } VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(expected, vector); } @Test public void testNullableUnionType() throws Exception { Schema schema = getSchema("test_nullable_union.avsc"); ArrayList data = new ArrayList<>(); ArrayList expected = new ArrayList<>(); for (int i = 0; i < 5; i++) { GenericRecord record = new GenericData.Record(schema); if (i % 3 == 0) { record.put(0, "test" + i); expected.add("test" + i); data.add(record); } else if (i % 3 == 1) { record.put(0, i); expected.add(i); data.add(record); } else { record.put(0, null); expected.add(null); data.add(record); } } VectorSchemaRoot root = writeAndRead(schema, data); FieldVector vector = root.getFieldVectors().get(0); checkPrimitiveResult(expected, vector); } } RoundTripDataTest.java000066400000000000000000001735501515451615200340640ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } } RoundTripSchemaTest.java000066400000000000000000000447121515451615200344100ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } TestWriteReadAvroRecord.java000066400000000000000000000163341515451615200352150ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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; import org.apache.avro.generic.GenericData; 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 { @TempDir public static File TMP; @Test public void testWriteAndRead() throws Exception { File dataFile = new File(TMP, "test.avro"); Schema schema = AvroTestBase.getSchema("test.avsc"); // write data to disk GenericRecord user1 = new GenericData.Record(schema); user1.put("name", "Alyssa"); user1.put("favorite_number", 256); GenericRecord user2 = new GenericData.Record(schema); user2.put("name", "Ben"); user2.put("favorite_number", 7); user2.put("favorite_color", "red"); DatumWriter datumWriter = new GenericDatumWriter(schema); try (DataFileWriter dataFileWriter = new DataFileWriter(datumWriter)) { dataFileWriter.create(schema, dataFile); dataFileWriter.append(user1); dataFileWriter.append(user2); } // read data from disk DatumReader datumReader = new GenericDatumReader(schema); List result = new ArrayList<>(); try (DataFileReader dataFileReader = new DataFileReader(dataFile, datumReader)) { while (dataFileReader.hasNext()) { GenericRecord user = dataFileReader.next(); result.add(user); } } assertEquals(2, result.size()); GenericRecord deUser1 = result.get(0); assertEquals("Alyssa", deUser1.get("name").toString()); assertEquals(256, deUser1.get("favorite_number")); assertEquals(null, deUser1.get("favorite_color")); GenericRecord deUser2 = result.get(1); assertEquals("Ben", deUser2.get("name").toString()); 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)); } } } } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/000077500000000000000000000000001515451615200232455ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/resources/logback.xml000066400000000000000000000025221515451615200253720ustar00rootroot00000000000000 %d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/000077500000000000000000000000001515451615200245055ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/attrs/000077500000000000000000000000001515451615200256425ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/attrs/test_enum_attrs.avsc000066400000000000000000000016731515451615200317470ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ { "type": "enum", "symbols" : ["SPADES", "HEARTS", "DIAMONDS", "CLUBS"], "name": "testEnum", "doc" : "enum doc", "aliases" : ["alias1", "alias2"] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/attrs/test_fixed_attr.avsc000066400000000000000000000016221515451615200317110ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ { "type": "fixed", "size": 6, "name": "testFixed", "doc" : "fixed doc", "aliases" : ["alias1", "alias2"] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/attrs/test_record_attrs.avsc000066400000000000000000000024101515451615200322470ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testAttrs", "fields": [ { "name" : "f0", "type" : { "type" : "record", "name" : "nestedInRecord", "doc" : "f0 doc", "aliases" : ["f0.a1"], "fields": [ {"name": "f1", "type": "string", "doc": "f1 doc", "aliases" : ["f1.a1", "f1.a2"]}, {"name": "f2", "type": "int", "doc": "f2 doc", "aliases" : ["f2.a1", "f2.a2"]} ] } } ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/logical/000077500000000000000000000000001515451615200261175ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/logical/test_date.avsc000066400000000000000000000016061515451615200307540ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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": "int", "logicalType" : "date" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/logical/test_decimal_invalid1.avsc000066400000000000000000000016521515451615200332250ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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": "bytes", "logicalType" : "decimal", "precision": 77, "scale": 2 } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/logical/test_decimal_invalid2.avsc000066400000000000000000000016531515451615200332270ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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": "bytes", "logicalType" : "decimal", "precision": 20, "scale": -1 } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/logical/test_decimal_invalid3.avsc000066400000000000000000000016531515451615200332300ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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": "bytes", "logicalType" : "decimal", "precision": 20, "scale": 40 } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/logical/test_decimal_invalid4.avsc000066400000000000000000000016671515451615200332360ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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": "fixed", "size" : 1, "logicalType" : "decimal", "precision": 30, "scale": 2 } test_decimal_with_original_bytes.avsc000066400000000000000000000016521515451615200355040ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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": "bytes", "logicalType" : "decimal", "precision": 10, "scale": 2 } test_decimal_with_original_fixed.avsc000066400000000000000000000016701515451615200354550ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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": "fixed", "size" : 10, "logicalType" : "decimal", "precision": 10, "scale": 2 } test_local_timestamp_micros.avsc000066400000000000000000000016311515451615200345070ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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" } test_local_timestamp_millis.avsc000066400000000000000000000016311515451615200345040ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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" } test_local_timestamp_nanos.avsc000066400000000000000000000016301515451615200343300ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/logical/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/logical/test_time_micros.avsc000066400000000000000000000016161515451615200323520ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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" : "time-micros" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/logical/test_time_millis.avsc000066400000000000000000000016151515451615200323460ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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": "int", "logicalType" : "time-millis" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/logical/test_timestamp_micros.avsc000066400000000000000000000016231515451615200334150ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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-micros" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/logical/test_timestamp_millis.avsc000066400000000000000000000016231515451615200334120ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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-millis" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/logical/test_timestamp_nanos.avsc000066400000000000000000000016221515451615200332360ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/000077500000000000000000000000001515451615200254535ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_array_before.avsc000066400000000000000000000020231515451615200330530ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "test", "fields": [ {"name": "f0", "type": "string"}, {"name": "f1", "type": {"type" : "array", "items": "string"}}, {"name": "f2", "type": "boolean"} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_array_expected.avsc000066400000000000000000000017171515451615200334230ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "test", "fields": [ {"name": "f0", "type": "string"}, {"name": "f2", "type": "boolean"} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_base1.avsc000066400000000000000000000021761515451615200314170ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testRecord", "fields": [ {"name": "f0", "type": {"type" : "fixed", "size":5, "name" : "fix"}}, {"name": "f1", "type": {"type" : "enum", "name" : "enum", "symbols": ["TEST0", "TEST1"]}}, {"name": "f2", "type": "string"}, {"name": "f3", "type": "bytes"} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_base2.avsc000066400000000000000000000021041515451615200314070ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testRecord", "fields": [ {"name": "f0", "type": "boolean"}, {"name": "f1", "type": "int"}, {"name": "f2", "type": "long"}, {"name": "f3", "type": "float"}, {"name": "f4", "type": "double"} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_boolean_expected.avsc000066400000000000000000000020341515451615200337150ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testRecord", "fields": [ {"name": "f1", "type": "int"}, {"name": "f2", "type": "long"}, {"name": "f3", "type": "float"}, {"name": "f4", "type": "double"} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_bytes_expected.avsc000066400000000000000000000021301515451615200334210ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testRecord", "fields": [ {"name": "f0", "type": {"type" : "fixed", "size":5, "name" : "fix"}}, {"name": "f1", "type": {"type" : "enum", "name" : "enum", "symbols": ["TEST0", "TEST1"]}}, {"name": "f2", "type": "string"} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_double_expected.avsc000066400000000000000000000020351515451615200335510ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testRecord", "fields": [ {"name": "f0", "type": "boolean"}, {"name": "f1", "type": "int"}, {"name": "f2", "type": "long"}, {"name": "f3", "type": "float"} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_enum_expected.avsc000066400000000000000000000020361515451615200332440ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testRecord", "fields": [ {"name": "f0", "type": {"type" : "fixed", "size":5, "name" : "fix"}}, {"name": "f2", "type": "string"}, {"name": "f3", "type": "bytes"} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_fixed_expected.avsc000066400000000000000000000020631515451615200333770ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testRecord", "fields": [ {"name": "f1", "type": {"type" : "enum", "name" : "enum", "symbols": ["TEST0", "TEST1"]}}, {"name": "f2", "type": "string"}, {"name": "f3", "type": "bytes"} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_float_expected.avsc000066400000000000000000000020361515451615200334050ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testRecord", "fields": [ {"name": "f0", "type": "boolean"}, {"name": "f1", "type": "int"}, {"name": "f2", "type": "long"}, {"name": "f4", "type": "double"} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_int_expected.avsc000066400000000000000000000020401515451615200330650ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testRecord", "fields": [ {"name": "f0", "type": "boolean"}, {"name": "f2", "type": "long"}, {"name": "f3", "type": "float"}, {"name": "f4", "type": "double"} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_long_expected.avsc000066400000000000000000000020371515451615200332400ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testRecord", "fields": [ {"name": "f0", "type": "boolean"}, {"name": "f1", "type": "int"}, {"name": "f3", "type": "float"}, {"name": "f4", "type": "double"} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_map_before.avsc000066400000000000000000000020221515451615200325110ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "test", "fields": [ {"name": "f0", "type": "string"}, {"name": "f1", "type": {"type" : "map", "values": "string"}}, {"name": "f2", "type": "boolean"} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_map_expected.avsc000066400000000000000000000017171515451615200330620ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "test", "fields": [ {"name": "f0", "type": "string"}, {"name": "f2", "type": "boolean"} ] } test_skip_multi_fields_expected.avsc000066400000000000000000000016531515451615200347050ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testSkip", "fields": [ {"name": "f0", "type": "string"} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_record_before.avsc000066400000000000000000000022251515451615200332170ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "test", "fields": [ { "name" : "f0", "type" : { "type" : "record", "name" : "nestedInRecord", "fields": [ {"name": "f00", "type": "string"}, {"name": "f01", "type": "int"} ] } }, { "name" : "f1", "type" : "int" } ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_record_expected.avsc000066400000000000000000000016471515451615200335650ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "test", "fields": [ { "name" : "f1", "type" : "int"} ] } test_skip_second_level_expected.avsc000066400000000000000000000021071515451615200346620ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testSkipNested", "fields": [ { "name" : "nested", "type" : { "type" : "record", "name" : "nestedInRecord", "fields": [ {"name": "f1", "type": "int"} ] } } ] } test_skip_single_field_expected.avsc000066400000000000000000000017231515451615200346470ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testSkip", "fields": [ {"name": "f0", "type": "string"}, {"name": "f2", "type": "boolean"} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_string_expected.avsc000066400000000000000000000021271515451615200336070ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testRecord", "fields": [ {"name": "f0", "type": {"type" : "fixed", "size":5, "name" : "fix"}}, {"name": "f1", "type": {"type" : "enum", "name" : "enum", "symbols": ["TEST0", "TEST1"]}}, {"name": "f3", "type": "bytes"} ] } test_skip_third_level_expected.avsc000066400000000000000000000026001515451615200345170ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "firstLevel", "fields": [ { "name" : "f0", "type" : { "type" : "record", "name" : "secondLevel", "fields": [ { "name" : "f0", "type" : { "type" : "record", "name" : "thirdLevel", "fields" : [ {"name": "f1", "type": "int"}, {"name": "f0", "type": "string"}, {"name": "f2", "type": "boolean"} ] } } ] } } ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/test_skip_union_before.avsc000066400000000000000000000020561515451615200330730ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "test", "fields": [ {"name": "f0", "type": ["string"]}, {"name": "f1", "type": ["string", "null"]}, {"name": "f2", "type": ["string", "int"]}, {"name": "f3", "type": "int"} ] } test_skip_union_multi_fields_expected.avsc000066400000000000000000000017761515451615200361230ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "test", "fields": [ {"name": "f0", "type": ["string"]}, {"name": "f1", "type": ["string", "null"]}, {"name": "f3", "type": "int"} ] } test_skip_union_nullable_field_expected.avsc000066400000000000000000000017751515451615200364030ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "test", "fields": [ {"name": "f0", "type": ["string"]}, {"name": "f2", "type": ["string", "int"]}, {"name": "f3", "type": "int"} ] } test_skip_union_one_field_expected.avsc000066400000000000000000000020211515451615200353470ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/skip/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "test", "fields": [ {"name": "f1", "type": ["string", "null"]}, {"name": "f2", "type": ["string", "int"]}, {"name": "f3", "type": ["string", "int"]} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test.avsc000066400000000000000000000020421515451615200263400ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "User", "fields": [ {"name": "name", "type": "string"}, {"name": "favorite_number", "type": ["int", "null"]}, {"name": "favorite_color", "type": ["string", "null"]} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_array.avsc000066400000000000000000000016101515451615200275360ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "array", "items": "string", "name": "testArray" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_fixed.avsc000066400000000000000000000016001515451615200275160ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "fixed", "size": 6, "name": "testFixed" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_large_data.avsc000066400000000000000000000040331515451615200305050ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testLargeData", "fields": [ { "name": "f0", "type": { "name" : "f0", "type" : "enum", "symbols" : ["value1", "value2", "value3", "value4", "value5"] } }, { "name" : "f1", "type" : { "type" : "record", "name" : "nestedRecord", "fields": [ {"name": "f1_0", "type": "string"}, {"name": "f1_1", "type": "int"} ] } }, {"name": "f2", "type": "string"}, {"name": "f3", "type": "int"}, {"name": "f4", "type": "boolean"}, {"name": "f5", "type": "float"}, {"name": "f6", "type": "double"}, {"name": "f7", "type": "bytes"}, {"name": "f8", "type": ["string", "int"]}, { "name": "f9", "type": { "name" : "f9", "type" : "array", "items" : "string" } }, { "name": "f10", "type": { "name" : "f10", "type" : "map", "values" : "string" } }, { "name": "f11", "type": { "type" : "fixed", "name" : "f11", "size" : 5 } } ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_map.avsc000066400000000000000000000016051515451615200272010ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "map", "values": "string", "name": "testMap" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_nested_record.avsc000066400000000000000000000021641515451615200312450ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testNestedRecord", "fields": [ { "name" : "f0", "type" : { "type" : "record", "name" : "nestedInRecord", "fields": [ {"name": "f0", "type": "string"}, {"name": "f1", "type": "int"} ] } } ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_nullable_boolean.avsc000066400000000000000000000016751515451615200317300ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "nullableBoolean", "fields": [ {"name": "f0", "type": ["null", "boolean"]} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_nullable_bytes.avsc000066400000000000000000000016711515451615200314330ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "nullableBytes", "fields": [ {"name": "f0", "type": ["null", "bytes"]} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_nullable_double.avsc000066400000000000000000000016731515451615200315610ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "nullableDouble", "fields": [ {"name": "f0", "type": ["null", "double"]} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_nullable_float.avsc000066400000000000000000000016711515451615200314120ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "nullableFloat", "fields": [ {"name": "f0", "type": ["null", "float"]} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_nullable_int.avsc000066400000000000000000000016651515451615200311020ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "nullableInt", "fields": [ {"name": "f0", "type": ["null", "int"]} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_nullable_long.avsc000066400000000000000000000016671515451615200312510ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "nullableLong", "fields": [ {"name": "f0", "type": ["null", "long"]} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_nullable_string.avsc000066400000000000000000000016731515451615200316150ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "nullableString", "fields": [ {"name": "f0", "type": ["null", "string"]} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_nullable_union.avsc000066400000000000000000000017061515451615200314340ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testNullableUnions", "fields": [ {"name": "f0", "type": ["string", "int", "null"]} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_primitive_boolean.avsc000066400000000000000000000015701515451615200321340ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "boolean", "name": "TestBoolean" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_primitive_bytes.avsc000066400000000000000000000015641515451615200316460ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "bytes", "name": "TestBytes" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_primitive_double.avsc000066400000000000000000000015661515451615200317740ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "double", "name": "TestDouble" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_primitive_enum.avsc000066400000000000000000000016521515451615200314620ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "enum", "symbols" : ["SPADES", "HEARTS", "DIAMONDS", "CLUBS"], "name": "testEnum" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_primitive_float.avsc000066400000000000000000000015641515451615200316250ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "float", "name": "TestFloat" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_primitive_int.avsc000066400000000000000000000015601515451615200313060ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "int", "name": "TestInt" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_primitive_long.avsc000066400000000000000000000015621515451615200314550ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "long", "name": "TestLong" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_primitive_string.avsc000066400000000000000000000015661515451615200320300ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "string", "name": "TestString" } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_record.avsc000066400000000000000000000017711515451615200277060ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testRecord", "fields": [ {"name": "f0", "type": "string"}, {"name": "f1", "type": "int"}, {"name": "f2", "type": "boolean"} ] } apache-arrow-java-19.0.0/adapter/avro/src/test/resources/schema/test_union.avsc000066400000000000000000000016661515451615200275630ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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", "type": "record", "name": "testUnions", "fields": [ {"name": "f0", "type": ["string", "int"]} ] } apache-arrow-java-19.0.0/adapter/jdbc/000077500000000000000000000000001515451615200174205ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/pom.xml000066400000000000000000000103241515451615200207350ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-java-root 19.0.0 ../../pom.xml arrow-jdbc Arrow JDBC Adapter (Contrib/Experimental)A library for converting JDBC data to Arrow data. http://maven.apache.org org.apache.arrow arrow-memory-core org.apache.arrow arrow-memory-netty runtime org.apache.arrow arrow-vector ${arrow.vector.classifier} org.immutables value-annotations com.h2database h2 2.3.232 test com.fasterxml.jackson.dataformat jackson-dataformat-yaml test com.fasterxml.jackson.core jackson-databind com.fasterxml.jackson.core jackson-core com.fasterxml.jackson.core jackson-annotations org.assertj assertj-core test org.apache.maven.plugins maven-dependency-plugin analyze verify com.fasterxml.jackson.core:jackson-annotations org.apache.maven.plugins maven-surefire-plugin --add-reads=org.apache.arrow.adapter.jdbc=com.fasterxml.jackson.dataformat.yaml --add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED -Duser.timezone=UTC org.apache.maven.plugins maven-compiler-plugin -Werror apache-arrow-java-19.0.0/adapter/jdbc/src/000077500000000000000000000000001515451615200202075ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/000077500000000000000000000000001515451615200211335ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/000077500000000000000000000000001515451615200220545ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/module-info.java000066400000000000000000000022551515451615200251410ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ module org.apache.arrow.adapter.jdbc { exports org.apache.arrow.adapter.jdbc.consumer; exports org.apache.arrow.adapter.jdbc; exports org.apache.arrow.adapter.jdbc.binder; requires com.fasterxml.jackson.annotation; requires com.fasterxml.jackson.databind; requires java.sql; requires jdk.unsupported; requires org.apache.arrow.memory.core; requires org.apache.arrow.vector; } apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/000077500000000000000000000000001515451615200226435ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/000077500000000000000000000000001515451615200240645ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200252165ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/000077500000000000000000000000001515451615200266365ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/000077500000000000000000000000001515451615200275405ustar00rootroot00000000000000ArrowVectorIterator.java000066400000000000000000000161231515451615200343160ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; import static org.apache.arrow.adapter.jdbc.JdbcToArrowUtils.isColumnNullable; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.util.Iterator; import org.apache.arrow.adapter.jdbc.consumer.CompositeJdbcConsumer; import org.apache.arrow.adapter.jdbc.consumer.JdbcConsumer; import org.apache.arrow.adapter.jdbc.consumer.exceptions.JdbcConsumerException; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.ValueVectorUtility; /** VectorSchemaRoot iterator for partially converting JDBC data. */ public class ArrowVectorIterator implements Iterator, AutoCloseable { private final ResultSet resultSet; private final JdbcToArrowConfig config; private final Schema schema; private final ResultSetMetaData rsmd; private final JdbcConsumer[] consumers; final CompositeJdbcConsumer compositeConsumer; // this is used only if resuing vector schema root is enabled. private VectorSchemaRoot nextBatch; private final int targetBatchSize; // This is used to track whether the ResultSet has been fully read, and is needed specifically for // cases where there // is a ResultSet having zero rows (empty): private boolean readComplete = false; /** Construct an instance. */ private ArrowVectorIterator(ResultSet resultSet, JdbcToArrowConfig config) throws SQLException { this.resultSet = resultSet; this.config = config; this.schema = JdbcToArrowUtils.jdbcToArrowSchema(resultSet.getMetaData(), config); this.targetBatchSize = config.getTargetBatchSize(); rsmd = resultSet.getMetaData(); consumers = new JdbcConsumer[rsmd.getColumnCount()]; this.compositeConsumer = new CompositeJdbcConsumer(consumers); this.nextBatch = config.isReuseVectorSchemaRoot() ? createVectorSchemaRoot() : null; } /** Create a ArrowVectorIterator to partially convert data. */ public static ArrowVectorIterator create(ResultSet resultSet, JdbcToArrowConfig config) throws SQLException { ArrowVectorIterator iterator = null; try { iterator = new ArrowVectorIterator(resultSet, config); } catch (Throwable e) { AutoCloseables.close(e, iterator); throw new RuntimeException("Error occurred while creating iterator.", e); } return iterator; } private void consumeData(VectorSchemaRoot root) { // consume data try { int readRowCount = 0; if (targetBatchSize == JdbcToArrowConfig.NO_LIMIT_BATCH_SIZE) { while (resultSet.next()) { ValueVectorUtility.ensureCapacity(root, readRowCount + 1); compositeConsumer.consume(resultSet); readRowCount++; } readComplete = true; } else { while ((readRowCount < targetBatchSize) && !readComplete) { if (resultSet.next()) { compositeConsumer.consume(resultSet); readRowCount++; } else { readComplete = true; } } } root.setRowCount(readRowCount); } catch (Throwable e) { compositeConsumer.close(); if (e instanceof JdbcConsumerException) { throw (JdbcConsumerException) e; } else { throw new RuntimeException("Error occurred while consuming data.", e); } } } private VectorSchemaRoot createVectorSchemaRoot() throws SQLException { VectorSchemaRoot root = null; try { root = VectorSchemaRoot.create(schema, config.getAllocator()); if (config.getTargetBatchSize() != JdbcToArrowConfig.NO_LIMIT_BATCH_SIZE) { ValueVectorUtility.preAllocate(root, config.getTargetBatchSize()); } } catch (Throwable e) { if (root != null) { root.close(); } throw new RuntimeException("Error occurred while creating schema root.", e); } initialize(root); return root; } private void initialize(VectorSchemaRoot root) throws SQLException { for (int i = 1; i <= consumers.length; i++) { final JdbcFieldInfo columnFieldInfo = JdbcToArrowUtils.getJdbcFieldInfoForColumn(rsmd, i, config); ArrowType arrowType = config.getJdbcToArrowTypeConverter().apply(columnFieldInfo); consumers[i - 1] = config .getJdbcConsumerGetter() .apply( arrowType, i, isColumnNullable(resultSet.getMetaData(), i, columnFieldInfo), root.getVector(i - 1), config); } } // Loads the next schema root or null if no more rows are available. private void load(VectorSchemaRoot root) { for (int i = 0; i < consumers.length; i++) { FieldVector vec = root.getVector(i); if (config.isReuseVectorSchemaRoot()) { // if we are reusing the vector schema root, // we must reset the vector before populating it with data. vec.reset(); } consumers[i].resetValueVector(vec); } consumeData(root); } @Override public boolean hasNext() { return !readComplete; } /** * Gets the next vector. If {@link JdbcToArrowConfig#isReuseVectorSchemaRoot()} is false, the * client is responsible for freeing its resources. * * @throws JdbcConsumerException on error from VectorConsumer */ @Override public VectorSchemaRoot next() { Preconditions.checkArgument(hasNext()); try { VectorSchemaRoot ret = config.isReuseVectorSchemaRoot() ? nextBatch : createVectorSchemaRoot(); load(ret); return ret; } catch (Exception e) { close(); if (e instanceof JdbcConsumerException) { throw (JdbcConsumerException) e; } else { throw new RuntimeException("Error occurred while getting next schema root.", e); } } } /** * Clean up resources ONLY WHEN THE {@link VectorSchemaRoot} HOLDING EACH BATCH IS REUSED. If a * new VectorSchemaRoot is created for each batch, each root must be closed manually by the client * code. */ @Override public void close() { if (config.isReuseVectorSchemaRoot()) { nextBatch.close(); compositeConsumer.close(); } } } apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/Constants.java000066400000000000000000000024151515451615200323610ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; /** String constants used for metadata returned on Vectors. */ public class Constants { private Constants() {} public static final String SQL_CATALOG_NAME_KEY = "SQL_CATALOG_NAME"; public static final String SQL_SCHEMA_NAME_KEY = "SQL_SCHEMA_NAME"; public static final String SQL_TABLE_NAME_KEY = "SQL_TABLE_NAME"; public static final String SQL_COLUMN_NAME_KEY = "SQL_COLUMN_NAME"; public static final String SQL_TYPE_KEY = "SQL_TYPE"; } apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcFieldInfo.java000066400000000000000000000155171515451615200330360ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Types; import org.apache.arrow.util.Preconditions; /** * This class represents the information about a JDBC ResultSet Field that is needed to construct an * {@link org.apache.arrow.vector.types.pojo.ArrowType}. Currently, this is: * *
    *
  • The JDBC {@link java.sql.Types} type. *
  • The nullability. *
  • The field's precision (used for {@link java.sql.Types#DECIMAL} and {@link * java.sql.Types#NUMERIC} types). *
  • The field's scale (used for {@link java.sql.Types#DECIMAL} and {@link * java.sql.Types#NUMERIC} types). *
*/ public class JdbcFieldInfo { private final int column; private final int jdbcType; private final int nullability; private final int precision; private final int scale; private final String typeName; private final int displaySize; /** * Builds a JdbcFieldInfo using only the {@link java.sql.Types} type. Do not use this * constructor if the field type is {@link java.sql.Types#DECIMAL} or {@link * java.sql.Types#NUMERIC}; the precision and scale will be set to 0. * * @param jdbcType The {@link java.sql.Types} type. * @throws IllegalArgumentException if jdbcType is {@link java.sql.Types#DECIMAL} or {@link * java.sql.Types#NUMERIC}. */ public JdbcFieldInfo(int jdbcType) { Preconditions.checkArgument( (jdbcType != Types.DECIMAL && jdbcType != Types.NUMERIC), "DECIMAL and NUMERIC types require a precision and scale; please use another constructor."); this.column = 0; this.jdbcType = jdbcType; this.nullability = ResultSetMetaData.columnNullableUnknown; this.precision = 0; this.scale = 0; this.typeName = ""; this.displaySize = 0; } /** * Builds a JdbcFieldInfo from the {@link java.sql.Types} type, precision, and scale. * Use this constructor for {@link java.sql.Types#DECIMAL} and {@link java.sql.Types#NUMERIC} * types. * * @param jdbcType The {@link java.sql.Types} type. * @param precision The field's numeric precision. * @param scale The field's numeric scale. */ public JdbcFieldInfo(int jdbcType, int precision, int scale) { this.column = 0; this.jdbcType = jdbcType; this.nullability = ResultSetMetaData.columnNullableUnknown; this.precision = precision; this.scale = scale; this.typeName = ""; this.displaySize = 0; } /** * Builds a JdbcFieldInfo from the {@link java.sql.Types} type, nullability, * precision, and scale. * * @param jdbcType The {@link java.sql.Types} type. * @param nullability The nullability. Must be one of {@link ResultSetMetaData#columnNoNulls}, * {@link ResultSetMetaData#columnNullable}, or {@link * ResultSetMetaData#columnNullableUnknown}. * @param precision The field's numeric precision. * @param scale The field's numeric scale. */ public JdbcFieldInfo(int jdbcType, int nullability, int precision, int scale) { this.column = 0; this.jdbcType = jdbcType; this.nullability = nullability; this.precision = precision; this.scale = scale; this.typeName = ""; this.displaySize = 0; } /** * Builds a JdbcFieldInfo from the corresponding {@link java.sql.ResultSetMetaData} * column. * * @param rsmd The {@link java.sql.ResultSetMetaData} to get the field information from. * @param column The column to get the field information for (on a 1-based index). * @throws SQLException If the column information cannot be retrieved. * @throws NullPointerException if rsmd is null. * @throws IllegalArgumentException if column is out of bounds. */ public JdbcFieldInfo(ResultSetMetaData rsmd, int column) throws SQLException { Preconditions.checkNotNull(rsmd, "ResultSetMetaData cannot be null."); Preconditions.checkArgument( column > 0, "ResultSetMetaData columns have indices starting at 1."); Preconditions.checkArgument( column <= rsmd.getColumnCount(), "The index must be within the number of columns (1 to %s, inclusive)", rsmd.getColumnCount()); this.column = column; this.jdbcType = rsmd.getColumnType(column); this.nullability = rsmd.isNullable(column); this.precision = rsmd.getPrecision(column); this.scale = rsmd.getScale(column); this.typeName = rsmd.getColumnTypeName(column); this.displaySize = rsmd.getColumnDisplaySize(column); } /** * Builds a JdbcFieldInfo from the corresponding row from a {@link * java.sql.DatabaseMetaData#getColumns} ResultSet. * * @param rs The {@link java.sql.ResultSet} to get the field information from. * @throws SQLException If the column information cannot be retrieved. */ public JdbcFieldInfo(ResultSet rs) throws SQLException { this.column = rs.getInt("ORDINAL_POSITION"); this.jdbcType = rs.getInt("DATA_TYPE"); this.nullability = rs.getInt("NULLABLE"); this.precision = rs.getInt("COLUMN_SIZE"); this.scale = rs.getInt("DECIMAL_DIGITS"); this.typeName = rs.getString("TYPE_NAME"); this.displaySize = rs.getInt("CHAR_OCTET_LENGTH"); } /** The {@link java.sql.Types} type. */ public int getJdbcType() { return jdbcType; } /** The nullability. */ public int isNullable() { return nullability; } /** * The numeric precision, for {@link java.sql.Types#NUMERIC} and {@link java.sql.Types#DECIMAL} * types. */ public int getPrecision() { return precision; } /** * The numeric scale, for {@link java.sql.Types#NUMERIC} and {@link java.sql.Types#DECIMAL} types. */ public int getScale() { return scale; } /** The column index for query column. */ public int getColumn() { return column; } /** The type name as reported by the database. */ public String getTypeName() { return typeName; } /** The max number of characters for the column. */ public int getDisplaySize() { return displaySize; } } JdbcParameterBinder.java000066400000000000000000000126361515451615200341630ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; import java.sql.PreparedStatement; import java.sql.SQLException; import java.util.HashMap; import java.util.Map; import org.apache.arrow.adapter.jdbc.binder.ColumnBinder; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.VectorSchemaRoot; /** * A binder binds JDBC prepared statement parameters to rows of Arrow data from a VectorSchemaRoot. * *

Each row of the VectorSchemaRoot will be bound to the configured parameters of the * PreparedStatement. One row of data is bound at a time. */ public class JdbcParameterBinder { private final PreparedStatement statement; private final VectorSchemaRoot root; private final ColumnBinder[] binders; private final int[] parameterIndices; private int nextRowIndex; /** * Create a new parameter binder. * * @param statement The statement to bind parameters to. * @param root The VectorSchemaRoot to pull data from. * @param binders Column binders to translate from Arrow data to JDBC parameters, one per * parameter. * @param parameterIndices For each binder in binders, the index of the parameter to bind * to. */ private JdbcParameterBinder( final PreparedStatement statement, final VectorSchemaRoot root, final ColumnBinder[] binders, int[] parameterIndices) { Preconditions.checkArgument( binders.length == parameterIndices.length, "Number of column binders (%s) must equal number of parameter indices (%s)", binders.length, parameterIndices.length); this.statement = statement; this.root = root; this.binders = binders; this.parameterIndices = parameterIndices; this.nextRowIndex = 0; } /** * Initialize a binder with a builder. * * @param statement The statement to bind to. The binder does not maintain ownership of the * statement. * @param root The {@link VectorSchemaRoot} to pull data from. The binder does not maintain * ownership of the vector schema root. */ public static Builder builder(final PreparedStatement statement, final VectorSchemaRoot root) { return new Builder(statement, root); } /** Reset the binder (so the root can be updated with new data). */ public void reset() { nextRowIndex = 0; } /** * Bind the next row of data to the parameters of the statement. * *

After this, the application should call the desired method on the prepared statement, such * as {@link PreparedStatement#executeUpdate()}, or {@link PreparedStatement#addBatch()}. * * @return true if a row was bound, false if rows were exhausted */ public boolean next() throws SQLException { if (nextRowIndex >= root.getRowCount()) { return false; } for (int i = 0; i < parameterIndices.length; i++) { final int parameterIndex = parameterIndices[i]; binders[i].bind(statement, parameterIndex, nextRowIndex); } nextRowIndex++; return true; } /** A builder for a {@link JdbcParameterBinder}. */ public static class Builder { private final PreparedStatement statement; private final VectorSchemaRoot root; private final Map bindings; Builder(PreparedStatement statement, VectorSchemaRoot root) { this.statement = statement; this.root = root; this.bindings = new HashMap<>(); } /** 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); } return this; } /** Bind the given parameter to the given column using the default binder. */ public Builder bind(int parameterIndex, int columnIndex) { return bind(parameterIndex, ColumnBinder.forVector(root.getVector(columnIndex))); } /** Bind the given parameter using the given binder. */ public Builder bind(int parameterIndex, ColumnBinder binder) { Preconditions.checkArgument( parameterIndex > 0, "parameterIndex %d must be positive", parameterIndex); bindings.put(parameterIndex, binder); return this; } /** Build the binder. */ public JdbcParameterBinder build() { ColumnBinder[] binders = new ColumnBinder[bindings.size()]; int[] parameterIndices = new int[bindings.size()]; int index = 0; for (Map.Entry entry : bindings.entrySet()) { binders[index] = entry.getValue(); parameterIndices[index] = entry.getKey(); index++; } return new JdbcParameterBinder(statement, root, binders, parameterIndices); } } } apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrow.java000066400000000000000000000101671515451615200325700ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; import java.io.IOException; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; /** * Utility class to convert JDBC objects to columnar Arrow format objects. * *

This utility uses following data mapping to map JDBC/SQL datatype to Arrow data types. * *

CHAR --> ArrowType.Utf8 NCHAR --> ArrowType.Utf8 VARCHAR --> ArrowType.Utf8 NVARCHAR --> * ArrowType.Utf8 LONGVARCHAR --> ArrowType.Utf8 LONGNVARCHAR --> ArrowType.Utf8 NUMERIC --> * ArrowType.Decimal(precision, scale) DECIMAL --> ArrowType.Decimal(precision, scale) BIT --> * ArrowType.Bool TINYINT --> ArrowType.Int(8, signed) SMALLINT --> ArrowType.Int(16, signed) * INTEGER --> ArrowType.Int(32, signed) BIGINT --> ArrowType.Int(64, signed) REAL --> * ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE) FLOAT --> * ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE) DOUBLE --> * ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE) BINARY --> ArrowType.Binary VARBINARY --> * ArrowType.Binary LONGVARBINARY --> ArrowType.Binary DATE --> ArrowType.Date(DateUnit.MILLISECOND) * TIME --> ArrowType.Time(TimeUnit.MILLISECOND, 32) TIMESTAMP --> * ArrowType.Timestamp(TimeUnit.MILLISECOND, timezone=null) CLOB --> ArrowType.Utf8 BLOB --> * ArrowType.Binary * * @since 0.10.0 */ public class JdbcToArrow { /*----------------------------------------------------------------* | | | Partial Convert API | | | *----------------------------------------------------------------*/ /** * For the given JDBC {@link ResultSet}, fetch the data from Relational DB and convert it to Arrow * objects. Note here uses the default targetBatchSize = 1024. * * @param resultSet ResultSet to use to fetch the data from underlying database * @param allocator Memory allocator * @return Arrow Data Objects {@link ArrowVectorIterator} * @throws SQLException on error */ public static ArrowVectorIterator sqlToArrowVectorIterator( ResultSet resultSet, BufferAllocator allocator) throws SQLException, IOException { Preconditions.checkNotNull(allocator, "Memory Allocator object cannot be null"); JdbcToArrowConfig config = new JdbcToArrowConfig(allocator, JdbcToArrowUtils.getUtcCalendar()); return sqlToArrowVectorIterator(resultSet, config); } /** * For the given JDBC {@link ResultSet}, fetch the data from Relational DB and convert it to Arrow * objects. Note if not specify {@link JdbcToArrowConfig#targetBatchSize}, will use default value * 1024. * * @param resultSet ResultSet to use to fetch the data from underlying database * @param config Configuration of the conversion from JDBC to Arrow. * @return Arrow Data Objects {@link ArrowVectorIterator} * @throws SQLException on error */ public static ArrowVectorIterator sqlToArrowVectorIterator( ResultSet resultSet, JdbcToArrowConfig config) throws SQLException, IOException { Preconditions.checkNotNull(resultSet, "JDBC ResultSet object cannot be null"); Preconditions.checkNotNull(config, "The configuration cannot be null"); return ArrowVectorIterator.create(resultSet, config); } } JdbcToArrowConfig.java000066400000000000000000000375321515451615200336440ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; import java.math.RoundingMode; import java.util.Calendar; import java.util.Map; import java.util.function.Function; import org.apache.arrow.adapter.jdbc.consumer.JdbcConsumer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.types.pojo.ArrowType; /** * This class configures the JDBC-to-Arrow conversion process. * *

The allocator is used to construct the {@link org.apache.arrow.vector.VectorSchemaRoot}, and * the calendar is used to define the time zone of any {@link * org.apache.arrow.vector.types.pojo.ArrowType.Timestamp} fields that are created during the * conversion. Neither field may be null. * *

If the includeMetadata flag is set, the Arrow field metadata will contain * information from the corresponding {@link java.sql.ResultSetMetaData} that was used to create the * {@link org.apache.arrow.vector.types.pojo.FieldType} of the corresponding {@link * org.apache.arrow.vector.FieldVector}. * *

If there are any {@link java.sql.Types#ARRAY} fields in the {@link java.sql.ResultSet}, the * corresponding {@link JdbcFieldInfo} for the array's contents must be defined here. Unfortunately, * the sub-type information cannot be retrieved from all JDBC implementations (H2 for example, * returns {@link java.sql.Types#NULL} for the array sub-type), so it must be configured here. The * column index or name can be used to map to a {@link JdbcFieldInfo}, and that will be used for the * conversion. */ public final class JdbcToArrowConfig { public static final int DEFAULT_TARGET_BATCH_SIZE = 1024; public static final int NO_LIMIT_BATCH_SIZE = -1; private final Calendar calendar; private final BufferAllocator allocator; private final boolean includeMetadata; private final boolean reuseVectorSchemaRoot; private final Map arraySubTypesByColumnIndex; private final Map arraySubTypesByColumnName; private final Map explicitTypesByColumnIndex; private final Map explicitTypesByColumnName; 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 * flag only useful for {@link JdbcToArrow#sqlToArrowVectorIterator} 1) if targetBatchSize != -1, * it will convert full data into multiple vectors with valueCount no more than targetBatchSize. * 2) if targetBatchSize == -1, it will convert full data into a single vector in {@link * ArrowVectorIterator} */ private final int targetBatchSize; private final Function jdbcToArrowTypeConverter; private final JdbcConsumerFactory jdbcConsumerGetter; /** * Constructs a new configuration from the provided allocator and calendar. The allocator * is used when constructing the Arrow vectors from the ResultSet, and the calendar is * used to define Arrow Timestamp fields, and to read time-based fields from the JDBC * ResultSet. * * @param allocator The memory allocator to construct the Arrow vectors with. * @param calendar The calendar to use when constructing Timestamp fields and reading time-based * results. */ JdbcToArrowConfig(BufferAllocator allocator, Calendar calendar) { this( allocator, calendar, /* include metadata */ false, /* reuse vector schema root */ false, /* array sub-types by column index */ null, /* array sub-types by column name */ null, DEFAULT_TARGET_BATCH_SIZE, null, null); } JdbcToArrowConfig( BufferAllocator allocator, Calendar calendar, boolean includeMetadata, boolean reuseVectorSchemaRoot, Map arraySubTypesByColumnIndex, Map arraySubTypesByColumnName, int targetBatchSize, Function jdbcToArrowTypeConverter) { this( allocator, calendar, includeMetadata, reuseVectorSchemaRoot, arraySubTypesByColumnIndex, arraySubTypesByColumnName, targetBatchSize, jdbcToArrowTypeConverter, null); } /** * Constructs a new configuration from the provided allocator and calendar. The allocator * is used when constructing the Arrow vectors from the ResultSet, and the calendar is * used to define Arrow Timestamp fields, and to read time-based fields from the JDBC * ResultSet. * * @param allocator The memory allocator to construct the Arrow vectors with. * @param calendar The calendar to use when constructing Timestamp fields and reading time-based * results. * @param includeMetadata Whether to include JDBC field metadata in the Arrow Schema Field * metadata. * @param reuseVectorSchemaRoot Whether to reuse the vector schema root for each data load. * @param arraySubTypesByColumnIndex The type of the JDBC array at the column index (1-based). * @param arraySubTypesByColumnName The type of the JDBC array at the column name. * @param targetBatchSize The target batch size to be used in preallocation of the resulting * vectors. * @param jdbcToArrowTypeConverter The function that maps JDBC field type information to arrow * type. If set to null, the default mapping will be used, which is defined as: *

    *
  • CHAR --> ArrowType.Utf8 *
  • NCHAR --> ArrowType.Utf8 *
  • VARCHAR --> ArrowType.Utf8 *
  • NVARCHAR --> ArrowType.Utf8 *
  • LONGVARCHAR --> ArrowType.Utf8 *
  • LONGNVARCHAR --> ArrowType.Utf8 *
  • NUMERIC --> ArrowType.Decimal(precision, scale) *
  • DECIMAL --> ArrowType.Decimal(precision, scale) *
  • BIT --> ArrowType.Bool *
  • TINYINT --> ArrowType.Int(8, signed) *
  • SMALLINT --> ArrowType.Int(16, signed) *
  • INTEGER --> ArrowType.Int(32, signed) *
  • BIGINT --> ArrowType.Int(64, signed) *
  • REAL --> ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE) *
  • FLOAT --> ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE) *
  • DOUBLE --> ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE) *
  • BINARY --> ArrowType.Binary *
  • VARBINARY --> ArrowType.Binary *
  • LONGVARBINARY --> ArrowType.Binary *
  • DATE --> ArrowType.Date(DateUnit.DAY) *
  • TIME --> ArrowType.Time(TimeUnit.MILLISECOND, 32) *
  • TIMESTAMP --> ArrowType.Timestamp(TimeUnit.MILLISECOND, calendar timezone) *
  • CLOB --> ArrowType.Utf8 *
  • BLOB --> ArrowType.Binary *
  • ARRAY --> ArrowType.List *
  • STRUCT --> ArrowType.Struct *
  • NULL --> ArrowType.Null *
* * @param bigDecimalRoundingMode The java.math.RoundingMode to be used in coercion of a BigDecimal * from a ResultSet having a scale which does not match that of the target vector. Use null * (default value) to require strict scale matching. */ JdbcToArrowConfig( BufferAllocator allocator, Calendar calendar, boolean includeMetadata, boolean reuseVectorSchemaRoot, Map arraySubTypesByColumnIndex, Map arraySubTypesByColumnName, int targetBatchSize, Function jdbcToArrowTypeConverter, RoundingMode bigDecimalRoundingMode) { this( allocator, calendar, includeMetadata, reuseVectorSchemaRoot, arraySubTypesByColumnIndex, arraySubTypesByColumnName, targetBatchSize, jdbcToArrowTypeConverter, null, null, null, null, bigDecimalRoundingMode); } JdbcToArrowConfig( BufferAllocator allocator, Calendar calendar, boolean includeMetadata, boolean reuseVectorSchemaRoot, Map arraySubTypesByColumnIndex, Map arraySubTypesByColumnName, int targetBatchSize, Function jdbcToArrowTypeConverter, Map explicitTypesByColumnIndex, Map explicitTypesByColumnName, Map schemaMetadata, Map> columnMetadataByColumnIndex, RoundingMode bigDecimalRoundingMode) { this( allocator, calendar, includeMetadata, reuseVectorSchemaRoot, arraySubTypesByColumnIndex, arraySubTypesByColumnName, targetBatchSize, jdbcToArrowTypeConverter, null, explicitTypesByColumnIndex, explicitTypesByColumnName, schemaMetadata, columnMetadataByColumnIndex, bigDecimalRoundingMode); } JdbcToArrowConfig( BufferAllocator allocator, Calendar calendar, boolean includeMetadata, boolean reuseVectorSchemaRoot, Map arraySubTypesByColumnIndex, Map arraySubTypesByColumnName, int targetBatchSize, Function jdbcToArrowTypeConverter, JdbcConsumerFactory jdbcConsumerGetter, Map explicitTypesByColumnIndex, Map explicitTypesByColumnName, Map schemaMetadata, Map> columnMetadataByColumnIndex, RoundingMode bigDecimalRoundingMode) { Preconditions.checkNotNull(allocator, "Memory allocator cannot be null"); this.allocator = allocator; this.calendar = calendar; this.includeMetadata = includeMetadata; this.reuseVectorSchemaRoot = reuseVectorSchemaRoot; this.arraySubTypesByColumnIndex = arraySubTypesByColumnIndex; this.arraySubTypesByColumnName = arraySubTypesByColumnName; this.targetBatchSize = targetBatchSize; this.explicitTypesByColumnIndex = explicitTypesByColumnIndex; this.explicitTypesByColumnName = explicitTypesByColumnName; this.schemaMetadata = schemaMetadata; this.columnMetadataByColumnIndex = columnMetadataByColumnIndex; this.bigDecimalRoundingMode = bigDecimalRoundingMode; // set up type converter this.jdbcToArrowTypeConverter = jdbcToArrowTypeConverter != null ? jdbcToArrowTypeConverter : (jdbcFieldInfo) -> JdbcToArrowUtils.getArrowTypeFromJdbcType(jdbcFieldInfo, calendar); this.jdbcConsumerGetter = jdbcConsumerGetter != null ? jdbcConsumerGetter : JdbcToArrowUtils::getConsumer; } /** * The calendar to use when defining Arrow Timestamp fields and retrieving {@link java.sql.Date}, * {@link java.sql.Time}, or {@link java.sql.Timestamp} data types from the {@link * java.sql.ResultSet}, or null if not converting. * * @return the calendar. */ public Calendar getCalendar() { return calendar; } /** * The Arrow memory allocator. * * @return the allocator. */ public BufferAllocator getAllocator() { return allocator; } /** * Whether to include JDBC ResultSet field metadata in the Arrow Schema field metadata. * * @return true to include field metadata, false to exclude it. */ public boolean shouldIncludeMetadata() { return includeMetadata; } /** Get the target batch size for partial read. */ public int getTargetBatchSize() { return targetBatchSize; } /** Get whether it is allowed to reuse the vector schema root. */ public boolean isReuseVectorSchemaRoot() { return reuseVectorSchemaRoot; } /** Gets the mapping between JDBC type information to Arrow type. */ public Function getJdbcToArrowTypeConverter() { return jdbcToArrowTypeConverter; } /** Gets the JDBC consumer getter. */ public JdbcConsumerFactory getJdbcConsumerGetter() { return jdbcConsumerGetter; } /** * Returns the array sub-type {@link JdbcFieldInfo} defined for the provided column index. * * @param index The {@link java.sql.ResultSetMetaData} column index of an {@link * java.sql.Types#ARRAY} type. * @return The {@link JdbcFieldInfo} for that array's sub-type, or null if not * defined. */ public JdbcFieldInfo getArraySubTypeByColumnIndex(int index) { if (arraySubTypesByColumnIndex == null) { return null; } else { return arraySubTypesByColumnIndex.get(index); } } /** * Returns the array sub-type {@link JdbcFieldInfo} defined for the provided column name. * * @param name The {@link java.sql.ResultSetMetaData} column name of an {@link * java.sql.Types#ARRAY} type. * @return The {@link JdbcFieldInfo} for that array's sub-type, or null if not * defined. */ public JdbcFieldInfo getArraySubTypeByColumnName(String name) { if (arraySubTypesByColumnName == null) { return null; } else { return arraySubTypesByColumnName.get(name); } } /** * Returns the type {@link JdbcFieldInfo} explicitly defined for the provided column index. * * @param index The {@link java.sql.ResultSetMetaData} column index to evaluate for explicit type * mapping. * @return The {@link JdbcFieldInfo} defined for the column, or null if not defined. */ public JdbcFieldInfo getExplicitTypeByColumnIndex(int index) { if (explicitTypesByColumnIndex == null) { return null; } else { return explicitTypesByColumnIndex.get(index); } } /** * Returns the type {@link JdbcFieldInfo} explicitly defined for the provided column name. * * @param name The {@link java.sql.ResultSetMetaData} column name to evaluate for explicit type * mapping. * @return The {@link JdbcFieldInfo} defined for the column, or null if not defined. */ public JdbcFieldInfo getExplicitTypeByColumnName(String name) { if (explicitTypesByColumnName == null) { return null; } else { return explicitTypesByColumnName.get(name); } } /** Return schema level metadata or null if not provided. */ public Map getSchemaMetadata() { return schemaMetadata; } /** Return metadata from columnIndex->meta map on per field basis or null if not provided. */ public Map> getColumnMetadataByColumnIndex() { return columnMetadataByColumnIndex; } public RoundingMode getBigDecimalRoundingMode() { return bigDecimalRoundingMode; } /** Interface for a function that gets a JDBC consumer for the given values. */ @FunctionalInterface public interface JdbcConsumerFactory { JdbcConsumer apply( ArrowType arrowType, int columnIndex, boolean nullable, FieldVector vector, JdbcToArrowConfig config); } } JdbcToArrowConfigBuilder.java000066400000000000000000000256021515451615200351460ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; import static org.apache.arrow.adapter.jdbc.JdbcToArrowConfig.DEFAULT_TARGET_BATCH_SIZE; import java.math.RoundingMode; import java.util.Calendar; import java.util.Map; import java.util.function.Function; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.types.pojo.ArrowType; /** This class builds {@link JdbcToArrowConfig}s. */ public class JdbcToArrowConfigBuilder { private Calendar calendar; private BufferAllocator allocator; private boolean includeMetadata; private boolean reuseVectorSchemaRoot; private Map arraySubTypesByColumnIndex; private Map arraySubTypesByColumnName; private Map explicitTypesByColumnIndex; private Map explicitTypesByColumnName; private Map schemaMetadata; private Map> columnMetadataByColumnIndex; private int targetBatchSize; private Function jdbcToArrowTypeConverter; private JdbcToArrowConfig.JdbcConsumerFactory jdbcConsumerGetter; private RoundingMode bigDecimalRoundingMode; /** * Default constructor for the JdbcToArrowConfigBuilder}. Use the setter methods for * the allocator and calendar; the allocator must be set. Otherwise, {@link #build()} will throw a * {@link NullPointerException}. */ public JdbcToArrowConfigBuilder() { this.allocator = null; this.calendar = null; this.includeMetadata = false; this.reuseVectorSchemaRoot = false; this.arraySubTypesByColumnIndex = null; this.arraySubTypesByColumnName = null; this.explicitTypesByColumnIndex = null; this.explicitTypesByColumnName = null; this.schemaMetadata = null; this.columnMetadataByColumnIndex = null; this.bigDecimalRoundingMode = null; } /** * Constructor for the JdbcToArrowConfigBuilder. The allocator is required, and a * {@link NullPointerException} will be thrown if it is null. * *

The allocator is used to construct Arrow vectors from the JDBC ResultSet. The calendar is * used to determine the time zone of {@link java.sql.Timestamp} fields and convert {@link * java.sql.Date}, {@link java.sql.Time}, and {@link java.sql.Timestamp} fields to a single, * common time zone when reading from the result set. * * @param allocator The Arrow Vector memory allocator. * @param calendar The calendar to use when constructing timestamp fields. */ public JdbcToArrowConfigBuilder(BufferAllocator allocator, Calendar calendar) { this(); Preconditions.checkNotNull(allocator, "Memory allocator cannot be null"); this.allocator = allocator; this.calendar = calendar; this.includeMetadata = false; this.reuseVectorSchemaRoot = false; this.targetBatchSize = DEFAULT_TARGET_BATCH_SIZE; } /** * Constructor for the JdbcToArrowConfigBuilder. Both the allocator and calendar are * required. A {@link NullPointerException} will be thrown if either of those arguments is * null. * *

The allocator is used to construct Arrow vectors from the JDBC ResultSet. The calendar is * used to determine the time zone of {@link java.sql.Timestamp} fields and convert {@link * java.sql.Date}, {@link java.sql.Time}, and {@link java.sql.Timestamp} fields to a single, * common time zone when reading from the result set. * *

The includeMetadata argument, if true will cause various * information about each database field to be added to the Vector Schema's field metadata. * * @param allocator The Arrow Vector memory allocator. * @param calendar The calendar to use when constructing timestamp fields. */ public JdbcToArrowConfigBuilder( BufferAllocator allocator, Calendar calendar, boolean includeMetadata) { this(allocator, calendar); this.includeMetadata = includeMetadata; } /** * Sets the memory allocator to use when constructing the Arrow vectors from the ResultSet. * * @param allocator the allocator to set. * @exception NullPointerException if allocator is null. */ public JdbcToArrowConfigBuilder setAllocator(BufferAllocator allocator) { Preconditions.checkNotNull(allocator, "Memory allocator cannot be null"); this.allocator = allocator; return this; } /** * Sets the {@link Calendar} to use when constructing timestamp fields in the Arrow schema, and * reading time-based fields from the JDBC ResultSet. * * @param calendar the calendar to set. */ public JdbcToArrowConfigBuilder setCalendar(Calendar calendar) { this.calendar = calendar; return this; } /** * Sets whether to include JDBC ResultSet field metadata in the Arrow Schema field metadata. * * @param includeMetadata Whether to include or exclude JDBC metadata in the Arrow Schema field * metadata. * @return This instance of the JdbcToArrowConfig, for chaining. */ public JdbcToArrowConfigBuilder setIncludeMetadata(boolean includeMetadata) { this.includeMetadata = includeMetadata; return this; } /** * Sets the mapping of column-index-to-{@link JdbcFieldInfo} used for columns of type {@link * java.sql.Types#ARRAY}. The column index is 1-based, to match the JDBC column index. * * @param map The mapping. * @return This instance of the JdbcToArrowConfig, for chaining. */ public JdbcToArrowConfigBuilder setArraySubTypeByColumnIndexMap(Map map) { this.arraySubTypesByColumnIndex = map; return this; } /** * Sets the mapping of column-name-to-{@link JdbcFieldInfo} used for columns of type {@link * java.sql.Types#ARRAY}. * * @param map The mapping. * @return This instance of the JdbcToArrowConfig, for chaining. */ public JdbcToArrowConfigBuilder setArraySubTypeByColumnNameMap(Map map) { this.arraySubTypesByColumnName = map; return this; } /** * Sets the mapping of column-index-to-{@link JdbcFieldInfo} used for column types. * *

This can be useful to override type information from JDBC drivers that provide incomplete * type info, e.g. DECIMAL with precision = scale = 0. * *

The column index is 1-based, to match the JDBC column index. * * @param map The mapping. */ public JdbcToArrowConfigBuilder setExplicitTypesByColumnIndex(Map map) { this.explicitTypesByColumnIndex = map; return this; } /** * Sets the mapping of column-name-to-{@link JdbcFieldInfo} used for column types. * *

This can be useful to override type information from JDBC drivers that provide incomplete * type info, e.g. DECIMAL with precision = scale = 0. * * @param map The mapping. */ public JdbcToArrowConfigBuilder setExplicitTypesByColumnName(Map map) { this.explicitTypesByColumnName = map; return this; } /** * Set the target number of rows to convert at once. * *

Use {@link JdbcToArrowConfig#NO_LIMIT_BATCH_SIZE} to read all rows at once. */ public JdbcToArrowConfigBuilder setTargetBatchSize(int targetBatchSize) { this.targetBatchSize = targetBatchSize; return this; } /** * Set the function used to convert JDBC types to Arrow types. * *

Defaults to wrapping {@link JdbcToArrowUtils#getArrowTypeFromJdbcType(JdbcFieldInfo, * Calendar)}. * * @see JdbcToArrowUtils#reportUnsupportedTypesAsUnknown(Function) */ public JdbcToArrowConfigBuilder setJdbcToArrowTypeConverter( Function jdbcToArrowTypeConverter) { this.jdbcToArrowTypeConverter = jdbcToArrowTypeConverter; return this; } /** * Set the function used to get a JDBC consumer for a given type. * *

Defaults to wrapping {@link JdbcToArrowUtils#getConsumer(ArrowType, Integer, Boolean, * FieldVector, JdbcToArrowConfig)}. */ public JdbcToArrowConfigBuilder setJdbcConsumerGetter( JdbcToArrowConfig.JdbcConsumerFactory jdbcConsumerGetter) { this.jdbcConsumerGetter = jdbcConsumerGetter; return this; } /** * Set whether to use the same {@link org.apache.arrow.vector.VectorSchemaRoot} instance on each * iteration, or to allocate a new one. */ public JdbcToArrowConfigBuilder setReuseVectorSchemaRoot(boolean reuseVectorSchemaRoot) { this.reuseVectorSchemaRoot = reuseVectorSchemaRoot; return this; } /** Set metadata for schema. */ public JdbcToArrowConfigBuilder setSchemaMetadata(Map schemaMetadata) { this.schemaMetadata = schemaMetadata; return this; } /** Set metadata from columnIndex->meta map on per field basis. */ public JdbcToArrowConfigBuilder setColumnMetadataByColumnIndex( Map> columnMetadataByColumnIndex) { this.columnMetadataByColumnIndex = columnMetadataByColumnIndex; return this; } /** * Set the rounding mode used when the scale of the actual value does not match the declared * scale. * *

By default, an error is raised in such cases. */ public JdbcToArrowConfigBuilder setBigDecimalRoundingMode(RoundingMode bigDecimalRoundingMode) { this.bigDecimalRoundingMode = bigDecimalRoundingMode; return this; } /** * This builds the {@link JdbcToArrowConfig} from the provided {@link BufferAllocator} and {@link * Calendar}. * * @return The built {@link JdbcToArrowConfig} * @throws NullPointerException if either the allocator or calendar was not set. */ public JdbcToArrowConfig build() { return new JdbcToArrowConfig( allocator, calendar, includeMetadata, reuseVectorSchemaRoot, arraySubTypesByColumnIndex, arraySubTypesByColumnName, targetBatchSize, jdbcToArrowTypeConverter, jdbcConsumerGetter, explicitTypesByColumnIndex, explicitTypesByColumnName, schemaMetadata, columnMetadataByColumnIndex, bigDecimalRoundingMode); } } JdbcToArrowUtils.java000066400000000000000000000574341515451615200335420ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; import static org.apache.arrow.vector.types.FloatingPointPrecision.DOUBLE; import static org.apache.arrow.vector.types.FloatingPointPrecision.SINGLE; import static org.apache.arrow.vector.types.Types.MinorType; import java.io.IOException; import java.math.RoundingMode; import java.sql.Date; import java.sql.ParameterMetaData; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Time; import java.sql.Timestamp; import java.sql.Types; import java.util.ArrayList; import java.util.Arrays; import java.util.Calendar; import java.util.HashMap; import java.util.List; import java.util.Locale; import java.util.Map; import java.util.TimeZone; import java.util.function.Function; import org.apache.arrow.adapter.jdbc.consumer.ArrayConsumer; import org.apache.arrow.adapter.jdbc.consumer.BigIntConsumer; import org.apache.arrow.adapter.jdbc.consumer.BinaryConsumer; import org.apache.arrow.adapter.jdbc.consumer.BitConsumer; import org.apache.arrow.adapter.jdbc.consumer.CompositeJdbcConsumer; import org.apache.arrow.adapter.jdbc.consumer.DateConsumer; import org.apache.arrow.adapter.jdbc.consumer.Decimal256Consumer; import org.apache.arrow.adapter.jdbc.consumer.DecimalConsumer; import org.apache.arrow.adapter.jdbc.consumer.DoubleConsumer; import org.apache.arrow.adapter.jdbc.consumer.FloatConsumer; import org.apache.arrow.adapter.jdbc.consumer.IntConsumer; import org.apache.arrow.adapter.jdbc.consumer.JdbcConsumer; import org.apache.arrow.adapter.jdbc.consumer.MapConsumer; import org.apache.arrow.adapter.jdbc.consumer.NullConsumer; import org.apache.arrow.adapter.jdbc.consumer.SmallIntConsumer; import org.apache.arrow.adapter.jdbc.consumer.TimeConsumer; import org.apache.arrow.adapter.jdbc.consumer.TimestampConsumer; import org.apache.arrow.adapter.jdbc.consumer.TimestampTZConsumer; import org.apache.arrow.adapter.jdbc.consumer.TinyIntConsumer; import org.apache.arrow.adapter.jdbc.consumer.VarCharConsumer; import org.apache.arrow.adapter.jdbc.consumer.exceptions.JdbcConsumerException; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.Preconditions; 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.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.TimeMilliVector; import org.apache.arrow.vector.TimeStampMilliTZVector; import org.apache.arrow.vector.TimeStampMilliVector; 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.extension.OpaqueType; import org.apache.arrow.vector.types.DateUnit; 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.FieldType; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.ValueVectorUtility; /** * Class that does most of the work to convert JDBC ResultSet data into Arrow columnar format Vector * objects. * * @since 0.10.0 */ public class JdbcToArrowUtils { private static final int JDBC_ARRAY_VALUE_COLUMN = 2; /** Returns the instance of a {java.util.Calendar} with the UTC time zone and root locale. */ public static Calendar getUtcCalendar() { return Calendar.getInstance(TimeZone.getTimeZone("UTC"), Locale.ROOT); } /** * Create Arrow {@link Schema} object for the given JDBC {@link ResultSetMetaData}. * * @param rsmd The ResultSetMetaData containing the results, to read the JDBC metadata from. * @param calendar The calendar to use the time zone field of, to construct Timestamp fields from. * @return {@link Schema} * @throws SQLException on error */ public static Schema jdbcToArrowSchema(ResultSetMetaData rsmd, Calendar calendar) throws SQLException { Preconditions.checkNotNull(calendar, "Calendar object can't be null"); return jdbcToArrowSchema(rsmd, new JdbcToArrowConfig(new RootAllocator(0), calendar)); } /** * Create Arrow {@link Schema} object for the given JDBC {@link ResultSetMetaData}. * * @param parameterMetaData The ResultSetMetaData containing the results, to read the JDBC * metadata from. * @param calendar The calendar to use the time zone field of, to construct Timestamp fields from. * @return {@link Schema} * @throws SQLException on error */ public static Schema jdbcToArrowSchema( final ParameterMetaData parameterMetaData, final Calendar calendar) throws SQLException { Preconditions.checkNotNull(calendar, "Calendar object can't be null"); Preconditions.checkNotNull(parameterMetaData); final List parameterFields = new ArrayList<>(parameterMetaData.getParameterCount()); for (int parameterCounter = 1; parameterCounter <= parameterMetaData.getParameterCount(); parameterCounter++) { final int jdbcDataType = parameterMetaData.getParameterType(parameterCounter); final int jdbcIsNullable = parameterMetaData.isNullable(parameterCounter); final boolean arrowIsNullable = jdbcIsNullable != ParameterMetaData.parameterNoNulls; final int precision = parameterMetaData.getPrecision(parameterCounter); 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); parameterFields.add(new Field(null, fieldType, null)); } return new Schema(parameterFields); } /** * Converts the provided JDBC type to its respective {@link ArrowType} counterpart. * * @param fieldInfo the {@link JdbcFieldInfo} with information about the original JDBC type. * @param calendar the {@link Calendar} to use for datetime data types. * @return a new {@link ArrowType}. */ public static ArrowType getArrowTypeFromJdbcType( final JdbcFieldInfo fieldInfo, final Calendar calendar) { switch (fieldInfo.getJdbcType()) { case Types.BOOLEAN: case Types.BIT: return new ArrowType.Bool(); case Types.TINYINT: return new ArrowType.Int(8, true); case Types.SMALLINT: return new ArrowType.Int(16, true); case Types.INTEGER: return new ArrowType.Int(32, true); case Types.BIGINT: return new ArrowType.Int(64, true); case Types.NUMERIC: case Types.DECIMAL: int precision = fieldInfo.getPrecision(); int scale = fieldInfo.getScale(); if (precision > 38) { return new ArrowType.Decimal(precision, scale, 256); } else { return new ArrowType.Decimal(precision, scale, 128); } case Types.REAL: case Types.FLOAT: return new ArrowType.FloatingPoint(SINGLE); case Types.DOUBLE: return new ArrowType.FloatingPoint(DOUBLE); case Types.CHAR: case Types.NCHAR: case Types.VARCHAR: case Types.NVARCHAR: case Types.LONGVARCHAR: case Types.LONGNVARCHAR: case Types.CLOB: return new ArrowType.Utf8(); case Types.DATE: return new ArrowType.Date(DateUnit.DAY); case Types.TIME: return new ArrowType.Time(TimeUnit.MILLISECOND, 32); case Types.TIMESTAMP: final String timezone; if (calendar != null) { timezone = calendar.getTimeZone().getID(); } else { timezone = null; } return new ArrowType.Timestamp(TimeUnit.MILLISECOND, timezone); case Types.BINARY: case Types.VARBINARY: case Types.LONGVARBINARY: case Types.BLOB: return new ArrowType.Binary(); case Types.ARRAY: return new ArrowType.List(); case Types.NULL: return new ArrowType.Null(); case Types.STRUCT: return new ArrowType.Struct(); default: throw new UnsupportedOperationException("Unmapped JDBC type: " + fieldInfo.getJdbcType()); } } /** * Wrap a JDBC to Arrow type converter such that {@link UnsupportedOperationException} becomes * {@link OpaqueType}. * * @param typeConverter The type converter to wrap. * @param vendorName The database name to report as the Opaque type's vendor name. */ public static Function reportUnsupportedTypesAsOpaque( Function typeConverter, String vendorName) { return (final JdbcFieldInfo fieldInfo) -> { try { return typeConverter.apply(fieldInfo); } catch (UnsupportedOperationException e) { return new OpaqueType(MinorType.NULL.getType(), fieldInfo.getTypeName(), vendorName); } }; } /** * Create Arrow {@link Schema} object for the given JDBC {@link java.sql.ResultSetMetaData}. * *

If {@link JdbcToArrowConfig#shouldIncludeMetadata()} returns true, the * following fields will be added to the {@link FieldType#getMetadata()}: * *

    *
  • {@link Constants#SQL_CATALOG_NAME_KEY} representing {@link * ResultSetMetaData#getCatalogName(int)} *
  • {@link Constants#SQL_TABLE_NAME_KEY} representing {@link * ResultSetMetaData#getTableName(int)} *
  • {@link Constants#SQL_COLUMN_NAME_KEY} representing {@link * ResultSetMetaData#getColumnLabel(int)} *
  • {@link Constants#SQL_TYPE_KEY} representing {@link * ResultSetMetaData#getColumnTypeName(int)} *
* *

If any columns are of type {@link java.sql.Types#ARRAY}, the configuration object will be * used to look up the array sub-type field. The {@link * JdbcToArrowConfig#getArraySubTypeByColumnIndex(int)} method will be checked first, followed by * the {@link JdbcToArrowConfig#getArraySubTypeByColumnName(String)} method. * * @param rsmd The ResultSetMetaData containing the results, to read the JDBC metadata from. * @param config The configuration to use when constructing the schema. * @return {@link Schema} * @throws SQLException on error * @throws IllegalArgumentException if rsmd contains an {@link java.sql.Types#ARRAY} * but the config does not have a sub-type definition for it. */ public static Schema jdbcToArrowSchema(ResultSetMetaData rsmd, JdbcToArrowConfig config) throws SQLException { Preconditions.checkNotNull(rsmd, "JDBC ResultSetMetaData object can't be null"); Preconditions.checkNotNull(config, "The configuration object must not be null"); List fields = new ArrayList<>(); int columnCount = rsmd.getColumnCount(); for (int i = 1; i <= columnCount; i++) { final String columnName = rsmd.getColumnLabel(i); final Map columnMetadata = config.getColumnMetadataByColumnIndex() != null ? config.getColumnMetadataByColumnIndex().get(i) : null; final Map metadata; if (config.shouldIncludeMetadata()) { metadata = new HashMap<>(); metadata.put(Constants.SQL_CATALOG_NAME_KEY, rsmd.getCatalogName(i)); metadata.put(Constants.SQL_SCHEMA_NAME_KEY, rsmd.getSchemaName(i)); metadata.put(Constants.SQL_TABLE_NAME_KEY, rsmd.getTableName(i)); metadata.put(Constants.SQL_COLUMN_NAME_KEY, columnName); metadata.put(Constants.SQL_TYPE_KEY, rsmd.getColumnTypeName(i)); if (columnMetadata != null && !columnMetadata.isEmpty()) { metadata.putAll(columnMetadata); } } else { if (columnMetadata != null && !columnMetadata.isEmpty()) { metadata = columnMetadata; } else { metadata = null; } } final JdbcFieldInfo columnFieldInfo = getJdbcFieldInfoForColumn(rsmd, i, config); final ArrowType arrowType = config.getJdbcToArrowTypeConverter().apply(columnFieldInfo); if (arrowType != null) { final FieldType fieldType = new FieldType( isColumnNullable(rsmd, i, columnFieldInfo), arrowType, /* dictionary encoding */ null, metadata); List children = null; if (arrowType.getTypeID() == ArrowType.List.TYPE_TYPE) { final JdbcFieldInfo arrayFieldInfo = getJdbcFieldInfoForArraySubType(rsmd, i, config); if (arrayFieldInfo == null) { throw new IllegalArgumentException( "Configuration does not provide a mapping for array column " + i); } children = new ArrayList(); final ArrowType childType = config.getJdbcToArrowTypeConverter().apply(arrayFieldInfo); children.add(new Field("child", FieldType.nullable(childType), null)); } else if (arrowType.getTypeID() == ArrowType.ArrowTypeID.Map) { FieldType mapType = new FieldType(false, ArrowType.Struct.INSTANCE, null, null); FieldType keyType = new FieldType(false, new ArrowType.Utf8(), null, null); FieldType valueType = new FieldType(false, new ArrowType.Utf8(), null, null); children = new ArrayList<>(); children.add( new Field( "child", mapType, Arrays.asList( new Field(MapVector.KEY_NAME, keyType, null), new Field(MapVector.VALUE_NAME, valueType, null)))); } fields.add(new Field(columnName, fieldType, children)); } } return new Schema(fields, config.getSchemaMetadata()); } static JdbcFieldInfo getJdbcFieldInfoForColumn( ResultSetMetaData rsmd, int arrayColumn, JdbcToArrowConfig config) throws SQLException { Preconditions.checkNotNull(rsmd, "ResultSet MetaData object cannot be null"); Preconditions.checkNotNull(config, "Configuration must not be null"); Preconditions.checkArgument( arrayColumn > 0, "ResultSetMetaData columns start with 1; column cannot be less than 1"); Preconditions.checkArgument( arrayColumn <= rsmd.getColumnCount(), "Column number cannot be more than the number of columns"); JdbcFieldInfo fieldInfo = config.getExplicitTypeByColumnIndex(arrayColumn); if (fieldInfo == null) { fieldInfo = config.getExplicitTypeByColumnName(rsmd.getColumnLabel(arrayColumn)); } if (fieldInfo != null) { return fieldInfo; } return new JdbcFieldInfo(rsmd, arrayColumn); } /* Uses the configuration to determine what the array sub-type JdbcFieldInfo is. * If no sub-type can be found, returns null. */ private static JdbcFieldInfo getJdbcFieldInfoForArraySubType( ResultSetMetaData rsmd, int arrayColumn, JdbcToArrowConfig config) throws SQLException { Preconditions.checkNotNull(rsmd, "ResultSet MetaData object cannot be null"); Preconditions.checkNotNull(config, "Configuration must not be null"); Preconditions.checkArgument( arrayColumn > 0, "ResultSetMetaData columns start with 1; column cannot be less than 1"); Preconditions.checkArgument( arrayColumn <= rsmd.getColumnCount(), "Column number cannot be more than the number of columns"); JdbcFieldInfo fieldInfo = config.getArraySubTypeByColumnIndex(arrayColumn); if (fieldInfo == null) { fieldInfo = config.getArraySubTypeByColumnName(rsmd.getColumnLabel(arrayColumn)); } return fieldInfo; } /** * Iterate the given JDBC {@link ResultSet} object to fetch the data and transpose it to populate * the given Arrow Vector objects. * * @param rs ResultSet to use to fetch the data from underlying database * @param root Arrow {@link VectorSchemaRoot} object to populate * @param calendar The calendar to use when reading {@link Date}, {@link Time}, or {@link * Timestamp} data types from the {@link ResultSet}, or null if not converting. * @throws SQLException on error */ public static void jdbcToArrowVectors(ResultSet rs, VectorSchemaRoot root, Calendar calendar) throws SQLException, IOException { Preconditions.checkNotNull(calendar, "Calendar object can't be null"); jdbcToArrowVectors(rs, root, new JdbcToArrowConfig(new RootAllocator(0), calendar)); } static boolean isColumnNullable( ResultSetMetaData resultSetMetadata, int index, JdbcFieldInfo info) throws SQLException { int nullableValue; if (info != null && info.isNullable() != ResultSetMetaData.columnNullableUnknown) { nullableValue = info.isNullable(); } else { nullableValue = resultSetMetadata.isNullable(index); } return nullableValue == ResultSetMetaData.columnNullable || nullableValue == ResultSetMetaData.columnNullableUnknown; } /** * Iterate the given JDBC {@link ResultSet} object to fetch the data and transpose it to populate * the given Arrow Vector objects. * * @param rs ResultSet to use to fetch the data from underlying database * @param root Arrow {@link VectorSchemaRoot} object to populate * @param config The configuration to use when reading the data. * @throws SQLException on error * @throws JdbcConsumerException on error from VectorConsumer */ public static void jdbcToArrowVectors( ResultSet rs, VectorSchemaRoot root, JdbcToArrowConfig config) throws SQLException, IOException { ResultSetMetaData rsmd = rs.getMetaData(); int columnCount = rsmd.getColumnCount(); JdbcConsumer[] consumers = new JdbcConsumer[columnCount]; for (int i = 1; i <= columnCount; i++) { FieldVector vector = root.getVector(rsmd.getColumnLabel(i)); final JdbcFieldInfo columnFieldInfo = getJdbcFieldInfoForColumn(rsmd, i, config); consumers[i - 1] = getConsumer( vector.getField().getType(), i, isColumnNullable(rsmd, i, columnFieldInfo), vector, config); } CompositeJdbcConsumer compositeConsumer = null; // Only clean resources when occurs error, // vectors within consumers are useful and users are responsible for its close. try { compositeConsumer = new CompositeJdbcConsumer(consumers); int readRowCount = 0; if (config.getTargetBatchSize() == JdbcToArrowConfig.NO_LIMIT_BATCH_SIZE) { while (rs.next()) { ValueVectorUtility.ensureCapacity(root, readRowCount + 1); compositeConsumer.consume(rs); readRowCount++; } } else { while (readRowCount < config.getTargetBatchSize() && rs.next()) { compositeConsumer.consume(rs); readRowCount++; } } root.setRowCount(readRowCount); } catch (Exception e) { // error occurs and clean up resources. if (compositeConsumer != null) { compositeConsumer.close(); } throw e; } } /** * Default function used for JdbcConsumerFactory. This function gets a JdbcConsumer for the given * column based on the Arrow type and provided vector. * * @param arrowType Arrow type for the column. * @param columnIndex Column index to fetch from the ResultSet * @param nullable Whether the value is nullable or not * @param vector Vector to store the consumed value * @param config Associated JdbcToArrowConfig, used mainly for the Calendar. * @return {@link JdbcConsumer} */ public static JdbcConsumer getConsumer( ArrowType arrowType, int columnIndex, boolean nullable, FieldVector vector, JdbcToArrowConfig config) { final Calendar calendar = config.getCalendar(); switch (arrowType.getTypeID()) { case Bool: return BitConsumer.createConsumer((BitVector) vector, columnIndex, nullable); case Int: switch (((ArrowType.Int) arrowType).getBitWidth()) { case 8: return TinyIntConsumer.createConsumer((TinyIntVector) vector, columnIndex, nullable); case 16: return SmallIntConsumer.createConsumer((SmallIntVector) vector, columnIndex, nullable); case 32: return IntConsumer.createConsumer((IntVector) vector, columnIndex, nullable); case 64: return BigIntConsumer.createConsumer((BigIntVector) vector, columnIndex, nullable); default: return null; } case Decimal: final RoundingMode bigDecimalRoundingMode = config.getBigDecimalRoundingMode(); if (((ArrowType.Decimal) arrowType).getBitWidth() == 256) { return Decimal256Consumer.createConsumer( (Decimal256Vector) vector, columnIndex, nullable, bigDecimalRoundingMode); } else { return DecimalConsumer.createConsumer( (DecimalVector) vector, columnIndex, nullable, bigDecimalRoundingMode); } case FloatingPoint: switch (((ArrowType.FloatingPoint) arrowType).getPrecision()) { case SINGLE: return FloatConsumer.createConsumer((Float4Vector) vector, columnIndex, nullable); case DOUBLE: return DoubleConsumer.createConsumer((Float8Vector) vector, columnIndex, nullable); default: return null; } case Utf8: case LargeUtf8: return VarCharConsumer.createConsumer((VarCharVector) vector, columnIndex, nullable); case Binary: case LargeBinary: return BinaryConsumer.createConsumer((VarBinaryVector) vector, columnIndex, nullable); case Date: return DateConsumer.createConsumer((DateDayVector) vector, columnIndex, nullable, calendar); case Time: return TimeConsumer.createConsumer( (TimeMilliVector) vector, columnIndex, nullable, calendar); case Timestamp: if (config.getCalendar() == null) { return TimestampConsumer.createConsumer( (TimeStampMilliVector) vector, columnIndex, nullable); } else { return TimestampTZConsumer.createConsumer( (TimeStampMilliTZVector) vector, columnIndex, nullable, calendar); } case List: FieldVector childVector = ((ListVector) vector).getDataVector(); JdbcConsumer delegate = getConsumer( childVector.getField().getType(), JDBC_ARRAY_VALUE_COLUMN, childVector.getField().isNullable(), childVector, config); return ArrayConsumer.createConsumer((ListVector) vector, delegate, columnIndex, nullable); case Map: return MapConsumer.createConsumer((MapVector) vector, columnIndex, nullable); case Null: return new NullConsumer((NullVector) vector); default: // no-op, shouldn't get here throw new UnsupportedOperationException("No consumer for Arrow type: " + arrowType); } } } apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/000077500000000000000000000000001515451615200310035ustar00rootroot00000000000000BaseColumnBinder.java000066400000000000000000000025571515451615200347540ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import org.apache.arrow.vector.FieldVector; /** * Base class for ColumnBinder implementations. * * @param The concrete FieldVector subtype. */ public abstract class BaseColumnBinder implements ColumnBinder { protected final V vector; protected final int jdbcType; public BaseColumnBinder(V vector, int jdbcType) { this.vector = vector; this.jdbcType = jdbcType; } @Override public int getJdbcType() { return jdbcType; } @Override public V getVector() { return vector; } } BigIntBinder.java000066400000000000000000000030031515451615200340630ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.PreparedStatement; import java.sql.SQLException; import java.sql.Types; import org.apache.arrow.vector.BigIntVector; /** A column binder for 8-bit integers. */ public class BigIntBinder extends BaseColumnBinder { public BigIntBinder(BigIntVector vector) { this(vector, Types.BIGINT); } public BigIntBinder(BigIntVector vector, int jdbcType) { super(vector, jdbcType); } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { final long value = vector.getDataBuffer().getLong((long) rowIndex * BigIntVector.TYPE_WIDTH); statement.setLong(parameterIndex, value); } } BitBinder.java000066400000000000000000000031471515451615200334360ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.PreparedStatement; import java.sql.SQLException; import java.sql.Types; import org.apache.arrow.vector.BitVector; /** A column binder for booleans. */ public class BitBinder extends BaseColumnBinder { public BitBinder(BitVector vector) { this(vector, Types.BOOLEAN); } public BitBinder(BitVector vector, int jdbcType) { super(vector, jdbcType); } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { // See BitVector#getBit final int byteIndex = rowIndex >> 3; final byte b = vector.getDataBuffer().getByte(byteIndex); final int bitIndex = rowIndex & 7; final int value = (b >> bitIndex) & 0x01; statement.setBoolean(parameterIndex, value != 0); } } ColumnBinder.java000066400000000000000000000046211515451615200341530ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.PreparedStatement; import java.sql.SQLException; import org.apache.arrow.vector.FieldVector; /** A helper to bind values from a wrapped Arrow vector to a JDBC PreparedStatement. */ public interface ColumnBinder { /** * Bind the given row to the given parameter. * * @param statement The statement to bind to. * @param parameterIndex The parameter to bind to (1-indexed) * @param rowIndex The row to bind values from (0-indexed) * @throws SQLException if an error occurs */ void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException; /** * Get the JDBC type code used by this binder. * * @return A type code from {@link java.sql.Types}. */ int getJdbcType(); /** Get the vector used by this binder. */ FieldVector getVector(); /** Create a column binder for a vector, using the default JDBC type code for null values. */ static ColumnBinder forVector(FieldVector vector) { return forVector(vector, /*jdbcType*/ null); } /** * Create a column binder for a vector, overriding the JDBC type code used for null values. * * @param vector The vector that the column binder will wrap. * @param jdbcType The JDBC type code to use (or null to use the default). */ static ColumnBinder forVector(FieldVector vector, Integer jdbcType) { final ColumnBinder binder = vector.getField().getType().accept(new ColumnBinderArrowTypeVisitor(vector, jdbcType)); if (vector.getField().isNullable()) { return new NullableColumnBinder(binder); } return binder; } } ColumnBinderArrowTypeVisitor.java000066400000000000000000000250701515451615200374110ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.Types; import java.time.ZoneId; import java.util.Calendar; import java.util.TimeZone; 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.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.LargeVarCharVector; 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.TimeStampVector; import org.apache.arrow.vector.TinyIntVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.types.pojo.ArrowType; /** * Visitor to create the base ColumnBinder for a vector. * *

To handle null values, wrap the returned binder in a {@link NullableColumnBinder}. */ public class ColumnBinderArrowTypeVisitor implements ArrowType.ArrowTypeVisitor { private final FieldVector vector; private final Integer jdbcType; /** * Create a binder using a custom JDBC type code. * * @param vector The vector that the binder will wrap. * @param jdbcType The JDBC type code (or null to use the default). */ public ColumnBinderArrowTypeVisitor(FieldVector vector, Integer jdbcType) { this.vector = vector; this.jdbcType = jdbcType; } @Override public ColumnBinder visit(ArrowType.Null type) { throw new UnsupportedOperationException("No column binder implemented for type " + type); } @Override public ColumnBinder visit(ArrowType.Struct type) { throw new UnsupportedOperationException("No column binder implemented for type " + type); } @Override public ColumnBinder visit(ArrowType.List type) { return new ListBinder((ListVector) vector); } @Override public ColumnBinder visit(ArrowType.LargeList type) { throw new UnsupportedOperationException("No column binder implemented for type " + type); } @Override public ColumnBinder visit(ArrowType.FixedSizeList type) { throw new UnsupportedOperationException("No column binder implemented for type " + type); } @Override public ColumnBinder visit(ArrowType.Union type) { throw new UnsupportedOperationException("No column binder implemented for type " + type); } @Override public ColumnBinder visit(ArrowType.RunEndEncoded type) { throw new UnsupportedOperationException("No column binder implemented for type " + type); } @Override public ColumnBinder visit(ArrowType.Map type) { return new MapBinder((MapVector) vector); } @Override public ColumnBinder visit(ArrowType.Int type) { if (!type.getIsSigned()) { throw new UnsupportedOperationException( "No column binder implemented for unsigned type " + type); } switch (type.getBitWidth()) { case 8: return jdbcType == null ? new TinyIntBinder((TinyIntVector) vector) : new TinyIntBinder((TinyIntVector) vector, jdbcType); case 16: return jdbcType == null ? new SmallIntBinder((SmallIntVector) vector) : new SmallIntBinder((SmallIntVector) vector, jdbcType); case 32: return jdbcType == null ? new IntBinder((IntVector) vector) : new IntBinder((IntVector) vector, jdbcType); case 64: return jdbcType == null ? new BigIntBinder((BigIntVector) vector) : new BigIntBinder((BigIntVector) vector, jdbcType); default: throw new UnsupportedOperationException("No column binder implemented for type " + type); } } @Override public ColumnBinder visit(ArrowType.FloatingPoint type) { switch (type.getPrecision()) { case SINGLE: return jdbcType == null ? new Float4Binder((Float4Vector) vector) : new Float4Binder((Float4Vector) vector, jdbcType); case DOUBLE: return jdbcType == null ? new Float8Binder((Float8Vector) vector) : new Float8Binder((Float8Vector) vector, jdbcType); default: throw new UnsupportedOperationException("No column binder implemented for type " + type); } } @Override public ColumnBinder visit(ArrowType.Utf8 type) { VarCharVector varChar = (VarCharVector) vector; return jdbcType == null ? new VarCharBinder<>(varChar, Types.VARCHAR) : new VarCharBinder<>(varChar, jdbcType); } @Override public ColumnBinder visit(ArrowType.Utf8View type) { throw new UnsupportedOperationException( "Column binder implemented for type " + type + " is not supported"); } @Override public ColumnBinder visit(ArrowType.LargeUtf8 type) { LargeVarCharVector varChar = (LargeVarCharVector) vector; return jdbcType == null ? new VarCharBinder<>(varChar, Types.LONGVARCHAR) : new VarCharBinder<>(varChar, jdbcType); } @Override public ColumnBinder visit(ArrowType.Binary type) { VarBinaryVector varBinary = (VarBinaryVector) vector; return jdbcType == null ? new VarBinaryBinder<>(varBinary, Types.VARBINARY) : new VarBinaryBinder<>(varBinary, jdbcType); } @Override public ColumnBinder visit(ArrowType.BinaryView type) { throw new UnsupportedOperationException( "Column binder implemented for type " + type + " is not supported"); } @Override public ColumnBinder visit(ArrowType.LargeBinary type) { LargeVarBinaryVector varBinary = (LargeVarBinaryVector) vector; return jdbcType == null ? new VarBinaryBinder<>(varBinary, Types.LONGVARBINARY) : new VarBinaryBinder<>(varBinary, jdbcType); } @Override public ColumnBinder visit(ArrowType.FixedSizeBinary type) { FixedSizeBinaryVector binary = (FixedSizeBinaryVector) vector; return jdbcType == null ? new FixedSizeBinaryBinder(binary, Types.BINARY) : new FixedSizeBinaryBinder(binary, jdbcType); } @Override public ColumnBinder visit(ArrowType.Bool type) { return jdbcType == null ? new BitBinder((BitVector) vector) : new BitBinder((BitVector) vector, jdbcType); } @Override public ColumnBinder visit(ArrowType.Decimal type) { if (type.getBitWidth() == 128) { DecimalVector decimalVector = (DecimalVector) vector; return jdbcType == null ? new Decimal128Binder(decimalVector) : new Decimal128Binder(decimalVector, jdbcType); } else if (type.getBitWidth() == 256) { Decimal256Vector decimalVector = (Decimal256Vector) vector; return jdbcType == null ? new Decimal256Binder(decimalVector) : new Decimal256Binder(decimalVector, jdbcType); } throw new UnsupportedOperationException("No column binder implemented for type " + type); } @Override public ColumnBinder visit(ArrowType.Date type) { switch (type.getUnit()) { case DAY: return jdbcType == null ? new DateDayBinder((DateDayVector) vector) : new DateDayBinder((DateDayVector) vector, /*calendar*/ null, jdbcType); case MILLISECOND: return jdbcType == null ? new DateMilliBinder((DateMilliVector) vector) : new DateMilliBinder((DateMilliVector) vector, /*calendar*/ null, jdbcType); default: throw new UnsupportedOperationException("No column binder implemented for type " + type); } } @Override public ColumnBinder visit(ArrowType.Time type) { switch (type.getUnit()) { case SECOND: return jdbcType == null ? new Time32Binder((TimeSecVector) vector) : new Time32Binder((TimeSecVector) vector, jdbcType); case MILLISECOND: return jdbcType == null ? new Time32Binder((TimeMilliVector) vector) : new Time32Binder((TimeMilliVector) vector, jdbcType); case MICROSECOND: return jdbcType == null ? new Time64Binder((TimeMicroVector) vector) : new Time64Binder((TimeMicroVector) vector, jdbcType); case NANOSECOND: return jdbcType == null ? new Time64Binder((TimeNanoVector) vector) : new Time64Binder((TimeNanoVector) vector, jdbcType); default: throw new UnsupportedOperationException("No column binder implemented for type " + type); } } @Override public ColumnBinder visit(ArrowType.Timestamp type) { Calendar calendar = null; final String timezone = type.getTimezone(); if (timezone != null && !timezone.isEmpty()) { calendar = Calendar.getInstance(TimeZone.getTimeZone(ZoneId.of(timezone))); } return new TimeStampBinder((TimeStampVector) vector, calendar); } @Override public ColumnBinder visit(ArrowType.Interval type) { throw new UnsupportedOperationException("No column binder implemented for type " + type); } @Override public ColumnBinder visit(ArrowType.Duration type) { throw new UnsupportedOperationException("No column binder implemented for type " + type); } @Override public ColumnBinder visit(ArrowType.ListView type) { throw new UnsupportedOperationException("No column binder implemented for type " + type); } @Override public ColumnBinder visit(ArrowType.LargeListView type) { throw new UnsupportedOperationException("No column binder implemented for type " + type); } } DateDayBinder.java000066400000000000000000000040021515451615200342220ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.Date; import java.sql.PreparedStatement; import java.sql.SQLException; import java.sql.Types; import java.util.Calendar; import org.apache.arrow.vector.DateDayVector; /** A column binder for 32-bit dates. */ public class DateDayBinder extends BaseColumnBinder { private static final long MILLIS_PER_DAY = 86_400_000; private final Calendar calendar; public DateDayBinder(DateDayVector vector) { this(vector, null, Types.DATE); } public DateDayBinder(DateDayVector vector, Calendar calendar) { this(vector, calendar, Types.DATE); } public DateDayBinder(DateDayVector vector, Calendar calendar, int jdbcType) { super(vector, jdbcType); this.calendar = calendar; } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { // TODO: multiply with overflow final long index = (long) rowIndex * DateDayVector.TYPE_WIDTH; final Date value = new Date(vector.getDataBuffer().getInt(index) * MILLIS_PER_DAY); if (calendar == null) { statement.setDate(parameterIndex, value); } else { statement.setDate(parameterIndex, value, calendar); } } } DateMilliBinder.java000066400000000000000000000036511515451615200345640ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.Date; import java.sql.PreparedStatement; import java.sql.SQLException; import java.sql.Types; import java.util.Calendar; import org.apache.arrow.vector.DateMilliVector; /** A column binder for 64-bit dates. */ public class DateMilliBinder extends BaseColumnBinder { private final Calendar calendar; public DateMilliBinder(DateMilliVector vector) { this(vector, null, Types.DATE); } public DateMilliBinder(DateMilliVector vector, Calendar calendar) { this(vector, calendar, Types.DATE); } public DateMilliBinder(DateMilliVector vector, Calendar calendar, int jdbcType) { super(vector, jdbcType); this.calendar = calendar; } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { final long index = (long) rowIndex * DateMilliVector.TYPE_WIDTH; final Date value = new Date(vector.getDataBuffer().getLong(index)); if (calendar == null) { statement.setDate(parameterIndex, value); } else { statement.setDate(parameterIndex, value, calendar); } } } Decimal128Binder.java000066400000000000000000000032571515451615200345130ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.math.BigDecimal; import java.sql.PreparedStatement; import java.sql.SQLException; import java.sql.Types; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.util.DecimalUtility; /** A binder for 128-bit decimals. */ public class Decimal128Binder extends BaseColumnBinder { public Decimal128Binder(DecimalVector vector) { this(vector, Types.DECIMAL); } public Decimal128Binder(DecimalVector vector, int jdbcType) { super(vector, jdbcType); } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { final BigDecimal value = DecimalUtility.getBigDecimalFromArrowBuf( vector.getDataBuffer(), rowIndex, vector.getScale(), DecimalVector.TYPE_WIDTH); statement.setBigDecimal(parameterIndex, value); } } Decimal256Binder.java000066400000000000000000000032761515451615200345160ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.math.BigDecimal; import java.sql.PreparedStatement; import java.sql.SQLException; import java.sql.Types; import org.apache.arrow.vector.Decimal256Vector; import org.apache.arrow.vector.util.DecimalUtility; /** A binder for 256-bit decimals. */ public class Decimal256Binder extends BaseColumnBinder { public Decimal256Binder(Decimal256Vector vector) { this(vector, Types.DECIMAL); } public Decimal256Binder(Decimal256Vector vector, int jdbcType) { super(vector, jdbcType); } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { final BigDecimal value = DecimalUtility.getBigDecimalFromArrowBuf( vector.getDataBuffer(), rowIndex, vector.getScale(), Decimal256Vector.TYPE_WIDTH); statement.setBigDecimal(parameterIndex, value); } } FixedSizeBinaryBinder.java000066400000000000000000000033471515451615200357610ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.PreparedStatement; import java.sql.SQLException; import org.apache.arrow.vector.FixedSizeBinaryVector; /** A binder for fixed-width binary types. */ public class FixedSizeBinaryBinder extends BaseColumnBinder { /** * Create a binder for the given vector using the given JDBC type for null values. * * @param vector The vector to draw values from. * @param jdbcType The JDBC type code. */ public FixedSizeBinaryBinder(FixedSizeBinaryVector vector, int jdbcType) { super(vector, jdbcType); } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { byte[] binaryData = new byte[vector.getByteWidth()]; vector .getDataBuffer() .getBytes((long) rowIndex * binaryData.length, binaryData, 0, binaryData.length); statement.setBytes(parameterIndex, binaryData); } } Float4Binder.java000066400000000000000000000027741515451615200340560ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.PreparedStatement; import java.sql.SQLException; import java.sql.Types; import org.apache.arrow.vector.Float4Vector; /** A binder for 32-bit floats. */ public class Float4Binder extends BaseColumnBinder { public Float4Binder(Float4Vector vector) { this(vector, Types.REAL); } public Float4Binder(Float4Vector vector, int jdbcType) { super(vector, jdbcType); } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { final float value = vector.getDataBuffer().getFloat((long) rowIndex * Float4Vector.TYPE_WIDTH); statement.setFloat(parameterIndex, value); } } Float8Binder.java000066400000000000000000000030111515451615200340430ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.PreparedStatement; import java.sql.SQLException; import java.sql.Types; import org.apache.arrow.vector.Float8Vector; /** A binder for 64-bit floats. */ public class Float8Binder extends BaseColumnBinder { public Float8Binder(Float8Vector vector) { this(vector, Types.DOUBLE); } public Float8Binder(Float8Vector vector, int jdbcType) { super(vector, jdbcType); } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { final double value = vector.getDataBuffer().getDouble((long) rowIndex * Float8Vector.TYPE_WIDTH); statement.setDouble(parameterIndex, value); } } IntBinder.java000066400000000000000000000027521515451615200334530ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.PreparedStatement; import java.sql.SQLException; import java.sql.Types; import org.apache.arrow.vector.IntVector; /** A column binder for 32-bit integers. */ public class IntBinder extends BaseColumnBinder { public IntBinder(IntVector vector) { this(vector, Types.INTEGER); } public IntBinder(IntVector vector, int jdbcType) { super(vector, jdbcType); } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { final int value = vector.getDataBuffer().getInt((long) rowIndex * IntVector.TYPE_WIDTH); statement.setInt(parameterIndex, value); } } ListBinder.java000066400000000000000000000056521515451615200336360ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.lang.reflect.Array; import java.util.ArrayList; import java.util.Arrays; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.impl.UnionListReader; import org.apache.arrow.vector.util.Text; /** A column binder for list of primitive values. */ public class ListBinder extends BaseColumnBinder { private final UnionListReader listReader; private final Class arrayElementClass; private final boolean isTextColumn; public ListBinder(ListVector vector) { this(vector, java.sql.Types.ARRAY); } /** * Init ListBinder and determine type of data vector. * * @param vector corresponding data vector from arrow buffer for binding * @param jdbcType parameter jdbc type */ public ListBinder(ListVector vector, int jdbcType) { super(vector, jdbcType); listReader = vector.getReader(); Class dataVectorClass = vector.getDataVector().getClass(); try { arrayElementClass = dataVectorClass.getMethod("getObject", Integer.TYPE).getReturnType(); } catch (NoSuchMethodException e) { final String message = String.format( "Issue to determine type for getObject method of data vector class %s ", dataVectorClass.getName()); throw new RuntimeException(message); } isTextColumn = arrayElementClass.isAssignableFrom(Text.class); } @Override public void bind(java.sql.PreparedStatement statement, int parameterIndex, int rowIndex) throws java.sql.SQLException { listReader.setPosition(rowIndex); ArrayList sourceArray = (ArrayList) listReader.readObject(); Object array; if (!isTextColumn) { array = Array.newInstance(arrayElementClass, sourceArray.size()); Arrays.setAll((Object[]) array, sourceArray::get); } else { array = new String[sourceArray.size()]; Arrays.setAll( (Object[]) array, idx -> sourceArray.get(idx) != null ? sourceArray.get(idx).toString() : null); } statement.setObject(parameterIndex, array); } } MapBinder.java000066400000000000000000000064621515451615200334400ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.PreparedStatement; import java.sql.SQLException; import java.sql.Types; import java.util.LinkedHashMap; import java.util.List; import java.util.Objects; import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.impl.UnionMapReader; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.util.JsonStringHashMap; /** A column binder for map of primitive values. */ public class MapBinder extends BaseColumnBinder { private UnionMapReader reader; private final boolean isTextKey; private final boolean isTextValue; public MapBinder(MapVector vector) { this(vector, Types.VARCHAR); } /** * Init MapBinder and determine type of data vector. * * @param vector corresponding data vector from arrow buffer for binding * @param jdbcType parameter jdbc type */ public MapBinder(MapVector vector, int jdbcType) { super(vector, jdbcType); reader = vector.getReader(); List structField = Objects.requireNonNull(vector.getField()).getChildren(); if (structField.size() != 1) { throw new IllegalArgumentException("Expected Struct field metadata inside Map field"); } List keyValueFields = Objects.requireNonNull(structField.get(0)).getChildren(); if (keyValueFields.size() != 2) { throw new IllegalArgumentException( "Expected two children fields " + "inside nested Struct field in Map"); } ArrowType keyType = Objects.requireNonNull(keyValueFields.get(0)).getType(); ArrowType valueType = Objects.requireNonNull(keyValueFields.get(1)).getType(); isTextKey = ArrowType.Utf8.INSTANCE.equals(keyType); isTextValue = ArrowType.Utf8.INSTANCE.equals(valueType); } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { reader.setPosition(rowIndex); LinkedHashMap tags = new JsonStringHashMap<>(); while (reader.next()) { Object key = reader.key().readObject(); Object value = reader.value().readObject(); tags.put( isTextKey && key != null ? key.toString() : key, isTextValue && value != null ? value.toString() : value); } switch (jdbcType) { case Types.VARCHAR: statement.setString(parameterIndex, tags.toString()); break; case Types.OTHER: default: statement.setObject(parameterIndex, tags); } } } NullableColumnBinder.java000066400000000000000000000032661515451615200356360ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.PreparedStatement; import java.sql.SQLException; import org.apache.arrow.vector.FieldVector; /** A ColumnBinder that checks for nullability before deferring to a type-specific binder. */ public class NullableColumnBinder implements ColumnBinder { private final ColumnBinder wrapped; public NullableColumnBinder(ColumnBinder wrapped) { this.wrapped = wrapped; } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { if (wrapped.getVector().isNull(rowIndex)) { statement.setNull(parameterIndex, wrapped.getJdbcType()); } else { wrapped.bind(statement, parameterIndex, rowIndex); } } @Override public int getJdbcType() { return wrapped.getJdbcType(); } @Override public FieldVector getVector() { return wrapped.getVector(); } } SmallIntBinder.java000066400000000000000000000030411515451615200344340ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.PreparedStatement; import java.sql.SQLException; import java.sql.Types; import org.apache.arrow.vector.SmallIntVector; /** A column binder for 8-bit integers. */ public class SmallIntBinder extends BaseColumnBinder { public SmallIntBinder(SmallIntVector vector) { this(vector, Types.SMALLINT); } public SmallIntBinder(SmallIntVector vector, int jdbcType) { super(vector, jdbcType); } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { final short value = vector.getDataBuffer().getShort((short) rowIndex * SmallIntVector.TYPE_WIDTH); statement.setShort(parameterIndex, value); } } Time32Binder.java000066400000000000000000000041571515451615200337650ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.PreparedStatement; import java.sql.SQLException; import java.sql.Time; import java.sql.Types; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.TimeMilliVector; import org.apache.arrow.vector.TimeSecVector; /** A binder for 32-bit time types. */ public class Time32Binder extends BaseColumnBinder { private static final long TYPE_WIDTH = 4; private final long factor; public Time32Binder(TimeSecVector vector) { this(vector, Types.TIME); } public Time32Binder(TimeMilliVector vector) { this(vector, Types.TIME); } public Time32Binder(TimeSecVector vector, int jdbcType) { this(vector, /*factor*/ 1_000, jdbcType); } public Time32Binder(TimeMilliVector vector, int jdbcType) { this(vector, /*factor*/ 1, jdbcType); } Time32Binder(BaseFixedWidthVector vector, long factor, int jdbcType) { super(vector, jdbcType); this.factor = factor; } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { // TODO: multiply with overflow // TODO: take in a Calendar as well? final Time value = new Time(vector.getDataBuffer().getInt(rowIndex * TYPE_WIDTH) * factor); statement.setTime(parameterIndex, value); } } Time64Binder.java000066400000000000000000000041711515451615200337660ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.PreparedStatement; import java.sql.SQLException; import java.sql.Time; import java.sql.Types; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.TimeMicroVector; import org.apache.arrow.vector.TimeNanoVector; /** A binder for 64-bit time types. */ public class Time64Binder extends BaseColumnBinder { private static final long TYPE_WIDTH = 8; private final long factor; public Time64Binder(TimeMicroVector vector) { this(vector, Types.TIME); } public Time64Binder(TimeNanoVector vector) { this(vector, Types.TIME); } public Time64Binder(TimeMicroVector vector, int jdbcType) { this(vector, /*factor*/ 1_000, jdbcType); } public Time64Binder(TimeNanoVector vector, int jdbcType) { this(vector, /*factor*/ 1_000_000, jdbcType); } Time64Binder(BaseFixedWidthVector vector, long factor, int jdbcType) { super(vector, jdbcType); this.factor = factor; } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { // TODO: option to throw on truncation (vendor Guava IntMath#multiply) final Time value = new Time(vector.getDataBuffer().getLong(rowIndex * TYPE_WIDTH) / factor); statement.setTime(parameterIndex, value); } } TimeStampBinder.java000066400000000000000000000076021515451615200346230ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.PreparedStatement; import java.sql.SQLException; import java.sql.Timestamp; import java.sql.Types; import java.util.Calendar; import org.apache.arrow.vector.TimeStampVector; import org.apache.arrow.vector.types.pojo.ArrowType; /** A column binder for timestamps. */ public class TimeStampBinder extends BaseColumnBinder { private final Calendar calendar; private final long unitsPerSecond; private final long nanosPerUnit; /** Create a binder for a timestamp vector using the default JDBC type code. */ public TimeStampBinder(TimeStampVector vector, Calendar calendar) { this( vector, calendar, isZoned(vector.getField().getType()) ? Types.TIMESTAMP_WITH_TIMEZONE : Types.TIMESTAMP); } /** * Create a binder for a timestamp vector. * * @param vector The vector to pull values from. * @param calendar Optionally, the calendar to pass to JDBC. * @param jdbcType The JDBC type code to use for null values. */ public TimeStampBinder(TimeStampVector vector, Calendar calendar, int jdbcType) { super(vector, jdbcType); this.calendar = calendar; final ArrowType.Timestamp type = (ArrowType.Timestamp) vector.getField().getType(); switch (type.getUnit()) { case SECOND: this.unitsPerSecond = 1; this.nanosPerUnit = 1_000_000_000; break; case MILLISECOND: this.unitsPerSecond = 1_000; this.nanosPerUnit = 1_000_000; break; case MICROSECOND: this.unitsPerSecond = 1_000_000; this.nanosPerUnit = 1_000; break; case NANOSECOND: this.unitsPerSecond = 1_000_000_000; this.nanosPerUnit = 1; break; default: throw new IllegalArgumentException("Invalid time unit in " + type); } } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { // TODO: option to throw on truncation (vendor Guava IntMath#multiply) or overflow final long rawValue = vector.getDataBuffer().getLong((long) rowIndex * TimeStampVector.TYPE_WIDTH); final long seconds = rawValue / unitsPerSecond; final int nanos = (int) ((rawValue - (seconds * unitsPerSecond)) * nanosPerUnit); final Timestamp value = new Timestamp(seconds * 1_000); value.setNanos(nanos); if (calendar != null) { // Timestamp == Date == UTC timestamp (confusingly). Arrow's timestamp with timezone is a UTC // value with a // zone offset, so we don't need to do any conversion. statement.setTimestamp(parameterIndex, value, calendar); } else { // Arrow timestamp without timezone isn't strictly convertible to any timezone. So this is // technically wrong, // but there is no 'correct' interpretation here. The application should provide a calendar. statement.setTimestamp(parameterIndex, value); } } private static boolean isZoned(ArrowType type) { final String timezone = ((ArrowType.Timestamp) type).getTimezone(); return timezone != null && !timezone.isEmpty(); } } TinyIntBinder.java000066400000000000000000000030141515451615200343070ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.PreparedStatement; import java.sql.SQLException; import java.sql.Types; import org.apache.arrow.vector.TinyIntVector; /** A column binder for 8-bit integers. */ public class TinyIntBinder extends BaseColumnBinder { public TinyIntBinder(TinyIntVector vector) { this(vector, Types.TINYINT); } public TinyIntBinder(TinyIntVector vector, int jdbcType) { super(vector, jdbcType); } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { final byte value = vector.getDataBuffer().getByte((long) rowIndex * TinyIntVector.TYPE_WIDTH); statement.setByte(parameterIndex, value); } } VarBinaryBinder.java000066400000000000000000000045061515451615200346150ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.sql.PreparedStatement; import java.sql.SQLException; import org.apache.arrow.memory.util.ArrowBufPointer; import org.apache.arrow.vector.ElementAddressableVector; import org.apache.arrow.vector.FieldVector; /** * A binder for variable-width binary types. * * @param The binary vector. */ public class VarBinaryBinder extends BaseColumnBinder { private final ArrowBufPointer element; /** * Create a binder for the given vector using the given JDBC type for null values. * * @param vector The vector to draw values from. * @param jdbcType The JDBC type code. */ public VarBinaryBinder(T vector, int jdbcType) { super(vector, jdbcType); this.element = new ArrowBufPointer(); } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { vector.getDataPointer(rowIndex, element); if (element.getBuf() == null) { statement.setNull(parameterIndex, jdbcType); return; } if (element.getLength() > (long) Integer.MAX_VALUE) { final String message = String.format( "Length of value at index %d (%d) exceeds Integer.MAX_VALUE", rowIndex, element.getLength()); throw new RuntimeException(message); } byte[] binaryData = new byte[(int) element.getLength()]; element.getBuf().getBytes(element.getOffset(), binaryData); statement.setBytes(parameterIndex, binaryData); } } VarCharBinder.java000066400000000000000000000046021515451615200342430ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.binder; import java.nio.charset.StandardCharsets; import java.sql.PreparedStatement; import java.sql.SQLException; import org.apache.arrow.memory.util.ArrowBufPointer; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VariableWidthVector; /** * A binder for variable-width string types. * * @param The text vector. */ public class VarCharBinder extends BaseColumnBinder { private final ArrowBufPointer element; /** * Create a binder for the given vector using the given JDBC type for null values. * * @param vector The vector to draw values from. * @param jdbcType The JDBC type code. */ public VarCharBinder(T vector, int jdbcType) { super(vector, jdbcType); this.element = new ArrowBufPointer(); } @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { vector.getDataPointer(rowIndex, element); if (element.getBuf() == null) { statement.setNull(parameterIndex, jdbcType); return; } if (element.getLength() > (long) Integer.MAX_VALUE) { final String message = String.format( "Length of value at index %d (%d) exceeds Integer.MAX_VALUE", rowIndex, element.getLength()); throw new RuntimeException(message); } byte[] utf8Bytes = new byte[(int) element.getLength()]; element.getBuf().getBytes(element.getOffset(), utf8Bytes); statement.setString(parameterIndex, new String(utf8Bytes, StandardCharsets.UTF_8)); } } package-info.java000066400000000000000000000016331515451615200341160ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ /** Utilities to bind Arrow data as JDBC prepared statement parameters. */ package org.apache.arrow.adapter.jdbc.binder; apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/000077500000000000000000000000001515451615200313735ustar00rootroot00000000000000ArrayConsumer.java000066400000000000000000000103401515451615200347470ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.io.IOException; import java.sql.Array; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.ListVector; /** * Consumer which consume array type values from {@link ResultSet}. Write the data to {@link * org.apache.arrow.vector.complex.ListVector}. */ public abstract class ArrayConsumer extends BaseConsumer { /** Creates a consumer for {@link ListVector}. */ public static ArrayConsumer createConsumer( ListVector vector, JdbcConsumer delegate, int index, boolean nullable) { if (nullable) { return new ArrayConsumer.NullableArrayConsumer(vector, delegate, index); } else { return new ArrayConsumer.NonNullableArrayConsumer(vector, delegate, index); } } protected final JdbcConsumer delegate; private final ValueVector innerVector; protected int innerVectorIndex = 0; /** Instantiate a ArrayConsumer. */ public ArrayConsumer(ListVector vector, JdbcConsumer delegate, int index) { super(vector, index); this.delegate = delegate; this.innerVector = vector.getDataVector(); } @Override public void close() throws Exception { this.vector.close(); this.delegate.close(); } @Override public void resetValueVector(ListVector vector) { super.resetValueVector(vector); FieldVector childVector = vector.getDataVector(); this.delegate.resetValueVector(childVector); innerVectorIndex = 0; } void ensureInnerVectorCapacity(int targetCapacity) { while (innerVector.getValueCapacity() < targetCapacity) { innerVector.reAlloc(); } } /** Nullable consumer for {@link ListVector}. */ static class NullableArrayConsumer extends ArrayConsumer { /** Instantiate a nullable array consumer. */ public NullableArrayConsumer(ListVector vector, JdbcConsumer delegate, int index) { super(vector, delegate, index); } @Override public void consume(ResultSet resultSet) throws SQLException, IOException { final Array array = resultSet.getArray(columnIndexInResultSet); if (!resultSet.wasNull()) { vector.startNewValue(currentIndex); int count = 0; try (ResultSet rs = array.getResultSet()) { while (rs.next()) { ensureInnerVectorCapacity(innerVectorIndex + count + 1); delegate.consume(rs); count++; } } vector.endValue(currentIndex, count); innerVectorIndex += count; } currentIndex++; } } /** Non-nullable consumer for {@link ListVector}. */ static class NonNullableArrayConsumer extends ArrayConsumer { /** Instantiate a nullable array consumer. */ public NonNullableArrayConsumer(ListVector vector, JdbcConsumer delegate, int index) { super(vector, delegate, index); } @Override public void consume(ResultSet resultSet) throws SQLException, IOException { final Array array = resultSet.getArray(columnIndexInResultSet); vector.startNewValue(currentIndex); int count = 0; try (ResultSet rs = array.getResultSet()) { while (rs.next()) { ensureInnerVectorCapacity(innerVectorIndex + count + 1); delegate.consume(rs); count++; } } vector.endValue(currentIndex, count); innerVectorIndex += count; currentIndex++; } } } BaseConsumer.java000066400000000000000000000031331515451615200345450ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import org.apache.arrow.vector.ValueVector; /** * Base class for all consumers. * * @param vector type. */ public abstract class BaseConsumer implements JdbcConsumer { protected V vector; protected final int columnIndexInResultSet; protected int currentIndex; /** * Constructs a new consumer. * * @param vector the underlying vector for the consumer. * @param index the column id for the consumer. */ public BaseConsumer(V vector, int index) { this.vector = vector; this.columnIndexInResultSet = index; } @Override public void close() throws Exception { this.vector.close(); } @Override public void resetValueVector(V vector) { this.vector = vector; this.currentIndex = 0; } } BigIntConsumer.java000066400000000000000000000052751515451615200350600ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.vector.BigIntVector; /** * Consumer which consume bigint type values from {@link ResultSet}. Write the data to {@link * org.apache.arrow.vector.BigIntVector}. */ public class BigIntConsumer { /** Creates a consumer for {@link BigIntVector}. */ public static JdbcConsumer createConsumer( BigIntVector vector, int index, boolean nullable) { if (nullable) { return new NullableBigIntConsumer(vector, index); } else { return new NonNullableBigIntConsumer(vector, index); } } /** Nullable consumer for big int. */ static class NullableBigIntConsumer extends BaseConsumer { /** Instantiate a BigIntConsumer. */ public NullableBigIntConsumer(BigIntVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { long value = resultSet.getLong(columnIndexInResultSet); if (!resultSet.wasNull()) { // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, value); } currentIndex++; } } /** Non-nullable consumer for big int. */ static class NonNullableBigIntConsumer extends BaseConsumer { /** Instantiate a BigIntConsumer. */ public NonNullableBigIntConsumer(BigIntVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { long value = resultSet.getLong(columnIndexInResultSet); // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, value); currentIndex++; } } } BinaryConsumer.java000066400000000000000000000076761515451615200351370ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.io.IOException; import java.io.InputStream; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.VarBinaryVector; /** * Consumer which consume binary type values from {@link ResultSet}. Write the data to {@link * org.apache.arrow.vector.VarBinaryVector}. */ public abstract class BinaryConsumer extends BaseConsumer { /** Creates a consumer for {@link VarBinaryVector}. */ public static BinaryConsumer createConsumer(VarBinaryVector vector, int index, boolean nullable) { if (nullable) { return new NullableBinaryConsumer(vector, index); } else { return new NonNullableBinaryConsumer(vector, index); } } private final byte[] reuseBytes = new byte[1024]; /** Instantiate a BinaryConsumer. */ public BinaryConsumer(VarBinaryVector vector, int index) { super(vector, index); if (vector != null) { vector.allocateNewSafe(); } } /** 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) { int read; while ((read = is.read(reuseBytes)) != -1) { while (vector.getDataBuffer().capacity() < (startOffset + dataLength + read)) { vector.reallocDataBuffer(); } vector.getDataBuffer().setBytes(startOffset + dataLength, reuseBytes, 0, read); dataLength += read; } BitVectorHelper.setBit(vector.getValidityBuffer(), currentIndex); } offsetBuffer.setInt( (currentIndex + 1) * ((long) VarBinaryVector.OFFSET_WIDTH), startOffset + dataLength); vector.setLastSet(currentIndex); } public void moveWriterPosition() { currentIndex++; } @Override public void resetValueVector(VarBinaryVector vector) { this.vector = vector; this.vector.allocateNewSafe(); this.currentIndex = 0; } /** Consumer for nullable binary data. */ static class NullableBinaryConsumer extends BinaryConsumer { /** Instantiate a BinaryConsumer. */ public NullableBinaryConsumer(VarBinaryVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException, IOException { InputStream is = resultSet.getBinaryStream(columnIndexInResultSet); consume(is); moveWriterPosition(); } } /** Consumer for non-nullable binary data. */ static class NonNullableBinaryConsumer extends BinaryConsumer { /** Instantiate a BinaryConsumer. */ public NonNullableBinaryConsumer(VarBinaryVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException, IOException { InputStream is = resultSet.getBinaryStream(columnIndexInResultSet); consume(is); moveWriterPosition(); } } } BitConsumer.java000066400000000000000000000052341515451615200344150ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.vector.BitVector; /** * Consumer which consume bit type values from {@link ResultSet}. Write the data to {@link * BitVector}. */ public class BitConsumer { /** Creates a consumer for {@link BitVector}. */ public static JdbcConsumer createConsumer( BitVector vector, int index, boolean nullable) { if (nullable) { return new NullableBitConsumer(vector, index); } else { return new NonNullableBitConsumer(vector, index); } } /** Nullable consumer for {@link BitVector}. */ static class NullableBitConsumer extends BaseConsumer { /** Instantiate a BitConsumer. */ public NullableBitConsumer(BitVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { boolean value = resultSet.getBoolean(columnIndexInResultSet); if (!resultSet.wasNull()) { // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, value ? 1 : 0); } currentIndex++; } } /** Non-nullable consumer for {@link BitVector}. */ static class NonNullableBitConsumer extends BaseConsumer { /** Instantiate a BitConsumer. */ public NonNullableBitConsumer(BitVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { boolean value = resultSet.getBoolean(columnIndexInResultSet); // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, value ? 1 : 0); currentIndex++; } } } BlobConsumer.java000066400000000000000000000042251515451615200345540ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.io.IOException; import java.sql.Blob; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.vector.VarBinaryVector; /** * Consumer which consume blob type values from {@link ResultSet}. Write the data to {@link * VarBinaryVector}. */ public class BlobConsumer extends BaseConsumer { private BinaryConsumer delegate; private final boolean nullable; /** Creates a consumer for {@link VarBinaryVector}. */ public static BlobConsumer createConsumer(BinaryConsumer delegate, int index, boolean nullable) { return new BlobConsumer(delegate, index, nullable); } /** Instantiate a BlobConsumer. */ public BlobConsumer(BinaryConsumer delegate, int index, boolean nullable) { super(null, index); this.delegate = delegate; this.nullable = nullable; } @Override public void consume(ResultSet resultSet) throws SQLException, IOException { Blob blob = resultSet.getBlob(columnIndexInResultSet); if (blob != null) { delegate.consume(blob.getBinaryStream()); } delegate.moveWriterPosition(); } @Override public void close() throws Exception { delegate.close(); } @Override public void resetValueVector(VarBinaryVector vector) { delegate = BinaryConsumer.createConsumer(vector, columnIndexInResultSet, nullable); } } ClobConsumer.java000066400000000000000000000120271515451615200345540ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.nio.charset.StandardCharsets; import java.sql.Clob; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.util.MemoryUtil; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.VarCharVector; /** * Consumer which consume clob type values from {@link ResultSet}. Write the data to {@link * org.apache.arrow.vector.VarCharVector}. */ public abstract class ClobConsumer extends BaseConsumer { /** Creates a consumer for {@link VarCharVector}. */ public static ClobConsumer createConsumer(VarCharVector vector, int index, boolean nullable) { if (nullable) { return new NullableClobConsumer(vector, index); } else { return new NonNullableClobConsumer(vector, index); } } private static final int BUFFER_SIZE = 256; /** Instantiate a ClobConsumer. */ public ClobConsumer(VarCharVector vector, int index) { super(vector, index); if (vector != null) { vector.allocateNewSafe(); } } @Override public void resetValueVector(VarCharVector vector) { this.vector = vector; this.vector.allocateNewSafe(); this.currentIndex = 0; } /** Nullable consumer for clob data. */ static class NullableClobConsumer extends ClobConsumer { /** Instantiate a ClobConsumer. */ public NullableClobConsumer(VarCharVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { Clob clob = resultSet.getClob(columnIndexInResultSet); if (!resultSet.wasNull()) { if (clob != null) { long length = clob.length(); int read = 1; int readSize = length < BUFFER_SIZE ? (int) length : BUFFER_SIZE; int totalBytes = 0; ArrowBuf dataBuffer = vector.getDataBuffer(); ArrowBuf offsetBuffer = vector.getOffsetBuffer(); int startIndex = offsetBuffer.getInt(currentIndex * 4L); while (read <= length) { String str = clob.getSubString(read, readSize); byte[] bytes = str.getBytes(StandardCharsets.UTF_8); while ((dataBuffer.writerIndex() + bytes.length) > dataBuffer.capacity()) { vector.reallocDataBuffer(); } MemoryUtil.copyToMemory( bytes, 0, dataBuffer.memoryAddress() + startIndex + totalBytes, bytes.length); totalBytes += bytes.length; read += readSize; } offsetBuffer.setInt((currentIndex + 1) * 4L, startIndex + totalBytes); BitVectorHelper.setBit(vector.getValidityBuffer(), currentIndex); vector.setLastSet(currentIndex); } } currentIndex++; } } /** Non-nullable consumer for clob data. */ static class NonNullableClobConsumer extends ClobConsumer { /** Instantiate a ClobConsumer. */ public NonNullableClobConsumer(VarCharVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { Clob clob = resultSet.getClob(columnIndexInResultSet); if (clob != null) { long length = clob.length(); int read = 1; int readSize = length < BUFFER_SIZE ? (int) length : BUFFER_SIZE; int totalBytes = 0; ArrowBuf dataBuffer = vector.getDataBuffer(); ArrowBuf offsetBuffer = vector.getOffsetBuffer(); int startIndex = offsetBuffer.getInt(currentIndex * 4L); while (read <= length) { String str = clob.getSubString(read, readSize); byte[] bytes = str.getBytes(StandardCharsets.UTF_8); while ((dataBuffer.writerIndex() + bytes.length) > dataBuffer.capacity()) { vector.reallocDataBuffer(); } MemoryUtil.copyToMemory( bytes, 0, dataBuffer.memoryAddress() + startIndex + totalBytes, bytes.length); totalBytes += bytes.length; read += readSize; } offsetBuffer.setInt((currentIndex + 1) * 4L, startIndex + totalBytes); BitVectorHelper.setBit(vector.getValidityBuffer(), currentIndex); vector.setLastSet(currentIndex); } currentIndex++; } } } CompositeJdbcConsumer.java000066400000000000000000000053461515451615200364300ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.io.IOException; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.adapter.jdbc.JdbcFieldInfo; import org.apache.arrow.adapter.jdbc.consumer.exceptions.JdbcConsumerException; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VectorSchemaRoot; /** Composite consumer which hold all consumers. It manages the consume and cleanup process. */ public class CompositeJdbcConsumer implements JdbcConsumer { private final JdbcConsumer[] consumers; /** Construct an instance. */ public CompositeJdbcConsumer(JdbcConsumer[] consumers) { this.consumers = consumers; } @Override public void consume(ResultSet rs) throws SQLException, IOException { for (int i = 0; i < consumers.length; i++) { try { consumers[i].consume(rs); } catch (Exception e) { if (consumers[i] instanceof BaseConsumer) { BaseConsumer consumer = (BaseConsumer) consumers[i]; JdbcFieldInfo fieldInfo = new JdbcFieldInfo(rs.getMetaData(), consumer.columnIndexInResultSet); throw new JdbcConsumerException( "Exception while consuming JDBC value", e, fieldInfo, consumer.vector.getField()); } else { throw e; } } } } @Override public void close() { try { // clean up AutoCloseables.close(consumers); } catch (Exception e) { throw new RuntimeException("Error occurred while releasing resources.", e); } } @Override public void resetValueVector(ValueVector vector) {} /** Reset inner consumers through vectors in the vector schema root. */ public void resetVectorSchemaRoot(VectorSchemaRoot root) { assert root.getFieldVectors().size() == consumers.length; for (int i = 0; i < consumers.length; i++) { consumers[i].resetValueVector(root.getVector(i)); } } } DateConsumer.java000066400000000000000000000072431515451615200345560ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.sql.Date; import java.sql.ResultSet; import java.sql.SQLException; import java.util.Calendar; import java.util.concurrent.TimeUnit; import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DateMilliVector; /** * Consumer which consume date type values from {@link ResultSet}. Write the data to {@link * org.apache.arrow.vector.DateDayVector}. */ public class DateConsumer { /** Creates a consumer for {@link DateMilliVector}. */ public static JdbcConsumer createConsumer( DateDayVector vector, int index, boolean nullable, Calendar calendar) { if (nullable) { return new NullableDateConsumer(vector, index, calendar); } else { return new NonNullableDateConsumer(vector, index, calendar); } } /** Nullable consumer for date. */ static class NullableDateConsumer extends BaseConsumer { protected final Calendar calendar; /** Instantiate a DateConsumer. */ public NullableDateConsumer(DateDayVector vector, int index) { this(vector, index, /* calendar */ null); } /** Instantiate a DateConsumer. */ public NullableDateConsumer(DateDayVector vector, int index, Calendar calendar) { super(vector, index); this.calendar = calendar; } @Override public void consume(ResultSet resultSet) throws SQLException { Date date = calendar == null ? resultSet.getDate(columnIndexInResultSet) : resultSet.getDate(columnIndexInResultSet, calendar); if (!resultSet.wasNull()) { // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, Math.toIntExact(TimeUnit.MILLISECONDS.toDays(date.getTime()))); } currentIndex++; } } /** Non-nullable consumer for date. */ static class NonNullableDateConsumer extends BaseConsumer { protected final Calendar calendar; /** Instantiate a DateConsumer. */ public NonNullableDateConsumer(DateDayVector vector, int index) { this(vector, index, /* calendar */ null); } /** Instantiate a DateConsumer. */ public NonNullableDateConsumer(DateDayVector vector, int index, Calendar calendar) { super(vector, index); this.calendar = calendar; } @Override public void consume(ResultSet resultSet) throws SQLException { Date date = calendar == null ? resultSet.getDate(columnIndexInResultSet) : resultSet.getDate(columnIndexInResultSet, calendar); // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, Math.toIntExact(TimeUnit.MILLISECONDS.toDays(date.getTime()))); currentIndex++; } } } Decimal256Consumer.java000066400000000000000000000104021515451615200354630ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.math.BigDecimal; import java.math.RoundingMode; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.vector.Decimal256Vector; /** * Consumer which consume decimal type values from {@link ResultSet}. Write the data to {@link * org.apache.arrow.vector.Decimal256Vector}. */ public abstract class Decimal256Consumer extends BaseConsumer { private final RoundingMode bigDecimalRoundingMode; private final int scale; /** * Constructs a new consumer. * * @param vector the underlying vector for the consumer. * @param index the column id for the consumer. */ public Decimal256Consumer(Decimal256Vector vector, int index) { this(vector, index, null); } /** * Constructs a new consumer, with optional coercibility. * * @param vector the underlying vector for the consumer. * @param index the column index for the consumer. * @param bigDecimalRoundingMode java.math.RoundingMode to be applied if the BigDecimal scale does * not match that of the target vector. Set to null to retain strict matching behavior (scale * of source and target vector must match exactly). */ public Decimal256Consumer( Decimal256Vector vector, int index, RoundingMode bigDecimalRoundingMode) { super(vector, index); this.bigDecimalRoundingMode = bigDecimalRoundingMode; this.scale = vector.getScale(); } /** Creates a consumer for {@link Decimal256Vector}. */ public static JdbcConsumer createConsumer( Decimal256Vector vector, int index, boolean nullable, RoundingMode bigDecimalRoundingMode) { if (nullable) { return new NullableDecimal256Consumer(vector, index, bigDecimalRoundingMode); } else { return new NonNullableDecimal256Consumer(vector, index, bigDecimalRoundingMode); } } protected void set(BigDecimal value) { if (bigDecimalRoundingMode != null && value.scale() != scale) { value = value.setScale(scale, bigDecimalRoundingMode); } vector.set(currentIndex, value); } /** Consumer for nullable decimal. */ static class NullableDecimal256Consumer extends Decimal256Consumer { /** Instantiate a Decimal256Consumer. */ public NullableDecimal256Consumer( Decimal256Vector vector, int index, RoundingMode bigDecimalRoundingMode) { super(vector, index, bigDecimalRoundingMode); } @Override public void consume(ResultSet resultSet) throws SQLException { BigDecimal value = resultSet.getBigDecimal(columnIndexInResultSet); if (!resultSet.wasNull()) { // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. set(value); } currentIndex++; } } /** Consumer for non-nullable decimal. */ static class NonNullableDecimal256Consumer extends Decimal256Consumer { /** Instantiate a Decimal256Consumer. */ public NonNullableDecimal256Consumer( Decimal256Vector vector, int index, RoundingMode bigDecimalRoundingMode) { super(vector, index, bigDecimalRoundingMode); } @Override public void consume(ResultSet resultSet) throws SQLException { BigDecimal value = resultSet.getBigDecimal(columnIndexInResultSet); // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. set(value); currentIndex++; } } } DecimalConsumer.java000066400000000000000000000102661515451615200352360ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.math.BigDecimal; import java.math.RoundingMode; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.vector.DecimalVector; /** * Consumer which consume decimal type values from {@link ResultSet}. Write the data to {@link * org.apache.arrow.vector.DecimalVector}. */ public abstract class DecimalConsumer extends BaseConsumer { private final RoundingMode bigDecimalRoundingMode; private final int scale; /** * Constructs a new consumer. * * @param vector the underlying vector for the consumer. * @param index the column id for the consumer. */ public DecimalConsumer(DecimalVector vector, int index) { this(vector, index, null); } /** * Constructs a new consumer, with optional coercibility. * * @param vector the underlying vector for the consumer. * @param index the column index for the consumer. * @param bigDecimalRoundingMode java.math.RoundingMode to be applied if the BigDecimal scale does * not match that of the target vector. Set to null to retain strict matching behavior (scale * of source and target vector must match exactly). */ public DecimalConsumer(DecimalVector vector, int index, RoundingMode bigDecimalRoundingMode) { super(vector, index); this.bigDecimalRoundingMode = bigDecimalRoundingMode; this.scale = vector.getScale(); } /** Creates a consumer for {@link DecimalVector}. */ public static JdbcConsumer createConsumer( DecimalVector vector, int index, boolean nullable, RoundingMode bigDecimalRoundingMode) { if (nullable) { return new NullableDecimalConsumer(vector, index, bigDecimalRoundingMode); } else { return new NonNullableDecimalConsumer(vector, index, bigDecimalRoundingMode); } } protected void set(BigDecimal value) { if (bigDecimalRoundingMode != null && value.scale() != scale) { value = value.setScale(scale, bigDecimalRoundingMode); } vector.set(currentIndex, value); } /** Consumer for nullable decimal. */ static class NullableDecimalConsumer extends DecimalConsumer { /** Instantiate a DecimalConsumer. */ public NullableDecimalConsumer( DecimalVector vector, int index, RoundingMode bigDecimalRoundingMode) { super(vector, index, bigDecimalRoundingMode); } @Override public void consume(ResultSet resultSet) throws SQLException { BigDecimal value = resultSet.getBigDecimal(columnIndexInResultSet); if (!resultSet.wasNull()) { // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. set(value); } currentIndex++; } } /** Consumer for non-nullable decimal. */ static class NonNullableDecimalConsumer extends DecimalConsumer { /** Instantiate a DecimalConsumer. */ public NonNullableDecimalConsumer( DecimalVector vector, int index, RoundingMode bigDecimalRoundingMode) { super(vector, index, bigDecimalRoundingMode); } @Override public void consume(ResultSet resultSet) throws SQLException { BigDecimal value = resultSet.getBigDecimal(columnIndexInResultSet); // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. set(value); currentIndex++; } } } DoubleConsumer.java000066400000000000000000000052731515451615200351140ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.vector.Float8Vector; /** * Consumer which consume double type values from {@link ResultSet}. Write the data to {@link * org.apache.arrow.vector.Float8Vector}. */ public class DoubleConsumer { /** Creates a consumer for {@link Float8Vector}. */ public static JdbcConsumer createConsumer( Float8Vector vector, int index, boolean nullable) { if (nullable) { return new NullableDoubleConsumer(vector, index); } else { return new NonNullableDoubleConsumer(vector, index); } } /** Nullable double consumer. */ static class NullableDoubleConsumer extends BaseConsumer { /** Instantiate a DoubleConsumer. */ public NullableDoubleConsumer(Float8Vector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { double value = resultSet.getDouble(columnIndexInResultSet); if (!resultSet.wasNull()) { // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, value); } currentIndex++; } } /** Non-nullable double consumer. */ static class NonNullableDoubleConsumer extends BaseConsumer { /** Instantiate a DoubleConsumer. */ public NonNullableDoubleConsumer(Float8Vector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { double value = resultSet.getDouble(columnIndexInResultSet); // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, value); currentIndex++; } } } FloatConsumer.java000066400000000000000000000052531515451615200347450ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.vector.Float4Vector; /** * Consumer which consume float type values from {@link ResultSet}. Write the data to {@link * org.apache.arrow.vector.Float4Vector}. */ public class FloatConsumer { /** Creates a consumer for {@link Float4Vector}. */ public static JdbcConsumer createConsumer( Float4Vector vector, int index, boolean nullable) { if (nullable) { return new NullableFloatConsumer(vector, index); } else { return new NonNullableFloatConsumer(vector, index); } } /** Nullable float consumer. */ static class NullableFloatConsumer extends BaseConsumer { /** Instantiate a FloatConsumer. */ public NullableFloatConsumer(Float4Vector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { float value = resultSet.getFloat(columnIndexInResultSet); if (!resultSet.wasNull()) { // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, value); } currentIndex++; } } /** Non-nullable float consumer. */ static class NonNullableFloatConsumer extends BaseConsumer { /** Instantiate a FloatConsumer. */ public NonNullableFloatConsumer(Float4Vector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { float value = resultSet.getFloat(columnIndexInResultSet); // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, value); currentIndex++; } } } IntConsumer.java000066400000000000000000000051701515451615200344300ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.vector.IntVector; /** * Consumer which consume int type values from {@link ResultSet}. Write the data to {@link * org.apache.arrow.vector.IntVector}. */ public class IntConsumer { /** Creates a consumer for {@link IntVector}. */ public static JdbcConsumer createConsumer( IntVector vector, int index, boolean nullable) { if (nullable) { return new NullableIntConsumer(vector, index); } else { return new NonNullableIntConsumer(vector, index); } } /** Nullable consumer for int. */ static class NullableIntConsumer extends BaseConsumer { /** Instantiate a IntConsumer. */ public NullableIntConsumer(IntVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { int value = resultSet.getInt(columnIndexInResultSet); if (!resultSet.wasNull()) { // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, value); } currentIndex++; } } /** Non-nullable consumer for int. */ static class NonNullableIntConsumer extends BaseConsumer { /** Instantiate a IntConsumer. */ public NonNullableIntConsumer(IntVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { int value = resultSet.getInt(columnIndexInResultSet); // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, value); currentIndex++; } } } JdbcConsumer.java000066400000000000000000000031271515451615200345400ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.io.IOException; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.vector.ValueVector; /** * An abstraction that is used to consume values from {@link ResultSet}. * * @param The vector within consumer or its delegate, used for partially consume purpose. */ public interface JdbcConsumer extends AutoCloseable { /** Consume a specific type value from {@link ResultSet} and write it to vector. */ void consume(ResultSet resultSet) throws SQLException, IOException; /** Close this consumer, do some clean work such as clear reuse ArrowBuf. */ @Override void close() throws Exception; /** Reset the vector within consumer for partial read purpose. */ void resetValueVector(T vector); } MapConsumer.java000066400000000000000000000074041515451615200344150ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.sql.ResultSet; import java.sql.SQLException; import java.util.Map; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.impl.UnionMapWriter; import org.apache.arrow.vector.util.ObjectMapperFactory; /** * Consumer which consume map type values from {@link ResultSet}. Write the data into {@link * org.apache.arrow.vector.complex.MapVector}. */ public class MapConsumer extends BaseConsumer { private final UnionMapWriter writer; private final ObjectMapper objectMapper = ObjectMapperFactory.newObjectMapper(); private final TypeReference> typeReference = new TypeReference>() {}; private int currentRow; /** Creates a consumer for {@link MapVector}. */ public static MapConsumer createConsumer(MapVector mapVector, int index, boolean nullable) { return new MapConsumer(mapVector, index); } /** Instantiate a MapConsumer. */ public MapConsumer(MapVector vector, int index) { super(vector, index); writer = vector.getWriter(); } @Override public void consume(ResultSet resultSet) throws SQLException, IOException { Object map = resultSet.getObject(columnIndexInResultSet); writer.setPosition(currentRow++); if (map != null) { if (map instanceof String) { writeJavaMapIntoVector(objectMapper.readValue((String) map, typeReference)); } else if (map instanceof Map) { writeJavaMapIntoVector((Map) map); } else { throw new IllegalArgumentException( "Unknown type of map type column from JDBC " + map.getClass().getName()); } } else { writer.writeNull(); } } private void writeJavaMapIntoVector(Map map) { BufferAllocator allocator = vector.getAllocator(); writer.startMap(); map.forEach( (key, value) -> { byte[] keyBytes = key.getBytes(StandardCharsets.UTF_8); byte[] valueBytes = value != null ? value.getBytes(StandardCharsets.UTF_8) : null; try (ArrowBuf keyBuf = allocator.buffer(keyBytes.length); ArrowBuf valueBuf = valueBytes != null ? allocator.buffer(valueBytes.length) : null; ) { writer.startEntry(); keyBuf.writeBytes(keyBytes); writer.key().varChar().writeVarChar(0, keyBytes.length, keyBuf); if (valueBytes != null) { valueBuf.writeBytes(valueBytes); writer.value().varChar().writeVarChar(0, valueBytes.length, valueBuf); } else { writer.value().varChar().writeNull(); } writer.endEntry(); } }); writer.endMap(); } } NullConsumer.java000066400000000000000000000024161515451615200346100ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.vector.NullVector; /** * Consumer which consume null type values from ResultSet. Corresponding to {@link * org.apache.arrow.vector.NullVector}. */ public class NullConsumer extends BaseConsumer { public NullConsumer(NullVector vector) { super(vector, 0); } @Override public void consume(ResultSet resultSet) throws SQLException {} } SmallIntConsumer.java000066400000000000000000000053531515451615200354240ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.vector.SmallIntVector; /** * Consumer which consume smallInt type values from {@link ResultSet}. Write the data to {@link * org.apache.arrow.vector.SmallIntVector}. */ public class SmallIntConsumer { /** Creates a consumer for {@link SmallIntVector}. */ public static BaseConsumer createConsumer( SmallIntVector vector, int index, boolean nullable) { if (nullable) { return new NullableSmallIntConsumer(vector, index); } else { return new NonNullableSmallIntConsumer(vector, index); } } /** Nullable consumer for small int. */ static class NullableSmallIntConsumer extends BaseConsumer { /** Instantiate a SmallIntConsumer. */ public NullableSmallIntConsumer(SmallIntVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { short value = resultSet.getShort(columnIndexInResultSet); if (!resultSet.wasNull()) { // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, value); } currentIndex++; } } /** Non-nullable consumer for small int. */ static class NonNullableSmallIntConsumer extends BaseConsumer { /** Instantiate a SmallIntConsumer. */ public NonNullableSmallIntConsumer(SmallIntVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { short value = resultSet.getShort(columnIndexInResultSet); // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, value); currentIndex++; } } } TimeConsumer.java000066400000000000000000000070761515451615200346030ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Time; import java.util.Calendar; import org.apache.arrow.vector.TimeMilliVector; /** * Consumer which consume time type values from {@link ResultSet}. Write the data to {@link * org.apache.arrow.vector.TimeMilliVector}. */ public abstract class TimeConsumer { /** Creates a consumer for {@link TimeMilliVector}. */ public static JdbcConsumer createConsumer( TimeMilliVector vector, int index, boolean nullable, Calendar calendar) { if (nullable) { return new NullableTimeConsumer(vector, index, calendar); } else { return new NonNullableTimeConsumer(vector, index, calendar); } } /** Nullable consumer for {@link TimeMilliVector}. */ static class NullableTimeConsumer extends BaseConsumer { protected final Calendar calendar; /** Instantiate a TimeConsumer. */ public NullableTimeConsumer(TimeMilliVector vector, int index) { this(vector, index, /* calendar */ null); } /** Instantiate a TimeConsumer. */ public NullableTimeConsumer(TimeMilliVector vector, int index, Calendar calendar) { super(vector, index); this.calendar = calendar; } @Override public void consume(ResultSet resultSet) throws SQLException { Time time = calendar == null ? resultSet.getTime(columnIndexInResultSet) : resultSet.getTime(columnIndexInResultSet, calendar); if (!resultSet.wasNull()) { // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, (int) time.getTime()); } currentIndex++; } } /** Non-nullable consumer for {@link TimeMilliVector}. */ static class NonNullableTimeConsumer extends BaseConsumer { protected final Calendar calendar; /** Instantiate a TimeConsumer. */ public NonNullableTimeConsumer(TimeMilliVector vector, int index) { this(vector, index, /* calendar */ null); } /** Instantiate a TimeConsumer. */ public NonNullableTimeConsumer(TimeMilliVector vector, int index, Calendar calendar) { super(vector, index); this.calendar = calendar; } @Override public void consume(ResultSet resultSet) throws SQLException { Time time = calendar == null ? resultSet.getTime(columnIndexInResultSet) : resultSet.getTime(columnIndexInResultSet, calendar); // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, (int) time.getTime()); currentIndex++; } } } TimestampConsumer.java000066400000000000000000000055531515451615200356460ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Timestamp; import org.apache.arrow.vector.TimeStampMilliVector; /** * Consumer which consume timestamp type values from {@link ResultSet}. Write the data to {@link * TimeStampMilliVector}. */ public abstract class TimestampConsumer { /** Creates a consumer for {@link TimeStampMilliVector}. */ public static JdbcConsumer createConsumer( TimeStampMilliVector vector, int index, boolean nullable) { if (nullable) { return new NullableTimestampConsumer(vector, index); } else { return new NonNullableTimestampConsumer(vector, index); } } /** Nullable consumer for timestamp. */ static class NullableTimestampConsumer extends BaseConsumer { /** Instantiate a TimestampConsumer. */ public NullableTimestampConsumer(TimeStampMilliVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { Timestamp timestamp = resultSet.getTimestamp(columnIndexInResultSet); if (!resultSet.wasNull()) { // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, timestamp.getTime()); } currentIndex++; } } /** Non-nullable consumer for timestamp. */ static class NonNullableTimestampConsumer extends BaseConsumer { /** Instantiate a TimestampConsumer. */ public NonNullableTimestampConsumer(TimeStampMilliVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { Timestamp timestamp = resultSet.getTimestamp(columnIndexInResultSet); // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, timestamp.getTime()); currentIndex++; } } } TimestampTZConsumer.java000066400000000000000000000065561515451615200361300ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Timestamp; import java.util.Calendar; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.TimeStampMilliTZVector; /** * Consumer which consume timestamp (with time zone) type values from {@link ResultSet}. Write the * data to {@link TimeStampMilliTZVector}. */ public class TimestampTZConsumer { /** Creates a consumer for {@link TimeStampMilliTZVector}. */ public static JdbcConsumer createConsumer( TimeStampMilliTZVector vector, int index, boolean nullable, Calendar calendar) { Preconditions.checkArgument(calendar != null, "Calendar cannot be null"); if (nullable) { return new TimestampTZConsumer.NullableTimestampTZConsumer(vector, index, calendar); } else { return new TimestampTZConsumer.NonNullableTimestampConsumer(vector, index, calendar); } } /** Nullable consumer for timestamp (with time zone). */ static class NullableTimestampTZConsumer extends BaseConsumer { protected final Calendar calendar; /** Instantiate a TimestampConsumer. */ public NullableTimestampTZConsumer( TimeStampMilliTZVector vector, int index, Calendar calendar) { super(vector, index); this.calendar = calendar; } @Override public void consume(ResultSet resultSet) throws SQLException { Timestamp timestamp = resultSet.getTimestamp(columnIndexInResultSet, calendar); if (!resultSet.wasNull()) { // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, timestamp.getTime()); } currentIndex++; } } /** Non-nullable consumer for timestamp (with time zone). */ static class NonNullableTimestampConsumer extends BaseConsumer { protected final Calendar calendar; /** Instantiate a TimestampConsumer. */ public NonNullableTimestampConsumer( TimeStampMilliTZVector vector, int index, Calendar calendar) { super(vector, index); this.calendar = calendar; } @Override public void consume(ResultSet resultSet) throws SQLException { Timestamp timestamp = resultSet.getTimestamp(columnIndexInResultSet, calendar); // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, timestamp.getTime()); currentIndex++; } } } TinyIntConsumer.java000066400000000000000000000053331515451615200352750ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.vector.TinyIntVector; /** * Consumer which consume tinyInt type values from {@link ResultSet}. Write the data to {@link * org.apache.arrow.vector.TinyIntVector}. */ public abstract class TinyIntConsumer { /** Creates a consumer for {@link TinyIntVector}. */ public static JdbcConsumer createConsumer( TinyIntVector vector, int index, boolean nullable) { if (nullable) { return new NullableTinyIntConsumer(vector, index); } else { return new NonNullableTinyIntConsumer(vector, index); } } /** Nullable consumer for tiny int. */ static class NullableTinyIntConsumer extends BaseConsumer { /** Instantiate a TinyIntConsumer. */ public NullableTinyIntConsumer(TinyIntVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { byte value = resultSet.getByte(columnIndexInResultSet); if (!resultSet.wasNull()) { // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, value); } currentIndex++; } } /** Non-nullable consumer for tiny int. */ static class NonNullableTinyIntConsumer extends BaseConsumer { /** Instantiate a TinyIntConsumer. */ public NonNullableTinyIntConsumer(TinyIntVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { byte value = resultSet.getByte(columnIndexInResultSet); // for fixed width vectors, we have allocated enough memory proactively, // so there is no need to call the setSafe method here. vector.set(currentIndex, value); currentIndex++; } } } VarCharConsumer.java000066400000000000000000000051631515451615200352260ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import java.nio.charset.StandardCharsets; import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.vector.VarCharVector; /** * Consumer which consume varchar type values from {@link ResultSet}. Write the data to {@link * org.apache.arrow.vector.VarCharVector}. */ public abstract class VarCharConsumer { /** Creates a consumer for {@link VarCharVector}. */ public static JdbcConsumer createConsumer( VarCharVector vector, int index, boolean nullable) { if (nullable) { return new NullableVarCharConsumer(vector, index); } else { return new NonNullableVarCharConsumer(vector, index); } } /** Nullable consumer for var char. */ static class NullableVarCharConsumer extends BaseConsumer { /** Instantiate a VarCharConsumer. */ public NullableVarCharConsumer(VarCharVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { String value = resultSet.getString(columnIndexInResultSet); if (!resultSet.wasNull()) { byte[] bytes = value.getBytes(StandardCharsets.UTF_8); vector.setSafe(currentIndex, bytes); } currentIndex++; } } /** Non-nullable consumer for var char. */ static class NonNullableVarCharConsumer extends BaseConsumer { /** Instantiate a VarCharConsumer. */ public NonNullableVarCharConsumer(VarCharVector vector, int index) { super(vector, index); } @Override public void consume(ResultSet resultSet) throws SQLException { String value = resultSet.getString(columnIndexInResultSet); byte[] bytes = value.getBytes(StandardCharsets.UTF_8); vector.setSafe(currentIndex, bytes); currentIndex++; } } } exceptions/000077500000000000000000000000001515451615200334755ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumerJdbcConsumerException.java000066400000000000000000000034511515451615200406000ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/exceptions/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer.exceptions; import org.apache.arrow.adapter.jdbc.JdbcFieldInfo; import org.apache.arrow.vector.types.pojo.Field; /** * Exception while consuming JDBC data. This exception stores the JdbcFieldInfo for the column and * the ArrowType for the corresponding vector for easier debugging. */ public class JdbcConsumerException extends RuntimeException { final JdbcFieldInfo fieldInfo; final Field field; /** * Construct JdbcConsumerException with all fields. * * @param message error message * @param cause original exception * @param fieldInfo JdbcFieldInfo for the column * @param field ArrowType for the corresponding vector */ public JdbcConsumerException( String message, Throwable cause, JdbcFieldInfo fieldInfo, Field field) { super(message, cause); this.fieldInfo = fieldInfo; this.field = field; } public Field getField() { return this.field; } public JdbcFieldInfo getFieldInfo() { return this.fieldInfo; } } apache-arrow-java-19.0.0/adapter/jdbc/src/test/000077500000000000000000000000001515451615200211665ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/000077500000000000000000000000001515451615200221075ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/000077500000000000000000000000001515451615200226765ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/000077500000000000000000000000001515451615200241175ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200252515ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/000077500000000000000000000000001515451615200266715ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/000077500000000000000000000000001515451615200275735ustar00rootroot00000000000000AbstractJdbcToArrowTest.java000066400000000000000000000371171515451615200350740ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; import java.io.IOException; import java.sql.Connection; import java.sql.DriverManager; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Statement; import java.sql.Types; import java.util.Calendar; import java.util.HashMap; import java.util.Map; import java.util.TimeZone; import java.util.function.Function; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.util.ValueVectorUtility; import org.junit.jupiter.api.AfterEach; /** Class to abstract out some common test functionality for testing JDBC to Arrow. */ public abstract class AbstractJdbcToArrowTest { protected static final String BIGINT = "BIGINT_FIELD5"; protected static final String BINARY = "BINARY_FIELD12"; protected static final String BIT = "BIT_FIELD17"; protected static final String BLOB = "BLOB_FIELD14"; protected static final String BOOL = "BOOL_FIELD2"; protected static final String CHAR = "CHAR_FIELD16"; protected static final String CLOB = "CLOB_FIELD15"; protected static final String DATE = "DATE_FIELD10"; protected static final String DECIMAL = "DECIMAL_FIELD6"; protected static final String DOUBLE = "DOUBLE_FIELD7"; protected static final String INT = "INT_FIELD1"; protected static final String LIST = "LIST_FIELD19"; protected static final String MAP = "MAP_FIELD20"; protected static final String REAL = "REAL_FIELD8"; protected static final String SMALLINT = "SMALLINT_FIELD4"; protected static final String TIME = "TIME_FIELD9"; protected static final String TIMESTAMP = "TIMESTAMP_FIELD11"; protected static final String TINYINT = "TINYINT_FIELD3"; protected static final String VARCHAR = "VARCHAR_FIELD13"; protected static final String NULL = "NULL_FIELD18"; protected static final Map ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP = new HashMap<>(); static { ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP.put(LIST, new JdbcFieldInfo(Types.INTEGER)); } protected Connection conn = null; protected Table table; protected boolean reuseVectorSchemaRoot; /** * This method creates Table object after reading YAML file. * * @param ymlFilePath path to file * @return Table object * @throws IOException on error */ protected static Table getTable(String ymlFilePath, @SuppressWarnings("rawtypes") Class clss) throws IOException { return new ObjectMapper(new YAMLFactory()) .readValue(clss.getClassLoader().getResourceAsStream(ymlFilePath), Table.class); } /** * This method creates Connection object and DB table and also populate data into table for test. * * @throws SQLException on error * @throws ClassNotFoundException on error */ protected void initializeDatabase(Table table) throws SQLException, ClassNotFoundException { this.table = table; TimeZone.setDefault(TimeZone.getTimeZone("UTC")); String url = "jdbc:h2:mem:JdbcToArrowTest"; String driver = "org.h2.Driver"; Class.forName(driver); conn = DriverManager.getConnection(url); try (Statement stmt = conn.createStatement(); ) { stmt.executeUpdate(table.getCreate()); for (String insert : table.getData()) { stmt.executeUpdate(insert); } } } /** * Clean up method to close connection after test completes. * * @throws SQLException on error */ @AfterEach public void destroy() throws SQLException { if (conn != null) { conn.close(); conn = null; } } /** * Prepares test data and returns collection of Table object for each test iteration. * * @param testFiles files for test * @param clss Class type * @return Collection of Table objects * @throws SQLException on error * @throws ClassNotFoundException on error * @throws IOException on error */ public static Object[][] prepareTestData( String[] testFiles, @SuppressWarnings("rawtypes") Class clss) throws SQLException, ClassNotFoundException, IOException { Object[][] tableArr = new Object[testFiles.length][]; int i = 0; for (String testFile : testFiles) { tableArr[i++] = new Object[] {getTable(testFile, clss)}; } return tableArr; } /** * Abstract method to implement test Functionality to test JdbcToArrow methods. * * @param table Table object * @throws SQLException on error * @throws IOException on error */ public abstract void testJdbcToArrowValues(Table table) throws SQLException, IOException, ClassNotFoundException; /** * Abstract method to implement logic to assert test various datatype values. * * @param root VectorSchemaRoot for test * @param isIncludeMapVector is this dataset checks includes map column. Jdbc type to 'map' * mapping declared in configuration only manually */ public abstract void testDataSets(VectorSchemaRoot root, boolean isIncludeMapVector); /** * For the given SQL query, execute and fetch the data from Relational DB and convert it to Arrow * objects. This method uses the default Calendar instance with default TimeZone and Locale as * returned by the JVM. If you wish to use specific TimeZone or Locale for any Date, Time and * Timestamp datasets, you may want use overloaded API that taken Calendar object instance. * *

This method is for test only. * * @param connection Database connection to be used. This method will not close the passed * connection object. Since the caller has passed the connection object it's the * responsibility of the caller to close or return the connection to the pool. * @param query The DB Query to fetch the data. * @param allocator Memory allocator * @return Arrow Data Objects {@link VectorSchemaRoot} * @throws SQLException Propagate any SQL Exceptions to the caller after closing any resources * opened such as ResultSet and Statement objects. */ public VectorSchemaRoot sqlToArrow(Connection connection, String query, BufferAllocator allocator) throws SQLException, IOException { Preconditions.checkNotNull(allocator, "Memory allocator object cannot be null"); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(allocator, JdbcToArrowUtils.getUtcCalendar()) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .build(); return sqlToArrow(connection, query, config); } /** * For the given SQL query, execute and fetch the data from Relational DB and convert it to Arrow * objects. * *

This method is for test only. * * @param connection Database connection to be used. This method will not close the passed * connection object. Since the caller has passed the connection object it's the * responsibility of the caller to close or return the connection to the pool. * @param query The DB Query to fetch the data. * @param allocator Memory allocator * @param calendar Calendar object to use to handle Date, Time and Timestamp datasets. * @return Arrow Data Objects {@link VectorSchemaRoot} * @throws SQLException Propagate any SQL Exceptions to the caller after closing any resources * opened such as ResultSet and Statement objects. */ public VectorSchemaRoot sqlToArrow( Connection connection, String query, BufferAllocator allocator, Calendar calendar) throws SQLException, IOException { Preconditions.checkNotNull(allocator, "Memory allocator object cannot be null"); Preconditions.checkNotNull(calendar, "Calendar object cannot be null"); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(allocator, calendar) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .build(); return sqlToArrow(connection, query, config); } /** * For the given SQL query, execute and fetch the data from Relational DB and convert it to Arrow * objects. * *

This method is for test only. * * @param connection Database connection to be used. This method will not close the passed * connection object. Since the caller has passed the connection object it's the * responsibility of the caller to close or return the connection to the pool. * @param query The DB Query to fetch the data. * @param config Configuration * @return Arrow Data Objects {@link VectorSchemaRoot} * @throws SQLException Propagate any SQL Exceptions to the caller after closing any resources * opened such as ResultSet and Statement objects. */ public static VectorSchemaRoot sqlToArrow( Connection connection, String query, JdbcToArrowConfig config) throws SQLException, IOException { Preconditions.checkNotNull(connection, "JDBC connection object cannot be null"); Preconditions.checkArgument( query != null && query.length() > 0, "SQL query cannot be null or empty"); try (Statement stmt = connection.createStatement()) { return sqlToArrow(stmt.executeQuery(query), config); } } /** * For the given JDBC {@link ResultSet}, fetch the data from Relational DB and convert it to Arrow * objects. This method uses the default RootAllocator and Calendar object. * *

This method is for test only. * * @param resultSet ResultSet to use to fetch the data from underlying database * @return Arrow Data Objects {@link VectorSchemaRoot} * @throws SQLException on error */ public static VectorSchemaRoot sqlToArrow(ResultSet resultSet) throws SQLException, IOException { Preconditions.checkNotNull(resultSet, "JDBC ResultSet object cannot be null"); return sqlToArrow(resultSet, JdbcToArrowUtils.getUtcCalendar()); } /** * For the given JDBC {@link ResultSet}, fetch the data from Relational DB and convert it to Arrow * objects. * *

This method is for test only. * * @param resultSet ResultSet to use to fetch the data from underlying database * @param allocator Memory allocator * @return Arrow Data Objects {@link VectorSchemaRoot} * @throws SQLException on error */ public static VectorSchemaRoot sqlToArrow(ResultSet resultSet, BufferAllocator allocator) throws SQLException, IOException { Preconditions.checkNotNull(allocator, "Memory Allocator object cannot be null"); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(allocator, JdbcToArrowUtils.getUtcCalendar()) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .build(); return sqlToArrow(resultSet, config); } /** * For the given JDBC {@link ResultSet}, fetch the data from Relational DB and convert it to Arrow * objects. * *

This method is for test only. * * @param resultSet ResultSet to use to fetch the data from underlying database * @param calendar Calendar instance to use for Date, Time and Timestamp datasets, or null * if none. * @return Arrow Data Objects {@link VectorSchemaRoot} * @throws SQLException on error */ public static VectorSchemaRoot sqlToArrow(ResultSet resultSet, Calendar calendar) throws SQLException, IOException { Preconditions.checkNotNull(resultSet, "JDBC ResultSet object cannot be null"); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(new RootAllocator(Integer.MAX_VALUE), calendar) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .build(); return sqlToArrow(resultSet, config); } /** * For the given JDBC {@link ResultSet}, fetch the data from Relational DB and convert it to Arrow * objects. * *

This method is for test only. * * @param resultSet ResultSet to use to fetch the data from underlying database * @param allocator Memory allocator to use. * @param calendar Calendar instance to use for Date, Time and Timestamp datasets, or null * if none. * @return Arrow Data Objects {@link VectorSchemaRoot} * @throws SQLException on error */ public static VectorSchemaRoot sqlToArrow( ResultSet resultSet, BufferAllocator allocator, Calendar calendar) throws SQLException, IOException { Preconditions.checkNotNull(allocator, "Memory Allocator object cannot be null"); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(allocator, calendar) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .build(); return sqlToArrow(resultSet, config); } /** * For the given JDBC {@link ResultSet}, fetch the data from Relational DB and convert it to Arrow * objects. * *

This method is for test only. * * @param resultSet ResultSet to use to fetch the data from underlying database * @param config Configuration of the conversion from JDBC to Arrow. * @return Arrow Data Objects {@link VectorSchemaRoot} * @throws SQLException on error */ public static VectorSchemaRoot sqlToArrow(ResultSet resultSet, JdbcToArrowConfig config) throws SQLException, IOException { Preconditions.checkNotNull(resultSet, "JDBC ResultSet object cannot be null"); Preconditions.checkNotNull(config, "The configuration cannot be null"); VectorSchemaRoot root = VectorSchemaRoot.create( JdbcToArrowUtils.jdbcToArrowSchema(resultSet.getMetaData(), config), config.getAllocator()); if (config.getTargetBatchSize() != JdbcToArrowConfig.NO_LIMIT_BATCH_SIZE) { ValueVectorUtility.preAllocate(root, config.getTargetBatchSize()); } JdbcToArrowUtils.jdbcToArrowVectors(resultSet, root, config); return root; } /** * Register MAP_FIELD20 as ArrowType.Map * * @param calendar Calendar instance to use for Date, Time and Timestamp datasets, or null * if none. * @param rsmd ResultSetMetaData to lookup column name from result set metadata * @return typeConverter instance with mapping column to Map type */ protected Function jdbcToArrowTypeConverter( Calendar calendar, ResultSetMetaData rsmd) { return (jdbcFieldInfo) -> { String columnLabel = null; try { int columnIndex = jdbcFieldInfo.getColumn(); if (columnIndex != 0) { columnLabel = rsmd.getColumnLabel(columnIndex); } } catch (SQLException e) { throw new RuntimeException(e); } if (MAP.equals(columnLabel)) { return new ArrowType.Map(false); } else { return JdbcToArrowUtils.getArrowTypeFromJdbcType(jdbcFieldInfo, calendar); } }; } protected ResultSetMetaData getQueryMetaData(String query) throws SQLException { return conn.createStatement().executeQuery(query).getMetaData(); } } JdbcFieldInfoTest.java000066400000000000000000000030401515451615200336360ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; import static org.junit.jupiter.api.Assertions.assertEquals; import java.sql.Types; import org.junit.jupiter.api.Test; public class JdbcFieldInfoTest { @Test public void testCreateJdbcFieldInfoWithJdbcType() { JdbcFieldInfo fieldInfo = new JdbcFieldInfo(Types.BLOB); assertEquals(Types.BLOB, fieldInfo.getJdbcType()); assertEquals(0, fieldInfo.getPrecision()); assertEquals(0, fieldInfo.getScale()); } @Test public void testCreateJdbcFieldInfoWithJdbcTypePrecisionAndScale() { JdbcFieldInfo fieldInfo = new JdbcFieldInfo(Types.BLOB, 1, 2); assertEquals(Types.BLOB, fieldInfo.getJdbcType()); assertEquals(1, fieldInfo.getPrecision()); assertEquals(2, fieldInfo.getScale()); } } JdbcParameterBinderTest.java000066400000000000000000001166271515451615200350630ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; import static org.assertj.core.api.Assertions.assertThat; import java.math.BigDecimal; import java.nio.charset.StandardCharsets; import java.sql.Date; import java.sql.PreparedStatement; import java.sql.SQLException; import java.sql.Time; import java.sql.Timestamp; import java.sql.Types; import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.Map; import java.util.function.BiConsumer; import org.apache.arrow.adapter.jdbc.binder.ColumnBinder; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; 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.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.LargeVarBinaryVector; import org.apache.arrow.vector.LargeVarCharVector; 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.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.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.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.JsonStringHashMap; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class JdbcParameterBinderTest { private static final long MILLIS_PER_DAY = 86_400_000; BufferAllocator allocator; @BeforeEach void beforeEach() { allocator = new RootAllocator(); } @AfterEach void afterEach() { allocator.close(); } @Test void bindOrder() throws SQLException { final Schema schema = new Schema( Arrays.asList( Field.nullable("ints0", new ArrowType.Int(32, true)), Field.nullable("ints1", new ArrowType.Int(32, true)), Field.nullable("ints2", new ArrowType.Int(32, true)))); try (final MockPreparedStatement statement = new MockPreparedStatement(); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final JdbcParameterBinder binder = JdbcParameterBinder.builder(statement, root) .bind(/* parameterIndex= */ 1, /* columnIndex= */ 2) .bind(/* parameterIndex= */ 2, /* columnIndex= */ 0) .build(); assertThat(binder.next()).isFalse(); final IntVector ints0 = (IntVector) root.getVector(0); final IntVector ints1 = (IntVector) root.getVector(1); final IntVector ints2 = (IntVector) root.getVector(2); ints0.setSafe(0, 4); ints0.setNull(1); ints1.setNull(0); ints1.setSafe(1, -8); ints2.setNull(0); ints2.setSafe(1, 12); root.setRowCount(2); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isNull(); assertThat(statement.getParamType(1)).isEqualTo(Types.INTEGER); assertThat(statement.getParamValue(2)).isEqualTo(4); assertThat(statement.getParam(3)).isNull(); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(12); assertThat(statement.getParamValue(2)).isNull(); assertThat(statement.getParamType(2)).isEqualTo(Types.INTEGER); assertThat(statement.getParam(3)).isNull(); assertThat(binder.next()).isFalse(); binder.reset(); ints0.setNull(0); ints0.setSafe(1, -2); ints2.setNull(0); ints2.setSafe(1, 6); root.setRowCount(2); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isNull(); assertThat(statement.getParamType(1)).isEqualTo(Types.INTEGER); assertThat(statement.getParamValue(2)).isNull(); assertThat(statement.getParamType(2)).isEqualTo(Types.INTEGER); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(6); assertThat(statement.getParamValue(2)).isEqualTo(-2); assertThat(statement.getParam(3)).isNull(); assertThat(binder.next()).isFalse(); } } @Test void customBinder() throws SQLException { final Schema schema = new Schema(Collections.singletonList(Field.nullable("ints0", new ArrowType.Int(32, true)))); try (final MockPreparedStatement statement = new MockPreparedStatement(); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final JdbcParameterBinder binder = JdbcParameterBinder.builder(statement, root) .bind( /* parameterIndex= */ 1, new ColumnBinder() { private final IntVector vector = (IntVector) root.getVector(0); @Override public void bind(PreparedStatement statement, int parameterIndex, int rowIndex) throws SQLException { Integer value = vector.getObject(rowIndex); if (value == null) { statement.setString(parameterIndex, "null"); } else { statement.setString(parameterIndex, Integer.toString(value)); } } @Override public int getJdbcType() { return Types.INTEGER; } @Override public FieldVector getVector() { return vector; } }) .build(); assertThat(binder.next()).isFalse(); final IntVector ints = (IntVector) root.getVector(0); ints.setSafe(0, 4); ints.setNull(1); root.setRowCount(2); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo("4"); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo("null"); assertThat(binder.next()).isFalse(); } } @Test void bool() throws SQLException { testSimpleType( ArrowType.Bool.INSTANCE, Types.BOOLEAN, (BitVector vector, Integer index, Boolean value) -> vector.setSafe(index, value ? 1 : 0), BitVector::setNull, Arrays.asList(true, false, true)); } @Test void int8() throws SQLException { testSimpleType( new ArrowType.Int(8, true), Types.TINYINT, TinyIntVector::setSafe, TinyIntVector::setNull, Arrays.asList(Byte.MAX_VALUE, Byte.MIN_VALUE, (byte) 42)); } @Test void int16() throws SQLException { testSimpleType( new ArrowType.Int(16, true), Types.SMALLINT, SmallIntVector::setSafe, SmallIntVector::setNull, Arrays.asList(Short.MAX_VALUE, Short.MIN_VALUE, (short) 42)); } @Test void int32() throws SQLException { testSimpleType( new ArrowType.Int(32, true), Types.INTEGER, IntVector::setSafe, IntVector::setNull, Arrays.asList(Integer.MAX_VALUE, Integer.MIN_VALUE, 42)); } @Test void int64() throws SQLException { testSimpleType( new ArrowType.Int(64, true), Types.BIGINT, BigIntVector::setSafe, BigIntVector::setNull, Arrays.asList(Long.MAX_VALUE, Long.MIN_VALUE, 42L)); } @Test void float32() throws SQLException { testSimpleType( new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE), Types.REAL, Float4Vector::setSafe, Float4Vector::setNull, Arrays.asList(Float.MIN_VALUE, Float.MAX_VALUE, Float.POSITIVE_INFINITY)); } @Test void float64() throws SQLException { testSimpleType( new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE), Types.DOUBLE, Float8Vector::setSafe, Float8Vector::setNull, Arrays.asList(Double.MIN_VALUE, Double.MAX_VALUE, Double.POSITIVE_INFINITY)); } @Test void time32() throws SQLException { testSimpleType( new ArrowType.Time(TimeUnit.SECOND, 32), Types.TIME, (valueVectors, index, value) -> valueVectors.setSafe(index, (int) (value.getTime() / 1_000)), TimeSecVector::setNull, Arrays.asList(new Time(-128_000), new Time(104_000), new Time(-42_000))); testSimpleType( new ArrowType.Time(TimeUnit.MILLISECOND, 32), Types.TIME, (valueVectors, index, value) -> valueVectors.setSafe(index, (int) value.getTime()), TimeMilliVector::setNull, Arrays.asList(new Time(-128_000), new Time(104_000), new Time(-42_000))); } @Test void time64() throws SQLException { testSimpleType( new ArrowType.Time(TimeUnit.MICROSECOND, 64), Types.TIME, (valueVectors, index, value) -> valueVectors.setSafe(index, (value.getTime() * 1_000)), TimeMicroVector::setNull, Arrays.asList(new Time(-128_000), new Time(104_000), new Time(-42_000))); testSimpleType( new ArrowType.Time(TimeUnit.NANOSECOND, 64), Types.TIME, (valueVectors, index, value) -> valueVectors.setSafe(index, (value.getTime() * 1_000_000)), TimeNanoVector::setNull, Arrays.asList(new Time(-128), new Time(104), new Time(-42))); } @Test void date32() throws SQLException { testSimpleType( new ArrowType.Date(DateUnit.DAY), Types.DATE, (valueVectors, index, value) -> valueVectors.setSafe(index, (int) (value.getTime() / MILLIS_PER_DAY)), DateDayVector::setNull, Arrays.asList( new Date(-5 * MILLIS_PER_DAY), new Date(2 * MILLIS_PER_DAY), new Date(MILLIS_PER_DAY))); } @Test void date64() throws SQLException { testSimpleType( new ArrowType.Date(DateUnit.MILLISECOND), Types.DATE, (valueVectors, index, value) -> valueVectors.setSafe(index, value.getTime()), DateMilliVector::setNull, Arrays.asList( new Date(-5 * MILLIS_PER_DAY), new Date(2 * MILLIS_PER_DAY), new Date(MILLIS_PER_DAY))); } @Test void timestamp() throws SQLException { List values = Arrays.asList(new Timestamp(-128_000), new Timestamp(104_000), new Timestamp(-42_000)); testSimpleType( new ArrowType.Timestamp(TimeUnit.SECOND, null), Types.TIMESTAMP, (valueVectors, index, value) -> valueVectors.setSafe(index, value.getTime() / 1_000), TimeStampSecVector::setNull, values); testSimpleType( new ArrowType.Timestamp(TimeUnit.MILLISECOND, null), Types.TIMESTAMP, (valueVectors, index, value) -> valueVectors.setSafe(index, value.getTime()), TimeStampMilliVector::setNull, values); testSimpleType( new ArrowType.Timestamp(TimeUnit.MICROSECOND, null), Types.TIMESTAMP, (valueVectors, index, value) -> valueVectors.setSafe(index, value.getTime() * 1_000), TimeStampMicroVector::setNull, values); testSimpleType( new ArrowType.Timestamp(TimeUnit.NANOSECOND, null), Types.TIMESTAMP, (valueVectors, index, value) -> valueVectors.setSafe(index, value.getTime() * 1_000_000), TimeStampNanoVector::setNull, values); } @Test void timestampTz() throws SQLException { List values = Arrays.asList(new Timestamp(-128_000), new Timestamp(104_000), new Timestamp(-42_000)); testSimpleType( new ArrowType.Timestamp(TimeUnit.SECOND, "UTC"), Types.TIMESTAMP_WITH_TIMEZONE, (valueVectors, index, value) -> valueVectors.setSafe(index, value.getTime() / 1_000), TimeStampSecTZVector::setNull, values); testSimpleType( new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC"), Types.TIMESTAMP_WITH_TIMEZONE, (valueVectors, index, value) -> valueVectors.setSafe(index, value.getTime()), TimeStampMilliTZVector::setNull, values); testSimpleType( new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC"), Types.TIMESTAMP_WITH_TIMEZONE, (valueVectors, index, value) -> valueVectors.setSafe(index, value.getTime() * 1_000), TimeStampMicroTZVector::setNull, values); testSimpleType( new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC"), Types.TIMESTAMP_WITH_TIMEZONE, (valueVectors, index, value) -> valueVectors.setSafe(index, value.getTime() * 1_000_000), TimeStampNanoTZVector::setNull, values); } @Test void utf8() throws SQLException { testSimpleType( ArrowType.Utf8.INSTANCE, Types.VARCHAR, (VarCharVector vector, Integer index, String value) -> vector.setSafe(index, value.getBytes(StandardCharsets.UTF_8)), BaseVariableWidthVector::setNull, Arrays.asList("", "foobar", "abc")); } @Test void largeUtf8() throws SQLException { testSimpleType( ArrowType.LargeUtf8.INSTANCE, Types.LONGVARCHAR, (LargeVarCharVector vector, Integer index, String value) -> vector.setSafe(index, value.getBytes(StandardCharsets.UTF_8)), BaseLargeVariableWidthVector::setNull, Arrays.asList("", "foobar", "abc")); } @Test void binary() throws SQLException { testSimpleType( ArrowType.Binary.INSTANCE, Types.VARBINARY, (VarBinaryVector vector, Integer index, byte[] value) -> vector.setSafe(index, value), BaseVariableWidthVector::setNull, Arrays.asList(new byte[0], new byte[] {2, -4}, new byte[] {0, -1, 127, -128})); } @Test void largeBinary() throws SQLException { testSimpleType( ArrowType.LargeBinary.INSTANCE, Types.LONGVARBINARY, (LargeVarBinaryVector vector, Integer index, byte[] value) -> vector.setSafe(index, value), BaseLargeVariableWidthVector::setNull, Arrays.asList(new byte[0], new byte[] {2, -4}, new byte[] {0, -1, 127, -128})); } @Test void fixedSizeBinary() throws SQLException { testSimpleType( new ArrowType.FixedSizeBinary(3), Types.BINARY, FixedSizeBinaryVector::setSafe, FixedSizeBinaryVector::setNull, Arrays.asList(new byte[3], new byte[] {1, 2, -4}, new byte[] {-1, 127, -128})); } @Test void decimal128() throws SQLException { testSimpleType( new ArrowType.Decimal(/*precision*/ 12, /*scale*/ 3, 128), Types.DECIMAL, DecimalVector::setSafe, DecimalVector::setNull, Arrays.asList( new BigDecimal("120.429"), new BigDecimal("-10590.123"), new BigDecimal("0.000"))); } @Test void decimal256() throws SQLException { testSimpleType( new ArrowType.Decimal(/*precision*/ 12, /*scale*/ 3, 256), Types.DECIMAL, Decimal256Vector::setSafe, Decimal256Vector::setNull, Arrays.asList( new BigDecimal("120.429"), new BigDecimal("-10590.123"), new BigDecimal("0.000"))); } @Test void listOfDouble() throws SQLException { TriConsumer setValue = (listVector, index, values) -> { org.apache.arrow.vector.complex.impl.UnionListWriter writer = listVector.getWriter(); writer.setPosition(index); writer.startList(); Arrays.stream(values).forEach(doubleValue -> writer.float8().writeFloat8(doubleValue)); writer.endList(); listVector.setLastSet(index); }; List values = Arrays.asList( new Double[] {0.0, Math.PI}, new Double[] {1.1, -352346.2, 2355.6}, new Double[] {-1024.3}, new Double[] {}); testListType( new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE), setValue, ListVector::setNull, values); } @Test void listOfInt64() throws SQLException { TriConsumer setValue = (listVector, index, values) -> { org.apache.arrow.vector.complex.impl.UnionListWriter writer = listVector.getWriter(); writer.setPosition(index); writer.startList(); Arrays.stream(values).forEach(longValue -> writer.bigInt().writeBigInt(longValue)); writer.endList(); listVector.setLastSet(index); }; List values = Arrays.asList( new Long[] {1L, 2L, 3L}, new Long[] {4L, 5L}, new Long[] {512L, 1024L, 2048L, 4096L}, new Long[] {}); testListType((ArrowType) new ArrowType.Int(64, true), setValue, ListVector::setNull, values); } @Test void listOfInt32() throws SQLException { TriConsumer setValue = (listVector, index, values) -> { org.apache.arrow.vector.complex.impl.UnionListWriter writer = listVector.getWriter(); writer.setPosition(index); writer.startList(); Arrays.stream(values).forEach(integerValue -> writer.integer().writeInt(integerValue)); writer.endList(); listVector.setLastSet(index); }; List values = Arrays.asList( new Integer[] {1, 2, 3}, new Integer[] {4, 5}, new Integer[] {512, 1024, 2048, 4096}, new Integer[] {}); testListType((ArrowType) new ArrowType.Int(32, true), setValue, ListVector::setNull, values); } @Test void listOfBoolean() throws SQLException { TriConsumer setValue = (listVector, index, values) -> { org.apache.arrow.vector.complex.impl.UnionListWriter writer = listVector.getWriter(); writer.setPosition(index); writer.startList(); Arrays.stream(values) .forEach(booleanValue -> writer.bit().writeBit(booleanValue ? 1 : 0)); writer.endList(); listVector.setLastSet(index); }; List values = Arrays.asList( new Boolean[] {true, false}, new Boolean[] {false, false}, new Boolean[] {true, true, false, true}, new Boolean[] {}); testListType((ArrowType) new ArrowType.Bool(), setValue, ListVector::setNull, values); } @Test void listOfString() throws SQLException { TriConsumer setValue = (listVector, index, values) -> { org.apache.arrow.vector.complex.impl.UnionListWriter writer = listVector.getWriter(); writer.setPosition(index); writer.startList(); Arrays.stream(values) .forEach( stringValue -> { if (stringValue != null) { byte[] stringValueBytes = stringValue.getBytes(StandardCharsets.UTF_8); try (ArrowBuf stringBuffer = allocator.buffer(stringValueBytes.length)) { stringBuffer.writeBytes(stringValueBytes); writer.varChar().writeVarChar(0, stringValueBytes.length, stringBuffer); } } else { writer.varChar().writeNull(); } }); writer.endList(); listVector.setLastSet(index); }; List values = Arrays.asList( new String[] {"aaaa", "b1"}, new String[] {"c", null, "d"}, new String[] {"e", "f", "g", "h"}, new String[] {}); testListType((ArrowType) new ArrowType.Utf8(), setValue, ListVector::setNull, values); } @Test void mapOfString() throws SQLException { TriConsumer> setValue = (mapVector, index, values) -> { org.apache.arrow.vector.complex.impl.UnionMapWriter mapWriter = mapVector.getWriter(); mapWriter.setPosition(index); mapWriter.startMap(); values .entrySet() .forEach( mapValue -> { if (mapValue != null) { byte[] keyBytes = mapValue.getKey().getBytes(StandardCharsets.UTF_8); byte[] valueBytes = mapValue.getValue().getBytes(StandardCharsets.UTF_8); try (ArrowBuf keyBuf = allocator.buffer(keyBytes.length); ArrowBuf valueBuf = allocator.buffer(valueBytes.length); ) { mapWriter.startEntry(); keyBuf.writeBytes(keyBytes); valueBuf.writeBytes(valueBytes); mapWriter.key().varChar().writeVarChar(0, keyBytes.length, keyBuf); mapWriter.value().varChar().writeVarChar(0, valueBytes.length, valueBuf); mapWriter.endEntry(); } } else { mapWriter.writeNull(); } }); mapWriter.endMap(); }; JsonStringHashMap value1 = new JsonStringHashMap(); value1.put("a", "b"); value1.put("c", "d"); JsonStringHashMap value2 = new JsonStringHashMap(); value2.put("d", "e"); value2.put("f", "g"); value2.put("k", "l"); JsonStringHashMap value3 = new JsonStringHashMap(); value3.put("y", "z"); value3.put("arrow", "cool"); List> values = Arrays.asList(value1, value2, value3, Collections.emptyMap()); testMapType( new ArrowType.Map(true), setValue, MapVector::setNull, values, new ArrowType.Utf8()); } @Test void mapOfInteger() throws SQLException { TriConsumer> setValue = (mapVector, index, values) -> { org.apache.arrow.vector.complex.impl.UnionMapWriter mapWriter = mapVector.getWriter(); mapWriter.setPosition(index); mapWriter.startMap(); values .entrySet() .forEach( mapValue -> { if (mapValue != null) { mapWriter.startEntry(); mapWriter.key().integer().writeInt(mapValue.getKey()); mapWriter.value().integer().writeInt(mapValue.getValue()); mapWriter.endEntry(); } else { mapWriter.writeNull(); } }); mapWriter.endMap(); }; JsonStringHashMap value1 = new JsonStringHashMap(); value1.put(1, 2); value1.put(3, 4); JsonStringHashMap value2 = new JsonStringHashMap(); value2.put(5, 6); value2.put(7, 8); value2.put(9, 1024); JsonStringHashMap value3 = new JsonStringHashMap(); value3.put(Integer.MIN_VALUE, Integer.MAX_VALUE); value3.put(0, 4096); List> values = Arrays.asList(value1, value2, value3, Collections.emptyMap()); testMapType( new ArrowType.Map(true), setValue, MapVector::setNull, values, new ArrowType.Int(32, true)); } @FunctionalInterface interface TriConsumer { void accept(T value1, U value2, V value3); } void testSimpleType( ArrowType arrowType, int jdbcType, TriConsumer setValue, BiConsumer setNull, List values) throws SQLException { Schema schema = new Schema(Collections.singletonList(Field.nullable("field", arrowType))); try (final MockPreparedStatement statement = new MockPreparedStatement(); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final JdbcParameterBinder binder = JdbcParameterBinder.builder(statement, root).bindAll().build(); assertThat(binder.next()).isFalse(); @SuppressWarnings("unchecked") final V vector = (V) root.getVector(0); final ColumnBinder columnBinder = ColumnBinder.forVector(vector); assertThat(columnBinder.getJdbcType()).isEqualTo(jdbcType); setValue.accept(vector, 0, values.get(0)); setValue.accept(vector, 1, values.get(1)); setNull.accept(vector, 2); root.setRowCount(3); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(0)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(1)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isNull(); assertThat(statement.getParamType(1)).isEqualTo(jdbcType); assertThat(binder.next()).isFalse(); binder.reset(); setNull.accept(vector, 0); setValue.accept(vector, 1, values.get(2)); setValue.accept(vector, 2, values.get(0)); setValue.accept(vector, 3, values.get(2)); setValue.accept(vector, 4, values.get(1)); root.setRowCount(5); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isNull(); assertThat(statement.getParamType(1)).isEqualTo(jdbcType); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(2)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(0)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(2)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(1)); assertThat(binder.next()).isFalse(); } // Non-nullable (since some types have a specialized binder) schema = new Schema(Collections.singletonList(Field.notNullable("field", arrowType))); try (final MockPreparedStatement statement = new MockPreparedStatement(); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final JdbcParameterBinder binder = JdbcParameterBinder.builder(statement, root).bindAll().build(); assertThat(binder.next()).isFalse(); @SuppressWarnings("unchecked") final V vector = (V) root.getVector(0); setValue.accept(vector, 0, values.get(0)); setValue.accept(vector, 1, values.get(1)); root.setRowCount(2); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(0)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(1)); assertThat(binder.next()).isFalse(); binder.reset(); setValue.accept(vector, 0, values.get(0)); setValue.accept(vector, 1, values.get(2)); setValue.accept(vector, 2, values.get(0)); setValue.accept(vector, 3, values.get(2)); setValue.accept(vector, 4, values.get(1)); root.setRowCount(5); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(0)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(2)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(0)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(2)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(1)); assertThat(binder.next()).isFalse(); } } void testListType( ArrowType arrowType, TriConsumer setValue, BiConsumer setNull, List values) throws SQLException { int jdbcType = Types.ARRAY; Schema schema = new Schema( Collections.singletonList( new Field( "field", FieldType.nullable(new ArrowType.List()), Collections.singletonList( new Field("element", FieldType.notNullable(arrowType), null))))); try (final MockPreparedStatement statement = new MockPreparedStatement(); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final JdbcParameterBinder binder = JdbcParameterBinder.builder(statement, root).bindAll().build(); assertThat(binder.next()).isFalse(); @SuppressWarnings("unchecked") final V vector = (V) root.getVector(0); final ColumnBinder columnBinder = ColumnBinder.forVector(vector); assertThat(columnBinder.getJdbcType()).isEqualTo(jdbcType); setValue.accept(vector, 0, values.get(0)); setValue.accept(vector, 1, values.get(1)); setNull.accept(vector, 2); root.setRowCount(3); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(0)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(1)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isNull(); assertThat(statement.getParamType(1)).isEqualTo(jdbcType); assertThat(binder.next()).isFalse(); binder.reset(); setNull.accept(vector, 0); setValue.accept(vector, 1, values.get(3)); setValue.accept(vector, 2, values.get(0)); setValue.accept(vector, 3, values.get(2)); setValue.accept(vector, 4, values.get(1)); root.setRowCount(5); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isNull(); assertThat(statement.getParamType(1)).isEqualTo(jdbcType); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(3)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(0)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(2)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(1)); assertThat(binder.next()).isFalse(); } // Non-nullable (since some types have a specialized binder) schema = new Schema( Collections.singletonList( new Field( "field", FieldType.notNullable(new ArrowType.List()), Collections.singletonList( new Field("element", FieldType.notNullable(arrowType), null))))); try (final MockPreparedStatement statement = new MockPreparedStatement(); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final JdbcParameterBinder binder = JdbcParameterBinder.builder(statement, root).bindAll().build(); assertThat(binder.next()).isFalse(); @SuppressWarnings("unchecked") final V vector = (V) root.getVector(0); setValue.accept(vector, 0, values.get(0)); setValue.accept(vector, 1, values.get(1)); root.setRowCount(2); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(0)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(1)); assertThat(binder.next()).isFalse(); binder.reset(); setValue.accept(vector, 0, values.get(0)); setValue.accept(vector, 1, values.get(2)); setValue.accept(vector, 2, values.get(0)); setValue.accept(vector, 3, values.get(2)); setValue.accept(vector, 4, values.get(1)); root.setRowCount(5); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(0)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(2)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(0)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(2)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(1)); assertThat(binder.next()).isFalse(); } } void testMapType( ArrowType arrowType, TriConsumer setValue, BiConsumer setNull, List values, ArrowType elementType) throws SQLException { int jdbcType = Types.VARCHAR; FieldType keyType = new FieldType(false, elementType, null, null); FieldType mapType = new FieldType(false, ArrowType.Struct.INSTANCE, null, null); Schema schema = new Schema( Collections.singletonList( new Field( "field", FieldType.nullable(arrowType), Collections.singletonList( new Field( MapVector.KEY_NAME, mapType, Arrays.asList( new Field(MapVector.KEY_NAME, keyType, null), new Field(MapVector.VALUE_NAME, keyType, null))))))); try (final MockPreparedStatement statement = new MockPreparedStatement(); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final JdbcParameterBinder binder = JdbcParameterBinder.builder(statement, root).bindAll().build(); assertThat(binder.next()).isFalse(); @SuppressWarnings("unchecked") final V vector = (V) root.getVector(0); final ColumnBinder columnBinder = ColumnBinder.forVector(vector); assertThat(columnBinder.getJdbcType()).isEqualTo(jdbcType); setValue.accept(vector, 0, values.get(0)); setValue.accept(vector, 1, values.get(1)); setNull.accept(vector, 2); root.setRowCount(3); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(0).toString()); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(1).toString()); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isNull(); assertThat(statement.getParamType(1)).isEqualTo(jdbcType); assertThat(binder.next()).isFalse(); binder.reset(); setNull.accept(vector, 0); setValue.accept(vector, 1, values.get(3)); setValue.accept(vector, 2, values.get(0)); setValue.accept(vector, 3, values.get(2)); setValue.accept(vector, 4, values.get(1)); root.setRowCount(5); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isNull(); assertThat(statement.getParamType(1)).isEqualTo(jdbcType); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(3).toString()); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(0).toString()); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(2).toString()); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(1).toString()); assertThat(binder.next()).isFalse(); } // Non-nullable (since some types have a specialized binder) schema = new Schema( Collections.singletonList( new Field( "field", FieldType.notNullable(arrowType), Collections.singletonList( new Field( MapVector.KEY_NAME, mapType, Arrays.asList( new Field(MapVector.KEY_NAME, keyType, null), new Field(MapVector.VALUE_NAME, keyType, null))))))); try (final MockPreparedStatement statement = new MockPreparedStatement(); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { @SuppressWarnings("unchecked") final V vector = (V) root.getVector(0); final JdbcParameterBinder binder = JdbcParameterBinder.builder(statement, root) .bind( 1, new org.apache.arrow.adapter.jdbc.binder.MapBinder( (MapVector) vector, Types.OTHER)) .build(); assertThat(binder.next()).isFalse(); setValue.accept(vector, 0, values.get(0)); setValue.accept(vector, 1, values.get(1)); root.setRowCount(2); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(0)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(1)); assertThat(binder.next()).isFalse(); binder.reset(); setValue.accept(vector, 0, values.get(0)); setValue.accept(vector, 1, values.get(2)); setValue.accept(vector, 2, values.get(0)); setValue.accept(vector, 3, values.get(2)); setValue.accept(vector, 4, values.get(1)); root.setRowCount(5); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(0)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(2)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(0)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(2)); assertThat(binder.next()).isTrue(); assertThat(statement.getParamValue(1)).isEqualTo(values.get(1)); assertThat(binder.next()).isFalse(); } } } JdbcToArrowCommentMetadataTest.java000066400000000000000000000242411515451615200363660ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; import static org.assertj.core.api.Assertions.assertThat; import java.sql.Connection; import java.sql.DatabaseMetaData; import java.sql.DriverManager; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Statement; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.Map; import java.util.Set; import org.apache.arrow.memory.RootAllocator; 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.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class JdbcToArrowCommentMetadataTest { private static final String COMMENT = "comment"; // use this metadata key for interoperability with Spark StructType private Connection conn = null; /** * This method creates Connection object and DB table and also populate data into table for test. * * @throws SQLException on error * @throws ClassNotFoundException on error */ @BeforeEach public void setUp() throws SQLException, ClassNotFoundException { String url = "jdbc:h2:mem:JdbcToArrowTest?characterEncoding=UTF-8;INIT=runscript from 'classpath:/h2/comment.sql'"; String driver = "org.h2.Driver"; Class.forName(driver); conn = DriverManager.getConnection(url); } @AfterEach public void tearDown() throws SQLException { if (conn != null) { conn.close(); conn = null; } } private static Field field( String name, boolean nullable, ArrowType type, Map metadata) { return new Field(name, new FieldType(nullable, type, null, metadata), Collections.emptyList()); } private static Map metadata(String... entries) { if (entries.length % 2 != 0) { throw new IllegalArgumentException("Map must have equal number of keys and values"); } final Map result = new HashMap<>(); for (int i = 0; i < entries.length; i += 2) { result.put(entries[i], entries[i + 1]); } return result; } @Test public void schemaComment() throws Exception { boolean includeMetadata = false; Schema schema = getSchemaWithCommentFromQuery(includeMetadata); Schema expectedSchema = new Schema( Arrays.asList( field( "ID", false, Types.MinorType.BIGINT.getType(), metadata("comment", "Record identifier")), field( "NAME", true, Types.MinorType.VARCHAR.getType(), metadata("comment", "Name of record")), field("COLUMN1", true, Types.MinorType.BIT.getType(), metadata()), field( "COLUMNN", true, Types.MinorType.INT.getType(), metadata("comment", "Informative description of columnN"))), metadata("comment", "This is super special table with valuable data")); assertThat(schema).isEqualTo(expectedSchema); } @Test public void schemaCommentWithDatabaseMetadata() throws Exception { boolean includeMetadata = true; Schema schema = getSchemaWithCommentFromQuery(includeMetadata); Schema expectedSchema = new Schema( Arrays.asList( field( "ID", false, Types.MinorType.BIGINT.getType(), metadata( "SQL_CATALOG_NAME", "JDBCTOARROWTEST?CHARACTERENCODING=UTF-8", "SQL_SCHEMA_NAME", "PUBLIC", "SQL_TABLE_NAME", "TABLE1", "SQL_COLUMN_NAME", "ID", "SQL_TYPE", "BIGINT", "comment", "Record identifier")), field( "NAME", true, Types.MinorType.VARCHAR.getType(), metadata( "SQL_CATALOG_NAME", "JDBCTOARROWTEST?CHARACTERENCODING=UTF-8", "SQL_SCHEMA_NAME", "PUBLIC", "SQL_TABLE_NAME", "TABLE1", "SQL_COLUMN_NAME", "NAME", "SQL_TYPE", "CHARACTER VARYING", "comment", "Name of record")), field( "COLUMN1", true, Types.MinorType.BIT.getType(), metadata( "SQL_CATALOG_NAME", "JDBCTOARROWTEST?CHARACTERENCODING=UTF-8", "SQL_SCHEMA_NAME", "PUBLIC", "SQL_TABLE_NAME", "TABLE1", "SQL_COLUMN_NAME", "COLUMN1", "SQL_TYPE", "BOOLEAN")), field( "COLUMNN", true, Types.MinorType.INT.getType(), metadata( "SQL_CATALOG_NAME", "JDBCTOARROWTEST?CHARACTERENCODING=UTF-8", "SQL_SCHEMA_NAME", "PUBLIC", "SQL_TABLE_NAME", "TABLE1", "SQL_COLUMN_NAME", "COLUMNN", "SQL_TYPE", "INTEGER", "comment", "Informative description of columnN"))), metadata("comment", "This is super special table with valuable data")); assertThat(schema).isEqualTo(expectedSchema); /* corresponding Apache Spark DDL after conversion: ID BIGINT NOT NULL COMMENT 'Record identifier', NAME STRING COMMENT 'Name of record', COLUMN1 BOOLEAN, COLUMNN INT COMMENT 'Informative description of columnN' */ assertThat(schema).isEqualTo(expectedSchema); } private Schema getSchemaWithCommentFromQuery(boolean includeMetadata) throws SQLException { DatabaseMetaData metaData = conn.getMetaData(); try (Statement statement = conn.createStatement()) { try (ResultSet resultSet = statement.executeQuery("select * from table1")) { ResultSetMetaData resultSetMetaData = resultSet.getMetaData(); Map> columnCommentByColumnIndex = getColumnComments(metaData, resultSetMetaData); String tableName = getTableNameFromResultSetMetaData(resultSetMetaData); String tableComment = getTableComment(metaData, tableName); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder() .setAllocator(new RootAllocator()) .setSchemaMetadata(Collections.singletonMap(COMMENT, tableComment)) .setColumnMetadataByColumnIndex(columnCommentByColumnIndex) .setIncludeMetadata(includeMetadata) .build(); return JdbcToArrowUtils.jdbcToArrowSchema(resultSetMetaData, config); } } } private String getTableNameFromResultSetMetaData(ResultSetMetaData resultSetMetaData) throws SQLException { Set tablesFromQuery = new HashSet<>(); for (int idx = 1, columnCount = resultSetMetaData.getColumnCount(); idx <= columnCount; idx++) { String tableName = resultSetMetaData.getTableName(idx); if (tableName != null && !tableName.isEmpty()) { tablesFromQuery.add(tableName); } } if (tablesFromQuery.size() == 1) { return tablesFromQuery.iterator().next(); } throw new RuntimeException("Table metadata is absent or ambiguous"); } private Map> getColumnComments( DatabaseMetaData metaData, ResultSetMetaData resultSetMetaData) throws SQLException { Map> columnCommentByColumnIndex = new HashMap<>(); for (int columnIdx = 1, columnCount = resultSetMetaData.getColumnCount(); columnIdx <= columnCount; columnIdx++) { String columnComment = getColumnComment( metaData, resultSetMetaData.getTableName(columnIdx), resultSetMetaData.getColumnName(columnIdx)); if (columnComment != null && !columnComment.isEmpty()) { columnCommentByColumnIndex.put(columnIdx, Collections.singletonMap(COMMENT, columnComment)); } } return columnCommentByColumnIndex; } private String getTableComment(DatabaseMetaData metaData, String tableName) throws SQLException { if (tableName == null || tableName.isEmpty()) { return null; } String comment = null; int rowCount = 0; try (ResultSet tableMetadata = metaData.getTables(null, null, tableName, null)) { if (tableMetadata.next()) { comment = tableMetadata.getString("REMARKS"); rowCount++; } } if (rowCount == 1) { return comment; } if (rowCount > 1) { throw new RuntimeException("Multiple tables found for table name"); } throw new RuntimeException("Table comment not found"); } private String getColumnComment(DatabaseMetaData metaData, String tableName, String columnName) throws SQLException { try (ResultSet tableMetadata = metaData.getColumns(null, null, tableName, columnName)) { if (tableMetadata.next()) { return tableMetadata.getString("REMARKS"); } } return null; } } JdbcToArrowConfigTest.java000066400000000000000000000147221515451615200345330ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.sql.Types; import java.util.Calendar; import java.util.HashMap; import java.util.Locale; import java.util.TimeZone; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.Test; public class JdbcToArrowConfigTest { private static final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); private static final Calendar calendar = Calendar.getInstance(TimeZone.getTimeZone("UTC"), Locale.ROOT); @Test public void testConfigNullArguments() { assertThrows( NullPointerException.class, () -> { new JdbcToArrowConfig(null, null); }); } @Test public void testBuilderNullArguments() { assertThrows( NullPointerException.class, () -> { new JdbcToArrowConfigBuilder(null, null); }); } @Test public void testConfigNullCalendar() { JdbcToArrowConfig config = new JdbcToArrowConfig(allocator, null); assertNull(config.getCalendar()); } @Test public void testBuilderNullCalendar() { JdbcToArrowConfigBuilder builder = new JdbcToArrowConfigBuilder(allocator, null); JdbcToArrowConfig config = builder.build(); assertNull(config.getCalendar()); } @Test public void testConfigNullAllocator() { assertThrows( NullPointerException.class, () -> { new JdbcToArrowConfig(null, calendar); }); } @Test public void testBuilderNullAllocator() { assertThrows( NullPointerException.class, () -> { new JdbcToArrowConfigBuilder(null, calendar); }); } @Test public void testSetNullAllocator() { assertThrows( NullPointerException.class, () -> { JdbcToArrowConfigBuilder builder = new JdbcToArrowConfigBuilder(allocator, calendar); builder.setAllocator(null); }); } @Test public void testSetNullCalendar() { JdbcToArrowConfigBuilder builder = new JdbcToArrowConfigBuilder(allocator, calendar); JdbcToArrowConfig config = builder.setCalendar(null).build(); assertNull(config.getCalendar()); } @Test public void testConfig() { JdbcToArrowConfigBuilder builder = new JdbcToArrowConfigBuilder(allocator, calendar); JdbcToArrowConfig config = builder.build(); assertEquals(allocator, config.getAllocator()); assertEquals(calendar, config.getCalendar()); Calendar newCalendar = Calendar.getInstance(); BufferAllocator newAllocator = new RootAllocator(Integer.SIZE); builder.setAllocator(newAllocator).setCalendar(newCalendar); config = builder.build(); assertEquals(newAllocator, config.getAllocator()); assertEquals(newCalendar, config.getCalendar()); } @Test public void testIncludeMetadata() { JdbcToArrowConfigBuilder builder = new JdbcToArrowConfigBuilder(allocator, calendar, false); JdbcToArrowConfig config = builder.build(); assertFalse(config.shouldIncludeMetadata()); builder.setIncludeMetadata(true); config = builder.build(); assertTrue(config.shouldIncludeMetadata()); config = new JdbcToArrowConfigBuilder(allocator, calendar, true).build(); assertTrue(config.shouldIncludeMetadata()); config = new JdbcToArrowConfig( allocator, calendar, /* include metadata */ true, /* reuse vector schema root */ true, null, null, JdbcToArrowConfig.NO_LIMIT_BATCH_SIZE, null); assertTrue(config.shouldIncludeMetadata()); assertTrue(config.isReuseVectorSchemaRoot()); config = new JdbcToArrowConfig( allocator, calendar, /* include metadata */ false, /* reuse vector schema root */ false, null, null, JdbcToArrowConfig.NO_LIMIT_BATCH_SIZE, null); assertFalse(config.shouldIncludeMetadata()); assertFalse(config.isReuseVectorSchemaRoot()); } @Test public void testArraySubTypes() { JdbcToArrowConfigBuilder builder = new JdbcToArrowConfigBuilder(allocator, calendar, false); JdbcToArrowConfig config = builder.build(); final int columnIndex = 1; final String columnName = "COLUMN"; assertNull(config.getArraySubTypeByColumnIndex(columnIndex)); assertNull(config.getArraySubTypeByColumnName(columnName)); final HashMap indexMapping = new HashMap(); indexMapping.put(2, new JdbcFieldInfo(Types.BIGINT)); final HashMap fieldMapping = new HashMap(); fieldMapping.put("NEW_COLUMN", new JdbcFieldInfo(Types.BINARY)); builder.setArraySubTypeByColumnIndexMap(indexMapping); builder.setArraySubTypeByColumnNameMap(fieldMapping); config = builder.build(); assertNull(config.getArraySubTypeByColumnIndex(columnIndex)); assertNull(config.getArraySubTypeByColumnName(columnName)); indexMapping.put(columnIndex, new JdbcFieldInfo(Types.BIT)); fieldMapping.put(columnName, new JdbcFieldInfo(Types.BLOB)); assertNotNull(config.getArraySubTypeByColumnIndex(columnIndex)); assertEquals(Types.BIT, config.getArraySubTypeByColumnIndex(columnIndex).getJdbcType()); assertEquals(Types.BLOB, config.getArraySubTypeByColumnName(columnName).getJdbcType()); } } JdbcToArrowTestHelper.java000066400000000000000000000413451515451615200345460ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import java.math.BigDecimal; import java.nio.charset.Charset; import java.nio.charset.StandardCharsets; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.util.AbstractMap; import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; import org.apache.arrow.vector.BaseValueVector; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DecimalVector; 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.TimeMilliVector; import org.apache.arrow.vector.TimeStampVector; 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.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.JsonStringArrayList; import org.apache.arrow.vector.util.JsonStringHashMap; import org.apache.arrow.vector.util.ObjectMapperFactory; import org.apache.arrow.vector.util.Text; /** * This is a Helper class which has functionalities to read and assert the values from the given * FieldVector object. */ public class JdbcToArrowTestHelper { public static void assertIntVectorValues(IntVector intVector, int rowCount, Integer[] values) { assertEquals(rowCount, intVector.getValueCount()); for (int j = 0; j < intVector.getValueCount(); j++) { if (values[j] == null) { assertTrue(intVector.isNull(j)); } else { assertEquals(values[j].intValue(), intVector.get(j)); } } } public static void assertBooleanVectorValues( BitVector bitVector, int rowCount, Boolean[] values) { assertEquals(rowCount, bitVector.getValueCount()); for (int j = 0; j < bitVector.getValueCount(); j++) { if (values[j] == null) { assertTrue(bitVector.isNull(j)); } else { assertEquals(values[j].booleanValue(), bitVector.get(j) == 1); } } } public static void assertBitVectorValues(BitVector bitVector, int rowCount, Integer[] values) { assertEquals(rowCount, bitVector.getValueCount()); for (int j = 0; j < bitVector.getValueCount(); j++) { if (values[j] == null) { assertTrue(bitVector.isNull(j)); } else { assertEquals(values[j].intValue(), bitVector.get(j)); } } } public static void assertTinyIntVectorValues( TinyIntVector tinyIntVector, int rowCount, Integer[] values) { assertEquals(rowCount, tinyIntVector.getValueCount()); for (int j = 0; j < tinyIntVector.getValueCount(); j++) { if (values[j] == null) { assertTrue(tinyIntVector.isNull(j)); } else { assertEquals(values[j].intValue(), tinyIntVector.get(j)); } } } public static void assertSmallIntVectorValues( SmallIntVector smallIntVector, int rowCount, Integer[] values) { assertEquals(rowCount, smallIntVector.getValueCount()); for (int j = 0; j < smallIntVector.getValueCount(); j++) { if (values[j] == null) { assertTrue(smallIntVector.isNull(j)); } else { assertEquals(values[j].intValue(), smallIntVector.get(j)); } } } public static void assertBigIntVectorValues( BigIntVector bigIntVector, int rowCount, Long[] values) { assertEquals(rowCount, bigIntVector.getValueCount()); for (int j = 0; j < bigIntVector.getValueCount(); j++) { if (values[j] == null) { assertTrue(bigIntVector.isNull(j)); } else { assertEquals(values[j].longValue(), bigIntVector.get(j)); } } } public static void assertDecimalVectorValues( DecimalVector decimalVector, int rowCount, BigDecimal[] values) { assertEquals(rowCount, decimalVector.getValueCount()); for (int j = 0; j < decimalVector.getValueCount(); j++) { if (values[j] == null) { assertTrue(decimalVector.isNull(j)); } else { assertEquals(values[j].doubleValue(), decimalVector.getObject(j).doubleValue(), 0); } } } public static void assertFloat8VectorValues( Float8Vector float8Vector, int rowCount, Double[] values) { assertEquals(rowCount, float8Vector.getValueCount()); for (int j = 0; j < float8Vector.getValueCount(); j++) { if (values[j] == null) { assertTrue(float8Vector.isNull(j)); } else { assertEquals(values[j], float8Vector.get(j), 0.01); } } } public static void assertFloat4VectorValues( Float4Vector float4Vector, int rowCount, Float[] values) { assertEquals(rowCount, float4Vector.getValueCount()); for (int j = 0; j < float4Vector.getValueCount(); j++) { if (values[j] == null) { assertTrue(float4Vector.isNull(j)); } else { assertEquals(values[j], float4Vector.get(j), 0.01); } } } public static void assertTimeVectorValues( TimeMilliVector timeMilliVector, int rowCount, Long[] values) { assertEquals(rowCount, timeMilliVector.getValueCount()); for (int j = 0; j < timeMilliVector.getValueCount(); j++) { if (values[j] == null) { assertTrue(timeMilliVector.isNull(j)); } else { assertEquals(values[j].longValue(), timeMilliVector.get(j)); } } } public static void assertDateVectorValues( DateDayVector dateDayVector, int rowCount, Integer[] values) { assertEquals(rowCount, dateDayVector.getValueCount()); for (int j = 0; j < dateDayVector.getValueCount(); j++) { if (values[j] == null) { assertTrue(dateDayVector.isNull(j)); } else { assertEquals(values[j].longValue(), dateDayVector.get(j)); } } } public static void assertTimeStampVectorValues( TimeStampVector timeStampVector, int rowCount, Long[] values) { assertEquals(rowCount, timeStampVector.getValueCount()); for (int j = 0; j < timeStampVector.getValueCount(); j++) { if (values[j] == null) { assertTrue(timeStampVector.isNull(j)); } else { assertEquals(values[j].longValue(), timeStampVector.get(j)); } } } public static void assertVarBinaryVectorValues( VarBinaryVector varBinaryVector, int rowCount, byte[][] values) { assertEquals(rowCount, varBinaryVector.getValueCount()); for (int j = 0; j < varBinaryVector.getValueCount(); j++) { if (values[j] == null) { assertTrue(varBinaryVector.isNull(j)); } else { assertArrayEquals(values[j], varBinaryVector.get(j)); } } } public static void assertVarcharVectorValues( VarCharVector varCharVector, int rowCount, byte[][] values) { assertEquals(rowCount, varCharVector.getValueCount()); for (int j = 0; j < varCharVector.getValueCount(); j++) { if (values[j] == null) { assertTrue(varCharVector.isNull(j)); } else { assertArrayEquals(values[j], varCharVector.get(j)); } } } public static void assertNullVectorValues(NullVector vector, int rowCount) { assertEquals(rowCount, vector.getValueCount()); } public static void assertListVectorValues( ListVector listVector, int rowCount, Integer[][] values) { assertEquals(rowCount, listVector.getValueCount()); for (int j = 0; j < listVector.getValueCount(); j++) { if (values[j] == null) { assertTrue(listVector.isNull(j)); } else { List list = (List) listVector.getObject(j); assertEquals(Arrays.asList(values[j]), list); } } } public static void assertMapVectorValues( MapVector mapVector, int rowCount, Map[] values) { assertEquals(rowCount, mapVector.getValueCount()); for (int j = 0; j < mapVector.getValueCount(); j++) { if (values[j] == null) { assertTrue(mapVector.isNull(j)); } else { JsonStringArrayList> actualSource = (JsonStringArrayList>) mapVector.getObject(j); Map actualMap = null; if (actualSource != null && !actualSource.isEmpty()) { actualMap = actualSource.stream() .map( entry -> new AbstractMap.SimpleEntry<>( entry.get("key").toString(), entry.get("value") != null ? entry.get("value").toString() : null)) .collect( HashMap::new, (collector, val) -> collector.put(val.getKey(), val.getValue()), HashMap::putAll); } assertEquals(values[j], actualMap); } } } public static Map[] getMapValues(String[] values, String dataType) { String[] dataArr = getValues(values, dataType); Map[] maps = new Map[dataArr.length]; ObjectMapper objectMapper = ObjectMapperFactory.newObjectMapper(); TypeReference> typeReference = new TypeReference>() {}; for (int idx = 0; idx < dataArr.length; idx++) { String jsonString = dataArr[idx].replace("|", ","); if (!jsonString.isEmpty()) { try { maps[idx] = objectMapper.readValue(jsonString, typeReference); } catch (JsonProcessingException e) { throw new RuntimeException(e); } } } return maps; } public static void assertNullValues(BaseValueVector vector, int rowCount) { assertEquals(rowCount, vector.getValueCount()); for (int j = 0; j < vector.getValueCount(); j++) { assertTrue(vector.isNull(j)); } } public static void assertFieldMetadataIsEmpty(VectorSchemaRoot schema) { assertNotNull(schema); assertNotNull(schema.getSchema()); assertNotNull(schema.getSchema().getFields()); for (Field field : schema.getSchema().getFields()) { assertNotNull(field.getMetadata()); assertEquals(0, field.getMetadata().size()); } } public static void assertFieldMetadataMatchesResultSetMetadata( ResultSetMetaData rsmd, Schema schema) throws SQLException { assertNotNull(schema); assertNotNull(schema.getFields()); assertNotNull(rsmd); List fields = schema.getFields(); assertEquals(rsmd.getColumnCount(), fields.size()); // Vector columns are created in the same order as ResultSet columns. for (int i = 1; i <= rsmd.getColumnCount(); ++i) { Map metadata = fields.get(i - 1).getMetadata(); assertNotNull(metadata); assertEquals(5, metadata.size()); assertEquals(rsmd.getCatalogName(i), metadata.get(Constants.SQL_CATALOG_NAME_KEY)); assertEquals(rsmd.getSchemaName(i), metadata.get(Constants.SQL_SCHEMA_NAME_KEY)); assertEquals(rsmd.getTableName(i), metadata.get(Constants.SQL_TABLE_NAME_KEY)); assertEquals(rsmd.getColumnLabel(i), metadata.get(Constants.SQL_COLUMN_NAME_KEY)); assertEquals(rsmd.getColumnTypeName(i), metadata.get(Constants.SQL_TYPE_KEY)); } } public static Integer[] getIntValues(String[] values, String dataType) { String[] dataArr = getValues(values, dataType); Integer[] valueArr = new Integer[dataArr.length]; int i = 0; for (String data : dataArr) { valueArr[i++] = "null".equals(data.trim()) ? null : Integer.parseInt(data); } return valueArr; } public static Boolean[] getBooleanValues(String[] values, String dataType) { String[] dataArr = getValues(values, dataType); Boolean[] valueArr = new Boolean[dataArr.length]; int i = 0; for (String data : dataArr) { valueArr[i++] = "null".equals(data.trim()) ? null : data.trim().equals("1"); } return valueArr; } public static BigDecimal[] getDecimalValues(String[] values, String dataType) { String[] dataArr = getValues(values, dataType); BigDecimal[] valueArr = new BigDecimal[dataArr.length]; int i = 0; for (String data : dataArr) { valueArr[i++] = "null".equals(data.trim()) ? null : new BigDecimal(data); } return valueArr; } public static Double[] getDoubleValues(String[] values, String dataType) { String[] dataArr = getValues(values, dataType); Double[] valueArr = new Double[dataArr.length]; int i = 0; for (String data : dataArr) { valueArr[i++] = "null".equals(data.trim()) ? null : Double.parseDouble(data); } return valueArr; } public static Float[] getFloatValues(String[] values, String dataType) { String[] dataArr = getValues(values, dataType); Float[] valueArr = new Float[dataArr.length]; int i = 0; for (String data : dataArr) { valueArr[i++] = "null".equals(data.trim()) ? null : Float.parseFloat(data); } return valueArr; } public static Long[] getLongValues(String[] values, String dataType) { String[] dataArr = getValues(values, dataType); Long[] valueArr = new Long[dataArr.length]; int i = 0; for (String data : dataArr) { valueArr[i++] = "null".equals(data.trim()) ? null : Long.parseLong(data); } return valueArr; } public static byte[][] getCharArray(String[] values, String dataType) { String[] dataArr = getValues(values, dataType); byte[][] valueArr = new byte[dataArr.length][]; int i = 0; for (String data : dataArr) { valueArr[i++] = "null".equals(data.trim()) ? null : data.trim().getBytes(StandardCharsets.UTF_8); } return valueArr; } public static byte[][] getCharArrayWithCharSet( String[] values, String dataType, Charset charSet) { String[] dataArr = getValues(values, dataType); byte[][] valueArr = new byte[dataArr.length][]; int i = 0; for (String data : dataArr) { valueArr[i++] = "null".equals(data.trim()) ? null : data.trim().getBytes(charSet); } return valueArr; } public static byte[][] getBinaryValues(String[] values, String dataType) { String[] dataArr = getValues(values, dataType); byte[][] valueArr = new byte[dataArr.length][]; int i = 0; for (String data : dataArr) { valueArr[i++] = "null".equals(data.trim()) ? null : data.trim().getBytes(StandardCharsets.UTF_8); } return valueArr; } @SuppressWarnings("StringSplitter") public static String[] getValues(String[] values, String dataType) { String value = ""; for (String val : values) { if (val.startsWith(dataType)) { value = val.split("=")[1]; break; } } return value.split(","); } public static Integer[][] getListValues(String[] values, String dataType) { String[] dataArr = getValues(values, dataType); return getListValues(dataArr); } @SuppressWarnings("StringSplitter") public static Integer[][] getListValues(String[] dataArr) { Integer[][] valueArr = new Integer[dataArr.length][]; int i = 0; for (String data : dataArr) { if ("null".equals(data.trim())) { valueArr[i++] = null; } else if ("()".equals(data.trim())) { valueArr[i++] = new Integer[0]; } else { String[] row = data.replace("(", "").replace(")", "").split(";"); Integer[] arr = new Integer[row.length]; for (int j = 0; j < arr.length; j++) { arr[j] = "null".equals(row[j]) ? null : Integer.parseInt(row[j]); } valueArr[i++] = arr; } } return valueArr; } } MockPreparedStatement.java000066400000000000000000000356371515451615200346360ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; import java.io.InputStream; import java.io.Reader; import java.math.BigDecimal; import java.net.URL; import java.sql.Array; import java.sql.Blob; import java.sql.Clob; import java.sql.Connection; import java.sql.Date; import java.sql.NClob; import java.sql.ParameterMetaData; import java.sql.PreparedStatement; import java.sql.Ref; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.RowId; import java.sql.SQLException; import java.sql.SQLWarning; import java.sql.SQLXML; import java.sql.Time; import java.sql.Timestamp; import java.util.Calendar; import java.util.HashMap; import java.util.Map; /** A PreparedStatement that just stores parameters set on it. */ public final class MockPreparedStatement implements PreparedStatement { static class ParameterHolder { final Object value; final Integer sqlType; public Calendar calendar; ParameterHolder(Object value, Integer sqlType) { this.value = value; this.sqlType = sqlType; } } private final Map parameters; MockPreparedStatement() { parameters = new HashMap<>(); } ParameterHolder getParam(int index) { return parameters.get(index); } Object getParamValue(int index) { return parameters.get(index).value; } Integer getParamType(int index) { return parameters.get(index).sqlType; } @Override public ResultSet executeQuery() throws SQLException { throw new UnsupportedOperationException(); } @Override public int executeUpdate() throws SQLException { throw new UnsupportedOperationException(); } @Override public void setNull(int parameterIndex, int sqlType) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(null, sqlType)); } @Override public void setBoolean(int parameterIndex, boolean x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setByte(int parameterIndex, byte x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setShort(int parameterIndex, short x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setInt(int parameterIndex, int x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setLong(int parameterIndex, long x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setFloat(int parameterIndex, float x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setDouble(int parameterIndex, double x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setBigDecimal(int parameterIndex, BigDecimal x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setString(int parameterIndex, String x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setBytes(int parameterIndex, byte[] x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setDate(int parameterIndex, Date x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setTime(int parameterIndex, Time x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setTimestamp(int parameterIndex, Timestamp x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setAsciiStream(int parameterIndex, InputStream x, int length) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override @Deprecated public void setUnicodeStream(int parameterIndex, InputStream x, int length) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setBinaryStream(int parameterIndex, InputStream x, int length) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void clearParameters() throws SQLException { parameters.clear(); } @Override public void setObject(int parameterIndex, Object x, int targetSqlType) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, targetSqlType)); } @Override public void setObject(int parameterIndex, Object x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public boolean execute() throws SQLException { throw new UnsupportedOperationException(); } @Override public void addBatch() throws SQLException { throw new UnsupportedOperationException(); } @Override public void setCharacterStream(int parameterIndex, Reader reader, int length) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(reader, null)); } @Override public void setRef(int parameterIndex, Ref x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setBlob(int parameterIndex, Blob x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setClob(int parameterIndex, Clob x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setArray(int parameterIndex, Array x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public ResultSetMetaData getMetaData() throws SQLException { throw new UnsupportedOperationException(); } @Override public void setDate(int parameterIndex, Date x, Calendar cal) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setTime(int parameterIndex, Time x, Calendar cal) throws SQLException {} @Override public void setTimestamp(int parameterIndex, Timestamp x, Calendar cal) throws SQLException { ParameterHolder value = new ParameterHolder(x, null); value.calendar = cal; parameters.put(parameterIndex, value); } @Override public void setNull(int parameterIndex, int sqlType, String typeName) throws SQLException {} @Override public void setURL(int parameterIndex, URL x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public ParameterMetaData getParameterMetaData() throws SQLException { throw new UnsupportedOperationException(); } @Override public void setRowId(int parameterIndex, RowId x) throws SQLException { parameters.put(parameterIndex, new ParameterHolder(x, null)); } @Override public void setNString(int parameterIndex, String value) throws SQLException {} @Override public void setNCharacterStream(int parameterIndex, Reader value, long length) throws SQLException {} @Override public void setNClob(int parameterIndex, NClob value) throws SQLException {} @Override public void setClob(int parameterIndex, Reader reader, long length) throws SQLException {} @Override public void setBlob(int parameterIndex, InputStream inputStream, long length) throws SQLException {} @Override public void setNClob(int parameterIndex, Reader reader, long length) throws SQLException {} @Override public void setSQLXML(int parameterIndex, SQLXML xmlObject) throws SQLException {} @Override public void setObject(int parameterIndex, Object x, int targetSqlType, int scaleOrLength) throws SQLException {} @Override public void setAsciiStream(int parameterIndex, InputStream x, long length) throws SQLException {} @Override public void setBinaryStream(int parameterIndex, InputStream x, long length) throws SQLException {} @Override public void setCharacterStream(int parameterIndex, Reader reader, long length) throws SQLException {} @Override public void setAsciiStream(int parameterIndex, InputStream x) throws SQLException {} @Override public void setBinaryStream(int parameterIndex, InputStream x) throws SQLException {} @Override public void setCharacterStream(int parameterIndex, Reader reader) throws SQLException {} @Override public void setNCharacterStream(int parameterIndex, Reader value) throws SQLException {} @Override public void setClob(int parameterIndex, Reader reader) throws SQLException {} @Override public void setBlob(int parameterIndex, InputStream inputStream) throws SQLException {} @Override public void setNClob(int parameterIndex, Reader reader) throws SQLException {} @Override public ResultSet executeQuery(String sql) throws SQLException { throw new UnsupportedOperationException(); } @Override public int executeUpdate(String sql) throws SQLException { throw new UnsupportedOperationException(); } @Override public void close() throws SQLException {} @Override public int getMaxFieldSize() throws SQLException { throw new UnsupportedOperationException(); } @Override public void setMaxFieldSize(int max) throws SQLException { throw new UnsupportedOperationException(); } @Override public int getMaxRows() throws SQLException { throw new UnsupportedOperationException(); } @Override public void setMaxRows(int max) throws SQLException { throw new UnsupportedOperationException(); } @Override public void setEscapeProcessing(boolean enable) throws SQLException { throw new UnsupportedOperationException(); } @Override public int getQueryTimeout() throws SQLException { throw new UnsupportedOperationException(); } @Override public void setQueryTimeout(int seconds) throws SQLException { throw new UnsupportedOperationException(); } @Override public void cancel() throws SQLException { throw new UnsupportedOperationException(); } @Override public SQLWarning getWarnings() throws SQLException { throw new UnsupportedOperationException(); } @Override public void clearWarnings() throws SQLException { throw new UnsupportedOperationException(); } @Override public void setCursorName(String name) throws SQLException { throw new UnsupportedOperationException(); } @Override public boolean execute(String sql) throws SQLException { throw new UnsupportedOperationException(); } @Override public ResultSet getResultSet() throws SQLException { throw new UnsupportedOperationException(); } @Override public int getUpdateCount() throws SQLException { throw new UnsupportedOperationException(); } @Override public boolean getMoreResults() throws SQLException { throw new UnsupportedOperationException(); } @Override public void setFetchDirection(int direction) throws SQLException { throw new UnsupportedOperationException(); } @Override public int getFetchDirection() throws SQLException { throw new UnsupportedOperationException(); } @Override public void setFetchSize(int rows) throws SQLException { throw new UnsupportedOperationException(); } @Override public int getFetchSize() throws SQLException { throw new UnsupportedOperationException(); } @Override public int getResultSetConcurrency() throws SQLException { throw new UnsupportedOperationException(); } @Override public int getResultSetType() throws SQLException { throw new UnsupportedOperationException(); } @Override public void addBatch(String sql) throws SQLException { throw new UnsupportedOperationException(); } @Override public void clearBatch() throws SQLException { throw new UnsupportedOperationException(); } @Override public int[] executeBatch() throws SQLException { throw new UnsupportedOperationException(); } @Override public Connection getConnection() throws SQLException { throw new UnsupportedOperationException(); } @Override public boolean getMoreResults(int current) throws SQLException { throw new UnsupportedOperationException(); } @Override public ResultSet getGeneratedKeys() throws SQLException { throw new UnsupportedOperationException(); } @Override public int executeUpdate(String sql, int autoGeneratedKeys) throws SQLException { throw new UnsupportedOperationException(); } @Override public int executeUpdate(String sql, int[] columnIndexes) throws SQLException { throw new UnsupportedOperationException(); } @Override public int executeUpdate(String sql, String[] columnNames) throws SQLException { throw new UnsupportedOperationException(); } @Override public boolean execute(String sql, int autoGeneratedKeys) throws SQLException { throw new UnsupportedOperationException(); } @Override public boolean execute(String sql, int[] columnIndexes) throws SQLException { throw new UnsupportedOperationException(); } @Override public boolean execute(String sql, String[] columnNames) throws SQLException { throw new UnsupportedOperationException(); } @Override public int getResultSetHoldability() throws SQLException { throw new UnsupportedOperationException(); } @Override public boolean isClosed() throws SQLException { throw new UnsupportedOperationException(); } @Override public void setPoolable(boolean poolable) throws SQLException { throw new UnsupportedOperationException(); } @Override public boolean isPoolable() throws SQLException { throw new UnsupportedOperationException(); } @Override public void closeOnCompletion() throws SQLException { throw new UnsupportedOperationException(); } @Override public boolean isCloseOnCompletion() throws SQLException { throw new UnsupportedOperationException(); } @Override public T unwrap(Class iface) throws SQLException { throw new UnsupportedOperationException(); } @Override public boolean isWrapperFor(Class iface) throws SQLException { throw new UnsupportedOperationException(); } } ResultSetUtility.java000066400000000000000000001316501515451615200337030ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; import java.io.InputStream; import java.io.Reader; import java.math.BigDecimal; import java.net.URL; import java.sql.Array; import java.sql.Blob; import java.sql.Clob; import java.sql.Date; import java.sql.NClob; import java.sql.Ref; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.RowId; import java.sql.SQLException; import java.sql.SQLType; import java.sql.SQLWarning; import java.sql.SQLXML; import java.sql.Statement; import java.sql.Time; import java.sql.Timestamp; import java.sql.Types; import java.util.ArrayList; import java.util.Calendar; import java.util.List; import java.util.Map; public class ResultSetUtility { public static ResultSet generateEmptyResultSet() throws SQLException { MockDataElement element = new MockDataElement("string_example"); MockResultSetMetaData.MockColumnMetaData columnMetaData = MockResultSetMetaData.MockColumnMetaData.fromDataElement(element, 1); ArrayList cols = new ArrayList<>(); cols.add(columnMetaData); ResultSetMetaData metadata = new MockResultSetMetaData(cols); return MockResultSet.builder().setMetaData(metadata).build(); } public static MockResultSet generateBasicResultSet(int rows) throws SQLException { MockResultSet.Builder builder = MockResultSet.builder(); for (int i = 0; i < rows; i++) { builder.addDataElement("row number: " + (i + 1)).addDataElement("data").finishRow(); } return builder.build(); } public static class MockResultSet extends ThrowingResultSet { private final List rows; private int index = 0; private boolean isClosed = false; private ResultSetMetaData metadata; private boolean wasNull; public MockResultSet(List rows) throws SQLException { this(rows, MockResultSetMetaData.fromRows(rows)); } public MockResultSet(List rows, ResultSetMetaData metadata) { this.rows = rows; this.metadata = metadata; this.wasNull = false; } public static Builder builder() { return new Builder(); } private void throwIfClosed() throws SQLException { if (isClosed) { throw new SQLException("ResultSet is already closed!"); } } private void setWasNull(MockDataElement element) { wasNull = element.isNull(); } @Override public boolean next() throws SQLException { throwIfClosed(); index++; return index <= rows.size(); } @Override public void close() throws SQLException { throwIfClosed(); isClosed = true; } @Override public boolean isBeforeFirst() throws SQLException { throwIfClosed(); return index == 0; } @Override public boolean isAfterLast() throws SQLException { return index > rows.size(); } @Override public boolean isFirst() throws SQLException { return index == 1; } @Override public boolean isLast() throws SQLException { return index == rows.size(); } @Override public void beforeFirst() throws SQLException { index = 0; } @Override public void afterLast() throws SQLException { index = rows.size(); } private MockRow getCurrentRow() throws SQLException { throwIfClosed(); if (index == 0) { throw new SQLException("Index is before first element!"); } if (index <= rows.size()) { return rows.get(index - 1); } throw new SQLException("Unable to fetch row at index: " + index); } private MockDataElement getDataElementAtCol(int idx) throws SQLException { MockRow row = getCurrentRow(); MockDataElement element = row.getDataElementAtIndex(idx - 1); setWasNull(element); return element; } @Override public String getString(int idx) throws SQLException { return getDataElementAtCol(idx).getString(); } @Override public boolean getBoolean(int idx) throws SQLException { return getDataElementAtCol(idx).getBoolean(); } @Override public short getShort(int idx) throws SQLException { return getDataElementAtCol(idx).getShort(); } @Override public int getInt(int idx) throws SQLException { return getDataElementAtCol(idx).getInt(); } @Override public long getLong(int idx) throws SQLException { return getDataElementAtCol(idx).getLong(); } @Override public float getFloat(int idx) throws SQLException { return getDataElementAtCol(idx).getFloat(); } @Override public double getDouble(int idx) throws SQLException { return getDataElementAtCol(idx).getDouble(); } @Override public BigDecimal getBigDecimal(int idx) throws SQLException { return getDataElementAtCol(idx).getBigDecimal(); } @Override public Date getDate(int idx) throws SQLException { return getDataElementAtCol(idx).getDate(); } @Override public Time getTime(int idx) throws SQLException { return getDataElementAtCol(idx).getTime(); } @Override public Timestamp getTimestamp(int idx) throws SQLException { return getDataElementAtCol(idx).getTimestamp(); } @Override public ResultSetMetaData getMetaData() throws SQLException { return metadata; } @Override public boolean wasNull() throws SQLException { return wasNull; } public static class Builder { private final ArrayList rows; private ArrayList bufferedElements; private ResultSetMetaData metadata; Builder() { this.rows = new ArrayList<>(); this.bufferedElements = new ArrayList<>(); } public Builder finishRow() { rows.add(new MockRow(this.bufferedElements)); this.bufferedElements = new ArrayList<>(); return this; } public Builder addDataElement(MockDataElement element) { this.bufferedElements.add(element); return this; } public Builder addDataElement(String str) { return this.addDataElement(new MockDataElement(str)); } public Builder addDataElement(Object val, int sqlType) { return this.addDataElement(new MockDataElement(val, sqlType)); } public Builder setMetaData(ResultSetMetaData metadata) { this.metadata = metadata; return this; } public MockResultSet build() throws SQLException { if (this.metadata == null) { return new MockResultSet(this.rows); } return new MockResultSet(this.rows, this.metadata); } } } public static class MockResultSetMetaData extends ThrowingResultSetMetaData { private final List columns; public MockResultSetMetaData(List columns) { this.columns = columns; } @Override public int getColumnCount() throws SQLException { return columns.size(); } @Override public String getColumnLabel(int column) throws SQLException { return columns.get(column - 1).getLabel(); } @Override public String getColumnName(int column) throws SQLException { return columns.get(column - 1).getName(); } @Override public int getColumnType(int column) throws SQLException { return columns.get(column - 1).getType(); } @Override public int getPrecision(int column) throws SQLException { return columns.get(column - 1).getPrecision(); } @Override public int getScale(int column) throws SQLException { return columns.get(column - 1).getScale(); } @Override public int isNullable(int column) throws SQLException { return columns.get(column - 1).isNullable(); } @Override public int getColumnDisplaySize(int column) throws SQLException { return columns.get(column - 1).getDisplaySize(); } @Override public String getColumnTypeName(int column) throws SQLException { return columns.get(column - 1).getTypeName(); } public static MockResultSetMetaData fromRows(List rows) throws SQLException { // Note: This attempts to dynamically construct ResultSetMetaData from the first row in a // given result set. // If there are now rows, or the result set contains no columns, this cannot be dynamically // generated and // an exception will be thrown. if (rows.size() == 0) { throw new SQLException( "Unable to dynamically generate ResultSetMetaData because row count is zero!"); } MockRow firstRow = rows.get(0); if (firstRow.dataElements.size() == 0) { throw new SQLException( "Unable to dynamically generate ResultSetMetaData because column count is zero!"); } ArrayList columns = new ArrayList<>(); for (int i = 0; i < firstRow.dataElements.size(); i++) { MockDataElement element = firstRow.getDataElementAtIndex(i); columns.add(MockColumnMetaData.fromDataElement(element, i)); } return new MockResultSetMetaData(columns); } public static class MockColumnMetaData { private int sqlType; private int precision; private int scale; private int nullable; private String label; private String typeName; private int displaySize; private MockColumnMetaData() {} private String getLabel() { return label; } private String getName() { return getLabel(); } private int getType() { return sqlType; } private int getPrecision() { return precision; } private int getScale() { return scale; } private int isNullable() { return nullable; } private String getTypeName() { return typeName; } private int getDisplaySize() { return displaySize; } public static MockColumnMetaData fromDataElement(MockDataElement element, int i) throws SQLException { return MockColumnMetaData.builder() .sqlType(element.getSqlType()) .precision(element.getPrecision()) .scale(element.getScale()) .nullable(element.isNullable()) .setTypeName("TYPE") .setDisplaySize(420) .label("col_" + i) .build(); } public static Builder builder() { return new Builder(); } public static class Builder { private MockColumnMetaData columnMetaData = new MockColumnMetaData(); public Builder label(String label) { this.columnMetaData.label = label; return this; } public Builder sqlType(int sqlType) { this.columnMetaData.sqlType = sqlType; return this; } public Builder precision(int precision) { this.columnMetaData.precision = precision; return this; } public Builder scale(int scale) { this.columnMetaData.scale = scale; return this; } public Builder nullable(int nullable) { this.columnMetaData.nullable = nullable; return this; } public Builder setTypeName(String typeName) { this.columnMetaData.typeName = typeName; return this; } public Builder setDisplaySize(int displaySize) { this.columnMetaData.displaySize = displaySize; return this; } public MockColumnMetaData build() { return this.columnMetaData; } } } } public static class MockRow { private final List dataElements; public MockRow(List elements) { this.dataElements = elements; } public MockDataElement getDataElementAtIndex(int idx) throws SQLException { if (idx > dataElements.size()) { throw new SQLException("Unable to find data element at position: " + idx); } return dataElements.get(idx); } } public static class MockDataElement { private final Object value; private final int sqlType; public MockDataElement(String val) { this(val, Types.VARCHAR); } public MockDataElement(Object val, int sqlType) { this.value = val; this.sqlType = sqlType; } private boolean isNull() { return value == null; } private String getValueAsString() { return value.toString(); } private int getPrecision() throws SQLException { switch (this.sqlType) { case Types.VARCHAR: return getValueAsString().length(); case Types.DECIMAL: return getBigDecimal().precision(); default: throw getExceptionToThrow("Unable to determine precision for data type: " + sqlType); } } private int getScale() throws SQLException { switch (this.sqlType) { case Types.VARCHAR: return 0; case Types.DECIMAL: return getBigDecimal().scale(); default: throw getExceptionToThrow("Unable to determine scale for data type!"); } } private int isNullable() throws SQLException { switch (this.sqlType) { case Types.VARCHAR: case Types.DECIMAL: return ResultSetMetaData.columnNullable; default: return ResultSetMetaData.columnNullableUnknown; } } private int getSqlType() throws SQLException { return this.sqlType; } public BigDecimal getBigDecimal() throws SQLException { if (value == null) { return null; } try { return new BigDecimal(getValueAsString()); } catch (Exception ex) { throw new SQLException(ex); } } public String getString() throws SQLException { if (value == null) { return null; } return getValueAsString(); } public boolean getBoolean() throws SQLException { if (value == null) { return false; } try { return (boolean) value; } catch (Exception ex) { throw new SQLException(ex); } } public int getInt() throws SQLException { if (value == null) { return 0; } try { return Integer.parseInt(getValueAsString()); } catch (Exception ex) { throw new SQLException(ex); } } public long getLong() throws SQLException { if (value == null) { return 0L; } try { return Long.parseLong(getValueAsString()); } catch (Exception ex) { throw new SQLException(ex); } } public double getDouble() throws SQLException { if (value == null) { return 0.0; } try { return Double.parseDouble(getValueAsString()); } catch (Exception ex) { throw new SQLException(ex); } } public Date getDate() throws SQLException { if (value == null) { return null; } try { return Date.valueOf(getValueAsString()); } catch (Exception ex) { throw new SQLException(ex); } } public Time getTime() throws SQLException { if (value == null) { return null; } try { return Time.valueOf(getValueAsString()); } catch (Exception ex) { throw new SQLException(ex); } } public Timestamp getTimestamp() throws SQLException { if (value == null) { return null; } try { return Timestamp.valueOf(getValueAsString()); } catch (Exception ex) { throw new SQLException(ex); } } public float getFloat() throws SQLException { if (value == null) { return 0.0f; } try { return Float.parseFloat(getValueAsString()); } catch (Exception ex) { throw new SQLException(ex); } } public short getShort() throws SQLException { if (value == null) { return 0; } try { return Short.parseShort(getValueAsString()); } catch (Exception ex) { throw new SQLException(ex); } } } public static class ThrowingResultSet implements ResultSet { @Override public boolean next() throws SQLException { throw getExceptionToThrow(); } @Override public void close() throws SQLException { throw getExceptionToThrow(); } @Override public boolean wasNull() throws SQLException { throw getExceptionToThrow(); } @Override public String getString(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public boolean getBoolean(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public byte getByte(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public short getShort(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public int getInt(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public long getLong(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public float getFloat(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public double getDouble(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public BigDecimal getBigDecimal(int columnIndex, int scale) throws SQLException { throw getExceptionToThrow(); } @Override public byte[] getBytes(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public Date getDate(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public Time getTime(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public Timestamp getTimestamp(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public InputStream getAsciiStream(int columnIndex) throws SQLException { return null; } @Override public InputStream getUnicodeStream(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public InputStream getBinaryStream(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public String getString(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public boolean getBoolean(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public byte getByte(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public short getShort(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public int getInt(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public long getLong(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public float getFloat(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public double getDouble(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public BigDecimal getBigDecimal(String columnLabel, int scale) throws SQLException { throw getExceptionToThrow(); } @Override public byte[] getBytes(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public Date getDate(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public Time getTime(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public Timestamp getTimestamp(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public InputStream getAsciiStream(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public InputStream getUnicodeStream(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public InputStream getBinaryStream(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public SQLWarning getWarnings() throws SQLException { throw getExceptionToThrow(); } @Override public void clearWarnings() throws SQLException { throw getExceptionToThrow(); } @Override public String getCursorName() throws SQLException { throw getExceptionToThrow(); } @Override public ResultSetMetaData getMetaData() throws SQLException { throw getExceptionToThrow(); } @Override public Object getObject(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public Object getObject(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public int findColumn(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public Reader getCharacterStream(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public Reader getCharacterStream(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public BigDecimal getBigDecimal(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public BigDecimal getBigDecimal(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public boolean isBeforeFirst() throws SQLException { throw getExceptionToThrow(); } @Override public boolean isAfterLast() throws SQLException { throw getExceptionToThrow(); } @Override public boolean isFirst() throws SQLException { throw getExceptionToThrow(); } @Override public boolean isLast() throws SQLException { throw getExceptionToThrow(); } @Override public void beforeFirst() throws SQLException { throw getExceptionToThrow(); } @Override public void afterLast() throws SQLException { throw getExceptionToThrow(); } @Override public boolean first() throws SQLException { throw getExceptionToThrow(); } @Override public boolean last() throws SQLException { throw getExceptionToThrow(); } @Override public int getRow() throws SQLException { throw getExceptionToThrow(); } @Override public boolean absolute(int row) throws SQLException { throw getExceptionToThrow(); } @Override public boolean relative(int rows) throws SQLException { throw getExceptionToThrow(); } @Override public boolean previous() throws SQLException { throw getExceptionToThrow(); } @Override public int getFetchDirection() throws SQLException { throw getExceptionToThrow(); } @Override public void setFetchDirection(int direction) throws SQLException { throw getExceptionToThrow(); } @Override public int getFetchSize() throws SQLException { throw getExceptionToThrow(); } @Override public void setFetchSize(int rows) throws SQLException { throw getExceptionToThrow(); } @Override public int getType() throws SQLException { throw getExceptionToThrow(); } @Override public int getConcurrency() throws SQLException { throw getExceptionToThrow(); } @Override public boolean rowUpdated() throws SQLException { throw getExceptionToThrow(); } @Override public boolean rowInserted() throws SQLException { throw getExceptionToThrow(); } @Override public boolean rowDeleted() throws SQLException { throw getExceptionToThrow(); } @Override public void updateNull(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBoolean(int columnIndex, boolean x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateByte(int columnIndex, byte x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateShort(int columnIndex, short x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateInt(int columnIndex, int x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateLong(int columnIndex, long x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateFloat(int columnIndex, float x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateDouble(int columnIndex, double x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBigDecimal(int columnIndex, BigDecimal x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateString(int columnIndex, String x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBytes(int columnIndex, byte[] x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateDate(int columnIndex, Date x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateTime(int columnIndex, Time x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateTimestamp(int columnIndex, Timestamp x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateAsciiStream(int columnIndex, InputStream x, int length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBinaryStream(int columnIndex, InputStream x, int length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateCharacterStream(int columnIndex, Reader x, int length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateObject(int columnIndex, Object x, int scaleOrLength) throws SQLException { throw getExceptionToThrow(); } @Override public void updateObject(int columnIndex, Object x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateNull(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBoolean(String columnLabel, boolean x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateByte(String columnLabel, byte x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateShort(String columnLabel, short x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateInt(String columnLabel, int x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateLong(String columnLabel, long x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateFloat(String columnLabel, float x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateDouble(String columnLabel, double x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBigDecimal(String columnLabel, BigDecimal x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateString(String columnLabel, String x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBytes(String columnLabel, byte[] x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateDate(String columnLabel, Date x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateTime(String columnLabel, Time x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateTimestamp(String columnLabel, Timestamp x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateAsciiStream(String columnLabel, InputStream x, int length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBinaryStream(String columnLabel, InputStream x, int length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateCharacterStream(String columnLabel, Reader reader, int length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateObject(String columnLabel, Object x, int scaleOrLength) throws SQLException { throw getExceptionToThrow(); } @Override public void updateObject(String columnLabel, Object x) throws SQLException { throw getExceptionToThrow(); } @Override public void insertRow() throws SQLException { throw getExceptionToThrow(); } @Override public void updateRow() throws SQLException { throw getExceptionToThrow(); } @Override public void deleteRow() throws SQLException { throw getExceptionToThrow(); } @Override public void refreshRow() throws SQLException { throw getExceptionToThrow(); } @Override public void cancelRowUpdates() throws SQLException { throw getExceptionToThrow(); } @Override public void moveToInsertRow() throws SQLException { throw getExceptionToThrow(); } @Override public void moveToCurrentRow() throws SQLException { throw getExceptionToThrow(); } @Override public Statement getStatement() throws SQLException { throw getExceptionToThrow(); } @Override public Object getObject(int columnIndex, Map> map) throws SQLException { throw getExceptionToThrow(); } @Override public Ref getRef(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public Blob getBlob(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public Clob getClob(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public Array getArray(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public Object getObject(String columnLabel, Map> map) throws SQLException { throw getExceptionToThrow(); } @Override public Ref getRef(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public Blob getBlob(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public Clob getClob(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public Array getArray(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public Date getDate(int columnIndex, Calendar cal) throws SQLException { throw getExceptionToThrow(); } @Override public Date getDate(String columnLabel, Calendar cal) throws SQLException { throw getExceptionToThrow(); } @Override public Time getTime(int columnIndex, Calendar cal) throws SQLException { throw getExceptionToThrow(); } @Override public Time getTime(String columnLabel, Calendar cal) throws SQLException { throw getExceptionToThrow(); } @Override public Timestamp getTimestamp(int columnIndex, Calendar cal) throws SQLException { throw getExceptionToThrow(); } @Override public Timestamp getTimestamp(String columnLabel, Calendar cal) throws SQLException { throw getExceptionToThrow(); } @Override public URL getURL(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public URL getURL(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public void updateRef(int columnIndex, Ref x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateRef(String columnLabel, Ref x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBlob(int columnIndex, Blob x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBlob(String columnLabel, Blob x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateClob(int columnIndex, Clob x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateClob(String columnLabel, Clob x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateArray(int columnIndex, Array x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateArray(String columnLabel, Array x) throws SQLException { throw getExceptionToThrow(); } @Override public RowId getRowId(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public RowId getRowId(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public void updateRowId(int columnIndex, RowId x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateRowId(String columnLabel, RowId x) throws SQLException { throw getExceptionToThrow(); } @Override public int getHoldability() throws SQLException { throw getExceptionToThrow(); } @Override public boolean isClosed() throws SQLException { throw getExceptionToThrow(); } @Override public void updateNString(int columnIndex, String nString) throws SQLException { throw getExceptionToThrow(); } @Override public void updateNString(String columnLabel, String nString) throws SQLException { throw getExceptionToThrow(); } @Override public void updateNClob(int columnIndex, NClob nClob) throws SQLException { throw getExceptionToThrow(); } @Override public void updateNClob(String columnLabel, NClob nClob) throws SQLException { throw getExceptionToThrow(); } @Override public NClob getNClob(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public NClob getNClob(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public SQLXML getSQLXML(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public SQLXML getSQLXML(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public void updateSQLXML(int columnIndex, SQLXML xmlObject) throws SQLException { throw getExceptionToThrow(); } @Override public void updateSQLXML(String columnLabel, SQLXML xmlObject) throws SQLException { throw getExceptionToThrow(); } @Override public String getNString(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public String getNString(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public Reader getNCharacterStream(int columnIndex) throws SQLException { throw getExceptionToThrow(); } @Override public Reader getNCharacterStream(String columnLabel) throws SQLException { throw getExceptionToThrow(); } @Override public void updateNCharacterStream(int columnIndex, Reader x, long length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateNCharacterStream(String columnLabel, Reader reader, long length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateAsciiStream(int columnIndex, InputStream x, long length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBinaryStream(int columnIndex, InputStream x, long length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateCharacterStream(int columnIndex, Reader x, long length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateAsciiStream(String columnLabel, InputStream x, long length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBinaryStream(String columnLabel, InputStream x, long length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateCharacterStream(String columnLabel, Reader reader, long length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBlob(int columnIndex, InputStream inputStream, long length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBlob(String columnLabel, InputStream inputStream, long length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateClob(int columnIndex, Reader reader, long length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateClob(String columnLabel, Reader reader, long length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateNClob(int columnIndex, Reader reader, long length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateNClob(String columnLabel, Reader reader, long length) throws SQLException { throw getExceptionToThrow(); } @Override public void updateNCharacterStream(int columnIndex, Reader x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateNCharacterStream(String columnLabel, Reader reader) throws SQLException { throw getExceptionToThrow(); } @Override public void updateAsciiStream(int columnIndex, InputStream x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBinaryStream(int columnIndex, InputStream x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateCharacterStream(int columnIndex, Reader x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateAsciiStream(String columnLabel, InputStream x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBinaryStream(String columnLabel, InputStream x) throws SQLException { throw getExceptionToThrow(); } @Override public void updateCharacterStream(String columnLabel, Reader reader) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBlob(int columnIndex, InputStream inputStream) throws SQLException { throw getExceptionToThrow(); } @Override public void updateBlob(String columnLabel, InputStream inputStream) throws SQLException { throw getExceptionToThrow(); } @Override public void updateClob(int columnIndex, Reader reader) throws SQLException { throw getExceptionToThrow(); } @Override public void updateClob(String columnLabel, Reader reader) throws SQLException { throw getExceptionToThrow(); } @Override public void updateNClob(int columnIndex, Reader reader) throws SQLException { throw getExceptionToThrow(); } @Override public void updateNClob(String columnLabel, Reader reader) throws SQLException { throw getExceptionToThrow(); } @Override public T getObject(int columnIndex, Class type) throws SQLException { throw getExceptionToThrow(); } @Override public T getObject(String columnLabel, Class type) throws SQLException { throw getExceptionToThrow(); } @Override public void updateObject(int columnIndex, Object x, SQLType targetSqlType, int scaleOrLength) throws SQLException { throw getExceptionToThrow(); } @Override public void updateObject(String columnLabel, Object x, SQLType targetSqlType, int scaleOrLength) throws SQLException { throw getExceptionToThrow(); } @Override public void updateObject(int columnIndex, Object x, SQLType targetSqlType) throws SQLException { throw getExceptionToThrow(); } @Override public void updateObject(String columnLabel, Object x, SQLType targetSqlType) throws SQLException { throw getExceptionToThrow(); } @Override public T unwrap(Class iface) throws SQLException { throw getExceptionToThrow(); } @Override public boolean isWrapperFor(Class iface) throws SQLException { throw getExceptionToThrow(); } } private static SQLException getExceptionToThrow() { return getExceptionToThrow("Method is not implemented!"); } private static SQLException getExceptionToThrow(String message) { return new SQLException(message); } public static class ThrowingResultSetMetaData implements ResultSetMetaData { @Override public int getColumnCount() throws SQLException { throw getExceptionToThrow(); } @Override public boolean isAutoIncrement(int column) throws SQLException { throw getExceptionToThrow(); } @Override public boolean isCaseSensitive(int column) throws SQLException { throw getExceptionToThrow(); } @Override public boolean isSearchable(int column) throws SQLException { throw getExceptionToThrow(); } @Override public boolean isCurrency(int column) throws SQLException { throw getExceptionToThrow(); } @Override public int isNullable(int column) throws SQLException { throw getExceptionToThrow(); } @Override public boolean isSigned(int column) throws SQLException { throw getExceptionToThrow(); } @Override public int getColumnDisplaySize(int column) throws SQLException { throw getExceptionToThrow(); } @Override public String getColumnLabel(int column) throws SQLException { throw getExceptionToThrow(); } @Override public String getColumnName(int column) throws SQLException { throw getExceptionToThrow(); } @Override public String getSchemaName(int column) throws SQLException { throw getExceptionToThrow(); } @Override public int getPrecision(int column) throws SQLException { throw getExceptionToThrow(); } @Override public int getScale(int column) throws SQLException { throw getExceptionToThrow(); } @Override public String getTableName(int column) throws SQLException { throw getExceptionToThrow(); } @Override public String getCatalogName(int column) throws SQLException { throw getExceptionToThrow(); } @Override public int getColumnType(int column) throws SQLException { throw getExceptionToThrow(); } @Override public String getColumnTypeName(int column) throws SQLException { throw getExceptionToThrow(); } @Override public boolean isReadOnly(int column) throws SQLException { throw getExceptionToThrow(); } @Override public boolean isWritable(int column) throws SQLException { throw getExceptionToThrow(); } @Override public boolean isDefinitelyWritable(int column) throws SQLException { throw getExceptionToThrow(); } @Override public String getColumnClassName(int column) throws SQLException { throw getExceptionToThrow(); } @Override public T unwrap(Class iface) throws SQLException { throw getExceptionToThrow(); } @Override public boolean isWrapperFor(Class iface) throws SQLException { throw getExceptionToThrow(); } } } ResultSetUtilityTest.java000066400000000000000000000107111515451615200345350ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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.sql.ResultSet; import java.sql.SQLException; import java.sql.Types; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.VectorSchemaRoot; import org.junit.jupiter.api.Test; /** Tests of the ResultSetUtility. */ public class ResultSetUtilityTest { @Test public void testZeroRowResultSet() throws Exception { for (boolean reuseVectorSchemaRoot : new boolean[] {false, true}) { try (BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) { ResultSet rs = ResultSetUtility.generateEmptyResultSet(); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder( allocator, JdbcToArrowUtils.getUtcCalendar(), /* include metadata */ false) .setReuseVectorSchemaRoot(reuseVectorSchemaRoot) .build(); 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(); } } } } } @Test public void testBasicResultSet() throws Exception { try (ResultSetUtility.MockResultSet resultSet = ResultSetUtility.generateBasicResultSet(3)) { // Before row 1: assertTrue(resultSet.isBeforeFirst()); assertFalse(resultSet.isFirst()); assertFalse(resultSet.isLast()); assertFalse(resultSet.isAfterLast()); assertThrows(SQLException.class, () -> resultSet.getString(1)); // Row 1: assertTrue(resultSet.next()); assertFalse(resultSet.isBeforeFirst()); assertTrue(resultSet.isFirst()); assertFalse(resultSet.isLast()); assertFalse(resultSet.isAfterLast()); assertEquals("row number: 1", resultSet.getString(1)); // Row 2: assertTrue(resultSet.next()); assertFalse(resultSet.isBeforeFirst()); assertFalse(resultSet.isFirst()); assertFalse(resultSet.isLast()); assertFalse(resultSet.isAfterLast()); assertEquals("row number: 2", resultSet.getString(1)); // Row 3: assertTrue(resultSet.next()); assertFalse(resultSet.isBeforeFirst()); assertFalse(resultSet.isFirst()); assertTrue(resultSet.isLast()); assertFalse(resultSet.isAfterLast()); assertEquals("row number: 3", resultSet.getString(1)); // After row 3: assertFalse(resultSet.next()); assertFalse(resultSet.isBeforeFirst()); assertFalse(resultSet.isFirst()); assertFalse(resultSet.isLast()); assertTrue(resultSet.isAfterLast()); } } @Test public void testMockDataTypes() throws SQLException { ResultSetUtility.MockDataElement element = new ResultSetUtility.MockDataElement(1L, Types.NUMERIC); assertEquals(1L, element.getLong()); assertEquals(1, element.getInt()); assertEquals("1", element.getString()); } } apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/Table.java000066400000000000000000000112511515451615200314650ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; import com.fasterxml.jackson.annotation.JsonIgnoreProperties; import java.math.BigDecimal; import java.nio.charset.StandardCharsets; /** POJO to handle the YAML data from the test YAML file. */ @JsonIgnoreProperties(ignoreUnknown = true) public class Table { private String name; private String type; private String vector; private String timezone; private String create; private String[] data; private String query; private String drop; private String[] values; private String[] vectors; private int rowCount; public Table() {} public String getName() { return name; } public void setName(String name) { this.name = name; } public String getType() { return type; } public void setType(String type) { this.type = type; } public String getVector() { return vector; } public void setVector(String vector) { this.vector = vector; } public String[] getValues() { return values; } public void setValues(String[] values) { this.values = values; } public Long[] getLongValues() { Long[] arr = new Long[values.length]; int i = 0; for (String str : values) { arr[i++] = Long.parseLong(str); } return arr; } public Integer[] getIntValues() { Integer[] arr = new Integer[values.length]; int i = 0; for (String str : values) { arr[i++] = Integer.parseInt(str); } return arr; } public Boolean[] getBoolValues() { Boolean[] arr = new Boolean[values.length]; int i = 0; for (String str : values) { arr[i++] = Boolean.parseBoolean(str); } return arr; } public BigDecimal[] getBigDecimalValues() { BigDecimal[] arr = new BigDecimal[values.length]; int i = 0; for (String str : values) { arr[i++] = new BigDecimal(str); } return arr; } public Double[] getDoubleValues() { Double[] arr = new Double[values.length]; int i = 0; for (String str : values) { arr[i++] = Double.parseDouble(str); } return arr; } public Float[] getFloatValues() { Float[] arr = new Float[values.length]; int i = 0; for (String str : values) { arr[i++] = Float.parseFloat(str); } return arr; } public byte[][] getBinaryValues() { return getByteArray(values); } public byte[][] getVarCharValues() { return getByteArray(values); } public byte[][] getBlobValues() { return getByteArray(values); } public byte[][] getClobValues() { return getByteArray(values); } public byte[][] getCharValues() { return getByteArray(values); } public Integer[][] getListValues() { return JdbcToArrowTestHelper.getListValues(values); } public String getCreate() { return create; } public void setCreate(String create) { this.create = create; } public String[] getData() { return data; } public void setData(String[] data) { this.data = data; } public String getQuery() { return query; } public void setQuery(String query) { this.query = query; } public String getDrop() { return drop; } public void setDrop(String drop) { this.drop = drop; } public String getTimezone() { return timezone; } public void setTimezone(String timezone) { this.timezone = timezone; } public String[] getVectors() { return vectors; } public void setVectors(String[] vectors) { this.vectors = vectors; } public int getRowCount() { return rowCount; } public void setRowCount(int rowCount) { this.rowCount = rowCount; } @Override public String toString() { return "Table{name='" + name + "', type='" + type + "'}"; } static byte[][] getByteArray(String[] data) { byte[][] byteArr = new byte[data.length][]; for (int i = 0; i < data.length; i++) { byteArr[i] = data[i].getBytes(StandardCharsets.UTF_8); } return byteArr; } } UnreliableMetaDataTest.java000066400000000000000000000321471515451615200347110ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; 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.math.BigDecimal; import java.math.RoundingMode; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Types; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.Map; import java.util.stream.Stream; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; 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.Schema; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; /** Test options for dealing with unreliable ResultSetMetaData from JDBC drivers. */ public class UnreliableMetaDataTest { private BufferAllocator allocator; @BeforeEach public void beforeEach() { allocator = new RootAllocator(); } @AfterEach public void afterEach() { allocator.close(); } public static Stream getTestData() { return Arrays.stream(new Object[][] {{false}, {true}}).map(Arguments::of); } @ParameterizedTest @MethodSource("getTestData") public void testUnreliableMetaDataPrecisionAndScale(boolean reuseVectorSchemaRoot) throws Exception { ResultSet rs = buildIncorrectPrecisionAndScaleMetaDataResultSet(); ResultSetMetaData rsmd = rs.getMetaData(); assertEquals(Types.DECIMAL, rsmd.getColumnType(1), "Column type should be Types.DECIMAL"); assertEquals(0, rsmd.getScale(1), "Column scale should be zero"); assertEquals(0, rsmd.getPrecision(1), "Column precision should be zero"); rs.next(); BigDecimal bd1 = rs.getBigDecimal(1); assertEquals(new BigDecimal("1000000000000000.01"), bd1, "Value should be 1000000000000000.01"); assertEquals(2, bd1.scale(), "Value scale should be 2"); assertEquals(18, bd1.precision(), "Value precision should be 18"); assertFalse(rs.next(), "No more rows!"); // reset the ResultSet: rs.beforeFirst(); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder( allocator, JdbcToArrowUtils.getUtcCalendar(), /* include metadata */ false) .setReuseVectorSchemaRoot(reuseVectorSchemaRoot) .build(); try (ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config)) { assertTrue(iter.hasNext()); assertThrows( RuntimeException.class, iter::next, "Expected to fail due to mismatched metadata!"); } // reset the ResultSet: rs.beforeFirst(); JdbcFieldInfo explicitMappingField = new JdbcFieldInfo(Types.DECIMAL, 18, 2); Map explicitMapping = new HashMap<>(); explicitMapping.put(1, explicitMappingField); config = new JdbcToArrowConfigBuilder( allocator, JdbcToArrowUtils.getUtcCalendar(), /* include metadata */ false) .setReuseVectorSchemaRoot(reuseVectorSchemaRoot) .setExplicitTypesByColumnIndex(explicitMapping) .build(); try (ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config)) { while (iter.hasNext()) { VectorSchemaRoot root = iter.next(); root.close(); } } } @ParameterizedTest @MethodSource("getTestData") public void testInconsistentPrecisionAndScale(boolean reuseVectorSchemaRoot) throws Exception { ResultSet rs = buildVaryingPrecisionAndScaleResultSet(); ResultSetMetaData rsmd = rs.getMetaData(); assertEquals(Types.DECIMAL, rsmd.getColumnType(1), "Column type should be Types.DECIMAL"); assertEquals(0, rsmd.getScale(1), "Column scale should be zero"); assertEquals(0, rsmd.getPrecision(1), "Column precision should be zero"); rs.next(); BigDecimal bd1 = rs.getBigDecimal(1); assertEquals(new BigDecimal("1000000000000000.01"), bd1, "Value should be 1000000000000000.01"); assertEquals(2, bd1.scale(), "Value scale should be 2"); assertEquals(18, bd1.precision(), "Value precision should be 18"); rs.next(); BigDecimal bd2 = rs.getBigDecimal(1); assertEquals( new BigDecimal("1000000000300.0000001"), bd2, "Value should be 1000000000300.0000001"); assertEquals(7, bd2.scale(), "Value scale should be 7"); assertEquals(20, bd2.precision(), "Value precision should be 20"); rs.beforeFirst(); JdbcFieldInfo explicitMappingField = new JdbcFieldInfo(Types.DECIMAL, 20, 7); Map explicitMapping = new HashMap<>(); explicitMapping.put(1, explicitMappingField); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder( allocator, JdbcToArrowUtils.getUtcCalendar(), /* include metadata */ false) .setReuseVectorSchemaRoot(reuseVectorSchemaRoot) .setExplicitTypesByColumnIndex(explicitMapping) .build(); try (ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config)) { assertTrue(iter.hasNext()); assertThrows( RuntimeException.class, iter::next, "This is expected to fail due to inconsistent BigDecimal scales, while strict matching is enabled."); } // Reuse same ResultSet, with RoundingMode.UNNECESSARY set to coerce BigDecimal scale as needed: config = new JdbcToArrowConfigBuilder( allocator, JdbcToArrowUtils.getUtcCalendar(), /* include metadata */ false) .setReuseVectorSchemaRoot(reuseVectorSchemaRoot) .setExplicitTypesByColumnIndex(explicitMapping) .setBigDecimalRoundingMode(RoundingMode.UNNECESSARY) .build(); try (ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config)) { while (iter.hasNext()) { VectorSchemaRoot root = iter.next(); root.close(); } } } @ParameterizedTest @MethodSource("getTestData") public void testIncorrectNullability(boolean reuseVectorSchemaRoot) throws Exception { // ARROW-17005: ResultSetMetaData may indicate a field is non-nullable even when there are nulls ResultSetUtility.MockResultSetMetaData.MockColumnMetaData columnMetaData = ResultSetUtility.MockResultSetMetaData.MockColumnMetaData.builder() .sqlType(Types.INTEGER) .nullable(ResultSetMetaData.columnNoNulls) .build(); ResultSetMetaData metadata = new ResultSetUtility.MockResultSetMetaData(Collections.singletonList(columnMetaData)); final ResultSetUtility.MockResultSet.Builder resultSetBuilder = ResultSetUtility.MockResultSet.builder() .setMetaData(metadata) .addDataElement(new ResultSetUtility.MockDataElement(1024, Types.INTEGER)) .finishRow() .addDataElement(new ResultSetUtility.MockDataElement(null, Types.INTEGER)) .finishRow(); final Schema notNullSchema = new Schema( Collections.singletonList( 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)))); try (final ResultSet rs = resultSetBuilder.build()) { JdbcToArrowConfig config = new JdbcToArrowConfigBuilder( allocator, JdbcToArrowUtils.getUtcCalendar(), /* include metadata */ false) .setReuseVectorSchemaRoot(reuseVectorSchemaRoot) .build(); try (ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config)) { assertTrue(iter.hasNext()); final VectorSchemaRoot root = iter.next(); // The wrong data is returned here assertEquals(notNullSchema, root.getSchema()); assertEquals(2, root.getRowCount()); final IntVector ints = (IntVector) root.getVector(0); assertEquals(1024, ints.get(0)); assertFalse(ints.isNull(1)); assertFalse(iter.hasNext()); root.close(); } rs.beforeFirst(); // Override the nullability to get the correct result final Map typeMapping = new HashMap<>(); JdbcFieldInfo realFieldInfo = new JdbcFieldInfo( Types.INTEGER, ResultSetMetaData.columnNullable, /*precision*/ 0, /*scale*/ 0); typeMapping.put(1, realFieldInfo); config = new JdbcToArrowConfigBuilder( allocator, JdbcToArrowUtils.getUtcCalendar(), /* include metadata */ false) .setReuseVectorSchemaRoot(reuseVectorSchemaRoot) .setExplicitTypesByColumnIndex(typeMapping) .build(); try (ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config)) { assertTrue(iter.hasNext()); final VectorSchemaRoot root = iter.next(); assertEquals(nullSchema, root.getSchema()); assertEquals(2, root.getRowCount()); final IntVector ints = (IntVector) root.getVector(0); assertEquals(1024, ints.get(0)); assertTrue(ints.isNull(1)); assertFalse(iter.hasNext()); root.close(); } rs.beforeFirst(); // columnNullableUnknown won't override the metadata realFieldInfo = new JdbcFieldInfo( Types.INTEGER, ResultSetMetaData.columnNullableUnknown, /*precision*/ 0, /*scale*/ 0); typeMapping.put(1, realFieldInfo); config = new JdbcToArrowConfigBuilder( allocator, JdbcToArrowUtils.getUtcCalendar(), /* include metadata */ false) .setReuseVectorSchemaRoot(reuseVectorSchemaRoot) .setExplicitTypesByColumnIndex(typeMapping) .build(); try (ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config)) { assertTrue(iter.hasNext()); final VectorSchemaRoot root = iter.next(); assertEquals(notNullSchema, root.getSchema()); assertEquals(2, root.getRowCount()); final IntVector ints = (IntVector) root.getVector(0); assertEquals(1024, ints.get(0)); assertFalse(ints.isNull(1)); assertFalse(iter.hasNext()); root.close(); } } } private ResultSet buildIncorrectPrecisionAndScaleMetaDataResultSet() throws SQLException { ResultSetUtility.MockResultSetMetaData.MockColumnMetaData columnMetaData = ResultSetUtility.MockResultSetMetaData.MockColumnMetaData.builder() .sqlType(Types.DECIMAL) .precision(0) .scale(0) .build(); ArrayList cols = new ArrayList<>(); cols.add(columnMetaData); ResultSetMetaData metadata = new ResultSetUtility.MockResultSetMetaData(cols); return ResultSetUtility.MockResultSet.builder() .setMetaData(metadata) .addDataElement( new ResultSetUtility.MockDataElement( new BigDecimal("1000000000000000.01"), Types.DECIMAL)) .finishRow() .build(); } private ResultSet buildVaryingPrecisionAndScaleResultSet() throws SQLException { ResultSetUtility.MockResultSetMetaData.MockColumnMetaData columnMetaData = ResultSetUtility.MockResultSetMetaData.MockColumnMetaData.builder() .sqlType(Types.DECIMAL) .precision(0) .scale(0) .build(); ArrayList cols = new ArrayList<>(); cols.add(columnMetaData); ResultSetMetaData metadata = new ResultSetUtility.MockResultSetMetaData(cols); return ResultSetUtility.MockResultSet.builder() .setMetaData(metadata) .addDataElement( new ResultSetUtility.MockDataElement( new BigDecimal("1000000000000000.01"), Types.DECIMAL)) .finishRow() .addDataElement( new ResultSetUtility.MockDataElement( new BigDecimal("1000000000300.0000001"), Types.DECIMAL)) .finishRow() .build(); } } apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/consumer/000077500000000000000000000000001515451615200314265ustar00rootroot00000000000000AbstractConsumerTest.java000066400000000000000000000023731515451615200363360ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; public abstract class AbstractConsumerTest { protected BufferAllocator allocator; @BeforeEach public void setUp() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void tearDown() { allocator.close(); } } BinaryConsumerTest.java000066400000000000000000000110641515451615200360140ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/consumer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.consumer; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; 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; public class BinaryConsumerTest extends AbstractConsumerTest { private static final int INITIAL_VALUE_ALLOCATION = BaseValueVector.INITIAL_VALUE_ALLOCATION; private static final int DEFAULT_RECORD_BYTE_COUNT = 8; interface InputStreamConsumer { void consume(BinaryConsumer consumer) throws IOException; } protected void assertConsume(boolean nullable, InputStreamConsumer dataConsumer, byte[][] expect) throws IOException { try (final VarBinaryVector vector = new VarBinaryVector("binary", allocator)) { BinaryConsumer consumer = BinaryConsumer.createConsumer(vector, 0, nullable); dataConsumer.consume(consumer); assertEquals(expect.length - 1, vector.getLastSet()); for (int i = 0; i < expect.length; i++) { byte[] value = expect[i]; if (value == null) { assertTrue(vector.isNull(i)); } else { assertArrayEquals(expect[i], vector.get(i)); } } } } private byte[] createBytes(int length) { byte[] bytes = new byte[length]; for (int i = 0; i < length; i++) { bytes[i] = (byte) (i % 1024); } return bytes; } public void testConsumeInputStream(byte[][] values, boolean nullable) throws IOException { assertConsume( nullable, binaryConsumer -> { for (byte[] value : values) { if (value != null) { binaryConsumer.consume(new ByteArrayInputStream(value)); } else { binaryConsumer.consume((InputStream) null); } binaryConsumer.moveWriterPosition(); } }, values); } @Test public void testConsumeInputStream() throws IOException { testConsumeInputStream(new byte[][] {createBytes(DEFAULT_RECORD_BYTE_COUNT)}, false); testConsumeInputStream( new byte[][] { createBytes(DEFAULT_RECORD_BYTE_COUNT), createBytes(DEFAULT_RECORD_BYTE_COUNT) }, false); testConsumeInputStream( new byte[][] { createBytes(DEFAULT_RECORD_BYTE_COUNT * 2), createBytes(DEFAULT_RECORD_BYTE_COUNT), createBytes(DEFAULT_RECORD_BYTE_COUNT) }, false); testConsumeInputStream( new byte[][] {createBytes(INITIAL_VALUE_ALLOCATION * DEFAULT_RECORD_BYTE_COUNT)}, false); testConsumeInputStream( new byte[][] { createBytes(INITIAL_VALUE_ALLOCATION * DEFAULT_RECORD_BYTE_COUNT * 10), }, false); testConsumeInputStream( new byte[][] { createBytes(INITIAL_VALUE_ALLOCATION * DEFAULT_RECORD_BYTE_COUNT), createBytes(INITIAL_VALUE_ALLOCATION * DEFAULT_RECORD_BYTE_COUNT) }, false); testConsumeInputStream( new byte[][] { createBytes(INITIAL_VALUE_ALLOCATION * DEFAULT_RECORD_BYTE_COUNT), createBytes(DEFAULT_RECORD_BYTE_COUNT), createBytes(INITIAL_VALUE_ALLOCATION * DEFAULT_RECORD_BYTE_COUNT) }, false); byte[][] testRecords = new byte[INITIAL_VALUE_ALLOCATION * 2][]; for (int i = 0; i < testRecords.length; i++) { 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); } } apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/h2/000077500000000000000000000000001515451615200301045ustar00rootroot00000000000000JdbcAliasToArrowTest.java000066400000000000000000000112621515451615200346640ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/h2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.h2; import static org.apache.arrow.adapter.jdbc.AbstractJdbcToArrowTest.sqlToArrow; import static org.junit.jupiter.api.Assertions.assertEquals; 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.Statement; import java.util.List; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Field; 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; public class JdbcAliasToArrowTest { private Connection conn = null; private static final String CREATE_STATEMENT = "CREATE TABLE example_table (id INTEGER);"; private static final String INSERT_STATEMENT = "INSERT INTO example_table (id) VALUES (?);"; private static final String QUERY = "SELECT id as a, id as b FROM example_table;"; private static final String DROP_STATEMENT = "DROP TABLE example_table;"; private static final String ORIGINAL_COLUMN_NAME = "ID"; private static final String COLUMN_A = "A"; private static final String COLUMN_B = "B"; @BeforeEach public void setUp() throws Exception { String url = "jdbc:h2:mem:JdbcAliasToArrowTest"; String driver = "org.h2.Driver"; Class.forName(driver); conn = DriverManager.getConnection(url); try (Statement stmt = conn.createStatement()) { stmt.executeUpdate(CREATE_STATEMENT); } } /** * Test h2 database query with alias for column name and column label. To verify reading field * alias from an H2 database works as expected. If this test fails, something is either wrong with * the setup, or the H2 SQL behavior changed. */ @Test public void testReadH2Alias() throws Exception { // insert rows int rowCount = 4; insertRows(rowCount); try (ResultSet resultSet = conn.createStatement().executeQuery(QUERY)) { ResultSetMetaData rsmd = resultSet.getMetaData(); assertEquals(2, rsmd.getColumnCount()); // check column name and column label assertEquals(ORIGINAL_COLUMN_NAME, rsmd.getColumnName(1)); assertEquals(COLUMN_A, rsmd.getColumnLabel(1)); assertEquals(ORIGINAL_COLUMN_NAME, rsmd.getColumnName(2)); assertEquals(COLUMN_B, rsmd.getColumnLabel(2)); int rowNum = 0; while (resultSet.next()) { assertEquals(rowNum, resultSet.getInt(COLUMN_A)); assertEquals(rowNum, resultSet.getInt(COLUMN_B)); ++rowNum; } assertEquals(rowCount, rowNum); } } /** * Test jdbc query results with alias to arrow works expected. Arrow result schema name should be * field alias name. */ @Test public void testJdbcAliasToArrow() throws Exception { int rowCount = 4; insertRows(rowCount); try (ResultSet resultSet = conn.createStatement().executeQuery(QUERY)) { final VectorSchemaRoot vector = sqlToArrow(resultSet, new RootAllocator(Integer.MAX_VALUE)); assertEquals(rowCount, vector.getRowCount()); Schema vectorSchema = vector.getSchema(); List vectorFields = vectorSchema.getFields(); assertEquals(COLUMN_A, vectorFields.get(0).getName()); assertEquals(COLUMN_B, vectorFields.get(1).getName()); } } @AfterEach public void tearDown() throws SQLException { try (Statement stmt = conn.createStatement()) { stmt.executeUpdate(DROP_STATEMENT); } finally { if (conn != null) { conn.close(); conn = null; } } } private void insertRows(int numRows) throws SQLException { // Insert [numRows] Rows try (PreparedStatement stmt = conn.prepareStatement(INSERT_STATEMENT)) { for (int i = 0; i < numRows; ++i) { stmt.setInt(1, i); stmt.executeUpdate(); } } } } JdbcToArrowArrayTest.java000066400000000000000000000317371515451615200347220ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/h2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.h2; import static org.apache.arrow.adapter.jdbc.AbstractJdbcToArrowTest.sqlToArrow; 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 java.nio.charset.StandardCharsets; import java.sql.Array; 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.Statement; import java.sql.Types; import java.util.HashMap; import java.util.Map; import org.apache.arrow.adapter.jdbc.JdbcFieldInfo; import org.apache.arrow.adapter.jdbc.JdbcToArrowConfig; import org.apache.arrow.adapter.jdbc.JdbcToArrowConfigBuilder; import org.apache.arrow.adapter.jdbc.JdbcToArrowUtils; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.complex.ListVector; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class JdbcToArrowArrayTest { private Connection conn = null; private static final String CREATE_STATEMENT = "CREATE TABLE array_table (id INTEGER, int_array INTEGER ARRAY, float_array REAL ARRAY, " + "string_array VARCHAR ARRAY);"; private static final String INSERT_STATEMENT = "INSERT INTO array_table (id, int_array, float_array, string_array) VALUES (?, ?, ?, ?);"; private static final String QUERY = "SELECT int_array, float_array, string_array FROM array_table ORDER BY id;"; private static final String DROP_STATEMENT = "DROP TABLE array_table;"; private static Map arrayFieldMapping; private static final String INT_ARRAY_FIELD_NAME = "INT_ARRAY"; private static final String FLOAT_ARRAY_FIELD_NAME = "FLOAT_ARRAY"; private static final String STRING_ARRAY_FIELD_NAME = "STRING_ARRAY"; @BeforeEach public void setUp() throws Exception { String url = "jdbc:h2:mem:JdbcToArrowTest"; String driver = "org.h2.Driver"; Class.forName(driver); conn = DriverManager.getConnection(url); try (Statement stmt = conn.createStatement()) { stmt.executeUpdate(CREATE_STATEMENT); } arrayFieldMapping = new HashMap(); arrayFieldMapping.put(INT_ARRAY_FIELD_NAME, new JdbcFieldInfo(Types.INTEGER)); arrayFieldMapping.put(FLOAT_ARRAY_FIELD_NAME, new JdbcFieldInfo(Types.REAL)); arrayFieldMapping.put(STRING_ARRAY_FIELD_NAME, new JdbcFieldInfo(Types.VARCHAR)); } // This test verifies reading an array field from an H2 database // works as expected. If this test fails, something is either wrong // with the setup, or the H2 SQL behavior changed. @Test public void testReadH2Array() throws Exception { int rowCount = 4; Integer[][] intArrays = generateIntegerArrayField(rowCount); Float[][] floatArrays = generateFloatArrayField(rowCount); String[][] strArrays = generateStringArrayField(rowCount); insertRows(rowCount, intArrays, floatArrays, strArrays); try (ResultSet resultSet = conn.createStatement().executeQuery(QUERY)) { ResultSetMetaData rsmd = resultSet.getMetaData(); assertEquals(3, rsmd.getColumnCount()); for (int i = 1; i <= rsmd.getColumnCount(); ++i) { assertEquals(Types.ARRAY, rsmd.getColumnType(i)); } int rowNum = 0; while (resultSet.next()) { Array intArray = resultSet.getArray(INT_ARRAY_FIELD_NAME); assertFalse(resultSet.wasNull()); try (ResultSet rs = intArray.getResultSet()) { int arrayIndex = 0; while (rs.next()) { assertEquals(intArrays[rowNum][arrayIndex].intValue(), rs.getInt(2)); ++arrayIndex; } assertEquals(intArrays[rowNum].length, arrayIndex); } Array floatArray = resultSet.getArray(FLOAT_ARRAY_FIELD_NAME); assertFalse(resultSet.wasNull()); try (ResultSet rs = floatArray.getResultSet()) { int arrayIndex = 0; while (rs.next()) { assertEquals(floatArrays[rowNum][arrayIndex].floatValue(), rs.getFloat(2), 0.001); ++arrayIndex; } assertEquals(floatArrays[rowNum].length, arrayIndex); } Array strArray = resultSet.getArray(STRING_ARRAY_FIELD_NAME); assertFalse(resultSet.wasNull()); try (ResultSet rs = strArray.getResultSet()) { int arrayIndex = 0; while (rs.next()) { assertEquals(strArrays[rowNum][arrayIndex], rs.getString(2)); ++arrayIndex; } assertEquals(strArrays[rowNum].length, arrayIndex); } ++rowNum; } assertEquals(rowCount, rowNum); } } @Test public void testJdbcToArrow() throws Exception { int rowCount = 4; Integer[][] intArrays = generateIntegerArrayField(rowCount); Float[][] floatArrays = generateFloatArrayField(rowCount); String[][] strArrays = generateStringArrayField(rowCount); insertRows(rowCount, intArrays, floatArrays, strArrays); final JdbcToArrowConfigBuilder builder = new JdbcToArrowConfigBuilder( new RootAllocator(Integer.MAX_VALUE), JdbcToArrowUtils.getUtcCalendar(), false); builder.setArraySubTypeByColumnNameMap(arrayFieldMapping); final JdbcToArrowConfig config = builder.build(); try (ResultSet resultSet = conn.createStatement().executeQuery(QUERY)) { final VectorSchemaRoot vector = sqlToArrow(resultSet, config); assertEquals(rowCount, vector.getRowCount()); assertIntegerVectorEquals( (ListVector) vector.getVector(INT_ARRAY_FIELD_NAME), rowCount, intArrays); assertFloatVectorEquals( (ListVector) vector.getVector(FLOAT_ARRAY_FIELD_NAME), rowCount, floatArrays); assertStringVectorEquals( (ListVector) vector.getVector(STRING_ARRAY_FIELD_NAME), rowCount, strArrays); } } @Test public void testJdbcToArrowWithNulls() throws Exception { int rowCount = 4; Integer[][] intArrays = { null, {0}, {1}, {}, }; Float[][] floatArrays = { {2.0f}, null, {3.0f}, {}, }; String[][] stringArrays = { {"4"}, null, {"5"}, {}, }; insertRows(rowCount, intArrays, floatArrays, stringArrays); final JdbcToArrowConfigBuilder builder = new JdbcToArrowConfigBuilder( new RootAllocator(Integer.MAX_VALUE), JdbcToArrowUtils.getUtcCalendar(), false); builder.setArraySubTypeByColumnNameMap(arrayFieldMapping); final JdbcToArrowConfig config = builder.build(); try (ResultSet resultSet = conn.createStatement().executeQuery(QUERY)) { final VectorSchemaRoot vector = sqlToArrow(resultSet, config); assertEquals(rowCount, vector.getRowCount()); assertIntegerVectorEquals( (ListVector) vector.getVector(INT_ARRAY_FIELD_NAME), rowCount, intArrays); assertFloatVectorEquals( (ListVector) vector.getVector(FLOAT_ARRAY_FIELD_NAME), rowCount, floatArrays); assertStringVectorEquals( (ListVector) vector.getVector(STRING_ARRAY_FIELD_NAME), rowCount, stringArrays); } } private void assertIntegerVectorEquals( ListVector listVector, int rowCount, Integer[][] expectedValues) { IntVector vector = (IntVector) listVector.getDataVector(); ArrowBuf offsetBuffer = listVector.getOffsetBuffer(); int prevOffset = 0; for (int row = 0; row < rowCount; ++row) { int offset = offsetBuffer.getInt((row + 1) * ListVector.OFFSET_WIDTH); if (expectedValues[row] == null) { assertEquals(0, listVector.isSet(row)); assertEquals(0, offset - prevOffset); continue; } assertEquals(1, listVector.isSet(row)); assertEquals(expectedValues[row].length, offset - prevOffset); for (int i = prevOffset; i < offset; ++i) { assertEquals(expectedValues[row][i - prevOffset].intValue(), vector.get(i)); } prevOffset = offset; } } private void assertFloatVectorEquals( ListVector listVector, int rowCount, Float[][] expectedValues) { Float4Vector vector = (Float4Vector) listVector.getDataVector(); ArrowBuf offsetBuffer = listVector.getOffsetBuffer(); int prevOffset = 0; for (int row = 0; row < rowCount; ++row) { int offset = offsetBuffer.getInt((row + 1) * ListVector.OFFSET_WIDTH); if (expectedValues[row] == null) { assertEquals(0, listVector.isSet(row)); assertEquals(0, offset - prevOffset); continue; } assertEquals(1, listVector.isSet(row)); assertEquals(expectedValues[row].length, offset - prevOffset); for (int i = prevOffset; i < offset; ++i) { assertEquals(expectedValues[row][i - prevOffset].floatValue(), vector.get(i), 0); } prevOffset = offset; } } private void assertStringVectorEquals( ListVector listVector, int rowCount, String[][] expectedValues) { VarCharVector vector = (VarCharVector) listVector.getDataVector(); ArrowBuf offsetBuffer = listVector.getOffsetBuffer(); int prevOffset = 0; for (int row = 0; row < rowCount; ++row) { int offset = offsetBuffer.getInt((row + 1) * ListVector.OFFSET_WIDTH); if (expectedValues[row] == null) { assertEquals(0, listVector.isSet(row)); assertEquals(0, offset - prevOffset); continue; } assertEquals(1, listVector.isSet(row)); assertEquals(expectedValues[row].length, offset - prevOffset); for (int i = prevOffset; i < offset; ++i) { assertArrayEquals( expectedValues[row][i - prevOffset].getBytes(StandardCharsets.UTF_8), vector.get(i)); } prevOffset = offset; } } @AfterEach public void tearDown() throws SQLException { try (Statement stmt = conn.createStatement()) { stmt.executeUpdate(DROP_STATEMENT); } finally { if (conn != null) { conn.close(); conn = null; } } } private Integer[][] generateIntegerArrayField(int numRows) { Integer[][] result = new Integer[numRows][]; for (int i = 0; i < numRows; ++i) { int val = i * 4; result[i] = new Integer[] {val, val + 1, val + 2, val + 3}; } return result; } private Float[][] generateFloatArrayField(int numRows) { Float[][] result = new Float[numRows][]; for (int i = 0; i < numRows; ++i) { int val = i * 4; result[i] = new Float[] {(float) val, (float) val + 1, (float) val + 2, (float) val + 3}; } return result; } private String[][] generateStringArrayField(int numRows) { String[][] result = new String[numRows][]; for (int i = 0; i < numRows; ++i) { int val = i * 4; result[i] = new String[] { String.valueOf(val), String.valueOf(val + 1), String.valueOf(val + 2), String.valueOf(val + 3) }; } return result; } private void insertRows( int numRows, Integer[][] integerArrays, Float[][] floatArrays, String[][] strArrays) throws SQLException { // Insert 4 Rows try (PreparedStatement stmt = conn.prepareStatement(INSERT_STATEMENT)) { for (int i = 0; i < numRows; ++i) { Integer[] integerArray = integerArrays[i]; Float[] floatArray = floatArrays[i]; String[] strArray = strArrays[i]; Array intArray = integerArray != null ? conn.createArrayOf("INT", integerArray) : null; Array realArray = floatArray != null ? conn.createArrayOf("REAL", floatArray) : null; Array varcharArray = strArray != null ? conn.createArrayOf("VARCHAR", strArray) : null; // Insert Arrays of 4 Values in Each Row stmt.setInt(1, i); stmt.setArray(2, intArray); stmt.setArray(3, realArray); stmt.setArray(4, varcharArray); stmt.executeUpdate(); if (intArray != null) { intArray.free(); } if (realArray != null) { realArray.free(); } if (varcharArray != null) { varcharArray.free(); } } } } } JdbcToArrowCharSetTest.java000066400000000000000000000151651515451615200351720ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/h2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.h2; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertVarcharVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getCharArrayWithCharSet; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.sql.DriverManager; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Statement; import java.util.Arrays; import java.util.Calendar; import java.util.stream.Stream; import org.apache.arrow.adapter.jdbc.AbstractJdbcToArrowTest; import org.apache.arrow.adapter.jdbc.JdbcToArrowConfig; import org.apache.arrow.adapter.jdbc.JdbcToArrowConfigBuilder; import org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper; import org.apache.arrow.adapter.jdbc.JdbcToArrowUtils; import org.apache.arrow.adapter.jdbc.Table; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; /** * JUnit Test Class which contains methods to test JDBC to Arrow data conversion functionality with * UTF-8 Charset, including the multi-byte CJK characters for H2 database. */ public class JdbcToArrowCharSetTest extends AbstractJdbcToArrowTest { private static final String[] testFiles = { "h2/test1_charset_h2.yml", "h2/test1_charset_ch_h2.yml", "h2/test1_charset_jp_h2.yml", "h2/test1_charset_kr_h2.yml" }; @Override public void initializeDatabase(Table table) throws SQLException, ClassNotFoundException { this.table = table; String url = "jdbc:h2:mem:JdbcToArrowTest?characterEncoding=UTF-8"; String driver = "org.h2.Driver"; Class.forName(driver); conn = DriverManager.getConnection(url); try (Statement stmt = conn.createStatement(); ) { stmt.executeUpdate(table.getCreate()); for (String insert : table.getData()) { stmt.executeUpdate(insert); } } } /** * Get the test data as a collection of Table objects for each test iteration. * * @return Collection of Table objects * @throws SQLException on error * @throws ClassNotFoundException on error * @throws IOException on error */ public static Stream getTestData() throws SQLException, ClassNotFoundException, IOException { return Arrays.stream(prepareTestData(testFiles, JdbcToArrowCharSetTest.class)) .map(Arguments::of); } /** * Test Method to test JdbcToArrow Functionality for various H2 DB based datatypes with UTF-8 * Charset, including the multi-byte CJK characters. */ @ParameterizedTest @MethodSource("getTestData") @Override public void testJdbcToArrowValues(Table table) throws SQLException, IOException, ClassNotFoundException { this.initializeDatabase(table); testDataSets( sqlToArrow( conn, table.getQuery(), new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()), false); testDataSets(sqlToArrow(conn, table.getQuery(), new RootAllocator(Integer.MAX_VALUE)), false); testDataSets( sqlToArrow( conn.createStatement().executeQuery(table.getQuery()), new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()), false); testDataSets(sqlToArrow(conn.createStatement().executeQuery(table.getQuery())), false); testDataSets( sqlToArrow( conn.createStatement().executeQuery(table.getQuery()), new RootAllocator(Integer.MAX_VALUE)), false); testDataSets( sqlToArrow(conn.createStatement().executeQuery(table.getQuery()), Calendar.getInstance()), false); testDataSets( sqlToArrow( conn.createStatement().executeQuery(table.getQuery()), new JdbcToArrowConfigBuilder( new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()) .build()), false); testDataSets( sqlToArrow( conn, table.getQuery(), new JdbcToArrowConfigBuilder( new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()) .build()), false); } @ParameterizedTest @MethodSource("getTestData") public void testJdbcSchemaMetadata(Table table) throws SQLException, ClassNotFoundException { this.initializeDatabase(table); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(new RootAllocator(0), Calendar.getInstance(), true).build(); ResultSetMetaData rsmd = conn.createStatement().executeQuery(table.getQuery()).getMetaData(); Schema schema = JdbcToArrowUtils.jdbcToArrowSchema(rsmd, config); JdbcToArrowTestHelper.assertFieldMetadataMatchesResultSetMetadata(rsmd, schema); } /** * This method calls the assert methods for various DataSets. * * @param root VectorSchemaRoot for test * @param isIncludeMapVector is this dataset checks includes map column. Jdbc type to 'map' * mapping declared in configuration only manually */ @Override public void testDataSets(VectorSchemaRoot root, boolean isIncludeMapVector) { JdbcToArrowTestHelper.assertFieldMetadataIsEmpty(root); assertVarcharVectorValues( (VarCharVector) root.getVector(CLOB), table.getRowCount(), getCharArrayWithCharSet(table.getValues(), CLOB, StandardCharsets.UTF_8)); assertVarcharVectorValues( (VarCharVector) root.getVector(VARCHAR), table.getRowCount(), getCharArrayWithCharSet(table.getValues(), VARCHAR, StandardCharsets.UTF_8)); assertVarcharVectorValues( (VarCharVector) root.getVector(CHAR), table.getRowCount(), getCharArrayWithCharSet(table.getValues(), CHAR, StandardCharsets.UTF_8)); } } JdbcToArrowDataTypesTest.java000066400000000000000000000342051515451615200355330ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/h2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.h2; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertBigIntVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertBitVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertBooleanVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertDateVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertDecimalVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertFloat4VectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertFloat8VectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertIntVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertListVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertNullVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertSmallIntVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertTimeStampVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertTimeVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertTinyIntVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertVarBinaryVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertVarcharVectorValues; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; import java.io.IOException; import java.sql.DriverManager; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Statement; import java.util.Arrays; import java.util.Calendar; import java.util.function.Function; import java.util.stream.Stream; import org.apache.arrow.adapter.jdbc.AbstractJdbcToArrowTest; import org.apache.arrow.adapter.jdbc.JdbcFieldInfo; import org.apache.arrow.adapter.jdbc.JdbcToArrowConfig; import org.apache.arrow.adapter.jdbc.JdbcToArrowConfigBuilder; import org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper; import org.apache.arrow.adapter.jdbc.JdbcToArrowUtils; import org.apache.arrow.adapter.jdbc.Table; 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.DecimalVector; 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.TimeMilliVector; import org.apache.arrow.vector.TimeStampVector; 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.extension.OpaqueType; 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.Schema; 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; /** * JUnit Test Class which contains methods to test JDBC to Arrow data conversion functionality with * various data types for H2 database using multiple test data files. */ public class JdbcToArrowDataTypesTest extends AbstractJdbcToArrowTest { private static final String BIGINT = "big_int"; private static final String BINARY = "binary"; private static final String BIT = "bit"; private static final String BLOB = "blob"; private static final String BOOL = "bool"; private static final String CHAR = "char"; private static final String CLOB = "clob"; private static final String DATE = "date"; private static final String DECIMAL = "decimal"; private static final String DOUBLE = "double"; private static final String INT = "int"; private static final String LIST = "list"; private static final String REAL = "real"; private static final String SMALLINT = "small_int"; private static final String TIME = "time"; private static final String TIMESTAMP = "timestamp"; private static final String TINYINT = "tiny_int"; private static final String VARCHAR = "varchar"; private static final String NULL = "null"; private static final String[] testFiles = { "h2/test1_bigint_h2.yml", "h2/test1_binary_h2.yml", "h2/test1_bit_h2.yml", "h2/test1_blob_h2.yml", "h2/test1_bool_h2.yml", "h2/test1_char_h2.yml", "h2/test1_clob_h2.yml", "h2/test1_date_h2.yml", "h2/test1_decimal_h2.yml", "h2/test1_double_h2.yml", "h2/test1_int_h2.yml", "h2/test1_list_h2.yml", "h2/test1_real_h2.yml", "h2/test1_smallint_h2.yml", "h2/test1_time_h2.yml", "h2/test1_timestamp_h2.yml", "h2/test1_tinyint_h2.yml", "h2/test1_varchar_h2.yml", "h2/test1_null_h2.yml" }; /** * Get the test data as a collection of Table objects for each test iteration. * * @return Collection of Table objects * @throws SQLException on error * @throws ClassNotFoundException on error * @throws IOException on error */ public static Stream getTestData() throws SQLException, ClassNotFoundException, IOException { return Arrays.stream(prepareTestData(testFiles, JdbcToArrowMapDataTypeTest.class)) .map(Arguments::of); } /** Test Method to test JdbcToArrow Functionality for various H2 DB based datatypes. */ @ParameterizedTest @MethodSource("getTestData") @Override public void testJdbcToArrowValues(Table table) throws SQLException, IOException, ClassNotFoundException { this.initializeDatabase(table); testDataSets( sqlToArrow( conn, table.getQuery(), new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()), false); testDataSets(sqlToArrow(conn, table.getQuery(), new RootAllocator(Integer.MAX_VALUE)), false); testDataSets( sqlToArrow( conn.createStatement().executeQuery(table.getQuery()), new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()), false); testDataSets(sqlToArrow(conn.createStatement().executeQuery(table.getQuery())), false); testDataSets( sqlToArrow( conn.createStatement().executeQuery(table.getQuery()), new RootAllocator(Integer.MAX_VALUE)), false); testDataSets( sqlToArrow(conn.createStatement().executeQuery(table.getQuery()), Calendar.getInstance()), false); testDataSets( sqlToArrow( conn.createStatement().executeQuery(table.getQuery()), new JdbcToArrowConfigBuilder( new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .build()), false); testDataSets( sqlToArrow( conn, table.getQuery(), new JdbcToArrowConfigBuilder( new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .build()), false); } @ParameterizedTest @MethodSource("getTestData") public void testJdbcSchemaMetadata(Table table) throws SQLException, ClassNotFoundException { this.initializeDatabase(table); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(new RootAllocator(0), Calendar.getInstance(), true) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .build(); ResultSetMetaData rsmd = conn.createStatement().executeQuery(table.getQuery()).getMetaData(); Schema schema = JdbcToArrowUtils.jdbcToArrowSchema(rsmd, config); JdbcToArrowTestHelper.assertFieldMetadataMatchesResultSetMetadata(rsmd, schema); } @Test void testOpaqueType() throws SQLException, ClassNotFoundException { try (BufferAllocator allocator = new RootAllocator()) { String url = "jdbc:h2:mem:JdbcToArrowTest"; String driver = "org.h2.Driver"; Class.forName(driver); conn = DriverManager.getConnection(url); try (Statement stmt = conn.createStatement()) { stmt.executeUpdate("CREATE TABLE unknowntype (a GEOMETRY, b INT)"); } String query = "SELECT * FROM unknowntype"; Calendar calendar = Calendar.getInstance(); Function typeConverter = (field) -> JdbcToArrowUtils.getArrowTypeFromJdbcType(field, calendar); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder() .setAllocator(allocator) .setJdbcToArrowTypeConverter( JdbcToArrowUtils.reportUnsupportedTypesAsOpaque(typeConverter, "H2")) .build(); Schema schema; try (Statement stmt = conn.createStatement(); ResultSet rs = stmt.executeQuery(query)) { schema = assertDoesNotThrow(() -> JdbcToArrowUtils.jdbcToArrowSchema(rs.getMetaData(), config)); } Schema expected = new Schema( Arrays.asList( Field.nullable( "A", new OpaqueType(Types.MinorType.NULL.getType(), "GEOMETRY", "H2")), Field.nullable("B", Types.MinorType.INT.getType()))); assertEquals(expected, schema); } } /** * This method calls the assert methods for various DataSets. * * @param root VectorSchemaRoot for test * @param isIncludeMapVector is this dataset checks includes map column. Jdbc type to 'map' * mapping declared in configuration only manually */ @Override public void testDataSets(VectorSchemaRoot root, boolean isIncludeMapVector) { JdbcToArrowTestHelper.assertFieldMetadataIsEmpty(root); switch (table.getType()) { case BIGINT: assertBigIntVectorValues( (BigIntVector) root.getVector(table.getVector()), table.getValues().length, table.getLongValues()); break; case BINARY: case BLOB: assertVarBinaryVectorValues( (VarBinaryVector) root.getVector(table.getVector()), table.getValues().length, table.getBinaryValues()); break; case BIT: assertBitVectorValues( (BitVector) root.getVector(table.getVector()), table.getValues().length, table.getIntValues()); break; case BOOL: assertBooleanVectorValues( (BitVector) root.getVector(table.getVector()), table.getValues().length, table.getBoolValues()); break; case CHAR: case VARCHAR: case CLOB: assertVarcharVectorValues( (VarCharVector) root.getVector(table.getVector()), table.getValues().length, table.getCharValues()); break; case DATE: assertDateVectorValues( (DateDayVector) root.getVector(table.getVector()), table.getValues().length, table.getIntValues()); break; case TIME: assertTimeVectorValues( (TimeMilliVector) root.getVector(table.getVector()), table.getValues().length, table.getLongValues()); break; case TIMESTAMP: assertTimeStampVectorValues( (TimeStampVector) root.getVector(table.getVector()), table.getValues().length, table.getLongValues()); break; case DECIMAL: assertDecimalVectorValues( (DecimalVector) root.getVector(table.getVector()), table.getValues().length, table.getBigDecimalValues()); break; case DOUBLE: assertFloat8VectorValues( (Float8Vector) root.getVector(table.getVector()), table.getValues().length, table.getDoubleValues()); break; case INT: assertIntVectorValues( (IntVector) root.getVector(table.getVector()), table.getValues().length, table.getIntValues()); break; case SMALLINT: assertSmallIntVectorValues( (SmallIntVector) root.getVector(table.getVector()), table.getValues().length, table.getIntValues()); break; case TINYINT: assertTinyIntVectorValues( (TinyIntVector) root.getVector(table.getVector()), table.getValues().length, table.getIntValues()); break; case REAL: assertFloat4VectorValues( (Float4Vector) root.getVector(table.getVector()), table.getValues().length, table.getFloatValues()); break; case NULL: assertNullVectorValues((NullVector) root.getVector(table.getVector()), table.getRowCount()); break; case LIST: assertListVectorValues( (ListVector) root.getVector(table.getVector()), table.getValues().length, table.getListValues()); break; default: // do nothing break; } } } JdbcToArrowMapDataTypeTest.java000066400000000000000000000067551515451615200360170ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/h2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.h2; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertMapVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getMapValues; import java.io.IOException; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.util.Calendar; import java.util.stream.Stream; import org.apache.arrow.adapter.jdbc.AbstractJdbcToArrowTest; import org.apache.arrow.adapter.jdbc.JdbcToArrowConfigBuilder; import org.apache.arrow.adapter.jdbc.Table; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.complex.MapVector; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; /** Test MapConsumer with OTHER jdbc type. */ public class JdbcToArrowMapDataTypeTest extends AbstractJdbcToArrowTest { public static Stream getTestData() throws IOException { return Stream.of( Arguments.of(getTable("h2/test1_map_h2.yml", JdbcToArrowMapDataTypeTest.class))); } /** Test Method to test JdbcToArrow Functionality for Map form Types.OTHER column. */ @ParameterizedTest @MethodSource("getTestData") @Override public void testJdbcToArrowValues(Table table) throws SQLException, IOException, ClassNotFoundException { this.initializeDatabase(table); Calendar calendar = Calendar.getInstance(); ResultSetMetaData rsmd = getQueryMetaData(table.getQuery()); testDataSets( sqlToArrow( conn.createStatement().executeQuery(table.getQuery()), new JdbcToArrowConfigBuilder( new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()) .setJdbcToArrowTypeConverter(jdbcToArrowTypeConverter(calendar, rsmd)) .build()), true); testDataSets( sqlToArrow( conn, table.getQuery(), new JdbcToArrowConfigBuilder( new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()) .setJdbcToArrowTypeConverter(jdbcToArrowTypeConverter(calendar, rsmd)) .build()), true); } /** * This method calls the assert methods for various DataSets. * * @param root VectorSchemaRoot for test * @param isIncludeMapVector is this dataset checks includes map column. Jdbc type to 'map' * mapping declared in configuration only manually */ @Override public void testDataSets(VectorSchemaRoot root, boolean isIncludeMapVector) { assertMapVectorValues( (MapVector) root.getVector(MAP), table.getRowCount(), getMapValues(table.getValues(), MAP)); } } JdbcToArrowNullTest.java000066400000000000000000000373001515451615200345460ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/h2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.h2; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertBigIntVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertBitVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertBooleanVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertDateVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertDecimalVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertFloat4VectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertFloat8VectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertIntVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertListVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertMapVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertNullValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertSmallIntVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertTimeStampVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertTimeVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertTinyIntVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertVarBinaryVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertVarcharVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getBinaryValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getBooleanValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getCharArray; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getDecimalValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getDoubleValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getFloatValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getIntValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getListValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getLongValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getMapValues; import java.io.IOException; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.util.Arrays; import java.util.Calendar; import java.util.stream.Stream; import org.apache.arrow.adapter.jdbc.AbstractJdbcToArrowTest; import org.apache.arrow.adapter.jdbc.JdbcToArrowConfig; import org.apache.arrow.adapter.jdbc.JdbcToArrowConfigBuilder; import org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper; import org.apache.arrow.adapter.jdbc.JdbcToArrowUtils; import org.apache.arrow.adapter.jdbc.Table; 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.DecimalVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.SmallIntVector; import org.apache.arrow.vector.TimeMilliVector; import org.apache.arrow.vector.TimeStampVector; 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.types.pojo.Schema; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; /** * JUnit Test Class which contains methods to test JDBC to Arrow data conversion functionality with * null values for H2 database. */ public class JdbcToArrowNullTest extends AbstractJdbcToArrowTest { private static final String NULL = "null"; private static final String SELECTED_NULL_ROW = "selected_null_row"; private static final String SELECTED_NULL_COLUMN = "selected_null_column"; private static final String[] testFiles = { "h2/test1_all_datatypes_null_h2.yml", "h2/test1_selected_datatypes_null_h2.yml", "h2/test1_all_datatypes_selected_null_rows_h2.yml" }; /** * Get the test data as a collection of Table objects for each test iteration. * * @return Collection of Table objects * @throws SQLException on error * @throws ClassNotFoundException on error * @throws IOException on error */ public static Stream getTestData() throws SQLException, ClassNotFoundException, IOException { return Arrays.stream(prepareTestData(testFiles, JdbcToArrowNullTest.class)).map(Arguments::of); } /** * Test Method to test JdbcToArrow Functionality for various H2 DB based datatypes with null * values. */ @ParameterizedTest @MethodSource("getTestData") @Override public void testJdbcToArrowValues(Table table) throws SQLException, IOException, ClassNotFoundException { this.initializeDatabase(table); testDataSets( sqlToArrow( conn, table.getQuery(), new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()), false); testDataSets(sqlToArrow(conn, table.getQuery(), new RootAllocator(Integer.MAX_VALUE)), false); testDataSets( sqlToArrow( conn.createStatement().executeQuery(table.getQuery()), new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()), false); testDataSets(sqlToArrow(conn.createStatement().executeQuery(table.getQuery())), false); testDataSets( sqlToArrow( conn.createStatement().executeQuery(table.getQuery()), new RootAllocator(Integer.MAX_VALUE)), false); testDataSets( sqlToArrow(conn.createStatement().executeQuery(table.getQuery()), Calendar.getInstance()), false); Calendar calendar = Calendar.getInstance(); ResultSetMetaData rsmd = getQueryMetaData(table.getQuery()); testDataSets( sqlToArrow( conn.createStatement().executeQuery(table.getQuery()), new JdbcToArrowConfigBuilder( new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .setJdbcToArrowTypeConverter(jdbcToArrowTypeConverter(calendar, rsmd)) .build()), true); testDataSets( sqlToArrow( conn, table.getQuery(), new JdbcToArrowConfigBuilder( new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .setJdbcToArrowTypeConverter(jdbcToArrowTypeConverter(calendar, rsmd)) .build()), true); } @ParameterizedTest @MethodSource("getTestData") public void testJdbcSchemaMetadata(Table table) throws SQLException, ClassNotFoundException { this.initializeDatabase(table); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(new RootAllocator(0), Calendar.getInstance(), true) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .build(); ResultSetMetaData rsmd = conn.createStatement().executeQuery(table.getQuery()).getMetaData(); Schema schema = JdbcToArrowUtils.jdbcToArrowSchema(rsmd, config); JdbcToArrowTestHelper.assertFieldMetadataMatchesResultSetMetadata(rsmd, schema); } /** * This method calls the assert methods for various DataSets. * * @param root VectorSchemaRoot for test * @param isIncludeMapVector is this dataset checks includes map column. Jdbc type to 'map' * mapping declared in configuration only manually */ @Override public void testDataSets(VectorSchemaRoot root, boolean isIncludeMapVector) { JdbcToArrowTestHelper.assertFieldMetadataIsEmpty(root); switch (table.getType()) { case NULL: sqlToArrowTestNullValues(table.getVectors(), root, table.getRowCount()); break; case SELECTED_NULL_COLUMN: sqlToArrowTestSelectedNullColumnsValues( table.getVectors(), root, table.getRowCount(), isIncludeMapVector); break; case SELECTED_NULL_ROW: testAllVectorValues(root, isIncludeMapVector); break; default: // do nothing break; } } private void testAllVectorValues(VectorSchemaRoot root, boolean isIncludeMapVector) { JdbcToArrowTestHelper.assertFieldMetadataIsEmpty(root); assertBigIntVectorValues( (BigIntVector) root.getVector(BIGINT), table.getRowCount(), getLongValues(table.getValues(), BIGINT)); assertTinyIntVectorValues( (TinyIntVector) root.getVector(TINYINT), table.getRowCount(), getIntValues(table.getValues(), TINYINT)); assertSmallIntVectorValues( (SmallIntVector) root.getVector(SMALLINT), table.getRowCount(), getIntValues(table.getValues(), SMALLINT)); assertVarBinaryVectorValues( (VarBinaryVector) root.getVector(BINARY), table.getRowCount(), getBinaryValues(table.getValues(), BINARY)); assertVarBinaryVectorValues( (VarBinaryVector) root.getVector(BLOB), table.getRowCount(), getBinaryValues(table.getValues(), BLOB)); assertVarcharVectorValues( (VarCharVector) root.getVector(CLOB), table.getRowCount(), getCharArray(table.getValues(), CLOB)); assertVarcharVectorValues( (VarCharVector) root.getVector(VARCHAR), table.getRowCount(), getCharArray(table.getValues(), VARCHAR)); assertVarcharVectorValues( (VarCharVector) root.getVector(CHAR), table.getRowCount(), getCharArray(table.getValues(), CHAR)); assertIntVectorValues( (IntVector) root.getVector(INT), table.getRowCount(), getIntValues(table.getValues(), INT)); assertBitVectorValues( (BitVector) root.getVector(BIT), table.getRowCount(), getIntValues(table.getValues(), BIT)); assertBooleanVectorValues( (BitVector) root.getVector(BOOL), table.getRowCount(), getBooleanValues(table.getValues(), BOOL)); assertDateVectorValues( (DateDayVector) root.getVector(DATE), table.getRowCount(), getIntValues(table.getValues(), DATE)); assertTimeVectorValues( (TimeMilliVector) root.getVector(TIME), table.getRowCount(), getLongValues(table.getValues(), TIME)); assertTimeStampVectorValues( (TimeStampVector) root.getVector(TIMESTAMP), table.getRowCount(), getLongValues(table.getValues(), TIMESTAMP)); assertDecimalVectorValues( (DecimalVector) root.getVector(DECIMAL), table.getRowCount(), getDecimalValues(table.getValues(), DECIMAL)); assertFloat8VectorValues( (Float8Vector) root.getVector(DOUBLE), table.getRowCount(), getDoubleValues(table.getValues(), DOUBLE)); assertFloat4VectorValues( (Float4Vector) root.getVector(REAL), table.getRowCount(), getFloatValues(table.getValues(), REAL)); assertListVectorValues( (ListVector) root.getVector(LIST), table.getRowCount(), getListValues(table.getValues(), LIST)); if (isIncludeMapVector) { assertMapVectorValues( (MapVector) root.getVector(MAP), table.getRowCount(), getMapValues(table.getValues(), MAP)); } } /** * This method assert tests null values in vectors for all the datatypes. * * @param vectors Vectors to test * @param root VectorSchemaRoot for test * @param rowCount number of rows */ public void sqlToArrowTestNullValues(String[] vectors, VectorSchemaRoot root, int rowCount) { assertNullValues((IntVector) root.getVector(vectors[0]), rowCount); assertNullValues((BitVector) root.getVector(vectors[1]), rowCount); assertNullValues((TinyIntVector) root.getVector(vectors[2]), rowCount); assertNullValues((SmallIntVector) root.getVector(vectors[3]), rowCount); assertNullValues((BigIntVector) root.getVector(vectors[4]), rowCount); assertNullValues((DecimalVector) root.getVector(vectors[5]), rowCount); assertNullValues((Float8Vector) root.getVector(vectors[6]), rowCount); assertNullValues((Float4Vector) root.getVector(vectors[7]), rowCount); assertNullValues((TimeMilliVector) root.getVector(vectors[8]), rowCount); assertNullValues((DateDayVector) root.getVector(vectors[9]), rowCount); assertNullValues((TimeStampVector) root.getVector(vectors[10]), rowCount); assertNullValues((VarBinaryVector) root.getVector(vectors[11]), rowCount); assertNullValues((VarCharVector) root.getVector(vectors[12]), rowCount); assertNullValues((VarBinaryVector) root.getVector(vectors[13]), rowCount); assertNullValues((VarCharVector) root.getVector(vectors[14]), rowCount); assertNullValues((VarCharVector) root.getVector(vectors[15]), rowCount); assertNullValues((BitVector) root.getVector(vectors[16]), rowCount); assertNullValues((ListVector) root.getVector(vectors[17]), rowCount); } /** * This method assert tests null values in vectors for some selected datatypes. * * @param vectors Vectors to test * @param root VectorSchemaRoot for test * @param rowCount number of rows * @param isIncludeMapVector is this dataset checks includes map column. Jdbc type to 'map' * mapping declared in configuration only manually */ public void sqlToArrowTestSelectedNullColumnsValues( String[] vectors, VectorSchemaRoot root, int rowCount, boolean isIncludeMapVector) { assertNullValues((BigIntVector) root.getVector(vectors[0]), rowCount); assertNullValues((DecimalVector) root.getVector(vectors[1]), rowCount); assertNullValues((Float8Vector) root.getVector(vectors[2]), rowCount); assertNullValues((Float4Vector) root.getVector(vectors[3]), rowCount); assertNullValues((TimeMilliVector) root.getVector(vectors[4]), rowCount); assertNullValues((DateDayVector) root.getVector(vectors[5]), rowCount); assertNullValues((TimeStampVector) root.getVector(vectors[6]), rowCount); assertNullValues((VarBinaryVector) root.getVector(vectors[7]), rowCount); assertNullValues((VarCharVector) root.getVector(vectors[8]), rowCount); assertNullValues((VarBinaryVector) root.getVector(vectors[9]), rowCount); assertNullValues((VarCharVector) root.getVector(vectors[10]), rowCount); assertNullValues((VarCharVector) root.getVector(vectors[11]), rowCount); assertNullValues((BitVector) root.getVector(vectors[12]), rowCount); assertNullValues((ListVector) root.getVector(vectors[13]), rowCount); if (isIncludeMapVector) { assertNullValues((MapVector) root.getVector(vectors[14]), rowCount); } } } JdbcToArrowOptionalColumnsTest.java000066400000000000000000000066521515451615200367700ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/h2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.h2; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.IOException; import java.sql.SQLException; import java.util.Arrays; import java.util.stream.Stream; import org.apache.arrow.adapter.jdbc.AbstractJdbcToArrowTest; import org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper; import org.apache.arrow.adapter.jdbc.Table; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.VectorSchemaRoot; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; /** * JUnit Test Class which contains methods to test JDBC to Arrow data conversion functionality for * (non-)optional columns, in particular with regard to the ensuing VectorSchemaRoot's schema. */ public class JdbcToArrowOptionalColumnsTest extends AbstractJdbcToArrowTest { private static final String[] testFiles = {"h2/test1_null_and_notnull.yml"}; /** * Get the test data as a collection of Table objects for each test iteration. * * @return Collection of Table objects * @throws SQLException on error * @throws ClassNotFoundException on error * @throws IOException on error */ public static Stream getTestData() throws SQLException, ClassNotFoundException, IOException { return Arrays.stream(prepareTestData(testFiles, JdbcToArrowOptionalColumnsTest.class)) .map(Arguments::of); } /** * Test Method to test JdbcToArrow Functionality for dealing with nullable and non-nullable * columns. */ @ParameterizedTest @MethodSource("getTestData") @Override public void testJdbcToArrowValues(Table table) throws SQLException, IOException, ClassNotFoundException { this.initializeDatabase(table); testDataSets(sqlToArrow(conn, table.getQuery(), new RootAllocator(Integer.MAX_VALUE)), false); } /** * This method calls the assert methods for various DataSets. We verify that a SQL `NULL` column * becomes nullable in the VectorSchemaRoot, and that a SQL `NOT NULL` column becomes * non-nullable. * * @param root VectorSchemaRoot for test * @param isIncludeMapVector is this dataset checks includes map column. Jdbc type to 'map' * mapping declared in configuration only manually */ @Override public void testDataSets(VectorSchemaRoot root, boolean isIncludeMapVector) { JdbcToArrowTestHelper.assertFieldMetadataIsEmpty(root); assertTrue(root.getSchema().getFields().get(0).isNullable()); assertFalse(root.getSchema().getFields().get(1).isNullable()); } } JdbcToArrowTest.java000066400000000000000000000257601515451615200337220ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/h2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.h2; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.*; import static org.junit.jupiter.api.Assertions.assertEquals; import java.io.IOException; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.util.Arrays; import java.util.Calendar; import java.util.stream.Stream; import org.apache.arrow.adapter.jdbc.AbstractJdbcToArrowTest; import org.apache.arrow.adapter.jdbc.ArrowVectorIterator; import org.apache.arrow.adapter.jdbc.JdbcToArrow; import org.apache.arrow.adapter.jdbc.JdbcToArrowConfig; import org.apache.arrow.adapter.jdbc.JdbcToArrowConfigBuilder; import org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper; import org.apache.arrow.adapter.jdbc.JdbcToArrowUtils; import org.apache.arrow.adapter.jdbc.ResultSetUtility; import org.apache.arrow.adapter.jdbc.Table; 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.DecimalVector; 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.TimeMilliVector; import org.apache.arrow.vector.TimeStampVector; 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.types.pojo.Schema; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; /** * JUnit Test Class which contains methods to test JDBC to Arrow data conversion functionality with * various data types for H2 database using single test data file. */ public class JdbcToArrowTest extends AbstractJdbcToArrowTest { private static final String[] testFiles = {"h2/test1_all_datatypes_h2.yml"}; /** * Get the test data as a collection of Table objects for each test iteration. * * @return Collection of Table objects * @throws SQLException on error * @throws ClassNotFoundException on error * @throws IOException on error */ public static Stream getTestData() throws SQLException, ClassNotFoundException, IOException { return Arrays.stream(prepareTestData(testFiles, JdbcToArrowTest.class)) .flatMap(row -> Stream.of(Arguments.of(row[0], true), Arguments.of(row[0], false))); } /** * Test Method to test JdbcToArrow Functionality for various H2 DB based datatypes with only one * test data file. */ @ParameterizedTest @MethodSource("getTestData") @Override public void testJdbcToArrowValues(Table table) throws SQLException, IOException, ClassNotFoundException { this.initializeDatabase(table); testDataSets( sqlToArrow( conn, table.getQuery(), new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()), false); testDataSets(sqlToArrow(conn, table.getQuery(), new RootAllocator(Integer.MAX_VALUE)), false); testDataSets( sqlToArrow( conn.createStatement().executeQuery(table.getQuery()), new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()), false); testDataSets(sqlToArrow(conn.createStatement().executeQuery(table.getQuery())), false); testDataSets( sqlToArrow( conn.createStatement().executeQuery(table.getQuery()), new RootAllocator(Integer.MAX_VALUE)), false); testDataSets( sqlToArrow(conn.createStatement().executeQuery(table.getQuery()), Calendar.getInstance()), false); Calendar calendar = Calendar.getInstance(); ResultSetMetaData rsmd = getQueryMetaData(table.getQuery()); testDataSets( sqlToArrow( conn.createStatement().executeQuery(table.getQuery()), new JdbcToArrowConfigBuilder( new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .setJdbcToArrowTypeConverter(jdbcToArrowTypeConverter(calendar, rsmd)) .build()), true); testDataSets( sqlToArrow( conn, table.getQuery(), new JdbcToArrowConfigBuilder(new RootAllocator(Integer.MAX_VALUE), calendar) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .setJdbcToArrowTypeConverter(jdbcToArrowTypeConverter(calendar, rsmd)) .build()), true); } @ParameterizedTest @MethodSource("getTestData") public void testJdbcSchemaMetadata(Table table, boolean reuseVectorSchemaRoot) throws SQLException, ClassNotFoundException { this.initializeDatabase(table); Calendar calendar = Calendar.getInstance(); ResultSetMetaData rsmd = getQueryMetaData(table.getQuery()); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(new RootAllocator(0), calendar, true) .setReuseVectorSchemaRoot(reuseVectorSchemaRoot) .setJdbcToArrowTypeConverter(jdbcToArrowTypeConverter(calendar, rsmd)) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .build(); Schema schema = JdbcToArrowUtils.jdbcToArrowSchema(rsmd, config); JdbcToArrowTestHelper.assertFieldMetadataMatchesResultSetMetadata(rsmd, schema); } /** * This method calls the assert methods for various DataSets. * * @param root VectorSchemaRoot for test * @param isIncludeMapVector is this dataset checks includes map column. Jdbc type to 'map' * mapping declared in configuration only manually */ @Override public void testDataSets(VectorSchemaRoot root, boolean isIncludeMapVector) { JdbcToArrowTestHelper.assertFieldMetadataIsEmpty(root); assertBigIntVectorValues( (BigIntVector) root.getVector(BIGINT), table.getRowCount(), getLongValues(table.getValues(), BIGINT)); assertTinyIntVectorValues( (TinyIntVector) root.getVector(TINYINT), table.getRowCount(), getIntValues(table.getValues(), TINYINT)); assertSmallIntVectorValues( (SmallIntVector) root.getVector(SMALLINT), table.getRowCount(), getIntValues(table.getValues(), SMALLINT)); assertVarBinaryVectorValues( (VarBinaryVector) root.getVector(BINARY), table.getRowCount(), getBinaryValues(table.getValues(), BINARY)); assertVarBinaryVectorValues( (VarBinaryVector) root.getVector(BLOB), table.getRowCount(), getBinaryValues(table.getValues(), BLOB)); assertVarcharVectorValues( (VarCharVector) root.getVector(CLOB), table.getRowCount(), getCharArray(table.getValues(), CLOB)); assertVarcharVectorValues( (VarCharVector) root.getVector(VARCHAR), table.getRowCount(), getCharArray(table.getValues(), VARCHAR)); assertVarcharVectorValues( (VarCharVector) root.getVector(CHAR), table.getRowCount(), getCharArray(table.getValues(), CHAR)); assertIntVectorValues( (IntVector) root.getVector(INT), table.getRowCount(), getIntValues(table.getValues(), INT)); assertBitVectorValues( (BitVector) root.getVector(BIT), table.getRowCount(), getIntValues(table.getValues(), BIT)); assertBooleanVectorValues( (BitVector) root.getVector(BOOL), table.getRowCount(), getBooleanValues(table.getValues(), BOOL)); assertDateVectorValues( (DateDayVector) root.getVector(DATE), table.getRowCount(), getIntValues(table.getValues(), DATE)); assertTimeVectorValues( (TimeMilliVector) root.getVector(TIME), table.getRowCount(), getLongValues(table.getValues(), TIME)); assertTimeStampVectorValues( (TimeStampVector) root.getVector(TIMESTAMP), table.getRowCount(), getLongValues(table.getValues(), TIMESTAMP)); assertDecimalVectorValues( (DecimalVector) root.getVector(DECIMAL), table.getRowCount(), getDecimalValues(table.getValues(), DECIMAL)); assertFloat8VectorValues( (Float8Vector) root.getVector(DOUBLE), table.getRowCount(), getDoubleValues(table.getValues(), DOUBLE)); assertFloat4VectorValues( (Float4Vector) root.getVector(REAL), table.getRowCount(), getFloatValues(table.getValues(), REAL)); assertNullVectorValues((NullVector) root.getVector(NULL), table.getRowCount()); assertListVectorValues( (ListVector) root.getVector(LIST), table.getRowCount(), getListValues(table.getValues(), LIST)); if (isIncludeMapVector) { assertMapVectorValues( (MapVector) root.getVector(MAP), table.getRowCount(), getMapValues(table.getValues(), MAP)); } } @ParameterizedTest @MethodSource("getTestData") public void runLargeNumberOfRows(Table table, boolean reuseVectorSchemaRoot) throws IOException, SQLException, ClassNotFoundException { this.initializeDatabase(table); BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); int x = 0; final int targetRows = 600000; ResultSet rs = ResultSetUtility.generateBasicResultSet(targetRows); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder( allocator, JdbcToArrowUtils.getUtcCalendar(), /* include metadata */ false) .setReuseVectorSchemaRoot(reuseVectorSchemaRoot) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .build(); try (ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config)) { while (iter.hasNext()) { VectorSchemaRoot root = iter.next(); x += root.getRowCount(); if (!reuseVectorSchemaRoot) { root.close(); } } } finally { allocator.close(); } assertEquals(targetRows, x); } } JdbcToArrowTimeZoneTest.java000066400000000000000000000162561515451615200353750ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/h2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.h2; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertDateVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertTimeStampVectorValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertTimeVectorValues; import java.io.IOException; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.util.Arrays; import java.util.Calendar; import java.util.TimeZone; import java.util.stream.Stream; import org.apache.arrow.adapter.jdbc.AbstractJdbcToArrowTest; import org.apache.arrow.adapter.jdbc.JdbcToArrowConfig; import org.apache.arrow.adapter.jdbc.JdbcToArrowConfigBuilder; import org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper; import org.apache.arrow.adapter.jdbc.JdbcToArrowUtils; import org.apache.arrow.adapter.jdbc.Table; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.TimeMilliVector; import org.apache.arrow.vector.TimeStampVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; /** * JUnit Test Class which contains methods to test JDBC to Arrow data conversion functionality with * TimeZone based Date, Time and Timestamp datatypes for H2 database. */ public class JdbcToArrowTimeZoneTest extends AbstractJdbcToArrowTest { private static final String EST_DATE = "est_date"; private static final String EST_TIME = "est_time"; private static final String EST_TIMESTAMP = "est_timestamp"; private static final String GMT_DATE = "gmt_date"; private static final String GMT_TIME = "gmt_time"; private static final String GMT_TIMESTAMP = "gmt_timestamp"; private static final String PST_DATE = "pst_date"; private static final String PST_TIME = "pst_time"; private static final String PST_TIMESTAMP = "pst_timestamp"; private static final String[] testFiles = { "h2/test1_est_date_h2.yml", "h2/test1_est_time_h2.yml", "h2/test1_est_timestamp_h2.yml", "h2/test1_gmt_date_h2.yml", "h2/test1_gmt_time_h2.yml", "h2/test1_gmt_timestamp_h2.yml", "h2/test1_pst_date_h2.yml", "h2/test1_pst_time_h2.yml", "h2/test1_pst_timestamp_h2.yml" }; /** * Get the test data as a collection of Table objects for each test iteration. * * @return Collection of Table objects * @throws SQLException on error * @throws ClassNotFoundException on error * @throws IOException on error */ public static Stream getTestData() throws SQLException, ClassNotFoundException, IOException { return Arrays.stream(prepareTestData(testFiles, JdbcToArrowTimeZoneTest.class)) .map(Arguments::of); } /** * Test Method to test JdbcToArrow Functionality for various H2 DB based datatypes with TimeZone * based Date, Time and Timestamp datatype. */ @ParameterizedTest @MethodSource("getTestData") @Override public void testJdbcToArrowValues(Table table) throws SQLException, IOException, ClassNotFoundException { this.initializeDatabase(table); testDataSets( sqlToArrow( conn, table.getQuery(), new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance(TimeZone.getTimeZone(table.getTimezone()))), false); testDataSets( sqlToArrow( conn.createStatement().executeQuery(table.getQuery()), new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance(TimeZone.getTimeZone(table.getTimezone()))), false); testDataSets( sqlToArrow( conn.createStatement().executeQuery(table.getQuery()), Calendar.getInstance(TimeZone.getTimeZone(table.getTimezone()))), false); testDataSets( sqlToArrow( conn.createStatement().executeQuery(table.getQuery()), new JdbcToArrowConfigBuilder( new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance(TimeZone.getTimeZone(table.getTimezone()))) .build()), false); testDataSets( sqlToArrow( conn, table.getQuery(), new JdbcToArrowConfigBuilder( new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance(TimeZone.getTimeZone(table.getTimezone()))) .build()), false); } @ParameterizedTest @MethodSource("getTestData") public void testJdbcSchemaMetadata(Table table) throws SQLException, ClassNotFoundException { this.initializeDatabase(table); Calendar calendar = Calendar.getInstance(TimeZone.getTimeZone(table.getTimezone())); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(new RootAllocator(0), calendar, true).build(); ResultSetMetaData rsmd = conn.createStatement().executeQuery(table.getQuery()).getMetaData(); Schema schema = JdbcToArrowUtils.jdbcToArrowSchema(rsmd, config); JdbcToArrowTestHelper.assertFieldMetadataMatchesResultSetMetadata(rsmd, schema); } /** * This method calls the assert methods for various DataSets. * * @param root VectorSchemaRoot for test * @param isIncludeMapVector is this dataset checks includes map column. Jdbc type to 'map' * mapping declared in configuration only manually */ @Override public void testDataSets(VectorSchemaRoot root, boolean isIncludeMapVector) { JdbcToArrowTestHelper.assertFieldMetadataIsEmpty(root); switch (table.getType()) { case EST_DATE: case GMT_DATE: case PST_DATE: assertDateVectorValues( (DateDayVector) root.getVector(table.getVector()), table.getValues().length, table.getIntValues()); break; case EST_TIME: case GMT_TIME: case PST_TIME: assertTimeVectorValues( (TimeMilliVector) root.getVector(table.getVector()), table.getValues().length, table.getLongValues()); break; case EST_TIMESTAMP: case GMT_TIMESTAMP: case PST_TIMESTAMP: assertTimeStampVectorValues( (TimeStampVector) root.getVector(table.getVector()), table.getValues().length, table.getLongValues()); break; default: // do nothing break; } } } JdbcToArrowVectorIteratorTest.java000066400000000000000000000553761515451615200366250ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/h2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc.h2; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getBinaryValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getBooleanValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getCharArray; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getDecimalValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getDoubleValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getFloatValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getIntValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getListValues; import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getLongValues; 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.io.IOException; import java.math.BigDecimal; import java.sql.SQLException; import java.sql.Types; import java.util.ArrayList; import java.util.Arrays; import java.util.Calendar; import java.util.List; import org.apache.arrow.adapter.jdbc.ArrowVectorIterator; import org.apache.arrow.adapter.jdbc.JdbcToArrow; import org.apache.arrow.adapter.jdbc.JdbcToArrowConfig; import org.apache.arrow.adapter.jdbc.JdbcToArrowConfigBuilder; import org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper; import org.apache.arrow.adapter.jdbc.Table; 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.DecimalVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.SmallIntVector; import org.apache.arrow.vector.TimeMilliVector; import org.apache.arrow.vector.TimeStampMilliTZVector; import org.apache.arrow.vector.TimeStampMilliVector; import org.apache.arrow.vector.TimeStampVector; import org.apache.arrow.vector.TinyIntVector; import org.apache.arrow.vector.ValueVector; 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.types.FloatingPointPrecision; import org.apache.arrow.vector.types.pojo.ArrowType; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.MethodSource; public class JdbcToArrowVectorIteratorTest extends JdbcToArrowTest { @ParameterizedTest @MethodSource("getTestData") @Override public void testJdbcToArrowValues(Table table) throws SQLException, IOException, ClassNotFoundException { this.initializeDatabase(table); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()) .setTargetBatchSize(3) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .build(); ArrowVectorIterator iterator = JdbcToArrow.sqlToArrowVectorIterator( conn.createStatement().executeQuery(table.getQuery()), config); validate(iterator); } @ParameterizedTest @MethodSource("getTestData") public void testVectorSchemaRootReuse(Table table, boolean reuseVectorSchemaRoot) throws SQLException, IOException, ClassNotFoundException { this.initializeDatabase(table); Integer[][] intValues = { {101, 102, 103}, {104, null, null}, {107, 108, 109}, {110} }; Integer[][][] listValues = { {{1, 2, 3}, {1, 2}, {1}}, {{2, 3, 4}, {2, 3}, {2}}, {{3, 4, 5}, {3, 4}, {3}}, {{}} }; JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()) .setTargetBatchSize(3) .setReuseVectorSchemaRoot(reuseVectorSchemaRoot) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .build(); ArrowVectorIterator iterator = JdbcToArrow.sqlToArrowVectorIterator( conn.createStatement().executeQuery(table.getQuery()), config); int batchCount = 0; VectorSchemaRoot prev = null; VectorSchemaRoot cur = null; while (iterator.hasNext()) { cur = iterator.next(); assertNotNull(cur); // verify the first column, with may contain nulls. List intVectors = new ArrayList<>(); intVectors.add((IntVector) cur.getVector(0)); assertIntVectorValues(intVectors, intValues[batchCount].length, intValues[batchCount]); // verify arrays are handled correctly List listVectors = new ArrayList<>(); listVectors.add((ListVector) cur.getVector(18)); assertListVectorValues(listVectors, listValues[batchCount].length, listValues[batchCount]); if (prev != null) { // skip the first iteration if (reuseVectorSchemaRoot) { // when reuse is enabled, different iterations are based on the same vector schema root. assertTrue(prev == cur); } else { // when reuse is enabled, a new vector schema root is created in each iteration. assertFalse(prev == cur); if (batchCount < 3) { cur.close(); } } } prev = cur; batchCount += 1; } iterator.close(); if (!reuseVectorSchemaRoot) { assertNotNull(cur); // test that closing the iterator does not close the vectors held by the consumers assertNotEquals(cur.getVector(0).getValueCount(), 0); cur.close(); } // make sure we have at least two batches, so the above test paths are actually covered assertTrue(batchCount > 1); } @ParameterizedTest @MethodSource("getTestData") public void testJdbcToArrowValuesNoLimit(Table table) throws SQLException, IOException, ClassNotFoundException { this.initializeDatabase(table); JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()) .setTargetBatchSize(JdbcToArrowConfig.NO_LIMIT_BATCH_SIZE) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .build(); ArrowVectorIterator iterator = JdbcToArrow.sqlToArrowVectorIterator( conn.createStatement().executeQuery(table.getQuery()), config); validate(iterator); } @ParameterizedTest @MethodSource("getTestData") public void testTimeStampConsumer(Table table, boolean reuseVectorSchemaRoot) throws SQLException, IOException, ClassNotFoundException { this.initializeDatabase(table); final String sql = "select timestamp_field11 from table1"; // first experiment, with calendar and time zone. JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()) .setTargetBatchSize(3) .setReuseVectorSchemaRoot(reuseVectorSchemaRoot) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .build(); assertNotNull(config.getCalendar()); try (ArrowVectorIterator iterator = JdbcToArrow.sqlToArrowVectorIterator(conn.createStatement().executeQuery(sql), config)) { VectorSchemaRoot root = iterator.next(); assertEquals(1, root.getFieldVectors().size()); // vector with time zone info. assertTrue(root.getVector(0) instanceof TimeStampMilliTZVector); } // second experiment, without calendar and time zone. config = new JdbcToArrowConfigBuilder(new RootAllocator(Integer.MAX_VALUE), null) .setTargetBatchSize(3) .setReuseVectorSchemaRoot(reuseVectorSchemaRoot) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP) .build(); assertNull(config.getCalendar()); try (ArrowVectorIterator iterator = JdbcToArrow.sqlToArrowVectorIterator(conn.createStatement().executeQuery(sql), config)) { VectorSchemaRoot root = iterator.next(); assertEquals(1, root.getFieldVectors().size()); // vector without time zone info. assertTrue(root.getVector(0) instanceof TimeStampMilliVector); } } private void validate(ArrowVectorIterator iterator) throws SQLException, IOException { List bigIntVectors = new ArrayList<>(); List tinyIntVectors = new ArrayList<>(); List intVectors = new ArrayList<>(); List smallIntVectors = new ArrayList<>(); List vectorsForBinary = new ArrayList<>(); List vectorsForBlob = new ArrayList<>(); List vectorsForClob = new ArrayList<>(); List vectorsForVarChar = new ArrayList<>(); List vectorsForChar = new ArrayList<>(); List vectorsForBit = new ArrayList<>(); List vectorsForBool = new ArrayList<>(); List dateDayVectors = new ArrayList<>(); List timeMilliVectors = new ArrayList<>(); List timeStampVectors = new ArrayList<>(); List decimalVectors = new ArrayList<>(); List float4Vectors = new ArrayList<>(); List float8Vectors = new ArrayList<>(); List listVectors = new ArrayList<>(); List roots = new ArrayList<>(); while (iterator.hasNext()) { VectorSchemaRoot root = iterator.next(); roots.add(root); JdbcToArrowTestHelper.assertFieldMetadataIsEmpty(root); bigIntVectors.add((BigIntVector) root.getVector(BIGINT)); tinyIntVectors.add((TinyIntVector) root.getVector(TINYINT)); intVectors.add((IntVector) root.getVector(INT)); smallIntVectors.add((SmallIntVector) root.getVector(SMALLINT)); vectorsForBinary.add((VarBinaryVector) root.getVector(BINARY)); vectorsForBlob.add((VarBinaryVector) root.getVector(BLOB)); vectorsForClob.add((VarCharVector) root.getVector(CLOB)); vectorsForVarChar.add((VarCharVector) root.getVector(VARCHAR)); vectorsForChar.add((VarCharVector) root.getVector(CHAR)); vectorsForBit.add((BitVector) root.getVector(BIT)); vectorsForBool.add((BitVector) root.getVector(BOOL)); dateDayVectors.add((DateDayVector) root.getVector(DATE)); timeMilliVectors.add((TimeMilliVector) root.getVector(TIME)); timeStampVectors.add((TimeStampVector) root.getVector(TIMESTAMP)); decimalVectors.add((DecimalVector) root.getVector(DECIMAL)); float4Vectors.add((Float4Vector) root.getVector(REAL)); float8Vectors.add((Float8Vector) root.getVector(DOUBLE)); listVectors.add((ListVector) root.getVector(LIST)); } assertBigIntVectorValues( bigIntVectors, table.getRowCount(), getLongValues(table.getValues(), BIGINT)); assertTinyIntVectorValues( tinyIntVectors, table.getRowCount(), getIntValues(table.getValues(), TINYINT)); assertIntVectorValues(intVectors, table.getRowCount(), getIntValues(table.getValues(), INT)); assertSmallIntVectorValues( smallIntVectors, table.getRowCount(), getIntValues(table.getValues(), SMALLINT)); assertBinaryVectorValues( vectorsForBinary, table.getRowCount(), getBinaryValues(table.getValues(), BINARY)); assertBinaryVectorValues( vectorsForBlob, table.getRowCount(), getBinaryValues(table.getValues(), BLOB)); assertVarCharVectorValues( vectorsForClob, table.getRowCount(), getCharArray(table.getValues(), CLOB)); assertVarCharVectorValues( vectorsForVarChar, table.getRowCount(), getCharArray(table.getValues(), VARCHAR)); assertVarCharVectorValues( vectorsForChar, table.getRowCount(), getCharArray(table.getValues(), CHAR)); assertBitVectorValues(vectorsForBit, table.getRowCount(), getIntValues(table.getValues(), BIT)); assertBooleanVectorValues( vectorsForBool, table.getRowCount(), getBooleanValues(table.getValues(), BOOL)); assertDateDayVectorValues( dateDayVectors, table.getRowCount(), getLongValues(table.getValues(), DATE)); assertTimeMilliVectorValues( timeMilliVectors, table.getRowCount(), getLongValues(table.getValues(), TIME)); assertTimeStampVectorValues( timeStampVectors, table.getRowCount(), getLongValues(table.getValues(), TIMESTAMP)); assertDecimalVectorValues( decimalVectors, table.getRowCount(), getDecimalValues(table.getValues(), DECIMAL)); assertFloat4VectorValues( float4Vectors, table.getRowCount(), getFloatValues(table.getValues(), REAL)); assertFloat8VectorValues( float8Vectors, table.getRowCount(), getDoubleValues(table.getValues(), DOUBLE)); assertListVectorValues( listVectors, table.getRowCount(), getListValues(table.getValues(), LIST)); roots.forEach(root -> root.close()); } private void assertFloat8VectorValues(List vectors, int rowCount, Double[] values) { int valueCount = vectors.stream().mapToInt(ValueVector::getValueCount).sum(); assertEquals(rowCount, valueCount); int index = 0; for (Float8Vector vector : vectors) { for (int i = 0; i < vector.getValueCount(); i++) { assertEquals(values[index++].doubleValue(), vector.get(i), 0.01); } } } private void assertFloat4VectorValues(List vectors, int rowCount, Float[] values) { int valueCount = vectors.stream().mapToInt(ValueVector::getValueCount).sum(); assertEquals(rowCount, valueCount); int index = 0; for (Float4Vector vector : vectors) { for (int i = 0; i < vector.getValueCount(); i++) { assertEquals(values[index++].floatValue(), vector.get(i), 0.01); } } } private void assertDecimalVectorValues( List vectors, int rowCount, BigDecimal[] values) { int valueCount = vectors.stream().mapToInt(ValueVector::getValueCount).sum(); assertEquals(rowCount, valueCount); int index = 0; for (DecimalVector vector : vectors) { for (int i = 0; i < vector.getValueCount(); i++) { assertNotNull(vector.getObject(i)); assertEquals(values[index++].doubleValue(), vector.getObject(i).doubleValue(), 0); } } } private void assertTimeStampVectorValues( List vectors, int rowCount, Long[] values) { int valueCount = vectors.stream().mapToInt(ValueVector::getValueCount).sum(); assertEquals(rowCount, valueCount); int index = 0; for (TimeStampVector vector : vectors) { for (int i = 0; i < vector.getValueCount(); i++) { assertEquals(values[index++].longValue(), vector.get(i)); } } } private void assertTimeMilliVectorValues( List vectors, int rowCount, Long[] values) { int valueCount = vectors.stream().mapToInt(ValueVector::getValueCount).sum(); assertEquals(rowCount, valueCount); int index = 0; for (TimeMilliVector vector : vectors) { for (int i = 0; i < vector.getValueCount(); i++) { assertEquals(values[index++].longValue(), vector.get(i)); } } } private void assertDateDayVectorValues(List vectors, int rowCount, Long[] values) { int valueCount = vectors.stream().mapToInt(ValueVector::getValueCount).sum(); assertEquals(rowCount, valueCount); int index = 0; for (DateDayVector vector : vectors) { for (int i = 0; i < vector.getValueCount(); i++) { assertEquals(values[index++].longValue(), vector.get(i)); } } } private void assertBitVectorValues(List vectors, int rowCount, Integer[] values) { int valueCount = vectors.stream().mapToInt(ValueVector::getValueCount).sum(); assertEquals(rowCount, valueCount); int index = 0; for (BitVector vector : vectors) { for (int i = 0; i < vector.getValueCount(); i++) { assertEquals(values[index++].intValue(), vector.get(i)); } } } private void assertBooleanVectorValues(List vectors, int rowCount, Boolean[] values) { int valueCount = vectors.stream().mapToInt(ValueVector::getValueCount).sum(); assertEquals(rowCount, valueCount); int index = 0; for (BitVector vector : vectors) { for (int i = 0; i < vector.getValueCount(); i++) { assertEquals(values[index++], vector.get(i) == 1); } } } private void assertVarCharVectorValues( List vectors, int rowCount, byte[][] values) { int valueCount = vectors.stream().mapToInt(ValueVector::getValueCount).sum(); assertEquals(rowCount, valueCount); int index = 0; for (VarCharVector vector : vectors) { for (int i = 0; i < vector.getValueCount(); i++) { assertArrayEquals(values[index++], vector.get(i)); } } } private void assertBinaryVectorValues( List vectors, int rowCount, byte[][] values) { int valueCount = vectors.stream().mapToInt(ValueVector::getValueCount).sum(); assertEquals(rowCount, valueCount); int index = 0; for (VarBinaryVector vector : vectors) { for (int i = 0; i < vector.getValueCount(); i++) { assertArrayEquals(values[index++], vector.get(i)); } } } private void assertSmallIntVectorValues( List vectors, int rowCount, Integer[] values) { int valueCount = vectors.stream().mapToInt(ValueVector::getValueCount).sum(); assertEquals(rowCount, valueCount); int index = 0; for (SmallIntVector vector : vectors) { for (int i = 0; i < vector.getValueCount(); i++) { assertEquals(values[index++].intValue(), vector.get(i)); } } } private void assertTinyIntVectorValues( List vectors, int rowCount, Integer[] values) { int valueCount = vectors.stream().mapToInt(ValueVector::getValueCount).sum(); assertEquals(rowCount, valueCount); int index = 0; for (TinyIntVector vector : vectors) { for (int i = 0; i < vector.getValueCount(); i++) { assertEquals(values[index++].intValue(), vector.get(i)); } } } private void assertBigIntVectorValues(List vectors, int rowCount, Long[] values) { int valueCount = vectors.stream().mapToInt(ValueVector::getValueCount).sum(); assertEquals(rowCount, valueCount); int index = 0; for (BigIntVector vector : vectors) { for (int i = 0; i < vector.getValueCount(); i++) { assertEquals(values[index++].longValue(), vector.get(i)); } } } private void assertIntVectorValues(List vectors, int rowCount, Integer[] values) { int valueCount = vectors.stream().mapToInt(ValueVector::getValueCount).sum(); assertEquals(rowCount, valueCount); int index = 0; for (IntVector vector : vectors) { for (int i = 0; i < vector.getValueCount(); i++) { if (values[index] == null) { assertTrue(vector.isNull(i)); } else { assertEquals(values[index].longValue(), vector.get(i)); } index++; } } } public static void assertListVectorValues( List vectors, int rowCount, Integer[][] values) { int valueCount = vectors.stream().mapToInt(ValueVector::getValueCount).sum(); assertEquals(rowCount, valueCount); int index = 0; for (ListVector vector : vectors) { for (int i = 0; i < vector.getValueCount(); i++) { if (values[index] == null) { assertTrue(vector.isNull(i)); } else { List list = (List) vector.getObject(i); assertEquals(Arrays.asList(values[index]), list); } index++; } } } /** Runs a simple query, and encapsulates the result into a field vector. */ private FieldVector getQueryResult(JdbcToArrowConfig config) throws SQLException, IOException { ArrowVectorIterator iterator = JdbcToArrow.sqlToArrowVectorIterator( conn.createStatement().executeQuery("select real_field8 from table1"), config); VectorSchemaRoot root = iterator.next(); // only one vector, since there is one column in the select statement. assertEquals(1, root.getFieldVectors().size()); FieldVector result = root.getVector(0); // make sure some data is actually read assertTrue(result.getValueCount() > 0); return result; } @ParameterizedTest @MethodSource("getTestData") public void testJdbcToArrowCustomTypeConversion(Table table, boolean reuseVectorSchemaRoot) throws SQLException, IOException, ClassNotFoundException { this.initializeDatabase(table); JdbcToArrowConfigBuilder builder = new JdbcToArrowConfigBuilder(new RootAllocator(Integer.MAX_VALUE), Calendar.getInstance()) .setTargetBatchSize(JdbcToArrowConfig.NO_LIMIT_BATCH_SIZE) .setReuseVectorSchemaRoot(reuseVectorSchemaRoot) .setArraySubTypeByColumnNameMap(ARRAY_SUB_TYPE_BY_COLUMN_NAME_MAP); // first experiment, using default type converter JdbcToArrowConfig config = builder.build(); try (FieldVector vector = getQueryResult(config)) { // the default converter translates real to float4 assertTrue(vector instanceof Float4Vector); } // second experiment, using customized type converter builder.setJdbcToArrowTypeConverter( (fieldInfo) -> { switch (fieldInfo.getJdbcType()) { case Types.REAL: // this is different from the default type converter return new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE); default: return null; } }); config = builder.build(); try (FieldVector vector = getQueryResult(config)) { // the customized converter translates real to float8 assertTrue(vector instanceof Float8Vector); } } } apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/000077500000000000000000000000001515451615200232005ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/000077500000000000000000000000001515451615200235115ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/comment.sql000066400000000000000000000021471515451615200257000ustar00rootroot00000000000000--Licensed to the Apache Software Foundation (ASF) under one or more contributor --license agreements. See the NOTICE file distributed with this work for additional --information regarding copyright ownership. The ASF licenses this file to --You under the Apache License, Version 2.0 (the "License"); you may not use --this file except in compliance with the License. You may obtain a copy of --the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required --by applicable law or agreed to in writing, software distributed under the --License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS --OF ANY KIND, either express or implied. See the License for the specific --language governing permissions and limitations under the License. create table table1( id bigint primary key, name varchar(255), column1 boolean, columnN int ); COMMENT ON TABLE table1 IS 'This is super special table with valuable data'; COMMENT ON COLUMN table1.id IS 'Record identifier'; COMMENT ON COLUMN table1.name IS 'Name of record'; COMMENT ON COLUMN table1.columnN IS 'Informative description of columnN'; apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_all_datatypes_h2.yml000066400000000000000000000310511515451615200305730ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_all_datatypes_h2' create: 'CREATE TABLE table1 (int_field1 INT, bool_field2 BOOLEAN, tinyint_field3 TINYINT, smallint_field4 SMALLINT, bigint_field5 BIGINT, decimal_field6 DECIMAL(20,2), double_field7 DOUBLE, real_field8 REAL, time_field9 TIME, date_field10 DATE, timestamp_field11 TIMESTAMP, binary_field12 VARBINARY(100), varchar_field13 VARCHAR(256), blob_field14 BLOB, clob_field15 CLOB, char_field16 CHAR(14), bit_field17 BIT, null_field18 NULL, list_field19 INT ARRAY, map_field20 VARCHAR(256));' data: - 'INSERT INTO table1 VALUES (101, 1, 45, 12000, 92233720, 17345667789.23, 56478356785.345, 56478356785.345, PARSEDATETIME(''12:45:35 GMT'', ''HH:mm:ss z''), PARSEDATETIME(''2018-02-12 GMT'', ''yyyy-MM-dd z''), PARSEDATETIME(''2018-02-12 12:45:35 GMT'', ''yyyy-MM-dd HH:mm:ss z''), ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to varchar'', ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to clob'', ''some char text'', 1, null, ARRAY[1, 2, 3], ''{"a":"b","key":"12345"}'');' - 'INSERT INTO table1 VALUES (102, 1, 45, 12000, 92233720, 17345667789.23, 56478356785.345, 56478356785.345, PARSEDATETIME(''12:45:35 GMT'', ''HH:mm:ss z''), PARSEDATETIME(''2018-02-12 GMT'', ''yyyy-MM-dd z''), PARSEDATETIME(''2018-02-12 12:45:35 GMT'', ''yyyy-MM-dd HH:mm:ss z''), ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to varchar'', ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to clob'', ''some char text'', 1, null, ARRAY[1, 2],''{"c":"d"}'');' - 'INSERT INTO table1 VALUES (103, 1, 45, 12000, 92233720, 17345667789.23, 56478356785.345, 56478356785.345, PARSEDATETIME(''12:45:35 GMT'', ''HH:mm:ss z''), PARSEDATETIME(''2018-02-12 GMT'', ''yyyy-MM-dd z''), PARSEDATETIME(''2018-02-12 12:45:35 GMT'', ''yyyy-MM-dd HH:mm:ss z''), ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to varchar'', ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to clob'', ''some char text'', 1, null, ARRAY[1],''{"e":"f"}'');' - 'INSERT INTO table1 VALUES (104, 1, 45, 12000, 92233720, 17345667789.23, 56478356785.345, 56478356785.345, PARSEDATETIME(''12:45:35 GMT'', ''HH:mm:ss z''), PARSEDATETIME(''2018-02-12 GMT'', ''yyyy-MM-dd z''), PARSEDATETIME(''2018-02-12 12:45:35 GMT'', ''yyyy-MM-dd HH:mm:ss z''), ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to varchar'', ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to clob'', ''some char text'', 1, null, ARRAY[2, 3, 4],''{"g":"h"}'');' - 'INSERT INTO table1 VALUES (null, 1, 45, 12000, 92233720, 17345667789.23, 56478356785.345, 56478356785.345, PARSEDATETIME(''12:45:35 GMT'', ''HH:mm:ss z''), PARSEDATETIME(''2018-02-12 GMT'', ''yyyy-MM-dd z''), PARSEDATETIME(''2018-02-12 12:45:35 GMT'', ''yyyy-MM-dd HH:mm:ss z''), ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to varchar'', ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to clob'', ''some char text'', 1, null, ARRAY[2, 3],''{"i":"j"}'');' - 'INSERT INTO table1 VALUES (null, 1, 45, 12000, 92233720, 17345667789.23, 56478356785.345, 56478356785.345, PARSEDATETIME(''12:45:35 GMT'', ''HH:mm:ss z''), PARSEDATETIME(''2018-02-12 GMT'', ''yyyy-MM-dd z''), PARSEDATETIME(''2018-02-12 12:45:35 GMT'', ''yyyy-MM-dd HH:mm:ss z''), ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to varchar'', ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to clob'', ''some char text'', 1, null, ARRAY[2],''{"k":"l"}'');' - 'INSERT INTO table1 VALUES (107, 1, 45, 12000, 92233720, 17345667789.23, 56478356785.345, 56478356785.345, PARSEDATETIME(''12:45:35 GMT'', ''HH:mm:ss z''), PARSEDATETIME(''2018-02-12 GMT'', ''yyyy-MM-dd z''), PARSEDATETIME(''2018-02-12 12:45:35 GMT'', ''yyyy-MM-dd HH:mm:ss z''), ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to varchar'', ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to clob'', ''some char text'', 1, null, ARRAY[3, 4, 5],''{"m":"n"}'');' - 'INSERT INTO table1 VALUES (108, 1, 45, 12000, 92233720, 17345667789.23, 56478356785.345, 56478356785.345, PARSEDATETIME(''12:45:35 GMT'', ''HH:mm:ss z''), PARSEDATETIME(''2018-02-12 GMT'', ''yyyy-MM-dd z''), PARSEDATETIME(''2018-02-12 12:45:35 GMT'', ''yyyy-MM-dd HH:mm:ss z''), ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to varchar'', ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to clob'', ''some char text'', 1, null, ARRAY[3, 4],''{"o":"p"}'');' - 'INSERT INTO table1 VALUES (109, 1, 45, 12000, 92233720, 17345667789.23, 56478356785.345, 56478356785.345, PARSEDATETIME(''12:45:35 GMT'', ''HH:mm:ss z''), PARSEDATETIME(''2018-02-12 GMT'', ''yyyy-MM-dd z''), PARSEDATETIME(''2018-02-12 12:45:35 GMT'', ''yyyy-MM-dd HH:mm:ss z''), ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to varchar'', ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to clob'', ''some char text'', 1, null, ARRAY[3],''{"q":"r"}'');' - 'INSERT INTO table1 VALUES (110, 1, 45, 12000, 92233720, 17345667789.23, 56478356785.345, 56478356785.345, PARSEDATETIME(''12:45:35 GMT'', ''HH:mm:ss z''), PARSEDATETIME(''2018-02-12 GMT'', ''yyyy-MM-dd z''), PARSEDATETIME(''2018-02-12 12:45:35 GMT'', ''yyyy-MM-dd HH:mm:ss z''), ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to varchar'', ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to clob'', ''some char text'', 1, null, ARRAY[],''{"s":"t"}'');' query: 'select int_field1, bool_field2, tinyint_field3, smallint_field4, bigint_field5, decimal_field6, double_field7, real_field8, time_field9, date_field10, timestamp_field11, binary_field12, varchar_field13, blob_field14, clob_field15, char_field16, bit_field17, null_field18, list_field19, map_field20 from table1' drop: 'DROP table table1;' rowCount: '10' values: - 'INT_FIELD1=101,102,103,104,null,null,107,108,109,110' - 'BOOL_FIELD2=1,1,1,1,1,1,1,1,1,1' - 'BIT_FIELD17=1,1,1,1,1,1,1,1,1,1' - 'TINYINT_FIELD3=45,45,45,45,45,45,45,45,45,45' - 'SMALLINT_FIELD4=12000,12000,12000,12000,12000,12000,12000,12000,12000,12000' - 'BIGINT_FIELD5=92233720,92233720,92233720,92233720,92233720,92233720,92233720,92233720,92233720,92233720' - 'REAL_FIELD8=56478356785.345f,56478356785.345f,56478356785.345f,56478356785.345f,56478356785.345f,56478356785.345f,56478356785.345f,56478356785.345f,56478356785.345f,56478356785.345f' - 'DECIMAL_FIELD6=17345667789.23,17345667789.23,17345667789.23,17345667789.23,17345667789.23,17345667789.23,17345667789.23,17345667789.23,17345667789.23,17345667789.23' - 'DOUBLE_FIELD7=56478356785.345,56478356785.345,56478356785.345,56478356785.345,56478356785.345,56478356785.345,56478356785.345,56478356785.345,56478356785.345,56478356785.345' - 'TIME_FIELD9=45935000,45935000,45935000,45935000,45935000,45935000,45935000,45935000,45935000,45935000' - 'DATE_FIELD10=17574,17574,17574,17574,17574,17574,17574,17574,17574,17574' - 'TIMESTAMP_FIELD11=1518439535000,1518439535000,1518439535000,1518439535000,1518439535000,1518439535000,1518439535000,1518439535000,1518439535000,1518439535000' - 'CHAR_FIELD16=some char text,some char text,some char text,some char text,some char text, some char text,some char text,some char text,some char text,some char text' - 'VARCHAR_FIELD13=some text that needs to be converted to varchar,some text that needs to be converted to varchar, some text that needs to be converted to varchar,some text that needs to be converted to varchar, some text that needs to be converted to varchar,some text that needs to be converted to varchar, some text that needs to be converted to varchar,some text that needs to be converted to varchar, some text that needs to be converted to varchar,some text that needs to be converted to varchar' - 'BINARY_FIELD12=736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - 'BLOB_FIELD14=736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, 736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - 'CLOB_FIELD15=some text that needs to be converted to clob,some text that needs to be converted to clob, some text that needs to be converted to clob,some text that needs to be converted to clob, some text that needs to be converted to clob,some text that needs to be converted to clob, some text that needs to be converted to clob,some text that needs to be converted to clob, some text that needs to be converted to clob,some text that needs to be converted to clob' - 'LIST_FIELD19=(1;2;3),(1;2),(1),(2;3;4),(2;3),(2),(3;4;5),(3;4),(3),()' - 'MAP_FIELD20={"a":"b"|"key":"12345"},{"c":"d"},{"e":"f"},{"g":"h"},{"i":"j"},{"k":"l"},{"m":"n"},{"o":"p"},{"q":"r"},{"s":"t"}' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_all_datatypes_null_h2.yml000066400000000000000000000052761515451615200316370ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_all_datatypes_null_h2' type: 'null' vectors: - 'INT_FIELD1' - 'BOOL_FIELD2' - 'TINYINT_FIELD3' - 'SMALLINT_FIELD4' - 'BIGINT_FIELD5' - 'DECIMAL_FIELD6' - 'DOUBLE_FIELD7' - 'REAL_FIELD8' - 'TIME_FIELD9' - 'DATE_FIELD10' - 'TIMESTAMP_FIELD11' - 'BINARY_FIELD12' - 'VARCHAR_FIELD13' - 'BLOB_FIELD14' - 'CLOB_FIELD15' - 'CHAR_FIELD16' - 'BIT_FIELD17' - 'LIST_FIELD19' - 'MAP_FIELD20' rowCount: '5' create: 'CREATE TABLE table1 (int_field1 INT, bool_field2 BOOLEAN, tinyint_field3 TINYINT, smallint_field4 SMALLINT, bigint_field5 BIGINT, decimal_field6 DECIMAL(20,2), double_field7 DOUBLE, real_field8 REAL, time_field9 TIME, date_field10 DATE, timestamp_field11 TIMESTAMP, binary_field12 VARBINARY(100), varchar_field13 VARCHAR(256), blob_field14 BLOB, clob_field15 CLOB, char_field16 CHAR(14), bit_field17 BIT, list_field19 INT ARRAY, map_field20 VARCHAR(256));' data: - 'INSERT INTO table1 VALUES (null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null);' - 'INSERT INTO table1 VALUES (null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null);' - 'INSERT INTO table1 VALUES (null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null);' - 'INSERT INTO table1 VALUES (null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null);' - 'INSERT INTO table1 VALUES (null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null);' query: 'select int_field1, bool_field2, tinyint_field3, smallint_field4, bigint_field5, decimal_field6, double_field7, real_field8, time_field9, date_field10, timestamp_field11, binary_field12, varchar_field13, blob_field14, clob_field15, char_field16, bit_field17, list_field19 from table1' drop: 'DROP table table1;' test1_all_datatypes_selected_null_rows_h2.yml000066400000000000000000000122541515451615200344740ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_all_datatypes_selected_null_rows_h2' type: 'selected_null_row' vectors: - 'INT_FIELD1' - 'BOOL_FIELD2' - 'TINYINT_FIELD3' - 'SMALLINT_FIELD4' - 'BIGINT_FIELD5' - 'DECIMAL_FIELD6' - 'DOUBLE_FIELD7' - 'REAL_FIELD8' - 'TIME_FIELD9' - 'DATE_FIELD10' - 'TIMESTAMP_FIELD11' - 'BINARY_FIELD12' - 'VARCHAR_FIELD13' - 'BLOB_FIELD14' - 'CLOB_FIELD15' - 'CHAR_FIELD16' - 'BIT_FIELD17' - 'LIST_FIELD19' - 'MAP_FIELD20' create: 'CREATE TABLE table1 (int_field1 INT, bool_field2 BOOLEAN, tinyint_field3 TINYINT, smallint_field4 SMALLINT, bigint_field5 BIGINT, decimal_field6 DECIMAL(20,2), double_field7 DOUBLE, real_field8 REAL, time_field9 TIME, date_field10 DATE, timestamp_field11 TIMESTAMP, binary_field12 VARBINARY(100), varchar_field13 VARCHAR(256), blob_field14 BLOB, clob_field15 CLOB, char_field16 CHAR(14), bit_field17 BIT, list_field19 INT ARRAY, map_field20 VARCHAR(256));' data: - 'INSERT INTO table1 VALUES (null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null);' - 'INSERT INTO table1 VALUES (101, 1, 45, 12000, 92233720, 17345667789.23, 56478356785.345, 56478356785.345, PARSEDATETIME(''12:45:35 GMT'', ''HH:mm:ss z''), PARSEDATETIME(''2018-02-12 GMT'', ''yyyy-MM-dd z''), PARSEDATETIME(''2018-02-12 12:45:35 GMT'', ''yyyy-MM-dd HH:mm:ss z''), ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to varchar'', ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to clob'', ''some char text'', 1, ARRAY[1, 2, 3],''{"a":"b"}'');' - 'INSERT INTO table1 VALUES (null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null);' - 'INSERT INTO table1 VALUES (101, 1, 45, 12000, 92233720, 17345667789.23, 56478356785.345, 56478356785.345, PARSEDATETIME(''12:45:35 GMT'', ''HH:mm:ss z''), PARSEDATETIME(''2018-02-12 GMT'', ''yyyy-MM-dd z''), PARSEDATETIME(''2018-02-12 12:45:35 GMT'', ''yyyy-MM-dd HH:mm:ss z''), ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to varchar'', ''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'', ''some text that needs to be converted to clob'', ''some char text'', 1, ARRAY[1, 2, 3],''{"c":"d"}'');' - 'INSERT INTO table1 VALUES (null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null);' query: 'select int_field1, bool_field2, tinyint_field3, smallint_field4, bigint_field5, decimal_field6, double_field7, real_field8, time_field9, date_field10, timestamp_field11, binary_field12, varchar_field13, blob_field14, clob_field15, char_field16, bit_field17, list_field19, map_field20 from table1' drop: 'DROP table table1;' rowCount: '5' values: - 'INT_FIELD1=null,101,null,101,null' - 'BOOL_FIELD2=null,1,null,1,null' - 'BIT_FIELD17=null,1,null,1,null' - 'TINYINT_FIELD3=null,45,null,45,null' - 'SMALLINT_FIELD4=null,12000,null,12000,null' - 'BIGINT_FIELD5=null,92233720,null,92233720,null' - 'REAL_FIELD8=null,56478356785.345f,null,56478356785.345f,null' - 'DECIMAL_FIELD6=null,17345667789.23,null,17345667789.23,null' - 'DOUBLE_FIELD7=null,56478356785.345,null,56478356785.345,null' - 'TIME_FIELD9=null,45935000,null,45935000,null' - 'DATE_FIELD10=null,17574,null,17574,null' - 'TIMESTAMP_FIELD11=null,1518439535000,null,1518439535000,null' - 'CHAR_FIELD16=null,some char text,null,some char text,null' - 'VARCHAR_FIELD13=null,some text that needs to be converted to varchar,null, some text that needs to be converted to varchar,null' - 'BINARY_FIELD12=null,736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, null,736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279,null' - 'BLOB_FIELD14=null,736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279, null,736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279,null' - 'CLOB_FIELD15=null,some text that needs to be converted to clob,null,some text that needs to be converted to clob,null' - 'LIST_FIELD19=null,(1;2;3),null,(1;2;3),null' - 'MAP_FIELD20=null,{"a":"b"},null,{"c":"d"},null' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_bigint_h2.yml000066400000000000000000000030131515451615200272160ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_bigint_h2' type: 'big_int' vector: 'BIGINT_FIELD5' create: 'CREATE TABLE table1 (bigint_field5 BIGINT);' data: - 'INSERT INTO table1 VALUES (92233720);' - 'INSERT INTO table1 VALUES (92233720);' - 'INSERT INTO table1 VALUES (92233720);' - 'INSERT INTO table1 VALUES (92233720);' - 'INSERT INTO table1 VALUES (92233720);' - 'INSERT INTO table1 VALUES (92233720);' - 'INSERT INTO table1 VALUES (92233720);' - 'INSERT INTO table1 VALUES (92233720);' - 'INSERT INTO table1 VALUES (92233720);' - 'INSERT INTO table1 VALUES (92233720);' query: 'select bigint_field5 from table1;' drop: 'DROP table table1;' values: - '92233720' - '92233720' - '92233720' - '92233720' - '92233720' - '92233720' - '92233720' - '92233720' - '92233720' - '92233720' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_binary_h2.yml000066400000000000000000000063151515451615200272360ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_binary_h2' type: 'binary' vector: 'BINARY_FIELD12' create: 'CREATE TABLE table1 (binary_field12 VARBINARY(100));' data: - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' query: 'select binary_field12 from table1;' drop: 'DROP table table1;' values: - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_bit_h2.yml000066400000000000000000000025571515451615200265340ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_bit_h2' type: 'bit' vector: 'BIT_FIELD17' create: 'CREATE TABLE table1 (bit_field17 BIT);' data: - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' query: 'select bit_field17 from table1;' drop: 'DROP table table1;' values: - '1' - '1' - '1' - '1' - '1' - '1' - '1' - '1' - '1' - '1' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_blob_h2.yml000066400000000000000000000062711515451615200266710ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_blob_h2' type: 'blob' vector: 'BLOB_FIELD14' create: 'CREATE TABLE table1 (blob_field14 BLOB);' data: - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' - 'INSERT INTO table1 VALUES (''736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279'');' query: 'select blob_field14 from table1;' drop: 'DROP table table1;' values: - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' - '736f6d6520746578742074686174206e6565647320746f20626520636f6e76657274656420746f2062696e617279' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_bool_h2.yml000066400000000000000000000026231515451615200267030ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_bool_h2' type: 'bool' vector: 'BOOL_FIELD2' create: 'CREATE TABLE table1 (bool_field2 BOOLEAN);' data: - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' - 'INSERT INTO table1 VALUES (1);' query: 'select bool_field2 from table1;' drop: 'DROP table table1;' values: - 'true' - 'true' - 'true' - 'true' - 'true' - 'true' - 'true' - 'true' - 'true' - 'true' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_char_h2.yml000066400000000000000000000032451515451615200266660ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_char_h2' type: 'char' vector: 'CHAR_FIELD16' create: 'CREATE TABLE table1 (char_field16 CHAR(14));' data: - 'INSERT INTO table1 VALUES (''some char text'');' - 'INSERT INTO table1 VALUES (''some char text'');' - 'INSERT INTO table1 VALUES (''some char text'');' - 'INSERT INTO table1 VALUES (''some char text'');' - 'INSERT INTO table1 VALUES (''some char text'');' - 'INSERT INTO table1 VALUES (''some char text'');' - 'INSERT INTO table1 VALUES (''some char text'');' - 'INSERT INTO table1 VALUES (''some char text'');' - 'INSERT INTO table1 VALUES (''some char text'');' - 'INSERT INTO table1 VALUES (''some char text'');' query: 'select char_field16 from table1;' drop: 'DROP table table1;' values: - 'some char text' - 'some char text' - 'some char text' - 'some char text' - 'some char text' - 'some char text' - 'some char text' - 'some char text' - 'some char text' - 'some char text' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_charset_ch_h2.yml000066400000000000000000000054311515451615200300530ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_charset_ch_h2' type: 'charset' vectors: - 'VARCHAR_FIELD13' - 'CHAR_FIELD16' - 'CLOB_FIELD15' rowCount: '5' charSet: 'GBK' create: 'CREATE TABLE table1 (int_field1 INT, varchar_field13 VARCHAR(256), clob_field15 CLOB, char_field16 CHAR(13));' data: - 'INSERT INTO table1 VALUES (101,''一些帶有char編碼的文本需è¦è½‰æ›ç‚ºvarchar'', ''一些带有charç¼–ç çš„文本需è¦è½¬æ¢ä¸ºclob'', ''一些charç¼–ç çš„字符文本'');' - 'INSERT INTO table1 VALUES (101,''一些帶有char編碼的文本需è¦è½‰æ›ç‚ºvarchar'', ''一些带有charç¼–ç çš„文本需è¦è½¬æ¢ä¸ºclob'', ''一些charç¼–ç çš„字符文本'');' - 'INSERT INTO table1 VALUES (101,''一些帶有char編碼的文本需è¦è½‰æ›ç‚ºvarchar'', ''一些带有charç¼–ç çš„文本需è¦è½¬æ¢ä¸ºclob'', ''一些charç¼–ç çš„字符文本'');' - 'INSERT INTO table1 VALUES (101,''一些帶有char編碼的文本需è¦è½‰æ›ç‚ºvarchar'', ''一些带有charç¼–ç çš„文本需è¦è½¬æ¢ä¸ºclob'', ''一些charç¼–ç çš„字符文本'');' - 'INSERT INTO table1 VALUES (101,''一些帶有char編碼的文本需è¦è½‰æ›ç‚ºvarchar'', ''一些带有charç¼–ç çš„文本需è¦è½¬æ¢ä¸ºclob'', ''一些charç¼–ç çš„字符文本'');' query: 'select varchar_field13, clob_field15, char_field16 from table1' drop: 'DROP table table1;' values: - 'VARCHAR_FIELD13=一些帶有char編碼的文本需è¦è½‰æ›ç‚ºvarchar,一些帶有char編碼的文本需è¦è½‰æ›ç‚ºvarchar,一些帶有char編碼的文本需è¦è½‰æ›ç‚ºvarchar, 一些帶有char編碼的文本需è¦è½‰æ›ç‚ºvarchar,一些帶有char編碼的文本需è¦è½‰æ›ç‚ºvarchar' - 'CLOB_FIELD15=一些带有charç¼–ç çš„文本需è¦è½¬æ¢ä¸ºclob,一些带有charç¼–ç çš„文本需è¦è½¬æ¢ä¸ºclob,一些带有charç¼–ç çš„文本需è¦è½¬æ¢ä¸ºclob, 一些带有charç¼–ç çš„文本需è¦è½¬æ¢ä¸ºclob,一些带有charç¼–ç çš„文本需è¦è½¬æ¢ä¸ºclob' - 'CHAR_FIELD16=一些charç¼–ç çš„字符文本,一些charç¼–ç çš„字符文本,一些charç¼–ç çš„字符文本,一些charç¼–ç çš„字符文本,一些charç¼–ç çš„字符文本' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_charset_h2.yml000066400000000000000000000116751515451615200274100ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_charset_h2' type: 'charset' vectors: - 'VARCHAR_FIELD13' - 'CHAR_FIELD16' - 'CLOB_FIELD15' rowCount: '10' create: 'CREATE TABLE table1 (int_field1 INT, varchar_field13 VARCHAR(256), clob_field15 CLOB, char_field16 CHAR(33));' data: - 'INSERT INTO table1 VALUES (101,''some text with char encoding that needs to be converted to varchar'', ''some text with char encoding that needs to be converted to clob'', ''some char text with char encoding'');' - 'INSERT INTO table1 VALUES (101,''some text with char encoding that needs to be converted to varchar'', ''some text with char encoding that needs to be converted to clob'', ''some char text with char encoding'');' - 'INSERT INTO table1 VALUES (101,''some text with char encoding that needs to be converted to varchar'', ''some text with char encoding that needs to be converted to clob'', ''some char text with char encoding'');' - 'INSERT INTO table1 VALUES (101,''some text with char encoding that needs to be converted to varchar'', ''some text with char encoding that needs to be converted to clob'', ''some char text with char encoding'');' - 'INSERT INTO table1 VALUES (101,''some text with char encoding that needs to be converted to varchar'', ''some text with char encoding that needs to be converted to clob'', ''some char text with char encoding'');' - 'INSERT INTO table1 VALUES (101,''some text with char encoding that needs to be converted to varchar'', ''some text with char encoding that needs to be converted to clob'', ''some char text with char encoding'');' - 'INSERT INTO table1 VALUES (101,''some text with char encoding that needs to be converted to varchar'', ''some text with char encoding that needs to be converted to clob'', ''some char text with char encoding'');' - 'INSERT INTO table1 VALUES (101,''some text with char encoding that needs to be converted to varchar'', ''some text with char encoding that needs to be converted to clob'', ''some char text with char encoding'');' - 'INSERT INTO table1 VALUES (101,''some text with char encoding that needs to be converted to varchar'', ''some text with char encoding that needs to be converted to clob'', ''some char text with char encoding'');' - 'INSERT INTO table1 VALUES (101,''some text with char encoding that needs to be converted to varchar'', ''some text with char encoding that needs to be converted to clob'', ''some char text with char encoding'');' query: 'select varchar_field13, clob_field15, char_field16 from table1' drop: 'DROP table table1;' values: - 'VARCHAR_FIELD13=some text with char encoding that needs to be converted to varchar,some text with char encoding that needs to be converted to varchar, some text with char encoding that needs to be converted to varchar,some text with char encoding that needs to be converted to varchar, some text with char encoding that needs to be converted to varchar,some text with char encoding that needs to be converted to varchar, some text with char encoding that needs to be converted to varchar,some text with char encoding that needs to be converted to varchar, some text with char encoding that needs to be converted to varchar,some text with char encoding that needs to be converted to varchar' - 'CLOB_FIELD15=some text with char encoding that needs to be converted to clob,some text with char encoding that needs to be converted to clob, some text with char encoding that needs to be converted to clob,some text with char encoding that needs to be converted to clob, some text with char encoding that needs to be converted to clob,some text with char encoding that needs to be converted to clob, some text with char encoding that needs to be converted to clob,some text with char encoding that needs to be converted to clob, some text with char encoding that needs to be converted to clob,some text with char encoding that needs to be converted to clob' - 'CHAR_FIELD16=some char text with char encoding,some char text with char encoding,some char text with char encoding,some char text with char encoding,some char text with char encoding, some char text with char encoding,some char text with char encoding,some char text with char encoding,some char text with char encoding,some char text with char encoding' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_charset_jp_h2.yml000066400000000000000000000071541515451615200300760ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_charset_jp_h2' type: 'charset' vectors: - 'VARCHAR_FIELD13' - 'CHAR_FIELD16' - 'CLOB_FIELD15' rowCount: '5' charSet: 'SJIS' create: 'CREATE TABLE table1 (int_field1 INT, varchar_field13 VARCHAR(256), clob_field15 CLOB, char_field16 CHAR(23));' data: - 'INSERT INTO table1 VALUES (101,''varcharã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹charエンコーディングã®ãƒ†ã‚­ã‚¹ãƒˆ'', ''charエンコーディングã®ã‚るテキストをclobã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹'', ''charエンコーディングã®ã‚ã‚‹charテキスト'');' - 'INSERT INTO table1 VALUES (101,''varcharã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹charエンコーディングã®ãƒ†ã‚­ã‚¹ãƒˆ'', ''charエンコーディングã®ã‚るテキストをclobã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹'', ''charエンコーディングã®ã‚ã‚‹charテキスト'');' - 'INSERT INTO table1 VALUES (101,''varcharã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹charエンコーディングã®ãƒ†ã‚­ã‚¹ãƒˆ'', ''charエンコーディングã®ã‚るテキストをclobã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹'', ''charエンコーディングã®ã‚ã‚‹charテキスト'');' - 'INSERT INTO table1 VALUES (101,''varcharã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹charエンコーディングã®ãƒ†ã‚­ã‚¹ãƒˆ'', ''charエンコーディングã®ã‚るテキストをclobã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹'', ''charエンコーディングã®ã‚ã‚‹charテキスト'');' - 'INSERT INTO table1 VALUES (101,''varcharã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹charエンコーディングã®ãƒ†ã‚­ã‚¹ãƒˆ'', ''charエンコーディングã®ã‚るテキストをclobã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹'', ''charエンコーディングã®ã‚ã‚‹charテキスト'');' query: 'select varchar_field13, clob_field15, char_field16 from table1' drop: 'DROP table table1;' values: - 'VARCHAR_FIELD13=varcharã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹charエンコーディングã®ãƒ†ã‚­ã‚¹ãƒˆ,varcharã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹charエンコーディングã®ãƒ†ã‚­ã‚¹ãƒˆ,varcharã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹charエンコーディングã®ãƒ†ã‚­ã‚¹ãƒˆ, varcharã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹charエンコーディングã®ãƒ†ã‚­ã‚¹ãƒˆ,varcharã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹charエンコーディングã®ãƒ†ã‚­ã‚¹ãƒˆ' - 'CLOB_FIELD15=charエンコーディングã®ã‚るテキストをclobã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹,charエンコーディングã®ã‚るテキストをclobã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹,charエンコーディングã®ã‚るテキストをclobã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹, charエンコーディングã®ã‚るテキストをclobã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹,charエンコーディングã®ã‚るテキストをclobã«å¤‰æ›ã™ã‚‹å¿…è¦ãŒã‚ã‚‹' - 'CHAR_FIELD16=charエンコーディングã®ã‚ã‚‹charテキスト,charエンコーディングã®ã‚ã‚‹charテキスト,charエンコーディングã®ã‚ã‚‹charテキスト,charエンコーディングã®ã‚ã‚‹charテキスト,charエンコーディングã®ã‚ã‚‹charテキスト' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_charset_kr_h2.yml000066400000000000000000000064561515451615200301050ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_charset_kr_h2' type: 'charset' vectors: - 'VARCHAR_FIELD13' - 'CHAR_FIELD16' - 'CLOB_FIELD15' rowCount: '5' charSet: 'EUC-KR' create: 'CREATE TABLE table1 (int_field1 INT, varchar_field13 VARCHAR(256), clob_field15 CLOB, char_field16 CHAR(22));' data: - 'INSERT INTO table1 VALUES (101,''char ì¸ì½”ë”©ì„ ì‚¬ìš©í•˜ëŠ” ì¼ë¶€ í…스트를 varcharë¡œ 변환해야합니다.'', ''clobë¡œ 변환해야하는 char ì¸ì½”ë”©ì„ ê°€ì§„ í…스트'', ''char ì¸ì½”ë”©ì„ ì‚¬ìš©í•œ char í…스트'');' - 'INSERT INTO table1 VALUES (101,''char ì¸ì½”ë”©ì„ ì‚¬ìš©í•˜ëŠ” ì¼ë¶€ í…스트를 varcharë¡œ 변환해야합니다.'', ''clobë¡œ 변환해야하는 char ì¸ì½”ë”©ì„ ê°€ì§„ í…스트'', ''char ì¸ì½”ë”©ì„ ì‚¬ìš©í•œ char í…스트'');' - 'INSERT INTO table1 VALUES (101,''char ì¸ì½”ë”©ì„ ì‚¬ìš©í•˜ëŠ” ì¼ë¶€ í…스트를 varcharë¡œ 변환해야합니다.'', ''clobë¡œ 변환해야하는 char ì¸ì½”ë”©ì„ ê°€ì§„ í…스트'', ''char ì¸ì½”ë”©ì„ ì‚¬ìš©í•œ char í…스트'');' - 'INSERT INTO table1 VALUES (101,''char ì¸ì½”ë”©ì„ ì‚¬ìš©í•˜ëŠ” ì¼ë¶€ í…스트를 varcharë¡œ 변환해야합니다.'', ''clobë¡œ 변환해야하는 char ì¸ì½”ë”©ì„ ê°€ì§„ í…스트'', ''char ì¸ì½”ë”©ì„ ì‚¬ìš©í•œ char í…스트'');' - 'INSERT INTO table1 VALUES (101,''char ì¸ì½”ë”©ì„ ì‚¬ìš©í•˜ëŠ” ì¼ë¶€ í…스트를 varcharë¡œ 변환해야합니다.'', ''clobë¡œ 변환해야하는 char ì¸ì½”ë”©ì„ ê°€ì§„ í…스트'', ''char ì¸ì½”ë”©ì„ ì‚¬ìš©í•œ char í…스트'');' query: 'select varchar_field13, clob_field15, char_field16 from table1' drop: 'DROP table table1;' values: - 'VARCHAR_FIELD13=char ì¸ì½”ë”©ì„ ì‚¬ìš©í•˜ëŠ” ì¼ë¶€ í…스트를 varcharë¡œ 변환해야합니다.,char ì¸ì½”ë”©ì„ ì‚¬ìš©í•˜ëŠ” ì¼ë¶€ í…스트를 varcharë¡œ 변환해야합니다.,char ì¸ì½”ë”©ì„ ì‚¬ìš©í•˜ëŠ” ì¼ë¶€ í…스트를 varcharë¡œ 변환해야합니다., char ì¸ì½”ë”©ì„ ì‚¬ìš©í•˜ëŠ” ì¼ë¶€ í…스트를 varcharë¡œ 변환해야합니다.,char ì¸ì½”ë”©ì„ ì‚¬ìš©í•˜ëŠ” ì¼ë¶€ í…스트를 varcharë¡œ 변환해야합니다.' - 'CLOB_FIELD15=clobë¡œ 변환해야하는 char ì¸ì½”ë”©ì„ ê°€ì§„ í…스트,clobë¡œ 변환해야하는 char ì¸ì½”ë”©ì„ ê°€ì§„ í…스트,clobë¡œ 변환해야하는 char ì¸ì½”ë”©ì„ ê°€ì§„ í…스트, clobë¡œ 변환해야하는 char ì¸ì½”ë”©ì„ ê°€ì§„ í…스트,clobë¡œ 변환해야하는 char ì¸ì½”ë”©ì„ ê°€ì§„ í…스트' - 'CHAR_FIELD16=char ì¸ì½”ë”©ì„ ì‚¬ìš©í•œ char í…스트,char ì¸ì½”ë”©ì„ ì‚¬ìš©í•œ char í…스트,char ì¸ì½”ë”©ì„ ì‚¬ìš©í•œ char í…스트,char ì¸ì½”ë”©ì„ ì‚¬ìš©í•œ char í…스트,char ì¸ì½”ë”©ì„ ì‚¬ìš©í•œ char í…스트' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_clob_h2.yml000066400000000000000000000043711515451615200266710ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_clob_h2' type: 'clob' vector: 'CLOB_FIELD15' create: 'CREATE TABLE table1 (clob_field15 CLOB);' data: - 'INSERT INTO table1 VALUES (''some text that needs to be converted to clob'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to clob'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to clob'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to clob'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to clob'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to clob'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to clob'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to clob'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to clob'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to clob'');' query: 'select CLOB_FIELD15 from table1;' drop: 'DROP table table1;' values: - 'some text that needs to be converted to clob' - 'some text that needs to be converted to clob' - 'some text that needs to be converted to clob' - 'some text that needs to be converted to clob' - 'some text that needs to be converted to clob' - 'some text that needs to be converted to clob' - 'some text that needs to be converted to clob' - 'some text that needs to be converted to clob' - 'some text that needs to be converted to clob' - 'some text that needs to be converted to clob' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_date_h2.yml000066400000000000000000000031331515451615200266620ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_date_h2' type: 'date' vector: 'DATE_FIELD10' create: 'CREATE TABLE table1 (date_field10 DATE);' data: - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''1969-01-01'');' query: 'select date_field10 from table1;' drop: 'DROP table table1;' values: - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' - '-365' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_decimal_h2.yml000066400000000000000000000032161515451615200273450ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_decimal_h2' type: 'decimal' vector: 'DECIMAL_FIELD6' create: 'CREATE TABLE table1 (decimal_field6 DECIMAL(20,2));' data: - 'INSERT INTO table1 VALUES (17345667789.23);' - 'INSERT INTO table1 VALUES (17345667789.23);' - 'INSERT INTO table1 VALUES (17345667789.23);' - 'INSERT INTO table1 VALUES (17345667789.23);' - 'INSERT INTO table1 VALUES (17345667789.23);' - 'INSERT INTO table1 VALUES (17345667789.23);' - 'INSERT INTO table1 VALUES (17345667789.23);' - 'INSERT INTO table1 VALUES (17345667789.23);' - 'INSERT INTO table1 VALUES (17345667789.23);' - 'INSERT INTO table1 VALUES (17345667789.23);' query: 'select decimal_field6 from table1;' drop: 'DROP table table1;' values: - '17345667789.23' - '17345667789.23' - '17345667789.23' - '17345667789.23' - '17345667789.23' - '17345667789.23' - '17345667789.23' - '17345667789.23' - '17345667789.23' - '17345667789.23' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_double_h2.yml000066400000000000000000000032261515451615200272220ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_double_h2' type: 'double' vector: 'DOUBLE_FIELD7' create: 'CREATE TABLE table1 (double_field7 DOUBLE);' data: - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' query: 'select double_field7 from table1;' drop: 'DROP table table1;' values: - '56478356785.345' - '56478356785.345' - '56478356785.345' - '56478356785.345' - '56478356785.345' - '56478356785.345' - '56478356785.345' - '56478356785.345' - '56478356785.345' - '56478356785.345' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_est_date_h2.yml000066400000000000000000000030701515451615200275350ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_est_date_h2' type: 'est_date' timezone: 'EST' vector: 'DATE_FIELD10' create: 'CREATE TABLE table1 (date_field10 DATE);' data: - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' query: 'select date_field10 from table1;' drop: 'DROP table table1;' values: - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_est_time_h2.yml000066400000000000000000000030771515451615200275650ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_est_time_h2' type: 'est_time' timezone: 'EST' vector: 'TIME_FIELD9' create: 'CREATE TABLE table1 (time_field9 TIME);' data: - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' query: 'select time_field9 from table1;' drop: 'DROP table table1;' values: - '63935000' - '63935000' - '63935000' - '63935000' - '63935000' - '63935000' - '63935000' - '63935000' - '63935000' - '63935000' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_est_timestamp_h2.yml000066400000000000000000000034001515451615200306200ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_est_timestamp_h2' type: 'est_timestamp' timezone: 'EST' vector: 'TIMESTAMP_FIELD11' create: 'CREATE TABLE table1 (timestamp_field11 TIMESTAMP);' data: - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' query: 'select timestamp_field11 from table1;' drop: 'DROP table table1;' values: - '1518457535000' - '1518457535000' - '1518457535000' - '1518457535000' - '1518457535000' - '1518457535000' - '1518457535000' - '1518457535000' - '1518457535000' - '1518457535000' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_gmt_date_h2.yml000066400000000000000000000030701515451615200275310ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_gmt_date_h2' type: 'gmt_date' timezone: 'GMT' vector: 'DATE_FIELD10' create: 'CREATE TABLE table1 (date_field10 DATE);' data: - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' query: 'select date_field10 from table1;' drop: 'DROP table table1;' values: - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_gmt_time_h2.yml000066400000000000000000000030771515451615200275610ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_gmt_time_h2' type: 'gmt_time' timezone: 'GMT' vector: 'TIME_FIELD9' create: 'CREATE TABLE table1 (time_field9 TIME);' data: - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' query: 'select time_field9 from table1;' drop: 'DROP table table1;' values: - '45935000' - '45935000' - '45935000' - '45935000' - '45935000' - '45935000' - '45935000' - '45935000' - '45935000' - '45935000' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_gmt_timestamp_h2.yml000066400000000000000000000034001515451615200306140ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_gmt_timestamp_h2' type: 'gmt_timestamp' timezone: 'GMT' vector: 'TIMESTAMP_FIELD11' create: 'CREATE TABLE table1 (timestamp_field11 TIMESTAMP);' data: - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' query: 'select timestamp_field11 from table1;' drop: 'DROP table table1;' values: - '1518439535000' - '1518439535000' - '1518439535000' - '1518439535000' - '1518439535000' - '1518439535000' - '1518439535000' - '1518439535000' - '1518439535000' - '1518439535000' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_int_h2.yml000066400000000000000000000026361515451615200265460ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_int_h2' type: 'int' vector: 'INT_FIELD1' create: 'CREATE TABLE table1 (int_field1 INT);' data: - 'INSERT INTO table1 VALUES (101);' - 'INSERT INTO table1 VALUES (101);' - 'INSERT INTO table1 VALUES (101);' - 'INSERT INTO table1 VALUES (101);' - 'INSERT INTO table1 VALUES (101);' - 'INSERT INTO table1 VALUES (101);' - 'INSERT INTO table1 VALUES (101);' - 'INSERT INTO table1 VALUES (101);' - 'INSERT INTO table1 VALUES (101);' - 'INSERT INTO table1 VALUES (101);' values: - '101' - '101' - '101' - '101' - '101' - '101' - '101' - '101' - '101' - '101' query: 'select int_field1 from table1;' drop: 'DROP table table1;' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_list_h2.yml000066400000000000000000000030701515451615200267200ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_list_h2' type: 'list' vector: 'LIST_FIELD19' create: 'CREATE TABLE table1 (list_field19 INT ARRAY);' data: - 'INSERT INTO table1 VALUES (ARRAY[1, 2, 3]);' - 'INSERT INTO table1 VALUES (ARRAY[1, 2, 3]);' - 'INSERT INTO table1 VALUES (ARRAY[1, 2, 3]);' - 'INSERT INTO table1 VALUES (ARRAY[1, 2, 3]);' - 'INSERT INTO table1 VALUES (ARRAY[1, 2, 3]);' - 'INSERT INTO table1 VALUES (ARRAY[1, 2, 3]);' - 'INSERT INTO table1 VALUES (ARRAY[1, 2, 3]);' - 'INSERT INTO table1 VALUES (ARRAY[1, 2, 3]);' - 'INSERT INTO table1 VALUES (ARRAY[1, 2, 3]);' - 'INSERT INTO table1 VALUES (ARRAY[1, 2, 3]);' query: 'select list_field19 from table1;' drop: 'DROP table table1;' values: - '(1;2;3)' - '(1;2;3)' - '(1;2;3)' - '(1;2;3)' - '(1;2;3)' - '(1;2;3)' - '(1;2;3)' - '(1;2;3)' - '(1;2;3)' - '(1;2;3)' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_map_h2.yml000066400000000000000000000057641515451615200265360ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_map_h2' type: 'map' vector: 'MAP_FIELD20' create: 'CREATE TABLE table1 (map_field20 OTHER);' rowCount: '4' data: - 'INSERT INTO table1 VALUES (X''aced00057372002e6f72672e6170616368652e6172726f772e766563746f722e7574696c2e4a736f6e537472696e67486173684d61709819d7169e7a2ecf020000787200176a6176612e7574696c2e4c696e6b6564486173684d617034c04e5c106cc0fb0200015a000b6163636573734f72646572787200116a6176612e7574696c2e486173684d61700507dac1c31660d103000246000a6c6f6164466163746f724900097468726573686f6c6478703f4000000000000c7708000000100000000374000161740001627400033132337400067177657274797400057a78637662740001217800'');' - 'INSERT INTO table1 VALUES (X''aced00057372002e6f72672e6170616368652e6172726f772e766563746f722e7574696c2e4a736f6e537472696e67486173684d61709819d7169e7a2ecf020000787200176a6176612e7574696c2e4c696e6b6564486173684d617034c04e5c106cc0fb0200015a000b6163636573734f72646572787200116a6176612e7574696c2e486173684d61700507dac1c31660d103000246000a6c6f6164466163746f724900097468726573686f6c6478703f4000000000000c77080000001000000003740001617400016274000163740001647400033132337400067177657274797800'');' - 'INSERT INTO table1 VALUES (X''aced00057372002e6f72672e6170616368652e6172726f772e766563746f722e7574696c2e4a736f6e537472696e67486173684d61709819d7169e7a2ecf020000787200176a6176612e7574696c2e4c696e6b6564486173684d617034c04e5c106cc0fb0200015a000b6163636573734f72646572787200116a6176612e7574696c2e486173684d61700507dac1c31660d103000246000a6c6f6164466163746f724900097468726573686f6c6478703f4000000000000c7708000000100000000174000074000576616c75657800'');' - 'INSERT INTO table1 VALUES (X''aced00057372002e6f72672e6170616368652e6172726f772e766563746f722e7574696c2e4a736f6e537472696e67486173684d61709819d7169e7a2ecf020000787200176a6176612e7574696c2e4c696e6b6564486173684d617034c04e5c106cc0fb0200015a000b6163636573734f72646572787200116a6176612e7574696c2e486173684d61700507dac1c31660d103000246000a6c6f6164466163746f724900097468726573686f6c6478703f4000000000000c7708000000100000000274000b6e6f6e456d7074794b65797074000c736f6d654f746865724b65797400007800'');' query: 'select map_field20 from table1;' drop: 'DROP table table1;' values: - 'MAP_FIELD20={"a":"b"|"123":"qwerty"|"zxcvb":"!"},{"a":"b"|"123":"qwerty"|"c":"d"},{"":"value"},{"nonEmptyKey":null|"someOtherKey":""}' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_null_and_notnull.yml000066400000000000000000000021011515451615200307150ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_null_and_notnull' type: 'nullnotnull' create: 'CREATE TABLE table1 (int_field_null INT, int_field_notnull INT NOT NULL);' data: - 'INSERT INTO table1 VALUES (0, 0);' - 'INSERT INTO table1 VALUES (1, 1);' rowCount: '2' query: 'select int_field_null, int_field_notnull from table1;' drop: 'DROP table table1;' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_null_h2.yml000066400000000000000000000025241515451615200267220ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_null_h2' type: 'null' vector: 'NULL_FIELD18' create: 'CREATE TABLE table1 (null_field18 NULL);' rowCount: '10' data: - 'INSERT INTO table1 VALUES (null);' - 'INSERT INTO table1 VALUES (null);' - 'INSERT INTO table1 VALUES (null);' - 'INSERT INTO table1 VALUES (null);' - 'INSERT INTO table1 VALUES (null);' - 'INSERT INTO table1 VALUES (null);' - 'INSERT INTO table1 VALUES (null);' - 'INSERT INTO table1 VALUES (null);' - 'INSERT INTO table1 VALUES (null);' - 'INSERT INTO table1 VALUES (null);' query: 'select null_field18 from table1;' drop: 'DROP table table1;' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_pst_date_h2.yml000066400000000000000000000030701515451615200275500ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_pst_date_h2' type: 'pst_date' timezone: 'PST' vector: 'DATE_FIELD10' create: 'CREATE TABLE table1 (date_field10 DATE);' data: - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' - 'INSERT INTO table1 VALUES (''2018-02-12'');' query: 'select date_field10 from table1;' drop: 'DROP table table1;' values: - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' - '17574' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_pst_time_h2.yml000066400000000000000000000030771515451615200276000ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_pst_time_h2' type: 'pst_time' timezone: 'PST' vector: 'TIME_FIELD9' create: 'CREATE TABLE table1 (time_field9 TIME);' data: - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' query: 'select time_field9 from table1;' drop: 'DROP table table1;' values: - '74735000' - '74735000' - '74735000' - '74735000' - '74735000' - '74735000' - '74735000' - '74735000' - '74735000' - '74735000' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_pst_timestamp_h2.yml000066400000000000000000000034001515451615200306330ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_pst_timestamp_h2' type: 'pst_timestamp' timezone: 'PST' vector: 'TIMESTAMP_FIELD11' create: 'CREATE TABLE table1 (timestamp_field11 TIMESTAMP);' data: - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' query: 'select timestamp_field11 from table1;' drop: 'DROP table table1;' values: - '1518468335000' - '1518468335000' - '1518468335000' - '1518468335000' - '1518468335000' - '1518468335000' - '1518468335000' - '1518468335000' - '1518468335000' - '1518468335000' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_real_h2.yml000066400000000000000000000032241515451615200266710ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_real_h2' type: 'real' vector: 'REAL_FIELD8' create: 'CREATE TABLE table1 (real_field8 REAL);' data: - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' - 'INSERT INTO table1 VALUES (56478356785.345);' query: 'select real_field8 from table1;' drop: 'DROP table table1;' values: - '56478356785.345f' - '56478356785.345f' - '56478356785.345f' - '56478356785.345f' - '56478356785.345f' - '56478356785.345f' - '56478356785.345f' - '56478356785.345f' - '56478356785.345f' - '56478356785.345f' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_selected_datatypes_null_h2.yml000066400000000000000000000046301515451615200326500ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_selected_datatypes_null_h2' type: 'selected_null_column' vectors: - 'BIGINT_FIELD5' - 'DECIMAL_FIELD6' - 'DOUBLE_FIELD7' - 'REAL_FIELD8' - 'TIME_FIELD9' - 'DATE_FIELD10' - 'TIMESTAMP_FIELD11' - 'BINARY_FIELD12' - 'VARCHAR_FIELD13' - 'BLOB_FIELD14' - 'CLOB_FIELD15' - 'CHAR_FIELD16' - 'BIT_FIELD17' - 'LIST_FIELD19' - 'MAP_FIELD20' rowCount: '5' create: 'CREATE TABLE table1 (int_field1 INT, bool_field2 BOOLEAN, tinyint_field3 TINYINT, smallint_field4 SMALLINT, bigint_field5 BIGINT, decimal_field6 DECIMAL(20,2), double_field7 DOUBLE, real_field8 REAL, time_field9 TIME, date_field10 DATE, timestamp_field11 TIMESTAMP, binary_field12 VARBINARY(100), varchar_field13 VARCHAR(256), blob_field14 BLOB, clob_field15 CLOB, char_field16 CHAR(14), bit_field17 BIT, list_field19 INT ARRAY, map_field20 VARCHAR(256));' data: - 'INSERT INTO table1 (int_field1, bool_field2, tinyint_field3, smallint_field4) VALUES (102, 0, 46, 12001);' - 'INSERT INTO table1 (int_field1, bool_field2, tinyint_field3, smallint_field4) VALUES (102, 0, 46, 12001);' - 'INSERT INTO table1 (int_field1, bool_field2, tinyint_field3, smallint_field4) VALUES (102, 0, 46, 12001);' - 'INSERT INTO table1 (int_field1, bool_field2, tinyint_field3, smallint_field4) VALUES (102, 0, 46, 12001);' - 'INSERT INTO table1 (int_field1, bool_field2, tinyint_field3, smallint_field4) VALUES (102, 0, 46, 12001);' query: 'select bigint_field5, decimal_field6, double_field7, real_field8, time_field9, date_field10, timestamp_field11, binary_field12, varchar_field13, blob_field14, clob_field15, char_field16, bit_field17, list_field19, map_field20 from table1' drop: 'DROP table table1;' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_smallint_h2.yml000066400000000000000000000027331515451615200275750ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_smallint_h2' type: 'small_int' vector: 'SMALLINT_FIELD4' create: 'CREATE TABLE table1 (smallint_field4 SMALLINT);' data: - 'INSERT INTO table1 VALUES (12000);' - 'INSERT INTO table1 VALUES (12000);' - 'INSERT INTO table1 VALUES (12000);' - 'INSERT INTO table1 VALUES (12000);' - 'INSERT INTO table1 VALUES (12000);' - 'INSERT INTO table1 VALUES (12000);' - 'INSERT INTO table1 VALUES (12000);' - 'INSERT INTO table1 VALUES (12000);' - 'INSERT INTO table1 VALUES (12000);' - 'INSERT INTO table1 VALUES (12000);' query: 'select smallint_field4 from table1;' drop: 'DROP table table1;' values: - '12000' - '12000' - '12000' - '12000' - '12000' - '12000' - '12000' - '12000' - '12000' - '12000' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_time_h2.yml000066400000000000000000000030461515451615200267060ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_time_h2' type: 'time' vector: 'TIME_FIELD9' create: 'CREATE TABLE table1 (time_field9 TIME);' data: - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' - 'INSERT INTO table1 VALUES (''12:45:35'');' query: 'select time_field9 from table1;' drop: 'DROP table table1;' values: - '45935000' - '45935000' - '45935000' - '45935000' - '45935000' - '45935000' - '45935000' - '45935000' - '45935000' - '45935000' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_timestamp_h2.yml000066400000000000000000000033471515451615200277570ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_timestamp_h2' type: 'timestamp' vector: 'TIMESTAMP_FIELD11' create: 'CREATE TABLE table1 (timestamp_field11 TIMESTAMP);' data: - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' - 'INSERT INTO table1 VALUES (''2018-02-12 12:45:35'');' query: 'select timestamp_field11 from table1;' drop: 'DROP table table1;' values: - '1518439535000' - '1518439535000' - '1518439535000' - '1518439535000' - '1518439535000' - '1518439535000' - '1518439535000' - '1518439535000' - '1518439535000' - '1518439535000' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_tinyint_h2.yml000066400000000000000000000026301515451615200274440ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_tinyint_h2' type: 'tinyint' vector: 'TINYINT_FIELD3' create: 'CREATE TABLE table1 (tinyint_field3 TINYINT);' data: - 'INSERT INTO table1 VALUES (45);' - 'INSERT INTO table1 VALUES (45);' - 'INSERT INTO table1 VALUES (45);' - 'INSERT INTO table1 VALUES (45);' - 'INSERT INTO table1 VALUES (45);' - 'INSERT INTO table1 VALUES (45);' - 'INSERT INTO table1 VALUES (45);' - 'INSERT INTO table1 VALUES (45);' - 'INSERT INTO table1 VALUES (45);' - 'INSERT INTO table1 VALUES (45);' query: 'select tinyint_field3 from table1;' drop: 'DROP table table1;' values: - '45' - '45' - '45' - '45' - '45' - '45' - '45' - '45' - '45' - '45' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/h2/test1_varchar_h2.yml000066400000000000000000000045141515451615200273770ustar00rootroot00000000000000#Licensed to the Apache Software Foundation (ASF) under one or more contributor #license agreements. See the NOTICE file distributed with this work for additional #information regarding copyright ownership. The ASF licenses this file to #You under the Apache License, Version 2.0 (the "License"); you may not use #this file except in compliance with the License. You may obtain a copy of #the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required #by applicable law or agreed to in writing, software distributed under the #License is distributed on an "AS IS" BASIS, WITHOUT 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: 'test1_varchar_h2' type: 'varchar' vector: 'VARCHAR_FIELD13' create: 'CREATE TABLE table1 (varchar_field13 VARCHAR(256));' data: - 'INSERT INTO table1 VALUES (''some text that needs to be converted to varchar'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to varchar'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to varchar'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to varchar'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to varchar'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to varchar'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to varchar'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to varchar'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to varchar'');' - 'INSERT INTO table1 VALUES (''some text that needs to be converted to varchar'');' query: 'select varchar_field13 from table1;' drop: 'DROP table table1;' values: - 'some text that needs to be converted to varchar' - 'some text that needs to be converted to varchar' - 'some text that needs to be converted to varchar' - 'some text that needs to be converted to varchar' - 'some text that needs to be converted to varchar' - 'some text that needs to be converted to varchar' - 'some text that needs to be converted to varchar' - 'some text that needs to be converted to varchar' - 'some text that needs to be converted to varchar' - 'some text that needs to be converted to varchar' apache-arrow-java-19.0.0/adapter/jdbc/src/test/resources/logback.xml000066400000000000000000000025221515451615200253250ustar00rootroot00000000000000 %d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n apache-arrow-java-19.0.0/adapter/orc/000077500000000000000000000000001515451615200173015ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/CMakeLists.txt000066400000000000000000000041661515451615200220500ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. find_package(Arrow REQUIRED) include_directories(${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_SOURCE_DIR} ${JNI_INCLUDE_DIRS} ${JNI_HEADERS_DIR}) add_jar(arrow_java_jni_orc_jar src/main/java/org/apache/arrow/adapter/orc/OrcReaderJniWrapper.java src/main/java/org/apache/arrow/adapter/orc/OrcStripeReaderJniWrapper.java src/main/java/org/apache/arrow/adapter/orc/OrcMemoryJniWrapper.java src/main/java/org/apache/arrow/adapter/orc/OrcJniUtils.java src/main/java/org/apache/arrow/adapter/orc/OrcRecordBatch.java src/main/java/org/apache/arrow/adapter/orc/OrcFieldNode.java GENERATE_NATIVE_HEADERS arrow_java_jni_orc_headers DESTINATION ${JNI_HEADERS_DIR}) add_library(arrow_java_jni_orc SHARED src/main/cpp/jni_wrapper.cpp) set_property(TARGET arrow_java_jni_orc PROPERTY OUTPUT_NAME "arrow_orc_jni") target_link_libraries(arrow_java_jni_orc arrow_java_jni_orc_headers jni Arrow::arrow_static) set(ARROW_JAVA_JNI_ORC_LIBDIR "${CMAKE_INSTALL_PREFIX}/lib/arrow_orc_jni/${ARROW_JAVA_JNI_ARCH_DIR}") set(ARROW_JAVA_JNI_ORC_BINDIR "${CMAKE_INSTALL_PREFIX}/bin/arrow_orc_jni/${ARROW_JAVA_JNI_ARCH_DIR}") install(TARGETS arrow_java_jni_orc LIBRARY DESTINATION ${ARROW_JAVA_JNI_ORC_LIBDIR} RUNTIME DESTINATION ${ARROW_JAVA_JNI_ORC_BINDIR}) apache-arrow-java-19.0.0/adapter/orc/pom.xml000066400000000000000000000132011515451615200206130ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-java-root 19.0.0 ../../pom.xml org.apache.arrow.orc arrow-orc jar Arrow Orc Adapter (Experimental/Contrib)A JNI wrapper for the C++ ORC reader implementation. ../../../cpp/release-build/ org.apache.arrow arrow-memory-core compile org.apache.arrow arrow-memory-netty runtime org.apache.arrow arrow-vector ${arrow.vector.classifier} compile org.immutables value-annotations org.apache.orc orc-core 2.3.0 test log4j log4j org.slf4j slf4j-log4j12 commons-logging commons-logging org.apache.hadoop hadoop-client-runtime ${dep.hadoop.version} test commons-logging commons-logging org.apache.hadoop hadoop-client-api ${dep.hadoop.version} org.apache.hadoop hadoop-common ${dep.hadoop.version} test commons-logging commons-logging log4j log4j org.slf4j slf4j-log4j12 javax.servlet servlet-api org.slf4j slf4j-reload4j ch.qos.reload4j reload4j org.apache.hive hive-storage-api 4.0.1 test ${arrow.cpp.build.dir} **/*arrow_orc_jni.* org.apache.maven.plugins maven-dependency-plugin analyze org.apache.arrow:arrow-format org.apache.maven.plugins maven-compiler-plugin -Werror apache-arrow-java-19.0.0/adapter/orc/src/000077500000000000000000000000001515451615200200705ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/main/000077500000000000000000000000001515451615200210145ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/main/cpp/000077500000000000000000000000001515451615200215765ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/main/cpp/concurrent_map.h000066400000000000000000000041161515451615200247700ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and */ #pragma once #include #include #include #include #include "arrow/util/macros.h" namespace arrow { namespace jni { /** * An utility class that map module id to module pointers. * @tparam Holder class of the object to hold. */ template class ConcurrentMap { public: ConcurrentMap() : module_id_(init_module_id_) {} jlong Insert(Holder holder) { std::lock_guard lock(mtx_); jlong result = module_id_++; map_.insert(std::pair(result, holder)); return result; } void Erase(jlong module_id) { std::lock_guard lock(mtx_); map_.erase(module_id); } Holder Lookup(jlong module_id) { std::lock_guard lock(mtx_); auto it = map_.find(module_id); if (it != map_.end()) { return it->second; } return NULLPTR; } void Clear() { std::lock_guard lock(mtx_); map_.clear(); } private: // Initialize the module id starting value to a number greater than zero // to allow for easier debugging of uninitialized java variables. static constexpr int init_module_id_ = 4; int64_t module_id_; std::mutex mtx_; // map from module ids returned to Java and module pointers std::unordered_map map_; }; } // namespace jni } // namespace arrow apache-arrow-java-19.0.0/adapter/orc/src/main/cpp/jni_wrapper.cpp000066400000000000000000000254631515451615200246340ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #include #include #include #include #include #include #include #include #include #include #include "org_apache_arrow_adapter_orc_OrcMemoryJniWrapper.h" #include "org_apache_arrow_adapter_orc_OrcReaderJniWrapper.h" #include "org_apache_arrow_adapter_orc_OrcStripeReaderJniWrapper.h" #include "./concurrent_map.h" using ORCFileReader = arrow::adapters::orc::ORCFileReader; using RecordBatchReader = arrow::RecordBatchReader; static jclass io_exception_class; static jclass illegal_access_exception_class; static jclass illegal_argument_exception_class; static jclass orc_field_node_class; static jmethodID orc_field_node_constructor; static jclass orc_memory_class; static jmethodID orc_memory_constructor; static jclass record_batch_class; static jmethodID record_batch_constructor; static jint JNI_VERSION = JNI_VERSION_10; using arrow::internal::checked_cast; using arrow::jni::ConcurrentMap; static ConcurrentMap> buffer_holder_; static ConcurrentMap> orc_stripe_reader_holder_; static ConcurrentMap> orc_reader_holder_; jclass CreateGlobalClassReference(JNIEnv* env, const char* class_name) { jclass local_class = env->FindClass(class_name); jclass global_class = (jclass)env->NewGlobalRef(local_class); env->DeleteLocalRef(local_class); return global_class; } jmethodID GetMethodID(JNIEnv* env, jclass this_class, const char* name, const char* sig) { jmethodID ret = env->GetMethodID(this_class, name, sig); if (ret == nullptr) { std::string error_message = "Unable to find method " + std::string(name) + " within signature" + std::string(sig); env->ThrowNew(illegal_access_exception_class, error_message.c_str()); } return ret; } std::string JStringToCString(JNIEnv* env, jstring string) { int32_t jlen, clen; clen = env->GetStringUTFLength(string); jlen = env->GetStringLength(string); std::vector buffer(clen); env->GetStringUTFRegion(string, 0, jlen, buffer.data()); return std::string(buffer.data(), clen); } std::shared_ptr GetFileReader(JNIEnv* env, jlong id) { auto reader = orc_reader_holder_.Lookup(id); if (!reader) { std::string error_message = "invalid reader id " + std::to_string(id); env->ThrowNew(illegal_argument_exception_class, error_message.c_str()); } return reader; } std::shared_ptr GetStripeReader(JNIEnv* env, jlong id) { auto reader = orc_stripe_reader_holder_.Lookup(id); if (!reader) { std::string error_message = "invalid stripe reader id " + std::to_string(id); env->ThrowNew(illegal_argument_exception_class, error_message.c_str()); } return reader; } #ifdef __cplusplus extern "C" { #endif jint JNI_OnLoad(JavaVM* vm, void* reserved) { JNIEnv* env; if (vm->GetEnv(reinterpret_cast(&env), JNI_VERSION) != JNI_OK) { return JNI_ERR; } io_exception_class = CreateGlobalClassReference(env, "Ljava/io/IOException;"); illegal_access_exception_class = CreateGlobalClassReference(env, "Ljava/lang/IllegalAccessException;"); illegal_argument_exception_class = CreateGlobalClassReference(env, "Ljava/lang/IllegalArgumentException;"); orc_field_node_class = CreateGlobalClassReference(env, "Lorg/apache/arrow/adapter/orc/OrcFieldNode;"); orc_field_node_constructor = GetMethodID(env, orc_field_node_class, "", "(II)V"); orc_memory_class = CreateGlobalClassReference( env, "Lorg/apache/arrow/adapter/orc/OrcMemoryJniWrapper;"); orc_memory_constructor = GetMethodID(env, orc_memory_class, "", "(JJJJ)V"); record_batch_class = CreateGlobalClassReference(env, "Lorg/apache/arrow/adapter/orc/OrcRecordBatch;"); record_batch_constructor = GetMethodID(env, record_batch_class, "", "(I[Lorg/apache/arrow/adapter/orc/OrcFieldNode;" "[Lorg/apache/arrow/adapter/orc/OrcMemoryJniWrapper;)V"); env->ExceptionDescribe(); return JNI_VERSION; } void JNI_OnUnload(JavaVM* vm, void* reserved) { JNIEnv* env; vm->GetEnv(reinterpret_cast(&env), JNI_VERSION); env->DeleteGlobalRef(io_exception_class); env->DeleteGlobalRef(illegal_access_exception_class); env->DeleteGlobalRef(illegal_argument_exception_class); env->DeleteGlobalRef(orc_field_node_class); env->DeleteGlobalRef(orc_memory_class); env->DeleteGlobalRef(record_batch_class); buffer_holder_.Clear(); orc_stripe_reader_holder_.Clear(); orc_reader_holder_.Clear(); } JNIEXPORT jlong JNICALL Java_org_apache_arrow_adapter_orc_OrcReaderJniWrapper_open( JNIEnv* env, jobject this_obj, jstring file_path) { std::string path = JStringToCString(env, file_path); if (path.find("hdfs://") == 0) { env->ThrowNew(io_exception_class, "hdfs path not supported yet."); } auto maybe_file = arrow::io::ReadableFile::Open(path); if (!maybe_file.ok()) { return -static_cast(maybe_file.status().code()); } auto maybe_reader = ORCFileReader::Open(*maybe_file, arrow::default_memory_pool()); if (!maybe_reader.ok()) { env->ThrowNew(io_exception_class, std::string("Failed open file" + path).c_str()); } return orc_reader_holder_.Insert( std::shared_ptr(*std::move(maybe_reader))); } JNIEXPORT void JNICALL Java_org_apache_arrow_adapter_orc_OrcReaderJniWrapper_close( JNIEnv* env, jobject this_obj, jlong id) { orc_reader_holder_.Erase(id); } JNIEXPORT jboolean JNICALL Java_org_apache_arrow_adapter_orc_OrcReaderJniWrapper_seek( JNIEnv* env, jobject this_obj, jlong id, jint row_number) { auto reader = GetFileReader(env, id); return reader->Seek(row_number).ok(); } JNIEXPORT jint JNICALL Java_org_apache_arrow_adapter_orc_OrcReaderJniWrapper_getNumberOfStripes(JNIEnv* env, jobject this_obj, jlong id) { auto reader = GetFileReader(env, id); return reader->NumberOfStripes(); } JNIEXPORT jlong JNICALL Java_org_apache_arrow_adapter_orc_OrcReaderJniWrapper_nextStripeReader(JNIEnv* env, jobject this_obj, jlong id, jlong batch_size) { auto reader = GetFileReader(env, id); auto maybe_stripe_reader = reader->NextStripeReader(batch_size); if (!maybe_stripe_reader.ok()) { return static_cast(maybe_stripe_reader.status().code()) * -1; } if (*maybe_stripe_reader == nullptr) { return static_cast(arrow::StatusCode::Invalid) * -1; } return orc_stripe_reader_holder_.Insert(*maybe_stripe_reader); } JNIEXPORT jbyteArray JNICALL Java_org_apache_arrow_adapter_orc_OrcStripeReaderJniWrapper_getSchema(JNIEnv* env, jclass this_cls, jlong id) { auto stripe_reader = GetStripeReader(env, id); auto schema = stripe_reader->schema(); auto maybe_buffer = arrow::ipc::SerializeSchema(*schema, arrow::default_memory_pool()); if (!maybe_buffer.ok()) { return nullptr; } auto buffer = *std::move(maybe_buffer); jbyteArray ret = env->NewByteArray(buffer->size()); auto src = reinterpret_cast(buffer->data()); env->SetByteArrayRegion(ret, 0, buffer->size(), src); return ret; } JNIEXPORT jobject JNICALL Java_org_apache_arrow_adapter_orc_OrcStripeReaderJniWrapper_next(JNIEnv* env, jclass this_cls, jlong id) { auto stripe_reader = GetStripeReader(env, id); std::shared_ptr record_batch; auto status = stripe_reader->ReadNext(&record_batch); if (!status.ok() || !record_batch) { return nullptr; } auto schema = stripe_reader->schema(); // TODO: ARROW-4714 Ensure JVM has sufficient capacity to create local references // create OrcFieldNode[] jobjectArray field_array = env->NewObjectArray(schema->num_fields(), orc_field_node_class, nullptr); std::vector> buffers; for (int i = 0; i < schema->num_fields(); ++i) { auto column = record_batch->column(i); auto dataArray = column->data(); jobject field = env->NewObject(orc_field_node_class, orc_field_node_constructor, column->length(), column->null_count()); env->SetObjectArrayElement(field_array, i, field); for (auto& buffer : dataArray->buffers) { buffers.push_back(buffer); } } // create OrcMemoryJniWrapper[] jobjectArray memory_array = env->NewObjectArray(buffers.size(), orc_memory_class, nullptr); for (size_t j = 0; j < buffers.size(); ++j) { auto buffer = buffers[j]; uint8_t* data = nullptr; int size = 0; int64_t capacity = 0; if (buffer != nullptr) { data = (uint8_t*)buffer->data(); size = (int)buffer->size(); capacity = buffer->capacity(); } jobject memory = env->NewObject(orc_memory_class, orc_memory_constructor, buffer_holder_.Insert(buffer), data, size, capacity); env->SetObjectArrayElement(memory_array, j, memory); } // create OrcRecordBatch jobject ret = env->NewObject(record_batch_class, record_batch_constructor, record_batch->num_rows(), field_array, memory_array); return ret; } JNIEXPORT void JNICALL Java_org_apache_arrow_adapter_orc_OrcStripeReaderJniWrapper_close( JNIEnv* env, jclass this_cls, jlong id) { orc_stripe_reader_holder_.Erase(id); } JNIEXPORT void JNICALL Java_org_apache_arrow_adapter_orc_OrcMemoryJniWrapper_release( JNIEnv* env, jobject this_obj, jlong id) { buffer_holder_.Erase(id); } #ifdef __cplusplus } #endif apache-arrow-java-19.0.0/adapter/orc/src/main/java/000077500000000000000000000000001515451615200217355ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/main/java/module-info.java000066400000000000000000000017431515451615200250230ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ open module org.apache.arrow.adapter.orc { exports org.apache.arrow.adapter.orc; requires hadoop.client.api; requires org.apache.arrow.memory.core; requires org.apache.arrow.vector; } apache-arrow-java-19.0.0/adapter/orc/src/main/java/org/000077500000000000000000000000001515451615200225245ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/main/java/org/apache/000077500000000000000000000000001515451615200237455ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200250775ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/main/java/org/apache/arrow/adapter/000077500000000000000000000000001515451615200265175ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/main/java/org/apache/arrow/adapter/orc/000077500000000000000000000000001515451615200273025ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/main/java/org/apache/arrow/adapter/orc/OrcFieldNode.java000066400000000000000000000025041515451615200324430ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.orc; /** Metadata about Vectors/Arrays that is passed via JNI interface. */ class OrcFieldNode { private final int length; private final int nullCount; /** * Construct a new instance. * * @param length the number of values written. * @param nullCount the number of null values. */ public OrcFieldNode(int length, int nullCount) { this.length = length; this.nullCount = nullCount; } int getLength() { return length; } int getNullCount() { return nullCount; } } apache-arrow-java-19.0.0/adapter/orc/src/main/java/org/apache/arrow/adapter/orc/OrcJniUtils.java000066400000000000000000000051131515451615200323520ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.orc; import java.io.File; import java.io.FileNotFoundException; import java.io.IOException; import java.io.InputStream; import java.nio.file.Files; import java.nio.file.StandardCopyOption; import java.util.Locale; /** Helper class for JNI related operations. */ class OrcJniUtils { private static final String LIBRARY_NAME = "arrow_orc_jni"; private static boolean isLoaded = false; private OrcJniUtils() {} static void loadOrcAdapterLibraryFromJar() throws IOException, IllegalAccessException { synchronized (OrcJniUtils.class) { if (!isLoaded) { final String libraryToLoad = LIBRARY_NAME + "/" + getNormalizedArch() + "/" + System.mapLibraryName(LIBRARY_NAME); final File libraryFile = moveFileFromJarToTemp( System.getProperty("java.io.tmpdir"), libraryToLoad, LIBRARY_NAME); System.load(libraryFile.getAbsolutePath()); isLoaded = true; } } } private static String getNormalizedArch() { String arch = System.getProperty("os.arch").toLowerCase(Locale.US); switch (arch) { case "amd64": arch = "x86_64"; break; case "aarch64": arch = "aarch_64"; break; default: break; } return arch; } private static File moveFileFromJarToTemp( final String tmpDir, String libraryToLoad, String libraryName) throws IOException { final File temp = File.createTempFile(tmpDir, libraryName); try (final InputStream is = OrcReaderJniWrapper.class.getClassLoader().getResourceAsStream(libraryToLoad)) { if (is == null) { throw new FileNotFoundException(libraryToLoad); } else { Files.copy(is, temp.toPath(), StandardCopyOption.REPLACE_EXISTING); } } return temp; } } OrcMemoryJniWrapper.java000066400000000000000000000042621515451615200340100ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/main/java/org/apache/arrow/adapter/orc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.orc; /** Wrapper for orc memory allocated by native code. */ class OrcMemoryJniWrapper implements AutoCloseable { private final long nativeInstanceId; private final long memoryAddress; private final long size; private final long capacity; /** * Construct a new instance. * * @param nativeInstanceId unique id of the underlying memory. * @param memoryAddress starting memory address of the underlying memory. * @param size size of the valid data. * @param capacity allocated memory size. */ OrcMemoryJniWrapper(long nativeInstanceId, long memoryAddress, long size, long capacity) { this.nativeInstanceId = nativeInstanceId; this.memoryAddress = memoryAddress; this.size = size; this.capacity = capacity; } /** * Return the size of underlying chunk of memory that has valid data. * * @return valid data size */ long getSize() { return size; } /** * Return the size of underlying chunk of memory managed by this OrcMemoryJniWrapper. * * @return underlying memory size */ long getCapacity() { return capacity; } /** * Return the memory address of underlying chunk of memory. * * @return memory address */ long getMemoryAddress() { return memoryAddress; } @Override public void close() { release(nativeInstanceId); } private native void release(long id); } apache-arrow-java-19.0.0/adapter/orc/src/main/java/org/apache/arrow/adapter/orc/OrcReader.java000066400000000000000000000061001515451615200320100ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.orc; import java.io.IOException; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.ipc.ArrowReader; /** * Orc Reader that allow accessing orc stripes in Orc file. This orc reader basically acts like an * ArrowReader iterator that iterate over orc stripes. Each stripe will be accessed via an * ArrowReader. */ public class OrcReader implements AutoCloseable { private final OrcReaderJniWrapper jniWrapper; private BufferAllocator allocator; /** reference to native reader instance. */ private final long nativeInstanceId; /** * Create an OrcReader that iterate over orc stripes. * * @param filePath file path to target file, currently only support local file. * @param allocator allocator provided to ArrowReader. * @throws IOException throws exception in case of file not found */ public OrcReader(String filePath, BufferAllocator allocator) throws IOException, IllegalAccessException { this.allocator = allocator; this.jniWrapper = OrcReaderJniWrapper.getInstance(); this.nativeInstanceId = jniWrapper.open(filePath); } /** * Seek to designated row. Invoke NextStripeReader() after seek will return stripe reader starting * from designated row. * * @param rowNumber the rows number to seek * @return true if seek operation is succeeded */ public boolean seek(int rowNumber) throws IllegalArgumentException { return jniWrapper.seek(nativeInstanceId, rowNumber); } /** * Get a stripe level ArrowReader with specified batchSize in each record batch. * * @param batchSize the number of rows loaded on each iteration * @return ArrowReader that iterate over current stripes */ public ArrowReader nextStripeReader(long batchSize) throws IllegalArgumentException { long stripeReaderId = jniWrapper.nextStripeReader(nativeInstanceId, batchSize); if (stripeReaderId < 0) { return null; } return new OrcStripeReader(stripeReaderId, allocator); } /** * The number of stripes in the file. * * @return number of stripes */ public int getNumberOfStripes() throws IllegalArgumentException { return jniWrapper.getNumberOfStripes(nativeInstanceId); } @Override public void close() { jniWrapper.close(nativeInstanceId); } } OrcReaderJniWrapper.java000066400000000000000000000051241515451615200337400ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/main/java/org/apache/arrow/adapter/orc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.orc; import java.io.IOException; /** JNI wrapper for Orc reader. */ class OrcReaderJniWrapper { private static volatile OrcReaderJniWrapper INSTANCE; static OrcReaderJniWrapper getInstance() throws IOException, IllegalAccessException { if (INSTANCE == null) { synchronized (OrcReaderJniWrapper.class) { if (INSTANCE == null) { OrcJniUtils.loadOrcAdapterLibraryFromJar(); INSTANCE = new OrcReaderJniWrapper(); } } } return INSTANCE; } /** * Construct a orc file reader over the target file. * * @param fileName absolute file path of target file * @return id of the orc reader instance if file opened successfully, otherwise return error code * * -1. */ native long open(String fileName); /** * Release resources associated with designated reader instance. * * @param readerId id of the reader instance. */ native void close(long readerId); /** * Seek to designated row. Invoke nextStripeReader() after seek will return id of stripe reader * starting from designated row. * * @param readerId id of the reader instance * @param rowNumber the rows number to seek * @return true if seek operation is succeeded */ native boolean seek(long readerId, int rowNumber); /** * The number of stripes in the file. * * @param readerId id of the reader instance * @return number of stripes */ native int getNumberOfStripes(long readerId); /** * Get a stripe level ArrowReader with specified batchSize in each record batch. * * @param readerId id of the reader instance * @param batchSize the number of rows loaded on each iteration * @return id of the stripe reader instance. */ native long nextStripeReader(long readerId, long batchSize); } apache-arrow-java-19.0.0/adapter/orc/src/main/java/org/apache/arrow/adapter/orc/OrcRecordBatch.java000066400000000000000000000030041515451615200327660ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.orc; import java.util.Arrays; import java.util.List; /** Wrapper for record batch meta and native memory. */ class OrcRecordBatch { final int length; /** Nodes correspond to the pre-ordered flattened logical schema. */ final List nodes; final List buffers; /** * Construct a new instance. * * @param length number of records included in current batch * @param nodes meta data for each fields * @param buffers buffers for underlying data */ OrcRecordBatch(int length, OrcFieldNode[] nodes, OrcMemoryJniWrapper[] buffers) { this.length = length; this.nodes = Arrays.asList(nodes); this.buffers = Arrays.asList(buffers); } } OrcReferenceManager.java000066400000000000000000000070321515451615200337250ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/main/java/org/apache/arrow/adapter/orc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.orc; import java.util.concurrent.atomic.AtomicInteger; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.OwnershipTransferResult; import org.apache.arrow.memory.ReferenceManager; import org.apache.arrow.util.Preconditions; /** * A simple reference manager implementation for memory allocated by native code. The underlying * memory will be released when reference count reach zero. */ public class OrcReferenceManager implements ReferenceManager { private final AtomicInteger bufRefCnt = new AtomicInteger(0); private OrcMemoryJniWrapper memory; OrcReferenceManager(OrcMemoryJniWrapper memory) { this.memory = memory; } @Override public int getRefCount() { return bufRefCnt.get(); } @Override public boolean release() { return release(1); } @Override public boolean release(int decrement) { Preconditions.checkState( decrement >= 1, "ref count decrement should be greater than or equal to 1"); // decrement the ref count final int refCnt; synchronized (this) { refCnt = bufRefCnt.addAndGet(-decrement); if (refCnt == 0) { // refcount of this reference manager has dropped to 0 // release the underlying memory memory.close(); } } // the new ref count should be >= 0 Preconditions.checkState(refCnt >= 0, "RefCnt has gone negative"); return refCnt == 0; } @Override public void retain() { retain(1); } @Override public void retain(int increment) { Preconditions.checkArgument(increment > 0, "retain(%s) argument is not positive", increment); bufRefCnt.addAndGet(increment); } @Override public ArrowBuf retain(ArrowBuf srcBuffer, BufferAllocator targetAllocator) { retain(); return srcBuffer; } @Override public ArrowBuf deriveBuffer(ArrowBuf sourceBuffer, long index, long length) { final long derivedBufferAddress = sourceBuffer.memoryAddress() + index; // create new ArrowBuf final ArrowBuf derivedBuf = new ArrowBuf( this, null, length, // length (in bytes) in the underlying memory chunk for this new ArrowBuf derivedBufferAddress // starting byte address in the underlying memory for this new // ArrowBuf, ); return derivedBuf; } @Override public OwnershipTransferResult transferOwnership( ArrowBuf sourceBuffer, BufferAllocator targetAllocator) { throw new UnsupportedOperationException(); } @Override public BufferAllocator getAllocator() { return null; } @Override public long getSize() { return memory.getSize(); } @Override public long getAccountedSize() { return 0; } } apache-arrow-java-19.0.0/adapter/orc/src/main/java/org/apache/arrow/adapter/orc/OrcStripeReader.java000066400000000000000000000072071515451615200332100ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.orc; import java.io.IOException; import java.util.ArrayList; import java.util.stream.Collectors; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.ipc.ReadChannel; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.ipc.message.MessageChannelReader; import org.apache.arrow.vector.ipc.message.MessageResult; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel; /** Orc stripe that load data into ArrowRecordBatch. */ public class OrcStripeReader extends ArrowReader { /** reference to native stripe reader instance. */ private final long nativeInstanceId; /** * Construct a new instance. * * @param nativeInstanceId nativeInstanceId of the stripe reader instance, obtained by calling * nextStripeReader from OrcReaderJniWrapper * @param allocator memory allocator for accounting. */ OrcStripeReader(long nativeInstanceId, BufferAllocator allocator) { super(allocator); this.nativeInstanceId = nativeInstanceId; } @Override public boolean loadNextBatch() throws IOException { OrcRecordBatch recordBatch = OrcStripeReaderJniWrapper.next(nativeInstanceId); if (recordBatch == null) { return false; } ArrayList buffers = new ArrayList<>(); for (OrcMemoryJniWrapper buffer : recordBatch.buffers) { buffers.add( new ArrowBuf( new OrcReferenceManager(buffer), null, (int) buffer.getSize(), buffer.getMemoryAddress())); } loadRecordBatch( new ArrowRecordBatch( recordBatch.length, recordBatch.nodes.stream() .map(buf -> new ArrowFieldNode(buf.getLength(), buf.getNullCount())) .collect(Collectors.toList()), buffers)); return true; } @Override public long bytesRead() { return 0; } @Override protected void closeReadSource() throws IOException { OrcStripeReaderJniWrapper.close(nativeInstanceId); } @Override protected Schema readSchema() throws IOException { byte[] schemaBytes = OrcStripeReaderJniWrapper.getSchema(nativeInstanceId); try (MessageChannelReader schemaReader = new MessageChannelReader( new ReadChannel(new ByteArrayReadableSeekableByteChannel(schemaBytes)), allocator)) { MessageResult result = schemaReader.readNext(); if (result == null) { throw new IOException("Unexpected end of input. Missing schema."); } return MessageSerializer.deserializeSchema(result.getMessage()); } } } OrcStripeReaderJniWrapper.java000066400000000000000000000027501515451615200351310ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/main/java/org/apache/arrow/adapter/orc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.orc; /** JNI wrapper for orc stripe reader. */ class OrcStripeReaderJniWrapper { /** * Get the schema of current stripe. * * @param readerId id of the stripe reader instance. * @return serialized schema. */ static native byte[] getSchema(long readerId); /** * Load next record batch. * * @param readerId id of the stripe reader instance. * @return loaded record batch, return null when reached the end of current stripe. */ static native OrcRecordBatch next(long readerId); /** * Release resources of underlying reader. * * @param readerId id of the stripe reader instance. */ static native void close(long readerId); } apache-arrow-java-19.0.0/adapter/orc/src/test/000077500000000000000000000000001515451615200210475ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/test/java/000077500000000000000000000000001515451615200217705ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/test/java/org/000077500000000000000000000000001515451615200225575ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/test/java/org/apache/000077500000000000000000000000001515451615200240005ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200251325ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/test/java/org/apache/arrow/adapter/000077500000000000000000000000001515451615200265525ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/test/java/org/apache/arrow/adapter/orc/000077500000000000000000000000001515451615200273355ustar00rootroot00000000000000apache-arrow-java-19.0.0/adapter/orc/src/test/java/org/apache/arrow/adapter/orc/OrcReaderTest.java000066400000000000000000000076111515451615200327130ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.orc; 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 java.io.File; import java.nio.charset.StandardCharsets; import java.util.List; 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.ipc.ArrowReader; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.Path; import org.apache.hadoop.hive.ql.exec.vector.BytesColumnVector; import org.apache.hadoop.hive.ql.exec.vector.LongColumnVector; import org.apache.hadoop.hive.ql.exec.vector.VectorizedRowBatch; import org.apache.orc.OrcFile; 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; public class OrcReaderTest { @TempDir public File testFolder; private static final int MAX_ALLOCATION = 8 * 1024; private static RootAllocator allocator; @BeforeAll 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"); File testFile = new File(testFolder, "test-orc"); Writer writer = OrcFile.createWriter( new Path(testFile.getAbsolutePath()), OrcFile.writerOptions(new Configuration()).setSchema(schema)); VectorizedRowBatch batch = schema.createRowBatch(); LongColumnVector longColumnVector = (LongColumnVector) batch.cols[0]; BytesColumnVector bytesColumnVector = (BytesColumnVector) batch.cols[1]; for (int r = 0; r < 1024; ++r) { int row = batch.size++; longColumnVector.vector[row] = r; byte[] buffer = ("Last-" + (r * 3)).getBytes(StandardCharsets.UTF_8); bytesColumnVector.setRef(row, buffer, 0, buffer.length); } writer.addRowBatch(batch); writer.close(); OrcReader reader = new OrcReader(testFile.getAbsolutePath(), allocator); assertEquals(1, reader.getNumberOfStripes()); ArrowReader stripeReader = reader.nextStripeReader(1024); VectorSchemaRoot schemaRoot = stripeReader.getVectorSchemaRoot(); stripeReader.loadNextBatch(); List fields = schemaRoot.getFieldVectors(); assertEquals(2, fields.size()); IntVector intVector = (IntVector) fields.get(0); VarCharVector varCharVector = (VarCharVector) fields.get(1); for (int i = 0; i < 1024; ++i) { assertEquals(i, intVector.get(i)); assertEquals("Last-" + (i * 3), new String(varCharVector.get(i), StandardCharsets.UTF_8)); } assertFalse(stripeReader.loadNextBatch()); assertNull(reader.nextStripeReader(1024)); stripeReader.close(); reader.close(); } } apache-arrow-java-19.0.0/algorithm/000077500000000000000000000000001515451615200170645ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/pom.xml000066400000000000000000000042521515451615200204040ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-java-root 19.0.0 arrow-algorithm Arrow Algorithms (Experimental/Contrib) A collection of algorithms for working with ValueVectors. org.apache.arrow arrow-vector ${arrow.vector.classifier} org.apache.arrow arrow-vector ${project.version} test-jar test org.apache.arrow arrow-memory-core org.apache.arrow arrow-memory-netty test org.immutables value-annotations apache-arrow-java-19.0.0/algorithm/src/000077500000000000000000000000001515451615200176535ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/000077500000000000000000000000001515451615200205775ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/000077500000000000000000000000001515451615200215205ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/module-info.java000066400000000000000000000023031515451615200245770ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ module org.apache.arrow.algorithm { exports org.apache.arrow.algorithm.search; exports org.apache.arrow.algorithm.deduplicate; exports org.apache.arrow.algorithm.dictionary; exports org.apache.arrow.algorithm.rank; exports org.apache.arrow.algorithm.misc; exports org.apache.arrow.algorithm.sort; requires jdk.unsupported; requires org.apache.arrow.memory.core; requires org.apache.arrow.vector; } apache-arrow-java-19.0.0/algorithm/src/main/java/org/000077500000000000000000000000001515451615200223075ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/000077500000000000000000000000001515451615200235305ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200246625ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/000077500000000000000000000000001515451615200266505ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/deduplicate/000077500000000000000000000000001515451615200311335ustar00rootroot00000000000000DeduplicationUtils.java000066400000000000000000000071151515451615200355300ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/deduplicate/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.deduplicate; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.compare.Range; import org.apache.arrow.vector.compare.RangeEqualsVisitor; import org.apache.arrow.vector.util.DataSizeRoundingUtil; /** Utilities for vector deduplication. */ class DeduplicationUtils { /** * Gets the start positions of the first distinct values in a vector. * * @param vector the target vector. * @param runStarts the bit set to hold the start positions. * @param vector type. */ public static void populateRunStartIndicators( V vector, ArrowBuf runStarts) { int bufSize = DataSizeRoundingUtil.divideBy8Ceil(vector.getValueCount()); Preconditions.checkArgument(runStarts.capacity() >= bufSize); runStarts.setZero(0, bufSize); BitVectorHelper.setBit(runStarts, 0); RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector, vector, null); Range range = new Range(0, 0, 1); for (int i = 1; i < vector.getValueCount(); i++) { range.setLeftStart(i).setRightStart(i - 1); if (!visitor.rangeEquals(range)) { BitVectorHelper.setBit(runStarts, i); } } } /** * Gets the run lengths, given the start positions. * * @param runStarts the bit set for start positions. * @param runLengths the run length vector to populate. * @param valueCount the number of values in the bit set. */ public static void populateRunLengths(ArrowBuf runStarts, IntVector runLengths, int valueCount) { int curStart = 0; int lengthIndex = 0; for (int i = 1; i < valueCount; i++) { if (BitVectorHelper.get(runStarts, i) != 0) { // we get a new distinct value runLengths.setSafe(lengthIndex++, i - curStart); curStart = i; } } // process the last value runLengths.setSafe(lengthIndex++, valueCount - curStart); runLengths.setValueCount(lengthIndex); } /** * Gets distinct values from the input vector by removing adjacent duplicated values. * * @param indicators the bit set containing the start positions of distinct values. * @param inputVector the input vector. * @param outputVector the output vector. * @param vector type. */ public static void populateDeduplicatedValues( ArrowBuf indicators, V inputVector, V outputVector) { int dstIdx = 0; for (int srcIdx = 0; srcIdx < inputVector.getValueCount(); srcIdx++) { if (BitVectorHelper.get(indicators, srcIdx) != 0) { outputVector.copyFromSafe(srcIdx, dstIdx++, inputVector); } } outputVector.setValueCount(dstIdx); } } VectorRunDeduplicator.java000066400000000000000000000071301515451615200362070ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/deduplicate/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.deduplicate; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.util.DataSizeRoundingUtil; /** * Remove adjacent equal elements from a vector. If the vector is sorted, it removes all duplicated * values in the vector. * * @param vector type. */ public class VectorRunDeduplicator implements AutoCloseable { /** * Bit set for distinct values. If the value at some index is not equal to the previous value, its * bit is set to 1, otherwise its bit is set to 0. */ private ArrowBuf distinctValueBuffer; /** The vector to deduplicate. */ private final V vector; private final BufferAllocator allocator; /** * Constructs a vector run deduplicator for a given vector. * * @param vector the vector to deduplicate. Ownership is NOT taken. * @param allocator the allocator used for allocating buffers for start indices. */ public VectorRunDeduplicator(V vector, BufferAllocator allocator) { this.vector = vector; this.allocator = allocator; } private void createDistinctValueBuffer() { Preconditions.checkArgument(distinctValueBuffer == null); int bufSize = DataSizeRoundingUtil.divideBy8Ceil(vector.getValueCount()); distinctValueBuffer = allocator.buffer(bufSize); DeduplicationUtils.populateRunStartIndicators(vector, distinctValueBuffer); } /** * Gets the number of values which are different from their predecessor. * * @return the run count. */ public int getRunCount() { if (distinctValueBuffer == null) { createDistinctValueBuffer(); } return vector.getValueCount() - BitVectorHelper.getNullCount(distinctValueBuffer, vector.getValueCount()); } /** * Gets the vector with deduplicated adjacent values removed. * * @param outVector the output vector. */ public void populateDeduplicatedValues(V outVector) { if (distinctValueBuffer == null) { createDistinctValueBuffer(); } DeduplicationUtils.populateDeduplicatedValues(distinctValueBuffer, vector, outVector); } /** * Gets the length of each distinct value. * * @param lengthVector the vector for holding length values. */ public void populateRunLengths(IntVector lengthVector) { if (distinctValueBuffer == null) { createDistinctValueBuffer(); } DeduplicationUtils.populateRunLengths( distinctValueBuffer, lengthVector, vector.getValueCount()); } @Override public void close() { if (distinctValueBuffer != null) { distinctValueBuffer.close(); distinctValueBuffer = null; } } } apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/dictionary/000077500000000000000000000000001515451615200310155ustar00rootroot00000000000000DictionaryBuilder.java000066400000000000000000000050541515451615200352210ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.dictionary; import org.apache.arrow.vector.ValueVector; /** * A dictionary builder is intended for the scenario frequently encountered in practice: the * dictionary is not known a priori, so it is generated dynamically. In particular, when a new value * arrives, it is tested to check if it is already in the dictionary. If so, it is simply neglected, * otherwise, it is added to the dictionary. * *

The dictionary builder is intended to build a single dictionary. So it cannot be used for * different dictionaries. * *

Below gives the sample code for using the dictionary builder * * 

{@code
 * DictionaryBuilder dictionaryBuilder = ...
 * ...
 * dictionaryBuild.addValue(newValue);
 * ...
 * }
* *

With the above code, the dictionary vector will be populated, and it can be retrieved by the * {@link DictionaryBuilder#getDictionary()} method. After that, dictionary encoding can proceed * with the populated dictionary.. * * @param the dictionary vector type. */ public interface DictionaryBuilder { /** * Try to add all values from the target vector to the dictionary. * * @param targetVector the target vector containing values to probe. * @return the number of values actually added to the dictionary. */ int addValues(V targetVector); /** * Try to add an element from the target vector to the dictionary. * * @param targetVector the target vector containing new element. * @param targetIndex the index of the new element in the target vector. * @return the index of the new element in the dictionary. */ int addValue(V targetVector, int targetIndex); /** * Gets the dictionary built. * * @return the dictionary. */ V getDictionary(); } DictionaryEncoder.java000066400000000000000000000032461515451615200352130ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.dictionary; import org.apache.arrow.vector.BaseIntVector; import org.apache.arrow.vector.ValueVector; /** * A dictionary encoder translates one vector into another one based on a dictionary vector. * According to Arrow specification, the encoded vector must be an integer based vector, which is * the index of the original vector element in the dictionary. * * @param type of the encoded vector. * @param type of the vector to encode. It is also the type of the dictionary vector. */ public interface DictionaryEncoder { /** * Translates an input vector into an output vector. * * @param input the input vector. * @param output the output vector. Note that it must be in a fresh state. At least, all its * validity bits should be clear. */ void encode(D input, E output); } HashTableBasedDictionaryBuilder.java000066400000000000000000000111741515451615200377340ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.dictionary; import java.util.HashMap; import org.apache.arrow.memory.util.ArrowBufPointer; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.memory.util.hash.SimpleHasher; import org.apache.arrow.vector.ElementAddressableVector; /** * This class builds the dictionary based on a hash table. Each add operation can be finished in * O(1) time, where n is the current dictionary size. * * @param the dictionary vector type. */ public class HashTableBasedDictionaryBuilder implements DictionaryBuilder { /** The dictionary to be built. */ private final V dictionary; /** If null should be encoded. */ private final boolean encodeNull; /** * The hash map for distinct dictionary entries. The key is the pointer to the dictionary element, * whereas the value is the index in the dictionary. */ private HashMap hashMap = new HashMap<>(); /** The hasher used for calculating the hash code. */ private final ArrowBufHasher hasher; /** Next pointer to try to add to the hash table. */ private ArrowBufPointer nextPointer; /** * Constructs a hash table based dictionary builder. * * @param dictionary the dictionary to populate. */ public HashTableBasedDictionaryBuilder(V dictionary) { this(dictionary, false); } /** * Constructs a hash table based dictionary builder. * * @param dictionary the dictionary to populate. * @param encodeNull if null values should be added to the dictionary. */ public HashTableBasedDictionaryBuilder(V dictionary, boolean encodeNull) { this(dictionary, encodeNull, SimpleHasher.INSTANCE); } /** * Constructs a hash table based dictionary builder. * * @param dictionary the dictionary to populate. * @param encodeNull if null values should be added to the dictionary. * @param hasher the hasher used to compute the hash code. */ public HashTableBasedDictionaryBuilder(V dictionary, boolean encodeNull, ArrowBufHasher hasher) { this.dictionary = dictionary; this.encodeNull = encodeNull; this.hasher = hasher; this.nextPointer = new ArrowBufPointer(hasher); } /** * Gets the dictionary built. * * @return the dictionary. */ @Override public V getDictionary() { return dictionary; } /** * Try to add all values from the target vector to the dictionary. * * @param targetVector the target vector containing values to probe. * @return the number of values actually added to the dictionary. */ @Override public int addValues(V targetVector) { int oldDictSize = dictionary.getValueCount(); for (int i = 0; i < targetVector.getValueCount(); i++) { if (!encodeNull && targetVector.isNull(i)) { continue; } addValue(targetVector, i); } return dictionary.getValueCount() - oldDictSize; } /** * Try to add an element from the target vector to the dictionary. * * @param targetVector the target vector containing new element. * @param targetIndex the index of the new element in the target vector. * @return the index of the new element in the dictionary. */ @Override public int addValue(V targetVector, int targetIndex) { targetVector.getDataPointer(targetIndex, nextPointer); Integer index = hashMap.get(nextPointer); if (index == null) { // a new dictionary element is found // insert it to the dictionary int dictSize = dictionary.getValueCount(); dictionary.copyFromSafe(targetIndex, dictSize, targetVector); dictionary.setValueCount(dictSize + 1); dictionary.getDataPointer(dictSize, nextPointer); // insert it to the hash map hashMap.put(nextPointer, dictSize); nextPointer = new ArrowBufPointer(hasher); return dictSize; } return index; } } HashTableDictionaryEncoder.java000066400000000000000000000116711515451615200367700ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.dictionary; import java.util.HashMap; import org.apache.arrow.memory.util.ArrowBufPointer; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.memory.util.hash.SimpleHasher; import org.apache.arrow.vector.BaseIntVector; import org.apache.arrow.vector.ElementAddressableVector; /** * Dictionary encoder based on hash table. * * @param encoded vector type. * @param decoded vector type, which is also the dictionary type. */ public class HashTableDictionaryEncoder implements DictionaryEncoder { /** The dictionary for encoding/decoding. It must be sorted. */ private final D dictionary; /** The hasher used to compute the hash code. */ private final ArrowBufHasher hasher; /** A flag indicating if null should be encoded. */ private final boolean encodeNull; /** * The hash map for distinct dictionary entries. The key is the pointer to the dictionary element, * whereas the value is the index in the dictionary. */ private HashMap hashMap = new HashMap<>(); /** The pointer used to probe each element to encode. */ private ArrowBufPointer reusablePointer; /** * Constructs a dictionary encoder. * * @param dictionary the dictionary. */ public HashTableDictionaryEncoder(D dictionary) { this(dictionary, false); } /** * Constructs a dictionary encoder. * * @param dictionary the dictionary. * @param encodeNull a flag indicating if null should be encoded. It determines the behaviors for * processing null values in the input during encoding/decoding. *

  • For encoding, when a null is encountered in the input, 1) If the flag is set to true, * the encoder searches for the value in the dictionary, and outputs the index in the * dictionary. 2) If the flag is set to false, the encoder simply produces a null in the * output. *
  • For decoding, when a null is encountered in the input, 1) If the flag is set to true, * the decoder should never expect a null in the input. 2) If set to false, the decoder * simply produces a null in the output. */ public HashTableDictionaryEncoder(D dictionary, boolean encodeNull) { this(dictionary, encodeNull, SimpleHasher.INSTANCE); } /** * Constructs a dictionary encoder. * * @param dictionary the dictionary. * @param encodeNull a flag indicating if null should be encoded. It determines the behaviors for * processing null values in the input during encoding. When a null is encountered in the * input, 1) If the flag is set to true, the encoder searches for the value in the dictionary, * and outputs the index in the dictionary. 2) If the flag is set to false, the encoder simply * produces a null in the output. * @param hasher the hasher used to calculate the hash code. */ public HashTableDictionaryEncoder(D dictionary, boolean encodeNull, ArrowBufHasher hasher) { this.dictionary = dictionary; this.hasher = hasher; this.encodeNull = encodeNull; reusablePointer = new ArrowBufPointer(hasher); buildHashMap(); } private void buildHashMap() { for (int i = 0; i < dictionary.getValueCount(); i++) { ArrowBufPointer pointer = new ArrowBufPointer(hasher); dictionary.getDataPointer(i, pointer); hashMap.put(pointer, i); } } /** * Encodes an input vector by a hash table. So the algorithm takes O(n) time, where n is the * length of the input vector. * * @param input the input vector. * @param output the output vector. */ @Override public void encode(D input, E output) { for (int i = 0; i < input.getValueCount(); i++) { if (!encodeNull && input.isNull(i)) { continue; } input.getDataPointer(i, reusablePointer); Integer index = hashMap.get(reusablePointer); if (index == null) { throw new IllegalArgumentException("The data element is not found in the dictionary"); } output.setWithPossibleTruncate(i, index); } output.setValueCount(input.getValueCount()); } } LinearDictionaryEncoder.java000066400000000000000000000102611515451615200363410ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.dictionary; import org.apache.arrow.vector.BaseIntVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.compare.Range; import org.apache.arrow.vector.compare.RangeEqualsVisitor; /** * Dictionary encoder based on linear search. * * @param encoded vector type. * @param decoded vector type, which is also the dictionary type. */ public class LinearDictionaryEncoder implements DictionaryEncoder { /** The dictionary for encoding. */ private final D dictionary; /** A flag indicating if null should be encoded. */ private final boolean encodeNull; private RangeEqualsVisitor equalizer; private Range range; /** * Constructs a dictionary encoder, with the encode null flag set to false. * * @param dictionary the dictionary. Its entries should be sorted in the non-increasing order of * their frequency. Otherwise, the encoder still produces correct results, but at the expense * of performance overhead. */ public LinearDictionaryEncoder(D dictionary) { this(dictionary, false); } /** * Constructs a dictionary encoder. * * @param dictionary the dictionary. Its entries should be sorted in the non-increasing order of * their frequency. Otherwise, the encoder still produces correct results, but at the expense * of performance overhead. * @param encodeNull a flag indicating if null should be encoded. It determines the behaviors for * processing null values in the input during encoding. When a null is encountered in the * input, 1) If the flag is set to true, the encoder searches for the value in the dictionary, * and outputs the index in the dictionary. 2) If the flag is set to false, the encoder simply * produces a null in the output. */ public LinearDictionaryEncoder(D dictionary, boolean encodeNull) { this.dictionary = dictionary; this.encodeNull = encodeNull; // temporarily set left and right vectors to dictionary equalizer = new RangeEqualsVisitor(dictionary, dictionary, null); range = new Range(0, 0, 1); } /** * Encodes an input vector by linear search. When the dictionary is sorted in the non-increasing * order of the entry frequency, it will have constant time complexity, with no extra memory * requirement. * * @param input the input vector. * @param output the output vector. Note that it must be in a fresh state. At least, all its * validity bits should be clear. */ @Override public void encode(D input, E output) { for (int i = 0; i < input.getValueCount(); i++) { if (!encodeNull && input.isNull(i)) { // for this case, we should simply output a null in the output. // by assuming the output vector is fresh, we do nothing here. continue; } int index = linearSearch(input, i); if (index == -1) { throw new IllegalArgumentException("The data element is not found in the dictionary: " + i); } output.setWithPossibleTruncate(i, index); } output.setValueCount(input.getValueCount()); } private int linearSearch(D input, int index) { range.setLeftStart(index); for (int i = 0; i < dictionary.getValueCount(); i++) { range.setRightStart(i); if (input.accept(equalizer, range)) { return i; } } return -1; } } SearchDictionaryEncoder.java000066400000000000000000000073121515451615200363370ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.dictionary; import org.apache.arrow.algorithm.search.VectorSearcher; import org.apache.arrow.algorithm.sort.VectorValueComparator; import org.apache.arrow.vector.BaseIntVector; import org.apache.arrow.vector.ValueVector; /** * Dictionary encoder based on searching. * * @param encoded vector type. * @param decoded vector type, which is also the dictionary type. */ public class SearchDictionaryEncoder implements DictionaryEncoder { /** The dictionary for encoding/decoding. It must be sorted. */ private final D dictionary; /** The criteria by which the dictionary is sorted. */ private final VectorValueComparator comparator; /** A flag indicating if null should be encoded. */ private final boolean encodeNull; /** * Constructs a dictionary encoder. * * @param dictionary the dictionary. It must be in sorted order. * @param comparator the criteria for sorting. */ public SearchDictionaryEncoder(D dictionary, VectorValueComparator comparator) { this(dictionary, comparator, false); } /** * Constructs a dictionary encoder. * * @param dictionary the dictionary. It must be in sorted order. * @param comparator the criteria for sorting. * @param encodeNull a flag indicating if null should be encoded. It determines the behaviors for * processing null values in the input during encoding. When a null is encountered in the * input, 1) If the flag is set to true, the encoder searches for the value in the dictionary, * and outputs the index in the dictionary. 2) If the flag is set to false, the encoder simply * produces a null in the output. */ public SearchDictionaryEncoder( D dictionary, VectorValueComparator comparator, boolean encodeNull) { this.dictionary = dictionary; this.comparator = comparator; this.encodeNull = encodeNull; } /** * Encodes an input vector by binary search. So the algorithm takes O(n * log(m)) time, where n is * the length of the input vector, and m is the length of the dictionary. * * @param input the input vector. * @param output the output vector. Note that it must be in a fresh state. At least, all its * validity bits should be clear. */ @Override public void encode(D input, E output) { for (int i = 0; i < input.getValueCount(); i++) { if (!encodeNull && input.isNull(i)) { // for this case, we should simply output a null in the output. // by assuming the output vector is fresh, we do nothing here. continue; } int index = VectorSearcher.binarySearch(dictionary, comparator, input, i); if (index == -1) { throw new IllegalArgumentException("The data element is not found in the dictionary: " + i); } output.setWithPossibleTruncate(i, index); } output.setValueCount(input.getValueCount()); } } SearchTreeBasedDictionaryBuilder.java000066400000000000000000000114601515451615200401240ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.dictionary; import java.util.TreeSet; import org.apache.arrow.algorithm.sort.VectorValueComparator; import org.apache.arrow.vector.ValueVector; /** * This class builds the dictionary based on a binary search tree. Each add operation can be * finished in O(log(n)) time, where n is the current dictionary size. * * @param the dictionary vector type. */ public class SearchTreeBasedDictionaryBuilder implements DictionaryBuilder { /** The dictionary to be built. */ private final V dictionary; /** The criteria for sorting in the search tree. */ protected final VectorValueComparator comparator; /** If null should be encoded. */ private final boolean encodeNull; /** The search tree for storing the value index. */ private TreeSet searchTree; /** * Construct a search tree-based dictionary builder. * * @param dictionary the dictionary vector. * @param comparator the criteria for value equality. */ public SearchTreeBasedDictionaryBuilder(V dictionary, VectorValueComparator comparator) { this(dictionary, comparator, false); } /** * Construct a search tree-based dictionary builder. * * @param dictionary the dictionary vector. * @param comparator the criteria for value equality. * @param encodeNull if null values should be added to the dictionary. */ public SearchTreeBasedDictionaryBuilder( V dictionary, VectorValueComparator comparator, boolean encodeNull) { this.dictionary = dictionary; this.comparator = comparator; this.encodeNull = encodeNull; this.comparator.attachVector(dictionary); searchTree = new TreeSet<>((index1, index2) -> comparator.compare(index1, index2)); } /** * Gets the dictionary built. Please note that the dictionary is not in sorted order. Instead, its * order is determined by the order of element insertion. To get the dictionary in sorted order, * please use {@link SearchTreeBasedDictionaryBuilder#populateSortedDictionary(ValueVector)}. * * @return the dictionary. */ @Override public V getDictionary() { return dictionary; } /** * Try to add all values from the target vector to the dictionary. * * @param targetVector the target vector containing values to probe. * @return the number of values actually added to the dictionary. */ @Override public int addValues(V targetVector) { int oldDictSize = dictionary.getValueCount(); for (int i = 0; i < targetVector.getValueCount(); i++) { if (!encodeNull && targetVector.isNull(i)) { continue; } addValue(targetVector, i); } return dictionary.getValueCount() - oldDictSize; } /** * Try to add an element from the target vector to the dictionary. * * @param targetVector the target vector containing new element. * @param targetIndex the index of the new element in the target vector. * @return the index of the new element in the dictionary. */ @Override public int addValue(V targetVector, int targetIndex) { // first copy the value to the end of the dictionary int dictSize = dictionary.getValueCount(); dictionary.copyFromSafe(targetIndex, dictSize, targetVector); // try to add the value to the dictionary, // if an equal element does not exist. // this operation can be done in O(log(n)) time. if (searchTree.add(dictSize)) { // the element is successfully added dictionary.setValueCount(dictSize + 1); return dictSize; } else { // the element is already in the dictionary // find its index in O(log(n)) time. return searchTree.ceiling(dictSize); } } /** * Gets the sorted dictionary. Note that given the binary search tree, the sort can finish in * O(n). */ public void populateSortedDictionary(V sortedDictionary) { int idx = 0; for (Integer dictIdx : searchTree) { sortedDictionary.copyFromSafe(dictIdx, idx++, dictionary); } sortedDictionary.setValueCount(dictionary.getValueCount()); } } apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/misc/000077500000000000000000000000001515451615200276035ustar00rootroot00000000000000PartialSumUtils.java000066400000000000000000000111751515451615200334760ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/misc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.misc; import org.apache.arrow.vector.BaseIntVector; /** Partial sum related utilities. */ public class PartialSumUtils { /** * Converts an input vector to a partial sum vector. This is an inverse operation of {@link * PartialSumUtils#toDeltaVector(BaseIntVector, BaseIntVector)}. Suppose we have input vector a * and output vector b. Then we have b(0) = sumBase; b(i + 1) = b(i) + a(i) (i = 0, 1, 2, ...). * * @param deltaVector the input vector. * @param partialSumVector the output vector. * @param sumBase the base of the partial sums. */ public static void toPartialSumVector( BaseIntVector deltaVector, BaseIntVector partialSumVector, long sumBase) { long sum = sumBase; partialSumVector.setWithPossibleTruncate(0, sumBase); for (int i = 0; i < deltaVector.getValueCount(); i++) { sum += deltaVector.getValueAsLong(i); partialSumVector.setWithPossibleTruncate(i + 1, sum); } partialSumVector.setValueCount(deltaVector.getValueCount() + 1); } /** * Converts an input vector to the delta vector. This is an inverse operation of {@link * PartialSumUtils#toPartialSumVector(BaseIntVector, BaseIntVector, long)}. Suppose we have input * vector a and output vector b. Then we have b(i) = a(i + 1) - a(i) (i = 0, 1, 2, ...). * * @param partialSumVector the input vector. * @param deltaVector the output vector. */ public static void toDeltaVector(BaseIntVector partialSumVector, BaseIntVector deltaVector) { for (int i = 0; i < partialSumVector.getValueCount() - 1; i++) { long delta = partialSumVector.getValueAsLong(i + 1) - partialSumVector.getValueAsLong(i); deltaVector.setWithPossibleTruncate(i, delta); } deltaVector.setValueCount(partialSumVector.getValueCount() - 1); } /** * Given a value and a partial sum vector, finds its position in the partial sum vector. In * particular, given an integer value a and partial sum vector v, we try to find a position i, so * that v(i) <= a < v(i + 1). The algorithm is based on binary search, so it takes O(log(n)) time, * where n is the length of the partial sum vector. * * @param partialSumVector the input partial sum vector. * @param value the value to search. * @return the position in the partial sum vector, if any, or -1, if none is found. */ public static int findPositionInPartialSumVector(BaseIntVector partialSumVector, long value) { if (value < partialSumVector.getValueAsLong(0) || value >= partialSumVector.getValueAsLong(partialSumVector.getValueCount() - 1)) { return -1; } int low = 0; int high = partialSumVector.getValueCount() - 1; while (low <= high) { int mid = low + (high - low) / 2; long midValue = partialSumVector.getValueAsLong(mid); if (midValue <= value) { if (mid == partialSumVector.getValueCount() - 1) { // the mid is the last element, we have found it return mid; } long nextMidValue = partialSumVector.getValueAsLong(mid + 1); if (value < nextMidValue) { // midValue <= value < nextMidValue // this is exactly what we want. return mid; } else { // value >= nextMidValue // continue to search from the next value on the right low = mid + 1; } } else { // midValue > value long prevMidValue = partialSumVector.getValueAsLong(mid - 1); if (prevMidValue <= value) { // prevMidValue <= value < midValue // this is exactly what we want return mid - 1; } else { // prevMidValue > value // continue to search from the previous value on the left high = mid - 1; } } } throw new IllegalStateException("Should never get here"); } private PartialSumUtils() {} } apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/rank/000077500000000000000000000000001515451615200276035ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/rank/VectorRank.java000066400000000000000000000057021515451615200325300ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.rank; import java.util.stream.IntStream; import org.apache.arrow.algorithm.sort.IndexSorter; import org.apache.arrow.algorithm.sort.VectorValueComparator; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.ValueVector; /** * Utility for calculating ranks of vector elements. * * @param the vector type */ public class VectorRank { private VectorValueComparator comparator; /** Vector indices. */ private IntVector indices; private final BufferAllocator allocator; /** * Constructs a vector rank utility. * * @param allocator the allocator to use. */ public VectorRank(BufferAllocator allocator) { this.allocator = allocator; } /** * Given a rank r, gets the index of the element that is the rth smallest in the vector. The * operation is performed without changing the vector, and takes O(n) time, where n is the length * of the vector. * * @param vector the vector from which to get the element index. * @param comparator the criteria for vector element comparison. * @param rank the rank to determine. * @return the element index with the given rank. */ public int indexAtRank(V vector, VectorValueComparator comparator, int rank) { Preconditions.checkArgument(rank >= 0 && rank < vector.getValueCount()); try { indices = new IntVector("index vector", allocator); indices.allocateNew(vector.getValueCount()); IntStream.range(0, vector.getValueCount()).forEach(i -> indices.set(i, i)); comparator.attachVector(vector); this.comparator = comparator; int pos = getRank(0, vector.getValueCount() - 1, rank); return indices.get(pos); } finally { indices.close(); } } private int getRank(int low, int high, int rank) { int mid = IndexSorter.partition(low, high, indices, comparator); if (mid < rank) { return getRank(mid + 1, high, rank); } else if (mid > rank) { return getRank(low, mid - 1, rank); } else { // mid == rank return mid; } } } apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/search/000077500000000000000000000000001515451615200301155ustar00rootroot00000000000000ParallelSearcher.java000066400000000000000000000167651515451615200341310ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/search/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.search; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; import org.apache.arrow.algorithm.sort.VectorValueComparator; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.compare.Range; import org.apache.arrow.vector.compare.RangeEqualsVisitor; /** * Search for a value in the vector by multiple threads. This is often used in scenarios where the * vector is large or low response time is required. * * @param the vector type. */ public class ParallelSearcher { /** The target vector to search. */ private final V vector; /** The thread pool. */ private final ExecutorService threadPool; /** The number of threads to use. */ private final int numThreads; /** The position of the key in the target vector, if any. */ private volatile int keyPosition = -1; /** * Constructs a parallel searcher. * * @param vector the vector to search. * @param threadPool the thread pool to use. * @param numThreads the number of threads to use. */ public ParallelSearcher(V vector, ExecutorService threadPool, int numThreads) { this.vector = vector; this.threadPool = threadPool; this.numThreads = numThreads; } private CompletableFuture[] initSearch() { keyPosition = -1; final CompletableFuture[] futures = new CompletableFuture[numThreads]; for (int i = 0; i < futures.length; i++) { futures[i] = new CompletableFuture<>(); } return futures; } /** * Search for the key in the target vector. The element-wise comparison is based on {@link * RangeEqualsVisitor}, so there are two possible results for each element-wise comparison: equal * and un-equal. * * @param keyVector the vector containing the search key. * @param keyIndex the index of the search key in the key vector. * @return the position of a matched value in the target vector, or -1 if none is found. Please * note that if there are multiple matches of the key in the target vector, this method makes * no guarantees about which instance is returned. For an alternative search implementation * that always finds the first match of the key, see {@link * VectorSearcher#linearSearch(ValueVector, VectorValueComparator, ValueVector, int)}. * @throws ExecutionException if an exception occurs in a thread. * @throws InterruptedException if a thread is interrupted. */ public int search(V keyVector, int keyIndex) throws ExecutionException, InterruptedException { final CompletableFuture[] futures = initSearch(); final int valueCount = vector.getValueCount(); for (int i = 0; i < numThreads; i++) { final int tid = i; Future unused = threadPool.submit( () -> { // convert to long to avoid overflow int start = (int) (((long) valueCount) * tid / numThreads); int end = (int) ((long) valueCount) * (tid + 1) / numThreads; if (start >= end) { // no data assigned to this task. futures[tid].complete(false); return; } RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector, keyVector, null); Range range = new Range(0, 0, 1); for (int pos = start; pos < end; pos++) { if (keyPosition != -1) { // the key has been found by another task futures[tid].complete(false); return; } range.setLeftStart(pos).setRightStart(keyIndex); if (visitor.rangeEquals(range)) { keyPosition = pos; futures[tid].complete(true); return; } } // no match value is found. futures[tid].complete(false); }); } CompletableFuture.allOf(futures).get(); return keyPosition; } /** * Search for the key in the target vector. The element-wise comparison is based on {@link * VectorValueComparator}, so there are three possible results for each element-wise comparison: * less than, equal to and greater than. * * @param keyVector the vector containing the search key. * @param keyIndex the index of the search key in the key vector. * @param comparator the comparator for comparing the key against vector elements. * @return the position of a matched value in the target vector, or -1 if none is found. Please * note that if there are multiple matches of the key in the target vector, this method makes * no guarantees about which instance is returned. For an alternative search implementation * that always finds the first match of the key, see {@link * VectorSearcher#linearSearch(ValueVector, VectorValueComparator, ValueVector, int)}. * @throws ExecutionException if an exception occurs in a thread. * @throws InterruptedException if a thread is interrupted. */ public int search(V keyVector, int keyIndex, VectorValueComparator comparator) throws ExecutionException, InterruptedException { final CompletableFuture[] futures = initSearch(); final int valueCount = vector.getValueCount(); for (int i = 0; i < numThreads; i++) { final int tid = i; Future unused = threadPool.submit( () -> { // convert to long to avoid overflow int start = (int) (((long) valueCount) * tid / numThreads); int end = (int) ((long) valueCount) * (tid + 1) / numThreads; if (start >= end) { // no data assigned to this task. futures[tid].complete(false); return; } VectorValueComparator localComparator = comparator.createNew(); localComparator.attachVectors(vector, keyVector); for (int pos = start; pos < end; pos++) { if (keyPosition != -1) { // the key has been found by another task futures[tid].complete(false); return; } if (localComparator.compare(pos, keyIndex) == 0) { keyPosition = pos; futures[tid].complete(true); return; } } // no match value is found. futures[tid].complete(false); }); } CompletableFuture.allOf(futures).get(); return keyPosition; } } VectorRangeSearcher.java000066400000000000000000000072271515451615200346050ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/search/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.search; import org.apache.arrow.algorithm.sort.VectorValueComparator; import org.apache.arrow.vector.ValueVector; /** Search for the range of a particular element in the target vector. */ public class VectorRangeSearcher { /** Result returned when a search fails. */ public static final int SEARCH_FAIL_RESULT = -1; /** * Search for the first occurrence of an element. The search is based on the binary search * algorithm. So the target vector must be sorted. * * @param targetVector the vector from which to perform the search. * @param comparator the criterion for the comparison. * @param keyVector the vector containing the element to search. * @param keyIndex the index of the search key in the key vector. * @param the vector type. * @return the index of the first matched element if any, and -1 otherwise. */ public static int getFirstMatch( V targetVector, VectorValueComparator comparator, V keyVector, int keyIndex) { comparator.attachVectors(keyVector, targetVector); int ret = SEARCH_FAIL_RESULT; int low = 0; int high = targetVector.getValueCount() - 1; while (low <= high) { int mid = low + (high - low) / 2; int result = comparator.compare(keyIndex, mid); if (result < 0) { // the key is smaller high = mid - 1; } else if (result > 0) { // the key is larger low = mid + 1; } else { // an equal element is found // continue to go left-ward ret = mid; high = mid - 1; } } return ret; } /** * Search for the last occurrence of an element. The search is based on the binary search * algorithm. So the target vector must be sorted. * * @param targetVector the vector from which to perform the search. * @param comparator the criterion for the comparison. * @param keyVector the vector containing the element to search. * @param keyIndex the index of the search key in the key vector. * @param the vector type. * @return the index of the last matched element if any, and -1 otherwise. */ public static int getLastMatch( V targetVector, VectorValueComparator comparator, V keyVector, int keyIndex) { comparator.attachVectors(keyVector, targetVector); int ret = SEARCH_FAIL_RESULT; int low = 0; int high = targetVector.getValueCount() - 1; while (low <= high) { int mid = low + (high - low) / 2; int result = comparator.compare(keyIndex, mid); if (result < 0) { // the key is smaller high = mid - 1; } else if (result > 0) { // the key is larger low = mid + 1; } else { // an equal element is found, // continue to go right-ward ret = mid; low = mid + 1; } } return ret; } } VectorSearcher.java000066400000000000000000000061331515451615200336230ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/search/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.search; import org.apache.arrow.algorithm.sort.VectorValueComparator; import org.apache.arrow.vector.ValueVector; /** Search for a particular element in the vector. */ public final class VectorSearcher { /** Result returned when a search fails. */ public static final int SEARCH_FAIL_RESULT = -1; /** * Search for a particular element from the key vector in the target vector by binary search. The * target vector must be sorted. * * @param targetVector the vector from which to perform the sort. * @param comparator the criterion for the sort. * @param keyVector the vector containing the element to search. * @param keyIndex the index of the search key in the key vector. * @param the vector type. * @return the index of a matched element if any, and -1 otherwise. */ public static int binarySearch( V targetVector, VectorValueComparator comparator, V keyVector, int keyIndex) { comparator.attachVectors(keyVector, targetVector); // perform binary search int low = 0; int high = targetVector.getValueCount() - 1; while (low <= high) { int mid = low + (high - low) / 2; int cmp = comparator.compare(keyIndex, mid); if (cmp < 0) { high = mid - 1; } else if (cmp > 0) { low = mid + 1; } else { return mid; } } return SEARCH_FAIL_RESULT; } /** * Search for a particular element from the key vector in the target vector by traversing the * vector in sequence. * * @param targetVector the vector from which to perform the search. * @param comparator the criterion for element equality. * @param keyVector the vector containing the element to search. * @param keyIndex the index of the search key in the key vector. * @param the vector type. * @return the index of a matched element if any, and -1 otherwise. */ public static int linearSearch( V targetVector, VectorValueComparator comparator, V keyVector, int keyIndex) { comparator.attachVectors(keyVector, targetVector); for (int i = 0; i < targetVector.getValueCount(); i++) { if (comparator.compare(keyIndex, i) == 0) { return i; } } return SEARCH_FAIL_RESULT; } private VectorSearcher() {} } apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/sort/000077500000000000000000000000001515451615200276375ustar00rootroot00000000000000CompositeVectorComparator.java000066400000000000000000000047231515451615200356060ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import org.apache.arrow.vector.ValueVector; /** * A composite vector comparator compares a number of vectors by a number of inner comparators. * *

    It works by first using the first comparator, if a non-zero value is returned, it simply * returns it. Otherwise, it uses the second comparator, and so on, until a non-zero value is * produced, or all inner comparators have been used. */ public class CompositeVectorComparator extends VectorValueComparator { private final VectorValueComparator[] innerComparators; public CompositeVectorComparator(VectorValueComparator[] innerComparators) { this.innerComparators = innerComparators; } @Override public int compareNotNull(int index1, int index2) { // short-cut for scenarios when the caller can be sure that the vectors are non-nullable. for (int i = 0; i < innerComparators.length; i++) { int result = innerComparators[i].compareNotNull(index1, index2); if (result != 0) { return result; } } return 0; } @Override public int compare(int index1, int index2) { for (int i = 0; i < innerComparators.length; i++) { int result = innerComparators[i].compare(index1, index2); if (result != 0) { return result; } } return 0; } @Override public VectorValueComparator createNew() { VectorValueComparator[] newInnerComparators = new VectorValueComparator[innerComparators.length]; for (int i = 0; i < innerComparators.length; i++) { newInnerComparators[i] = innerComparators[i].createNew(); } return new CompositeVectorComparator(newInnerComparators); } } DefaultVectorComparators.java000066400000000000000000000654371515451615200354240ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import static org.apache.arrow.vector.complex.BaseRepeatedValueVector.OFFSET_WIDTH; import java.math.BigDecimal; import java.time.Duration; import org.apache.arrow.memory.util.ArrowBufPointer; import org.apache.arrow.memory.util.ByteFunctionHelpers; import org.apache.arrow.vector.BaseFixedWidthVector; 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.DurationVector; 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.IntervalDayVector; import org.apache.arrow.vector.IntervalMonthDayNanoVector; 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.TimeStampVector; 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.ValueVector; import org.apache.arrow.vector.VariableWidthVector; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.RepeatedValueVector; import org.apache.arrow.vector.holders.NullableFixedSizeBinaryHolder; /** Default comparator implementations for different types of vectors. */ public class DefaultVectorComparators { /** * Create the default comparator for the vector. * * @param vector the vector. * @param the vector type. * @return the default comparator. */ public static VectorValueComparator createDefaultComparator(T vector) { if (vector instanceof BaseFixedWidthVector) { if (vector instanceof TinyIntVector) { return (VectorValueComparator) new ByteComparator(); } else if (vector instanceof SmallIntVector) { return (VectorValueComparator) new ShortComparator(); } else if (vector instanceof IntVector) { return (VectorValueComparator) new IntComparator(); } else if (vector instanceof BigIntVector) { return (VectorValueComparator) new LongComparator(); } else if (vector instanceof Float4Vector) { return (VectorValueComparator) new Float4Comparator(); } else if (vector instanceof Float8Vector) { return (VectorValueComparator) new Float8Comparator(); } else if (vector instanceof UInt1Vector) { return (VectorValueComparator) new UInt1Comparator(); } else if (vector instanceof UInt2Vector) { return (VectorValueComparator) new UInt2Comparator(); } else if (vector instanceof UInt4Vector) { return (VectorValueComparator) new UInt4Comparator(); } else if (vector instanceof UInt8Vector) { return (VectorValueComparator) new UInt8Comparator(); } else if (vector instanceof BitVector) { return (VectorValueComparator) new BitComparator(); } else if (vector instanceof DateDayVector) { return (VectorValueComparator) new DateDayComparator(); } else if (vector instanceof DateMilliVector) { return (VectorValueComparator) new DateMilliComparator(); } else if (vector instanceof Decimal256Vector) { return (VectorValueComparator) new Decimal256Comparator(); } else if (vector instanceof DecimalVector) { return (VectorValueComparator) new DecimalComparator(); } else if (vector instanceof DurationVector) { return (VectorValueComparator) new DurationComparator(); } else if (vector instanceof IntervalDayVector) { return (VectorValueComparator) new IntervalDayComparator(); } else if (vector instanceof IntervalMonthDayNanoVector) { throw new IllegalArgumentException( "No default comparator for " + vector.getClass().getCanonicalName()); } else if (vector instanceof TimeMicroVector) { return (VectorValueComparator) new TimeMicroComparator(); } else if (vector instanceof TimeMilliVector) { return (VectorValueComparator) new TimeMilliComparator(); } else if (vector instanceof TimeNanoVector) { return (VectorValueComparator) new TimeNanoComparator(); } else if (vector instanceof TimeSecVector) { return (VectorValueComparator) new TimeSecComparator(); } else if (vector instanceof TimeStampVector) { return (VectorValueComparator) new TimeStampComparator(); } else if (vector instanceof FixedSizeBinaryVector) { return (VectorValueComparator) new FixedSizeBinaryComparator(); } } else if (vector instanceof VariableWidthVector) { return (VectorValueComparator) new VariableWidthComparator(); } else if (vector instanceof RepeatedValueVector) { VectorValueComparator innerComparator = createDefaultComparator(((RepeatedValueVector) vector).getDataVector()); return new RepeatedValueComparator(innerComparator); } else if (vector instanceof FixedSizeListVector) { VectorValueComparator innerComparator = createDefaultComparator(((FixedSizeListVector) vector).getDataVector()); return new FixedSizeListComparator(innerComparator); } else if (vector instanceof NullVector) { return (VectorValueComparator) new NullComparator(); } throw new IllegalArgumentException( "No default comparator for " + vector.getClass().getCanonicalName()); } /** Default comparator for bytes. The comparison is based on values, with null comes first. */ public static class ByteComparator extends VectorValueComparator { public ByteComparator() { super(Byte.SIZE / 8); } @Override public int compareNotNull(int index1, int index2) { byte value1 = vector1.get(index1); byte value2 = vector2.get(index2); return value1 - value2; } @Override public VectorValueComparator createNew() { return new ByteComparator(); } } /** * Default comparator for short integers. The comparison is based on values, with null comes * first. */ public static class ShortComparator extends VectorValueComparator { public ShortComparator() { super(Short.SIZE / 8); } @Override public int compareNotNull(int index1, int index2) { short value1 = vector1.get(index1); short value2 = vector2.get(index2); return value1 - value2; } @Override public VectorValueComparator createNew() { return new ShortComparator(); } } /** * Default comparator for 32-bit integers. The comparison is based on int values, with null comes * first. */ public static class IntComparator extends VectorValueComparator { public IntComparator() { super(Integer.SIZE / 8); } @Override public int compareNotNull(int index1, int index2) { int value1 = vector1.get(index1); int value2 = vector2.get(index2); return Integer.compare(value1, value2); } @Override public VectorValueComparator createNew() { return new IntComparator(); } } /** * Default comparator for long integers. The comparison is based on values, with null comes first. */ public static class LongComparator extends VectorValueComparator { public LongComparator() { super(Long.SIZE / 8); } @Override public int compareNotNull(int index1, int index2) { long value1 = vector1.get(index1); long value2 = vector2.get(index2); return Long.compare(value1, value2); } @Override public VectorValueComparator createNew() { return new LongComparator(); } } /** * Default comparator for unsigned bytes. The comparison is based on values, with null comes * first. */ public static class UInt1Comparator extends VectorValueComparator { public UInt1Comparator() { super(1); } @Override public int compareNotNull(int index1, int index2) { byte value1 = vector1.get(index1); byte value2 = vector2.get(index2); return (value1 & 0xff) - (value2 & 0xff); } @Override public VectorValueComparator createNew() { return new UInt1Comparator(); } } /** * Default comparator for unsigned short integer. The comparison is based on values, with null * comes first. */ public static class UInt2Comparator extends VectorValueComparator { public UInt2Comparator() { super(2); } @Override public int compareNotNull(int index1, int index2) { char value1 = vector1.get(index1); char value2 = vector2.get(index2); // please note that we should not use the built-in // Character#compare method here, as that method // essentially compares char values as signed integers. return (value1 & 0xffff) - (value2 & 0xffff); } @Override public VectorValueComparator createNew() { return new UInt2Comparator(); } } /** * Default comparator for unsigned integer. The comparison is based on values, with null comes * first. */ public static class UInt4Comparator extends VectorValueComparator { public UInt4Comparator() { super(4); } @Override public int compareNotNull(int index1, int index2) { int value1 = vector1.get(index1); int value2 = vector2.get(index2); return ByteFunctionHelpers.unsignedIntCompare(value1, value2); } @Override public VectorValueComparator createNew() { return new UInt4Comparator(); } } /** * Default comparator for unsigned long integer. The comparison is based on values, with null * comes first. */ public static class UInt8Comparator extends VectorValueComparator { public UInt8Comparator() { super(8); } @Override public int compareNotNull(int index1, int index2) { long value1 = vector1.get(index1); long value2 = vector2.get(index2); return ByteFunctionHelpers.unsignedLongCompare(value1, value2); } @Override public VectorValueComparator createNew() { return new UInt8Comparator(); } } /** * Default comparator for float type. The comparison is based on values, with null comes first. */ public static class Float4Comparator extends VectorValueComparator { public Float4Comparator() { super(Float.SIZE / 8); } @Override public int compareNotNull(int index1, int index2) { float value1 = vector1.get(index1); float value2 = vector2.get(index2); boolean isNan1 = Float.isNaN(value1); boolean isNan2 = Float.isNaN(value2); if (isNan1 || isNan2) { if (isNan1 && isNan2) { return 0; } else if (isNan1) { // nan is greater than any normal value return 1; } else { return -1; } } return (int) Math.signum(value1 - value2); } @Override public VectorValueComparator createNew() { return new Float4Comparator(); } } /** * Default comparator for double type. The comparison is based on values, with null comes first. */ public static class Float8Comparator extends VectorValueComparator { public Float8Comparator() { super(Double.SIZE / 8); } @Override public int compareNotNull(int index1, int index2) { double value1 = vector1.get(index1); double value2 = vector2.get(index2); boolean isNan1 = Double.isNaN(value1); boolean isNan2 = Double.isNaN(value2); if (isNan1 || isNan2) { if (isNan1 && isNan2) { return 0; } else if (isNan1) { // nan is greater than any normal value return 1; } else { return -1; } } return (int) Math.signum(value1 - value2); } @Override public VectorValueComparator createNew() { return new Float8Comparator(); } } /** Default comparator for bit type. The comparison is based on values, with null comes first. */ public static class BitComparator extends VectorValueComparator { public BitComparator() { super(-1); } @Override public int compareNotNull(int index1, int index2) { boolean value1 = vector1.get(index1) != 0; boolean value2 = vector2.get(index2) != 0; return Boolean.compare(value1, value2); } @Override public VectorValueComparator createNew() { return new BitComparator(); } } /** * Default comparator for DateDay type. The comparison is based on values, with null comes first. */ public static class DateDayComparator extends VectorValueComparator { public DateDayComparator() { super(DateDayVector.TYPE_WIDTH); } @Override public int compareNotNull(int index1, int index2) { int value1 = vector1.get(index1); int value2 = vector2.get(index2); return Integer.compare(value1, value2); } @Override public VectorValueComparator createNew() { return new DateDayComparator(); } } /** * Default comparator for DateMilli type. The comparison is based on values, with null comes * first. */ public static class DateMilliComparator extends VectorValueComparator { public DateMilliComparator() { super(DateMilliVector.TYPE_WIDTH); } @Override public int compareNotNull(int index1, int index2) { long value1 = vector1.get(index1); long value2 = vector2.get(index2); return Long.compare(value1, value2); } @Override public VectorValueComparator createNew() { return new DateMilliComparator(); } } /** * Default comparator for Decimal256 type. The comparison is based on values, with null comes * first. */ public static class Decimal256Comparator extends VectorValueComparator { public Decimal256Comparator() { super(Decimal256Vector.TYPE_WIDTH); } @Override public int compareNotNull(int index1, int index2) { BigDecimal value1 = vector1.getObjectNotNull(index1); BigDecimal value2 = vector2.getObjectNotNull(index2); return value1.compareTo(value2); } @Override public VectorValueComparator createNew() { return new Decimal256Comparator(); } } /** * Default comparator for Decimal type. The comparison is based on values, with null comes first. */ public static class DecimalComparator extends VectorValueComparator { public DecimalComparator() { super(DecimalVector.TYPE_WIDTH); } @Override public int compareNotNull(int index1, int index2) { BigDecimal value1 = vector1.getObjectNotNull(index1); BigDecimal value2 = vector2.getObjectNotNull(index2); return value1.compareTo(value2); } @Override public VectorValueComparator createNew() { return new DecimalComparator(); } } /** * Default comparator for Duration type. The comparison is based on values, with null comes first. */ public static class DurationComparator extends VectorValueComparator { public DurationComparator() { super(DurationVector.TYPE_WIDTH); } @Override public int compareNotNull(int index1, int index2) { Duration value1 = vector1.getObjectNotNull(index1); Duration value2 = vector2.getObjectNotNull(index2); return value1.compareTo(value2); } @Override public VectorValueComparator createNew() { return new DurationComparator(); } } /** * Default comparator for IntervalDay type. The comparison is based on values, with null comes * first. */ public static class IntervalDayComparator extends VectorValueComparator { public IntervalDayComparator() { super(IntervalDayVector.TYPE_WIDTH); } @Override public int compareNotNull(int index1, int index2) { Duration value1 = vector1.getObjectNotNull(index1); Duration value2 = vector2.getObjectNotNull(index2); return value1.compareTo(value2); } @Override public VectorValueComparator createNew() { return new IntervalDayComparator(); } } /** * Default comparator for TimeMicro type. The comparison is based on values, with null comes * first. */ public static class TimeMicroComparator extends VectorValueComparator { public TimeMicroComparator() { super(TimeMicroVector.TYPE_WIDTH); } @Override public int compareNotNull(int index1, int index2) { long value1 = vector1.get(index1); long value2 = vector2.get(index2); return Long.compare(value1, value2); } @Override public VectorValueComparator createNew() { return new TimeMicroComparator(); } } /** * Default comparator for TimeMilli type. The comparison is based on values, with null comes * first. */ public static class TimeMilliComparator extends VectorValueComparator { public TimeMilliComparator() { super(TimeMilliVector.TYPE_WIDTH); } @Override public int compareNotNull(int index1, int index2) { int value1 = vector1.get(index1); int value2 = vector2.get(index2); return Integer.compare(value1, value2); } @Override public VectorValueComparator createNew() { return new TimeMilliComparator(); } } /** * Default comparator for TimeNano type. The comparison is based on values, with null comes first. */ public static class TimeNanoComparator extends VectorValueComparator { public TimeNanoComparator() { super(TimeNanoVector.TYPE_WIDTH); } @Override public int compareNotNull(int index1, int index2) { long value1 = vector1.get(index1); long value2 = vector2.get(index2); return Long.compare(value1, value2); } @Override public VectorValueComparator createNew() { return new TimeNanoComparator(); } } /** * Default comparator for TimeSec type. The comparison is based on values, with null comes first. */ public static class TimeSecComparator extends VectorValueComparator { public TimeSecComparator() { super(TimeSecVector.TYPE_WIDTH); } @Override public int compareNotNull(int index1, int index2) { int value1 = vector1.get(index1); int value2 = vector2.get(index2); return Integer.compare(value1, value2); } @Override public VectorValueComparator createNew() { return new TimeSecComparator(); } } /** * Default comparator for TimeSec type. The comparison is based on values, with null comes first. */ public static class TimeStampComparator extends VectorValueComparator { public TimeStampComparator() { super(TimeStampVector.TYPE_WIDTH); } @Override public int compareNotNull(int index1, int index2) { long value1 = vector1.get(index1); long value2 = vector2.get(index2); return Long.compare(value1, value2); } @Override public VectorValueComparator createNew() { return new TimeStampComparator(); } } /** * Default comparator for {@link org.apache.arrow.vector.FixedSizeBinaryVector}. The comparison is * in lexicographic order, with null comes first. */ public static class FixedSizeBinaryComparator extends VectorValueComparator { @Override public int compare(int index1, int index2) { NullableFixedSizeBinaryHolder holder1 = new NullableFixedSizeBinaryHolder(); NullableFixedSizeBinaryHolder holder2 = new NullableFixedSizeBinaryHolder(); vector1.get(index1, holder1); vector2.get(index2, holder2); return ByteFunctionHelpers.compare( holder1.buffer, 0, holder1.byteWidth, holder2.buffer, 0, holder2.byteWidth); } @Override public int compareNotNull(int index1, int index2) { NullableFixedSizeBinaryHolder holder1 = new NullableFixedSizeBinaryHolder(); NullableFixedSizeBinaryHolder holder2 = new NullableFixedSizeBinaryHolder(); vector1.get(index1, holder1); vector2.get(index2, holder2); return ByteFunctionHelpers.compare( holder1.buffer, 0, holder1.byteWidth, holder2.buffer, 0, holder2.byteWidth); } @Override public VectorValueComparator createNew() { return new FixedSizeBinaryComparator(); } } /** Default comparator for {@link org.apache.arrow.vector.NullVector}. */ public static class NullComparator extends VectorValueComparator { @Override public int compare(int index1, int index2) { // Values are always equal (and are always null). return 0; } @Override public int compareNotNull(int index1, int index2) { throw new AssertionError("Cannot compare non-null values in a NullVector."); } @Override public VectorValueComparator createNew() { return new NullComparator(); } } /** * Default comparator for {@link org.apache.arrow.vector.VariableWidthVector}. The comparison is * in lexicographic order, with null comes first. */ public static class VariableWidthComparator extends VectorValueComparator { private final ArrowBufPointer reusablePointer1 = new ArrowBufPointer(); private final ArrowBufPointer reusablePointer2 = new ArrowBufPointer(); @Override public int compare(int index1, int index2) { vector1.getDataPointer(index1, reusablePointer1); vector2.getDataPointer(index2, reusablePointer2); return reusablePointer1.compareTo(reusablePointer2); } @Override public int compareNotNull(int index1, int index2) { vector1.getDataPointer(index1, reusablePointer1); vector2.getDataPointer(index2, reusablePointer2); return reusablePointer1.compareTo(reusablePointer2); } @Override public VectorValueComparator createNew() { return new VariableWidthComparator(); } } /** * Default comparator for {@link RepeatedValueVector}. It works by comparing the underlying vector * in a lexicographic order. * * @param inner vector type. */ public static class RepeatedValueComparator extends VectorValueComparator { private final VectorValueComparator innerComparator; public RepeatedValueComparator(VectorValueComparator innerComparator) { this.innerComparator = innerComparator; } @Override public int compareNotNull(int index1, int index2) { int startIdx1 = vector1.getOffsetBuffer().getInt((long) index1 * OFFSET_WIDTH); int startIdx2 = vector2.getOffsetBuffer().getInt((long) index2 * OFFSET_WIDTH); int endIdx1 = vector1.getOffsetBuffer().getInt((long) (index1 + 1) * OFFSET_WIDTH); int endIdx2 = vector2.getOffsetBuffer().getInt((long) (index2 + 1) * OFFSET_WIDTH); int length1 = endIdx1 - startIdx1; int length2 = endIdx2 - startIdx2; int length = Math.min(length1, length2); for (int i = 0; i < length; i++) { int result = innerComparator.compare(startIdx1 + i, startIdx2 + i); if (result != 0) { return result; } } return length1 - length2; } @Override public VectorValueComparator createNew() { VectorValueComparator newInnerComparator = innerComparator.createNew(); return new RepeatedValueComparator<>(newInnerComparator); } @Override public void attachVectors(RepeatedValueVector vector1, RepeatedValueVector vector2) { this.vector1 = vector1; this.vector2 = vector2; innerComparator.attachVectors((T) vector1.getDataVector(), (T) vector2.getDataVector()); } } /** * Default comparator for {@link RepeatedValueVector}. It works by comparing the underlying vector * in a lexicographic order. * * @param inner vector type. */ public static class FixedSizeListComparator extends VectorValueComparator { private final VectorValueComparator innerComparator; public FixedSizeListComparator(VectorValueComparator innerComparator) { this.innerComparator = innerComparator; } @Override public int compareNotNull(int index1, int index2) { int length1 = vector1.getListSize(); int length2 = vector2.getListSize(); int length = Math.min(length1, length2); int startIdx1 = vector1.getElementStartIndex(index1); int startIdx2 = vector2.getElementStartIndex(index2); for (int i = 0; i < length; i++) { int result = innerComparator.compare(startIdx1 + i, startIdx2 + i); if (result != 0) { return result; } } return length1 - length2; } @Override public VectorValueComparator createNew() { VectorValueComparator newInnerComparator = innerComparator.createNew(); return new FixedSizeListComparator<>(newInnerComparator); } @Override public void attachVectors(FixedSizeListVector vector1, FixedSizeListVector vector2) { this.vector1 = vector1; this.vector2 = vector2; innerComparator.attachVectors((T) vector1.getDataVector(), (T) vector2.getDataVector()); } } private DefaultVectorComparators() {} } FixedWidthInPlaceVectorSorter.java000066400000000000000000000111731515451615200363030ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import org.apache.arrow.vector.BaseFixedWidthVector; /** * Default in-place sorter for fixed-width vectors. It is based on quick-sort, with average time * complexity O(n*log(n)). * * @param vector type. */ public class FixedWidthInPlaceVectorSorter implements InPlaceVectorSorter { /** * If the number of items is smaller than this threshold, we will use another algorithm to sort * the data. */ public static final int CHANGE_ALGORITHM_THRESHOLD = 15; static final int STOP_CHOOSING_PIVOT_THRESHOLD = 3; VectorValueComparator comparator; /** The vector to sort. */ V vec; /** The buffer to hold the pivot. It always has length 1. */ V pivotBuffer; @Override public void sortInPlace(V vec, VectorValueComparator comparator) { try { this.vec = vec; this.comparator = comparator; this.pivotBuffer = (V) vec.getField().createVector(vec.getAllocator()); this.pivotBuffer.allocateNew(1); this.pivotBuffer.setValueCount(1); comparator.attachVectors(vec, pivotBuffer); quickSort(); } finally { this.pivotBuffer.close(); } } private void quickSort() { try (OffHeapIntStack rangeStack = new OffHeapIntStack(vec.getAllocator())) { rangeStack.push(0); rangeStack.push(vec.getValueCount() - 1); while (!rangeStack.isEmpty()) { int high = rangeStack.pop(); int low = rangeStack.pop(); if (low < high) { if (high - low < CHANGE_ALGORITHM_THRESHOLD) { // switch to insertion sort InsertionSorter.insertionSort(vec, low, high, comparator, pivotBuffer); continue; } int mid = partition(low, high); // push the larger part to stack first, // to reduce the required stack size if (high - mid < mid - low) { rangeStack.push(low); rangeStack.push(mid - 1); rangeStack.push(mid + 1); rangeStack.push(high); } else { rangeStack.push(mid + 1); rangeStack.push(high); rangeStack.push(low); rangeStack.push(mid - 1); } } } } } /** Select the pivot as the median of 3 samples. */ void choosePivot(int low, int high) { // we need at least 3 items if (high - low + 1 < STOP_CHOOSING_PIVOT_THRESHOLD) { pivotBuffer.copyFrom(low, 0, vec); return; } comparator.attachVector(vec); int mid = low + (high - low) / 2; // find the median by at most 3 comparisons int medianIdx; if (comparator.compare(low, mid) < 0) { if (comparator.compare(mid, high) < 0) { medianIdx = mid; } else { if (comparator.compare(low, high) < 0) { medianIdx = high; } else { medianIdx = low; } } } else { if (comparator.compare(mid, high) > 0) { medianIdx = mid; } else { if (comparator.compare(low, high) < 0) { medianIdx = low; } else { medianIdx = high; } } } // move the pivot to the low position, if necessary if (medianIdx != low) { pivotBuffer.copyFrom(medianIdx, 0, vec); vec.copyFrom(low, medianIdx, vec); vec.copyFrom(0, low, pivotBuffer); } else { pivotBuffer.copyFrom(low, 0, vec); } comparator.attachVectors(vec, pivotBuffer); } private int partition(int low, int high) { choosePivot(low, high); while (low < high) { while (low < high && comparator.compare(high, 0) >= 0) { high -= 1; } vec.copyFrom(high, low, vec); while (low < high && comparator.compare(low, 0) <= 0) { low += 1; } vec.copyFrom(low, high, vec); } vec.copyFrom(0, low, pivotBuffer); return low; } } FixedWidthOutOfPlaceVectorSorter.java000066400000000000000000000071751515451615200370000ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; 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.BitVector; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.IntVector; /** * Default out-of-place sorter for fixed-width vectors. It is an out-of-place sort, with time * complexity O(n*log(n)). * * @param vector type. */ public class FixedWidthOutOfPlaceVectorSorter implements OutOfPlaceVectorSorter { protected IndexSorter indexSorter = new IndexSorter<>(); @Override public void sortOutOfPlace(V srcVector, V dstVector, VectorValueComparator comparator) { if (srcVector instanceof BitVector) { throw new IllegalArgumentException( "BitVector is not supported with FixedWidthOutOfPlaceVectorSorter."); } comparator.attachVector(srcVector); int valueWidth = comparator.getValueWidth(); // buffers referenced in the sort ArrowBuf srcValueBuffer = srcVector.getDataBuffer(); ArrowBuf dstValidityBuffer = dstVector.getValidityBuffer(); ArrowBuf dstValueBuffer = dstVector.getDataBuffer(); // check buffer size Preconditions.checkArgument( dstValidityBuffer.capacity() * 8 >= srcVector.getValueCount(), "Not enough capacity for the validity buffer of the dst vector. " + "Expected capacity %s, actual capacity %s", (srcVector.getValueCount() + 7) / 8, dstValidityBuffer.capacity()); Preconditions.checkArgument( dstValueBuffer.capacity() >= srcVector.getValueCount() * ((long) srcVector.getTypeWidth()), "Not enough capacity for the data buffer of the dst vector. " + "Expected capacity %s, actual capacity %s", srcVector.getValueCount() * srcVector.getTypeWidth(), dstValueBuffer.capacity()); // sort value indices try (IntVector sortedIndices = new IntVector("", srcVector.getAllocator())) { sortedIndices.allocateNew(srcVector.getValueCount()); sortedIndices.setValueCount(srcVector.getValueCount()); indexSorter.sort(srcVector, sortedIndices, comparator); // copy sorted values to the output vector for (int dstIndex = 0; dstIndex < sortedIndices.getValueCount(); dstIndex++) { int srcIndex = sortedIndices.get(dstIndex); if (srcVector.isNull(srcIndex)) { BitVectorHelper.unsetBit(dstValidityBuffer, dstIndex); } else { BitVectorHelper.setBit(dstValidityBuffer, dstIndex); MemoryUtil.copyMemory( srcValueBuffer.memoryAddress() + srcIndex * ((long) valueWidth), dstValueBuffer.memoryAddress() + dstIndex * ((long) valueWidth), valueWidth); } } } } } GeneralOutOfPlaceVectorSorter.java000066400000000000000000000047241515451615200363130ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.ValueVector; /** * An out-of-place sorter for vectors of arbitrary type, with time complexity O(n*log(n)). Since it * does not make any assumptions about the memory layout of the vector, its performance can be * sub-optimal. So if another sorter is applicable ({@link FixedWidthInPlaceVectorSorter}), it * should be used in preference. * * @param vector type. */ public class GeneralOutOfPlaceVectorSorter implements OutOfPlaceVectorSorter { @Override public void sortOutOfPlace(V srcVector, V dstVector, VectorValueComparator comparator) { comparator.attachVector(srcVector); // check vector capacity Preconditions.checkArgument( dstVector.getValueCapacity() >= srcVector.getValueCount(), "Not enough capacity for the target vector. " + "Expected capacity %s, actual capacity %s", srcVector.getValueCount(), dstVector.getValueCapacity()); // sort value indices try (IntVector sortedIndices = new IntVector("", srcVector.getAllocator())) { sortedIndices.allocateNew(srcVector.getValueCount()); sortedIndices.setValueCount(srcVector.getValueCount()); IndexSorter indexSorter = new IndexSorter<>(); indexSorter.sort(srcVector, sortedIndices, comparator); // copy sorted values to the output vector for (int dstIndex = 0; dstIndex < sortedIndices.getValueCount(); dstIndex++) { int srcIndex = sortedIndices.get(dstIndex); dstVector.copyFromSafe(srcIndex, dstIndex, srcVector); } } } } InPlaceVectorSorter.java000066400000000000000000000024771515451615200343320ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import org.apache.arrow.vector.ValueVector; /** * Basic interface for sorting a vector in-place. That is, the sorting is performed by modifying the * input vector, without creating a new sorted vector. * * @param the vector type. */ public interface InPlaceVectorSorter { /** * Sort a vector in-place. * * @param vec the vector to sort. * @param comparator the criteria for sort. */ void sortInPlace(V vec, VectorValueComparator comparator); } apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/sort/IndexSorter.java000066400000000000000000000127701515451615200327570ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import java.util.stream.IntStream; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.ValueVector; /** * Sorter for the indices of a vector. * * @param vector type. */ public class IndexSorter { /** * If the number of items is smaller than this threshold, we will use another algorithm to sort * the data. */ public static final int CHANGE_ALGORITHM_THRESHOLD = 15; /** Comparator for vector indices. */ private VectorValueComparator comparator; /** Vector indices to sort. */ private IntVector indices; /** * Sorts indices, by quick-sort. Suppose the vector is denoted by v. After calling this method, * the following relations hold: v(indices[0]) <= v(indices[1]) <= ... * * @param vector the vector whose indices need to be sorted. * @param indices the vector for storing the sorted indices. * @param comparator the comparator to sort indices. */ public void sort(V vector, IntVector indices, VectorValueComparator comparator) { comparator.attachVector(vector); this.indices = indices; IntStream.range(0, vector.getValueCount()).forEach(i -> indices.set(i, i)); this.comparator = comparator; quickSort(); } private void quickSort() { try (OffHeapIntStack rangeStack = new OffHeapIntStack(indices.getAllocator())) { rangeStack.push(0); rangeStack.push(indices.getValueCount() - 1); while (!rangeStack.isEmpty()) { int high = rangeStack.pop(); int low = rangeStack.pop(); if (low < high) { if (high - low < CHANGE_ALGORITHM_THRESHOLD) { InsertionSorter.insertionSort(indices, low, high, comparator); continue; } int mid = partition(low, high, indices, comparator); // push the larger part to stack first, // to reduce the required stack size if (high - mid < mid - low) { rangeStack.push(low); rangeStack.push(mid - 1); rangeStack.push(mid + 1); rangeStack.push(high); } else { rangeStack.push(mid + 1); rangeStack.push(high); rangeStack.push(low); rangeStack.push(mid - 1); } } } } } /** Select the pivot as the median of 3 samples. */ static int choosePivot( int low, int high, IntVector indices, VectorValueComparator comparator) { // we need at least 3 items if (high - low + 1 < FixedWidthInPlaceVectorSorter.STOP_CHOOSING_PIVOT_THRESHOLD) { return indices.get(low); } int mid = low + (high - low) / 2; // find the median by at most 3 comparisons int medianIdx; if (comparator.compare(indices.get(low), indices.get(mid)) < 0) { if (comparator.compare(indices.get(mid), indices.get(high)) < 0) { medianIdx = mid; } else { if (comparator.compare(indices.get(low), indices.get(high)) < 0) { medianIdx = high; } else { medianIdx = low; } } } else { if (comparator.compare(indices.get(mid), indices.get(high)) > 0) { medianIdx = mid; } else { if (comparator.compare(indices.get(low), indices.get(high)) < 0) { medianIdx = low; } else { medianIdx = high; } } } // move the pivot to the low position, if necessary if (medianIdx != low) { int tmp = indices.get(medianIdx); indices.set(medianIdx, indices.get(low)); indices.set(low, tmp); return tmp; } else { return indices.get(low); } } /** * Partition a range of values in a vector into two parts, with elements in one part smaller than * elements from the other part. The partition is based on the element indices, so it does not * modify the underlying vector. * * @param low the lower bound of the range. * @param high the upper bound of the range. * @param indices vector element indices. * @param comparator criteria for comparison. * @param the vector type. * @return the index of the split point. */ public static int partition( int low, int high, IntVector indices, VectorValueComparator comparator) { int pivotIndex = choosePivot(low, high, indices, comparator); while (low < high) { while (low < high && comparator.compare(indices.get(high), pivotIndex) >= 0) { high -= 1; } indices.set(low, indices.get(high)); while (low < high && comparator.compare(indices.get(low), pivotIndex) <= 0) { low += 1; } indices.set(high, indices.get(low)); } indices.set(low, pivotIndex); return low; } } InsertionSorter.java000066400000000000000000000052171515451615200336010ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.ValueVector; /** Insertion sorter. */ class InsertionSorter { /** * Sorts the range of a vector by insertion sort. * * @param vector the vector to be sorted. * @param startIdx the start index of the range (inclusive). * @param endIdx the end index of the range (inclusive). * @param buffer an extra buffer with capacity 1 to hold the current key. * @param comparator the criteria for vector element comparison. * @param the vector type. */ static void insertionSort( V vector, int startIdx, int endIdx, VectorValueComparator comparator, V buffer) { comparator.attachVectors(vector, buffer); for (int i = startIdx; i <= endIdx; i++) { buffer.copyFrom(i, 0, vector); int j = i - 1; while (j >= startIdx && comparator.compare(j, 0) > 0) { vector.copyFrom(j, j + 1, vector); j = j - 1; } vector.copyFrom(0, j + 1, buffer); } } /** * Sorts the range of vector indices by insertion sort. * * @param indices the vector indices. * @param startIdx the start index of the range (inclusive). * @param endIdx the end index of the range (inclusive). * @param comparator the criteria for vector element comparison. * @param the vector type. */ static void insertionSort( IntVector indices, int startIdx, int endIdx, VectorValueComparator comparator) { for (int i = startIdx; i <= endIdx; i++) { int key = indices.get(i); int j = i - 1; while (j >= startIdx && comparator.compare(indices.get(j), key) > 0) { indices.set(j + 1, indices.get(j)); j = j - 1; } indices.set(j + 1, key); } } } OffHeapIntStack.java000066400000000000000000000036571515451615200334070ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.IntVector; /** An off heap implementation of stack with int elements. */ class OffHeapIntStack implements AutoCloseable { private static final int INIT_SIZE = 128; private IntVector intVector; private int top = 0; public OffHeapIntStack(BufferAllocator allocator) { intVector = new IntVector("int stack inner vector", allocator); intVector.allocateNew(INIT_SIZE); intVector.setValueCount(INIT_SIZE); } public void push(int value) { if (top == intVector.getValueCount()) { int targetCapacity = intVector.getValueCount() * 2; while (intVector.getValueCapacity() < targetCapacity) { intVector.reAlloc(); } intVector.setValueCount(targetCapacity); } intVector.set(top++, value); } public int pop() { return intVector.get(--top); } public int getTop() { return intVector.get(top - 1); } public boolean isEmpty() { return top == 0; } public int getCount() { return top; } @Override public void close() { intVector.close(); } } OutOfPlaceVectorSorter.java000066400000000000000000000026531515451615200350140ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import org.apache.arrow.vector.ValueVector; /** * Basic interface for sorting a vector out-of-place. That is, the sorting is performed on a * newly-created vector, and the original vector is not modified. * * @param the vector type. */ public interface OutOfPlaceVectorSorter { /** * Sort a vector out-of-place. * * @param inVec the input vector. * @param outVec the output vector, which has the same size as the input vector. * @param comparator the criteria for sort. */ void sortOutOfPlace(V inVec, V outVec, VectorValueComparator comparator); } StableVectorComparator.java000066400000000000000000000045261515451615200350570ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.ValueVector; /** * Stable sorter. It compares values like ordinary comparators. However, when values are equal, it * breaks ties by the value indices. Therefore, sort algorithms using this comparator always produce * stable sort results. * * @param type of the vector. */ public class StableVectorComparator extends VectorValueComparator { private final VectorValueComparator innerComparator; /** * Constructs a stable comparator from a given comparator. * * @param innerComparator the comparator to convert to stable comparator.. */ public StableVectorComparator(VectorValueComparator innerComparator) { this.innerComparator = innerComparator; } @Override public void attachVector(V vector) { super.attachVector(vector); innerComparator.attachVector(vector); } @Override public void attachVectors(V vector1, V vector2) { Preconditions.checkArgument( vector1 == vector2, "Stable comparator only supports comparing values from the same vector"); super.attachVectors(vector1, vector2); innerComparator.attachVectors(vector1, vector2); } @Override public int compareNotNull(int index1, int index2) { int result = innerComparator.compare(index1, index2); return result != 0 ? result : index1 - index2; } @Override public VectorValueComparator createNew() { return new StableVectorComparator(innerComparator.createNew()); } } VariableWidthOutOfPlaceVectorSorter.java000066400000000000000000000106271515451615200374620ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.util.MemoryUtil; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.IntVector; /** * Default sorter for variable-width vectors. It is an out-of-place sort, with time complexity * O(n*log(n)). * * @param vector type. */ public class VariableWidthOutOfPlaceVectorSorter implements OutOfPlaceVectorSorter { protected IndexSorter indexSorter = new IndexSorter<>(); @Override public void sortOutOfPlace(V srcVector, V dstVector, VectorValueComparator comparator) { comparator.attachVector(srcVector); // buffers referenced in the sort ArrowBuf srcValueBuffer = srcVector.getDataBuffer(); ArrowBuf srcOffsetBuffer = srcVector.getOffsetBuffer(); ArrowBuf dstValidityBuffer = dstVector.getValidityBuffer(); ArrowBuf dstValueBuffer = dstVector.getDataBuffer(); ArrowBuf dstOffsetBuffer = dstVector.getOffsetBuffer(); // check buffer size Preconditions.checkArgument( dstValidityBuffer.capacity() * 8 >= srcVector.getValueCount(), "Not enough capacity for the validity buffer of the dst vector. " + "Expected capacity %s, actual capacity %s", (srcVector.getValueCount() + 7) / 8, dstValidityBuffer.capacity()); Preconditions.checkArgument( dstOffsetBuffer.capacity() >= (srcVector.getValueCount() + 1) * ((long) BaseVariableWidthVector.OFFSET_WIDTH), "Not enough capacity for the offset buffer of the dst vector. " + "Expected capacity %s, actual capacity %s", (srcVector.getValueCount() + 1) * BaseVariableWidthVector.OFFSET_WIDTH, dstOffsetBuffer.capacity()); long dataSize = srcVector .getOffsetBuffer() .getInt(srcVector.getValueCount() * ((long) BaseVariableWidthVector.OFFSET_WIDTH)); Preconditions.checkArgument( dstValueBuffer.capacity() >= dataSize, "No enough capacity for the data buffer of the dst vector. " + "Expected capacity %s, actual capacity %s", dataSize, dstValueBuffer.capacity()); // sort value indices try (IntVector sortedIndices = new IntVector("", srcVector.getAllocator())) { sortedIndices.allocateNew(srcVector.getValueCount()); sortedIndices.setValueCount(srcVector.getValueCount()); indexSorter.sort(srcVector, sortedIndices, comparator); int dstOffset = 0; dstOffsetBuffer.setInt(0, 0); // copy sorted values to the output vector for (int dstIndex = 0; dstIndex < sortedIndices.getValueCount(); dstIndex++) { int srcIndex = sortedIndices.get(dstIndex); if (srcVector.isNull(srcIndex)) { BitVectorHelper.unsetBit(dstValidityBuffer, dstIndex); } else { BitVectorHelper.setBit(dstValidityBuffer, dstIndex); int srcOffset = srcOffsetBuffer.getInt(srcIndex * ((long) BaseVariableWidthVector.OFFSET_WIDTH)); int valueLength = srcOffsetBuffer.getInt((srcIndex + 1) * ((long) BaseVariableWidthVector.OFFSET_WIDTH)) - srcOffset; MemoryUtil.copyMemory( srcValueBuffer.memoryAddress() + srcOffset, dstValueBuffer.memoryAddress() + dstOffset, valueLength); dstOffset += valueLength; } dstOffsetBuffer.setInt( (dstIndex + 1) * ((long) BaseVariableWidthVector.OFFSET_WIDTH), dstOffset); } } } } VectorValueComparator.java000066400000000000000000000106301515451615200347120ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/main/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import org.apache.arrow.vector.ValueVector; /** * Compare two values at the given indices in the vectors. This is used for vector sorting. * * @param type of the vector. */ public abstract class VectorValueComparator { /** The first vector to compare. */ protected V vector1; /** The second vector to compare. */ protected V vector2; /** Width of the vector value. For variable-length vectors, this value makes no sense. */ protected int valueWidth; private boolean checkNullsOnCompare = true; /** * This value is true by default and re-computed when vectors are attached to the comparator. If * both vectors cannot contain nulls then this value is {@code false} and calls to {@code * compare(i1, i2)} are short-circuited to {@code compareNotNull(i1, i2)} thereby speeding up * comparisons resulting in faster sorts etc. */ public boolean checkNullsOnCompare() { return this.checkNullsOnCompare; } /** Constructor for variable-width vectors. */ protected VectorValueComparator() {} /** * Constructor for fixed-width vectors. * * @param valueWidth the record width (in bytes). */ protected VectorValueComparator(int valueWidth) { this.valueWidth = valueWidth; } public int getValueWidth() { return valueWidth; } /** * Attach both vectors to compare to the same input vector. * * @param vector the vector to attach. */ public void attachVector(V vector) { attachVectors(vector, vector); } /** * Attach vectors to compare. * * @param vector1 the first vector to compare. * @param vector2 the second vector to compare. */ public void attachVectors(V vector1, V vector2) { this.vector1 = vector1; this.vector2 = vector2; final boolean v1MayHaveNulls = mayHaveNulls(vector1); final boolean v2MayHaveNulls = mayHaveNulls(vector2); this.checkNullsOnCompare = v1MayHaveNulls || v2MayHaveNulls; } private boolean mayHaveNulls(V v) { if (v.getValueCount() == 0) { return true; } if (!v.getField().isNullable()) { return false; } return v.getNullCount() > 0; } /** * Compare two values, given their indices. * * @param index1 index of the first value to compare. * @param index2 index of the second value to compare. * @return an integer greater than 0, if the first value is greater; an integer smaller than 0, if * the first value is smaller; or 0, if both values are equal. */ public int compare(int index1, int index2) { if (checkNullsOnCompare) { boolean isNull1 = vector1.isNull(index1); boolean isNull2 = vector2.isNull(index2); if (isNull1 || isNull2) { if (isNull1 && isNull2) { return 0; } else if (isNull1) { // null is smaller return -1; } else { return 1; } } } return compareNotNull(index1, index2); } /** * Compare two values, given their indices. This is a fast path for comparing non-null values, so * the caller must make sure that values at both indices are not null. * * @param index1 index of the first value to compare. * @param index2 index of the second value to compare. * @return an integer greater than 0, if the first value is greater; an integer smaller than 0, if * the first value is smaller; or 0, if both values are equal. */ public abstract int compareNotNull(int index1, int index2); /** * Creates a comparator of the same type. * * @return the newly created comparator. */ public abstract VectorValueComparator createNew(); } apache-arrow-java-19.0.0/algorithm/src/test/000077500000000000000000000000001515451615200206325ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/000077500000000000000000000000001515451615200215535ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/000077500000000000000000000000001515451615200223425ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/000077500000000000000000000000001515451615200235635ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200247155ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/000077500000000000000000000000001515451615200267035ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/deduplicate/000077500000000000000000000000001515451615200311665ustar00rootroot00000000000000TestDeduplicationUtils.java000066400000000000000000000121151515451615200364170ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/deduplicate/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.deduplicate; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import java.nio.charset.StandardCharsets; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.util.DataSizeRoundingUtil; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link DeduplicationUtils}. */ public class TestDeduplicationUtils { private static final int VECTOR_LENGTH = 100; private static final int REPETITION_COUNT = 3; private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testDeduplicateFixedWidth() { try (IntVector origVec = new IntVector("original vec", allocator); IntVector dedupVec = new IntVector("deduplicated vec", allocator); IntVector lengthVec = new IntVector("length vec", allocator); ArrowBuf distinctBuf = allocator.buffer( DataSizeRoundingUtil.divideBy8Ceil(VECTOR_LENGTH * REPETITION_COUNT))) { origVec.allocateNew(VECTOR_LENGTH * REPETITION_COUNT); origVec.setValueCount(VECTOR_LENGTH * REPETITION_COUNT); lengthVec.allocateNew(); // prepare data for (int i = 0; i < VECTOR_LENGTH; i++) { for (int j = 0; j < REPETITION_COUNT; j++) { origVec.set(i * REPETITION_COUNT + j, i); } } DeduplicationUtils.populateRunStartIndicators(origVec, distinctBuf); assertEquals( VECTOR_LENGTH, VECTOR_LENGTH * REPETITION_COUNT - BitVectorHelper.getNullCount(distinctBuf, VECTOR_LENGTH * REPETITION_COUNT)); DeduplicationUtils.populateDeduplicatedValues(distinctBuf, origVec, dedupVec); assertEquals(VECTOR_LENGTH, dedupVec.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { assertEquals(i, dedupVec.get(i)); } DeduplicationUtils.populateRunLengths( distinctBuf, lengthVec, VECTOR_LENGTH * REPETITION_COUNT); assertEquals(VECTOR_LENGTH, lengthVec.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { assertEquals(REPETITION_COUNT, lengthVec.get(i)); } } } @Test public void testDeduplicateVariableWidth() { try (VarCharVector origVec = new VarCharVector("original vec", allocator); VarCharVector dedupVec = new VarCharVector("deduplicated vec", allocator); IntVector lengthVec = new IntVector("length vec", allocator); ArrowBuf distinctBuf = allocator.buffer( DataSizeRoundingUtil.divideBy8Ceil(VECTOR_LENGTH * REPETITION_COUNT))) { origVec.allocateNew(VECTOR_LENGTH * REPETITION_COUNT * 10, VECTOR_LENGTH * REPETITION_COUNT); origVec.setValueCount(VECTOR_LENGTH * REPETITION_COUNT); lengthVec.allocateNew(); // prepare data for (int i = 0; i < VECTOR_LENGTH; i++) { String str = String.valueOf(i * i); for (int j = 0; j < REPETITION_COUNT; j++) { origVec.set(i * REPETITION_COUNT + j, str.getBytes(StandardCharsets.UTF_8)); } } DeduplicationUtils.populateRunStartIndicators(origVec, distinctBuf); assertEquals( VECTOR_LENGTH, VECTOR_LENGTH * REPETITION_COUNT - BitVectorHelper.getNullCount(distinctBuf, VECTOR_LENGTH * REPETITION_COUNT)); DeduplicationUtils.populateDeduplicatedValues(distinctBuf, origVec, dedupVec); assertEquals(VECTOR_LENGTH, dedupVec.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { assertArrayEquals(String.valueOf(i * i).getBytes(StandardCharsets.UTF_8), dedupVec.get(i)); } DeduplicationUtils.populateRunLengths( distinctBuf, lengthVec, VECTOR_LENGTH * REPETITION_COUNT); assertEquals(VECTOR_LENGTH, lengthVec.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { assertEquals(REPETITION_COUNT, lengthVec.get(i)); } } } } TestVectorRunDeduplicator.java000066400000000000000000000110111515451615200370730ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/deduplicate/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.deduplicate; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import java.nio.charset.StandardCharsets; 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.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link VectorRunDeduplicator}. */ public class TestVectorRunDeduplicator { private static final int VECTOR_LENGTH = 100; private static final int REPETITION_COUNT = 3; private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testDeduplicateFixedWidth() { try (IntVector origVec = new IntVector("original vec", allocator); IntVector dedupVec = new IntVector("deduplicated vec", allocator); IntVector lengthVec = new IntVector("length vec", allocator); VectorRunDeduplicator deduplicator = new VectorRunDeduplicator<>(origVec, allocator)) { origVec.allocateNew(VECTOR_LENGTH * REPETITION_COUNT); origVec.setValueCount(VECTOR_LENGTH * REPETITION_COUNT); lengthVec.allocateNew(); // prepare data for (int i = 0; i < VECTOR_LENGTH; i++) { for (int j = 0; j < REPETITION_COUNT; j++) { origVec.set(i * REPETITION_COUNT + j, i); } } int distinctCount = deduplicator.getRunCount(); assertEquals(VECTOR_LENGTH, distinctCount); dedupVec.allocateNew(distinctCount); deduplicator.populateDeduplicatedValues(dedupVec); assertEquals(VECTOR_LENGTH, dedupVec.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { assertEquals(i, dedupVec.get(i)); } deduplicator.populateRunLengths(lengthVec); assertEquals(VECTOR_LENGTH, lengthVec.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { assertEquals(REPETITION_COUNT, lengthVec.get(i)); } } } @Test public void testDeduplicateVariableWidth() { try (VarCharVector origVec = new VarCharVector("original vec", allocator); VarCharVector dedupVec = new VarCharVector("deduplicated vec", allocator); IntVector lengthVec = new IntVector("length vec", allocator); VectorRunDeduplicator deduplicator = new VectorRunDeduplicator<>(origVec, allocator)) { origVec.allocateNew(VECTOR_LENGTH * REPETITION_COUNT * 10, VECTOR_LENGTH * REPETITION_COUNT); origVec.setValueCount(VECTOR_LENGTH * REPETITION_COUNT); lengthVec.allocateNew(); // prepare data for (int i = 0; i < VECTOR_LENGTH; i++) { String str = String.valueOf(i * i); for (int j = 0; j < REPETITION_COUNT; j++) { origVec.set(i * REPETITION_COUNT + j, str.getBytes(StandardCharsets.UTF_8)); } } int distinctCount = deduplicator.getRunCount(); assertEquals(VECTOR_LENGTH, distinctCount); dedupVec.allocateNew(distinctCount * 10L, distinctCount); deduplicator.populateDeduplicatedValues(dedupVec); assertEquals(VECTOR_LENGTH, dedupVec.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { assertArrayEquals(String.valueOf(i * i).getBytes(StandardCharsets.UTF_8), dedupVec.get(i)); } deduplicator.populateRunLengths(lengthVec); assertEquals(VECTOR_LENGTH, lengthVec.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { assertEquals(REPETITION_COUNT, lengthVec.get(i)); } } } } apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/dictionary/000077500000000000000000000000001515451615200310505ustar00rootroot00000000000000TestHashTableBasedDictionaryBuilder.java000066400000000000000000000163741515451615200406360ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.dictionary; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.charset.StandardCharsets; import java.util.Objects; 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.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link HashTableBasedDictionaryBuilder}. */ public class TestHashTableBasedDictionaryBuilder { private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testBuildVariableWidthDictionaryWithNull() { try (VarCharVector vec = new VarCharVector("", allocator); VarCharVector dictionary = new VarCharVector("", allocator)) { vec.allocateNew(100, 10); vec.setValueCount(10); dictionary.allocateNew(); // fill data vec.set(0, "hello".getBytes(StandardCharsets.UTF_8)); vec.set(1, "abc".getBytes(StandardCharsets.UTF_8)); vec.setNull(2); vec.set(3, "world".getBytes(StandardCharsets.UTF_8)); vec.set(4, "12".getBytes(StandardCharsets.UTF_8)); vec.set(5, "dictionary".getBytes(StandardCharsets.UTF_8)); vec.setNull(6); vec.set(7, "hello".getBytes(StandardCharsets.UTF_8)); vec.set(8, "good".getBytes(StandardCharsets.UTF_8)); vec.set(9, "abc".getBytes(StandardCharsets.UTF_8)); HashTableBasedDictionaryBuilder dictionaryBuilder = new HashTableBasedDictionaryBuilder<>(dictionary, true); int result = dictionaryBuilder.addValues(vec); assertEquals(7, result); assertEquals(7, dictionary.getValueCount()); assertEquals( "hello", new String(Objects.requireNonNull(dictionary.get(0)), StandardCharsets.UTF_8)); assertEquals( "abc", new String(Objects.requireNonNull(dictionary.get(1)), StandardCharsets.UTF_8)); assertNull(dictionary.get(2)); assertEquals( "world", new String(Objects.requireNonNull(dictionary.get(3)), StandardCharsets.UTF_8)); assertEquals( "12", new String(Objects.requireNonNull(dictionary.get(4)), StandardCharsets.UTF_8)); assertEquals( "dictionary", new String(Objects.requireNonNull(dictionary.get(5)), StandardCharsets.UTF_8)); assertEquals( "good", new String(Objects.requireNonNull(dictionary.get(6)), StandardCharsets.UTF_8)); } } @Test public void testBuildVariableWidthDictionaryWithoutNull() { try (VarCharVector vec = new VarCharVector("", allocator); VarCharVector dictionary = new VarCharVector("", allocator)) { vec.allocateNew(100, 10); vec.setValueCount(10); dictionary.allocateNew(); // fill data vec.set(0, "hello".getBytes(StandardCharsets.UTF_8)); vec.set(1, "abc".getBytes(StandardCharsets.UTF_8)); vec.setNull(2); vec.set(3, "world".getBytes(StandardCharsets.UTF_8)); vec.set(4, "12".getBytes(StandardCharsets.UTF_8)); vec.set(5, "dictionary".getBytes(StandardCharsets.UTF_8)); vec.setNull(6); vec.set(7, "hello".getBytes(StandardCharsets.UTF_8)); vec.set(8, "good".getBytes(StandardCharsets.UTF_8)); vec.set(9, "abc".getBytes(StandardCharsets.UTF_8)); HashTableBasedDictionaryBuilder dictionaryBuilder = new HashTableBasedDictionaryBuilder<>(dictionary, false); int result = dictionaryBuilder.addValues(vec); assertEquals(6, result); assertEquals(6, dictionary.getValueCount()); assertEquals( "hello", new String(Objects.requireNonNull(dictionary.get(0)), StandardCharsets.UTF_8)); assertEquals( "abc", new String(Objects.requireNonNull(dictionary.get(1)), StandardCharsets.UTF_8)); assertEquals( "world", new String(Objects.requireNonNull(dictionary.get(2)), StandardCharsets.UTF_8)); assertEquals( "12", new String(Objects.requireNonNull(dictionary.get(3)), StandardCharsets.UTF_8)); assertEquals( "dictionary", new String(Objects.requireNonNull(dictionary.get(4)), StandardCharsets.UTF_8)); assertEquals( "good", new String(Objects.requireNonNull(dictionary.get(5)), StandardCharsets.UTF_8)); } } @Test public void testBuildFixedWidthDictionaryWithNull() { try (IntVector vec = new IntVector("", allocator); IntVector dictionary = new IntVector("", allocator)) { vec.allocateNew(10); vec.setValueCount(10); dictionary.allocateNew(); // fill data vec.set(0, 4); vec.set(1, 8); vec.set(2, 32); vec.set(3, 8); vec.set(4, 16); vec.set(5, 32); vec.setNull(6); vec.set(7, 4); vec.set(8, 4); vec.setNull(9); HashTableBasedDictionaryBuilder dictionaryBuilder = new HashTableBasedDictionaryBuilder<>(dictionary, true); int result = dictionaryBuilder.addValues(vec); assertEquals(5, result); assertEquals(5, dictionary.getValueCount()); assertEquals(4, dictionary.get(0)); assertEquals(8, dictionary.get(1)); assertEquals(32, dictionary.get(2)); assertEquals(16, dictionary.get(3)); assertTrue(dictionary.isNull(4)); } } @Test public void testBuildFixedWidthDictionaryWithoutNull() { try (IntVector vec = new IntVector("", allocator); IntVector dictionary = new IntVector("", allocator)) { vec.allocateNew(10); vec.setValueCount(10); dictionary.allocateNew(); // fill data vec.set(0, 4); vec.set(1, 8); vec.set(2, 32); vec.set(3, 8); vec.set(4, 16); vec.set(5, 32); vec.setNull(6); vec.set(7, 4); vec.set(8, 4); vec.setNull(9); HashTableBasedDictionaryBuilder dictionaryBuilder = new HashTableBasedDictionaryBuilder<>(dictionary, false); int result = dictionaryBuilder.addValues(vec); assertEquals(4, result); assertEquals(4, dictionary.getValueCount()); assertEquals(4, dictionary.get(0)); assertEquals(8, dictionary.get(1)); assertEquals(32, dictionary.get(2)); assertEquals(16, dictionary.get(3)); } } } TestHashTableDictionaryEncoder.java000066400000000000000000000323351515451615200376630ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.dictionary; import static org.junit.jupiter.api.Assertions.assertArrayEquals; 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.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.Random; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryEncoder; import org.apache.arrow.vector.types.pojo.DictionaryEncoding; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link HashTableDictionaryEncoder}. */ public class TestHashTableDictionaryEncoder { private final int VECTOR_LENGTH = 50; private final int DICTIONARY_LENGTH = 10; private BufferAllocator allocator; byte[] zero = "000".getBytes(StandardCharsets.UTF_8); byte[] one = "111".getBytes(StandardCharsets.UTF_8); byte[] two = "222".getBytes(StandardCharsets.UTF_8); byte[][] data = new byte[][] {zero, one, two}; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testEncodeAndDecode() { Random random = new Random(); try (VarCharVector rawVector = new VarCharVector("original vector", allocator); IntVector encodedVector = new IntVector("encoded vector", allocator); VarCharVector dictionary = new VarCharVector("dictionary", allocator)) { // set up dictionary dictionary.allocateNew(); for (int i = 0; i < DICTIONARY_LENGTH; i++) { // encode "i" as i dictionary.setSafe(i, String.valueOf(i).getBytes(StandardCharsets.UTF_8)); } dictionary.setValueCount(DICTIONARY_LENGTH); // set up raw vector rawVector.allocateNew(10 * VECTOR_LENGTH, VECTOR_LENGTH); for (int i = 0; i < VECTOR_LENGTH; i++) { int val = (random.nextInt() & Integer.MAX_VALUE) % DICTIONARY_LENGTH; rawVector.set(i, String.valueOf(val).getBytes(StandardCharsets.UTF_8)); } rawVector.setValueCount(VECTOR_LENGTH); HashTableDictionaryEncoder encoder = new HashTableDictionaryEncoder<>(dictionary, false); // perform encoding encodedVector.allocateNew(); encoder.encode(rawVector, encodedVector); // verify encoding results assertEquals(rawVector.getValueCount(), encodedVector.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { assertArrayEquals( rawVector.get(i), String.valueOf(encodedVector.get(i)).getBytes(StandardCharsets.UTF_8)); } // perform decoding Dictionary dict = new Dictionary(dictionary, new DictionaryEncoding(1L, false, null)); try (VarCharVector decodedVector = (VarCharVector) DictionaryEncoder.decode(encodedVector, dict)) { // verify decoding results assertEquals(encodedVector.getValueCount(), decodedVector.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { assertArrayEquals( String.valueOf(encodedVector.get(i)).getBytes(StandardCharsets.UTF_8), decodedVector.get(i)); } } } } @Test public void testEncodeAndDecodeWithNull() { Random random = new Random(); try (VarCharVector rawVector = new VarCharVector("original vector", allocator); IntVector encodedVector = new IntVector("encoded vector", allocator); VarCharVector dictionary = new VarCharVector("dictionary", allocator)) { // set up dictionary dictionary.allocateNew(); dictionary.setNull(0); for (int i = 1; i < DICTIONARY_LENGTH; i++) { // encode "i" as i dictionary.setSafe(i, String.valueOf(i).getBytes(StandardCharsets.UTF_8)); } dictionary.setValueCount(DICTIONARY_LENGTH); // set up raw vector rawVector.allocateNew(10 * VECTOR_LENGTH, VECTOR_LENGTH); for (int i = 0; i < VECTOR_LENGTH; i++) { if (i % 10 == 0) { rawVector.setNull(i); } else { int val = (random.nextInt() & Integer.MAX_VALUE) % (DICTIONARY_LENGTH - 1) + 1; rawVector.set(i, String.valueOf(val).getBytes(StandardCharsets.UTF_8)); } } rawVector.setValueCount(VECTOR_LENGTH); HashTableDictionaryEncoder encoder = new HashTableDictionaryEncoder<>(dictionary, true); // perform encoding encodedVector.allocateNew(); encoder.encode(rawVector, encodedVector); // verify encoding results assertEquals(rawVector.getValueCount(), encodedVector.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { if (i % 10 == 0) { assertEquals(0, encodedVector.get(i)); } else { assertArrayEquals( rawVector.get(i), String.valueOf(encodedVector.get(i)).getBytes(StandardCharsets.UTF_8)); } } // perform decoding Dictionary dict = new Dictionary(dictionary, new DictionaryEncoding(1L, false, null)); try (VarCharVector decodedVector = (VarCharVector) DictionaryEncoder.decode(encodedVector, dict)) { // verify decoding results assertEquals(encodedVector.getValueCount(), decodedVector.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { if (i % 10 == 0) { assertTrue(decodedVector.isNull(i)); } else { assertArrayEquals( String.valueOf(encodedVector.get(i)).getBytes(StandardCharsets.UTF_8), decodedVector.get(i)); } } } } } @Test public void testEncodeNullWithoutNullInDictionary() { try (VarCharVector rawVector = new VarCharVector("original vector", allocator); IntVector encodedVector = new IntVector("encoded vector", allocator); VarCharVector dictionary = new VarCharVector("dictionary", allocator)) { // set up dictionary, with no null in it. dictionary.allocateNew(); for (int i = 0; i < DICTIONARY_LENGTH; i++) { // encode "i" as i dictionary.setSafe(i, String.valueOf(i).getBytes(StandardCharsets.UTF_8)); } dictionary.setValueCount(DICTIONARY_LENGTH); // the vector to encode has a null inside. rawVector.allocateNew(1); rawVector.setNull(0); rawVector.setValueCount(1); encodedVector.allocateNew(); HashTableDictionaryEncoder encoder = new HashTableDictionaryEncoder<>(dictionary, true); // the encoder should encode null, but no null in the dictionary, // so an exception should be thrown. assertThrows( IllegalArgumentException.class, () -> { encoder.encode(rawVector, encodedVector); }); } } @Test public void testEncodeStrings() { // Create a new value vector try (final VarCharVector vector = new VarCharVector("foo", allocator); final IntVector encoded = new IntVector("encoded", allocator); final VarCharVector dictionaryVector = new VarCharVector("dict", allocator)) { vector.allocateNew(512, 5); encoded.allocateNew(); // set some values vector.setSafe(0, zero, 0, zero.length); vector.setSafe(1, one, 0, one.length); vector.setSafe(2, one, 0, one.length); vector.setSafe(3, two, 0, two.length); vector.setSafe(4, zero, 0, zero.length); vector.setValueCount(5); // set some dictionary values dictionaryVector.allocateNew(512, 3); dictionaryVector.setSafe(0, zero, 0, one.length); dictionaryVector.setSafe(1, one, 0, two.length); dictionaryVector.setSafe(2, two, 0, zero.length); dictionaryVector.setValueCount(3); HashTableDictionaryEncoder encoder = new HashTableDictionaryEncoder<>(dictionaryVector); encoder.encode(vector, encoded); // verify indices assertEquals(5, encoded.getValueCount()); assertEquals(0, encoded.get(0)); assertEquals(1, encoded.get(1)); assertEquals(1, encoded.get(2)); assertEquals(2, encoded.get(3)); assertEquals(0, encoded.get(4)); // now run through the decoder and verify we get the original back Dictionary dict = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (VarCharVector decoded = (VarCharVector) DictionaryEncoder.decode(encoded, dict)) { assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < 5; i++) { assertEquals(vector.getObject(i), decoded.getObject(i)); } } } } @Test public void testEncodeLargeVector() { // Create a new value vector try (final VarCharVector vector = new VarCharVector("foo", allocator); final IntVector encoded = new IntVector("encoded", allocator); final VarCharVector dictionaryVector = new VarCharVector("dict", allocator)) { vector.allocateNew(); encoded.allocateNew(); int count = 10000; for (int i = 0; i < 10000; ++i) { vector.setSafe(i, data[i % 3], 0, data[i % 3].length); } vector.setValueCount(count); dictionaryVector.allocateNew(512, 3); dictionaryVector.setSafe(0, zero, 0, one.length); dictionaryVector.setSafe(1, one, 0, two.length); dictionaryVector.setSafe(2, two, 0, zero.length); dictionaryVector.setValueCount(3); HashTableDictionaryEncoder encoder = new HashTableDictionaryEncoder<>(dictionaryVector); encoder.encode(vector, encoded); assertEquals(count, encoded.getValueCount()); for (int i = 0; i < count; ++i) { assertEquals(i % 3, encoded.get(i)); } // now run through the decoder and verify we get the original back Dictionary dict = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (VarCharVector decoded = (VarCharVector) DictionaryEncoder.decode(encoded, dict)) { assertEquals(vector.getClass(), decoded.getClass()); assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < count; ++i) { assertEquals(vector.getObject(i), decoded.getObject(i)); } } } } @Test public void testEncodeBinaryVector() { // Create a new value vector try (final VarBinaryVector vector = new VarBinaryVector("foo", allocator); final VarBinaryVector dictionaryVector = new VarBinaryVector("dict", allocator); final IntVector encoded = new IntVector("encoded", allocator)) { vector.allocateNew(512, 5); vector.allocateNew(); encoded.allocateNew(); // set some values vector.setSafe(0, zero, 0, zero.length); vector.setSafe(1, one, 0, one.length); vector.setSafe(2, one, 0, one.length); vector.setSafe(3, two, 0, two.length); vector.setSafe(4, zero, 0, zero.length); vector.setValueCount(5); // set some dictionary values dictionaryVector.allocateNew(512, 3); dictionaryVector.setSafe(0, zero, 0, one.length); dictionaryVector.setSafe(1, one, 0, two.length); dictionaryVector.setSafe(2, two, 0, zero.length); dictionaryVector.setValueCount(3); HashTableDictionaryEncoder encoder = new HashTableDictionaryEncoder<>(dictionaryVector); encoder.encode(vector, encoded); assertEquals(5, encoded.getValueCount()); assertEquals(0, encoded.get(0)); assertEquals(1, encoded.get(1)); assertEquals(1, encoded.get(2)); assertEquals(2, encoded.get(3)); assertEquals(0, encoded.get(4)); // now run through the decoder and verify we get the original back Dictionary dict = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (VarBinaryVector decoded = (VarBinaryVector) DictionaryEncoder.decode(encoded, dict)) { assertEquals(vector.getClass(), decoded.getClass()); assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < 5; i++) { assertTrue(Arrays.equals(vector.getObject(i), decoded.getObject(i))); } } } } } TestLinearDictionaryEncoder.java000066400000000000000000000322621515451615200372410ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.dictionary; import static org.junit.jupiter.api.Assertions.assertArrayEquals; 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.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.Random; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryEncoder; import org.apache.arrow.vector.types.pojo.DictionaryEncoding; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link LinearDictionaryEncoder}. */ public class TestLinearDictionaryEncoder { private final int VECTOR_LENGTH = 50; private final int DICTIONARY_LENGTH = 10; private BufferAllocator allocator; byte[] zero = "000".getBytes(StandardCharsets.UTF_8); byte[] one = "111".getBytes(StandardCharsets.UTF_8); byte[] two = "222".getBytes(StandardCharsets.UTF_8); byte[][] data = new byte[][] {zero, one, two}; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testEncodeAndDecode() { Random random = new Random(); try (VarCharVector rawVector = new VarCharVector("original vector", allocator); IntVector encodedVector = new IntVector("encoded vector", allocator); VarCharVector dictionary = new VarCharVector("dictionary", allocator)) { // set up dictionary dictionary.allocateNew(); for (int i = 0; i < DICTIONARY_LENGTH; i++) { // encode "i" as i dictionary.setSafe(i, String.valueOf(i).getBytes(StandardCharsets.UTF_8)); } dictionary.setValueCount(DICTIONARY_LENGTH); // set up raw vector rawVector.allocateNew(10 * VECTOR_LENGTH, VECTOR_LENGTH); for (int i = 0; i < VECTOR_LENGTH; i++) { int val = (random.nextInt() & Integer.MAX_VALUE) % DICTIONARY_LENGTH; rawVector.set(i, String.valueOf(val).getBytes(StandardCharsets.UTF_8)); } rawVector.setValueCount(VECTOR_LENGTH); LinearDictionaryEncoder encoder = new LinearDictionaryEncoder<>(dictionary, false); // perform encoding encodedVector.allocateNew(); encoder.encode(rawVector, encodedVector); // verify encoding results assertEquals(rawVector.getValueCount(), encodedVector.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { assertArrayEquals( rawVector.get(i), String.valueOf(encodedVector.get(i)).getBytes(StandardCharsets.UTF_8)); } // perform decoding Dictionary dict = new Dictionary(dictionary, new DictionaryEncoding(1L, false, null)); try (VarCharVector decodedVector = (VarCharVector) DictionaryEncoder.decode(encodedVector, dict)) { // verify decoding results assertEquals(encodedVector.getValueCount(), decodedVector.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { assertArrayEquals( String.valueOf(encodedVector.get(i)).getBytes(StandardCharsets.UTF_8), decodedVector.get(i)); } } } } @Test public void testEncodeAndDecodeWithNull() { Random random = new Random(); try (VarCharVector rawVector = new VarCharVector("original vector", allocator); IntVector encodedVector = new IntVector("encoded vector", allocator); VarCharVector dictionary = new VarCharVector("dictionary", allocator)) { // set up dictionary dictionary.allocateNew(); dictionary.setNull(0); for (int i = 1; i < DICTIONARY_LENGTH; i++) { // encode "i" as i dictionary.setSafe(i, String.valueOf(i).getBytes(StandardCharsets.UTF_8)); } dictionary.setValueCount(DICTIONARY_LENGTH); // set up raw vector rawVector.allocateNew(10 * VECTOR_LENGTH, VECTOR_LENGTH); for (int i = 0; i < VECTOR_LENGTH; i++) { if (i % 10 == 0) { rawVector.setNull(i); } else { int val = (random.nextInt() & Integer.MAX_VALUE) % (DICTIONARY_LENGTH - 1) + 1; rawVector.set(i, String.valueOf(val).getBytes(StandardCharsets.UTF_8)); } } rawVector.setValueCount(VECTOR_LENGTH); LinearDictionaryEncoder encoder = new LinearDictionaryEncoder<>(dictionary, true); // perform encoding encodedVector.allocateNew(); encoder.encode(rawVector, encodedVector); // verify encoding results assertEquals(rawVector.getValueCount(), encodedVector.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { if (i % 10 == 0) { assertEquals(0, encodedVector.get(i)); } else { assertArrayEquals( rawVector.get(i), String.valueOf(encodedVector.get(i)).getBytes(StandardCharsets.UTF_8)); } } // perform decoding Dictionary dict = new Dictionary(dictionary, new DictionaryEncoding(1L, false, null)); try (VarCharVector decodedVector = (VarCharVector) DictionaryEncoder.decode(encodedVector, dict)) { // verify decoding results assertEquals(encodedVector.getValueCount(), decodedVector.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { if (i % 10 == 0) { assertTrue(decodedVector.isNull(i)); } else { assertArrayEquals( String.valueOf(encodedVector.get(i)).getBytes(StandardCharsets.UTF_8), decodedVector.get(i)); } } } } } @Test public void testEncodeNullWithoutNullInDictionary() { try (VarCharVector rawVector = new VarCharVector("original vector", allocator); IntVector encodedVector = new IntVector("encoded vector", allocator); VarCharVector dictionary = new VarCharVector("dictionary", allocator)) { // set up dictionary, with no null in it. dictionary.allocateNew(); for (int i = 0; i < DICTIONARY_LENGTH; i++) { // encode "i" as i dictionary.setSafe(i, String.valueOf(i).getBytes(StandardCharsets.UTF_8)); } dictionary.setValueCount(DICTIONARY_LENGTH); // the vector to encode has a null inside. rawVector.allocateNew(1); rawVector.setNull(0); rawVector.setValueCount(1); encodedVector.allocateNew(); LinearDictionaryEncoder encoder = new LinearDictionaryEncoder<>(dictionary, true); // the encoder should encode null, but no null in the dictionary, // so an exception should be thrown. assertThrows( IllegalArgumentException.class, () -> { encoder.encode(rawVector, encodedVector); }); } } @Test public void testEncodeStrings() { // Create a new value vector try (final VarCharVector vector = new VarCharVector("foo", allocator); final IntVector encoded = new IntVector("encoded", allocator); final VarCharVector dictionaryVector = new VarCharVector("dict", allocator)) { vector.allocateNew(512, 5); encoded.allocateNew(); // set some values vector.setSafe(0, zero, 0, zero.length); vector.setSafe(1, one, 0, one.length); vector.setSafe(2, one, 0, one.length); vector.setSafe(3, two, 0, two.length); vector.setSafe(4, zero, 0, zero.length); vector.setValueCount(5); // set some dictionary values dictionaryVector.allocateNew(512, 3); dictionaryVector.setSafe(0, zero, 0, one.length); dictionaryVector.setSafe(1, one, 0, two.length); dictionaryVector.setSafe(2, two, 0, zero.length); dictionaryVector.setValueCount(3); LinearDictionaryEncoder encoder = new LinearDictionaryEncoder<>(dictionaryVector); encoder.encode(vector, encoded); // verify indices assertEquals(5, encoded.getValueCount()); assertEquals(0, encoded.get(0)); assertEquals(1, encoded.get(1)); assertEquals(1, encoded.get(2)); assertEquals(2, encoded.get(3)); assertEquals(0, encoded.get(4)); // now run through the decoder and verify we get the original back Dictionary dict = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (VarCharVector decoded = (VarCharVector) DictionaryEncoder.decode(encoded, dict)) { assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < 5; i++) { assertEquals(vector.getObject(i), decoded.getObject(i)); } } } } @Test public void testEncodeLargeVector() { // Create a new value vector try (final VarCharVector vector = new VarCharVector("foo", allocator); final IntVector encoded = new IntVector("encoded", allocator); final VarCharVector dictionaryVector = new VarCharVector("dict", allocator)) { vector.allocateNew(); encoded.allocateNew(); int count = 10000; for (int i = 0; i < 10000; ++i) { vector.setSafe(i, data[i % 3], 0, data[i % 3].length); } vector.setValueCount(count); dictionaryVector.allocateNew(512, 3); dictionaryVector.setSafe(0, zero, 0, one.length); dictionaryVector.setSafe(1, one, 0, two.length); dictionaryVector.setSafe(2, two, 0, zero.length); dictionaryVector.setValueCount(3); LinearDictionaryEncoder encoder = new LinearDictionaryEncoder<>(dictionaryVector); encoder.encode(vector, encoded); assertEquals(count, encoded.getValueCount()); for (int i = 0; i < count; ++i) { assertEquals(i % 3, encoded.get(i)); } // now run through the decoder and verify we get the original back Dictionary dict = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (VarCharVector decoded = (VarCharVector) DictionaryEncoder.decode(encoded, dict)) { assertEquals(vector.getClass(), decoded.getClass()); assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < count; ++i) { assertEquals(vector.getObject(i), decoded.getObject(i)); } } } } @Test public void testEncodeBinaryVector() { // Create a new value vector try (final VarBinaryVector vector = new VarBinaryVector("foo", allocator); final VarBinaryVector dictionaryVector = new VarBinaryVector("dict", allocator); final IntVector encoded = new IntVector("encoded", allocator)) { vector.allocateNew(512, 5); vector.allocateNew(); encoded.allocateNew(); // set some values vector.setSafe(0, zero, 0, zero.length); vector.setSafe(1, one, 0, one.length); vector.setSafe(2, one, 0, one.length); vector.setSafe(3, two, 0, two.length); vector.setSafe(4, zero, 0, zero.length); vector.setValueCount(5); // set some dictionary values dictionaryVector.allocateNew(512, 3); dictionaryVector.setSafe(0, zero, 0, one.length); dictionaryVector.setSafe(1, one, 0, two.length); dictionaryVector.setSafe(2, two, 0, zero.length); dictionaryVector.setValueCount(3); LinearDictionaryEncoder encoder = new LinearDictionaryEncoder<>(dictionaryVector); encoder.encode(vector, encoded); assertEquals(5, encoded.getValueCount()); assertEquals(0, encoded.get(0)); assertEquals(1, encoded.get(1)); assertEquals(1, encoded.get(2)); assertEquals(2, encoded.get(3)); assertEquals(0, encoded.get(4)); // now run through the decoder and verify we get the original back Dictionary dict = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (VarBinaryVector decoded = (VarBinaryVector) DictionaryEncoder.decode(encoded, dict)) { assertEquals(vector.getClass(), decoded.getClass()); assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < 5; i++) { assertTrue(Arrays.equals(vector.getObject(i), decoded.getObject(i))); } } } } } TestSearchDictionaryEncoder.java000066400000000000000000000332621515451615200372350ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.dictionary; import static org.junit.jupiter.api.Assertions.assertArrayEquals; 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.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.Random; import org.apache.arrow.algorithm.sort.DefaultVectorComparators; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryEncoder; import org.apache.arrow.vector.types.pojo.DictionaryEncoding; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link SearchDictionaryEncoder}. */ public class TestSearchDictionaryEncoder { private final int VECTOR_LENGTH = 50; private final int DICTIONARY_LENGTH = 10; private BufferAllocator allocator; byte[] zero = "000".getBytes(StandardCharsets.UTF_8); byte[] one = "111".getBytes(StandardCharsets.UTF_8); byte[] two = "222".getBytes(StandardCharsets.UTF_8); byte[][] data = new byte[][] {zero, one, two}; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testEncodeAndDecode() { Random random = new Random(); try (VarCharVector rawVector = new VarCharVector("original vector", allocator); IntVector encodedVector = new IntVector("encoded vector", allocator); VarCharVector dictionary = new VarCharVector("dictionary", allocator)) { // set up dictionary dictionary.allocateNew(); for (int i = 0; i < DICTIONARY_LENGTH; i++) { // encode "i" as i dictionary.setSafe(i, String.valueOf(i).getBytes(StandardCharsets.UTF_8)); } dictionary.setValueCount(DICTIONARY_LENGTH); // set up raw vector rawVector.allocateNew(10 * VECTOR_LENGTH, VECTOR_LENGTH); for (int i = 0; i < VECTOR_LENGTH; i++) { int val = (random.nextInt() & Integer.MAX_VALUE) % DICTIONARY_LENGTH; rawVector.set(i, String.valueOf(val).getBytes(StandardCharsets.UTF_8)); } rawVector.setValueCount(VECTOR_LENGTH); SearchDictionaryEncoder encoder = new SearchDictionaryEncoder<>( dictionary, DefaultVectorComparators.createDefaultComparator(rawVector), false); // perform encoding encodedVector.allocateNew(); encoder.encode(rawVector, encodedVector); // verify encoding results assertEquals(rawVector.getValueCount(), encodedVector.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { assertArrayEquals( rawVector.get(i), String.valueOf(encodedVector.get(i)).getBytes(StandardCharsets.UTF_8)); } // perform decoding Dictionary dict = new Dictionary(dictionary, new DictionaryEncoding(1L, false, null)); try (VarCharVector decodedVector = (VarCharVector) DictionaryEncoder.decode(encodedVector, dict)) { // verify decoding results assertEquals(encodedVector.getValueCount(), decodedVector.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { assertArrayEquals( String.valueOf(encodedVector.get(i)).getBytes(StandardCharsets.UTF_8), decodedVector.get(i)); } } } } @Test public void testEncodeAndDecodeWithNull() { Random random = new Random(); try (VarCharVector rawVector = new VarCharVector("original vector", allocator); IntVector encodedVector = new IntVector("encoded vector", allocator); VarCharVector dictionary = new VarCharVector("dictionary", allocator)) { // set up dictionary dictionary.allocateNew(); dictionary.setNull(0); for (int i = 1; i < DICTIONARY_LENGTH; i++) { // encode "i" as i dictionary.setSafe(i, String.valueOf(i).getBytes(StandardCharsets.UTF_8)); } dictionary.setValueCount(DICTIONARY_LENGTH); // set up raw vector rawVector.allocateNew(10 * VECTOR_LENGTH, VECTOR_LENGTH); for (int i = 0; i < VECTOR_LENGTH; i++) { if (i % 10 == 0) { rawVector.setNull(i); } else { int val = (random.nextInt() & Integer.MAX_VALUE) % (DICTIONARY_LENGTH - 1) + 1; rawVector.set(i, String.valueOf(val).getBytes(StandardCharsets.UTF_8)); } } rawVector.setValueCount(VECTOR_LENGTH); SearchDictionaryEncoder encoder = new SearchDictionaryEncoder<>( dictionary, DefaultVectorComparators.createDefaultComparator(rawVector), true); // perform encoding encodedVector.allocateNew(); encoder.encode(rawVector, encodedVector); // verify encoding results assertEquals(rawVector.getValueCount(), encodedVector.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { if (i % 10 == 0) { assertEquals(0, encodedVector.get(i)); } else { assertArrayEquals( rawVector.get(i), String.valueOf(encodedVector.get(i)).getBytes(StandardCharsets.UTF_8)); } } // perform decoding Dictionary dict = new Dictionary(dictionary, new DictionaryEncoding(1L, false, null)); try (VarCharVector decodedVector = (VarCharVector) DictionaryEncoder.decode(encodedVector, dict)) { // verify decoding results assertEquals(encodedVector.getValueCount(), decodedVector.getValueCount()); for (int i = 0; i < VECTOR_LENGTH; i++) { if (i % 10 == 0) { assertTrue(decodedVector.isNull(i)); } else { assertArrayEquals( String.valueOf(encodedVector.get(i)).getBytes(StandardCharsets.UTF_8), decodedVector.get(i)); } } } } } @Test public void testEncodeNullWithoutNullInDictionary() { try (VarCharVector rawVector = new VarCharVector("original vector", allocator); IntVector encodedVector = new IntVector("encoded vector", allocator); VarCharVector dictionary = new VarCharVector("dictionary", allocator)) { // set up dictionary, with no null in it. dictionary.allocateNew(); for (int i = 0; i < DICTIONARY_LENGTH; i++) { // encode "i" as i dictionary.setSafe(i, String.valueOf(i).getBytes(StandardCharsets.UTF_8)); } dictionary.setValueCount(DICTIONARY_LENGTH); // the vector to encode has a null inside. rawVector.allocateNew(1); rawVector.setNull(0); rawVector.setValueCount(1); encodedVector.allocateNew(); SearchDictionaryEncoder encoder = new SearchDictionaryEncoder<>( dictionary, DefaultVectorComparators.createDefaultComparator(rawVector), true); // the encoder should encode null, but no null in the dictionary, // so an exception should be thrown. assertThrows( IllegalArgumentException.class, () -> { encoder.encode(rawVector, encodedVector); }); } } @Test public void testEncodeStrings() { // Create a new value vector try (final VarCharVector vector = new VarCharVector("foo", allocator); final IntVector encoded = new IntVector("encoded", allocator); final VarCharVector dictionaryVector = new VarCharVector("dict", allocator)) { vector.allocateNew(512, 5); encoded.allocateNew(); // set some values vector.setSafe(0, zero, 0, zero.length); vector.setSafe(1, one, 0, one.length); vector.setSafe(2, one, 0, one.length); vector.setSafe(3, two, 0, two.length); vector.setSafe(4, zero, 0, zero.length); vector.setValueCount(5); // set some dictionary values dictionaryVector.allocateNew(512, 3); dictionaryVector.setSafe(0, zero, 0, one.length); dictionaryVector.setSafe(1, one, 0, two.length); dictionaryVector.setSafe(2, two, 0, zero.length); dictionaryVector.setValueCount(3); SearchDictionaryEncoder encoder = new SearchDictionaryEncoder<>( dictionaryVector, DefaultVectorComparators.createDefaultComparator(vector)); encoder.encode(vector, encoded); // verify indices assertEquals(5, encoded.getValueCount()); assertEquals(0, encoded.get(0)); assertEquals(1, encoded.get(1)); assertEquals(1, encoded.get(2)); assertEquals(2, encoded.get(3)); assertEquals(0, encoded.get(4)); // now run through the decoder and verify we get the original back Dictionary dict = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (VarCharVector decoded = (VarCharVector) DictionaryEncoder.decode(encoded, dict)) { assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < 5; i++) { assertEquals(vector.getObject(i), decoded.getObject(i)); } } } } @Test public void testEncodeLargeVector() { // Create a new value vector try (final VarCharVector vector = new VarCharVector("foo", allocator); final IntVector encoded = new IntVector("encoded", allocator); final VarCharVector dictionaryVector = new VarCharVector("dict", allocator)) { vector.allocateNew(); encoded.allocateNew(); int count = 10000; for (int i = 0; i < 10000; ++i) { vector.setSafe(i, data[i % 3], 0, data[i % 3].length); } vector.setValueCount(count); dictionaryVector.allocateNew(512, 3); dictionaryVector.setSafe(0, zero, 0, one.length); dictionaryVector.setSafe(1, one, 0, two.length); dictionaryVector.setSafe(2, two, 0, zero.length); dictionaryVector.setValueCount(3); SearchDictionaryEncoder encoder = new SearchDictionaryEncoder<>( dictionaryVector, DefaultVectorComparators.createDefaultComparator(vector)); encoder.encode(vector, encoded); assertEquals(count, encoded.getValueCount()); for (int i = 0; i < count; ++i) { assertEquals(i % 3, encoded.get(i)); } // now run through the decoder and verify we get the original back Dictionary dict = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (VarCharVector decoded = (VarCharVector) DictionaryEncoder.decode(encoded, dict)) { assertEquals(vector.getClass(), decoded.getClass()); assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < count; ++i) { assertEquals(vector.getObject(i), decoded.getObject(i)); } } } } @Test public void testEncodeBinaryVector() { // Create a new value vector try (final VarBinaryVector vector = new VarBinaryVector("foo", allocator); final VarBinaryVector dictionaryVector = new VarBinaryVector("dict", allocator); final IntVector encoded = new IntVector("encoded", allocator)) { vector.allocateNew(512, 5); vector.allocateNew(); encoded.allocateNew(); // set some values vector.setSafe(0, zero, 0, zero.length); vector.setSafe(1, one, 0, one.length); vector.setSafe(2, one, 0, one.length); vector.setSafe(3, two, 0, two.length); vector.setSafe(4, zero, 0, zero.length); vector.setValueCount(5); // set some dictionary values dictionaryVector.allocateNew(512, 3); dictionaryVector.setSafe(0, zero, 0, one.length); dictionaryVector.setSafe(1, one, 0, two.length); dictionaryVector.setSafe(2, two, 0, zero.length); dictionaryVector.setValueCount(3); SearchDictionaryEncoder encoder = new SearchDictionaryEncoder<>( dictionaryVector, DefaultVectorComparators.createDefaultComparator(vector)); encoder.encode(vector, encoded); assertEquals(5, encoded.getValueCount()); assertEquals(0, encoded.get(0)); assertEquals(1, encoded.get(1)); assertEquals(1, encoded.get(2)); assertEquals(2, encoded.get(3)); assertEquals(0, encoded.get(4)); // now run through the decoder and verify we get the original back Dictionary dict = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (VarBinaryVector decoded = (VarBinaryVector) DictionaryEncoder.decode(encoded, dict)) { assertEquals(vector.getClass(), decoded.getClass()); assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < 5; i++) { assertTrue(Arrays.equals(vector.getObject(i), decoded.getObject(i))); } } } } } TestSearchTreeBasedDictionaryBuilder.java000066400000000000000000000214111515451615200410140ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.dictionary; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.charset.StandardCharsets; import java.util.Objects; import org.apache.arrow.algorithm.sort.DefaultVectorComparators; import org.apache.arrow.algorithm.sort.VectorValueComparator; 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.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link SearchTreeBasedDictionaryBuilder}. */ public class TestSearchTreeBasedDictionaryBuilder { private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testBuildVariableWidthDictionaryWithNull() { try (VarCharVector vec = new VarCharVector("", allocator); VarCharVector dictionary = new VarCharVector("", allocator); VarCharVector sortedDictionary = new VarCharVector("", allocator)) { vec.allocateNew(100, 10); vec.setValueCount(10); dictionary.allocateNew(); sortedDictionary.allocateNew(); // fill data vec.set(0, "hello".getBytes(StandardCharsets.UTF_8)); vec.set(1, "abc".getBytes(StandardCharsets.UTF_8)); vec.setNull(2); vec.set(3, "world".getBytes(StandardCharsets.UTF_8)); vec.set(4, "12".getBytes(StandardCharsets.UTF_8)); vec.set(5, "dictionary".getBytes(StandardCharsets.UTF_8)); vec.setNull(6); vec.set(7, "hello".getBytes(StandardCharsets.UTF_8)); vec.set(8, "good".getBytes(StandardCharsets.UTF_8)); vec.set(9, "abc".getBytes(StandardCharsets.UTF_8)); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); SearchTreeBasedDictionaryBuilder dictionaryBuilder = new SearchTreeBasedDictionaryBuilder<>(dictionary, comparator, true); int result = dictionaryBuilder.addValues(vec); assertEquals(7, result); assertEquals(7, dictionary.getValueCount()); dictionaryBuilder.populateSortedDictionary(sortedDictionary); assertTrue(sortedDictionary.isNull(0)); assertEquals( "12", new String(Objects.requireNonNull(sortedDictionary.get(1)), StandardCharsets.UTF_8)); assertEquals( "abc", new String(Objects.requireNonNull(sortedDictionary.get(2)), StandardCharsets.UTF_8)); assertEquals( "dictionary", new String(Objects.requireNonNull(sortedDictionary.get(3)), StandardCharsets.UTF_8)); assertEquals( "good", new String(Objects.requireNonNull(sortedDictionary.get(4)), StandardCharsets.UTF_8)); assertEquals( "hello", new String(Objects.requireNonNull(sortedDictionary.get(5)), StandardCharsets.UTF_8)); assertEquals( "world", new String(Objects.requireNonNull(sortedDictionary.get(6)), StandardCharsets.UTF_8)); } } @Test public void testBuildVariableWidthDictionaryWithoutNull() { try (VarCharVector vec = new VarCharVector("", allocator); VarCharVector dictionary = new VarCharVector("", allocator); VarCharVector sortedDictionary = new VarCharVector("", allocator)) { vec.allocateNew(100, 10); vec.setValueCount(10); dictionary.allocateNew(); sortedDictionary.allocateNew(); // fill data vec.set(0, "hello".getBytes(StandardCharsets.UTF_8)); vec.set(1, "abc".getBytes(StandardCharsets.UTF_8)); vec.setNull(2); vec.set(3, "world".getBytes(StandardCharsets.UTF_8)); vec.set(4, "12".getBytes(StandardCharsets.UTF_8)); vec.set(5, "dictionary".getBytes(StandardCharsets.UTF_8)); vec.setNull(6); vec.set(7, "hello".getBytes(StandardCharsets.UTF_8)); vec.set(8, "good".getBytes(StandardCharsets.UTF_8)); vec.set(9, "abc".getBytes(StandardCharsets.UTF_8)); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); SearchTreeBasedDictionaryBuilder dictionaryBuilder = new SearchTreeBasedDictionaryBuilder<>(dictionary, comparator, false); int result = dictionaryBuilder.addValues(vec); assertEquals(6, result); assertEquals(6, dictionary.getValueCount()); dictionaryBuilder.populateSortedDictionary(sortedDictionary); assertEquals( "12", new String(Objects.requireNonNull(sortedDictionary.get(0)), StandardCharsets.UTF_8)); assertEquals( "abc", new String(Objects.requireNonNull(sortedDictionary.get(1)), StandardCharsets.UTF_8)); assertEquals( "dictionary", new String(Objects.requireNonNull(sortedDictionary.get(2)), StandardCharsets.UTF_8)); assertEquals( "good", new String(Objects.requireNonNull(sortedDictionary.get(3)), StandardCharsets.UTF_8)); assertEquals( "hello", new String(Objects.requireNonNull(sortedDictionary.get(4)), StandardCharsets.UTF_8)); assertEquals( "world", new String(Objects.requireNonNull(sortedDictionary.get(5)), StandardCharsets.UTF_8)); } } @Test public void testBuildFixedWidthDictionaryWithNull() { try (IntVector vec = new IntVector("", allocator); IntVector dictionary = new IntVector("", allocator); IntVector sortedDictionary = new IntVector("", allocator)) { vec.allocateNew(10); vec.setValueCount(10); dictionary.allocateNew(); sortedDictionary.allocateNew(); // fill data vec.set(0, 4); vec.set(1, 8); vec.set(2, 32); vec.set(3, 8); vec.set(4, 16); vec.set(5, 32); vec.setNull(6); vec.set(7, 4); vec.set(8, 4); vec.setNull(9); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); SearchTreeBasedDictionaryBuilder dictionaryBuilder = new SearchTreeBasedDictionaryBuilder<>(dictionary, comparator, true); int result = dictionaryBuilder.addValues(vec); assertEquals(5, result); assertEquals(5, dictionary.getValueCount()); dictionaryBuilder.populateSortedDictionary(sortedDictionary); assertTrue(sortedDictionary.isNull(0)); assertEquals(4, sortedDictionary.get(1)); assertEquals(8, sortedDictionary.get(2)); assertEquals(16, sortedDictionary.get(3)); assertEquals(32, sortedDictionary.get(4)); } } @Test public void testBuildFixedWidthDictionaryWithoutNull() { try (IntVector vec = new IntVector("", allocator); IntVector dictionary = new IntVector("", allocator); IntVector sortedDictionary = new IntVector("", allocator)) { vec.allocateNew(10); vec.setValueCount(10); dictionary.allocateNew(); sortedDictionary.allocateNew(); // fill data vec.set(0, 4); vec.set(1, 8); vec.set(2, 32); vec.set(3, 8); vec.set(4, 16); vec.set(5, 32); vec.setNull(6); vec.set(7, 4); vec.set(8, 4); vec.setNull(9); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); SearchTreeBasedDictionaryBuilder dictionaryBuilder = new SearchTreeBasedDictionaryBuilder<>(dictionary, comparator, false); int result = dictionaryBuilder.addValues(vec); assertEquals(4, result); assertEquals(4, dictionary.getValueCount()); dictionaryBuilder.populateSortedDictionary(sortedDictionary); assertEquals(4, sortedDictionary.get(0)); assertEquals(8, sortedDictionary.get(1)); assertEquals(16, sortedDictionary.get(2)); assertEquals(32, sortedDictionary.get(3)); } } } apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/misc/000077500000000000000000000000001515451615200276365ustar00rootroot00000000000000TestPartialSumUtils.java000066400000000000000000000110351515451615200343640ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/misc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.misc; import static org.junit.jupiter.api.Assertions.assertEquals; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link PartialSumUtils}. */ public class TestPartialSumUtils { private static final int PARTIAL_SUM_VECTOR_LENGTH = 101; private static final int DELTA_VECTOR_LENGTH = 100; private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testToPartialSumVector() { try (IntVector delta = new IntVector("delta", allocator); IntVector partialSum = new IntVector("partial sum", allocator)) { delta.allocateNew(DELTA_VECTOR_LENGTH); delta.setValueCount(DELTA_VECTOR_LENGTH); partialSum.allocateNew(PARTIAL_SUM_VECTOR_LENGTH); // populate delta vector for (int i = 0; i < delta.getValueCount(); i++) { delta.set(i, 3); } final long sumBase = 10; PartialSumUtils.toPartialSumVector(delta, partialSum, sumBase); // verify results assertEquals(PARTIAL_SUM_VECTOR_LENGTH, partialSum.getValueCount()); for (int i = 0; i < partialSum.getValueCount(); i++) { assertEquals(i * 3L + sumBase, partialSum.get(i)); } } } @Test public void testToDeltaVector() { try (IntVector partialSum = new IntVector("partial sum", allocator); IntVector delta = new IntVector("delta", allocator)) { partialSum.allocateNew(PARTIAL_SUM_VECTOR_LENGTH); partialSum.setValueCount(PARTIAL_SUM_VECTOR_LENGTH); delta.allocateNew(DELTA_VECTOR_LENGTH); // populate delta vector final int sumBase = 10; for (int i = 0; i < partialSum.getValueCount(); i++) { partialSum.set(i, sumBase + 3 * i); } PartialSumUtils.toDeltaVector(partialSum, delta); // verify results assertEquals(DELTA_VECTOR_LENGTH, delta.getValueCount()); for (int i = 0; i < delta.getValueCount(); i++) { assertEquals(3, delta.get(i)); } } } @Test public void testFindPositionInPartialSumVector() { try (IntVector partialSum = new IntVector("partial sum", allocator)) { partialSum.allocateNew(PARTIAL_SUM_VECTOR_LENGTH); partialSum.setValueCount(PARTIAL_SUM_VECTOR_LENGTH); // populate delta vector final int sumBase = 10; for (int i = 0; i < partialSum.getValueCount(); i++) { partialSum.set(i, sumBase + 3 * i); } // search and verify results for (int i = 0; i < PARTIAL_SUM_VECTOR_LENGTH - 1; i++) { assertEquals( i, PartialSumUtils.findPositionInPartialSumVector(partialSum, sumBase + 3 * i + 1)); } } } @Test public void testFindPositionInPartialSumVectorNegative() { try (IntVector partialSum = new IntVector("partial sum", allocator)) { partialSum.allocateNew(PARTIAL_SUM_VECTOR_LENGTH); partialSum.setValueCount(PARTIAL_SUM_VECTOR_LENGTH); // populate delta vector final int sumBase = 10; for (int i = 0; i < partialSum.getValueCount(); i++) { partialSum.set(i, sumBase + 3 * i); } // search and verify results assertEquals(0, PartialSumUtils.findPositionInPartialSumVector(partialSum, sumBase)); assertEquals(-1, PartialSumUtils.findPositionInPartialSumVector(partialSum, sumBase - 1)); assertEquals( -1, PartialSumUtils.findPositionInPartialSumVector( partialSum, sumBase + 3 * (PARTIAL_SUM_VECTOR_LENGTH - 1))); } } } apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/rank/000077500000000000000000000000001515451615200276365ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/rank/TestVectorRank.java000066400000000000000000000130271515451615200334220ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.rank; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import java.nio.charset.StandardCharsets; import org.apache.arrow.algorithm.sort.DefaultVectorComparators; import org.apache.arrow.algorithm.sort.VectorValueComparator; 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.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link org.apache.arrow.algorithm.rank.VectorRank}. */ public class TestVectorRank { private BufferAllocator allocator; private static final int VECTOR_LENGTH = 10; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testFixedWidthRank() { VectorRank rank = new VectorRank<>(allocator); try (IntVector vector = new IntVector("int vec", allocator)) { vector.allocateNew(VECTOR_LENGTH); vector.setValueCount(VECTOR_LENGTH); vector.set(0, 1); vector.set(1, 5); vector.set(2, 3); vector.set(3, 7); vector.set(4, 9); vector.set(5, 8); vector.set(6, 2); vector.set(7, 0); vector.set(8, 4); vector.set(9, 6); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vector); assertEquals(7, rank.indexAtRank(vector, comparator, 0)); assertEquals(0, rank.indexAtRank(vector, comparator, 1)); assertEquals(6, rank.indexAtRank(vector, comparator, 2)); assertEquals(2, rank.indexAtRank(vector, comparator, 3)); assertEquals(8, rank.indexAtRank(vector, comparator, 4)); assertEquals(1, rank.indexAtRank(vector, comparator, 5)); assertEquals(9, rank.indexAtRank(vector, comparator, 6)); assertEquals(3, rank.indexAtRank(vector, comparator, 7)); assertEquals(5, rank.indexAtRank(vector, comparator, 8)); assertEquals(4, rank.indexAtRank(vector, comparator, 9)); } } @Test public void testVariableWidthRank() { VectorRank rank = new VectorRank<>(allocator); try (VarCharVector vector = new VarCharVector("varchar vec", allocator)) { vector.allocateNew(VECTOR_LENGTH * 5, VECTOR_LENGTH); vector.setValueCount(VECTOR_LENGTH); vector.set(0, String.valueOf(1).getBytes(StandardCharsets.UTF_8)); vector.set(1, String.valueOf(5).getBytes(StandardCharsets.UTF_8)); vector.set(2, String.valueOf(3).getBytes(StandardCharsets.UTF_8)); vector.set(3, String.valueOf(7).getBytes(StandardCharsets.UTF_8)); vector.set(4, String.valueOf(9).getBytes(StandardCharsets.UTF_8)); vector.set(5, String.valueOf(8).getBytes(StandardCharsets.UTF_8)); vector.set(6, String.valueOf(2).getBytes(StandardCharsets.UTF_8)); vector.set(7, String.valueOf(0).getBytes(StandardCharsets.UTF_8)); vector.set(8, String.valueOf(4).getBytes(StandardCharsets.UTF_8)); vector.set(9, String.valueOf(6).getBytes(StandardCharsets.UTF_8)); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vector); assertEquals(7, rank.indexAtRank(vector, comparator, 0)); assertEquals(0, rank.indexAtRank(vector, comparator, 1)); assertEquals(6, rank.indexAtRank(vector, comparator, 2)); assertEquals(2, rank.indexAtRank(vector, comparator, 3)); assertEquals(8, rank.indexAtRank(vector, comparator, 4)); assertEquals(1, rank.indexAtRank(vector, comparator, 5)); assertEquals(9, rank.indexAtRank(vector, comparator, 6)); assertEquals(3, rank.indexAtRank(vector, comparator, 7)); assertEquals(5, rank.indexAtRank(vector, comparator, 8)); assertEquals(4, rank.indexAtRank(vector, comparator, 9)); } } @Test public void testRankNegative() { VectorRank rank = new VectorRank<>(allocator); try (IntVector vector = new IntVector("int vec", allocator)) { vector.allocateNew(VECTOR_LENGTH); vector.setValueCount(VECTOR_LENGTH); vector.set(0, 1); vector.set(1, 5); vector.set(2, 3); vector.set(3, 7); vector.set(4, 9); vector.set(5, 8); vector.set(6, 2); vector.set(7, 0); vector.set(8, 4); vector.set(9, 6); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vector); assertThrows( IllegalArgumentException.class, () -> { rank.indexAtRank(vector, comparator, VECTOR_LENGTH + 1); }); } } } apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/search/000077500000000000000000000000001515451615200301505ustar00rootroot00000000000000TestParallelSearcher.java000066400000000000000000000131321515451615200350050ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/search/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.search; import static org.junit.jupiter.api.Assertions.assertEquals; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.stream.Stream; import org.apache.arrow.algorithm.sort.DefaultVectorComparators; import org.apache.arrow.algorithm.sort.VectorValueComparator; 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.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; /** Test cases for {@link ParallelSearcher}. */ public class TestParallelSearcher { private enum ComparatorType { EqualityComparator, OrderingComparator; } private static final int VECTOR_LENGTH = 10000; private BufferAllocator allocator; private ExecutorService threadPool; public static Stream getComparatorName() { List params = new ArrayList<>(); int[] threadCounts = {1, 2, 5, 10, 20, 50}; for (ComparatorType type : ComparatorType.values()) { for (int count : threadCounts) { params.add(Arguments.of(type, count)); } } return params.stream(); } @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); if (threadPool != null) { threadPool.shutdown(); } } @ParameterizedTest @MethodSource("getComparatorName") public void testParallelIntSearch(ComparatorType comparatorType, int threadCount) throws ExecutionException, InterruptedException { threadPool = Executors.newFixedThreadPool(threadCount); try (IntVector targetVector = new IntVector("targetVector", allocator); IntVector keyVector = new IntVector("keyVector", allocator)) { targetVector.allocateNew(VECTOR_LENGTH); keyVector.allocateNew(VECTOR_LENGTH); // if we are comparing elements using equality semantics, we do not need a comparator here. VectorValueComparator comparator = comparatorType == ComparatorType.EqualityComparator ? null : DefaultVectorComparators.createDefaultComparator(targetVector); for (int i = 0; i < VECTOR_LENGTH; i++) { targetVector.set(i, i); keyVector.set(i, i * 2); } targetVector.setValueCount(VECTOR_LENGTH); keyVector.setValueCount(VECTOR_LENGTH); ParallelSearcher searcher = new ParallelSearcher<>(targetVector, threadPool, threadCount); for (int i = 0; i < VECTOR_LENGTH; i++) { int pos = comparator == null ? searcher.search(keyVector, i) : searcher.search(keyVector, i, comparator); if (i * 2 < VECTOR_LENGTH) { assertEquals(i * 2, pos); } else { assertEquals(-1, pos); } } } } @ParameterizedTest @MethodSource("getComparatorName") public void testParallelStringSearch(ComparatorType comparatorType, int threadCount) throws ExecutionException, InterruptedException { threadPool = Executors.newFixedThreadPool(threadCount); try (VarCharVector targetVector = new VarCharVector("targetVector", allocator); VarCharVector keyVector = new VarCharVector("keyVector", allocator)) { targetVector.allocateNew(VECTOR_LENGTH); keyVector.allocateNew(VECTOR_LENGTH); // if we are comparing elements using equality semantics, we do not need a comparator here. VectorValueComparator comparator = comparatorType == ComparatorType.EqualityComparator ? null : DefaultVectorComparators.createDefaultComparator(targetVector); for (int i = 0; i < VECTOR_LENGTH; i++) { targetVector.setSafe(i, String.valueOf(i).getBytes(StandardCharsets.UTF_8)); keyVector.setSafe(i, String.valueOf(i * 2).getBytes(StandardCharsets.UTF_8)); } targetVector.setValueCount(VECTOR_LENGTH); keyVector.setValueCount(VECTOR_LENGTH); ParallelSearcher searcher = new ParallelSearcher<>(targetVector, threadPool, threadCount); for (int i = 0; i < VECTOR_LENGTH; i++) { int pos = comparator == null ? searcher.search(keyVector, i) : searcher.search(keyVector, i, comparator); if (i * 2 < VECTOR_LENGTH) { assertEquals(i * 2, pos); } else { assertEquals(-1, pos); } } } } } TestVectorRangeSearcher.java000066400000000000000000000142041515451615200354710ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/search/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.search; import static org.junit.jupiter.api.Assertions.assertEquals; import java.util.stream.Stream; import org.apache.arrow.algorithm.sort.DefaultVectorComparators; import org.apache.arrow.algorithm.sort.VectorValueComparator; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; /** Test cases for {@link VectorRangeSearcher}. */ public class TestVectorRangeSearcher { private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @ParameterizedTest @MethodSource("getRepeat") public void testGetLowerBounds(int repeat) { final int maxValue = 100; try (IntVector intVector = new IntVector("int vec", allocator)) { // allocate vector intVector.allocateNew(maxValue * repeat); intVector.setValueCount(maxValue * repeat); // prepare data in sorted order // each value is repeated some times for (int i = 0; i < maxValue; i++) { for (int j = 0; j < repeat; j++) { if (i == 0) { intVector.setNull(i * repeat + j); } else { intVector.set(i * repeat + j, i); } } } // do search VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(intVector); for (int i = 0; i < maxValue; i++) { int result = VectorRangeSearcher.getFirstMatch(intVector, comparator, intVector, i * repeat); assertEquals(i * ((long) repeat), result); } } } @ParameterizedTest @MethodSource("getRepeat") public void testGetLowerBoundsNegative(int repeat) { final int maxValue = 100; try (IntVector intVector = new IntVector("int vec", allocator); IntVector negVector = new IntVector("neg vec", allocator)) { // allocate vector intVector.allocateNew(maxValue * repeat); intVector.setValueCount(maxValue * repeat); negVector.allocateNew(maxValue); negVector.setValueCount(maxValue); // prepare data in sorted order // each value is repeated some times for (int i = 0; i < maxValue; i++) { for (int j = 0; j < repeat; j++) { if (i == 0) { intVector.setNull(i * repeat + j); } else { intVector.set(i * repeat + j, i); } } negVector.set(i, maxValue + i); } // do search VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(intVector); for (int i = 0; i < maxValue; i++) { int result = VectorRangeSearcher.getFirstMatch(intVector, comparator, negVector, i); assertEquals(-1, result); } } } @ParameterizedTest @MethodSource("getRepeat") public void testGetUpperBounds(int repeat) { final int maxValue = 100; try (IntVector intVector = new IntVector("int vec", allocator)) { // allocate vector intVector.allocateNew(maxValue * repeat); intVector.setValueCount(maxValue * repeat); // prepare data in sorted order // each value is repeated some times for (int i = 0; i < maxValue; i++) { for (int j = 0; j < repeat; j++) { if (i == 0) { intVector.setNull(i * repeat + j); } else { intVector.set(i * repeat + j, i); } } } // do search VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(intVector); for (int i = 0; i < maxValue; i++) { int result = VectorRangeSearcher.getLastMatch(intVector, comparator, intVector, i * repeat); assertEquals((i + 1) * repeat - 1, result); } } } @ParameterizedTest @MethodSource("getRepeat") public void testGetUpperBoundsNegative(int repeat) { final int maxValue = 100; try (IntVector intVector = new IntVector("int vec", allocator); IntVector negVector = new IntVector("neg vec", allocator)) { // allocate vector intVector.allocateNew(maxValue * repeat); intVector.setValueCount(maxValue * repeat); negVector.allocateNew(maxValue); negVector.setValueCount(maxValue); // prepare data in sorted order // each value is repeated some times for (int i = 0; i < maxValue; i++) { for (int j = 0; j < repeat; j++) { if (i == 0) { intVector.setNull(i * repeat + j); } else { intVector.set(i * repeat + j, i); } } negVector.set(i, maxValue + i); } // do search VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(intVector); for (int i = 0; i < maxValue; i++) { int result = VectorRangeSearcher.getLastMatch(intVector, comparator, negVector, i); assertEquals(-1, result); } } } public static Stream getRepeat() { return Stream.of(Arguments.of(1), Arguments.of(2), Arguments.of(5), Arguments.of(10)); } } TestVectorSearcher.java000066400000000000000000000233541515451615200345220ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/search/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.search; import static org.apache.arrow.vector.complex.BaseRepeatedValueVector.OFFSET_WIDTH; import static org.junit.jupiter.api.Assertions.assertEquals; import java.nio.charset.StandardCharsets; import org.apache.arrow.algorithm.sort.DefaultVectorComparators; import org.apache.arrow.algorithm.sort.VectorValueComparator; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.FieldType; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link org.apache.arrow.algorithm.search.VectorSearcher}. */ public class TestVectorSearcher { private final int VECTOR_LENGTH = 100; private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testBinarySearchInt() { try (IntVector rawVector = new IntVector("", allocator); IntVector negVector = new IntVector("", allocator)) { rawVector.allocateNew(VECTOR_LENGTH); rawVector.setValueCount(VECTOR_LENGTH); negVector.allocateNew(1); negVector.setValueCount(1); // prepare data in sorted order for (int i = 0; i < VECTOR_LENGTH; i++) { if (i == 0) { rawVector.setNull(i); } else { rawVector.set(i, i); } } negVector.set(0, -333); // do search VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(rawVector); for (int i = 0; i < VECTOR_LENGTH; i++) { int result = VectorSearcher.binarySearch(rawVector, comparator, rawVector, i); assertEquals(i, result); } // negative case assertEquals(-1, VectorSearcher.binarySearch(rawVector, comparator, negVector, 0)); } } @Test public void testLinearSearchInt() { try (IntVector rawVector = new IntVector("", allocator); IntVector negVector = new IntVector("", allocator)) { rawVector.allocateNew(VECTOR_LENGTH); rawVector.setValueCount(VECTOR_LENGTH); negVector.allocateNew(1); negVector.setValueCount(1); // prepare data in sorted order for (int i = 0; i < VECTOR_LENGTH; i++) { if (i == 0) { rawVector.setNull(i); } else { rawVector.set(i, i); } } negVector.set(0, -333); // do search VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(rawVector); for (int i = 0; i < VECTOR_LENGTH; i++) { int result = VectorSearcher.linearSearch(rawVector, comparator, rawVector, i); assertEquals(i, result); } // negative case assertEquals(-1, VectorSearcher.linearSearch(rawVector, comparator, negVector, 0)); } } @Test public void testBinarySearchVarChar() { try (VarCharVector rawVector = new VarCharVector("", allocator); VarCharVector negVector = new VarCharVector("", allocator)) { rawVector.allocateNew(VECTOR_LENGTH * 16, VECTOR_LENGTH); rawVector.setValueCount(VECTOR_LENGTH); negVector.allocateNew(VECTOR_LENGTH, 1); negVector.setValueCount(1); byte[] content = new byte[2]; // prepare data in sorted order for (int i = 0; i < VECTOR_LENGTH; i++) { if (i == 0) { rawVector.setNull(i); } else { int q = i / 10; int r = i % 10; content[0] = (byte) ('a' + q); content[1] = (byte) r; rawVector.set(i, content); } } negVector.set(0, "abcd".getBytes(StandardCharsets.UTF_8)); // do search VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(rawVector); for (int i = 0; i < VECTOR_LENGTH; i++) { int result = VectorSearcher.binarySearch(rawVector, comparator, rawVector, i); assertEquals(i, result); } // negative case assertEquals(-1, VectorSearcher.binarySearch(rawVector, comparator, negVector, 0)); } } @Test public void testLinearSearchVarChar() { try (VarCharVector rawVector = new VarCharVector("", allocator); VarCharVector negVector = new VarCharVector("", allocator)) { rawVector.allocateNew(VECTOR_LENGTH * 16, VECTOR_LENGTH); rawVector.setValueCount(VECTOR_LENGTH); negVector.allocateNew(VECTOR_LENGTH, 1); negVector.setValueCount(1); byte[] content = new byte[2]; // prepare data in sorted order for (int i = 0; i < VECTOR_LENGTH; i++) { if (i == 0) { rawVector.setNull(i); } else { int q = i / 10; int r = i % 10; content[0] = (byte) ('a' + q); content[1] = (byte) r; rawVector.set(i, content); } } negVector.set(0, "abcd".getBytes(StandardCharsets.UTF_8)); // do search VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(rawVector); for (int i = 0; i < VECTOR_LENGTH; i++) { int result = VectorSearcher.linearSearch(rawVector, comparator, rawVector, i); assertEquals(i, result); } // negative case assertEquals(-1, VectorSearcher.linearSearch(rawVector, comparator, negVector, 0)); } } private ListVector createListVector() { final int innerCount = 100; final int outerCount = 10; final int listLength = innerCount / outerCount; ListVector listVector = ListVector.empty("list vector", allocator); Types.MinorType type = Types.MinorType.INT; listVector.addOrGetVector(FieldType.nullable(type.getType())); listVector.allocateNew(); IntVector dataVector = (IntVector) listVector.getDataVector(); for (int i = 0; i < innerCount; i++) { dataVector.set(i, i); } dataVector.setValueCount(innerCount); for (int i = 0; i < outerCount; i++) { BitVectorHelper.setBit(listVector.getValidityBuffer(), i); listVector.getOffsetBuffer().setInt(i * OFFSET_WIDTH, i * listLength); listVector.getOffsetBuffer().setInt((i + 1) * OFFSET_WIDTH, (i + 1) * listLength); } listVector.setLastSet(outerCount - 1); listVector.setValueCount(outerCount); return listVector; } private ListVector createNegativeListVector() { final int innerCount = 100; final int outerCount = 10; final int listLength = innerCount / outerCount; ListVector listVector = ListVector.empty("list vector", allocator); Types.MinorType type = Types.MinorType.INT; listVector.addOrGetVector(FieldType.nullable(type.getType())); listVector.allocateNew(); IntVector dataVector = (IntVector) listVector.getDataVector(); for (int i = 0; i < innerCount; i++) { dataVector.set(i, i + 1000); } dataVector.setValueCount(innerCount); for (int i = 0; i < outerCount; i++) { BitVectorHelper.setBit(listVector.getValidityBuffer(), i); listVector.getOffsetBuffer().setInt(i * OFFSET_WIDTH, i * listLength); listVector.getOffsetBuffer().setInt((i + 1) * OFFSET_WIDTH, (i + 1) * listLength); } listVector.setValueCount(outerCount); return listVector; } @Test public void testBinarySearchList() { try (ListVector rawVector = createListVector(); ListVector negVector = createNegativeListVector()) { // do search VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(rawVector); for (int i = 0; i < rawVector.getValueCount(); i++) { int result = VectorSearcher.binarySearch(rawVector, comparator, rawVector, i); assertEquals(i, result); } // negative case for (int i = 0; i < rawVector.getValueCount(); i++) { int result = VectorSearcher.binarySearch(rawVector, comparator, negVector, i); assertEquals(-1, result); } } } @Test public void testLinearSearchList() { try (ListVector rawVector = createListVector(); ListVector negVector = createNegativeListVector()) { // do search VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(rawVector); for (int i = 0; i < rawVector.getValueCount(); i++) { int result = VectorSearcher.linearSearch(rawVector, comparator, rawVector, i); assertEquals(i, result); } // negative case for (int i = 0; i < rawVector.getValueCount(); i++) { int result = VectorSearcher.linearSearch(rawVector, comparator, negVector, i); assertEquals(-1, result); } } } } apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/sort/000077500000000000000000000000001515451615200276725ustar00rootroot00000000000000TestCompositeVectorComparator.java000066400000000000000000000102521515451615200364730ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.charset.StandardCharsets; import java.util.Arrays; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link CompositeVectorComparator}. */ public class TestCompositeVectorComparator { private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testCompareVectorSchemaRoot() { final int vectorLength = 10; IntVector intVec1 = new IntVector("int1", allocator); VarCharVector strVec1 = new VarCharVector("str1", allocator); IntVector intVec2 = new IntVector("int2", allocator); VarCharVector strVec2 = new VarCharVector("str2", allocator); try (VectorSchemaRoot batch1 = new VectorSchemaRoot(Arrays.asList(intVec1, strVec1)); VectorSchemaRoot batch2 = new VectorSchemaRoot(Arrays.asList(intVec2, strVec2))) { intVec1.allocateNew(vectorLength); strVec1.allocateNew(vectorLength * 10, vectorLength); intVec2.allocateNew(vectorLength); strVec2.allocateNew(vectorLength * 10, vectorLength); for (int i = 0; i < vectorLength; i++) { intVec1.set(i, i); strVec1.set(i, ("a" + i).getBytes(StandardCharsets.UTF_8)); intVec2.set(i, i); strVec2.set(i, "a5".getBytes(StandardCharsets.UTF_8)); } VectorValueComparator innerComparator1 = DefaultVectorComparators.createDefaultComparator(intVec1); innerComparator1.attachVectors(intVec1, intVec2); VectorValueComparator innerComparator2 = DefaultVectorComparators.createDefaultComparator(strVec1); innerComparator2.attachVectors(strVec1, strVec2); VectorValueComparator comparator = new CompositeVectorComparator( new VectorValueComparator[] {innerComparator1, innerComparator2}); // verify results // both elements are equal, the result is equal assertEquals(0, comparator.compare(5, 5)); // the first element being equal, the second is smaller, and the result is smaller assertTrue(comparator.compare(1, 1) < 0); assertTrue(comparator.compare(2, 2) < 0); assertTrue(comparator.compare(3, 3) < 0); // the first element being equal, the second is larger, and the result is larger assertTrue(comparator.compare(7, 7) > 0); assertTrue(comparator.compare(8, 8) > 0); assertTrue(comparator.compare(9, 9) > 0); // the first element is smaller, the result is always smaller assertTrue(comparator.compare(1, 2) < 0); assertTrue(comparator.compare(3, 7) < 0); assertTrue(comparator.compare(4, 9) < 0); // the first element is larger, the result is always larger assertTrue(comparator.compare(2, 0) > 0); assertTrue(comparator.compare(8, 7) > 0); assertTrue(comparator.compare(4, 1) > 0); } } } TestDefaultVectorComparator.java000066400000000000000000001120261515451615200361170ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import static org.apache.arrow.vector.complex.BaseRepeatedValueVector.OFFSET_WIDTH; 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 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.DateMilliVector; import org.apache.arrow.vector.Decimal256Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; 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.IntervalDayVector; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.LargeVarCharVector; 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.TimeStampMilliVector; import org.apache.arrow.vector.TimeStampVector; 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.ValueVector; 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.LargeListVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; 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.FieldType; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link DefaultVectorComparators}. */ public class TestDefaultVectorComparator { private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } private ListVector createListVector(int count) { ListVector listVector = ListVector.empty("list vector", allocator); Types.MinorType type = Types.MinorType.INT; listVector.addOrGetVector(FieldType.nullable(type.getType())); listVector.allocateNew(); IntVector dataVector = (IntVector) listVector.getDataVector(); for (int i = 0; i < count; i++) { dataVector.set(i, i); } dataVector.setValueCount(count); listVector.setNotNull(0); listVector.getOffsetBuffer().setInt(0, 0); listVector.getOffsetBuffer().setInt(OFFSET_WIDTH, count); listVector.setLastSet(0); listVector.setValueCount(1); return listVector; } @Test public void testCompareLists() { try (ListVector listVector1 = createListVector(10); ListVector listVector2 = createListVector(11)) { VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(listVector1); comparator.attachVectors(listVector1, listVector2); // prefix is smaller assertTrue(comparator.compare(0, 0) < 0); } try (ListVector listVector1 = createListVector(11); ListVector listVector2 = createListVector(11)) { ((IntVector) listVector2.getDataVector()).set(10, 110); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(listVector1); comparator.attachVectors(listVector1, listVector2); // breaking tie by the last element assertTrue(comparator.compare(0, 0) < 0); } try (ListVector listVector1 = createListVector(10); ListVector listVector2 = createListVector(10)) { VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(listVector1); comparator.attachVectors(listVector1, listVector2); // list vector elements equal assertTrue(comparator.compare(0, 0) == 0); } } @Test public void testCopiedComparatorForLists() { for (int i = 1; i < 10; i++) { for (int j = 1; j < 10; j++) { try (ListVector listVector1 = createListVector(10); ListVector listVector2 = createListVector(11)) { VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(listVector1); comparator.attachVectors(listVector1, listVector2); VectorValueComparator copyComparator = comparator.createNew(); copyComparator.attachVectors(listVector1, listVector2); assertEquals(comparator.compare(0, 0), copyComparator.compare(0, 0)); } } } } private FixedSizeListVector createFixedSizeListVector(int count) { FixedSizeListVector listVector = FixedSizeListVector.empty("list vector", count, allocator); Types.MinorType type = Types.MinorType.INT; listVector.addOrGetVector(FieldType.nullable(type.getType())); listVector.allocateNew(); IntVector dataVector = (IntVector) listVector.getDataVector(); for (int i = 0; i < count; i++) { dataVector.set(i, i); } dataVector.setValueCount(count); listVector.setNotNull(0); listVector.setValueCount(1); return listVector; } @Test public void testCompareFixedSizeLists() { try (FixedSizeListVector listVector1 = createFixedSizeListVector(10); FixedSizeListVector listVector2 = createFixedSizeListVector(11)) { VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(listVector1); comparator.attachVectors(listVector1, listVector2); // prefix is smaller assertTrue(comparator.compare(0, 0) < 0); } try (FixedSizeListVector listVector1 = createFixedSizeListVector(11); FixedSizeListVector listVector2 = createFixedSizeListVector(11)) { ((IntVector) listVector2.getDataVector()).set(10, 110); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(listVector1); comparator.attachVectors(listVector1, listVector2); // breaking tie by the last element assertTrue(comparator.compare(0, 0) < 0); } try (FixedSizeListVector listVector1 = createFixedSizeListVector(10); FixedSizeListVector listVector2 = createFixedSizeListVector(10)) { VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(listVector1); comparator.attachVectors(listVector1, listVector2); // list vector elements equal assertTrue(comparator.compare(0, 0) == 0); } } @Test public void testCompareUInt1() { try (UInt1Vector vec = new UInt1Vector("", allocator)) { vec.allocateNew(10); vec.setValueCount(10); vec.setNull(0); vec.set(1, -2); vec.set(2, -1); vec.set(3, 0); vec.set(4, 1); vec.set(5, 2); vec.set(6, -2); vec.setNull(7); vec.set(8, Byte.MAX_VALUE); vec.set(9, Byte.MIN_VALUE); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(1, 2) < 0); assertTrue(comparator.compare(1, 3) > 0); assertTrue(comparator.compare(2, 5) > 0); assertTrue(comparator.compare(4, 5) < 0); assertTrue(comparator.compare(1, 6) == 0); assertTrue(comparator.compare(0, 7) == 0); assertTrue(comparator.compare(8, 9) < 0); assertTrue(comparator.compare(4, 8) < 0); assertTrue(comparator.compare(5, 9) < 0); assertTrue(comparator.compare(2, 9) > 0); } } @Test public void testCompareUInt2() { try (UInt2Vector vec = new UInt2Vector("", allocator)) { vec.allocateNew(10); ValueVectorDataPopulator.setVector( vec, null, (char) (Character.MAX_VALUE - 1), Character.MAX_VALUE, (char) 0, (char) 1, (char) 2, (char) (Character.MAX_VALUE - 1), null, '\u7FFF', // value for the max 16-byte signed integer '\u8000' // value for the min 16-byte signed integer ); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(1, 2) < 0); assertTrue(comparator.compare(1, 3) > 0); assertTrue(comparator.compare(2, 5) > 0); assertTrue(comparator.compare(4, 5) < 0); assertEquals(0, comparator.compare(1, 6)); assertEquals(0, comparator.compare(0, 7)); assertTrue(comparator.compare(8, 9) < 0); assertTrue(comparator.compare(4, 8) < 0); assertTrue(comparator.compare(5, 9) < 0); assertTrue(comparator.compare(2, 9) > 0); } } @Test public void testCompareUInt4() { try (UInt4Vector vec = new UInt4Vector("", allocator)) { vec.allocateNew(10); vec.setValueCount(10); vec.setNull(0); vec.set(1, -2); vec.set(2, -1); vec.set(3, 0); vec.set(4, 1); vec.set(5, 2); vec.set(6, -2); vec.setNull(7); vec.set(8, Integer.MAX_VALUE); vec.set(9, Integer.MIN_VALUE); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(1, 2) < 0); assertTrue(comparator.compare(1, 3) > 0); assertTrue(comparator.compare(2, 5) > 0); assertTrue(comparator.compare(4, 5) < 0); assertTrue(comparator.compare(1, 6) == 0); assertTrue(comparator.compare(0, 7) == 0); assertTrue(comparator.compare(8, 9) < 0); assertTrue(comparator.compare(4, 8) < 0); assertTrue(comparator.compare(5, 9) < 0); assertTrue(comparator.compare(2, 9) > 0); } } @Test public void testCompareUInt8() { try (UInt8Vector vec = new UInt8Vector("", allocator)) { vec.allocateNew(10); vec.setValueCount(10); vec.setNull(0); vec.set(1, -2); vec.set(2, -1); vec.set(3, 0); vec.set(4, 1); vec.set(5, 2); vec.set(6, -2); vec.setNull(7); vec.set(8, Long.MAX_VALUE); vec.set(9, Long.MIN_VALUE); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(1, 2) < 0); assertTrue(comparator.compare(1, 3) > 0); assertTrue(comparator.compare(2, 5) > 0); assertTrue(comparator.compare(4, 5) < 0); assertTrue(comparator.compare(1, 6) == 0); assertTrue(comparator.compare(0, 7) == 0); assertTrue(comparator.compare(8, 9) < 0); assertTrue(comparator.compare(4, 8) < 0); assertTrue(comparator.compare(5, 9) < 0); assertTrue(comparator.compare(2, 9) > 0); } } @Test public void testCompareFloat4() { try (Float4Vector vec = new Float4Vector("", allocator)) { vec.allocateNew(9); ValueVectorDataPopulator.setVector( vec, -1.1f, 0.0f, 1.0f, null, 1.0f, 2.0f, Float.NaN, Float.NaN, Float.POSITIVE_INFINITY, Float.NEGATIVE_INFINITY); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(0, 2) < 0); assertTrue(comparator.compare(2, 1) > 0); // test equality assertTrue(comparator.compare(5, 5) == 0); assertTrue(comparator.compare(2, 4) == 0); // null first assertTrue(comparator.compare(3, 4) < 0); assertTrue(comparator.compare(5, 3) > 0); assertTrue(comparator.compare(8, 3) > 0); // NaN behavior. assertTrue(comparator.compare(6, 7) == 0); assertTrue(comparator.compare(7, 6) == 0); assertTrue(comparator.compare(7, 7) == 0); assertTrue(comparator.compare(6, 0) > 0); assertTrue(comparator.compare(6, 8) > 0); assertTrue(comparator.compare(6, 3) > 0); } } @Test public void testCompareFloat8() { try (Float8Vector vec = new Float8Vector("", allocator)) { vec.allocateNew(9); ValueVectorDataPopulator.setVector( vec, -1.1, 0.0, 1.0, null, 1.0, 2.0, Double.NaN, Double.NaN, Double.POSITIVE_INFINITY, Double.NEGATIVE_INFINITY); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(0, 2) < 0); assertTrue(comparator.compare(2, 1) > 0); // test equality assertTrue(comparator.compare(5, 5) == 0); assertTrue(comparator.compare(2, 4) == 0); // null first assertTrue(comparator.compare(3, 4) < 0); assertTrue(comparator.compare(5, 3) > 0); assertTrue(comparator.compare(8, 3) > 0); // NaN behavior. assertTrue(comparator.compare(6, 7) == 0); assertTrue(comparator.compare(7, 6) == 0); assertTrue(comparator.compare(7, 7) == 0); assertTrue(comparator.compare(6, 0) > 0); assertTrue(comparator.compare(6, 8) > 0); assertTrue(comparator.compare(6, 3) > 0); } } @Test public void testCompareLong() { try (BigIntVector vec = new BigIntVector("", allocator)) { vec.allocateNew(8); ValueVectorDataPopulator.setVector( vec, -1L, 0L, 1L, null, 1L, 5L, Long.MIN_VALUE + 1L, Long.MAX_VALUE); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(0, 2) < 0); assertTrue(comparator.compare(2, 1) > 0); // test equality assertTrue(comparator.compare(5, 5) == 0); assertTrue(comparator.compare(2, 4) == 0); // null first assertTrue(comparator.compare(3, 4) < 0); assertTrue(comparator.compare(5, 3) > 0); // potential overflow assertTrue(comparator.compare(6, 7) < 0); assertTrue(comparator.compare(7, 6) > 0); assertTrue(comparator.compare(7, 7) == 0); } } @Test public void testCompareInt() { try (IntVector vec = new IntVector("", allocator)) { vec.allocateNew(8); ValueVectorDataPopulator.setVector( vec, -1, 0, 1, null, 1, 5, Integer.MIN_VALUE + 1, Integer.MAX_VALUE); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(0, 2) < 0); assertTrue(comparator.compare(2, 1) > 0); // test equality assertTrue(comparator.compare(5, 5) == 0); assertTrue(comparator.compare(2, 4) == 0); // null first assertTrue(comparator.compare(3, 4) < 0); assertTrue(comparator.compare(5, 3) > 0); // potential overflow assertTrue(comparator.compare(6, 7) < 0); assertTrue(comparator.compare(7, 6) > 0); assertTrue(comparator.compare(7, 7) == 0); } } @Test public void testCompareShort() { try (SmallIntVector vec = new SmallIntVector("", allocator)) { vec.allocateNew(8); ValueVectorDataPopulator.setVector( vec, (short) -1, (short) 0, (short) 1, null, (short) 1, (short) 5, (short) (Short.MIN_VALUE + 1), Short.MAX_VALUE); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(0, 2) < 0); assertTrue(comparator.compare(2, 1) > 0); // test equality assertTrue(comparator.compare(5, 5) == 0); assertTrue(comparator.compare(2, 4) == 0); // null first assertTrue(comparator.compare(3, 4) < 0); assertTrue(comparator.compare(5, 3) > 0); // potential overflow assertTrue(comparator.compare(6, 7) < 0); assertTrue(comparator.compare(7, 6) > 0); assertTrue(comparator.compare(7, 7) == 0); } } @Test public void testCompareByte() { try (TinyIntVector vec = new TinyIntVector("", allocator)) { vec.allocateNew(8); ValueVectorDataPopulator.setVector( vec, (byte) -1, (byte) 0, (byte) 1, null, (byte) 1, (byte) 5, (byte) (Byte.MIN_VALUE + 1), Byte.MAX_VALUE); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(0, 2) < 0); assertTrue(comparator.compare(2, 1) > 0); // test equality assertTrue(comparator.compare(5, 5) == 0); assertTrue(comparator.compare(2, 4) == 0); // null first assertTrue(comparator.compare(3, 4) < 0); assertTrue(comparator.compare(5, 3) > 0); // potential overflow assertTrue(comparator.compare(6, 7) < 0); assertTrue(comparator.compare(7, 6) > 0); assertTrue(comparator.compare(7, 7) == 0); } } @Test public void testCompareBit() { try (BitVector vec = new BitVector("", allocator)) { vec.allocateNew(6); ValueVectorDataPopulator.setVector(vec, 1, 2, 0, 0, -1, null); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) == 0); assertTrue(comparator.compare(0, 2) > 0); assertTrue(comparator.compare(0, 4) == 0); assertTrue(comparator.compare(2, 1) < 0); assertTrue(comparator.compare(2, 4) < 0); // null first assertTrue(comparator.compare(5, 0) < 0); assertTrue(comparator.compare(5, 2) < 0); } } @Test public void testCompareDateDay() { try (DateDayVector vec = new DateDayVector("", allocator)) { vec.allocateNew(8); ValueVectorDataPopulator.setVector( vec, -1, 0, 1, null, 1, 5, Integer.MIN_VALUE + 1, Integer.MAX_VALUE); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(0, 2) < 0); assertTrue(comparator.compare(2, 1) > 0); // test equality assertTrue(comparator.compare(5, 5) == 0); assertTrue(comparator.compare(2, 4) == 0); // null first assertTrue(comparator.compare(3, 4) < 0); assertTrue(comparator.compare(5, 3) > 0); // potential overflow assertTrue(comparator.compare(6, 7) < 0); assertTrue(comparator.compare(7, 6) > 0); assertTrue(comparator.compare(7, 7) == 0); } } @Test public void testCompareDateMilli() { try (DateMilliVector vec = new DateMilliVector("", allocator)) { vec.allocateNew(8); ValueVectorDataPopulator.setVector( vec, -1L, 0L, 1L, null, 1L, 5L, Long.MIN_VALUE + 1L, Long.MAX_VALUE); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(0, 2) < 0); assertTrue(comparator.compare(2, 1) > 0); // test equality assertTrue(comparator.compare(5, 5) == 0); assertTrue(comparator.compare(2, 4) == 0); // null first assertTrue(comparator.compare(3, 4) < 0); assertTrue(comparator.compare(5, 3) > 0); // potential overflow assertTrue(comparator.compare(6, 7) < 0); assertTrue(comparator.compare(7, 6) > 0); assertTrue(comparator.compare(7, 7) == 0); } } @Test public void testCompareDecimal() { try (DecimalVector vec = new DecimalVector("", allocator, 10, 1)) { vec.allocateNew(8); ValueVectorDataPopulator.setVector( vec, -1L, 0L, 1L, null, 1L, 5L, Long.MIN_VALUE + 1L, Long.MAX_VALUE); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(0, 2) < 0); assertTrue(comparator.compare(2, 1) > 0); // test equality assertTrue(comparator.compare(5, 5) == 0); assertTrue(comparator.compare(2, 4) == 0); // null first assertTrue(comparator.compare(3, 4) < 0); assertTrue(comparator.compare(5, 3) > 0); // potential overflow assertTrue(comparator.compare(6, 7) < 0); assertTrue(comparator.compare(7, 6) > 0); assertTrue(comparator.compare(7, 7) == 0); } } @Test public void testCompareDecimal256() { try (Decimal256Vector vec = new Decimal256Vector("", allocator, 10, 1)) { vec.allocateNew(8); ValueVectorDataPopulator.setVector( vec, -1L, 0L, 1L, null, 1L, 5L, Long.MIN_VALUE + 1L, Long.MAX_VALUE); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(0, 2) < 0); assertTrue(comparator.compare(2, 1) > 0); // test equality assertTrue(comparator.compare(5, 5) == 0); assertTrue(comparator.compare(2, 4) == 0); // null first assertTrue(comparator.compare(3, 4) < 0); assertTrue(comparator.compare(5, 3) > 0); // potential overflow assertTrue(comparator.compare(6, 7) < 0); assertTrue(comparator.compare(7, 6) > 0); assertTrue(comparator.compare(7, 7) == 0); } } @Test public void testCompareDuration() { try (DurationVector vec = new DurationVector( "", FieldType.nullable(new ArrowType.Duration(TimeUnit.MILLISECOND)), allocator)) { vec.allocateNew(8); ValueVectorDataPopulator.setVector( vec, -1L, 0L, 1L, null, 1L, 5L, Long.MIN_VALUE + 1L, Long.MAX_VALUE); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(0, 2) < 0); assertTrue(comparator.compare(2, 1) > 0); // test equality assertTrue(comparator.compare(5, 5) == 0); assertTrue(comparator.compare(2, 4) == 0); // null first assertTrue(comparator.compare(3, 4) < 0); assertTrue(comparator.compare(5, 3) > 0); // potential overflow assertTrue(comparator.compare(6, 7) < 0); assertTrue(comparator.compare(7, 6) > 0); assertTrue(comparator.compare(7, 7) == 0); } } @Test public void testCompareIntervalDay() { try (IntervalDayVector vec = new IntervalDayVector( "", FieldType.nullable(new ArrowType.Duration(TimeUnit.MILLISECOND)), allocator)) { vec.allocateNew(8); vec.set(0, -1, 0); vec.set(1, 0, 0); vec.set(2, 1, 0); vec.setNull(3); vec.set(4, -1, -1); vec.set(5, 1, 1); vec.set(6, 1, 1); vec.set(7, -1, -1); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(0, 2) < 0); assertTrue(comparator.compare(2, 1) > 0); assertTrue(comparator.compare(2, 5) < 0); assertTrue(comparator.compare(0, 4) > 0); // test equality assertTrue(comparator.compare(5, 6) == 0); assertTrue(comparator.compare(4, 7) == 0); // null first assertTrue(comparator.compare(3, 4) < 0); assertTrue(comparator.compare(5, 3) > 0); } } @Test public void testCompareTimeMicro() { try (TimeMicroVector vec = new TimeMicroVector("", allocator)) { vec.allocateNew(8); ValueVectorDataPopulator.setVector( vec, -1L, 0L, 1L, null, 1L, 5L, Long.MIN_VALUE + 1L, Long.MAX_VALUE); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(0, 2) < 0); assertTrue(comparator.compare(2, 1) > 0); // test equality assertTrue(comparator.compare(5, 5) == 0); assertTrue(comparator.compare(2, 4) == 0); // null first assertTrue(comparator.compare(3, 4) < 0); assertTrue(comparator.compare(5, 3) > 0); // potential overflow assertTrue(comparator.compare(6, 7) < 0); assertTrue(comparator.compare(7, 6) > 0); assertTrue(comparator.compare(7, 7) == 0); } } @Test public void testCompareTimeMilli() { try (TimeMilliVector vec = new TimeMilliVector("", allocator)) { vec.allocateNew(8); ValueVectorDataPopulator.setVector( vec, -1, 0, 1, null, 1, 5, Integer.MIN_VALUE + 1, Integer.MAX_VALUE); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(0, 2) < 0); assertTrue(comparator.compare(2, 1) > 0); // test equality assertTrue(comparator.compare(5, 5) == 0); assertTrue(comparator.compare(2, 4) == 0); // null first assertTrue(comparator.compare(3, 4) < 0); assertTrue(comparator.compare(5, 3) > 0); // potential overflow assertTrue(comparator.compare(6, 7) < 0); assertTrue(comparator.compare(7, 6) > 0); assertTrue(comparator.compare(7, 7) == 0); } } @Test public void testCompareTimeNano() { try (TimeNanoVector vec = new TimeNanoVector("", allocator)) { vec.allocateNew(8); ValueVectorDataPopulator.setVector( vec, -1L, 0L, 1L, null, 1L, 5L, Long.MIN_VALUE + 1L, Long.MAX_VALUE); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(0, 2) < 0); assertTrue(comparator.compare(2, 1) > 0); // test equality assertTrue(comparator.compare(5, 5) == 0); assertTrue(comparator.compare(2, 4) == 0); // null first assertTrue(comparator.compare(3, 4) < 0); assertTrue(comparator.compare(5, 3) > 0); // potential overflow assertTrue(comparator.compare(6, 7) < 0); assertTrue(comparator.compare(7, 6) > 0); assertTrue(comparator.compare(7, 7) == 0); } } @Test public void testCompareTimeSec() { try (TimeSecVector vec = new TimeSecVector("", allocator)) { vec.allocateNew(8); ValueVectorDataPopulator.setVector( vec, -1, 0, 1, null, 1, 5, Integer.MIN_VALUE + 1, Integer.MAX_VALUE); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(0, 2) < 0); assertTrue(comparator.compare(2, 1) > 0); // test equality assertTrue(comparator.compare(5, 5) == 0); assertTrue(comparator.compare(2, 4) == 0); // null first assertTrue(comparator.compare(3, 4) < 0); assertTrue(comparator.compare(5, 3) > 0); // potential overflow assertTrue(comparator.compare(6, 7) < 0); assertTrue(comparator.compare(7, 6) > 0); assertTrue(comparator.compare(7, 7) == 0); } } @Test public void testCompareTimeStamp() { try (TimeStampMilliVector vec = new TimeStampMilliVector("", allocator)) { vec.allocateNew(8); ValueVectorDataPopulator.setVector( vec, -1L, 0L, 1L, null, 1L, 5L, Long.MIN_VALUE + 1L, Long.MAX_VALUE); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.compare(0, 1) < 0); assertTrue(comparator.compare(0, 2) < 0); assertTrue(comparator.compare(2, 1) > 0); // test equality assertTrue(comparator.compare(5, 5) == 0); assertTrue(comparator.compare(2, 4) == 0); // null first assertTrue(comparator.compare(3, 4) < 0); assertTrue(comparator.compare(5, 3) > 0); // potential overflow assertTrue(comparator.compare(6, 7) < 0); assertTrue(comparator.compare(7, 6) > 0); assertTrue(comparator.compare(7, 7) == 0); } } @Test public void testCompareFixedSizeBinary() { try (FixedSizeBinaryVector vector1 = new FixedSizeBinaryVector("test1", allocator, 2); FixedSizeBinaryVector vector2 = new FixedSizeBinaryVector("test1", allocator, 3)) { vector1.allocateNew(); vector2.allocateNew(); vector1.set(0, new byte[] {1, 1}); vector2.set(0, new byte[] {1, 1, 0}); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vector1); comparator.attachVectors(vector1, vector2); // prefix is smaller assertTrue(comparator.compare(0, 0) < 0); } try (FixedSizeBinaryVector vector1 = new FixedSizeBinaryVector("test1", allocator, 3); FixedSizeBinaryVector vector2 = new FixedSizeBinaryVector("test1", allocator, 3)) { vector1.allocateNew(); vector2.allocateNew(); vector1.set(0, new byte[] {1, 1, 0}); vector2.set(0, new byte[] {1, 1, 1}); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vector1); comparator.attachVectors(vector1, vector2); // breaking tie by the last element assertTrue(comparator.compare(0, 0) < 0); } try (FixedSizeBinaryVector vector1 = new FixedSizeBinaryVector("test1", allocator, 3); FixedSizeBinaryVector vector2 = new FixedSizeBinaryVector("test1", allocator, 3)) { vector1.allocateNew(); vector2.allocateNew(); vector1.set(0, new byte[] {1, 1, 1}); vector2.set(0, new byte[] {1, 1, 1}); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vector1); comparator.attachVectors(vector1, vector2); // list vector elements equal assertTrue(comparator.compare(0, 0) == 0); } } @Test public void testCompareNull() { try (NullVector vec = new NullVector("test", FieldType.notNullable(new ArrowType.Int(32, false)))) { vec.setValueCount(2); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertEquals(DefaultVectorComparators.NullComparator.class, comparator.getClass()); assertEquals(0, comparator.compare(0, 1)); } } @Test public void testCheckNullsOnCompareIsFalseForNonNullableVector() { try (IntVector vec = new IntVector( "not nullable", FieldType.notNullable(new ArrowType.Int(32, false)), allocator)) { ValueVectorDataPopulator.setVector(vec, 1, 2, 3, 4); final VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertFalse(comparator.checkNullsOnCompare()); } } @Test public void testCheckNullsOnCompareIsTrueForNullableVector() { try (IntVector vec = new IntVector("nullable", FieldType.nullable(new ArrowType.Int(32, false)), allocator); IntVector vec2 = new IntVector( "not-nullable", FieldType.notNullable(new ArrowType.Int(32, false)), allocator)) { ValueVectorDataPopulator.setVector(vec, 1, null, 3, 4); ValueVectorDataPopulator.setVector(vec2, 1, 2, 3, 4); final VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertTrue(comparator.checkNullsOnCompare()); comparator.attachVectors(vec, vec2); assertTrue(comparator.checkNullsOnCompare()); } } @Test public void testCheckNullsOnCompareIsFalseWithNoNulls() { try (IntVector vec = new IntVector("nullable", FieldType.nullable(new ArrowType.Int(32, false)), allocator); IntVector vec2 = new IntVector( "also-nullable", FieldType.nullable(new ArrowType.Int(32, false)), allocator)) { // no null values ValueVectorDataPopulator.setVector(vec, 1, 2, 3, 4); ValueVectorDataPopulator.setVector(vec2, 1, 2, 3, 4); final VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec); assertFalse(comparator.checkNullsOnCompare()); comparator.attachVectors(vec, vec2); assertFalse(comparator.checkNullsOnCompare()); } } @Test public void testCheckNullsOnCompareIsTrueWithEmptyVectors() { try (IntVector vec = new IntVector("nullable", FieldType.nullable(new ArrowType.Int(32, false)), allocator); IntVector vec2 = new IntVector( "also-nullable", FieldType.nullable(new ArrowType.Int(32, false)), allocator)) { final VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); comparator.attachVector(vec2); assertTrue(comparator.checkNullsOnCompare()); comparator.attachVectors(vec, vec2); assertTrue(comparator.checkNullsOnCompare()); } } @Test public void testVariableWidthDefaultComparators() { try (VarCharVector vec = new VarCharVector("test", allocator)) { verifyVariableWidthComparatorReturned(vec); } try (VarBinaryVector vec = new VarBinaryVector("test", allocator)) { verifyVariableWidthComparatorReturned(vec); } try (LargeVarCharVector vec = new LargeVarCharVector("test", allocator)) { verifyVariableWidthComparatorReturned(vec); } try (LargeVarBinaryVector vec = new LargeVarBinaryVector("test", allocator)) { verifyVariableWidthComparatorReturned(vec); } } private static void verifyVariableWidthComparatorReturned(V vec) { VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); assertEquals(DefaultVectorComparators.VariableWidthComparator.class, comparator.getClass()); } @Test public void testRepeatedDefaultComparators() { final FieldType type = FieldType.nullable(Types.MinorType.INT.getType()); try (final LargeListVector vector = new LargeListVector("list", allocator, type, null)) { vector.addOrGetVector(FieldType.nullable(type.getType())); verifyRepeatedComparatorReturned(vector); } } private static void verifyRepeatedComparatorReturned(V vec) { VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); assertEquals(DefaultVectorComparators.RepeatedValueComparator.class, comparator.getClass()); } } TestFixedWidthInPlaceVectorSorter.java000066400000000000000000000173671515451615200372110ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.stream.IntStream; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link FixedWidthInPlaceVectorSorter}. */ public class TestFixedWidthInPlaceVectorSorter { private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testSortInt() { try (IntVector vec = new IntVector("", allocator)) { vec.allocateNew(10); vec.setValueCount(10); // fill data to sort vec.set(0, 10); vec.set(1, 8); vec.setNull(2); vec.set(3, 10); vec.set(4, 12); vec.set(5, 17); vec.setNull(6); vec.set(7, 23); vec.set(8, 35); vec.set(9, 2); // sort the vector FixedWidthInPlaceVectorSorter sorter = new FixedWidthInPlaceVectorSorter(); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); sorter.sortInPlace(vec, comparator); // verify results assertEquals(10, vec.getValueCount()); assertTrue(vec.isNull(0)); assertTrue(vec.isNull(1)); assertEquals(2, vec.get(2)); assertEquals(8, vec.get(3)); assertEquals(10, vec.get(4)); assertEquals(10, vec.get(5)); assertEquals(12, vec.get(6)); assertEquals(17, vec.get(7)); assertEquals(23, vec.get(8)); assertEquals(35, vec.get(9)); } } /** * Tests the worst case for quick sort. It may cause stack overflow if the algorithm is * implemented as a recursive algorithm. */ @Test public void testSortLargeIncreasingInt() { final int vectorLength = 20000; try (IntVector vec = new IntVector("", allocator)) { vec.allocateNew(vectorLength); vec.setValueCount(vectorLength); // fill data to sort for (int i = 0; i < vectorLength; i++) { vec.set(i, i); } // sort the vector FixedWidthInPlaceVectorSorter sorter = new FixedWidthInPlaceVectorSorter(); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); sorter.sortInPlace(vec, comparator); // verify results assertEquals(vectorLength, vec.getValueCount()); for (int i = 0; i < vectorLength; i++) { assertEquals(i, vec.get(i)); } } } @Test public void testChoosePivot() { final int vectorLength = 100; try (IntVector vec = new IntVector("", allocator)) { vec.allocateNew(vectorLength); // the vector is sorted, so the pivot should be in the middle for (int i = 0; i < vectorLength; i++) { vec.set(i, i * 100); } vec.setValueCount(vectorLength); FixedWidthInPlaceVectorSorter sorter = new FixedWidthInPlaceVectorSorter(); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); try (IntVector pivotBuffer = (IntVector) vec.getField().createVector(allocator)) { // setup internal data structures pivotBuffer.allocateNew(1); sorter.pivotBuffer = pivotBuffer; sorter.comparator = comparator; sorter.vec = vec; comparator.attachVectors(vec, pivotBuffer); int low = 5; int high = 6; int pivotValue = vec.get(low); assertTrue(high - low + 1 < FixedWidthInPlaceVectorSorter.STOP_CHOOSING_PIVOT_THRESHOLD); // the range is small enough, so the pivot is simply selected as the low value sorter.choosePivot(low, high); assertEquals(pivotValue, vec.get(low)); low = 30; high = 80; pivotValue = vec.get((low + high) / 2); assertTrue(high - low + 1 >= FixedWidthInPlaceVectorSorter.STOP_CHOOSING_PIVOT_THRESHOLD); // the range is large enough, so the median is selected as the pivot sorter.choosePivot(low, high); assertEquals(pivotValue, vec.get(low)); } } } /** Evaluates choosing pivot for all possible permutations of 3 numbers. */ @Test public void testChoosePivotAllPermutes() { try (IntVector vec = new IntVector("", allocator)) { vec.allocateNew(3); FixedWidthInPlaceVectorSorter sorter = new FixedWidthInPlaceVectorSorter(); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); try (IntVector pivotBuffer = (IntVector) vec.getField().createVector(allocator)) { // setup internal data structures pivotBuffer.allocateNew(1); sorter.pivotBuffer = pivotBuffer; sorter.comparator = comparator; sorter.vec = vec; comparator.attachVectors(vec, pivotBuffer); int low = 0; int high = 2; ValueVectorDataPopulator.setVector(vec, 11, 22, 33); sorter.choosePivot(low, high); assertEquals(22, vec.get(0)); ValueVectorDataPopulator.setVector(vec, 11, 33, 22); sorter.choosePivot(low, high); assertEquals(22, vec.get(0)); ValueVectorDataPopulator.setVector(vec, 22, 11, 33); sorter.choosePivot(low, high); assertEquals(22, vec.get(0)); ValueVectorDataPopulator.setVector(vec, 22, 33, 11); sorter.choosePivot(low, high); assertEquals(22, vec.get(0)); ValueVectorDataPopulator.setVector(vec, 33, 11, 22); sorter.choosePivot(low, high); assertEquals(22, vec.get(0)); ValueVectorDataPopulator.setVector(vec, 33, 22, 11); sorter.choosePivot(low, high); assertEquals(22, vec.get(0)); } } } @Test public void testSortInt2() { try (IntVector vector = new IntVector("vector", allocator)) { ValueVectorDataPopulator.setVector( vector, 0, 1, 2, 3, 4, 5, 30, 31, 32, 33, 34, 35, 60, 61, 62, 63, 64, 65, 6, 7, 8, 9, 10, 11, 36, 37, 38, 39, 40, 41, 66, 67, 68, 69, 70, 71); FixedWidthInPlaceVectorSorter sorter = new FixedWidthInPlaceVectorSorter(); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vector); sorter.sortInPlace(vector, comparator); int[] actual = new int[vector.getValueCount()]; IntStream.range(0, vector.getValueCount()).forEach(i -> actual[i] = vector.get(i)); assertArrayEquals( new int[] { 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71 }, actual); } } } TestFixedWidthOutOfPlaceVectorSorter.java000066400000000000000000000314101515451615200376600ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.stream.IntStream; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.SmallIntVector; import org.apache.arrow.vector.TinyIntVector; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.MethodSource; /** Test cases for {@link FixedWidthOutOfPlaceVectorSorter}. */ public class TestFixedWidthOutOfPlaceVectorSorter extends TestOutOfPlaceVectorSorter { private BufferAllocator allocator; OutOfPlaceVectorSorter getSorter(boolean generalSorter) { return generalSorter ? new GeneralOutOfPlaceVectorSorter<>() : new FixedWidthOutOfPlaceVectorSorter<>(); } @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @ParameterizedTest @MethodSource("getParameter") public void testSortByte(boolean generalSorter) { try (TinyIntVector vec = new TinyIntVector("", allocator)) { vec.allocateNew(10); vec.setValueCount(10); // fill data to sort vec.set(0, 10); vec.set(1, 8); vec.setNull(2); vec.set(3, 10); vec.set(4, 12); vec.set(5, 17); vec.setNull(6); vec.set(7, 23); vec.set(8, 35); vec.set(9, 2); // sort the vector OutOfPlaceVectorSorter sorter = getSorter(generalSorter); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); TinyIntVector sortedVec = (TinyIntVector) vec.getField().getFieldType().createNewSingleVector("", allocator, null); sortedVec.allocateNew(vec.getValueCount()); sortedVec.setValueCount(vec.getValueCount()); sorter.sortOutOfPlace(vec, sortedVec, comparator); // verify results assertEquals(vec.getValueCount(), sortedVec.getValueCount()); assertTrue(sortedVec.isNull(0)); assertTrue(sortedVec.isNull(1)); assertEquals((byte) 2, sortedVec.get(2)); assertEquals((byte) 8, sortedVec.get(3)); assertEquals((byte) 10, sortedVec.get(4)); assertEquals((byte) 10, sortedVec.get(5)); assertEquals((byte) 12, sortedVec.get(6)); assertEquals((byte) 17, sortedVec.get(7)); assertEquals((byte) 23, sortedVec.get(8)); assertEquals((byte) 35, sortedVec.get(9)); sortedVec.close(); } } @ParameterizedTest @MethodSource("getParameter") public void testSortShort(boolean generalSorter) { try (SmallIntVector vec = new SmallIntVector("", allocator)) { vec.allocateNew(10); vec.setValueCount(10); // fill data to sort vec.set(0, 10); vec.set(1, 8); vec.setNull(2); vec.set(3, 10); vec.set(4, 12); vec.set(5, 17); vec.setNull(6); vec.set(7, 23); vec.set(8, 35); vec.set(9, 2); // sort the vector OutOfPlaceVectorSorter sorter = getSorter(generalSorter); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); SmallIntVector sortedVec = (SmallIntVector) vec.getField().getFieldType().createNewSingleVector("", allocator, null); sortedVec.allocateNew(vec.getValueCount()); sortedVec.setValueCount(vec.getValueCount()); sorter.sortOutOfPlace(vec, sortedVec, comparator); // verify results assertEquals(vec.getValueCount(), sortedVec.getValueCount()); assertTrue(sortedVec.isNull(0)); assertTrue(sortedVec.isNull(1)); assertEquals((short) 2, sortedVec.get(2)); assertEquals((short) 8, sortedVec.get(3)); assertEquals((short) 10, sortedVec.get(4)); assertEquals((short) 10, sortedVec.get(5)); assertEquals((short) 12, sortedVec.get(6)); assertEquals((short) 17, sortedVec.get(7)); assertEquals((short) 23, sortedVec.get(8)); assertEquals((short) 35, sortedVec.get(9)); sortedVec.close(); } } @ParameterizedTest @MethodSource("getParameter") public void testSortInt(boolean generalSorter) { try (IntVector vec = new IntVector("", allocator)) { vec.allocateNew(10); vec.setValueCount(10); // fill data to sort vec.set(0, 10); vec.set(1, 8); vec.setNull(2); vec.set(3, 10); vec.set(4, 12); vec.set(5, 17); vec.setNull(6); vec.set(7, 23); vec.set(8, 35); vec.set(9, 2); // sort the vector OutOfPlaceVectorSorter sorter = getSorter(generalSorter); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); IntVector sortedVec = (IntVector) vec.getField().getFieldType().createNewSingleVector("", allocator, null); sortedVec.allocateNew(vec.getValueCount()); sortedVec.setValueCount(vec.getValueCount()); sorter.sortOutOfPlace(vec, sortedVec, comparator); // verify results assertEquals(vec.getValueCount(), sortedVec.getValueCount()); assertTrue(sortedVec.isNull(0)); assertTrue(sortedVec.isNull(1)); assertEquals(2, sortedVec.get(2)); assertEquals(8, sortedVec.get(3)); assertEquals(10, sortedVec.get(4)); assertEquals(10, sortedVec.get(5)); assertEquals(12, sortedVec.get(6)); assertEquals(17, sortedVec.get(7)); assertEquals(23, sortedVec.get(8)); assertEquals(35, sortedVec.get(9)); sortedVec.close(); } } @ParameterizedTest @MethodSource("getParameter") public void testSortLong(boolean generalSorter) { try (BigIntVector vec = new BigIntVector("", allocator)) { vec.allocateNew(10); vec.setValueCount(10); // fill data to sort vec.set(0, 10L); vec.set(1, 8L); vec.setNull(2); vec.set(3, 10L); vec.set(4, 12L); vec.set(5, 17L); vec.setNull(6); vec.set(7, 23L); vec.set(8, 1L << 35L); vec.set(9, 2L); // sort the vector OutOfPlaceVectorSorter sorter = getSorter(generalSorter); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); BigIntVector sortedVec = (BigIntVector) vec.getField().getFieldType().createNewSingleVector("", allocator, null); sortedVec.allocateNew(vec.getValueCount()); sortedVec.setValueCount(vec.getValueCount()); sorter.sortOutOfPlace(vec, sortedVec, comparator); // verify results assertEquals(vec.getValueCount(), sortedVec.getValueCount()); assertTrue(sortedVec.isNull(0)); assertTrue(sortedVec.isNull(1)); assertEquals(2L, sortedVec.get(2)); assertEquals(8L, sortedVec.get(3)); assertEquals(10L, sortedVec.get(4)); assertEquals(10L, sortedVec.get(5)); assertEquals(12L, sortedVec.get(6)); assertEquals(17L, sortedVec.get(7)); assertEquals(23L, sortedVec.get(8)); assertEquals(1L << 35L, sortedVec.get(9)); sortedVec.close(); } } @ParameterizedTest @MethodSource("getParameter") public void testSortFloat(boolean generalSorter) { try (Float4Vector vec = new Float4Vector("", allocator)) { vec.allocateNew(10); vec.setValueCount(10); // fill data to sort vec.set(0, 10f); vec.set(1, 8f); vec.setNull(2); vec.set(3, 10f); vec.set(4, 12f); vec.set(5, 17f); vec.setNull(6); vec.set(7, 23f); vec.set(8, Float.NaN); vec.set(9, 2f); // sort the vector OutOfPlaceVectorSorter sorter = getSorter(generalSorter); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); Float4Vector sortedVec = (Float4Vector) vec.getField().getFieldType().createNewSingleVector("", allocator, null); sortedVec.allocateNew(vec.getValueCount()); sortedVec.setValueCount(vec.getValueCount()); sorter.sortOutOfPlace(vec, sortedVec, comparator); // verify results assertEquals(vec.getValueCount(), sortedVec.getValueCount()); assertTrue(sortedVec.isNull(0)); assertTrue(sortedVec.isNull(1)); assertEquals(2f, sortedVec.get(2), 0f); assertEquals(8f, sortedVec.get(3), 0f); assertEquals(10f, sortedVec.get(4), 0f); assertEquals(10f, sortedVec.get(5), 0f); assertEquals(12f, sortedVec.get(6), 0f); assertEquals(17f, sortedVec.get(7), 0f); assertEquals(23f, sortedVec.get(8), 0f); assertEquals(Float.NaN, sortedVec.get(9), 0f); sortedVec.close(); } } @ParameterizedTest @MethodSource("getParameter") public void testSortDouble(boolean generalSorter) { try (Float8Vector vec = new Float8Vector("", allocator)) { vec.allocateNew(10); vec.setValueCount(10); // fill data to sort vec.set(0, 10); vec.set(1, 8); vec.setNull(2); vec.set(3, 10); vec.set(4, 12); vec.set(5, 17); vec.setNull(6); vec.set(7, Double.NaN); vec.set(8, 35); vec.set(9, 2); // sort the vector OutOfPlaceVectorSorter sorter = getSorter(generalSorter); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); Float8Vector sortedVec = (Float8Vector) vec.getField().getFieldType().createNewSingleVector("", allocator, null); sortedVec.allocateNew(vec.getValueCount()); sortedVec.setValueCount(vec.getValueCount()); sorter.sortOutOfPlace(vec, sortedVec, comparator); // verify results assertEquals(vec.getValueCount(), sortedVec.getValueCount()); assertTrue(sortedVec.isNull(0)); assertTrue(sortedVec.isNull(1)); assertEquals(2, sortedVec.get(2), 0); assertEquals(8, sortedVec.get(3), 0); assertEquals(10, sortedVec.get(4), 0); assertEquals(10, sortedVec.get(5), 0); assertEquals(12, sortedVec.get(6), 0); assertEquals(17, sortedVec.get(7), 0); assertEquals(35, sortedVec.get(8), 0); assertEquals(Double.NaN, sortedVec.get(9), 0); sortedVec.close(); } } @ParameterizedTest @MethodSource("getParameter") public void testSortInt2(boolean generalSorter) { try (IntVector vec = new IntVector("", allocator)) { ValueVectorDataPopulator.setVector( vec, 0, 1, 2, 3, 4, 5, 30, 31, 32, 33, 34, 35, 60, 61, 62, 63, 64, 65, 6, 7, 8, 9, 10, 11, 36, 37, 38, 39, 40, 41, 66, 67, 68, 69, 70, 71); // sort the vector OutOfPlaceVectorSorter sorter = getSorter(generalSorter); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); try (IntVector sortedVec = (IntVector) vec.getField().getFieldType().createNewSingleVector("", allocator, null)) { sortedVec.allocateNew(vec.getValueCount()); sortedVec.setValueCount(vec.getValueCount()); sorter.sortOutOfPlace(vec, sortedVec, comparator); // verify results int[] actual = new int[sortedVec.getValueCount()]; IntStream.range(0, sortedVec.getValueCount()).forEach(i -> actual[i] = sortedVec.get(i)); assertArrayEquals( new int[] { 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71 }, actual); } } } } TestFixedWidthSorting.java000066400000000000000000000147731515451615200347370ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import java.util.ArrayList; import java.util.List; import java.util.function.Function; import java.util.stream.Stream; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.SmallIntVector; import org.apache.arrow.vector.TinyIntVector; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; /** Test sorting fixed width vectors with random data. */ public class TestFixedWidthSorting> { static final int[] VECTOR_LENGTHS = new int[] {2, 5, 10, 50, 100, 1000, 3000}; static final double[] NULL_FRACTIONS = {0, 0.1, 0.3, 0.5, 0.7, 0.9, 1}; private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(Integer.MAX_VALUE); } @AfterEach public void shutdown() { allocator.close(); } @ParameterizedTest @MethodSource("getParameters") public void testSort( boolean inPlace, int length, double nullFraction, String desc, Function vectorGenerator, TestSortingUtil.DataGenerator dataGenerator) { if (inPlace) { sortInPlace(length, nullFraction, vectorGenerator, dataGenerator); } else { sortOutOfPlace(length, nullFraction, vectorGenerator, dataGenerator); } } void sortInPlace( int length, double nullFraction, Function vectorGenerator, TestSortingUtil.DataGenerator dataGenerator) { try (V vector = vectorGenerator.apply(allocator)) { U[] array = dataGenerator.populate(vector, length, nullFraction); TestSortingUtil.sortArray(array); FixedWidthInPlaceVectorSorter sorter = new FixedWidthInPlaceVectorSorter(); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vector); sorter.sortInPlace(vector, comparator); TestSortingUtil.verifyResults(vector, array); } } void sortOutOfPlace( int length, double nullFraction, Function vectorGenerator, TestSortingUtil.DataGenerator dataGenerator) { try (V vector = vectorGenerator.apply(allocator)) { U[] array = dataGenerator.populate(vector, length, nullFraction); TestSortingUtil.sortArray(array); // sort the vector FixedWidthOutOfPlaceVectorSorter sorter = new FixedWidthOutOfPlaceVectorSorter(); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vector); try (V sortedVec = (V) vector.getField().getFieldType().createNewSingleVector("", allocator, null)) { sortedVec.allocateNew(vector.getValueCount()); sortedVec.setValueCount(vector.getValueCount()); sorter.sortOutOfPlace(vector, sortedVec, comparator); // verify results TestSortingUtil.verifyResults(sortedVec, array); } } } public static Stream getParameters() { List params = new ArrayList<>(); for (int length : VECTOR_LENGTHS) { for (double nullFrac : NULL_FRACTIONS) { for (boolean inPlace : new boolean[] {true, false}) { params.add( Arguments.of( inPlace, length, nullFrac, "TinyIntVector", (Function) allocator -> new TinyIntVector("vector", allocator), TestSortingUtil.TINY_INT_GENERATOR)); params.add( Arguments.of( inPlace, length, nullFrac, "SmallIntVector", (Function) allocator -> new SmallIntVector("vector", allocator), TestSortingUtil.SMALL_INT_GENERATOR)); params.add( Arguments.of( inPlace, length, nullFrac, "IntVector", (Function) allocator -> new IntVector("vector", allocator), TestSortingUtil.INT_GENERATOR)); params.add( Arguments.of( inPlace, length, nullFrac, "BigIntVector", (Function) allocator -> new BigIntVector("vector", allocator), TestSortingUtil.LONG_GENERATOR)); params.add( Arguments.of( inPlace, length, nullFrac, "Float4Vector", (Function) allocator -> new Float4Vector("vector", allocator), TestSortingUtil.FLOAT_GENERATOR)); params.add( Arguments.of( inPlace, length, nullFrac, "Float8Vector", (Function) allocator -> new Float8Vector("vector", allocator), TestSortingUtil.DOUBLE_GENERATOR)); } } } return params.stream(); } } TestGeneralOutOfPlaceVectorSorter.java000066400000000000000000000120721515451615200372010ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import static org.junit.jupiter.api.Assertions.assertEquals; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.FieldType; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link GeneralOutOfPlaceVectorSorter}. */ public class TestGeneralOutOfPlaceVectorSorter { private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } VectorValueComparator getComparator(StructVector structVector) { IntVector child0 = structVector.getChild("column0", IntVector.class); VectorValueComparator childComp0 = DefaultVectorComparators.createDefaultComparator(child0); childComp0.attachVector(child0); IntVector child1 = structVector.getChild("column1", IntVector.class); VectorValueComparator childComp1 = DefaultVectorComparators.createDefaultComparator(child1); childComp1.attachVector(child1); VectorValueComparator comp = new VectorValueComparator() { @Override public int compareNotNull(int index1, int index2) { // compare values by lexicographic order int result0 = childComp0.compare(index1, index2); if (result0 != 0) { return result0; } return childComp1.compare(index1, index2); } @Override public VectorValueComparator createNew() { return this; } }; return comp; } @Test public void testSortStructVector() { final int vectorLength = 7; try (StructVector srcVector = StructVector.empty("src struct", allocator); StructVector dstVector = StructVector.empty("dst struct", allocator)) { IntVector srcChild0 = srcVector.addOrGet( "column0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); IntVector srcChild1 = srcVector.addOrGet( "column1", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); IntVector dstChild0 = dstVector.addOrGet( "column0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); IntVector dstChild1 = dstVector.addOrGet( "column1", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); // src struct vector values: // [ // (2, 1) // (3, 4) // (5, 4) // (null, 3) // (7, null) // (null, null) // (6, 6) // ] ValueVectorDataPopulator.setVector(srcChild0, 2, 3, 5, null, 7, null, 6); ValueVectorDataPopulator.setVector(srcChild1, 1, 4, 4, 3, null, null, 6); srcVector.setIndexDefined(0); srcVector.setIndexDefined(1); srcVector.setIndexDefined(2); srcVector.setIndexDefined(3); srcVector.setIndexDefined(4); srcVector.setIndexDefined(6); srcVector.setValueCount(vectorLength); dstChild0.allocateNew(vectorLength); dstChild1.allocateNew(vectorLength); dstVector.setValueCount(vectorLength); // construct the comparator VectorValueComparator comp = getComparator(srcVector); // sort the vector GeneralOutOfPlaceVectorSorter sorter = new GeneralOutOfPlaceVectorSorter<>(); sorter.sortOutOfPlace(srcVector, dstVector, comp); // validate results assertEquals(vectorLength, dstVector.getValueCount()); assertEquals( "[" + "null, " + "{\"column1\":3}, " + "{\"column0\":2,\"column1\":1}, " + "{\"column0\":3,\"column1\":4}, " + "{\"column0\":5,\"column1\":4}, " + "{\"column0\":6,\"column1\":6}, " + "{\"column0\":7}" + "]", dstVector.toString()); } } } TestIndexSorter.java000066400000000000000000000161541515451615200335730ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link IndexSorter}. */ public class TestIndexSorter { private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testIndexSort() { try (IntVector vec = new IntVector("", allocator)) { vec.allocateNew(10); vec.setValueCount(10); // fill data to sort ValueVectorDataPopulator.setVector(vec, 11, 8, 33, 10, 12, 17, null, 23, 35, 2); // sort the index IndexSorter indexSorter = new IndexSorter<>(); DefaultVectorComparators.IntComparator intComparator = new DefaultVectorComparators.IntComparator(); intComparator.attachVector(vec); IntVector indices = new IntVector("", allocator); indices.setValueCount(10); indexSorter.sort(vec, indices, intComparator); int[] expected = new int[] {6, 9, 1, 3, 0, 4, 5, 7, 2, 8}; for (int i = 0; i < expected.length; i++) { assertTrue(!indices.isNull(i)); assertEquals(expected[i], indices.get(i)); } indices.close(); } } /** * Tests the worst case for quick sort. It may cause stack overflow if the algorithm is * implemented as a recursive algorithm. */ @Test public void testSortLargeIncreasingInt() { final int vectorLength = 20000; try (IntVector vec = new IntVector("", allocator)) { vec.allocateNew(vectorLength); vec.setValueCount(vectorLength); // fill data to sort for (int i = 0; i < vectorLength; i++) { vec.set(i, i); } // sort the vector IndexSorter indexSorter = new IndexSorter<>(); DefaultVectorComparators.IntComparator intComparator = new DefaultVectorComparators.IntComparator(); intComparator.attachVector(vec); try (IntVector indices = new IntVector("", allocator)) { indices.setValueCount(vectorLength); indexSorter.sort(vec, indices, intComparator); for (int i = 0; i < vectorLength; i++) { assertTrue(!indices.isNull(i)); assertEquals(i, indices.get(i)); } } } } @Test public void testChoosePivot() { final int vectorLength = 100; try (IntVector vec = new IntVector("vector", allocator); IntVector indices = new IntVector("indices", allocator)) { vec.allocateNew(vectorLength); indices.allocateNew(vectorLength); // the vector is sorted, so the pivot should be in the middle for (int i = 0; i < vectorLength; i++) { vec.set(i, i * 100); indices.set(i, i); } vec.setValueCount(vectorLength); indices.setValueCount(vectorLength); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); // setup internal data structures comparator.attachVector(vec); int low = 5; int high = 6; assertTrue(high - low + 1 < FixedWidthInPlaceVectorSorter.STOP_CHOOSING_PIVOT_THRESHOLD); // the range is small enough, so the pivot is simply selected as the low value int pivotIndex = IndexSorter.choosePivot(low, high, indices, comparator); assertEquals(pivotIndex, low); assertEquals(pivotIndex, indices.get(low)); low = 30; high = 80; assertTrue(high - low + 1 >= FixedWidthInPlaceVectorSorter.STOP_CHOOSING_PIVOT_THRESHOLD); // the range is large enough, so the median is selected as the pivot pivotIndex = IndexSorter.choosePivot(low, high, indices, comparator); assertEquals(pivotIndex, (low + high) / 2); assertEquals(pivotIndex, indices.get(low)); } } /** Evaluates choosing pivot for all possible permutations of 3 numbers. */ @Test public void testChoosePivotAllPermutes() { try (IntVector vec = new IntVector("vector", allocator); IntVector indices = new IntVector("indices", allocator)) { vec.allocateNew(); indices.allocateNew(); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); // setup internal data structures comparator.attachVector(vec); int low = 0; int high = 2; // test all the 6 permutations of 3 numbers ValueVectorDataPopulator.setVector(indices, 0, 1, 2); ValueVectorDataPopulator.setVector(vec, 11, 22, 33); int pivotIndex = IndexSorter.choosePivot(low, high, indices, comparator); assertEquals(1, pivotIndex); assertEquals(1, indices.get(low)); ValueVectorDataPopulator.setVector(indices, 0, 1, 2); ValueVectorDataPopulator.setVector(vec, 11, 33, 22); pivotIndex = IndexSorter.choosePivot(low, high, indices, comparator); assertEquals(2, pivotIndex); assertEquals(2, indices.get(low)); ValueVectorDataPopulator.setVector(indices, 0, 1, 2); ValueVectorDataPopulator.setVector(vec, 22, 11, 33); pivotIndex = IndexSorter.choosePivot(low, high, indices, comparator); assertEquals(0, pivotIndex); assertEquals(0, indices.get(low)); ValueVectorDataPopulator.setVector(indices, 0, 1, 2); ValueVectorDataPopulator.setVector(vec, 22, 33, 11); pivotIndex = IndexSorter.choosePivot(low, high, indices, comparator); assertEquals(0, pivotIndex); assertEquals(0, indices.get(low)); ValueVectorDataPopulator.setVector(indices, 0, 1, 2); ValueVectorDataPopulator.setVector(vec, 33, 11, 22); pivotIndex = IndexSorter.choosePivot(low, high, indices, comparator); assertEquals(2, pivotIndex); assertEquals(2, indices.get(low)); ValueVectorDataPopulator.setVector(indices, 0, 1, 2); ValueVectorDataPopulator.setVector(vec, 33, 22, 11); pivotIndex = IndexSorter.choosePivot(low, high, indices, comparator); assertEquals(1, pivotIndex); assertEquals(1, indices.get(low)); } } } TestInsertionSorter.java000066400000000000000000000106441515451615200344740ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link InsertionSorter}. */ public class TestInsertionSorter { private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } private static final int VECTOR_LENGTH = 10; private void testSortIntVectorRange(int start, int end, int[] expected) { try (IntVector vector = new IntVector("vector", allocator); IntVector buffer = new IntVector("buffer", allocator)) { buffer.allocateNew(1); ValueVectorDataPopulator.setVector(vector, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0); assertEquals(VECTOR_LENGTH, vector.getValueCount()); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vector); InsertionSorter.insertionSort(vector, start, end, comparator, buffer); assertEquals(VECTOR_LENGTH, expected.length); for (int i = 0; i < VECTOR_LENGTH; i++) { assertFalse(vector.isNull(i)); assertEquals(expected[i], vector.get(i)); } } } @Test public void testSortIntVector() { testSortIntVectorRange(2, 5, new int[] {9, 8, 4, 5, 6, 7, 3, 2, 1, 0}); testSortIntVectorRange(3, 7, new int[] {9, 8, 7, 2, 3, 4, 5, 6, 1, 0}); testSortIntVectorRange(3, 4, new int[] {9, 8, 7, 5, 6, 4, 3, 2, 1, 0}); testSortIntVectorRange(7, 7, new int[] {9, 8, 7, 6, 5, 4, 3, 2, 1, 0}); testSortIntVectorRange(0, 5, new int[] {4, 5, 6, 7, 8, 9, 3, 2, 1, 0}); testSortIntVectorRange(8, 9, new int[] {9, 8, 7, 6, 5, 4, 3, 2, 0, 1}); testSortIntVectorRange(0, 9, new int[] {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); } private void testSortIndicesRange(int start, int end, int[] expectedIndices) { try (IntVector vector = new IntVector("vector", allocator); IntVector indices = new IntVector("indices", allocator)) { ValueVectorDataPopulator.setVector(vector, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0); ValueVectorDataPopulator.setVector(indices, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9); assertEquals(VECTOR_LENGTH, vector.getValueCount()); assertEquals(VECTOR_LENGTH, indices.getValueCount()); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vector); comparator.attachVector(vector); InsertionSorter.insertionSort(indices, start, end, comparator); // verify results assertEquals(VECTOR_LENGTH, expectedIndices.length); for (int i = 0; i < VECTOR_LENGTH; i++) { assertFalse(indices.isNull(i)); assertEquals(expectedIndices[i], indices.get(i)); } } } @Test public void testSortIndices() { testSortIndicesRange(2, 5, new int[] {0, 1, 5, 4, 3, 2, 6, 7, 8, 9}); testSortIndicesRange(3, 7, new int[] {0, 1, 2, 7, 6, 5, 4, 3, 8, 9}); testSortIndicesRange(3, 4, new int[] {0, 1, 2, 4, 3, 5, 6, 7, 8, 9}); testSortIndicesRange(7, 7, new int[] {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); testSortIndicesRange(0, 5, new int[] {5, 4, 3, 2, 1, 0, 6, 7, 8, 9}); testSortIndicesRange(8, 9, new int[] {0, 1, 2, 3, 4, 5, 6, 7, 9, 8}); testSortIndicesRange(0, 9, new int[] {9, 8, 7, 6, 5, 4, 3, 2, 1, 0}); } } TestOffHeapIntStack.java000066400000000000000000000037311515451615200342730ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link OffHeapIntStack}. */ public class TestOffHeapIntStack { private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testPushPop() { try (OffHeapIntStack stack = new OffHeapIntStack(allocator)) { assertTrue(stack.isEmpty()); final int elemCount = 500; for (int i = 0; i < elemCount; i++) { stack.push(i); assertEquals(i, stack.getTop()); } assertEquals(elemCount, stack.getCount()); for (int i = 0; i < elemCount; i++) { assertEquals(elemCount - i - 1, stack.getTop()); assertEquals(elemCount - i - 1, stack.pop()); } assertTrue(stack.isEmpty()); } } } TestOutOfPlaceVectorSorter.java000066400000000000000000000024471515451615200357100ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import java.util.ArrayList; import java.util.List; import java.util.stream.Stream; import org.junit.jupiter.params.provider.Arguments; /** Test cases for out-of-place sorters. */ public abstract class TestOutOfPlaceVectorSorter { public static Stream getParameter() { List args = new ArrayList<>(); for (boolean generalSorter : new boolean[] {false, true}) { args.add(Arguments.of(generalSorter)); } return args.stream(); } } TestSortingUtil.java000066400000000000000000000136171515451615200336110ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import static org.junit.jupiter.api.Assertions.assertEquals; import java.lang.reflect.Array; import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.Random; import java.util.function.BiConsumer; import java.util.function.Supplier; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.SmallIntVector; import org.apache.arrow.vector.TinyIntVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.testing.RandomDataGenerator; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; /** Utilities for sorting related utilities. */ public class TestSortingUtil { static final Random random = new Random(0); static final DataGenerator TINY_INT_GENERATOR = new DataGenerator<>( RandomDataGenerator.TINY_INT_GENERATOR, (vector, array) -> ValueVectorDataPopulator.setVector(vector, array), Byte.class); static final DataGenerator SMALL_INT_GENERATOR = new DataGenerator<>( RandomDataGenerator.SMALL_INT_GENERATOR, (vector, array) -> ValueVectorDataPopulator.setVector(vector, array), Short.class); static final DataGenerator INT_GENERATOR = new DataGenerator<>( RandomDataGenerator.INT_GENERATOR, (vector, array) -> ValueVectorDataPopulator.setVector(vector, array), Integer.class); static final DataGenerator LONG_GENERATOR = new DataGenerator<>( RandomDataGenerator.LONG_GENERATOR, (vector, array) -> ValueVectorDataPopulator.setVector(vector, array), Long.class); static final DataGenerator FLOAT_GENERATOR = new DataGenerator<>( RandomDataGenerator.FLOAT_GENERATOR, (vector, array) -> ValueVectorDataPopulator.setVector(vector, array), Float.class); static final DataGenerator DOUBLE_GENERATOR = new DataGenerator<>( RandomDataGenerator.DOUBLE_GENERATOR, (vector, array) -> ValueVectorDataPopulator.setVector(vector, array), Double.class); static final DataGenerator STRING_GENERATOR = new DataGenerator<>( () -> { int strLength = random.nextInt(20) + 1; return generateRandomString(strLength); }, (vector, array) -> ValueVectorDataPopulator.setVector(vector, array), String.class); private TestSortingUtil() {} /** Verify that a vector is equal to an array. */ public static void verifyResults(V vector, U[] expected) { assertEquals(vector.getValueCount(), expected.length); for (int i = 0; i < expected.length; i++) { assertEquals(vector.getObject(i), expected[i]); } } /** Sort an array with null values come first. */ public static > void sortArray(U[] array) { Arrays.sort( array, (a, b) -> { if (a == null || b == null) { if (a == null && b == null) { return 0; } // exactly one is null if (a == null) { return -1; } else { return 1; } } return a.compareTo(b); }); } /** Generate a string with alphabetic characters only. */ static String generateRandomString(int length) { byte[] str = new byte[length]; final int lower = 'a'; final int upper = 'z'; for (int i = 0; i < length; i++) { // make r non-negative int r = random.nextInt() & Integer.MAX_VALUE; str[i] = (byte) (r % (upper - lower + 1) + lower); } return new String(str, StandardCharsets.UTF_8); } /** * Utility to generate data for testing. * * @param vector type. * @param data element type. */ static class DataGenerator> { final Supplier dataGenerator; final BiConsumer vectorPopulator; final Class clazz; DataGenerator(Supplier dataGenerator, BiConsumer vectorPopulator, Class clazz) { this.dataGenerator = dataGenerator; this.vectorPopulator = vectorPopulator; this.clazz = clazz; } /** * Populate the vector according to the specified parameters. * * @param vector the vector to populate. * @param length vector length. * @param nullFraction the fraction of null values. * @return An array with the same data as the vector. */ U[] populate(V vector, int length, double nullFraction) { U[] array = (U[]) Array.newInstance(clazz, length); for (int i = 0; i < length; i++) { double r = Math.random(); U value = r < nullFraction ? null : dataGenerator.get(); array[i] = value; } vectorPopulator.accept(vector, array); return array; } } } TestStableVectorComparator.java000066400000000000000000000134061515451615200357470ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.charset.StandardCharsets; import java.util.Objects; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.VarCharVector; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link StableVectorComparator}. */ public class TestStableVectorComparator { private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testCompare() { try (VarCharVector vec = new VarCharVector("", allocator)) { vec.allocateNew(100, 5); vec.setValueCount(10); // fill data to sort vec.set(0, "ba".getBytes(StandardCharsets.UTF_8)); vec.set(1, "abc".getBytes(StandardCharsets.UTF_8)); vec.set(2, "aa".getBytes(StandardCharsets.UTF_8)); vec.set(3, "abc".getBytes(StandardCharsets.UTF_8)); vec.set(4, "a".getBytes(StandardCharsets.UTF_8)); VectorValueComparator comparator = new TestVarCharSorter(); VectorValueComparator stableComparator = new StableVectorComparator<>(comparator); stableComparator.attachVector(vec); assertTrue(stableComparator.compare(0, 1) > 0); assertTrue(stableComparator.compare(1, 2) < 0); assertTrue(stableComparator.compare(2, 3) < 0); assertTrue(stableComparator.compare(1, 3) < 0); assertTrue(stableComparator.compare(3, 1) > 0); assertEquals(0, stableComparator.compare(3, 3)); } } @Test public void testStableSortString() { try (VarCharVector vec = new VarCharVector("", allocator)) { vec.allocateNew(100, 10); vec.setValueCount(10); // fill data to sort vec.set(0, "a".getBytes(StandardCharsets.UTF_8)); vec.set(1, "abc".getBytes(StandardCharsets.UTF_8)); vec.set(2, "aa".getBytes(StandardCharsets.UTF_8)); vec.set(3, "a1".getBytes(StandardCharsets.UTF_8)); vec.set(4, "abcdefg".getBytes(StandardCharsets.UTF_8)); vec.set(5, "accc".getBytes(StandardCharsets.UTF_8)); vec.set(6, "afds".getBytes(StandardCharsets.UTF_8)); vec.set(7, "0".getBytes(StandardCharsets.UTF_8)); vec.set(8, "01".getBytes(StandardCharsets.UTF_8)); vec.set(9, "0c".getBytes(StandardCharsets.UTF_8)); // sort the vector VariableWidthOutOfPlaceVectorSorter sorter = new VariableWidthOutOfPlaceVectorSorter(); VectorValueComparator comparator = new TestVarCharSorter(); VectorValueComparator stableComparator = new StableVectorComparator<>(comparator); try (VarCharVector sortedVec = (VarCharVector) vec.getField().getFieldType().createNewSingleVector("", allocator, null)) { sortedVec.allocateNew(vec.getByteCapacity(), vec.getValueCount()); sortedVec.setLastSet(vec.getValueCount() - 1); sortedVec.setValueCount(vec.getValueCount()); sorter.sortOutOfPlace(vec, sortedVec, stableComparator); // verify results // the results are stable assertEquals( "0", new String(Objects.requireNonNull(sortedVec.get(0)), StandardCharsets.UTF_8)); assertEquals( "01", new String(Objects.requireNonNull(sortedVec.get(1)), StandardCharsets.UTF_8)); assertEquals( "0c", new String(Objects.requireNonNull(sortedVec.get(2)), StandardCharsets.UTF_8)); assertEquals( "a", new String(Objects.requireNonNull(sortedVec.get(3)), StandardCharsets.UTF_8)); assertEquals( "abc", new String(Objects.requireNonNull(sortedVec.get(4)), StandardCharsets.UTF_8)); assertEquals( "aa", new String(Objects.requireNonNull(sortedVec.get(5)), StandardCharsets.UTF_8)); assertEquals( "a1", new String(Objects.requireNonNull(sortedVec.get(6)), StandardCharsets.UTF_8)); assertEquals( "abcdefg", new String(Objects.requireNonNull(sortedVec.get(7)), StandardCharsets.UTF_8)); assertEquals( "accc", new String(Objects.requireNonNull(sortedVec.get(8)), StandardCharsets.UTF_8)); assertEquals( "afds", new String(Objects.requireNonNull(sortedVec.get(9)), StandardCharsets.UTF_8)); } } } /** Utility comparator that compares varchars by the first character. */ private static class TestVarCharSorter extends VectorValueComparator { @Override public int compareNotNull(int index1, int index2) { byte b1 = vector1.get(index1)[0]; byte b2 = vector2.get(index2)[0]; return b1 - b2; } @Override public VectorValueComparator createNew() { return new TestVarCharSorter(); } } } TestVariableWidthOutOfPlaceVectorSorter.java000066400000000000000000000111321515451615200403450ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.charset.StandardCharsets; import java.util.Objects; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.VarCharVector; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.MethodSource; /** Test cases for {@link VariableWidthOutOfPlaceVectorSorter}. */ public class TestVariableWidthOutOfPlaceVectorSorter extends TestOutOfPlaceVectorSorter { private BufferAllocator allocator; OutOfPlaceVectorSorter getSorter(boolean generalSorter) { return generalSorter ? new GeneralOutOfPlaceVectorSorter<>() : new VariableWidthOutOfPlaceVectorSorter(); } @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @ParameterizedTest @MethodSource("getParameter") public void testSortString(boolean generalSorter) { try (VarCharVector vec = new VarCharVector("", allocator)) { vec.allocateNew(100, 10); vec.setValueCount(10); // fill data to sort vec.set(0, "hello".getBytes(StandardCharsets.UTF_8)); vec.set(1, "abc".getBytes(StandardCharsets.UTF_8)); vec.setNull(2); vec.set(3, "world".getBytes(StandardCharsets.UTF_8)); vec.set(4, "12".getBytes(StandardCharsets.UTF_8)); vec.set(5, "dictionary".getBytes(StandardCharsets.UTF_8)); vec.setNull(6); vec.set(7, "hello".getBytes(StandardCharsets.UTF_8)); vec.set(8, "good".getBytes(StandardCharsets.UTF_8)); vec.set(9, "yes".getBytes(StandardCharsets.UTF_8)); // sort the vector OutOfPlaceVectorSorter sorter = getSorter(generalSorter); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vec); VarCharVector sortedVec = (VarCharVector) vec.getField().getFieldType().createNewSingleVector("", allocator, null); sortedVec.allocateNew(vec.getByteCapacity(), vec.getValueCount()); sortedVec.setLastSet(vec.getValueCount() - 1); sortedVec.setValueCount(vec.getValueCount()); sorter.sortOutOfPlace(vec, sortedVec, comparator); // verify results assertEquals(vec.getValueCount(), sortedVec.getValueCount()); assertEquals(vec.getByteCapacity(), sortedVec.getByteCapacity()); assertEquals(vec.getLastSet(), sortedVec.getLastSet()); assertTrue(sortedVec.isNull(0)); assertTrue(sortedVec.isNull(1)); assertEquals( "12", new String(Objects.requireNonNull(sortedVec.get(2)), StandardCharsets.UTF_8)); assertEquals( "abc", new String(Objects.requireNonNull(sortedVec.get(3)), StandardCharsets.UTF_8)); assertEquals( "dictionary", new String(Objects.requireNonNull(sortedVec.get(4)), StandardCharsets.UTF_8)); assertEquals( "good", new String(Objects.requireNonNull(sortedVec.get(5)), StandardCharsets.UTF_8)); assertEquals( "hello", new String(Objects.requireNonNull(sortedVec.get(6)), StandardCharsets.UTF_8)); assertEquals( "hello", new String(Objects.requireNonNull(sortedVec.get(7)), StandardCharsets.UTF_8)); assertEquals( "world", new String(Objects.requireNonNull(sortedVec.get(8)), StandardCharsets.UTF_8)); assertEquals( "yes", new String(Objects.requireNonNull(sortedVec.get(9)), StandardCharsets.UTF_8)); sortedVec.close(); } } } TestVariableWidthSorting.java000066400000000000000000000125541515451615200354200ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/java/org/apache/arrow/algorithm/sort/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.sort; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; import java.util.Comparator; import java.util.List; import java.util.function.Function; import java.util.stream.Stream; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.util.Text; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; /** Test sorting variable width vectors with random data. */ public class TestVariableWidthSorting> { static final int[] VECTOR_LENGTHS = new int[] {2, 5, 10, 50, 100, 1000, 3000}; static final double[] NULL_FRACTIONS = {0, 0.1, 0.3, 0.5, 0.7, 0.9, 1}; private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(Integer.MAX_VALUE); } @AfterEach public void shutdown() { allocator.close(); } @ParameterizedTest @MethodSource("getParameters") public void testSort( int length, double nullFraction, Function vectorGenerator, TestSortingUtil.DataGenerator dataGenerator) { sortOutOfPlace(length, nullFraction, vectorGenerator, dataGenerator); } void sortOutOfPlace( int length, double nullFraction, Function vectorGenerator, TestSortingUtil.DataGenerator dataGenerator) { try (V vector = vectorGenerator.apply(allocator)) { U[] array = dataGenerator.populate(vector, length, nullFraction); Arrays.sort(array, (Comparator) new StringComparator()); // sort the vector VariableWidthOutOfPlaceVectorSorter sorter = new VariableWidthOutOfPlaceVectorSorter(); VectorValueComparator comparator = DefaultVectorComparators.createDefaultComparator(vector); try (V sortedVec = (V) vector.getField().getFieldType().createNewSingleVector("", allocator, null)) { int dataSize = vector.getOffsetBuffer().getInt(vector.getValueCount() * 4L); sortedVec.allocateNew(dataSize, vector.getValueCount()); sortedVec.setValueCount(vector.getValueCount()); sorter.sortOutOfPlace(vector, sortedVec, comparator); // verify results verifyResults(sortedVec, (String[]) array); } } } public static Stream getParameters() { List params = new ArrayList<>(); for (int length : VECTOR_LENGTHS) { for (double nullFrac : NULL_FRACTIONS) { params.add( Arguments.of( length, nullFrac, (Function) allocator -> new VarCharVector("vector", allocator), TestSortingUtil.STRING_GENERATOR)); } } return params.stream(); } /** Verify results as byte arrays. */ public static void verifyResults(V vector, String[] expected) { assertEquals(vector.getValueCount(), expected.length); for (int i = 0; i < expected.length; i++) { if (expected[i] == null) { assertTrue(vector.isNull(i)); } else { assertArrayEquals( ((Text) vector.getObject(i)).getBytes(), expected[i].getBytes(StandardCharsets.UTF_8)); } } } /** * String comparator with the same behavior as that of {@link * DefaultVectorComparators.VariableWidthComparator}. */ static class StringComparator implements Comparator { @Override public int compare(String str1, String str2) { if (str1 == null || str2 == null) { if (str1 == null && str2 == null) { return 0; } return str1 == null ? -1 : 1; } byte[] bytes1 = str1.getBytes(StandardCharsets.UTF_8); byte[] bytes2 = str2.getBytes(StandardCharsets.UTF_8); for (int i = 0; i < bytes1.length && i < bytes2.length; i++) { if (bytes1[i] != bytes2[i]) { return (bytes1[i] & 0xff) < (bytes2[i] & 0xff) ? -1 : 1; } } return bytes1.length - bytes2.length; } } } apache-arrow-java-19.0.0/algorithm/src/test/resources/000077500000000000000000000000001515451615200226445ustar00rootroot00000000000000apache-arrow-java-19.0.0/algorithm/src/test/resources/logback.xml000066400000000000000000000025221515451615200247710ustar00rootroot00000000000000 %d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n apache-arrow-java-19.0.0/api-changes.md000066400000000000000000000037521515451615200176060ustar00rootroot00000000000000 # Arrow Java API Changes This document tracks behavior changes to Java APIs, as listed below. - **[ARROW-5973](https://issues.apache.org/jira/browse/ARROW-5973)**: * **Start date**: 2019/07/18 * **Resolve date**: 2019/07/20 * **Brief description**: The semantics of the get methods for [VarCharVector](./vector/scr/main/org/apache/arrow/vector/VarCharVector.java), [VarBinaryVector](./vector/scr/main/org/apache/arrow/vector/VarBinaryVector.java), and [FixedSizeBinaryVector](./vector/scr/main/org/apache/arrow/vector/FixedSizeBinaryVector.java) changes. In the past, if the validity bit is clear, the methods throw throws an IllegalStateException when NULL_CHECKING_ENABLED is set, or returns an empty object when the flag is not set. Now, the get methods return a null if the validity bit is clear. - **[ARROW-5842](https://issues.apache.org/jira/browse/ARROW-5842)**: * **Start date**: 2019/07/04 * **Resolve date**: 2019/07/11 * **Brief description**: The semantics of lastSet member in class [ListVector](./vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java) changes. In the past, it refers to the next index that will be set. After the change it points to the last index that is actually set. apache-arrow-java-19.0.0/arrow-format/000077500000000000000000000000001515451615200175165ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-format/File.fbs000066400000000000000000000027501515451615200210750ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. include "Schema.fbs"; namespace org.apache.arrow.flatbuf; /// ---------------------------------------------------------------------- /// Arrow File metadata /// table Footer { version: org.apache.arrow.flatbuf.MetadataVersion; schema: org.apache.arrow.flatbuf.Schema; dictionaries: [ Block ]; recordBatches: [ Block ]; /// User-defined metadata custom_metadata: [ KeyValue ]; } struct Block { /// Index to the start of the RecordBlock (note this is past the Message header) offset: long; /// Length of the metadata metaDataLength: int; /// Length of the data (this is aligned so there can be a gap between this and /// the metadata). bodyLength: long; } root_type Footer; apache-arrow-java-19.0.0/arrow-format/Flight.proto000066400000000000000000000501141515451615200220210ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at *

    * http://www.apache.org/licenses/LICENSE-2.0 *

    * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ syntax = "proto3"; import "google/protobuf/timestamp.proto"; option java_package = "org.apache.arrow.flight.impl"; option go_package = "github.com/apache/arrow-go/arrow/flight/gen/flight"; option csharp_namespace = "Apache.Arrow.Flight.Protocol"; package arrow.flight.protocol; /* * A flight service is an endpoint for retrieving or storing Arrow data. A * flight service can expose one or more predefined endpoints that can be * accessed using the Arrow Flight Protocol. Additionally, a flight service * can expose a set of actions that are available. */ service FlightService { /* * Handshake between client and server. Depending on the server, the * handshake may be required to determine the token that should be used for * future operations. Both request and response are streams to allow multiple * round-trips depending on auth mechanism. */ rpc Handshake(stream HandshakeRequest) returns (stream HandshakeResponse) {} /* * Get a list of available streams given a particular criteria. Most flight * services will expose one or more streams that are readily available for * retrieval. This api allows listing the streams available for * consumption. A user can also provide a criteria. The criteria can limit * the subset of streams that can be listed via this interface. Each flight * service allows its own definition of how to consume criteria. */ rpc ListFlights(Criteria) returns (stream FlightInfo) {} /* * For a given FlightDescriptor, get information about how the flight can be * consumed. This is a useful interface if the consumer of the interface * already can identify the specific flight to consume. This interface can * also allow a consumer to generate a flight stream through a specified * descriptor. For example, a flight descriptor might be something that * includes a SQL statement or a Pickled Python operation that will be * executed. In those cases, the descriptor will not be previously available * within the list of available streams provided by ListFlights but will be * available for consumption for the duration defined by the specific flight * service. */ rpc GetFlightInfo(FlightDescriptor) returns (FlightInfo) {} /* * For a given FlightDescriptor, start a query and get information * to poll its execution status. This is a useful interface if the * query may be a long-running query. The first PollFlightInfo call * should return as quickly as possible. (GetFlightInfo doesn't * return until the query is complete.) * * A client can consume any available results before * the query is completed. See PollInfo.info for details. * * A client can poll the updated query status by calling * PollFlightInfo() with PollInfo.flight_descriptor. A server * should not respond until the result would be different from last * time. That way, the client can "long poll" for updates * without constantly making requests. Clients can set a short timeout * to avoid blocking calls if desired. * * A client can't use PollInfo.flight_descriptor after * PollInfo.expiration_time passes. A server might not accept the * retry descriptor anymore and the query may be cancelled. * * A client may use the CancelFlightInfo action with * PollInfo.info to cancel the running query. */ rpc PollFlightInfo(FlightDescriptor) returns (PollInfo) {} /* * For a given FlightDescriptor, get the Schema as described in Schema.fbs::Schema * This is used when a consumer needs the Schema of flight stream. Similar to * GetFlightInfo this interface may generate a new flight that was not previously * available in ListFlights. */ rpc GetSchema(FlightDescriptor) returns (SchemaResult) {} /* * Retrieve a single stream associated with a particular descriptor * associated with the referenced ticket. A Flight can be composed of one or * more streams where each stream can be retrieved using a separate opaque * ticket that the flight service uses for managing a collection of streams. */ rpc DoGet(Ticket) returns (stream FlightData) {} /* * Push a stream to the flight service associated with a particular * flight stream. This allows a client of a flight service to upload a stream * of data. Depending on the particular flight service, a client consumer * could be allowed to upload a single stream per descriptor or an unlimited * number. In the latter, the service might implement a 'seal' action that * can be applied to a descriptor once all streams are uploaded. */ rpc DoPut(stream FlightData) returns (stream PutResult) {} /* * Open a bidirectional data channel for a given descriptor. This * allows clients to send and receive arbitrary Arrow data and * application-specific metadata in a single logical stream. In * contrast to DoGet/DoPut, this is more suited for clients * offloading computation (rather than storage) to a Flight service. */ rpc DoExchange(stream FlightData) returns (stream FlightData) {} /* * Flight services can support an arbitrary number of simple actions in * addition to the possible ListFlights, GetFlightInfo, DoGet, DoPut * operations that are potentially available. DoAction allows a flight client * to do a specific action against a flight service. An action includes * opaque request and response objects that are specific to the type action * being undertaken. */ rpc DoAction(Action) returns (stream Result) {} /* * A flight service exposes all of the available action types that it has * along with descriptions. This allows different flight consumers to * understand the capabilities of the flight service. */ rpc ListActions(Empty) returns (stream ActionType) {} } /* * The request that a client provides to a server on handshake. */ message HandshakeRequest { /* * A defined protocol version */ uint64 protocol_version = 1; /* * Arbitrary auth/handshake info. */ bytes payload = 2; } message HandshakeResponse { /* * A defined protocol version */ uint64 protocol_version = 1; /* * Arbitrary auth/handshake info. */ bytes payload = 2; } /* * A message for doing simple auth. */ message BasicAuth { string username = 2; string password = 3; } message Empty {} /* * Describes an available action, including both the name used for execution * along with a short description of the purpose of the action. */ message ActionType { string type = 1; string description = 2; } /* * A service specific expression that can be used to return a limited set * of available Arrow Flight streams. */ message Criteria { bytes expression = 1; } /* * An opaque action specific for the service. */ message Action { string type = 1; bytes body = 2; } /* * An opaque result returned after executing an action. */ message Result { bytes body = 1; } /* * Wrap the result of a getSchema call */ message SchemaResult { // The schema of the dataset in its IPC form: // 4 bytes - an optional IPC_CONTINUATION_TOKEN prefix // 4 bytes - the byte length of the payload // a flatbuffer Message whose header is the Schema bytes schema = 1; } /* * The name or tag for a Flight. May be used as a way to retrieve or generate * a flight or be used to expose a set of previously defined flights. */ message FlightDescriptor { /* * Describes what type of descriptor is defined. */ enum DescriptorType { // Protobuf pattern, not used. UNKNOWN = 0; /* * A named path that identifies a dataset. A path is composed of a string * or list of strings describing a particular dataset. This is conceptually * similar to a path inside a filesystem. */ PATH = 1; /* * An opaque command to generate a dataset. */ CMD = 2; } DescriptorType type = 1; /* * Opaque value used to express a command. Should only be defined when * type = CMD. */ bytes cmd = 2; /* * List of strings identifying a particular dataset. Should only be defined * when type = PATH. */ repeated string path = 3; } /* * The access coordinates for retrieval of a dataset. With a FlightInfo, a * consumer is able to determine how to retrieve a dataset. */ message FlightInfo { // The schema of the dataset in its IPC form: // 4 bytes - an optional IPC_CONTINUATION_TOKEN prefix // 4 bytes - the byte length of the payload // a flatbuffer Message whose header is the Schema bytes schema = 1; /* * The descriptor associated with this info. */ FlightDescriptor flight_descriptor = 2; /* * A list of endpoints associated with the flight. To consume the * whole flight, all endpoints (and hence all Tickets) must be * consumed. Endpoints can be consumed in any order. * * In other words, an application can use multiple endpoints to * represent partitioned data. * * If the returned data has an ordering, an application can use * "FlightInfo.ordered = true" or should return the all data in a * single endpoint. Otherwise, there is no ordering defined on * endpoints or the data within. * * A client can read ordered data by reading data from returned * endpoints, in order, from front to back. * * Note that a client may ignore "FlightInfo.ordered = true". If an * ordering is important for an application, an application must * choose one of them: * * * An application requires that all clients must read data in * returned endpoints order. * * An application must return the all data in a single endpoint. */ repeated FlightEndpoint endpoint = 3; // Set these to -1 if unknown. int64 total_records = 4; int64 total_bytes = 5; /* * FlightEndpoints are in the same order as the data. */ bool ordered = 6; /* * Application-defined metadata. * * There is no inherent or required relationship between this * and the app_metadata fields in the FlightEndpoints or resulting * FlightData messages. Since this metadata is application-defined, * a given application could define there to be a relationship, * but there is none required by the spec. */ bytes app_metadata = 7; } /* * The information to process a long-running query. */ message PollInfo { /* * The currently available results. * * If "flight_descriptor" is not specified, the query is complete * and "info" specifies all results. Otherwise, "info" contains * partial query results. * * Note that each PollInfo response contains a complete * FlightInfo (not just the delta between the previous and current * FlightInfo). * * Subsequent PollInfo responses may only append new endpoints to * info. * * Clients can begin fetching results via DoGet(Ticket) with the * ticket in the info before the query is * completed. FlightInfo.ordered is also valid. */ FlightInfo info = 1; /* * The descriptor the client should use on the next try. * If unset, the query is complete. */ FlightDescriptor flight_descriptor = 2; /* * Query progress. If known, must be in [0.0, 1.0] but need not be * monotonic or nondecreasing. If unknown, do not set. */ optional double progress = 3; /* * Expiration time for this request. After this passes, the server * might not accept the retry descriptor anymore (and the query may * be cancelled). This may be updated on a call to PollFlightInfo. */ google.protobuf.Timestamp expiration_time = 4; } /* * The request of the CancelFlightInfo action. * * The request should be stored in Action.body. */ message CancelFlightInfoRequest { FlightInfo info = 1; } /* * The result of a cancel operation. * * This is used by CancelFlightInfoResult.status. */ enum CancelStatus { // The cancellation status is unknown. Servers should avoid using // this value (send a NOT_FOUND error if the requested query is // not known). Clients can retry the request. CANCEL_STATUS_UNSPECIFIED = 0; // The cancellation request is complete. Subsequent requests with // the same payload may return CANCELLED or a NOT_FOUND error. CANCEL_STATUS_CANCELLED = 1; // The cancellation request is in progress. The client may retry // the cancellation request. CANCEL_STATUS_CANCELLING = 2; // The query is not cancellable. The client should not retry the // cancellation request. CANCEL_STATUS_NOT_CANCELLABLE = 3; } /* * The result of the CancelFlightInfo action. * * The result should be stored in Result.body. */ message CancelFlightInfoResult { CancelStatus status = 1; } /* * An opaque identifier that the service can use to retrieve a particular * portion of a stream. * * Tickets are meant to be single use. It is an error/application-defined * behavior to reuse a ticket. */ message Ticket { bytes ticket = 1; } /* * A location where a Flight service will accept retrieval of a particular * stream given a ticket. */ message Location { string uri = 1; } /* * A particular stream or split associated with a flight. */ message FlightEndpoint { /* * Token used to retrieve this stream. */ Ticket ticket = 1; /* * A list of URIs where this ticket can be redeemed via DoGet(). * * If the list is empty, the expectation is that the ticket can only * be redeemed on the current service where the ticket was * generated. * * If the list is not empty, the expectation is that the ticket can be * redeemed at any of the locations, and that the data returned will be * equivalent. In this case, the ticket may only be redeemed at one of the * given locations, and not (necessarily) on the current service. If one * of the given locations is "arrow-flight-reuse-connection://?", the * client may redeem the ticket on the service where the ticket was * generated (i.e., the same as above), in addition to the other * locations. (This URI was chosen to maximize compatibility, as 'scheme:' * or 'scheme://' are not accepted by Java's java.net.URI.) * * In other words, an application can use multiple locations to * represent redundant and/or load balanced services. */ repeated Location location = 2; /* * Expiration time of this stream. If present, clients may assume * they can retry DoGet requests. Otherwise, it is * application-defined whether DoGet requests may be retried. */ google.protobuf.Timestamp expiration_time = 3; /* * Application-defined metadata. * * There is no inherent or required relationship between this * and the app_metadata fields in the FlightInfo or resulting * FlightData messages. Since this metadata is application-defined, * a given application could define there to be a relationship, * but there is none required by the spec. */ bytes app_metadata = 4; } /* * The request of the RenewFlightEndpoint action. * * The request should be stored in Action.body. */ message RenewFlightEndpointRequest { FlightEndpoint endpoint = 1; } /* * A batch of Arrow data as part of a stream of batches. */ message FlightData { /* * The descriptor of the data. This is only relevant when a client is * starting a new DoPut stream. */ FlightDescriptor flight_descriptor = 1; /* * Header for message data as described in Message.fbs::Message. */ bytes data_header = 2; /* * Application-defined metadata. */ bytes app_metadata = 3; /* * The actual batch of Arrow data. Preferably handled with minimal-copies * coming last in the definition to help with sidecar patterns (it is * expected that some implementations will fetch this field off the wire * with specialized code to avoid extra memory copies). */ bytes data_body = 1000; } /** * The response message associated with the submission of a DoPut. */ message PutResult { bytes app_metadata = 1; } /* * EXPERIMENTAL: Union of possible value types for a Session Option to be set to. * * By convention, an attempt to set a valueless SessionOptionValue should * attempt to unset or clear the named option value on the server. */ message SessionOptionValue { message StringListValue { repeated string values = 1; } oneof option_value { string string_value = 1; bool bool_value = 2; sfixed64 int64_value = 3; double double_value = 4; StringListValue string_list_value = 5; } } /* * EXPERIMENTAL: A request to set session options for an existing or new (implicit) * server session. * * Sessions are persisted and referenced via a transport-level state management, typically * RFC 6265 HTTP cookies when using an HTTP transport. The suggested cookie name or state * context key is 'arrow_flight_session_id', although implementations may freely choose their * own name. * * Session creation (if one does not already exist) is implied by this RPC request, however * server implementations may choose to initiate a session that also contains client-provided * session options at any other time, e.g. on authentication, or when any other call is made * and the server wishes to use a session to persist any state (or lack thereof). */ message SetSessionOptionsRequest { map session_options = 1; } /* * EXPERIMENTAL: The results (individually) of setting a set of session options. * * Option names should only be present in the response if they were not successfully * set on the server; that is, a response without an Error for a name provided in the * SetSessionOptionsRequest implies that the named option value was set successfully. */ message SetSessionOptionsResult { enum ErrorValue { // Protobuf deserialization fallback value: The status is unknown or unrecognized. // Servers should avoid using this value. The request may be retried by the client. UNSPECIFIED = 0; // The given session option name is invalid. INVALID_NAME = 1; // The session option value or type is invalid. INVALID_VALUE = 2; // The session option cannot be set. ERROR = 3; } message Error { ErrorValue value = 1; } map errors = 1; } /* * EXPERIMENTAL: A request to access the session options for the current server session. * * The existing session is referenced via a cookie header or similar (see * SetSessionOptionsRequest above); it is an error to make this request with a missing, * invalid, or expired session cookie header or other implementation-defined session * reference token. */ message GetSessionOptionsRequest { } /* * EXPERIMENTAL: The result containing the current server session options. */ message GetSessionOptionsResult { map session_options = 1; } /* * Request message for the "Close Session" action. * * The exiting session is referenced via a cookie header. */ message CloseSessionRequest { } /* * The result of closing a session. */ message CloseSessionResult { enum Status { // Protobuf deserialization fallback value: The session close status is unknown or // not recognized. Servers should avoid using this value (send a NOT_FOUND error if // the requested session is not known or expired). Clients can retry the request. UNSPECIFIED = 0; // The session close request is complete. Subsequent requests with // the same session produce a NOT_FOUND error. CLOSED = 1; // The session close request is in progress. The client may retry // the close request. CLOSING = 2; // The session is not closeable. The client should not retry the // close request. NOT_CLOSEABLE = 3; } Status status = 1; } apache-arrow-java-19.0.0/arrow-format/FlightSql.proto000066400000000000000000002306421515451615200225070ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at *

    * http://www.apache.org/licenses/LICENSE-2.0 *

    * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ syntax = "proto3"; import "google/protobuf/descriptor.proto"; option java_package = "org.apache.arrow.flight.sql.impl"; option go_package = "github.com/apache/arrow-go/arrow/flight/gen/flight"; package arrow.flight.protocol.sql; /* * Represents a metadata request. Used in the command member of FlightDescriptor * for the following RPC calls: * - GetSchema: return the Arrow schema of the query. * - GetFlightInfo: execute the metadata request. * * The returned Arrow schema will be: * < * info_name: uint32 not null, * value: dense_union< * string_value: utf8, * bool_value: bool, * bigint_value: int64, * int32_bitmask: int32, * string_list: list * int32_to_int32_list_map: map> * > * where there is one row per requested piece of metadata information. */ message CommandGetSqlInfo { /* * Values are modelled after ODBC's SQLGetInfo() function. This information is intended to provide * Flight SQL clients with basic, SQL syntax and SQL functions related information. * More information types can be added in future releases. * E.g. more SQL syntax support types, scalar functions support, type conversion support etc. * * Note that the set of metadata may expand. * * Initially, Flight SQL will support the following information types: * - Server Information - Range [0-500) * - Syntax Information - Range [500-1000) * Range [0-10,000) is reserved for defaults (see SqlInfo enum for default options). * Custom options should start at 10,000. * * If omitted, then all metadata will be retrieved. * Flight SQL Servers may choose to include additional metadata above and beyond the specified set, however they must * at least return the specified set. IDs ranging from 0 to 10,000 (exclusive) are reserved for future use. * If additional metadata is included, the metadata IDs should start from 10,000. */ repeated uint32 info = 1; } // Options for CommandGetSqlInfo. enum SqlInfo { // Server Information [0-500): Provides basic information about the Flight SQL Server. // Retrieves a UTF-8 string with the name of the Flight SQL Server. FLIGHT_SQL_SERVER_NAME = 0; // Retrieves a UTF-8 string with the native version of the Flight SQL Server. FLIGHT_SQL_SERVER_VERSION = 1; // Retrieves a UTF-8 string with the Arrow format version of the Flight SQL Server. FLIGHT_SQL_SERVER_ARROW_VERSION = 2; /* * Retrieves a boolean value indicating whether the Flight SQL Server is read only. * * Returns: * - false: if read-write * - true: if read only */ FLIGHT_SQL_SERVER_READ_ONLY = 3; /* * Retrieves a boolean value indicating whether the Flight SQL Server supports executing * SQL queries. * * Note that the absence of this info (as opposed to a false value) does not necessarily * mean that SQL is not supported, as this property was not originally defined. */ FLIGHT_SQL_SERVER_SQL = 4; /* * Retrieves a boolean value indicating whether the Flight SQL Server supports executing * Substrait plans. */ FLIGHT_SQL_SERVER_SUBSTRAIT = 5; /* * Retrieves a string value indicating the minimum supported Substrait version, or null * if Substrait is not supported. */ FLIGHT_SQL_SERVER_SUBSTRAIT_MIN_VERSION = 6; /* * Retrieves a string value indicating the maximum supported Substrait version, or null * if Substrait is not supported. */ FLIGHT_SQL_SERVER_SUBSTRAIT_MAX_VERSION = 7; /* * Retrieves an int32 indicating whether the Flight SQL Server supports the * BeginTransaction/EndTransaction/BeginSavepoint/EndSavepoint actions. * * Even if this is not supported, the database may still support explicit "BEGIN * TRANSACTION"/"COMMIT" SQL statements (see SQL_TRANSACTIONS_SUPPORTED); this property * is only about whether the server implements the Flight SQL API endpoints. * * The possible values are listed in `SqlSupportedTransaction`. */ FLIGHT_SQL_SERVER_TRANSACTION = 8; /* * Retrieves a boolean value indicating whether the Flight SQL Server supports explicit * query cancellation (the CancelQuery action). */ FLIGHT_SQL_SERVER_CANCEL = 9; /* * Retrieves a boolean value indicating whether the Flight SQL Server supports executing * bulk ingestion. */ FLIGHT_SQL_SERVER_BULK_INGESTION = 10; /* * Retrieves a boolean value indicating whether transactions are supported for bulk ingestion. If not, invoking * the method commit in the context of a bulk ingestion is a noop, and the isolation level is * `arrow.flight.protocol.sql.SqlTransactionIsolationLevel.TRANSACTION_NONE`. * * Returns: * - false: if bulk ingestion transactions are unsupported; * - true: if bulk ingestion transactions are supported. */ FLIGHT_SQL_SERVER_INGEST_TRANSACTIONS_SUPPORTED = 11; /* * Retrieves an int32 indicating the timeout (in milliseconds) for prepared statement handles. * * If 0, there is no timeout. Servers should reset the timeout when the handle is used in a command. */ FLIGHT_SQL_SERVER_STATEMENT_TIMEOUT = 100; /* * Retrieves an int32 indicating the timeout (in milliseconds) for transactions, since transactions are not tied to a connection. * * If 0, there is no timeout. Servers should reset the timeout when the handle is used in a command. */ FLIGHT_SQL_SERVER_TRANSACTION_TIMEOUT = 101; // SQL Syntax Information [500-1000): provides information about SQL syntax supported by the Flight SQL Server. /* * Retrieves a boolean value indicating whether the Flight SQL Server supports CREATE and DROP of catalogs. * * Returns: * - false: if it doesn't support CREATE and DROP of catalogs. * - true: if it supports CREATE and DROP of catalogs. */ SQL_DDL_CATALOG = 500; /* * Retrieves a boolean value indicating whether the Flight SQL Server supports CREATE and DROP of schemas. * * Returns: * - false: if it doesn't support CREATE and DROP of schemas. * - true: if it supports CREATE and DROP of schemas. */ SQL_DDL_SCHEMA = 501; /* * Indicates whether the Flight SQL Server supports CREATE and DROP of tables. * * Returns: * - false: if it doesn't support CREATE and DROP of tables. * - true: if it supports CREATE and DROP of tables. */ SQL_DDL_TABLE = 502; /* * Retrieves a int32 ordinal representing the case sensitivity of catalog, table, schema and table names. * * The possible values are listed in `arrow.flight.protocol.sql.SqlSupportedCaseSensitivity`. */ SQL_IDENTIFIER_CASE = 503; // Retrieves a UTF-8 string with the supported character(s) used to surround a delimited identifier. SQL_IDENTIFIER_QUOTE_CHAR = 504; /* * Retrieves a int32 describing the case sensitivity of quoted identifiers. * * The possible values are listed in `arrow.flight.protocol.sql.SqlSupportedCaseSensitivity`. */ SQL_QUOTED_IDENTIFIER_CASE = 505; /* * Retrieves a boolean value indicating whether all tables are selectable. * * Returns: * - false: if not all tables are selectable or if none are; * - true: if all tables are selectable. */ SQL_ALL_TABLES_ARE_SELECTABLE = 506; /* * Retrieves the null ordering. * * Returns a int32 ordinal for the null ordering being used, as described in * `arrow.flight.protocol.sql.SqlNullOrdering`. */ SQL_NULL_ORDERING = 507; // Retrieves a UTF-8 string list with values of the supported keywords. SQL_KEYWORDS = 508; // Retrieves a UTF-8 string list with values of the supported numeric functions. SQL_NUMERIC_FUNCTIONS = 509; // Retrieves a UTF-8 string list with values of the supported string functions. SQL_STRING_FUNCTIONS = 510; // Retrieves a UTF-8 string list with values of the supported system functions. SQL_SYSTEM_FUNCTIONS = 511; // Retrieves a UTF-8 string list with values of the supported datetime functions. SQL_DATETIME_FUNCTIONS = 512; /* * Retrieves the UTF-8 string that can be used to escape wildcard characters. * This is the string that can be used to escape '_' or '%' in the catalog search parameters that are a pattern * (and therefore use one of the wildcard characters). * The '_' character represents any single character; the '%' character represents any sequence of zero or more * characters. */ SQL_SEARCH_STRING_ESCAPE = 513; /* * Retrieves a UTF-8 string with all the "extra" characters that can be used in unquoted identifier names * (those beyond a-z, A-Z, 0-9 and _). */ SQL_EXTRA_NAME_CHARACTERS = 514; /* * Retrieves a boolean value indicating whether column aliasing is supported. * If so, the SQL AS clause can be used to provide names for computed columns or to provide alias names for columns * as required. * * Returns: * - false: if column aliasing is unsupported; * - true: if column aliasing is supported. */ SQL_SUPPORTS_COLUMN_ALIASING = 515; /* * Retrieves a boolean value indicating whether concatenations between null and non-null values being * null are supported. * * - Returns: * - false: if concatenations between null and non-null values being null are unsupported; * - true: if concatenations between null and non-null values being null are supported. */ SQL_NULL_PLUS_NULL_IS_NULL = 516; /* * Retrieves a map where the key is the type to convert from and the value is a list with the types to convert to, * indicating the supported conversions. Each key and each item on the list value is a value to a predefined type on * SqlSupportsConvert enum. * The returned map will be: map> */ SQL_SUPPORTS_CONVERT = 517; /* * Retrieves a boolean value indicating whether, when table correlation names are supported, * they are restricted to being different from the names of the tables. * * Returns: * - false: if table correlation names are unsupported; * - true: if table correlation names are supported. */ SQL_SUPPORTS_TABLE_CORRELATION_NAMES = 518; /* * Retrieves a boolean value indicating whether, when table correlation names are supported, * they are restricted to being different from the names of the tables. * * Returns: * - false: if different table correlation names are unsupported; * - true: if different table correlation names are supported */ SQL_SUPPORTS_DIFFERENT_TABLE_CORRELATION_NAMES = 519; /* * Retrieves a boolean value indicating whether expressions in ORDER BY lists are supported. * * Returns: * - false: if expressions in ORDER BY are unsupported; * - true: if expressions in ORDER BY are supported; */ SQL_SUPPORTS_EXPRESSIONS_IN_ORDER_BY = 520; /* * Retrieves a boolean value indicating whether using a column that is not in the SELECT statement in a GROUP BY * clause is supported. * * Returns: * - false: if using a column that is not in the SELECT statement in a GROUP BY clause is unsupported; * - true: if using a column that is not in the SELECT statement in a GROUP BY clause is supported. */ SQL_SUPPORTS_ORDER_BY_UNRELATED = 521; /* * Retrieves the supported GROUP BY commands; * * Returns an int32 bitmask value representing the supported commands. * The returned bitmask should be parsed in order to retrieve the supported commands. * * For instance: * - return 0 (\b0) => [] (GROUP BY is unsupported); * - return 1 (\b1) => [SQL_GROUP_BY_UNRELATED]; * - return 2 (\b10) => [SQL_GROUP_BY_BEYOND_SELECT]; * - return 3 (\b11) => [SQL_GROUP_BY_UNRELATED, SQL_GROUP_BY_BEYOND_SELECT]. * Valid GROUP BY types are described under `arrow.flight.protocol.sql.SqlSupportedGroupBy`. */ SQL_SUPPORTED_GROUP_BY = 522; /* * Retrieves a boolean value indicating whether specifying a LIKE escape clause is supported. * * Returns: * - false: if specifying a LIKE escape clause is unsupported; * - true: if specifying a LIKE escape clause is supported. */ SQL_SUPPORTS_LIKE_ESCAPE_CLAUSE = 523; /* * Retrieves a boolean value indicating whether columns may be defined as non-nullable. * * Returns: * - false: if columns cannot be defined as non-nullable; * - true: if columns may be defined as non-nullable. */ SQL_SUPPORTS_NON_NULLABLE_COLUMNS = 524; /* * Retrieves the supported SQL grammar level as per the ODBC specification. * * Returns an int32 bitmask value representing the supported SQL grammar level. * The returned bitmask should be parsed in order to retrieve the supported grammar levels. * * For instance: * - return 0 (\b0) => [] (SQL grammar is unsupported); * - return 1 (\b1) => [SQL_MINIMUM_GRAMMAR]; * - return 2 (\b10) => [SQL_CORE_GRAMMAR]; * - return 3 (\b11) => [SQL_MINIMUM_GRAMMAR, SQL_CORE_GRAMMAR]; * - return 4 (\b100) => [SQL_EXTENDED_GRAMMAR]; * - return 5 (\b101) => [SQL_MINIMUM_GRAMMAR, SQL_EXTENDED_GRAMMAR]; * - return 6 (\b110) => [SQL_CORE_GRAMMAR, SQL_EXTENDED_GRAMMAR]; * - return 7 (\b111) => [SQL_MINIMUM_GRAMMAR, SQL_CORE_GRAMMAR, SQL_EXTENDED_GRAMMAR]. * Valid SQL grammar levels are described under `arrow.flight.protocol.sql.SupportedSqlGrammar`. */ SQL_SUPPORTED_GRAMMAR = 525; /* * Retrieves the supported ANSI92 SQL grammar level. * * Returns an int32 bitmask value representing the supported ANSI92 SQL grammar level. * The returned bitmask should be parsed in order to retrieve the supported commands. * * For instance: * - return 0 (\b0) => [] (ANSI92 SQL grammar is unsupported); * - return 1 (\b1) => [ANSI92_ENTRY_SQL]; * - return 2 (\b10) => [ANSI92_INTERMEDIATE_SQL]; * - return 3 (\b11) => [ANSI92_ENTRY_SQL, ANSI92_INTERMEDIATE_SQL]; * - return 4 (\b100) => [ANSI92_FULL_SQL]; * - return 5 (\b101) => [ANSI92_ENTRY_SQL, ANSI92_FULL_SQL]; * - return 6 (\b110) => [ANSI92_INTERMEDIATE_SQL, ANSI92_FULL_SQL]; * - return 7 (\b111) => [ANSI92_ENTRY_SQL, ANSI92_INTERMEDIATE_SQL, ANSI92_FULL_SQL]. * Valid ANSI92 SQL grammar levels are described under `arrow.flight.protocol.sql.SupportedAnsi92SqlGrammarLevel`. */ SQL_ANSI92_SUPPORTED_LEVEL = 526; /* * Retrieves a boolean value indicating whether the SQL Integrity Enhancement Facility is supported. * * Returns: * - false: if the SQL Integrity Enhancement Facility is supported; * - true: if the SQL Integrity Enhancement Facility is supported. */ SQL_SUPPORTS_INTEGRITY_ENHANCEMENT_FACILITY = 527; /* * Retrieves the support level for SQL OUTER JOINs. * * Returns a int32 ordinal for the SQL ordering being used, as described in * `arrow.flight.protocol.sql.SqlOuterJoinsSupportLevel`. */ SQL_OUTER_JOINS_SUPPORT_LEVEL = 528; // Retrieves a UTF-8 string with the preferred term for "schema". SQL_SCHEMA_TERM = 529; // Retrieves a UTF-8 string with the preferred term for "procedure". SQL_PROCEDURE_TERM = 530; /* * Retrieves a UTF-8 string with the preferred term for "catalog". * If a empty string is returned its assumed that the server does NOT supports catalogs. */ SQL_CATALOG_TERM = 531; /* * Retrieves a boolean value indicating whether a catalog appears at the start of a fully qualified table name. * * - false: if a catalog does not appear at the start of a fully qualified table name; * - true: if a catalog appears at the start of a fully qualified table name. */ SQL_CATALOG_AT_START = 532; /* * Retrieves the supported actions for a SQL schema. * * Returns an int32 bitmask value representing the supported actions for a SQL schema. * The returned bitmask should be parsed in order to retrieve the supported actions for a SQL schema. * * For instance: * - return 0 (\b0) => [] (no supported actions for SQL schema); * - return 1 (\b1) => [SQL_ELEMENT_IN_PROCEDURE_CALLS]; * - return 2 (\b10) => [SQL_ELEMENT_IN_INDEX_DEFINITIONS]; * - return 3 (\b11) => [SQL_ELEMENT_IN_PROCEDURE_CALLS, SQL_ELEMENT_IN_INDEX_DEFINITIONS]; * - return 4 (\b100) => [SQL_ELEMENT_IN_PRIVILEGE_DEFINITIONS]; * - return 5 (\b101) => [SQL_ELEMENT_IN_PROCEDURE_CALLS, SQL_ELEMENT_IN_PRIVILEGE_DEFINITIONS]; * - return 6 (\b110) => [SQL_ELEMENT_IN_INDEX_DEFINITIONS, SQL_ELEMENT_IN_PRIVILEGE_DEFINITIONS]; * - return 7 (\b111) => [SQL_ELEMENT_IN_PROCEDURE_CALLS, SQL_ELEMENT_IN_INDEX_DEFINITIONS, SQL_ELEMENT_IN_PRIVILEGE_DEFINITIONS]. * Valid actions for a SQL schema described under `arrow.flight.protocol.sql.SqlSupportedElementActions`. */ SQL_SCHEMAS_SUPPORTED_ACTIONS = 533; /* * Retrieves the supported actions for a SQL schema. * * Returns an int32 bitmask value representing the supported actions for a SQL catalog. * The returned bitmask should be parsed in order to retrieve the supported actions for a SQL catalog. * * For instance: * - return 0 (\b0) => [] (no supported actions for SQL catalog); * - return 1 (\b1) => [SQL_ELEMENT_IN_PROCEDURE_CALLS]; * - return 2 (\b10) => [SQL_ELEMENT_IN_INDEX_DEFINITIONS]; * - return 3 (\b11) => [SQL_ELEMENT_IN_PROCEDURE_CALLS, SQL_ELEMENT_IN_INDEX_DEFINITIONS]; * - return 4 (\b100) => [SQL_ELEMENT_IN_PRIVILEGE_DEFINITIONS]; * - return 5 (\b101) => [SQL_ELEMENT_IN_PROCEDURE_CALLS, SQL_ELEMENT_IN_PRIVILEGE_DEFINITIONS]; * - return 6 (\b110) => [SQL_ELEMENT_IN_INDEX_DEFINITIONS, SQL_ELEMENT_IN_PRIVILEGE_DEFINITIONS]; * - return 7 (\b111) => [SQL_ELEMENT_IN_PROCEDURE_CALLS, SQL_ELEMENT_IN_INDEX_DEFINITIONS, SQL_ELEMENT_IN_PRIVILEGE_DEFINITIONS]. * Valid actions for a SQL catalog are described under `arrow.flight.protocol.sql.SqlSupportedElementActions`. */ SQL_CATALOGS_SUPPORTED_ACTIONS = 534; /* * Retrieves the supported SQL positioned commands. * * Returns an int32 bitmask value representing the supported SQL positioned commands. * The returned bitmask should be parsed in order to retrieve the supported SQL positioned commands. * * For instance: * - return 0 (\b0) => [] (no supported SQL positioned commands); * - return 1 (\b1) => [SQL_POSITIONED_DELETE]; * - return 2 (\b10) => [SQL_POSITIONED_UPDATE]; * - return 3 (\b11) => [SQL_POSITIONED_DELETE, SQL_POSITIONED_UPDATE]. * Valid SQL positioned commands are described under `arrow.flight.protocol.sql.SqlSupportedPositionedCommands`. */ SQL_SUPPORTED_POSITIONED_COMMANDS = 535; /* * Retrieves a boolean value indicating whether SELECT FOR UPDATE statements are supported. * * Returns: * - false: if SELECT FOR UPDATE statements are unsupported; * - true: if SELECT FOR UPDATE statements are supported. */ SQL_SELECT_FOR_UPDATE_SUPPORTED = 536; /* * Retrieves a boolean value indicating whether stored procedure calls that use the stored procedure escape syntax * are supported. * * Returns: * - false: if stored procedure calls that use the stored procedure escape syntax are unsupported; * - true: if stored procedure calls that use the stored procedure escape syntax are supported. */ SQL_STORED_PROCEDURES_SUPPORTED = 537; /* * Retrieves the supported SQL subqueries. * * Returns an int32 bitmask value representing the supported SQL subqueries. * The returned bitmask should be parsed in order to retrieve the supported SQL subqueries. * * For instance: * - return 0 (\b0) => [] (no supported SQL subqueries); * - return 1 (\b1) => [SQL_SUBQUERIES_IN_COMPARISONS]; * - return 2 (\b10) => [SQL_SUBQUERIES_IN_EXISTS]; * - return 3 (\b11) => [SQL_SUBQUERIES_IN_COMPARISONS, SQL_SUBQUERIES_IN_EXISTS]; * - return 4 (\b100) => [SQL_SUBQUERIES_IN_INS]; * - return 5 (\b101) => [SQL_SUBQUERIES_IN_COMPARISONS, SQL_SUBQUERIES_IN_INS]; * - return 6 (\b110) => [SQL_SUBQUERIES_IN_INS, SQL_SUBQUERIES_IN_EXISTS]; * - return 7 (\b111) => [SQL_SUBQUERIES_IN_COMPARISONS, SQL_SUBQUERIES_IN_EXISTS, SQL_SUBQUERIES_IN_INS]; * - return 8 (\b1000) => [SQL_SUBQUERIES_IN_QUANTIFIEDS]; * - return 9 (\b1001) => [SQL_SUBQUERIES_IN_COMPARISONS, SQL_SUBQUERIES_IN_QUANTIFIEDS]; * - return 10 (\b1010) => [SQL_SUBQUERIES_IN_EXISTS, SQL_SUBQUERIES_IN_QUANTIFIEDS]; * - return 11 (\b1011) => [SQL_SUBQUERIES_IN_COMPARISONS, SQL_SUBQUERIES_IN_EXISTS, SQL_SUBQUERIES_IN_QUANTIFIEDS]; * - return 12 (\b1100) => [SQL_SUBQUERIES_IN_INS, SQL_SUBQUERIES_IN_QUANTIFIEDS]; * - return 13 (\b1101) => [SQL_SUBQUERIES_IN_COMPARISONS, SQL_SUBQUERIES_IN_INS, SQL_SUBQUERIES_IN_QUANTIFIEDS]; * - return 14 (\b1110) => [SQL_SUBQUERIES_IN_EXISTS, SQL_SUBQUERIES_IN_INS, SQL_SUBQUERIES_IN_QUANTIFIEDS]; * - return 15 (\b1111) => [SQL_SUBQUERIES_IN_COMPARISONS, SQL_SUBQUERIES_IN_EXISTS, SQL_SUBQUERIES_IN_INS, SQL_SUBQUERIES_IN_QUANTIFIEDS]; * - ... * Valid SQL subqueries are described under `arrow.flight.protocol.sql.SqlSupportedSubqueries`. */ SQL_SUPPORTED_SUBQUERIES = 538; /* * Retrieves a boolean value indicating whether correlated subqueries are supported. * * Returns: * - false: if correlated subqueries are unsupported; * - true: if correlated subqueries are supported. */ SQL_CORRELATED_SUBQUERIES_SUPPORTED = 539; /* * Retrieves the supported SQL UNIONs. * * Returns an int32 bitmask value representing the supported SQL UNIONs. * The returned bitmask should be parsed in order to retrieve the supported SQL UNIONs. * * For instance: * - return 0 (\b0) => [] (no supported SQL positioned commands); * - return 1 (\b1) => [SQL_UNION]; * - return 2 (\b10) => [SQL_UNION_ALL]; * - return 3 (\b11) => [SQL_UNION, SQL_UNION_ALL]. * Valid SQL positioned commands are described under `arrow.flight.protocol.sql.SqlSupportedUnions`. */ SQL_SUPPORTED_UNIONS = 540; // Retrieves a int64 value representing the maximum number of hex characters allowed in an inline binary literal. SQL_MAX_BINARY_LITERAL_LENGTH = 541; // Retrieves a int64 value representing the maximum number of characters allowed for a character literal. SQL_MAX_CHAR_LITERAL_LENGTH = 542; // Retrieves a int64 value representing the maximum number of characters allowed for a column name. SQL_MAX_COLUMN_NAME_LENGTH = 543; // Retrieves a int64 value representing the maximum number of columns allowed in a GROUP BY clause. SQL_MAX_COLUMNS_IN_GROUP_BY = 544; // Retrieves a int64 value representing the maximum number of columns allowed in an index. SQL_MAX_COLUMNS_IN_INDEX = 545; // Retrieves a int64 value representing the maximum number of columns allowed in an ORDER BY clause. SQL_MAX_COLUMNS_IN_ORDER_BY = 546; // Retrieves a int64 value representing the maximum number of columns allowed in a SELECT list. SQL_MAX_COLUMNS_IN_SELECT = 547; // Retrieves a int64 value representing the maximum number of columns allowed in a table. SQL_MAX_COLUMNS_IN_TABLE = 548; // Retrieves a int64 value representing the maximum number of concurrent connections possible. SQL_MAX_CONNECTIONS = 549; // Retrieves a int64 value the maximum number of characters allowed in a cursor name. SQL_MAX_CURSOR_NAME_LENGTH = 550; /* * Retrieves a int64 value representing the maximum number of bytes allowed for an index, * including all of the parts of the index. */ SQL_MAX_INDEX_LENGTH = 551; // Retrieves a int64 value representing the maximum number of characters allowed in a schema name. SQL_DB_SCHEMA_NAME_LENGTH = 552; // Retrieves a int64 value representing the maximum number of characters allowed in a procedure name. SQL_MAX_PROCEDURE_NAME_LENGTH = 553; // Retrieves a int64 value representing the maximum number of characters allowed in a catalog name. SQL_MAX_CATALOG_NAME_LENGTH = 554; // Retrieves a int64 value representing the maximum number of bytes allowed in a single row. SQL_MAX_ROW_SIZE = 555; /* * Retrieves a boolean indicating whether the return value for the JDBC method getMaxRowSize includes the SQL * data types LONGVARCHAR and LONGVARBINARY. * * Returns: * - false: if return value for the JDBC method getMaxRowSize does * not include the SQL data types LONGVARCHAR and LONGVARBINARY; * - true: if return value for the JDBC method getMaxRowSize includes * the SQL data types LONGVARCHAR and LONGVARBINARY. */ SQL_MAX_ROW_SIZE_INCLUDES_BLOBS = 556; /* * Retrieves a int64 value representing the maximum number of characters allowed for an SQL statement; * a result of 0 (zero) means that there is no limit or the limit is not known. */ SQL_MAX_STATEMENT_LENGTH = 557; // Retrieves a int64 value representing the maximum number of active statements that can be open at the same time. SQL_MAX_STATEMENTS = 558; // Retrieves a int64 value representing the maximum number of characters allowed in a table name. SQL_MAX_TABLE_NAME_LENGTH = 559; // Retrieves a int64 value representing the maximum number of tables allowed in a SELECT statement. SQL_MAX_TABLES_IN_SELECT = 560; // Retrieves a int64 value representing the maximum number of characters allowed in a user name. SQL_MAX_USERNAME_LENGTH = 561; /* * Retrieves this database's default transaction isolation level as described in * `arrow.flight.protocol.sql.SqlTransactionIsolationLevel`. * * Returns a int32 ordinal for the SQL transaction isolation level. */ SQL_DEFAULT_TRANSACTION_ISOLATION = 562; /* * Retrieves a boolean value indicating whether transactions are supported. If not, invoking the method commit is a * noop, and the isolation level is `arrow.flight.protocol.sql.SqlTransactionIsolationLevel.TRANSACTION_NONE`. * * Returns: * - false: if transactions are unsupported; * - true: if transactions are supported. */ SQL_TRANSACTIONS_SUPPORTED = 563; /* * Retrieves the supported transactions isolation levels. * * Returns an int32 bitmask value representing the supported transactions isolation levels. * The returned bitmask should be parsed in order to retrieve the supported transactions isolation levels. * * For instance: * - return 0 (\b0) => [] (no supported SQL transactions isolation levels); * - return 1 (\b1) => [SQL_TRANSACTION_NONE]; * - return 2 (\b10) => [SQL_TRANSACTION_READ_UNCOMMITTED]; * - return 3 (\b11) => [SQL_TRANSACTION_NONE, SQL_TRANSACTION_READ_UNCOMMITTED]; * - return 4 (\b100) => [SQL_TRANSACTION_REPEATABLE_READ]; * - return 5 (\b101) => [SQL_TRANSACTION_NONE, SQL_TRANSACTION_REPEATABLE_READ]; * - return 6 (\b110) => [SQL_TRANSACTION_READ_UNCOMMITTED, SQL_TRANSACTION_REPEATABLE_READ]; * - return 7 (\b111) => [SQL_TRANSACTION_NONE, SQL_TRANSACTION_READ_UNCOMMITTED, SQL_TRANSACTION_REPEATABLE_READ]; * - return 8 (\b1000) => [SQL_TRANSACTION_REPEATABLE_READ]; * - return 9 (\b1001) => [SQL_TRANSACTION_NONE, SQL_TRANSACTION_REPEATABLE_READ]; * - return 10 (\b1010) => [SQL_TRANSACTION_READ_UNCOMMITTED, SQL_TRANSACTION_REPEATABLE_READ]; * - return 11 (\b1011) => [SQL_TRANSACTION_NONE, SQL_TRANSACTION_READ_UNCOMMITTED, SQL_TRANSACTION_REPEATABLE_READ]; * - return 12 (\b1100) => [SQL_TRANSACTION_REPEATABLE_READ, SQL_TRANSACTION_REPEATABLE_READ]; * - return 13 (\b1101) => [SQL_TRANSACTION_NONE, SQL_TRANSACTION_REPEATABLE_READ, SQL_TRANSACTION_REPEATABLE_READ]; * - return 14 (\b1110) => [SQL_TRANSACTION_READ_UNCOMMITTED, SQL_TRANSACTION_REPEATABLE_READ, SQL_TRANSACTION_REPEATABLE_READ]; * - return 15 (\b1111) => [SQL_TRANSACTION_NONE, SQL_TRANSACTION_READ_UNCOMMITTED, SQL_TRANSACTION_REPEATABLE_READ, SQL_TRANSACTION_REPEATABLE_READ]; * - return 16 (\b10000) => [SQL_TRANSACTION_SERIALIZABLE]; * - ... * Valid SQL positioned commands are described under `arrow.flight.protocol.sql.SqlTransactionIsolationLevel`. */ SQL_SUPPORTED_TRANSACTIONS_ISOLATION_LEVELS = 564; /* * Retrieves a boolean value indicating whether a data definition statement within a transaction forces * the transaction to commit. * * Returns: * - false: if a data definition statement within a transaction does not force the transaction to commit; * - true: if a data definition statement within a transaction forces the transaction to commit. */ SQL_DATA_DEFINITION_CAUSES_TRANSACTION_COMMIT = 565; /* * Retrieves a boolean value indicating whether a data definition statement within a transaction is ignored. * * Returns: * - false: if a data definition statement within a transaction is taken into account; * - true: a data definition statement within a transaction is ignored. */ SQL_DATA_DEFINITIONS_IN_TRANSACTIONS_IGNORED = 566; /* * Retrieves an int32 bitmask value representing the supported result set types. * The returned bitmask should be parsed in order to retrieve the supported result set types. * * For instance: * - return 0 (\b0) => [] (no supported result set types); * - return 1 (\b1) => [SQL_RESULT_SET_TYPE_UNSPECIFIED]; * - return 2 (\b10) => [SQL_RESULT_SET_TYPE_FORWARD_ONLY]; * - return 3 (\b11) => [SQL_RESULT_SET_TYPE_UNSPECIFIED, SQL_RESULT_SET_TYPE_FORWARD_ONLY]; * - return 4 (\b100) => [SQL_RESULT_SET_TYPE_SCROLL_INSENSITIVE]; * - return 5 (\b101) => [SQL_RESULT_SET_TYPE_UNSPECIFIED, SQL_RESULT_SET_TYPE_SCROLL_INSENSITIVE]; * - return 6 (\b110) => [SQL_RESULT_SET_TYPE_FORWARD_ONLY, SQL_RESULT_SET_TYPE_SCROLL_INSENSITIVE]; * - return 7 (\b111) => [SQL_RESULT_SET_TYPE_UNSPECIFIED, SQL_RESULT_SET_TYPE_FORWARD_ONLY, SQL_RESULT_SET_TYPE_SCROLL_INSENSITIVE]; * - return 8 (\b1000) => [SQL_RESULT_SET_TYPE_SCROLL_SENSITIVE]; * - ... * Valid result set types are described under `arrow.flight.protocol.sql.SqlSupportedResultSetType`. */ SQL_SUPPORTED_RESULT_SET_TYPES = 567; /* * Returns an int32 bitmask value concurrency types supported for * `arrow.flight.protocol.sql.SqlSupportedResultSetType.SQL_RESULT_SET_TYPE_UNSPECIFIED`. * * For instance: * - return 0 (\b0) => [] (no supported concurrency types for this result set type) * - return 1 (\b1) => [SQL_RESULT_SET_CONCURRENCY_UNSPECIFIED] * - return 2 (\b10) => [SQL_RESULT_SET_CONCURRENCY_READ_ONLY] * - return 3 (\b11) => [SQL_RESULT_SET_CONCURRENCY_UNSPECIFIED, SQL_RESULT_SET_CONCURRENCY_READ_ONLY] * - return 4 (\b100) => [SQL_RESULT_SET_CONCURRENCY_UPDATABLE] * - return 5 (\b101) => [SQL_RESULT_SET_CONCURRENCY_UNSPECIFIED, SQL_RESULT_SET_CONCURRENCY_UPDATABLE] * - return 6 (\b110) => [SQL_RESULT_SET_CONCURRENCY_READ_ONLY, SQL_RESULT_SET_CONCURRENCY_UPDATABLE] * - return 7 (\b111) => [SQL_RESULT_SET_CONCURRENCY_UNSPECIFIED, SQL_RESULT_SET_CONCURRENCY_READ_ONLY, SQL_RESULT_SET_CONCURRENCY_UPDATABLE] * Valid result set types are described under `arrow.flight.protocol.sql.SqlSupportedResultSetConcurrency`. */ SQL_SUPPORTED_CONCURRENCIES_FOR_RESULT_SET_UNSPECIFIED = 568; /* * Returns an int32 bitmask value concurrency types supported for * `arrow.flight.protocol.sql.SqlSupportedResultSetType.SQL_RESULT_SET_TYPE_FORWARD_ONLY`. * * For instance: * - return 0 (\b0) => [] (no supported concurrency types for this result set type) * - return 1 (\b1) => [SQL_RESULT_SET_CONCURRENCY_UNSPECIFIED] * - return 2 (\b10) => [SQL_RESULT_SET_CONCURRENCY_READ_ONLY] * - return 3 (\b11) => [SQL_RESULT_SET_CONCURRENCY_UNSPECIFIED, SQL_RESULT_SET_CONCURRENCY_READ_ONLY] * - return 4 (\b100) => [SQL_RESULT_SET_CONCURRENCY_UPDATABLE] * - return 5 (\b101) => [SQL_RESULT_SET_CONCURRENCY_UNSPECIFIED, SQL_RESULT_SET_CONCURRENCY_UPDATABLE] * - return 6 (\b110) => [SQL_RESULT_SET_CONCURRENCY_READ_ONLY, SQL_RESULT_SET_CONCURRENCY_UPDATABLE] * - return 7 (\b111) => [SQL_RESULT_SET_CONCURRENCY_UNSPECIFIED, SQL_RESULT_SET_CONCURRENCY_READ_ONLY, SQL_RESULT_SET_CONCURRENCY_UPDATABLE] * Valid result set types are described under `arrow.flight.protocol.sql.SqlSupportedResultSetConcurrency`. */ SQL_SUPPORTED_CONCURRENCIES_FOR_RESULT_SET_FORWARD_ONLY = 569; /* * Returns an int32 bitmask value concurrency types supported for * `arrow.flight.protocol.sql.SqlSupportedResultSetType.SQL_RESULT_SET_TYPE_SCROLL_SENSITIVE`. * * For instance: * - return 0 (\b0) => [] (no supported concurrency types for this result set type) * - return 1 (\b1) => [SQL_RESULT_SET_CONCURRENCY_UNSPECIFIED] * - return 2 (\b10) => [SQL_RESULT_SET_CONCURRENCY_READ_ONLY] * - return 3 (\b11) => [SQL_RESULT_SET_CONCURRENCY_UNSPECIFIED, SQL_RESULT_SET_CONCURRENCY_READ_ONLY] * - return 4 (\b100) => [SQL_RESULT_SET_CONCURRENCY_UPDATABLE] * - return 5 (\b101) => [SQL_RESULT_SET_CONCURRENCY_UNSPECIFIED, SQL_RESULT_SET_CONCURRENCY_UPDATABLE] * - return 6 (\b110) => [SQL_RESULT_SET_CONCURRENCY_READ_ONLY, SQL_RESULT_SET_CONCURRENCY_UPDATABLE] * - return 7 (\b111) => [SQL_RESULT_SET_CONCURRENCY_UNSPECIFIED, SQL_RESULT_SET_CONCURRENCY_READ_ONLY, SQL_RESULT_SET_CONCURRENCY_UPDATABLE] * Valid result set types are described under `arrow.flight.protocol.sql.SqlSupportedResultSetConcurrency`. */ SQL_SUPPORTED_CONCURRENCIES_FOR_RESULT_SET_SCROLL_SENSITIVE = 570; /* * Returns an int32 bitmask value concurrency types supported for * `arrow.flight.protocol.sql.SqlSupportedResultSetType.SQL_RESULT_SET_TYPE_SCROLL_INSENSITIVE`. * * For instance: * - return 0 (\b0) => [] (no supported concurrency types for this result set type) * - return 1 (\b1) => [SQL_RESULT_SET_CONCURRENCY_UNSPECIFIED] * - return 2 (\b10) => [SQL_RESULT_SET_CONCURRENCY_READ_ONLY] * - return 3 (\b11) => [SQL_RESULT_SET_CONCURRENCY_UNSPECIFIED, SQL_RESULT_SET_CONCURRENCY_READ_ONLY] * - return 4 (\b100) => [SQL_RESULT_SET_CONCURRENCY_UPDATABLE] * - return 5 (\b101) => [SQL_RESULT_SET_CONCURRENCY_UNSPECIFIED, SQL_RESULT_SET_CONCURRENCY_UPDATABLE] * - return 6 (\b110) => [SQL_RESULT_SET_CONCURRENCY_READ_ONLY, SQL_RESULT_SET_CONCURRENCY_UPDATABLE] * - return 7 (\b111) => [SQL_RESULT_SET_CONCURRENCY_UNSPECIFIED, SQL_RESULT_SET_CONCURRENCY_READ_ONLY, SQL_RESULT_SET_CONCURRENCY_UPDATABLE] * Valid result set types are described under `arrow.flight.protocol.sql.SqlSupportedResultSetConcurrency`. */ SQL_SUPPORTED_CONCURRENCIES_FOR_RESULT_SET_SCROLL_INSENSITIVE = 571; /* * Retrieves a boolean value indicating whether this database supports batch updates. * * - false: if this database does not support batch updates; * - true: if this database supports batch updates. */ SQL_BATCH_UPDATES_SUPPORTED = 572; /* * Retrieves a boolean value indicating whether this database supports savepoints. * * Returns: * - false: if this database does not support savepoints; * - true: if this database supports savepoints. */ SQL_SAVEPOINTS_SUPPORTED = 573; /* * Retrieves a boolean value indicating whether named parameters are supported in callable statements. * * Returns: * - false: if named parameters in callable statements are unsupported; * - true: if named parameters in callable statements are supported. */ SQL_NAMED_PARAMETERS_SUPPORTED = 574; /* * Retrieves a boolean value indicating whether updates made to a LOB are made on a copy or directly to the LOB. * * Returns: * - false: if updates made to a LOB are made directly to the LOB; * - true: if updates made to a LOB are made on a copy. */ SQL_LOCATORS_UPDATE_COPY = 575; /* * Retrieves a boolean value indicating whether invoking user-defined or vendor functions * using the stored procedure escape syntax is supported. * * Returns: * - false: if invoking user-defined or vendor functions using the stored procedure escape syntax is unsupported; * - true: if invoking user-defined or vendor functions using the stored procedure escape syntax is supported. */ SQL_STORED_FUNCTIONS_USING_CALL_SYNTAX_SUPPORTED = 576; } // The level of support for Flight SQL transaction RPCs. enum SqlSupportedTransaction { // Unknown/not indicated/no support SQL_SUPPORTED_TRANSACTION_NONE = 0; // Transactions, but not savepoints. // A savepoint is a mark within a transaction that can be individually // rolled back to. Not all databases support savepoints. SQL_SUPPORTED_TRANSACTION_TRANSACTION = 1; // Transactions and savepoints SQL_SUPPORTED_TRANSACTION_SAVEPOINT = 2; } enum SqlSupportedCaseSensitivity { SQL_CASE_SENSITIVITY_UNKNOWN = 0; SQL_CASE_SENSITIVITY_CASE_INSENSITIVE = 1; SQL_CASE_SENSITIVITY_UPPERCASE = 2; SQL_CASE_SENSITIVITY_LOWERCASE = 3; } enum SqlNullOrdering { SQL_NULLS_SORTED_HIGH = 0; SQL_NULLS_SORTED_LOW = 1; SQL_NULLS_SORTED_AT_START = 2; SQL_NULLS_SORTED_AT_END = 3; } enum SupportedSqlGrammar { SQL_MINIMUM_GRAMMAR = 0; SQL_CORE_GRAMMAR = 1; SQL_EXTENDED_GRAMMAR = 2; } enum SupportedAnsi92SqlGrammarLevel { ANSI92_ENTRY_SQL = 0; ANSI92_INTERMEDIATE_SQL = 1; ANSI92_FULL_SQL = 2; } enum SqlOuterJoinsSupportLevel { SQL_JOINS_UNSUPPORTED = 0; SQL_LIMITED_OUTER_JOINS = 1; SQL_FULL_OUTER_JOINS = 2; } enum SqlSupportedGroupBy { SQL_GROUP_BY_UNRELATED = 0; SQL_GROUP_BY_BEYOND_SELECT = 1; } enum SqlSupportedElementActions { SQL_ELEMENT_IN_PROCEDURE_CALLS = 0; SQL_ELEMENT_IN_INDEX_DEFINITIONS = 1; SQL_ELEMENT_IN_PRIVILEGE_DEFINITIONS = 2; } enum SqlSupportedPositionedCommands { SQL_POSITIONED_DELETE = 0; SQL_POSITIONED_UPDATE = 1; } enum SqlSupportedSubqueries { SQL_SUBQUERIES_IN_COMPARISONS = 0; SQL_SUBQUERIES_IN_EXISTS = 1; SQL_SUBQUERIES_IN_INS = 2; SQL_SUBQUERIES_IN_QUANTIFIEDS = 3; } enum SqlSupportedUnions { SQL_UNION = 0; SQL_UNION_ALL = 1; } enum SqlTransactionIsolationLevel { SQL_TRANSACTION_NONE = 0; SQL_TRANSACTION_READ_UNCOMMITTED = 1; SQL_TRANSACTION_READ_COMMITTED = 2; SQL_TRANSACTION_REPEATABLE_READ = 3; SQL_TRANSACTION_SERIALIZABLE = 4; } enum SqlSupportedTransactions { SQL_TRANSACTION_UNSPECIFIED = 0; SQL_DATA_DEFINITION_TRANSACTIONS = 1; SQL_DATA_MANIPULATION_TRANSACTIONS = 2; } enum SqlSupportedResultSetType { SQL_RESULT_SET_TYPE_UNSPECIFIED = 0; SQL_RESULT_SET_TYPE_FORWARD_ONLY = 1; SQL_RESULT_SET_TYPE_SCROLL_INSENSITIVE = 2; SQL_RESULT_SET_TYPE_SCROLL_SENSITIVE = 3; } enum SqlSupportedResultSetConcurrency { SQL_RESULT_SET_CONCURRENCY_UNSPECIFIED = 0; SQL_RESULT_SET_CONCURRENCY_READ_ONLY = 1; SQL_RESULT_SET_CONCURRENCY_UPDATABLE = 2; } enum SqlSupportsConvert { SQL_CONVERT_BIGINT = 0; SQL_CONVERT_BINARY = 1; SQL_CONVERT_BIT = 2; SQL_CONVERT_CHAR = 3; SQL_CONVERT_DATE = 4; SQL_CONVERT_DECIMAL = 5; SQL_CONVERT_FLOAT = 6; SQL_CONVERT_INTEGER = 7; SQL_CONVERT_INTERVAL_DAY_TIME = 8; SQL_CONVERT_INTERVAL_YEAR_MONTH = 9; SQL_CONVERT_LONGVARBINARY = 10; SQL_CONVERT_LONGVARCHAR = 11; SQL_CONVERT_NUMERIC = 12; SQL_CONVERT_REAL = 13; SQL_CONVERT_SMALLINT = 14; SQL_CONVERT_TIME = 15; SQL_CONVERT_TIMESTAMP = 16; SQL_CONVERT_TINYINT = 17; SQL_CONVERT_VARBINARY = 18; SQL_CONVERT_VARCHAR = 19; } /** * The JDBC/ODBC-defined type of any object. * All the values here are the same as in the JDBC and ODBC specs. */ enum XdbcDataType { XDBC_UNKNOWN_TYPE = 0; XDBC_CHAR = 1; XDBC_NUMERIC = 2; XDBC_DECIMAL = 3; XDBC_INTEGER = 4; XDBC_SMALLINT = 5; XDBC_FLOAT = 6; XDBC_REAL = 7; XDBC_DOUBLE = 8; XDBC_DATETIME = 9; XDBC_INTERVAL = 10; XDBC_VARCHAR = 12; XDBC_DATE = 91; XDBC_TIME = 92; XDBC_TIMESTAMP = 93; XDBC_LONGVARCHAR = -1; XDBC_BINARY = -2; XDBC_VARBINARY = -3; XDBC_LONGVARBINARY = -4; XDBC_BIGINT = -5; XDBC_TINYINT = -6; XDBC_BIT = -7; XDBC_WCHAR = -8; XDBC_WVARCHAR = -9; } /** * Detailed subtype information for XDBC_TYPE_DATETIME and XDBC_TYPE_INTERVAL. */ enum XdbcDatetimeSubcode { option allow_alias = true; XDBC_SUBCODE_UNKNOWN = 0; XDBC_SUBCODE_YEAR = 1; XDBC_SUBCODE_DATE = 1; XDBC_SUBCODE_TIME = 2; XDBC_SUBCODE_MONTH = 2; XDBC_SUBCODE_TIMESTAMP = 3; XDBC_SUBCODE_DAY = 3; XDBC_SUBCODE_TIME_WITH_TIMEZONE = 4; XDBC_SUBCODE_HOUR = 4; XDBC_SUBCODE_TIMESTAMP_WITH_TIMEZONE = 5; XDBC_SUBCODE_MINUTE = 5; XDBC_SUBCODE_SECOND = 6; XDBC_SUBCODE_YEAR_TO_MONTH = 7; XDBC_SUBCODE_DAY_TO_HOUR = 8; XDBC_SUBCODE_DAY_TO_MINUTE = 9; XDBC_SUBCODE_DAY_TO_SECOND = 10; XDBC_SUBCODE_HOUR_TO_MINUTE = 11; XDBC_SUBCODE_HOUR_TO_SECOND = 12; XDBC_SUBCODE_MINUTE_TO_SECOND = 13; XDBC_SUBCODE_INTERVAL_YEAR = 101; XDBC_SUBCODE_INTERVAL_MONTH = 102; XDBC_SUBCODE_INTERVAL_DAY = 103; XDBC_SUBCODE_INTERVAL_HOUR = 104; XDBC_SUBCODE_INTERVAL_MINUTE = 105; XDBC_SUBCODE_INTERVAL_SECOND = 106; XDBC_SUBCODE_INTERVAL_YEAR_TO_MONTH = 107; XDBC_SUBCODE_INTERVAL_DAY_TO_HOUR = 108; XDBC_SUBCODE_INTERVAL_DAY_TO_MINUTE = 109; XDBC_SUBCODE_INTERVAL_DAY_TO_SECOND = 110; XDBC_SUBCODE_INTERVAL_HOUR_TO_MINUTE = 111; XDBC_SUBCODE_INTERVAL_HOUR_TO_SECOND = 112; XDBC_SUBCODE_INTERVAL_MINUTE_TO_SECOND = 113; } enum Nullable { /** * Indicates that the fields does not allow the use of null values. */ NULLABILITY_NO_NULLS = 0; /** * Indicates that the fields allow the use of null values. */ NULLABILITY_NULLABLE = 1; /** * Indicates that nullability of the fields cannot be determined. */ NULLABILITY_UNKNOWN = 2; } enum Searchable { /** * Indicates that column cannot be used in a WHERE clause. */ SEARCHABLE_NONE = 0; /** * Indicates that the column can be used in a WHERE clause if it is using a * LIKE operator. */ SEARCHABLE_CHAR = 1; /** * Indicates that the column can be used In a WHERE clause with any * operator other than LIKE. * * - Allowed operators: comparison, quantified comparison, BETWEEN, * DISTINCT, IN, MATCH, and UNIQUE. */ SEARCHABLE_BASIC = 2; /** * Indicates that the column can be used in a WHERE clause using any operator. */ SEARCHABLE_FULL = 3; } /* * Represents a request to retrieve information about data type supported on a Flight SQL enabled backend. * Used in the command member of FlightDescriptor for the following RPC calls: * - GetSchema: return the schema of the query. * - GetFlightInfo: execute the catalog metadata request. * * The returned schema will be: * < * type_name: utf8 not null (The name of the data type, for example: VARCHAR, INTEGER, etc), * data_type: int32 not null (The SQL data type), * column_size: int32 (The maximum size supported by that column. * In case of exact numeric types, this represents the maximum precision. * In case of string types, this represents the character length. * In case of datetime data types, this represents the length in characters of the string representation. * NULL is returned for data types where column size is not applicable.), * literal_prefix: utf8 (Character or characters used to prefix a literal, NULL is returned for * data types where a literal prefix is not applicable.), * literal_suffix: utf8 (Character or characters used to terminate a literal, * NULL is returned for data types where a literal suffix is not applicable.), * create_params: list * (A list of keywords corresponding to which parameters can be used when creating * a column for that specific type. * NULL is returned if there are no parameters for the data type definition.), * nullable: int32 not null (Shows if the data type accepts a NULL value. The possible values can be seen in the * Nullable enum.), * case_sensitive: bool not null (Shows if a character data type is case-sensitive in collations and comparisons), * searchable: int32 not null (Shows how the data type is used in a WHERE clause. The possible values can be seen in the * Searchable enum.), * unsigned_attribute: bool (Shows if the data type is unsigned. NULL is returned if the attribute is * not applicable to the data type or the data type is not numeric.), * fixed_prec_scale: bool not null (Shows if the data type has predefined fixed precision and scale.), * auto_increment: bool (Shows if the data type is auto incremental. NULL is returned if the attribute * is not applicable to the data type or the data type is not numeric.), * local_type_name: utf8 (Localized version of the data source-dependent name of the data type. NULL * is returned if a localized name is not supported by the data source), * minimum_scale: int32 (The minimum scale of the data type on the data source. * If a data type has a fixed scale, the MINIMUM_SCALE and MAXIMUM_SCALE * columns both contain this value. NULL is returned if scale is not applicable.), * maximum_scale: int32 (The maximum scale of the data type on the data source. * NULL is returned if scale is not applicable.), * sql_data_type: int32 not null (The value of the SQL DATA TYPE which has the same values * as data_type value. Except for interval and datetime, which * uses generic values. More info about those types can be * obtained through datetime_subcode. The possible values can be seen * in the XdbcDataType enum.), * datetime_subcode: int32 (Only used when the SQL DATA TYPE is interval or datetime. It contains * its sub types. For type different from interval and datetime, this value * is NULL. The possible values can be seen in the XdbcDatetimeSubcode enum.), * num_prec_radix: int32 (If the data type is an approximate numeric type, this column contains * the value 2 to indicate that COLUMN_SIZE specifies a number of bits. For * exact numeric types, this column contains the value 10 to indicate that * column size specifies a number of decimal digits. Otherwise, this column is NULL.), * interval_precision: int32 (If the data type is an interval data type, then this column contains the value * of the interval leading precision. Otherwise, this column is NULL. This fields * is only relevant to be used by ODBC). * > * The returned data should be ordered by data_type and then by type_name. */ message CommandGetXdbcTypeInfo { /* * Specifies the data type to search for the info. */ optional int32 data_type = 1; } /* * Represents a request to retrieve the list of catalogs on a Flight SQL enabled backend. * The definition of a catalog depends on vendor/implementation. It is usually the database itself * Used in the command member of FlightDescriptor for the following RPC calls: * - GetSchema: return the Arrow schema of the query. * - GetFlightInfo: execute the catalog metadata request. * * The returned Arrow schema will be: * < * catalog_name: utf8 not null * > * The returned data should be ordered by catalog_name. */ message CommandGetCatalogs { } /* * Represents a request to retrieve the list of database schemas on a Flight SQL enabled backend. * The definition of a database schema depends on vendor/implementation. It is usually a collection of tables. * Used in the command member of FlightDescriptor for the following RPC calls: * - GetSchema: return the Arrow schema of the query. * - GetFlightInfo: execute the catalog metadata request. * * The returned Arrow schema will be: * < * catalog_name: utf8, * db_schema_name: utf8 not null * > * The returned data should be ordered by catalog_name, then db_schema_name. */ message CommandGetDbSchemas { /* * Specifies the Catalog to search for the tables. * An empty string retrieves those without a catalog. * If omitted the catalog name should not be used to narrow the search. */ optional string catalog = 1; /* * Specifies a filter pattern for schemas to search for. * When no db_schema_filter_pattern is provided, the pattern will not be used to narrow the search. * In the pattern string, two special characters can be used to denote matching rules: * - "%" means to match any substring with 0 or more characters. * - "_" means to match any one character. */ optional string db_schema_filter_pattern = 2; } /* * Represents a request to retrieve the list of tables, and optionally their schemas, on a Flight SQL enabled backend. * Used in the command member of FlightDescriptor for the following RPC calls: * - GetSchema: return the Arrow schema of the query. * - GetFlightInfo: execute the catalog metadata request. * * The returned Arrow schema will be: * < * catalog_name: utf8, * db_schema_name: utf8, * table_name: utf8 not null, * table_type: utf8 not null, * [optional] table_schema: bytes not null (schema of the table as described in Schema.fbs::Schema, * it is serialized as an IPC message.) * > * Fields on table_schema may contain the following metadata: * - ARROW:FLIGHT:SQL:CATALOG_NAME - Table's catalog name * - ARROW:FLIGHT:SQL:DB_SCHEMA_NAME - Database schema name * - ARROW:FLIGHT:SQL:TABLE_NAME - Table name * - ARROW:FLIGHT:SQL:TYPE_NAME - The data source-specific name for the data type of the column. * - ARROW:FLIGHT:SQL:PRECISION - Column precision/size * - ARROW:FLIGHT:SQL:SCALE - Column scale/decimal digits if applicable * - ARROW:FLIGHT:SQL:IS_AUTO_INCREMENT - "1" indicates if the column is auto incremented, "0" otherwise. * - 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 { /* * Specifies the Catalog to search for the tables. * An empty string retrieves those without a catalog. * If omitted the catalog name should not be used to narrow the search. */ optional string catalog = 1; /* * Specifies a filter pattern for schemas to search for. * When no db_schema_filter_pattern is provided, all schemas matching other filters are searched. * In the pattern string, two special characters can be used to denote matching rules: * - "%" means to match any substring with 0 or more characters. * - "_" means to match any one character. */ optional string db_schema_filter_pattern = 2; /* * Specifies a filter pattern for tables to search for. * When no table_name_filter_pattern is provided, all tables matching other filters are searched. * In the pattern string, two special characters can be used to denote matching rules: * - "%" means to match any substring with 0 or more characters. * - "_" means to match any one character. */ optional string table_name_filter_pattern = 3; /* * Specifies a filter of table types which must match. * The table types depend on vendor/implementation. It is usually used to separate tables from views or system tables. * TABLE, VIEW, and SYSTEM TABLE are commonly supported. */ repeated string table_types = 4; // Specifies if the Arrow schema should be returned for found tables. bool include_schema = 5; } /* * Represents a request to retrieve the list of table types on a Flight SQL enabled backend. * The table types depend on vendor/implementation. It is usually used to separate tables from views or system tables. * TABLE, VIEW, and SYSTEM TABLE are commonly supported. * Used in the command member of FlightDescriptor for the following RPC calls: * - GetSchema: return the Arrow schema of the query. * - GetFlightInfo: execute the catalog metadata request. * * The returned Arrow schema will be: * < * table_type: utf8 not null * > * The returned data should be ordered by table_type. */ message CommandGetTableTypes { } /* * Represents a request to retrieve the primary keys of a table on a Flight SQL enabled backend. * Used in the command member of FlightDescriptor for the following RPC calls: * - GetSchema: return the Arrow schema of the query. * - GetFlightInfo: execute the catalog metadata request. * * The returned Arrow schema will be: * < * catalog_name: utf8, * db_schema_name: utf8, * table_name: utf8 not null, * column_name: utf8 not null, * key_name: utf8, * key_sequence: int32 not null * > * The returned data should be ordered by catalog_name, db_schema_name, table_name, key_name, then key_sequence. */ message CommandGetPrimaryKeys { /* * Specifies the catalog to search for the table. * An empty string retrieves those without a catalog. * If omitted the catalog name should not be used to narrow the search. */ optional string catalog = 1; /* * Specifies the schema to search for the table. * An empty string retrieves those without a schema. * If omitted the schema name should not be used to narrow the search. */ optional string db_schema = 2; // Specifies the table to get the primary keys for. string table = 3; } enum UpdateDeleteRules { CASCADE = 0; RESTRICT = 1; SET_NULL = 2; NO_ACTION = 3; SET_DEFAULT = 4; } /* * Represents a request to retrieve a description of the foreign key columns that reference the given table's * primary key columns (the foreign keys exported by a table) of a table on a Flight SQL enabled backend. * Used in the command member of FlightDescriptor for the following RPC calls: * - GetSchema: return the Arrow schema of the query. * - GetFlightInfo: execute the catalog metadata request. * * The returned Arrow schema will be: * < * pk_catalog_name: utf8, * pk_db_schema_name: utf8, * pk_table_name: utf8 not null, * pk_column_name: utf8 not null, * fk_catalog_name: utf8, * fk_db_schema_name: utf8, * fk_table_name: utf8 not null, * fk_column_name: utf8 not null, * key_sequence: int32 not null, * fk_key_name: utf8, * pk_key_name: utf8, * update_rule: uint8 not null, * delete_rule: uint8 not null * > * The returned data should be ordered by fk_catalog_name, fk_db_schema_name, fk_table_name, fk_key_name, then key_sequence. * update_rule and delete_rule returns a byte that is equivalent to actions declared on UpdateDeleteRules enum. */ message CommandGetExportedKeys { /* * Specifies the catalog to search for the foreign key table. * An empty string retrieves those without a catalog. * If omitted the catalog name should not be used to narrow the search. */ optional string catalog = 1; /* * Specifies the schema to search for the foreign key table. * An empty string retrieves those without a schema. * If omitted the schema name should not be used to narrow the search. */ optional string db_schema = 2; // Specifies the foreign key table to get the foreign keys for. string table = 3; } /* * Represents a request to retrieve the foreign keys of a table on a Flight SQL enabled backend. * Used in the command member of FlightDescriptor for the following RPC calls: * - GetSchema: return the Arrow schema of the query. * - GetFlightInfo: execute the catalog metadata request. * * The returned Arrow schema will be: * < * pk_catalog_name: utf8, * pk_db_schema_name: utf8, * pk_table_name: utf8 not null, * pk_column_name: utf8 not null, * fk_catalog_name: utf8, * fk_db_schema_name: utf8, * fk_table_name: utf8 not null, * fk_column_name: utf8 not null, * key_sequence: int32 not null, * fk_key_name: utf8, * pk_key_name: utf8, * update_rule: uint8 not null, * delete_rule: uint8 not null * > * The returned data should be ordered by pk_catalog_name, pk_db_schema_name, pk_table_name, pk_key_name, then key_sequence. * update_rule and delete_rule returns a byte that is equivalent to actions: * - 0 = CASCADE * - 1 = RESTRICT * - 2 = SET NULL * - 3 = NO ACTION * - 4 = SET DEFAULT */ message CommandGetImportedKeys { /* * Specifies the catalog to search for the primary key table. * An empty string retrieves those without a catalog. * If omitted the catalog name should not be used to narrow the search. */ optional string catalog = 1; /* * Specifies the schema to search for the primary key table. * An empty string retrieves those without a schema. * If omitted the schema name should not be used to narrow the search. */ optional string db_schema = 2; // Specifies the primary key table to get the foreign keys for. string table = 3; } /* * Represents a request to retrieve a description of the foreign key columns in the given foreign key table that * reference the primary key or the columns representing a unique constraint of the parent table (could be the same * or a different table) on a Flight SQL enabled backend. * Used in the command member of FlightDescriptor for the following RPC calls: * - GetSchema: return the Arrow schema of the query. * - GetFlightInfo: execute the catalog metadata request. * * The returned Arrow schema will be: * < * pk_catalog_name: utf8, * pk_db_schema_name: utf8, * pk_table_name: utf8 not null, * pk_column_name: utf8 not null, * fk_catalog_name: utf8, * fk_db_schema_name: utf8, * fk_table_name: utf8 not null, * fk_column_name: utf8 not null, * key_sequence: int32 not null, * fk_key_name: utf8, * pk_key_name: utf8, * update_rule: uint8 not null, * delete_rule: uint8 not null * > * The returned data should be ordered by pk_catalog_name, pk_db_schema_name, pk_table_name, pk_key_name, then key_sequence. * update_rule and delete_rule returns a byte that is equivalent to actions: * - 0 = CASCADE * - 1 = RESTRICT * - 2 = SET NULL * - 3 = NO ACTION * - 4 = SET DEFAULT */ message CommandGetCrossReference { /** * The catalog name where the parent table is. * An empty string retrieves those without a catalog. * If omitted the catalog name should not be used to narrow the search. */ optional string pk_catalog = 1; /** * The Schema name where the parent table is. * An empty string retrieves those without a schema. * If omitted the schema name should not be used to narrow the search. */ optional string pk_db_schema = 2; /** * The parent table name. It cannot be null. */ string pk_table = 3; /** * The catalog name where the foreign table is. * An empty string retrieves those without a catalog. * If omitted the catalog name should not be used to narrow the search. */ optional string fk_catalog = 4; /** * The schema name where the foreign table is. * An empty string retrieves those without a schema. * If omitted the schema name should not be used to narrow the search. */ optional string fk_db_schema = 5; /** * The foreign table name. It cannot be null. */ string fk_table = 6; } // Query Execution Action Messages /* * Request message for the "CreatePreparedStatement" action on a Flight SQL enabled backend. */ message ActionCreatePreparedStatementRequest { // The valid SQL string to create a prepared statement for. string query = 1; // Create/execute the prepared statement as part of this transaction (if // unset, executions of the prepared statement will be auto-committed). optional bytes transaction_id = 2; } /* * An embedded message describing a Substrait plan to execute. */ message SubstraitPlan { // The serialized substrait.Plan to create a prepared statement for. // XXX(ARROW-16902): this is bytes instead of an embedded message // because Protobuf does not really support one DLL using Protobuf // definitions from another DLL. bytes plan = 1; // The Substrait release, e.g. "0.12.0". This information is not // tracked in the plan itself, so this is the only way for consumers // to potentially know if they can handle the plan. string version = 2; } /* * Request message for the "CreatePreparedSubstraitPlan" action on a Flight SQL enabled backend. */ message ActionCreatePreparedSubstraitPlanRequest { // The serialized substrait.Plan to create a prepared statement for. SubstraitPlan plan = 1; // Create/execute the prepared statement as part of this transaction (if // unset, executions of the prepared statement will be auto-committed). optional bytes transaction_id = 2; } /* * Wrap the result of a "CreatePreparedStatement" or "CreatePreparedSubstraitPlan" action. * * The resultant PreparedStatement can be closed either: * - Manually, through the "ClosePreparedStatement" action; * - Automatically, by a server timeout. * * The result should be wrapped in a google.protobuf.Any message. */ message ActionCreatePreparedStatementResult { // Opaque handle for the prepared statement on the server. bytes prepared_statement_handle = 1; // If a result set generating query was provided, dataset_schema contains the // schema of the result set. It should be an IPC-encapsulated Schema, as described in Schema.fbs. // For some queries, the schema of the results may depend on the schema of the parameters. The server // should provide its best guess as to the schema at this point. Clients must not assume that this // schema, if provided, will be accurate. bytes dataset_schema = 2; // 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; } /* * Request message for the "ClosePreparedStatement" action on a Flight SQL enabled backend. * Closes server resources associated with the prepared statement handle. */ message ActionClosePreparedStatementRequest { // Opaque handle for the prepared statement on the server. bytes prepared_statement_handle = 1; } /* * Request message for the "BeginTransaction" action. * Begins a transaction. */ message ActionBeginTransactionRequest { } /* * Request message for the "BeginSavepoint" action. * Creates a savepoint within a transaction. * * Only supported if FLIGHT_SQL_TRANSACTION is * FLIGHT_SQL_TRANSACTION_SUPPORT_SAVEPOINT. */ message ActionBeginSavepointRequest { // The transaction to which a savepoint belongs. bytes transaction_id = 1; // Name for the savepoint. string name = 2; } /* * The result of a "BeginTransaction" action. * * The transaction can be manipulated with the "EndTransaction" action, or * automatically via server timeout. If the transaction times out, then it is * automatically rolled back. * * The result should be wrapped in a google.protobuf.Any message. */ message ActionBeginTransactionResult { // Opaque handle for the transaction on the server. bytes transaction_id = 1; } /* * The result of a "BeginSavepoint" action. * * The transaction can be manipulated with the "EndSavepoint" action. * If the associated transaction is committed, rolled back, or times * out, then the savepoint is also invalidated. * * The result should be wrapped in a google.protobuf.Any message. */ message ActionBeginSavepointResult { // Opaque handle for the savepoint on the server. bytes savepoint_id = 1; } /* * Request message for the "EndTransaction" action. * * Commit (COMMIT) or rollback (ROLLBACK) the transaction. * * If the action completes successfully, the transaction handle is * invalidated, as are all associated savepoints. */ message ActionEndTransactionRequest { enum EndTransaction { END_TRANSACTION_UNSPECIFIED = 0; // Commit the transaction. END_TRANSACTION_COMMIT = 1; // Roll back the transaction. END_TRANSACTION_ROLLBACK = 2; } // Opaque handle for the transaction on the server. bytes transaction_id = 1; // Whether to commit/rollback the given transaction. EndTransaction action = 2; } /* * Request message for the "EndSavepoint" action. * * Release (RELEASE) the savepoint or rollback (ROLLBACK) to the * savepoint. * * Releasing a savepoint invalidates that savepoint. Rolling back to * a savepoint does not invalidate the savepoint, but invalidates all * savepoints created after the current savepoint. */ message ActionEndSavepointRequest { enum EndSavepoint { END_SAVEPOINT_UNSPECIFIED = 0; // Release the savepoint. END_SAVEPOINT_RELEASE = 1; // Roll back to a savepoint. END_SAVEPOINT_ROLLBACK = 2; } // Opaque handle for the savepoint on the server. bytes savepoint_id = 1; // Whether to rollback/release the given savepoint. EndSavepoint action = 2; } // Query Execution Messages. /* * Represents a SQL query. Used in the command member of FlightDescriptor * for the following RPC calls: * - GetSchema: return the Arrow schema of the query. * Fields on this schema may contain the following metadata: * - ARROW:FLIGHT:SQL:CATALOG_NAME - Table's catalog name * - ARROW:FLIGHT:SQL:DB_SCHEMA_NAME - Database schema name * - ARROW:FLIGHT:SQL:TABLE_NAME - Table name * - ARROW:FLIGHT:SQL:TYPE_NAME - The data source-specific name for the data type of the column. * - ARROW:FLIGHT:SQL:PRECISION - Column precision/size * - ARROW:FLIGHT:SQL:SCALE - Column scale/decimal digits if applicable * - ARROW:FLIGHT:SQL:IS_AUTO_INCREMENT - "1" indicates if the column is auto incremented, "0" otherwise. * - 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 { // The SQL syntax. string query = 1; // Include the query as part of this transaction (if unset, the query is auto-committed). optional bytes transaction_id = 2; } /* * Represents a Substrait plan. Used in the command member of FlightDescriptor * for the following RPC calls: * - GetSchema: return the Arrow schema of the query. * Fields on this schema may contain the following metadata: * - ARROW:FLIGHT:SQL:CATALOG_NAME - Table's catalog name * - ARROW:FLIGHT:SQL:DB_SCHEMA_NAME - Database schema name * - ARROW:FLIGHT:SQL:TABLE_NAME - Table name * - ARROW:FLIGHT:SQL:TYPE_NAME - The data source-specific name for the data type of the column. * - ARROW:FLIGHT:SQL:PRECISION - Column precision/size * - ARROW:FLIGHT:SQL:SCALE - Column scale/decimal digits if applicable * - ARROW:FLIGHT:SQL:IS_AUTO_INCREMENT - "1" indicates if the column is auto incremented, "0" otherwise. * - 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. */ message CommandStatementSubstraitPlan { // A serialized substrait.Plan SubstraitPlan plan = 1; // Include the query as part of this transaction (if unset, the query is auto-committed). optional bytes transaction_id = 2; } /** * Represents a ticket resulting from GetFlightInfo with a CommandStatementQuery. * This should be used only once and treated as an opaque value, that is, clients should not attempt to parse this. */ message TicketStatementQuery { // Unique identifier for the instance of the statement to execute. bytes statement_handle = 1; } /* * Represents an instance of executing a prepared statement. Used in the command member of FlightDescriptor for * the following RPC calls: * - GetSchema: return the Arrow schema of the query. * Fields on this schema may contain the following metadata: * - ARROW:FLIGHT:SQL:CATALOG_NAME - Table's catalog name * - ARROW:FLIGHT:SQL:DB_SCHEMA_NAME - Database schema name * - ARROW:FLIGHT:SQL:TABLE_NAME - Table name * - ARROW:FLIGHT:SQL:TYPE_NAME - The data source-specific name for the data type of the column. * - ARROW:FLIGHT:SQL:PRECISION - Column precision/size * - ARROW:FLIGHT:SQL:SCALE - Column scale/decimal digits if applicable * - ARROW:FLIGHT:SQL:IS_AUTO_INCREMENT - "1" indicates if the column is auto incremented, "0" otherwise. * - 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. * - DoPut: bind parameter values. All of the bound parameter sets will be executed as a single atomic execution. * - GetFlightInfo: execute the prepared statement instance. */ message CommandPreparedStatementQuery { // Opaque handle for the prepared statement on the server. bytes prepared_statement_handle = 1; } /* * Represents a SQL update query. Used in the command member of FlightDescriptor * for the RPC call DoPut to cause the server to execute the included SQL update. */ message CommandStatementUpdate { // The SQL syntax. string query = 1; // Include the query as part of this transaction (if unset, the query is auto-committed). optional bytes transaction_id = 2; } /* * Represents a SQL update query. Used in the command member of FlightDescriptor * for the RPC call DoPut to cause the server to execute the included * prepared statement handle as an update. */ message CommandPreparedStatementUpdate { // Opaque handle for the prepared statement on the server. bytes prepared_statement_handle = 1; } /* * Represents a bulk ingestion request. Used in the command member of FlightDescriptor * for the the RPC call DoPut to cause the server load the contents of the stream's * FlightData into the target destination. */ message CommandStatementIngest { // Options for table definition behavior message TableDefinitionOptions { // The action to take if the target table does not exist enum TableNotExistOption { // Do not use. Servers should error if this is specified by a client. TABLE_NOT_EXIST_OPTION_UNSPECIFIED = 0; // Create the table if it does not exist TABLE_NOT_EXIST_OPTION_CREATE = 1; // Fail if the table does not exist TABLE_NOT_EXIST_OPTION_FAIL = 2; } // The action to take if the target table already exists enum TableExistsOption { // Do not use. Servers should error if this is specified by a client. TABLE_EXISTS_OPTION_UNSPECIFIED = 0; // Fail if the table already exists TABLE_EXISTS_OPTION_FAIL = 1; // Append to the table if it already exists TABLE_EXISTS_OPTION_APPEND = 2; // Drop and recreate the table if it already exists TABLE_EXISTS_OPTION_REPLACE = 3; } TableNotExistOption if_not_exist = 1; TableExistsOption if_exists = 2; } // The behavior for handling the table definition. TableDefinitionOptions table_definition_options = 1; // The table to load data into. string table = 2; // The db_schema of the destination table to load data into. If unset, a backend-specific default may be used. optional string schema = 3; // The catalog of the destination table to load data into. If unset, a backend-specific default may be used. optional string catalog = 4; /* * Store ingested data in a temporary table. * The effect of setting temporary is to place the table in a backend-defined namespace, and to drop the table at the end of the session. * The namespacing may make use of a backend-specific schema and/or catalog. * The server should return an error if an explicit choice of schema or catalog is incompatible with the server's namespacing decision. */ bool temporary = 5; // Perform the ingestion as part of this transaction. If specified, results should not be committed in the event of an error/cancellation. optional bytes transaction_id = 6; // Future extensions to the parameters of CommandStatementIngest should be added here, at a lower index than the generic 'options' parameter. // Backend-specific options. map options = 1000; } /* * Returned from the RPC call DoPut when a CommandStatementUpdate, * CommandPreparedStatementUpdate, or CommandStatementIngest was * in the request, containing results from the update. */ message DoPutUpdateResult { // The number of records updated. A return value of -1 represents // an unknown updated record count. int64 record_count = 1; } /* An *optional* response returned when `DoPut` is called with `CommandPreparedStatementQuery`. * * *Note on legacy behavior*: previous versions of the protocol did not return any result for * this command, and that behavior should still be supported by clients. In that case, the client * can continue as though the fields in this message were not provided or set to sensible default values. */ message DoPutPreparedStatementResult { // Represents a (potentially updated) opaque handle for the prepared statement on the server. // Because the handle could potentially be updated, any previous handles for this prepared // statement should be considered invalid, and all subsequent requests for this prepared // statement must use this new handle. // The updated handle allows implementing query parameters with stateless services. // // When an updated handle is not provided by the server, clients should contiue // using the previous handle provided by `ActionCreatePreparedStatementResonse`. optional bytes prepared_statement_handle = 1; } /* * Request message for the "CancelQuery" action. * * Explicitly cancel a running query. * * This lets a single client explicitly cancel work, no matter how many clients * are involved/whether the query is distributed or not, given server support. * The transaction/statement is not rolled back; it is the application's job to * commit or rollback as appropriate. This only indicates the client no longer * wishes to read the remainder of the query results or continue submitting * data. * * This command is idempotent. * * This command is deprecated since 13.0.0. Use the "CancelFlightInfo" * action with DoAction instead. */ message ActionCancelQueryRequest { option deprecated = true; // The result of the GetFlightInfo RPC that initiated the query. // XXX(ARROW-16902): this must be a serialized FlightInfo, but is // rendered as bytes because Protobuf does not really support one // DLL using Protobuf definitions from another DLL. bytes info = 1; } /* * The result of cancelling a query. * * The result should be wrapped in a google.protobuf.Any message. * * This command is deprecated since 13.0.0. Use the "CancelFlightInfo" * action with DoAction instead. */ message ActionCancelQueryResult { option deprecated = true; enum CancelResult { // The cancellation status is unknown. Servers should avoid using // this value (send a NOT_FOUND error if the requested query is // not known). Clients can retry the request. CANCEL_RESULT_UNSPECIFIED = 0; // The cancellation request is complete. Subsequent requests with // the same payload may return CANCELLED or a NOT_FOUND error. CANCEL_RESULT_CANCELLED = 1; // The cancellation request is in progress. The client may retry // the cancellation request. CANCEL_RESULT_CANCELLING = 2; // The query is not cancellable. The client should not retry the // cancellation request. CANCEL_RESULT_NOT_CANCELLABLE = 3; } CancelResult result = 1; } extend google.protobuf.MessageOptions { bool experimental = 1000; } apache-arrow-java-19.0.0/arrow-format/Message.fbs000066400000000000000000000140631515451615200216020ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. include "Schema.fbs"; include "SparseTensor.fbs"; include "Tensor.fbs"; namespace org.apache.arrow.flatbuf; /// ---------------------------------------------------------------------- /// Data structures for describing a table row batch (a collection of /// equal-length Arrow arrays) /// Metadata about a field at some level of a nested type tree (but not /// its children). /// /// For example, a List with values `[[1, 2, 3], null, [4], [5, 6], null]` /// would have {length: 5, null_count: 2} for its List node, and {length: 6, /// null_count: 0} for its Int16 node, as separate FieldNode structs struct FieldNode { /// The number of value slots in the Arrow array at this level of a nested /// tree length: long; /// The number of observed nulls. Fields with null_count == 0 may choose not /// to write their physical validity bitmap out as a materialized buffer, /// instead setting the length of the bitmap buffer to 0. null_count: long; } enum CompressionType:byte { // LZ4 frame format, for portability, as provided by lz4frame.h or wrappers // thereof. Not to be confused with "raw" (also called "block") format // provided by lz4.h LZ4_FRAME, // Zstandard ZSTD } /// Provided for forward compatibility in case we need to support different /// strategies for compressing the IPC message body (like whole-body /// compression rather than buffer-level) in the future enum BodyCompressionMethod:byte { /// Each constituent buffer is first compressed with the indicated /// compressor, and then written with the uncompressed length in the first 8 /// bytes as a 64-bit little-endian signed integer followed by the compressed /// buffer bytes (and then padding as required by the protocol). The /// uncompressed length may be set to -1 to indicate that the data that /// follows is not compressed, which can be useful for cases where /// compression does not yield appreciable savings. BUFFER } /// Optional compression for the memory buffers constituting IPC message /// bodies. Intended for use with RecordBatch but could be used for other /// message types table BodyCompression { /// Compressor library. /// For LZ4_FRAME, each compressed buffer must consist of a single frame. codec: CompressionType = LZ4_FRAME; /// Indicates the way the record batch body was compressed method: BodyCompressionMethod = BUFFER; } /// A data header describing the shared memory layout of a "record" or "row" /// batch. Some systems call this a "row batch" internally and others a "record /// batch". table RecordBatch { /// number of records / rows. The arrays in the batch should all have this /// length length: long; /// Nodes correspond to the pre-ordered flattened logical schema nodes: [FieldNode]; /// Buffers correspond to the pre-ordered flattened buffer tree /// /// The number of buffers appended to this list depends on the schema. For /// example, most primitive arrays will have 2 buffers, 1 for the validity /// bitmap and 1 for the values. For struct arrays, there will only be a /// single buffer for the validity (nulls) bitmap buffers: [Buffer]; /// Optional compression of the message body compression: BodyCompression; /// Some types such as Utf8View are represented using a variable number of buffers. /// For each such Field in the pre-ordered flattened logical schema, there will be /// an entry in variadicBufferCounts to indicate the number of number of variadic /// buffers which belong to that Field in the current RecordBatch. /// /// For example, the schema /// col1: Struct /// col2: Utf8View /// contains two Fields with variadic buffers so variadicBufferCounts will have /// two entries, the first counting the variadic buffers of `col1.beta` and the /// second counting `col2`'s. /// /// This field may be omitted if and only if the schema contains no Fields with /// a variable number of buffers, such as BinaryView and Utf8View. variadicBufferCounts: [long]; } /// For sending dictionary encoding information. Any Field can be /// dictionary-encoded, but in this case none of its children may be /// dictionary-encoded. /// There is one vector / column per dictionary, but that vector / column /// may be spread across multiple dictionary batches by using the isDelta /// flag table DictionaryBatch { id: long; data: RecordBatch; /// If isDelta is true the values in the dictionary are to be appended to a /// dictionary with the indicated id. If isDelta is false this dictionary /// should replace the existing dictionary. isDelta: bool = false; } /// ---------------------------------------------------------------------- /// The root Message type /// This union enables us to easily send different message types without /// redundant storage, and in the future we can easily add new message types. /// /// Arrow implementations do not need to implement all of the message types, /// which may include experimental metadata types. For maximum compatibility, /// it is best to send data using RecordBatch union MessageHeader { Schema, DictionaryBatch, RecordBatch, Tensor, SparseTensor } table Message { version: org.apache.arrow.flatbuf.MetadataVersion; header: MessageHeader; bodyLength: long; custom_metadata: [ KeyValue ]; } root_type Message; apache-arrow-java-19.0.0/arrow-format/README.rst000066400000000000000000000016271515451615200212130ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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 is a copy of https://github.com/apache/arrow/tree/main/format at revision 2e6fe1a4ccfc4ffc63e5d0dc09f972ebbca12426. apache-arrow-java-19.0.0/arrow-format/Schema.fbs000066400000000000000000000524041515451615200214170ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. /// Logical types, vector layouts, and schemas /// Format Version History. /// Version 1.0 - Forward and backwards compatibility guaranteed. /// Version 1.1 - Add Decimal256. /// Version 1.2 - Add Interval MONTH_DAY_NANO. /// Version 1.3 - Add Run-End Encoded. /// Version 1.4 - Add BinaryView, Utf8View, variadicBufferCounts, ListView, and /// LargeListView. /// Version 1.5 - Add 32-bit and 64-bit as allowed bit widths for Decimal namespace org.apache.arrow.flatbuf; enum MetadataVersion:short { /// 0.1.0 (October 2016). V1, /// 0.2.0 (February 2017). Non-backwards compatible with V1. V2, /// 0.3.0 -> 0.7.1 (May - December 2017). Non-backwards compatible with V2. V3, /// >= 0.8.0 (December 2017). Non-backwards compatible with V3. V4, /// >= 1.0.0 (July 2020). Backwards compatible with V4 (V5 readers can read V4 /// metadata and IPC messages). Implementations are recommended to provide a /// V4 compatibility mode with V5 format changes disabled. /// /// Incompatible changes between V4 and V5: /// - Union buffer layout has changed. In V5, Unions don't have a validity /// bitmap buffer. V5, } /// Represents Arrow Features that might not have full support /// within implementations. This is intended to be used in /// two scenarios: /// 1. A mechanism for readers of Arrow Streams /// and files to understand that the stream or file makes /// use of a feature that isn't supported or unknown to /// the implementation (and therefore can meet the Arrow /// forward compatibility guarantees). /// 2. A means of negotiating between a client and server /// what features a stream is allowed to use. The enums /// values here are intented to represent higher level /// features, additional details maybe negotiated /// with key-value pairs specific to the protocol. /// /// Enums added to this list should be assigned power-of-two values /// to facilitate exchanging and comparing bitmaps for supported /// features. enum Feature : long { /// Needed to make flatbuffers happy. UNUSED = 0, /// The stream makes use of multiple full dictionaries with the /// same ID and assumes clients implement dictionary replacement /// correctly. DICTIONARY_REPLACEMENT = 1, /// The stream makes use of compressed bodies as described /// in Message.fbs. COMPRESSED_BODY = 2 } /// These are stored in the flatbuffer in the Type union below table Null { } /// A Struct_ in the flatbuffer metadata is the same as an Arrow Struct /// (according to the physical memory layout). We used Struct_ here as /// Struct is a reserved word in Flatbuffers table Struct_ { } table List { } /// Same as List, but with 64-bit offsets, allowing to represent /// extremely large data values. table LargeList { } /// Represents the same logical types that List can, but contains offsets and /// sizes allowing for writes in any order and sharing of child values among /// list values. table ListView { } /// Same as ListView, but with 64-bit offsets and sizes, allowing to represent /// extremely large data values. table LargeListView { } table FixedSizeList { /// Number of list items per value listSize: int; } /// A Map is a logical nested type that is represented as /// /// List> /// /// In this layout, the keys and values are each respectively contiguous. We do /// not constrain the key and value types, so the application is responsible /// for ensuring that the keys are hashable and unique. Whether the keys are sorted /// may be set in the metadata for this field. /// /// In a field with Map type, the field has a child Struct field, which then /// has two children: key type and the second the value type. The names of the /// child fields may be respectively "entries", "key", and "value", but this is /// not enforced. /// /// Map /// ```text /// - child[0] entries: Struct /// - child[0] key: K /// - child[1] value: V /// ``` /// Neither the "entries" field nor the "key" field may be nullable. /// /// The metadata is structured so that Arrow systems without special handling /// for Map can make Map an alias for List. The "layout" attribute for the Map /// field must have the same contents as a List. table Map { /// Set to true if the keys within each value are sorted keysSorted: bool; } enum UnionMode:short { Sparse, Dense } /// A union is a complex type with children in Field /// By default ids in the type vector refer to the offsets in the children /// optionally typeIds provides an indirection between the child offset and the type id /// for each child `typeIds[offset]` is the id used in the type vector table Union { mode: UnionMode; typeIds: [ int ]; // optional, describes typeid of each child. } table Int { bitWidth: int; // restricted to 8, 16, 32, and 64 in v1 is_signed: bool; } enum Precision:short {HALF, SINGLE, DOUBLE} table FloatingPoint { precision: Precision; } /// Unicode with UTF-8 encoding table Utf8 { } /// Opaque binary data table Binary { } /// Same as Utf8, but with 64-bit offsets, allowing to represent /// extremely large data values. table LargeUtf8 { } /// Same as Binary, but with 64-bit offsets, allowing to represent /// extremely large data values. table LargeBinary { } /// Logically the same as Utf8, but the internal representation uses a view /// struct that contains the string length and either the string's entire data /// inline (for small strings) or an inlined prefix, an index of another buffer, /// and an offset pointing to a slice in that buffer (for non-small strings). /// /// Since it uses a variable number of data buffers, each Field with this type /// must have a corresponding entry in `variadicBufferCounts`. table Utf8View { } /// Logically the same as Binary, but the internal representation uses a view /// struct that contains the string length and either the string's entire data /// inline (for small strings) or an inlined prefix, an index of another buffer, /// and an offset pointing to a slice in that buffer (for non-small strings). /// /// Since it uses a variable number of data buffers, each Field with this type /// must have a corresponding entry in `variadicBufferCounts`. table BinaryView { } table FixedSizeBinary { /// Number of bytes per value byteWidth: int; } table Bool { } /// Contains two child arrays, run_ends and values. /// The run_ends child array must be a 16/32/64-bit integer array /// which encodes the indices at which the run with the value in /// each corresponding index in the values child array ends. /// Like list/struct types, the value array can be of any type. table RunEndEncoded { } /// Exact decimal value represented as an integer value in two's /// complement. Currently 32-bit (4-byte), 64-bit (8-byte), /// 128-bit (16-byte) and 256-bit (32-byte) integers are used. /// The representation uses the endianness indicated in the Schema. table Decimal { /// Total number of decimal digits precision: int; /// Number of digits after the decimal point "." scale: int; /// Number of bits per value. The accepted widths are 32, 64, 128 and 256. /// We use bitWidth for consistency with Int::bitWidth. bitWidth: int = 128; } enum DateUnit: short { DAY, MILLISECOND } /// Date is either a 32-bit or 64-bit signed integer type representing an /// elapsed time since UNIX epoch (1970-01-01), stored in either of two units: /// /// * Milliseconds (64 bits) indicating UNIX time elapsed since the epoch (no /// leap seconds), where the values are evenly divisible by 86400000 /// * Days (32 bits) since the UNIX epoch table Date { unit: DateUnit = MILLISECOND; } enum TimeUnit: short { SECOND, MILLISECOND, MICROSECOND, NANOSECOND } /// Time is either a 32-bit or 64-bit signed integer type representing an /// elapsed time since midnight, stored in either of four units: seconds, /// milliseconds, microseconds or nanoseconds. /// /// The integer `bitWidth` depends on the `unit` and must be one of the following: /// * SECOND and MILLISECOND: 32 bits /// * MICROSECOND and NANOSECOND: 64 bits /// /// The allowed values are between 0 (inclusive) and 86400 (=24*60*60) seconds /// (exclusive), adjusted for the time unit (for example, up to 86400000 /// exclusive for the MILLISECOND unit). /// This definition doesn't allow for leap seconds. Time values from /// measurements with leap seconds will need to be corrected when ingesting /// into Arrow (for example by replacing the value 86400 with 86399). table Time { unit: TimeUnit = MILLISECOND; bitWidth: int = 32; } /// Timestamp is a 64-bit signed integer representing an elapsed time since a /// fixed epoch, stored in either of four units: seconds, milliseconds, /// microseconds or nanoseconds, and is optionally annotated with a timezone. /// /// Timestamp values do not include any leap seconds (in other words, all /// days are considered 86400 seconds long). /// /// Timestamps with a non-empty timezone /// ------------------------------------ /// /// If a Timestamp column has a non-empty timezone value, its epoch is /// 1970-01-01 00:00:00 (January 1st 1970, midnight) in the *UTC* timezone /// (the Unix epoch), regardless of the Timestamp's own timezone. /// /// Therefore, timestamp values with a non-empty timezone correspond to /// physical points in time together with some additional information about /// how the data was obtained and/or how to display it (the timezone). /// /// For example, the timestamp value 0 with the timezone string "Europe/Paris" /// corresponds to "January 1st 1970, 00h00" in the UTC timezone, but the /// application may prefer to display it as "January 1st 1970, 01h00" in /// the Europe/Paris timezone (which is the same physical point in time). /// /// One consequence is that timestamp values with a non-empty timezone /// can be compared and ordered directly, since they all share the same /// well-known point of reference (the Unix epoch). /// /// Timestamps with an unset / empty timezone /// ----------------------------------------- /// /// If a Timestamp column has no timezone value, its epoch is /// 1970-01-01 00:00:00 (January 1st 1970, midnight) in an *unknown* timezone. /// /// Therefore, timestamp values without a timezone cannot be meaningfully /// interpreted as physical points in time, but only as calendar / clock /// indications ("wall clock time") in an unspecified timezone. /// /// For example, the timestamp value 0 with an empty timezone string /// corresponds to "January 1st 1970, 00h00" in an unknown timezone: there /// is not enough information to interpret it as a well-defined physical /// point in time. /// /// One consequence is that timestamp values without a timezone cannot /// be reliably compared or ordered, since they may have different points of /// reference. In particular, it is *not* possible to interpret an unset /// or empty timezone as the same as "UTC". /// /// Conversion between timezones /// ---------------------------- /// /// If a Timestamp column has a non-empty timezone, changing the timezone /// to a different non-empty value is a metadata-only operation: /// the timestamp values need not change as their point of reference remains /// the same (the Unix epoch). /// /// However, if a Timestamp column has no timezone value, changing it to a /// non-empty value requires to think about the desired semantics. /// One possibility is to assume that the original timestamp values are /// relative to the epoch of the timezone being set; timestamp values should /// then adjusted to the Unix epoch (for example, changing the timezone from /// empty to "Europe/Paris" would require converting the timestamp values /// from "Europe/Paris" to "UTC", which seems counter-intuitive but is /// nevertheless correct). /// /// Guidelines for encoding data from external libraries /// ---------------------------------------------------- /// /// Date & time libraries often have multiple different data types for temporal /// data. In order to ease interoperability between different implementations the /// Arrow project has some recommendations for encoding these types into a Timestamp /// column. /// /// An "instant" represents a physical point in time that has no relevant timezone /// (for example, astronomical data). To encode an instant, use a Timestamp with /// the timezone string set to "UTC", and make sure the Timestamp values /// are relative to the UTC epoch (January 1st 1970, midnight). /// /// A "zoned date-time" represents a physical point in time annotated with an /// informative timezone (for example, the timezone in which the data was /// recorded). To encode a zoned date-time, use a Timestamp with the timezone /// string set to the name of the timezone, and make sure the Timestamp values /// are relative to the UTC epoch (January 1st 1970, midnight). /// /// (There is some ambiguity between an instant and a zoned date-time with the /// UTC timezone. Both of these are stored the same in Arrow. Typically, /// this distinction does not matter. If it does, then an application should /// use custom metadata or an extension type to distinguish between the two cases.) /// /// An "offset date-time" represents a physical point in time combined with an /// explicit offset from UTC. To encode an offset date-time, use a Timestamp /// with the timezone string set to the numeric timezone offset string /// (e.g. "+03:00"), and make sure the Timestamp values are relative to /// the UTC epoch (January 1st 1970, midnight). /// /// A "naive date-time" (also called "local date-time" in some libraries) /// represents a wall clock time combined with a calendar date, but with /// no indication of how to map this information to a physical point in time. /// Naive date-times must be handled with care because of this missing /// information, and also because daylight saving time (DST) may make /// some values ambiguous or nonexistent. A naive date-time may be /// stored as a struct with Date and Time fields. However, it may also be /// encoded into a Timestamp column with an empty timezone. The timestamp /// values should be computed "as if" the timezone of the date-time values /// was UTC; for example, the naive date-time "January 1st 1970, 00h00" would /// be encoded as timestamp value 0. table Timestamp { unit: TimeUnit; /// The timezone is an optional string indicating the name of a timezone, /// one of: /// /// * As used in the Olson timezone database (the "tz database" or /// "tzdata"), such as "America/New_York". /// * An absolute timezone offset of the form "+XX:XX" or "-XX:XX", /// such as "+07:30". /// /// Whether a timezone string is present indicates different semantics about /// the data (see above). timezone: string; } enum IntervalUnit: short { YEAR_MONTH, DAY_TIME, MONTH_DAY_NANO} // A "calendar" interval which models types that don't necessarily // have a precise duration without the context of a base timestamp (e.g. // days can differ in length during day light savings time transitions). // All integers in the types below are stored in the endianness indicated // by the schema. // // YEAR_MONTH - Indicates the number of elapsed whole months, stored as // 4-byte signed integers. // DAY_TIME - Indicates the number of elapsed days and milliseconds (no leap seconds), // stored as 2 contiguous 32-bit signed integers (8-bytes in total). Support // of this IntervalUnit is not required for full arrow compatibility. // MONTH_DAY_NANO - A triple of the number of elapsed months, days, and nanoseconds. // The values are stored contiguously in 16-byte blocks. Months and days are // encoded as 32-bit signed integers and nanoseconds is encoded as a 64-bit // signed integer. Nanoseconds does not allow for leap seconds. Each field is // independent (e.g. there is no constraint that nanoseconds have the same // sign as days or that the quantity of nanoseconds represents less than a // day's worth of time). table Interval { unit: IntervalUnit; } // An absolute length of time unrelated to any calendar artifacts. // // For the purposes of Arrow Implementations, adding this value to a Timestamp // ("t1") naively (i.e. simply summing the two number) is acceptable even // though in some cases the resulting Timestamp (t2) would not account for // leap-seconds during the elapsed time between "t1" and "t2". Similarly, // representing the difference between two Unix timestamp is acceptable, but // would yield a value that is possibly a few seconds off from the true elapsed // time. // // The resolution defaults to millisecond, but can be any of the other // supported TimeUnit values as with Timestamp and Time types. This type is // always represented as an 8-byte integer. table Duration { unit: TimeUnit = MILLISECOND; } /// ---------------------------------------------------------------------- /// Top-level Type value, enabling extensible type-specific metadata. We can /// add new logical types to Type without breaking backwards compatibility union Type { Null, Int, FloatingPoint, Binary, Utf8, Bool, Decimal, Date, Time, Timestamp, Interval, List, Struct_, Union, FixedSizeBinary, FixedSizeList, Map, Duration, LargeBinary, LargeUtf8, LargeList, RunEndEncoded, BinaryView, Utf8View, ListView, LargeListView, } /// ---------------------------------------------------------------------- /// user defined key value pairs to add custom metadata to arrow /// key namespacing is the responsibility of the user table KeyValue { key: string; value: string; } /// ---------------------------------------------------------------------- /// Dictionary encoding metadata /// Maintained for forwards compatibility, in the future /// Dictionaries might be explicit maps between integers and values /// allowing for non-contiguous index values enum DictionaryKind : short { DenseArray } table DictionaryEncoding { /// The known dictionary id in the application where this data is used. In /// the file or streaming formats, the dictionary ids are found in the /// DictionaryBatch messages id: long; /// The dictionary indices are constrained to be non-negative integers. If /// this field is null, the indices must be signed int32. To maximize /// cross-language compatibility and performance, implementations are /// recommended to prefer signed integer types over unsigned integer types /// and to avoid uint64 indices unless they are required by an application. indexType: Int; /// By default, dictionaries are not ordered, or the order does not have /// semantic meaning. In some statistical, applications, dictionary-encoding /// is used to represent ordered categorical data, and we provide a way to /// preserve that metadata here isOrdered: bool; dictionaryKind: DictionaryKind; } /// ---------------------------------------------------------------------- /// A field represents a named column in a record / row batch or child of a /// nested type. table Field { /// Name is not required, in i.e. a List name: string; /// Whether or not this field can contain nulls. Should be true in general. nullable: bool; /// This is the type of the decoded value if the field is dictionary encoded. type: Type; /// Present only if the field is dictionary encoded. dictionary: DictionaryEncoding; /// children apply only to nested data types like Struct, List and Union. For /// primitive types children will have length 0. children: [ Field ]; /// User-defined metadata custom_metadata: [ KeyValue ]; } /// ---------------------------------------------------------------------- /// Endianness of the platform producing the data enum Endianness:short { Little, Big } /// ---------------------------------------------------------------------- /// A Buffer represents a single contiguous memory segment struct Buffer { /// The relative offset into the shared memory page where the bytes for this /// buffer starts offset: long; /// The absolute length (in bytes) of the memory buffer. The memory is found /// from offset (inclusive) to offset + length (non-inclusive). When building /// messages using the encapsulated IPC message, padding bytes may be written /// after a buffer, but such padding bytes do not need to be accounted for in /// the size here. length: long; } /// ---------------------------------------------------------------------- /// A Schema describes the columns in a row batch table Schema { /// endianness of the buffer /// it is Little Endian by default /// if endianness doesn't match the underlying system then the vectors need to be converted endianness: Endianness=Little; fields: [Field]; // User-defined metadata custom_metadata: [ KeyValue ]; /// Features used in the stream/file. features : [ Feature ]; } root_type Schema; apache-arrow-java-19.0.0/arrow-format/SparseTensor.fbs000066400000000000000000000172011515451615200226430ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. /// EXPERIMENTAL: Metadata for n-dimensional sparse arrays, aka "sparse tensors". /// Arrow implementations in general are not required to implement this type include "Tensor.fbs"; namespace org.apache.arrow.flatbuf; /// ---------------------------------------------------------------------- /// EXPERIMENTAL: Data structures for sparse tensors /// Coordinate (COO) format of sparse tensor index. /// /// COO's index list are represented as a NxM matrix, /// where N is the number of non-zero values, /// and M is the number of dimensions of a sparse tensor. /// /// indicesBuffer stores the location and size of the data of this indices /// matrix. The value type and the stride of the indices matrix is /// specified in indicesType and indicesStrides fields. /// /// For example, let X be a 2x3x4x5 tensor, and it has the following /// 6 non-zero values: /// ```text /// X[0, 1, 2, 0] := 1 /// X[1, 1, 2, 3] := 2 /// X[0, 2, 1, 0] := 3 /// X[0, 1, 3, 0] := 4 /// X[0, 1, 2, 1] := 5 /// X[1, 2, 0, 4] := 6 /// ``` /// In COO format, the index matrix of X is the following 4x6 matrix: /// ```text /// [[0, 0, 0, 0, 1, 1], /// [1, 1, 1, 2, 1, 2], /// [2, 2, 3, 1, 2, 0], /// [0, 1, 0, 0, 3, 4]] /// ``` /// When isCanonical is true, the indices is sorted in lexicographical order /// (row-major order), and it does not have duplicated entries. Otherwise, /// the indices may not be sorted, or may have duplicated entries. table SparseTensorIndexCOO { /// The type of values in indicesBuffer indicesType: Int (required); /// Non-negative byte offsets to advance one value cell along each dimension /// If omitted, default to row-major order (C-like). indicesStrides: [long]; /// The location and size of the indices matrix's data indicesBuffer: Buffer (required); /// This flag is true if and only if the indices matrix is sorted in /// row-major order, and does not have duplicated entries. /// This sort order is the same as of Tensorflow's SparseTensor, /// but it is inverse order of SciPy's canonical coo_matrix /// (SciPy employs column-major order for its coo_matrix). isCanonical: bool; } enum SparseMatrixCompressedAxis: short { Row, Column } /// Compressed Sparse format, that is matrix-specific. table SparseMatrixIndexCSX { /// Which axis, row or column, is compressed compressedAxis: SparseMatrixCompressedAxis; /// The type of values in indptrBuffer indptrType: Int (required); /// indptrBuffer stores the location and size of indptr array that /// represents the range of the rows. /// The i-th row spans from `indptr[i]` to `indptr[i+1]` in the data. /// The length of this array is 1 + (the number of rows), and the type /// of index value is long. /// /// For example, let X be the following 6x4 matrix: /// ```text /// X := [[0, 1, 2, 0], /// [0, 0, 3, 0], /// [0, 4, 0, 5], /// [0, 0, 0, 0], /// [6, 0, 7, 8], /// [0, 9, 0, 0]]. /// ``` /// The array of non-zero values in X is: /// ```text /// values(X) = [1, 2, 3, 4, 5, 6, 7, 8, 9]. /// ``` /// And the indptr of X is: /// ```text /// indptr(X) = [0, 2, 3, 5, 5, 8, 10]. /// ``` indptrBuffer: Buffer (required); /// The type of values in indicesBuffer indicesType: Int (required); /// indicesBuffer stores the location and size of the array that /// contains the column indices of the corresponding non-zero values. /// The type of index value is long. /// /// For example, the indices of the above X is: /// ```text /// indices(X) = [1, 2, 2, 1, 3, 0, 2, 3, 1]. /// ``` /// Note that the indices are sorted in lexicographical order for each row. indicesBuffer: Buffer (required); } /// Compressed Sparse Fiber (CSF) sparse tensor index. table SparseTensorIndexCSF { /// CSF is a generalization of compressed sparse row (CSR) index. /// See [smith2017knl](http://shaden.io/pub-files/smith2017knl.pdf) /// /// CSF index recursively compresses each dimension of a tensor into a set /// of prefix trees. Each path from a root to leaf forms one tensor /// non-zero index. CSF is implemented with two arrays of buffers and one /// arrays of integers. /// /// For example, let X be a 2x3x4x5 tensor and let it have the following /// 8 non-zero values: /// ```text /// X[0, 0, 0, 1] := 1 /// X[0, 0, 0, 2] := 2 /// X[0, 1, 0, 0] := 3 /// X[0, 1, 0, 2] := 4 /// X[0, 1, 1, 0] := 5 /// X[1, 1, 1, 0] := 6 /// X[1, 1, 1, 1] := 7 /// X[1, 1, 1, 2] := 8 /// ``` /// As a prefix tree this would be represented as: /// ```text /// 0 1 /// / \ | /// 0 1 1 /// / / \ | /// 0 0 1 1 /// /| /| | /| | /// 1 2 0 2 0 0 1 2 /// ``` /// The type of values in indptrBuffers indptrType: Int (required); /// indptrBuffers stores the sparsity structure. /// Each two consecutive dimensions in a tensor correspond to a buffer in /// indptrBuffers. A pair of consecutive values at `indptrBuffers[dim][i]` /// and `indptrBuffers[dim][i + 1]` signify a range of nodes in /// `indicesBuffers[dim + 1]` who are children of `indicesBuffers[dim][i]` node. /// /// For example, the indptrBuffers for the above X is: /// ```text /// indptrBuffer(X) = [ /// [0, 2, 3], /// [0, 1, 3, 4], /// [0, 2, 4, 5, 8] /// ]. /// ``` indptrBuffers: [Buffer] (required); /// The type of values in indicesBuffers indicesType: Int (required); /// indicesBuffers stores values of nodes. /// Each tensor dimension corresponds to a buffer in indicesBuffers. /// For example, the indicesBuffers for the above X is: /// ```text /// indicesBuffer(X) = [ /// [0, 1], /// [0, 1, 1], /// [0, 0, 1, 1], /// [1, 2, 0, 2, 0, 0, 1, 2] /// ]. /// ``` indicesBuffers: [Buffer] (required); /// axisOrder stores the sequence in which dimensions were traversed to /// produce the prefix tree. /// For example, the axisOrder for the above X is: /// ```text /// axisOrder(X) = [0, 1, 2, 3]. /// ``` axisOrder: [int] (required); } union SparseTensorIndex { SparseTensorIndexCOO, SparseMatrixIndexCSX, SparseTensorIndexCSF } table SparseTensor { /// The type of data contained in a value cell. /// Currently only fixed-width value types are supported, /// no strings or nested types. type: Type (required); /// The dimensions of the tensor, optionally named. shape: [TensorDim] (required); /// The number of non-zero values in a sparse tensor. non_zero_length: long; /// Sparse tensor index sparseIndex: SparseTensorIndex (required); /// The location and size of the tensor's data data: Buffer (required); } root_type SparseTensor; apache-arrow-java-19.0.0/arrow-format/Tensor.fbs000066400000000000000000000034311515451615200214650ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. /// EXPERIMENTAL: Metadata for n-dimensional arrays, aka "tensors" or /// "ndarrays". Arrow implementations in general are not required to implement /// this type include "Schema.fbs"; namespace org.apache.arrow.flatbuf; /// ---------------------------------------------------------------------- /// Data structures for dense tensors /// Shape data for a single axis in a tensor table TensorDim { /// Length of dimension size: long; /// Name of the dimension, optional name: string; } table Tensor { /// The type of data contained in a value cell. Currently only fixed-width /// value types are supported, no strings or nested types type: Type (required); /// The dimensions of the tensor, optionally named shape: [TensorDim] (required); /// Non-negative byte offsets to advance one value cell along each dimension /// If omitted, default to row-major order (C-like). strides: [long]; /// The location and size of the tensor's data data: Buffer (required); } root_type Tensor; apache-arrow-java-19.0.0/arrow-format/substrait/000077500000000000000000000000001515451615200215365ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-format/substrait/extension_types.yaml000066400000000000000000000130641515451615200256660ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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::{ExtensionTypeVariation, ExtensionType}s # for wrapping types which appear in the arrow type system but # are not first-class in substrait. These include: # - null # - unsigned integers # - half-precision floating point numbers # - 32-bit times and 64-bit dates # - timestamps with units other than microseconds # - timestamps with timezones other than UTC # - 256-bit decimals # - sparse and dense unions # - dictionary encoded types # - durations # - string and binary with 64 bit offsets # - list with 64-bit offsets # - interval # - interval # - interval # - arrow::ExtensionTypes # # These types fall into several categories of behavior: # Certain Arrow data types are, from Substrait's point of view, encodings. # These include dictionary, the view types (e.g. binary view, list view), # and REE. # # These types are not logically distinct from the type they are encoding. # Specifically, the types meet the following criteria: # * There is no value in the decoded type that cannot be represented # as a value in the encoded type and vice versa. # * Functions have the same meaning when applied to the encoded type # # Note: if two types have a different range (e.g. string and large_string) then # they do not satisfy the above criteria and are not encodings. # # These types will never have a Substrait equivalent. In the Substrait point # of view these are execution details. # The following types are encodings: # binary_view # list_view # dictionary # ree # Arrow-cpp's Substrait serde does not yet handle parameterized UDTs. This means # the following types are not yet supported but may be supported in the future. # We define them below in case other implementations support them in the meantime. # decimal256 # large_list # fixed_size_list # duration # Other types are not encodings, but are not first-class in Substrait. These # types are often similar to existing Substrait types but define a different range # of values. For example, unsigned integer types are very similar to their integer # counterparts, but have a different range of values. These types are defined here # as extension types. # # A full description of the types, along with their specified range, can be found # in Schema.fbs # # Consumers should take care when supporting the below types. Should Substrait decide # later to support these types, the consumer will need to make sure to continue supporting # the extension type names as aliases for proper backwards compatibility. types: - name: "null" structure: {} - name: interval_month structure: months: i32 - name: interval_day_milli structure: days: i32 millis: i32 - name: interval_month_day_nano structure: months: i32 days: i32 nanos: i64 # All unsigned integer literals are encoded as user defined literals with # a google.protobuf.UInt64Value message. - name: u8 structure: {} - name: u16 structure: {} - name: u32 structure: {} - name: u64 structure: {} # fp16 literals are encoded as user defined literals with # a google.protobuf.UInt32Value message where the lower 16 bits are # the fp16 value. - name: fp16 structure: {} # 64-bit integers are big. Even though date64 stores ms and not days it # can still represent about 50x more dates than date32. Since it has a # different range of values, it is an extension type. # # date64 literals are encoded as user defined literals with # a google.protobuf.Int64Value message. - name: date_millis structure: {} # time literals are encoded as user defined literals with # a google.protobuf.Int32Value message (for time_seconds/time_millis) # or a google.protobuf.Int64Value message (for time_nanos). - name: time_seconds structure: {} - name: time_millis structure: {} - name: time_nanos structure: {} # Large string literals are encoded using a # google.protobuf.StringValue message. - name: large_string structure: {} # Large binary literals are encoded using a # google.protobuf.BytesValue message. - name: large_binary structure: {} # We cannot generate these today because they are parameterized UDTs and # substrait-cpp does not yet support parameterized UDTs. - name: decimal256 structure: {} parameters: - name: precision type: integer min: 0 max: 76 - name: scale type: integer min: 0 max: 76 - name: large_list structure: {} parameters: - name: value_type type: dataType - name: fixed_size_list structure: {} parameters: - name: value_type type: dataType - name: dimension type: integer min: 0 - name: duration structure: {} parameters: - name: unit type: string apache-arrow-java-19.0.0/arrow-variant/000077500000000000000000000000001515451615200176725ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/pom.xml000066400000000000000000000036161515451615200212150ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-java-root 19.0.0 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 apache-arrow-java-19.0.0/arrow-variant/src/000077500000000000000000000000001515451615200204615ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/main/000077500000000000000000000000001515451615200214055ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/main/java/000077500000000000000000000000001515451615200223265ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/main/java/module-info.java000066400000000000000000000021771515451615200254160ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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; } apache-arrow-java-19.0.0/arrow-variant/src/main/java/org/000077500000000000000000000000001515451615200231155ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/main/java/org/apache/000077500000000000000000000000001515451615200243365ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200254705ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/main/java/org/apache/arrow/variant/000077500000000000000000000000001515451615200271345ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/main/java/org/apache/arrow/variant/Variant.java000066400000000000000000000131541515451615200314070ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; } } } apache-arrow-java-19.0.0/arrow-variant/src/main/java/org/apache/arrow/variant/extension/000077500000000000000000000000001515451615200311505ustar00rootroot00000000000000VariantType.java000066400000000000000000000060001515451615200341760ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/main/java/org/apache/arrow/variant/extension/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } VariantVector.java000066400000000000000000000306671515451615200345370ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/main/java/org/apache/arrow/variant/extension/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } } apache-arrow-java-19.0.0/arrow-variant/src/main/java/org/apache/arrow/variant/holders/000077500000000000000000000000001515451615200305745ustar00rootroot00000000000000NullableVariantHolder.java000066400000000000000000000033361515451615200356060ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/main/java/org/apache/arrow/variant/holders/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; } } VariantHolder.java000066400000000000000000000033301515451615200341210ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/main/java/org/apache/arrow/variant/holders/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; } } apache-arrow-java-19.0.0/arrow-variant/src/main/java/org/apache/arrow/variant/impl/000077500000000000000000000000001515451615200300755ustar00rootroot00000000000000NullableVariantHolderReaderImpl.java000066400000000000000000000044161515451615200370540ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/main/java/org/apache/arrow/variant/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; } } VariantReaderImpl.java000066400000000000000000000044541515451615200342410ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/main/java/org/apache/arrow/variant/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } VariantWriterImpl.java000066400000000000000000000117161515451615200343120ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/main/java/org/apache/arrow/variant/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } apache-arrow-java-19.0.0/arrow-variant/src/test/000077500000000000000000000000001515451615200214405ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/test/java/000077500000000000000000000000001515451615200223615ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/test/java/org/000077500000000000000000000000001515451615200231505ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/test/java/org/apache/000077500000000000000000000000001515451615200243715ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200255235ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/test/java/org/apache/arrow/variant/000077500000000000000000000000001515451615200271675ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/test/java/org/apache/arrow/variant/TestVariant.java000066400000000000000000000330211515451615200322750ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } } apache-arrow-java-19.0.0/arrow-variant/src/test/java/org/apache/arrow/variant/extension/000077500000000000000000000000001515451615200312035ustar00rootroot00000000000000TestVariantExtensionType.java000066400000000000000000000241411515451615200367740ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/test/java/org/apache/arrow/variant/extension/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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)); } } } TestVariantInListVector.java000066400000000000000000000161271515451615200365500ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/test/java/org/apache/arrow/variant/extension/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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)); } } } } TestVariantInMapVector.java000066400000000000000000000107341515451615200363500ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/test/java/org/apache/arrow/variant/extension/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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)); } } } TestVariantType.java000066400000000000000000000256541515451615200351110ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/test/java/org/apache/arrow/variant/extension/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } } TestVariantVector.java000066400000000000000000000673461515451615200354360ustar00rootroot00000000000000apache-arrow-java-19.0.0/arrow-variant/src/test/java/org/apache/arrow/variant/extension/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } } apache-arrow-java-19.0.0/bom/000077500000000000000000000000001515451615200156535ustar00rootroot00000000000000apache-arrow-java-19.0.0/bom/pom.xml000066400000000000000000000237331515451615200172000ustar00rootroot00000000000000 4.0.0 org.apache apache 33 org.apache.arrow arrow-bom 19.0.0 pom Arrow Bill of Materials Arrow Bill of Materials https://arrow.apache.org/ Developer List dev-subscribe@arrow.apache.org dev-unsubscribe@arrow.apache.org dev@arrow.apache.org https://lists.apache.org/list.html?dev@arrow.apache.org Commits List commits-subscribe@arrow.apache.org commits-unsubscribe@arrow.apache.org commits@arrow.apache.org https://lists.apache.org/list.html?commits@arrow.apache.org Issues List issues-subscribe@arrow.apache.org issues-unsubscribe@arrow.apache.org https://lists.apache.org/list.html?issues@arrow.apache.org GitHub List github-subscribe@arrow.apache.org github-unsubscribe@arrow.apache.org https://lists.apache.org/list.html?github@arrow.apache.org scm:git:https://github.com/apache/arrow-java.git scm:git:https://github.com/apache/arrow-java.git v19.0.0 https://github.com/apache/arrow-java/tree/${project.scm.tag} GitHub https://github.com/apache/arrow-java/issues 11 11 11 11 org.apache.arrow arrow-vector ${project.version} org.apache.arrow arrow-vector ${project.version} ${arrow.vector.classifier} org.apache.arrow arrow-avro ${project.version} org.apache.arrow arrow-jdbc ${project.version} org.apache.arrow arrow-orc ${project.version} org.apache.arrow arrow-algorithm ${project.version} org.apache.arrow arrow-c-data ${project.version} org.apache.arrow arrow-compression ${project.version} org.apache.arrow arrow-dataset ${project.version} org.apache.arrow flight-core ${project.version} org.apache.arrow flight-integration-tests ${project.version} org.apache.arrow flight-sql ${project.version} org.apache.arrow flight-sql-jdbc-core ${project.version} org.apache.arrow flight-sql-jdbc-driver ${project.version} org.apache.arrow arrow-format ${project.version} org.apache.arrow.gandiva arrow-gandiva ${project.version} org.apache.arrow arrow-memory-core ${project.version} org.apache.arrow arrow-memory-netty ${project.version} org.apache.arrow arrow-memory-unsafe ${project.version} org.apache.arrow arrow-performance ${project.version} org.apache.arrow arrow-tools ${project.version} org.apache.arrow arrow-variant ${project.version} com.diffplug.spotless spotless-maven-plugin 2.44.4 org.codehaus.mojo versions-maven-plugin 2.21.0 org.apache.maven.plugins maven-project-info-reports-plugin org.apache.maven.plugins maven-site-plugin com.diffplug.spotless spotless-maven-plugin ${maven.multiModuleProjectDirectory}/dev/license/asf-xml.license (<configuration|<project) spotless-check check org.apache.maven.plugins maven-project-info-reports-plugin org.apache.maven.plugins maven-site-plugin apache-release org.apache.maven.plugins maven-assembly-plugin source-release-assembly true apache-arrow-java-19.0.0/c/000077500000000000000000000000001515451615200153205ustar00rootroot00000000000000apache-arrow-java-19.0.0/c/CMakeLists.txt000066400000000000000000000035131515451615200200620ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. include_directories(${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_SOURCE_DIR} ${JNI_INCLUDE_DIRS} ${JNI_HEADERS_DIR}) add_jar(arrow_java_jni_cdata_jar src/main/java/org/apache/arrow/c/jni/CDataJniException.java src/main/java/org/apache/arrow/c/jni/JniLoader.java src/main/java/org/apache/arrow/c/jni/JniWrapper.java src/main/java/org/apache/arrow/c/jni/PrivateData.java GENERATE_NATIVE_HEADERS arrow_java_jni_cdata_headers) add_library(arrow_java_jni_cdata SHARED src/main/cpp/jni_wrapper.cc) set_property(TARGET arrow_java_jni_cdata PROPERTY OUTPUT_NAME "arrow_cdata_jni") target_link_libraries(arrow_java_jni_cdata arrow_java_jni_cdata_headers jni) set(ARROW_JAVA_JNI_C_LIBDIR "${CMAKE_INSTALL_PREFIX}/lib/arrow_cdata_jni/${ARROW_JAVA_JNI_ARCH_DIR}") set(ARROW_JAVA_JNI_C_BINDIR "${CMAKE_INSTALL_PREFIX}/bin/arrow_cdata_jni/${ARROW_JAVA_JNI_ARCH_DIR}") install(TARGETS arrow_java_jni_cdata LIBRARY DESTINATION ${ARROW_JAVA_JNI_C_LIBDIR} RUNTIME DESTINATION ${ARROW_JAVA_JNI_C_BINDIR}) apache-arrow-java-19.0.0/c/README.md000066400000000000000000000024361515451615200166040ustar00rootroot00000000000000 # C Interfaces for Arrow Java ## Setup Build Environment install: - Java 8 or later - Maven 3.3 or later - A C++17-enabled compiler - CMake 3.11 or later - Make or ninja build utilities ## Building JNI wrapper shared library ``` mkdir -p build pushd build cmake .. cmake --build . popd ``` ## Building and running tests Run tests with ``` mvn test ``` To install Apache Arrow (Java) with this module enabled run the following from the project root directory: ``` cd java mvn -Parrow-c-data install ``` apache-arrow-java-19.0.0/c/pom.xml000066400000000000000000000066051515451615200166440ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-java-root 19.0.0 arrow-c-data jar Arrow Java C Data Interface Java implementation of C Data Interface ./build org.apache.arrow arrow-vector ${arrow.vector.classifier} compile org.apache.arrow arrow-vector ${project.version} test-jar test org.apache.arrow arrow-memory-core compile org.slf4j slf4j-api org.immutables value-annotations org.apache.arrow arrow-memory-unsafe test org.apache.arrow arrow-format test com.google.guava guava test org.assertj assertj-core test ${arrow.c.jni.dist.dir} **/*arrow_cdata_jni.* org.apache.maven.plugins maven-compiler-plugin -Werror apache-arrow-java-19.0.0/c/src/000077500000000000000000000000001515451615200161075ustar00rootroot00000000000000apache-arrow-java-19.0.0/c/src/main/000077500000000000000000000000001515451615200170335ustar00rootroot00000000000000apache-arrow-java-19.0.0/c/src/main/cpp/000077500000000000000000000000001515451615200176155ustar00rootroot00000000000000apache-arrow-java-19.0.0/c/src/main/cpp/abi.h000066400000000000000000000064141515451615200205260ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #pragma once #include #ifdef __cplusplus extern "C" { #endif #ifndef ARROW_C_DATA_INTERFACE #define ARROW_C_DATA_INTERFACE #define ARROW_FLAG_DICTIONARY_ORDERED 1 #define ARROW_FLAG_NULLABLE 2 #define ARROW_FLAG_MAP_KEYS_SORTED 4 struct ArrowSchema { // Array type description const char* format; const char* name; const char* metadata; int64_t flags; int64_t n_children; struct ArrowSchema** children; struct ArrowSchema* dictionary; // Release callback void (*release)(struct ArrowSchema*); // Opaque producer-specific data void* private_data; }; struct ArrowArray { // Array data description int64_t length; int64_t null_count; int64_t offset; int64_t n_buffers; int64_t n_children; const void** buffers; struct ArrowArray** children; struct ArrowArray* dictionary; // Release callback void (*release)(struct ArrowArray*); // Opaque producer-specific data void* private_data; }; #endif // ARROW_C_DATA_INTERFACE #ifndef ARROW_C_STREAM_INTERFACE #define ARROW_C_STREAM_INTERFACE struct ArrowArrayStream { // Callback to get the stream type // (will be the same for all arrays in the stream). // // Return value: 0 if successful, an `errno`-compatible error code otherwise. // // If successful, the ArrowSchema must be released independently from the stream. int (*get_schema)(struct ArrowArrayStream*, struct ArrowSchema* out); // Callback to get the next array // (if no error and the array is released, the stream has ended) // // Return value: 0 if successful, an `errno`-compatible error code otherwise. // // If successful, the ArrowArray must be released independently from the stream. int (*get_next)(struct ArrowArrayStream*, struct ArrowArray* out); // Callback to get optional detailed error information. // This must only be called if the last stream operation failed // with a non-0 return code. // // Return value: pointer to a null-terminated character array describing // the last error, or NULL if no description is available. // // The returned pointer is only valid until the next operation on this stream // (including release). const char* (*get_last_error)(struct ArrowArrayStream*); // Release callback: release the stream's own resources. // Note that arrays returned by `get_next` must be individually released. void (*release)(struct ArrowArrayStream*); // Opaque producer-specific data void* private_data; }; #endif // ARROW_C_STREAM_INTERFACE #ifdef __cplusplus } #endif apache-arrow-java-19.0.0/c/src/main/cpp/jni_wrapper.cc000066400000000000000000000413411515451615200224470ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #include #include #include #include #include #include #include #include "./abi.h" #include "org_apache_arrow_c_jni_JniWrapper.h" namespace { jclass kObjectClass; jclass kRuntimeExceptionClass; jclass kPrivateDataClass; jclass kCDataExceptionClass; jclass kStreamPrivateDataClass; jfieldID kPrivateDataLastErrorField; jmethodID kObjectToStringMethod; jmethodID kPrivateDataCloseMethod; jmethodID kPrivateDataGetNextMethod; jmethodID kPrivateDataGetSchemaMethod; jmethodID kCDataExceptionConstructor; jint JNI_VERSION = JNI_VERSION_10; class JniPendingException : public std::runtime_error { public: explicit JniPendingException(const std::string& arg) : std::runtime_error(arg) {} }; void ThrowPendingException(const std::string& message) { throw JniPendingException(message); } void JniThrow(std::string message) { ThrowPendingException(message); } jclass CreateGlobalClassReference(JNIEnv* env, const char* class_name) { jclass local_class = env->FindClass(class_name); if (!local_class) { std::string message = "Could not find class "; message += class_name; ThrowPendingException(message); } jclass global_class = (jclass)env->NewGlobalRef(local_class); if (!global_class) { std::string message = "Could not create global reference to class "; message += class_name; ThrowPendingException(message); } env->DeleteLocalRef(local_class); return global_class; } jmethodID GetMethodID(JNIEnv* env, jclass this_class, const char* name, const char* sig) { jmethodID ret = env->GetMethodID(this_class, name, sig); if (ret == nullptr) { std::string error_message = "Unable to find method " + std::string(name) + " with signature " + std::string(sig); ThrowPendingException(error_message); } return ret; } jfieldID GetFieldID(JNIEnv* env, jclass this_class, const char* name, const char* sig) { jfieldID fieldId = env->GetFieldID(this_class, name, sig); if (fieldId == nullptr) { std::string error_message = "Unable to find field " + std::string(name) + " with signature " + std::string(sig); ThrowPendingException(error_message); } return fieldId; } class InnerPrivateData { public: InnerPrivateData(JavaVM* vm, jobject private_data) : vm_(vm), j_private_data_(private_data) {} JavaVM* vm_; jobject j_private_data_; // Only for ArrowArrayStream std::string last_error_; }; class JNIEnvGuard { public: explicit JNIEnvGuard(JavaVM* vm) : vm_(vm), should_detach_(false) { JNIEnv* env; jint code = vm->GetEnv(reinterpret_cast(&env), JNI_VERSION); if (code == JNI_EDETACHED) { JavaVMAttachArgs args; args.version = JNI_VERSION; args.name = NULL; args.group = NULL; code = vm->AttachCurrentThread(reinterpret_cast(&env), &args); should_detach_ = (code == JNI_OK); } if (code != JNI_OK) { ThrowPendingException("Failed to attach the current thread to a Java VM"); } env_ = env; } JNIEnv* env() { return env_; } ~JNIEnvGuard() { if (should_detach_) { vm_->DetachCurrentThread(); should_detach_ = false; } } private: bool should_detach_; JavaVM* vm_; JNIEnv* env_; }; template void release_exported(T* base) { // This should not be called on already released structure assert(base->release != nullptr); // Release children for (int64_t i = 0; i < base->n_children; ++i) { T* child = base->children[i]; if (child->release != nullptr) { child->release(child); assert(child->release == nullptr); } } // Release dictionary T* dict = base->dictionary; if (dict != nullptr && dict->release != nullptr) { dict->release(dict); assert(dict->release == nullptr); } // Release all data directly owned by the struct InnerPrivateData* private_data = reinterpret_cast(base->private_data); // It is possible for the JVM to be shut down when this is called; // guard against that. Example: Python code using JPype may shut // down the JVM before releasing the stream. try { JNIEnvGuard guard(private_data->vm_); JNIEnv* env = guard.env(); env->CallObjectMethod(private_data->j_private_data_, kPrivateDataCloseMethod); if (env->ExceptionCheck()) { // Can't signal this to caller, so log and then try to free things // as best we can env->ExceptionDescribe(); env->ExceptionClear(); } env->DeleteGlobalRef(private_data->j_private_data_); } catch (const JniPendingException& e) { std::cerr << "WARNING: Failed to release Java C Data resource: " << e.what() << std::endl; } delete private_data; base->private_data = nullptr; // Mark released base->release = nullptr; } // Attempt to copy the JVM-side lastError to the C++ side void TryCopyLastError(JNIEnv* env, InnerPrivateData* private_data) { jobject error_data = env->GetObjectField(private_data->j_private_data_, kPrivateDataLastErrorField); if (!error_data) { private_data->last_error_.clear(); return; } auto arr = reinterpret_cast(error_data); jbyte* error_bytes = env->GetByteArrayElements(arr, nullptr); if (!error_bytes) { private_data->last_error_.clear(); return; } jsize error_bytes_len = env->GetArrayLength(arr); char* error_str = reinterpret_cast(error_bytes); private_data->last_error_ = std::string(error_str, error_bytes_len); env->ReleaseByteArrayElements(arr, error_bytes, JNI_ABORT); } // Normally the Java side catches all exceptions and populates // lastError. If that fails we check for an exception and try to // populate last_error_ ourselves. void TryHandleUncaughtException(JNIEnv* env, InnerPrivateData* private_data, jthrowable exc) { jstring message = reinterpret_cast(env->CallObjectMethod(exc, kObjectToStringMethod)); if (!message) { private_data->last_error_.clear(); return; } const char* str = env->GetStringUTFChars(message, 0); if (!str) { private_data->last_error_.clear(); return; } private_data->last_error_ = str; env->ReleaseStringUTFChars(message, 0); } int ArrowArrayStreamGetSchema(ArrowArrayStream* stream, ArrowSchema* out) { assert(stream->private_data != nullptr); InnerPrivateData* private_data = reinterpret_cast(stream->private_data); JNIEnvGuard guard(private_data->vm_); JNIEnv* env = guard.env(); const jlong out_addr = static_cast(reinterpret_cast(out)); const int err_code = env->CallIntMethod(private_data->j_private_data_, kPrivateDataGetSchemaMethod, out_addr); if (jthrowable exc = env->ExceptionOccurred()) { TryHandleUncaughtException(env, private_data, exc); env->ExceptionClear(); return EIO; } else if (err_code != 0) { TryCopyLastError(env, private_data); } return err_code; } int ArrowArrayStreamGetNext(ArrowArrayStream* stream, ArrowArray* out) { assert(stream->private_data != nullptr); InnerPrivateData* private_data = reinterpret_cast(stream->private_data); JNIEnvGuard guard(private_data->vm_); JNIEnv* env = guard.env(); const jlong out_addr = static_cast(reinterpret_cast(out)); const int err_code = env->CallIntMethod(private_data->j_private_data_, kPrivateDataGetNextMethod, out_addr); if (jthrowable exc = env->ExceptionOccurred()) { TryHandleUncaughtException(env, private_data, exc); env->ExceptionClear(); return EIO; } else if (err_code != 0) { TryCopyLastError(env, private_data); } return err_code; } const char* ArrowArrayStreamGetLastError(ArrowArrayStream* stream) { assert(stream->private_data != nullptr); InnerPrivateData* private_data = reinterpret_cast(stream->private_data); JNIEnvGuard guard(private_data->vm_); JNIEnv* env = guard.env(); if (private_data->last_error_.empty()) return nullptr; return private_data->last_error_.c_str(); } void ArrowArrayStreamRelease(ArrowArrayStream* stream) { // This should not be called on already released structure assert(stream->release != nullptr); // Release all data directly owned by the struct InnerPrivateData* private_data = reinterpret_cast(stream->private_data); // It is possible for the JVM to be shut down (see above) try { JNIEnvGuard guard(private_data->vm_); JNIEnv* env = guard.env(); env->CallObjectMethod(private_data->j_private_data_, kPrivateDataCloseMethod); if (env->ExceptionCheck()) { env->ExceptionDescribe(); env->ExceptionClear(); } env->DeleteGlobalRef(private_data->j_private_data_); } catch (const JniPendingException& e) { std::cerr << "WARNING: Failed to release Java ArrowArrayStream: " << e.what() << std::endl; } delete private_data; stream->private_data = nullptr; // Mark released stream->release = nullptr; } } // namespace #define JNI_METHOD_START try { // macro ended #define JNI_METHOD_END(fallback_expr) \ } \ catch (JniPendingException & e) { \ env->ThrowNew(kRuntimeExceptionClass, e.what()); \ return fallback_expr; \ } // macro ended jint JNI_OnLoad(JavaVM* vm, void* reserved) { JNIEnv* env; 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, "java/lang/Object"); kRuntimeExceptionClass = CreateGlobalClassReference(env, "java/lang/RuntimeException"); kPrivateDataClass = CreateGlobalClassReference(env, "org/apache/arrow/c/jni/PrivateData"); kCDataExceptionClass = CreateGlobalClassReference(env, "org/apache/arrow/c/jni/CDataJniException"); kStreamPrivateDataClass = CreateGlobalClassReference( env, "org/apache/arrow/c/ArrayStreamExporter$ExportedArrayStreamPrivateData"); kPrivateDataLastErrorField = GetFieldID(env, kStreamPrivateDataClass, "lastError", "[B"); kObjectToStringMethod = GetMethodID(env, kObjectClass, "toString", "()Ljava/lang/String;"); kPrivateDataCloseMethod = GetMethodID(env, kPrivateDataClass, "close", "()V"); kPrivateDataGetNextMethod = GetMethodID(env, kStreamPrivateDataClass, "getNext", "(J)I"); kPrivateDataGetSchemaMethod = GetMethodID(env, kStreamPrivateDataClass, "getSchema", "(J)I"); kCDataExceptionConstructor = GetMethodID(env, kCDataExceptionClass, "", "(ILjava/lang/String;)V"); return JNI_VERSION; JNI_METHOD_END(JNI_ERR) } void JNI_OnUnload(JavaVM* vm, void* reserved) { JNIEnv* env; vm->GetEnv(reinterpret_cast(&env), JNI_VERSION); env->DeleteGlobalRef(kObjectClass); env->DeleteGlobalRef(kRuntimeExceptionClass); env->DeleteGlobalRef(kPrivateDataClass); env->DeleteGlobalRef(kCDataExceptionClass); env->DeleteGlobalRef(kStreamPrivateDataClass); } /* * Class: org_apache_arrow_c_jni_JniWrapper * Method: releaseSchema * Signature: (J)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_c_jni_JniWrapper_releaseSchema( JNIEnv* env, jobject, jlong address) { JNI_METHOD_START ArrowSchema* schema = reinterpret_cast(address); if (schema->release != nullptr) { schema->release(schema); } JNI_METHOD_END() } /* * Class: org_apache_arrow_c_jni_JniWrapper * Method: releaseArray * Signature: (J)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_c_jni_JniWrapper_releaseArray(JNIEnv* env, jobject, jlong address) { JNI_METHOD_START ArrowArray* array = reinterpret_cast(address); if (array->release != nullptr) { array->release(array); } JNI_METHOD_END() } /* * Class: org_apache_arrow_c_jni_JniWrapper * Method: getNextArrayStream * Signature: (JJ)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_c_jni_JniWrapper_getNextArrayStream( JNIEnv* env, jobject, jlong address, jlong out_address) { JNI_METHOD_START auto* stream = reinterpret_cast(address); auto* out = reinterpret_cast(out_address); const int err_code = stream->get_next(stream, out); if (err_code != 0) { const char* message = stream->get_last_error(stream); if (!message) message = std::strerror(err_code); jstring java_message = env->NewStringUTF(message); jthrowable exception = static_cast(env->NewObject( kCDataExceptionClass, kCDataExceptionConstructor, err_code, java_message)); env->Throw(exception); } JNI_METHOD_END() } /* * Class: org_apache_arrow_c_jni_JniWrapper * Method: getSchemaArrayStream * Signature: (JJ)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_c_jni_JniWrapper_getSchemaArrayStream( JNIEnv* env, jobject, jlong address, jlong out_address) { JNI_METHOD_START auto* stream = reinterpret_cast(address); auto* out = reinterpret_cast(out_address); const int err_code = stream->get_schema(stream, out); if (err_code != 0) { const char* message = stream->get_last_error(stream); if (!message) message = std::strerror(err_code); jstring java_message = env->NewStringUTF(message); jthrowable exception = static_cast(env->NewObject( kCDataExceptionClass, kCDataExceptionConstructor, err_code, java_message)); env->Throw(exception); } JNI_METHOD_END() } /* * Class: org_apache_arrow_c_jni_JniWrapper * Method: releaseArrayStream * Signature: (J)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_c_jni_JniWrapper_releaseArrayStream( JNIEnv* env, jobject, jlong address) { JNI_METHOD_START auto* stream = reinterpret_cast(address); if (stream->release != nullptr) { stream->release(stream); } JNI_METHOD_END() } /* * Class: org_apache_arrow_c_jni_JniWrapper * Method: exportSchema * Signature: (JLorg/apache/arrow/c/jni/PrivateData;)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_c_jni_JniWrapper_exportSchema( JNIEnv* env, jobject, jlong address, jobject private_data) { JNI_METHOD_START ArrowSchema* schema = reinterpret_cast(address); JavaVM* vm; if (env->GetJavaVM(&vm) != JNI_OK) { JniThrow("Unable to get JavaVM instance"); } jobject private_data_ref = env->NewGlobalRef(private_data); schema->private_data = new InnerPrivateData(vm, private_data_ref); schema->release = &release_exported; JNI_METHOD_END() } /* * Class: org_apache_arrow_c_jni_JniWrapper * Method: exportArray * Signature: (JLorg/apache/arrow/c/jni/PrivateData;)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_c_jni_JniWrapper_exportArray( JNIEnv* env, jobject, jlong address, jobject private_data) { JNI_METHOD_START ArrowArray* array = reinterpret_cast(address); JavaVM* vm; if (env->GetJavaVM(&vm) != JNI_OK) { JniThrow("Unable to get JavaVM instance"); } jobject private_data_ref = env->NewGlobalRef(private_data); array->private_data = new InnerPrivateData(vm, private_data_ref); array->release = &release_exported; JNI_METHOD_END() } /* * Class: org_apache_arrow_c_jni_JniWrapper * Method: exportArrayStream * Signature: (JLorg/apache/arrow/c/jni/PrivateData;)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_c_jni_JniWrapper_exportArrayStream( JNIEnv* env, jobject, jlong address, jobject private_data) { JNI_METHOD_START auto* stream = reinterpret_cast(address); JavaVM* vm; if (env->GetJavaVM(&vm) != JNI_OK) { JniThrow("Unable to get JavaVM instance"); } jobject private_data_ref = env->NewGlobalRef(private_data); stream->get_schema = &ArrowArrayStreamGetSchema; stream->get_next = &ArrowArrayStreamGetNext; stream->get_last_error = &ArrowArrayStreamGetLastError; stream->release = &ArrowArrayStreamRelease; stream->private_data = new InnerPrivateData(vm, private_data_ref); JNI_METHOD_END() } apache-arrow-java-19.0.0/c/src/main/java/000077500000000000000000000000001515451615200177545ustar00rootroot00000000000000apache-arrow-java-19.0.0/c/src/main/java/module-info.java000066400000000000000000000020421515451615200230330ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ open module org.apache.arrow.c { exports org.apache.arrow.c; exports org.apache.arrow.c.jni; requires flatbuffers.java; requires jdk.unsupported; requires org.apache.arrow.memory.core; requires org.apache.arrow.vector; requires org.slf4j; } apache-arrow-java-19.0.0/c/src/main/java/org/000077500000000000000000000000001515451615200205435ustar00rootroot00000000000000apache-arrow-java-19.0.0/c/src/main/java/org/apache/000077500000000000000000000000001515451615200217645ustar00rootroot00000000000000apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200231165ustar00rootroot00000000000000apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/000077500000000000000000000000001515451615200233405ustar00rootroot00000000000000apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/ArrayExporter.java000066400000000000000000000114311515451615200270120ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.c.NativeUtil.NULL; import static org.apache.arrow.c.NativeUtil.addressOrNull; import static org.apache.arrow.util.Preconditions.checkNotNull; import java.util.ArrayList; import java.util.List; import org.apache.arrow.c.jni.JniWrapper; import org.apache.arrow.c.jni.PrivateData; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.types.pojo.DictionaryEncoding; /** Exporter for {@link ArrowArray}. */ final class ArrayExporter { private final BufferAllocator allocator; public ArrayExporter(BufferAllocator allocator) { this.allocator = allocator; } /** Private data structure for exported arrays. */ static class ExportedArrayPrivateData implements PrivateData { ArrowBuf buffers_ptrs; List buffers; ArrowBuf children_ptrs; List children; ArrowArray dictionary; @Override public void close() { NativeUtil.closeBuffer(buffers_ptrs); if (buffers != null) { for (ArrowBuf buffer : buffers) { NativeUtil.closeBuffer(buffer); } } NativeUtil.closeBuffer(children_ptrs); if (children != null) { for (ArrowArray child : children) { child.close(); } } if (dictionary != null) { dictionary.close(); } } } void export(ArrowArray array, FieldVector vector, DictionaryProvider dictionaryProvider) { List children = vector.getChildrenFromFields(); int valueCount = vector.getValueCount(); int nullCount = vector.getNullCount(); DictionaryEncoding dictionaryEncoding = vector.getField().getDictionary(); ExportedArrayPrivateData data = new ExportedArrayPrivateData(); try { if (children != null) { data.children = new ArrayList<>(children.size()); data.children_ptrs = allocator.buffer((long) children.size() * Long.BYTES); for (int i = 0; i < children.size(); i++) { ArrowArray child = ArrowArray.allocateNew(allocator); data.children.add(child); data.children_ptrs.writeLong(child.memoryAddress()); } } data.buffers = new ArrayList<>(vector.getExportedCDataBufferCount()); data.buffers_ptrs = allocator.buffer((long) vector.getExportedCDataBufferCount() * Long.BYTES); vector.exportCDataBuffers(data.buffers, data.buffers_ptrs, NULL); if (dictionaryEncoding != null) { Dictionary dictionary = dictionaryProvider.lookup(dictionaryEncoding.getId()); checkNotNull(dictionary, "Dictionary lookup failed on export of dictionary encoded array"); data.dictionary = ArrowArray.allocateNew(allocator); FieldVector dictionaryVector = dictionary.getVector(); export(data.dictionary, dictionaryVector, dictionaryProvider); } ArrowArray.Snapshot snapshot = new ArrowArray.Snapshot(); snapshot.length = valueCount; snapshot.null_count = nullCount; snapshot.offset = 0; snapshot.n_buffers = (data.buffers != null) ? data.buffers.size() : 0; snapshot.n_children = (data.children != null) ? data.children.size() : 0; snapshot.buffers = addressOrNull(data.buffers_ptrs); snapshot.children = addressOrNull(data.children_ptrs); snapshot.dictionary = addressOrNull(data.dictionary); snapshot.release = NULL; array.save(snapshot); // sets release and private data JniWrapper.get().exportArray(array.memoryAddress(), data); } catch (Exception e) { data.close(); throw e; } // Export children if (children != null) { for (int i = 0; i < children.size(); i++) { FieldVector childVector = children.get(i); ArrowArray child = data.children.get(i); export(child, childVector, dictionaryProvider); } } } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/ArrayImporter.java000066400000000000000000000133311515451615200270040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.c.NativeUtil.NULL; import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import static org.apache.arrow.util.Preconditions.checkNotNull; import static org.apache.arrow.util.Preconditions.checkState; import java.util.Collections; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.types.pojo.DictionaryEncoding; /** Importer for {@link ArrowArray}. */ final class ArrayImporter { private static final int MAX_IMPORT_RECURSION_LEVEL = 64; private final BufferAllocator allocator; private final FieldVector vector; private final DictionaryProvider dictionaryProvider; private ReferenceCountedArrowArray underlyingAllocation; private int recursionLevel; ArrayImporter( BufferAllocator allocator, FieldVector vector, DictionaryProvider dictionaryProvider) { this.allocator = Preconditions.checkNotNull(allocator); this.vector = Preconditions.checkNotNull(vector); this.dictionaryProvider = dictionaryProvider; } void importArray(ArrowArray src) { ArrowArray.Snapshot snapshot = src.snapshot(); checkState(snapshot.release != NULL, "Cannot import released ArrowArray"); // Move imported array ArrowArray ownedArray = ArrowArray.allocateNew(allocator); ownedArray.save(snapshot); src.markReleased(); recursionLevel = 0; // This keeps the array alive as long as there are any buffers that need it underlyingAllocation = new ReferenceCountedArrowArray(ownedArray); try { doImport(snapshot); } finally { underlyingAllocation.release(); } } private void importChild(ArrayImporter parent, ArrowArray src) { ArrowArray.Snapshot snapshot = src.snapshot(); checkState(snapshot.release != NULL, "Cannot import released ArrowArray"); recursionLevel = parent.recursionLevel + 1; checkState( recursionLevel <= MAX_IMPORT_RECURSION_LEVEL, "Recursion level in ArrowArray struct exceeded"); // Child buffers will keep the entire parent import alive. underlyingAllocation = parent.underlyingAllocation; doImport(snapshot); } private void doImport(ArrowArray.Snapshot snapshot) { // First import children (required for reconstituting parent array data) long[] children = NativeUtil.toJavaArray(snapshot.children, checkedCastToInt(snapshot.n_children)); if (children != null && children.length > 0) { List childVectors = vector.getChildrenFromFields(); checkState( children.length == childVectors.size(), "ArrowArray struct has %s children (expected %s)", children.length, childVectors.size()); for (int i = 0; i < children.length; i++) { checkState(children[i] != NULL, "ArrowArray struct has NULL child at position %s", i); ArrayImporter childImporter = new ArrayImporter(allocator, childVectors.get(i), dictionaryProvider); childImporter.importChild(this, ArrowArray.wrap(children[i])); } } // Handle import of a dictionary encoded vector if (snapshot.dictionary != NULL) { DictionaryEncoding encoding = vector.getField().getDictionary(); checkNotNull(encoding, "Missing encoding on import of ArrowArray with dictionary"); Dictionary dictionary = dictionaryProvider.lookup(encoding.getId()); checkNotNull(dictionary, "Dictionary lookup failed on import of ArrowArray with dictionary"); // reset the dictionary vector to the initial state dictionary.getVector().clear(); ArrayImporter dictionaryImporter = new ArrayImporter(allocator, dictionary.getVector(), dictionaryProvider); dictionaryImporter.importChild(this, ArrowArray.wrap(snapshot.dictionary)); } // Import main data ArrowFieldNode fieldNode = new ArrowFieldNode(snapshot.length, snapshot.null_count); long[] bufferPointers = NativeUtil.toJavaArray(snapshot.buffers, checkedCastToInt(snapshot.n_buffers)); try (final BufferImportTypeVisitor visitor = new BufferImportTypeVisitor(allocator, underlyingAllocation, fieldNode, bufferPointers)) { final List buffers; if (bufferPointers == null || bufferPointers.length == 0) { buffers = Collections.emptyList(); } else { buffers = vector.getField().getType().accept(visitor); } vector.loadFieldBuffers(fieldNode, buffers); } catch (Exception e) { throw new IllegalArgumentException( "Could not load buffers for field " + vector.getField() + ". error message: " + e.getMessage(), e); } } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/ArrayStreamExporter.java000066400000000000000000000072601515451615200301730ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.io.IOException; import java.io.PrintWriter; import java.io.StringWriter; import java.nio.charset.StandardCharsets; import org.apache.arrow.c.jni.JniWrapper; import org.apache.arrow.c.jni.PrivateData; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.types.pojo.Schema; /** Utility to export an {@link ArrowReader} as an ArrowArrayStream. */ final class ArrayStreamExporter { private final BufferAllocator allocator; ArrayStreamExporter(BufferAllocator allocator) { this.allocator = allocator; } /** Java-side state for the exported stream. */ static class ExportedArrayStreamPrivateData implements PrivateData { final BufferAllocator allocator; final ArrowReader reader; // Read by the JNI side for get_last_error byte[] lastError; ExportedArrayStreamPrivateData(BufferAllocator allocator, ArrowReader reader) { this.allocator = allocator; this.reader = reader; } private int setLastError(Throwable err) { // Do not let exceptions propagate up to JNI try { StringWriter buf = new StringWriter(); PrintWriter writer = new PrintWriter(buf); err.printStackTrace(writer); lastError = buf.toString().getBytes(StandardCharsets.UTF_8); } catch (Throwable e) { // Bail out of setting the error message - we'll still return an error code lastError = null; } return 5; // = EIO } @SuppressWarnings("unused") // Used by JNI int getNext(long arrayAddress) { try (ArrowArray out = ArrowArray.wrap(arrayAddress)) { if (reader.loadNextBatch()) { Data.exportVectorSchemaRoot(allocator, reader.getVectorSchemaRoot(), reader, out); } else { out.markReleased(); } return 0; } catch (Throwable e) { return setLastError(e); } } @SuppressWarnings("unused") // Used by JNI int getSchema(long schemaAddress) { try (ArrowSchema out = ArrowSchema.wrap(schemaAddress)) { final Schema schema = reader.getVectorSchemaRoot().getSchema(); Data.exportSchema(allocator, schema, reader, out); return 0; } catch (Throwable e) { return setLastError(e); } } @Override public void close() { try { reader.close(); } catch (IOException e) { // XXX: C Data Interface gives us no way to signal errors to the caller, // but the JNI side will catch this and log an error. throw new RuntimeException(e); } } } void export(ArrowArrayStream stream, ArrowReader reader) { ExportedArrayStreamPrivateData data = new ExportedArrayStreamPrivateData(allocator, reader); try { JniWrapper.get().exportArrayStream(stream.memoryAddress(), data); } catch (Exception e) { data.close(); throw e; } } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/ArrowArray.java000066400000000000000000000125371515451615200263040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.c.NativeUtil.NULL; import static org.apache.arrow.util.Preconditions.checkNotNull; import java.nio.ByteBuffer; import java.nio.ByteOrder; import org.apache.arrow.c.jni.JniWrapper; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReferenceManager; import org.apache.arrow.memory.util.MemoryUtil; import org.apache.arrow.util.VisibleForTesting; /** * C Data Interface ArrowArray. * *

    Represents a wrapper for the following C structure: * * 

     * struct ArrowArray {
 *     // Array data description
 *     int64_t length;
 *     int64_t null_count;
 *     int64_t offset;
 *     int64_t n_buffers;
 *     int64_t n_children;
 *     const void** buffers;
 *     struct ArrowArray** children;
 *     struct ArrowArray* dictionary;
 *
 *     // Release callback
 *     void (*release)(struct ArrowArray*);
 *     // Opaque producer-specific data
 *     void* private_data;
 * };
 *
    */ public class ArrowArray implements BaseStruct { private static final int SIZE_OF = 80; private static final int INDEX_RELEASE_CALLBACK = 64; private ArrowBuf data; /** Snapshot of the ArrowArray raw data. */ public static class Snapshot { public long length; public long null_count; public long offset; public long n_buffers; public long n_children; public long buffers; public long children; public long dictionary; public long release; public long private_data; /** Initialize empty ArrowArray snapshot. */ public Snapshot() { length = NULL; null_count = NULL; offset = NULL; n_buffers = NULL; n_children = NULL; buffers = NULL; children = NULL; dictionary = NULL; release = NULL; private_data = NULL; } } /** * Create ArrowArray from an existing memory address. * *

    The resulting ArrowArray does not own the memory. * * @param memoryAddress Memory address to wrap * @return A new ArrowArray instance */ public static ArrowArray wrap(long memoryAddress) { return new ArrowArray( new ArrowBuf(ReferenceManager.NO_OP, null, ArrowArray.SIZE_OF, memoryAddress)); } /** * Create ArrowArray by allocating memory. * *

    The resulting ArrowArray owns the memory. * * @param allocator Allocator for memory allocations * @return A new ArrowArray instance */ public static ArrowArray allocateNew(BufferAllocator allocator) { ArrowArray array = new ArrowArray(allocator.buffer(ArrowArray.SIZE_OF)); array.markReleased(); return array; } ArrowArray(ArrowBuf data) { checkNotNull(data, "ArrowArray initialized with a null buffer"); this.data = data; } /** Mark the array as released. */ public void markReleased() { directBuffer().putLong(INDEX_RELEASE_CALLBACK, NULL); } @Override public long memoryAddress() { checkNotNull(data, "ArrowArray is already closed"); return data.memoryAddress(); } @Override public void release() { long address = memoryAddress(); JniWrapper.get().releaseArray(address); } @Override public void close() { if (data != null) { data.close(); data = null; } } @VisibleForTesting boolean isClosed() { return data == null; } private ByteBuffer directBuffer() { return MemoryUtil.directBuffer(memoryAddress(), ArrowArray.SIZE_OF) .order(ByteOrder.nativeOrder()); } /** * Take a snapshot of the ArrowArray raw values. * * @return snapshot */ public Snapshot snapshot() { ByteBuffer data = directBuffer(); Snapshot snapshot = new Snapshot(); snapshot.length = data.getLong(); snapshot.null_count = data.getLong(); snapshot.offset = data.getLong(); snapshot.n_buffers = data.getLong(); snapshot.n_children = data.getLong(); snapshot.buffers = data.getLong(); snapshot.children = data.getLong(); snapshot.dictionary = data.getLong(); snapshot.release = data.getLong(); snapshot.private_data = data.getLong(); return snapshot; } /** Write values from Snapshot to the underlying ArrowArray memory buffer. */ public void save(Snapshot snapshot) { directBuffer() .putLong(snapshot.length) .putLong(snapshot.null_count) .putLong(snapshot.offset) .putLong(snapshot.n_buffers) .putLong(snapshot.n_children) .putLong(snapshot.buffers) .putLong(snapshot.children) .putLong(snapshot.dictionary) .putLong(snapshot.release) .putLong(snapshot.private_data); } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/ArrowArrayStream.java000066400000000000000000000133541515451615200274560ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.c.NativeUtil.NULL; import static org.apache.arrow.util.Preconditions.checkNotNull; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.ByteOrder; import org.apache.arrow.c.jni.CDataJniException; import org.apache.arrow.c.jni.JniWrapper; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReferenceManager; import org.apache.arrow.memory.util.MemoryUtil; /** * C Stream Interface ArrowArrayStream. * *

    Represents a wrapper for the following C structure: * * 

     * struct ArrowArrayStream {
 *   int (*get_schema)(struct ArrowArrayStream*, struct ArrowSchema* out);
 *   int (*get_next)(struct ArrowArrayStream*, struct ArrowArray* out);
 *   const char* (*get_last_error)(struct ArrowArrayStream*);
 *   void (*release)(struct ArrowArrayStream*);
 *   void* private_data;
 * };
 *
    */ public class ArrowArrayStream implements BaseStruct { private static final int SIZE_OF = 40; private static final int INDEX_RELEASE_CALLBACK = 24; private ArrowBuf data; /** Snapshot of the ArrowArrayStream raw data. */ public static class Snapshot { public long get_schema; public long get_next; public long get_last_error; public long release; public long private_data; /** Initialize empty ArrowArray snapshot. */ public Snapshot() { get_schema = NULL; get_next = NULL; get_last_error = NULL; release = NULL; private_data = NULL; } } /** * Create ArrowArrayStream from an existing memory address. * *

    The resulting ArrowArrayStream does not own the memory. * * @param memoryAddress Memory address to wrap * @return A new ArrowArrayStream instance */ public static ArrowArrayStream wrap(long memoryAddress) { return new ArrowArrayStream( new ArrowBuf(ReferenceManager.NO_OP, null, ArrowArrayStream.SIZE_OF, memoryAddress)); } /** * Create ArrowArrayStream by allocating memory. * *

    The resulting ArrowArrayStream owns the memory. * * @param allocator Allocator for memory allocations * @return A new ArrowArrayStream instance */ public static ArrowArrayStream allocateNew(BufferAllocator allocator) { ArrowArrayStream array = new ArrowArrayStream(allocator.buffer(ArrowArrayStream.SIZE_OF)); array.markReleased(); return array; } ArrowArrayStream(ArrowBuf data) { checkNotNull(data, "ArrowArrayStream initialized with a null buffer"); this.data = data; } /** Mark the array as released. */ public void markReleased() { directBuffer().putLong(INDEX_RELEASE_CALLBACK, NULL); } @Override public long memoryAddress() { checkNotNull(data, "ArrowArrayStream is already closed"); return data.memoryAddress(); } @Override public void release() { long address = memoryAddress(); JniWrapper.get().releaseArrayStream(address); } /** * Get the schema of the stream. * * @param schema The ArrowSchema struct to output to * @throws IOException if the stream returns an error */ public void getSchema(ArrowSchema schema) throws IOException { long address = memoryAddress(); try { JniWrapper.get().getSchemaArrayStream(address, schema.memoryAddress()); } catch (CDataJniException e) { throw new IOException("[errno " + e.getErrno() + "] " + e.getMessage()); } } /** * Get the next batch in the stream. * * @param array The ArrowArray struct to output to * @throws IOException if the stream returns an error */ public void getNext(ArrowArray array) throws IOException { long address = memoryAddress(); try { JniWrapper.get().getNextArrayStream(address, array.memoryAddress()); } catch (CDataJniException e) { throw new IOException("[errno " + e.getErrno() + "] " + e.getMessage()); } } @Override public void close() { if (data != null) { data.close(); data = null; } } private ByteBuffer directBuffer() { return MemoryUtil.directBuffer(memoryAddress(), ArrowArrayStream.SIZE_OF) .order(ByteOrder.nativeOrder()); } /** * Take a snapshot of the ArrowArrayStream raw values. * * @return snapshot */ public ArrowArrayStream.Snapshot snapshot() { ByteBuffer data = directBuffer(); ArrowArrayStream.Snapshot snapshot = new ArrowArrayStream.Snapshot(); snapshot.get_schema = data.getLong(); snapshot.get_next = data.getLong(); snapshot.get_last_error = data.getLong(); snapshot.release = data.getLong(); snapshot.private_data = data.getLong(); return snapshot; } /** Write values from Snapshot to the underlying ArrowArrayStream memory buffer. */ public void save(ArrowArrayStream.Snapshot snapshot) { directBuffer() .putLong(snapshot.get_schema) .putLong(snapshot.get_next) .putLong(snapshot.get_last_error) .putLong(snapshot.release) .putLong(snapshot.private_data); } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/ArrowArrayStreamReader.java000066400000000000000000000057341515451615200306040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.c.NativeUtil.NULL; import static org.apache.arrow.util.Preconditions.checkState; import java.io.IOException; import java.util.Map; import java.util.function.Function; import java.util.stream.Collectors; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.types.pojo.Schema; /** An implementation of an {@link ArrowReader} backed by an ArrowArrayStream. */ final class ArrowArrayStreamReader extends ArrowReader { private final ArrowArrayStream ownedStream; private final CDataDictionaryProvider provider; ArrowArrayStreamReader(BufferAllocator allocator, ArrowArrayStream stream) { super(allocator); this.provider = new CDataDictionaryProvider(); ArrowArrayStream.Snapshot snapshot = stream.snapshot(); checkState(snapshot.release != NULL, "Cannot import released ArrowArrayStream"); // Move imported stream this.ownedStream = ArrowArrayStream.allocateNew(allocator); this.ownedStream.save(snapshot); stream.markReleased(); } @Override public Map getDictionaryVectors() { return provider.getDictionaryIds().stream() .collect(Collectors.toMap(Function.identity(), provider::lookup)); } @Override public Dictionary lookup(long id) { return provider.lookup(id); } @Override public boolean loadNextBatch() throws IOException { try (ArrowArray array = ArrowArray.allocateNew(allocator)) { ownedStream.getNext(array); if (array.snapshot().release == NULL) { return false; } Data.importIntoVectorSchemaRoot(allocator, array, getVectorSchemaRoot(), provider); return true; } } @Override public long bytesRead() { return 0; } @Override protected void closeReadSource() { ownedStream.release(); ownedStream.close(); provider.close(); } @Override protected Schema readSchema() throws IOException { try (ArrowSchema schema = ArrowSchema.allocateNew(allocator)) { ownedStream.getSchema(schema); return Data.importSchema(allocator, schema, provider); } } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/ArrowSchema.java000066400000000000000000000121231515451615200264150ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.c.NativeUtil.NULL; import static org.apache.arrow.util.Preconditions.checkNotNull; import java.nio.ByteBuffer; import java.nio.ByteOrder; import org.apache.arrow.c.jni.JniWrapper; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReferenceManager; import org.apache.arrow.memory.util.MemoryUtil; /** * C Data Interface ArrowSchema. * *

    Represents a wrapper for the following C structure: * * 

     * struct ArrowSchema {
 *     // Array type description
 *     const char* format;
 *     const char* name;
 *     const char* metadata;
 *     int64_t flags;
 *     int64_t n_children;
 *     struct ArrowSchema** children;
 *     struct ArrowSchema* dictionary;
 *
 *     // Release callback
 *     void (*release)(struct ArrowSchema*);
 *     // Opaque producer-specific data
 *     void* private_data;
 * };
 *
    */ public class ArrowSchema implements BaseStruct { private static final int SIZE_OF = 72; private static final int INDEX_RELEASE_CALLBACK = 56; private ArrowBuf data; /** Snapshot of the ArrowSchema raw data. */ public static class Snapshot { public long format; public long name; public long metadata; public long flags; public long n_children; public long children; public long dictionary; public long release; public long private_data; /** Initialize empty ArrowSchema snapshot. */ public Snapshot() { format = NULL; name = NULL; metadata = NULL; flags = NULL; n_children = NULL; children = NULL; dictionary = NULL; release = NULL; private_data = NULL; } } /** * Create ArrowSchema from an existing memory address. * *

    The resulting ArrowSchema does not own the memory. * * @param memoryAddress Memory address to wrap * @return A new ArrowSchema instance */ public static ArrowSchema wrap(long memoryAddress) { return new ArrowSchema( new ArrowBuf(ReferenceManager.NO_OP, null, ArrowSchema.SIZE_OF, memoryAddress)); } /** * Create ArrowSchema by allocating memory. * *

    The resulting ArrowSchema owns the memory. * * @param allocator Allocator for memory allocations * @return A new ArrowSchema instance */ public static ArrowSchema allocateNew(BufferAllocator allocator) { ArrowSchema schema = new ArrowSchema(allocator.buffer(ArrowSchema.SIZE_OF)); schema.markReleased(); return schema; } ArrowSchema(ArrowBuf data) { checkNotNull(data, "ArrowSchema initialized with a null buffer"); this.data = data; } /** Mark the schema as released. */ public void markReleased() { directBuffer().putLong(INDEX_RELEASE_CALLBACK, NULL); } @Override public long memoryAddress() { checkNotNull(data, "ArrowSchema is already closed"); return data.memoryAddress(); } @Override public void release() { long address = memoryAddress(); JniWrapper.get().releaseSchema(address); } @Override public void close() { if (data != null) { data.close(); data = null; } } private ByteBuffer directBuffer() { return MemoryUtil.directBuffer(memoryAddress(), ArrowSchema.SIZE_OF) .order(ByteOrder.nativeOrder()); } /** * Take a snapshot of the ArrowSchema raw values. * * @return snapshot */ public Snapshot snapshot() { ByteBuffer data = directBuffer(); Snapshot snapshot = new Snapshot(); snapshot.format = data.getLong(); snapshot.name = data.getLong(); snapshot.metadata = data.getLong(); snapshot.flags = data.getLong(); snapshot.n_children = data.getLong(); snapshot.children = data.getLong(); snapshot.dictionary = data.getLong(); snapshot.release = data.getLong(); snapshot.private_data = data.getLong(); return snapshot; } /** Write values from Snapshot to the underlying ArrowSchema memory buffer. */ public void save(Snapshot snapshot) { directBuffer() .putLong(snapshot.format) .putLong(snapshot.name) .putLong(snapshot.metadata) .putLong(snapshot.flags) .putLong(snapshot.n_children) .putLong(snapshot.children) .putLong(snapshot.dictionary) .putLong(snapshot.release) .putLong(snapshot.private_data); } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/BaseStruct.java000066400000000000000000000023411515451615200262620ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** Base interface for C Data Interface structures. */ public interface BaseStruct extends AutoCloseable { /** * Get memory address. * * @return Memory address */ long memoryAddress(); /** * Call the release callback of an ArrowArray. * *

    This function must not be called for child arrays. */ void release(); /** Close to release the main buffer. */ @Override void close(); } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/BufferImportTypeVisitor.java000066400000000000000000000354311515451615200310370ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.c.NativeUtil.NULL; import static org.apache.arrow.util.Preconditions.checkState; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.VisibleForTesting; import org.apache.arrow.vector.BaseVariableWidthViewVector; import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DateMilliVector; import org.apache.arrow.vector.DurationVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.IntervalDayVector; import org.apache.arrow.vector.IntervalMonthDayNanoVector; import org.apache.arrow.vector.IntervalYearVector; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.LargeVarCharVector; 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.TimeStampVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.LargeListVector; 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.MapVector; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.util.DataSizeRoundingUtil; /** Import buffers from a C Data Interface struct. */ class BufferImportTypeVisitor implements ArrowType.ArrowTypeVisitor>, AutoCloseable { private final BufferAllocator allocator; private final ReferenceCountedArrowArray underlyingAllocation; private final ArrowFieldNode fieldNode; private final long[] buffers; private final List imported; BufferImportTypeVisitor( BufferAllocator allocator, ReferenceCountedArrowArray underlyingAllocation, ArrowFieldNode fieldNode, long[] buffers) { this.allocator = allocator; this.underlyingAllocation = underlyingAllocation; this.fieldNode = fieldNode; this.buffers = buffers; this.imported = new ArrayList<>(); } @Override public void close() throws Exception { AutoCloseables.close(imported); } @VisibleForTesting ArrowBuf importBuffer(ArrowType type, int index, long capacity) { checkState( buffers.length > index, "Expected at least %s buffers for type %s, but found %s", index + 1, type, buffers.length); long bufferPtr = buffers[index]; if (bufferPtr == NULL) { // C array may be NULL but only accept that if expected capacity is zero too if (capacity != 0) { throw new IllegalStateException( String.format("Buffer %s for type %s cannot be null", index, type)); } else { // no data in the C array, return an empty buffer return allocator.getEmpty(); } } ArrowBuf buf = underlyingAllocation.unsafeAssociateAllocation(allocator, capacity, bufferPtr); imported.add(buf); return buf; } private ArrowBuf importFixedBits(ArrowType type, int index, long bitsPerSlot) { final long capacity = DataSizeRoundingUtil.divideBy8Ceil(bitsPerSlot * fieldNode.getLength()); return importBuffer(type, index, capacity); } private ArrowBuf importFixedBytes(ArrowType type, int index, long bytesPerSlot) { final long capacity = bytesPerSlot * fieldNode.getLength(); return importBuffer(type, index, capacity); } private ArrowBuf importOffsets(ArrowType type, long bytesPerSlot) { final long capacity = bytesPerSlot * (fieldNode.getLength() + 1); return importBuffer(type, 1, capacity); } private ArrowBuf importData(ArrowType type, long capacity) { return importBuffer(type, 2, capacity); } private ArrowBuf maybeImportBitmap(ArrowType type) { checkState( buffers.length > 0, "Expected at least %s buffers for type %s, but found %s", 1, type, buffers.length); if (buffers[0] == NULL) { return null; } return importFixedBits(type, 0, /* bitsPerSlot= */ 1); } @Override public List visit(ArrowType.Null type) { checkState( buffers.length == 0, "Expected %s buffers for type %s, but found %s", 0, type, buffers.length); return Collections.emptyList(); } @Override public List visit(ArrowType.Struct type) { return Collections.singletonList(maybeImportBitmap(type)); } @Override public List visit(ArrowType.List type) { return Arrays.asList(maybeImportBitmap(type), importOffsets(type, ListVector.OFFSET_WIDTH)); } @Override public List visit(ArrowType.LargeList type) { return Arrays.asList( maybeImportBitmap(type), importOffsets(type, LargeListVector.OFFSET_WIDTH)); } @Override public List visit(ArrowType.FixedSizeList type) { return Collections.singletonList(maybeImportBitmap(type)); } @Override public List visit(ArrowType.Union type) { switch (type.getMode()) { case Sparse: return Collections.singletonList(importFixedBytes(type, 0, UnionVector.TYPE_WIDTH)); case Dense: return Arrays.asList( importFixedBytes(type, 0, DenseUnionVector.TYPE_WIDTH), importFixedBytes(type, 1, DenseUnionVector.OFFSET_WIDTH)); default: throw new UnsupportedOperationException("Importing buffers for union type: " + type); } } @Override public List visit(ArrowType.RunEndEncoded type) { return List.of(); } @Override public List visit(ArrowType.Map type) { return Arrays.asList(maybeImportBitmap(type), importOffsets(type, MapVector.OFFSET_WIDTH)); } @Override public List visit(ArrowType.Int type) { return Arrays.asList(maybeImportBitmap(type), importFixedBits(type, 1, type.getBitWidth())); } @Override public List visit(ArrowType.FloatingPoint type) { switch (type.getPrecision()) { case HALF: return Arrays.asList( maybeImportBitmap(type), importFixedBytes(type, 1, /* bytesPerSlot= */ 2)); case SINGLE: return Arrays.asList( maybeImportBitmap(type), importFixedBytes(type, 1, Float4Vector.TYPE_WIDTH)); case DOUBLE: return Arrays.asList( maybeImportBitmap(type), importFixedBytes(type, 1, Float8Vector.TYPE_WIDTH)); default: throw new UnsupportedOperationException("Importing buffers for type: " + type); } } @Override public List visit(ArrowType.Utf8 type) { try (ArrowBuf offsets = importOffsets(type, VarCharVector.OFFSET_WIDTH)) { final int start = offsets.getInt(0); final int end = offsets.getInt(fieldNode.getLength() * (long) VarCharVector.OFFSET_WIDTH); checkState( end >= start, "Offset buffer for type %s is malformed: start: %s, end: %s", type, start, end); offsets.getReferenceManager().retain(); return Arrays.asList(maybeImportBitmap(type), offsets, importData(type, end)); } } private List visitVariableWidthView(ArrowType type) { final int viewBufferIndex = 1; final int variadicSizeBufferIndex = this.buffers.length - 1; final long numOfVariadicBuffers = this.buffers.length - 3L; final long variadicSizeBufferCapacity = numOfVariadicBuffers * Long.BYTES; List buffers = new ArrayList<>(); ArrowBuf variadicSizeBuffer = importBuffer(type, variadicSizeBufferIndex, variadicSizeBufferCapacity); ArrowBuf view = importFixedBytes(type, viewBufferIndex, BaseVariableWidthViewVector.ELEMENT_SIZE); buffers.add(maybeImportBitmap(type)); buffers.add(view); // 0th buffer is validity buffer // 1st buffer is view buffer // 2nd buffer onwards are variadic buffer // N-1 (this.buffers.length - 1) buffer is variadic size buffer final int variadicBufferReadOffset = 2; for (int i = 0; i < numOfVariadicBuffers; i++) { long size = variadicSizeBuffer.getLong((long) i * Long.BYTES); buffers.add(importBuffer(type, i + variadicBufferReadOffset, size)); } return buffers; } @Override public List visit(ArrowType.Utf8View type) { return visitVariableWidthView(type); } @Override public List visit(ArrowType.LargeUtf8 type) { try (ArrowBuf offsets = importOffsets(type, LargeVarCharVector.OFFSET_WIDTH)) { final long start = offsets.getLong(0); final long end = offsets.getLong(fieldNode.getLength() * (long) LargeVarCharVector.OFFSET_WIDTH); checkState( end >= start, "Offset buffer for type %s is malformed: start: %s, end: %s", type, start, end); offsets.getReferenceManager().retain(); return Arrays.asList(maybeImportBitmap(type), offsets, importData(type, end)); } } @Override public List visit(ArrowType.Binary type) { try (ArrowBuf offsets = importOffsets(type, VarBinaryVector.OFFSET_WIDTH)) { final int start = offsets.getInt(0); final int end = offsets.getInt(fieldNode.getLength() * (long) VarBinaryVector.OFFSET_WIDTH); checkState( end >= start, "Offset buffer for type %s is malformed: start: %s, end: %s", type, start, end); offsets.getReferenceManager().retain(); return Arrays.asList(maybeImportBitmap(type), offsets, importData(type, end)); } } @Override public List visit(ArrowType.BinaryView type) { return visitVariableWidthView(type); } @Override public List visit(ArrowType.LargeBinary type) { try (ArrowBuf offsets = importOffsets(type, LargeVarBinaryVector.OFFSET_WIDTH)) { final long start = offsets.getLong(0); // TODO: need better tests to cover the failure when I forget to multiply by offset width final long end = offsets.getLong(fieldNode.getLength() * (long) LargeVarBinaryVector.OFFSET_WIDTH); checkState( end >= start, "Offset buffer for type %s is malformed: start: %s, end: %s", type, start, end); offsets.getReferenceManager().retain(); return Arrays.asList(maybeImportBitmap(type), offsets, importData(type, end)); } } @Override public List visit(ArrowType.FixedSizeBinary type) { return Arrays.asList(maybeImportBitmap(type), importFixedBytes(type, 1, type.getByteWidth())); } @Override public List visit(ArrowType.Bool type) { return Arrays.asList(maybeImportBitmap(type), importFixedBits(type, 1, /* bitsPerSlot= */ 1)); } @Override public List visit(ArrowType.Decimal type) { return Arrays.asList(maybeImportBitmap(type), importFixedBits(type, 1, type.getBitWidth())); } @Override public List visit(ArrowType.Date type) { switch (type.getUnit()) { case DAY: return Arrays.asList( maybeImportBitmap(type), importFixedBytes(type, 1, DateDayVector.TYPE_WIDTH)); case MILLISECOND: return Arrays.asList( maybeImportBitmap(type), importFixedBytes(type, 1, DateMilliVector.TYPE_WIDTH)); default: throw new UnsupportedOperationException("Importing buffers for type: " + type); } } @Override public List visit(ArrowType.Time type) { switch (type.getUnit()) { case SECOND: return Arrays.asList( maybeImportBitmap(type), importFixedBytes(type, 1, TimeSecVector.TYPE_WIDTH)); case MILLISECOND: return Arrays.asList( maybeImportBitmap(type), importFixedBytes(type, 1, TimeMilliVector.TYPE_WIDTH)); case MICROSECOND: return Arrays.asList( maybeImportBitmap(type), importFixedBytes(type, 1, TimeMicroVector.TYPE_WIDTH)); case NANOSECOND: return Arrays.asList( maybeImportBitmap(type), importFixedBytes(type, 1, TimeNanoVector.TYPE_WIDTH)); default: throw new UnsupportedOperationException("Importing buffers for type: " + type); } } @Override public List visit(ArrowType.Timestamp type) { return Arrays.asList( maybeImportBitmap(type), importFixedBytes(type, 1, TimeStampVector.TYPE_WIDTH)); } @Override public List visit(ArrowType.Interval type) { switch (type.getUnit()) { case YEAR_MONTH: return Arrays.asList( maybeImportBitmap(type), importFixedBytes(type, 1, IntervalYearVector.TYPE_WIDTH)); case DAY_TIME: return Arrays.asList( maybeImportBitmap(type), importFixedBytes(type, 1, IntervalDayVector.TYPE_WIDTH)); case MONTH_DAY_NANO: return Arrays.asList( maybeImportBitmap(type), importFixedBytes(type, 1, IntervalMonthDayNanoVector.TYPE_WIDTH)); default: throw new UnsupportedOperationException("Importing buffers for type: " + type); } } @Override public List visit(ArrowType.Duration type) { return Arrays.asList( maybeImportBitmap(type), importFixedBytes(type, 1, DurationVector.TYPE_WIDTH)); } @Override public List visit(ArrowType.ListView type) { return Arrays.asList( maybeImportBitmap(type), importFixedBytes(type, 1, ListViewVector.OFFSET_WIDTH), importFixedBytes(type, 2, ListViewVector.SIZE_WIDTH)); } @Override public List visit(ArrowType.LargeListView type) { return Arrays.asList( maybeImportBitmap(type), importFixedBytes(type, 1, LargeListViewVector.OFFSET_WIDTH), importFixedBytes(type, 2, LargeListViewVector.SIZE_WIDTH)); } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/CDataDictionaryProvider.java000066400000000000000000000044731515451615200307300ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.HashMap; import java.util.Map; import java.util.Set; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; /** * A DictionaryProvider that is used in C Data Interface for imports. * *

    CDataDictionaryProvider is similar to {@link DictionaryProvider.MapDictionaryProvider} with a * key difference that the dictionaries are owned by the provider so it must eventually be closed. * *

    The typical usage is to create the CDataDictionaryProvider and pass it to {@link * Data#importField} or {@link Data#importSchema} to allocate empty dictionaries based on the * information in {@link ArrowSchema}. Then you can re-use the same dictionary provider in any * function that imports an {@link ArrowArray} that has the same schema. */ public class CDataDictionaryProvider implements DictionaryProvider, AutoCloseable { private final Map map; public CDataDictionaryProvider() { this.map = new HashMap<>(); } void put(Dictionary dictionary) { Dictionary previous = map.put(dictionary.getEncoding().getId(), dictionary); if (previous != null) { previous.getVector().close(); } } @Override public final Set getDictionaryIds() { return map.keySet(); } @Override public Dictionary lookup(long id) { return map.get(id); } @Override public void close() { for (Dictionary dictionary : map.values()) { dictionary.getVector().close(); } map.clear(); } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/Data.java000066400000000000000000000613011515451615200250550ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; import org.apache.arrow.vector.complex.StructVector; 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.table.Table; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeID; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; /** * Functions for working with the C data interface. * *

    This API is EXPERIMENTAL. Note that currently only 64bit systems are supported. */ public final class Data { private Data() {} /** * Export Java Field using the C data interface format. * * @param allocator Buffer allocator for allocating C data interface fields * @param field Field object to export * @param provider Dictionary provider for dictionary encoded fields (optional) * @param out C struct where to export the field */ public static void exportField( BufferAllocator allocator, Field field, DictionaryProvider provider, ArrowSchema out) { SchemaExporter exporter = new SchemaExporter(allocator); exporter.export(out, field, provider); } /** * Export Java Schema using the C data interface format. * * @param allocator Buffer allocator for allocating C data interface fields * @param schema Schema object to export * @param provider Dictionary provider for dictionary encoded fields (optional) * @param out C struct where to export the field */ public static void exportSchema( BufferAllocator allocator, Schema schema, DictionaryProvider provider, ArrowSchema out) { // Convert to a struct field equivalent to the input schema FieldType fieldType = new FieldType(false, new ArrowType.Struct(), null, schema.getCustomMetadata()); Field field = new Field("", fieldType, schema.getFields()); exportField(allocator, field, provider, out); } /** * Export Java FieldVector using the C data interface format. * *

    The resulting ArrowArray struct keeps the array data and buffers alive until its release * callback is called by the consumer. * * @param allocator Buffer allocator for allocating C data interface fields * @param vector Vector object to export * @param provider Dictionary provider for dictionary encoded vectors (optional) * @param out C struct where to export the array */ public static void exportVector( BufferAllocator allocator, FieldVector vector, DictionaryProvider provider, ArrowArray out) { exportVector(allocator, vector, provider, out, null); } /** * Export Java FieldVector using the C data interface format. * *

    The resulting ArrowArray struct keeps the array data and buffers alive until its release * callback is called by the consumer. * * @param allocator Buffer allocator for allocating C data interface fields * @param vector Vector object to export * @param provider Dictionary provider for dictionary encoded vectors (optional) * @param out C struct where to export the array * @param outSchema C struct where to export the array type (optional) */ public static void exportVector( BufferAllocator allocator, FieldVector vector, DictionaryProvider provider, ArrowArray out, ArrowSchema outSchema) { if (outSchema != null) { exportField(allocator, vector.getField(), provider, outSchema); } ArrayExporter exporter = new ArrayExporter(allocator); exporter.export(out, vector, provider); } /** * Export the current contents of a Java Table using the C data interface format. * *

    The table is exported as if it were a struct array. The resulting ArrowArray struct keeps * the record batch data and buffers alive until its release callback is called by the consumer. * * @param allocator Buffer allocator for allocating C data interface fields * @param table Table to export * @param out C struct where to export the record batch */ public static void exportTable(BufferAllocator allocator, Table table, ArrowArray out) { exportTable(allocator, table, table.getDictionaryProvider(), out, null); } /** * Export the current contents of a Java Table using the C data interface format. * *

    The table is exported as if it were a struct array. The resulting ArrowArray struct keeps * the record batch data and buffers alive until its release callback is called by the consumer. * * @param allocator Buffer allocator for allocating C data interface fields * @param table Table to export * @param provider Dictionary provider for dictionary encoded vectors (optional) * @param out C struct where to export the record batch */ public static void exportTable( BufferAllocator allocator, Table table, DictionaryProvider provider, ArrowArray out) { exportTable(allocator, table, provider, out, null); } /** * Export the current contents of a Java Table using the C data interface format. * *

    The table is exported as if it were a struct array. The resulting ArrowArray struct keeps * the record batch data and buffers alive until its release callback is called by the consumer. * * @param allocator Buffer allocator for allocating C data interface fields * @param table Table to export * @param provider Dictionary provider for dictionary encoded vectors (optional) * @param out C struct where to export the record batch * @param outSchema C struct where to export the record batch schema (optional) */ public static void exportTable( BufferAllocator allocator, Table table, DictionaryProvider provider, ArrowArray out, ArrowSchema outSchema) { try (VectorSchemaRoot root = table.toVectorSchemaRoot()) { exportVectorSchemaRoot(allocator, root, provider, out, outSchema); } } /** * Export the current contents of a Java VectorSchemaRoot using the C data interface format. * *

    The vector schema root is exported as if it were a struct array. The resulting ArrowArray * struct keeps the record batch data and buffers alive until its release callback is called by * the consumer. * * @param allocator Buffer allocator for allocating C data interface fields * @param vsr Vector schema root to export * @param provider Dictionary provider for dictionary encoded vectors (optional) * @param out C struct where to export the record batch */ public static void exportVectorSchemaRoot( BufferAllocator allocator, VectorSchemaRoot vsr, DictionaryProvider provider, ArrowArray out) { exportVectorSchemaRoot(allocator, vsr, provider, out, null); } /** * Export the current contents of a Java VectorSchemaRoot using the C data interface format. * *

    The vector schema root is exported as if it were a struct array. The resulting ArrowArray * struct keeps the record batch data and buffers alive until its release callback is called by * the consumer. * * @param allocator Buffer allocator for allocating C data interface fields * @param vsr Vector schema root to export * @param provider Dictionary provider for dictionary encoded vectors (optional) * @param out C struct where to export the record batch * @param outSchema C struct where to export the record batch schema (optional) */ public static void exportVectorSchemaRoot( BufferAllocator allocator, VectorSchemaRoot vsr, DictionaryProvider provider, ArrowArray out, ArrowSchema outSchema) { if (outSchema != null) { exportSchema(allocator, vsr.getSchema(), provider, outSchema); } VectorUnloader unloader = new VectorUnloader(vsr); try (ArrowRecordBatch recordBatch = unloader.getRecordBatch()) { StructVectorLoader loader = new StructVectorLoader(vsr.getSchema()); try (StructVector vector = loader.load(allocator, recordBatch)) { exportVector(allocator, vector, provider, out); } } } /** * Export a reader as an ArrowArrayStream using the C Stream Interface. * * @param allocator Buffer allocator for allocating C data interface fields * @param reader Reader to export * @param out C struct to export the stream */ public static void exportArrayStream( BufferAllocator allocator, ArrowReader reader, ArrowArrayStream out) { 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. * *

    The given ArrowSchema struct is released (as per the C data interface specification), even * if this function fails. * * @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) * @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, boolean closeImportedStructs) { try { SchemaImporter importer = new SchemaImporter(allocator); return importer.importField(schema, provider); } finally { schema.release(); 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. * *

    The given ArrowSchema struct is released (as per the C data interface specification), even * if this function fails. * * @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) * @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, 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"); } 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. * *

    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, 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); } /** * Import Java vector and its type 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 vector. The ArrowSchema struct is released, even * if this function fails. * * @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) * @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, boolean closeImportedStructs) { Field field = importField(allocator, schema, provider, closeImportedStructs); FieldVector vector = field.createVector(allocator); 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. * *

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

    The schema of the vector schema root must match the input array (undefined behavior * otherwise). * * @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) * @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, boolean closeImportedStructs) { try (StructVector structVector = StructVector.emptyWithDuplicates("", allocator)) { structVector.initializeChildrenFromFields(root.getSchema().getFields()); importIntoVector(allocator, array, structVector, provider, closeImportedStructs); StructVectorUnloader unloader = new StructVectorUnloader(structVector); VectorLoader loader = new VectorLoader(root); try (ArrowRecordBatch recordBatch = unloader.getRecordBatch()) { loader.load(recordBatch); } } } /** * 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. * *

    The type represented by the ArrowSchema struct must be a struct type array. * *

    The ArrowSchema struct is released, even if this function fails. * * @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, 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); } /** * Import Java vector schema root from the C data interface. * *

    The type represented by the ArrowSchema struct must be a struct type array. * *

    The ArrowArray struct has its contents moved (as per the C data interface specification) to * a private object held alive by the resulting record batch. The ArrowSchema struct is released, * even if this function fails. * *

    Prefer {@link #importIntoVectorSchemaRoot} for loading array data while reusing the same * vector schema root. * * @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) * @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, boolean closeImportedStructs) { VectorSchemaRoot vsr = VectorSchemaRoot.create( importSchema(allocator, schema, provider, closeImportedStructs), allocator); if (array != null) { importIntoVectorSchemaRoot(allocator, array, vsr, provider, closeImportedStructs); } return vsr; } /** * 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 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; } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/Flags.java000066400000000000000000000033231515451615200252400ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeID; import org.apache.arrow.vector.types.pojo.Field; /** Flags as defined in the C data interface specification. */ final class Flags { static final int ARROW_FLAG_DICTIONARY_ORDERED = 1; static final int ARROW_FLAG_NULLABLE = 2; static final int ARROW_FLAG_MAP_KEYS_SORTED = 4; private Flags() {} static long forField(Field field) { long flags = 0L; if (field.isNullable()) { flags |= ARROW_FLAG_NULLABLE; } if (field.getDictionary() != null && field.getDictionary().isOrdered()) { flags |= ARROW_FLAG_DICTIONARY_ORDERED; } if (field.getType().getTypeID() == ArrowTypeID.Map) { ArrowType.Map map = (ArrowType.Map) field.getType(); if (map.getKeysSorted()) { flags |= ARROW_FLAG_MAP_KEYS_SORTED; } } return flags; } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/Format.java000066400000000000000000000321121515451615200254320ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.Arrays; import java.util.Locale; import java.util.stream.Collectors; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.IntervalUnit; 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.ArrowType.ExtensionType; /** * Conversion between {@link ArrowType} and string formats, as per C data interface specification. */ final class Format { private Format() {} static String asString(ArrowType arrowType) { if (arrowType instanceof ExtensionType) { ArrowType innerType = ((ExtensionType) arrowType).storageType(); return asString(innerType); } switch (arrowType.getTypeID()) { case Binary: return "z"; case Bool: return "b"; case Date: { ArrowType.Date type = (ArrowType.Date) arrowType; switch (type.getUnit()) { case DAY: return "tdD"; case MILLISECOND: return "tdm"; default: throw new UnsupportedOperationException( String.format("Date type with unit %s is unsupported", type.getUnit())); } } case Decimal: { ArrowType.Decimal type = (ArrowType.Decimal) arrowType; if (type.getBitWidth() == 128) { return String.format("d:%d,%d", type.getPrecision(), type.getScale()); } return String.format( "d:%d,%d,%d", type.getPrecision(), type.getScale(), type.getBitWidth()); } case Duration: { ArrowType.Duration type = (ArrowType.Duration) arrowType; switch (type.getUnit()) { case SECOND: return "tDs"; case MILLISECOND: return "tDm"; case MICROSECOND: return "tDu"; case NANOSECOND: return "tDn"; default: throw new UnsupportedOperationException( String.format("Duration type with unit %s is unsupported", type.getUnit())); } } case FixedSizeBinary: { ArrowType.FixedSizeBinary type = (ArrowType.FixedSizeBinary) arrowType; return String.format("w:%d", type.getByteWidth()); } case FixedSizeList: { ArrowType.FixedSizeList type = (ArrowType.FixedSizeList) arrowType; return String.format("+w:%d", type.getListSize()); } case FloatingPoint: { ArrowType.FloatingPoint type = (ArrowType.FloatingPoint) arrowType; switch (type.getPrecision()) { case HALF: return "e"; case SINGLE: return "f"; case DOUBLE: return "g"; default: throw new UnsupportedOperationException( String.format( "FloatingPoint type with precision %s is unsupported", type.getPrecision())); } } case Int: { String format; ArrowType.Int type = (ArrowType.Int) arrowType; switch (type.getBitWidth()) { case Byte.SIZE: format = "C"; break; case Short.SIZE: format = "S"; break; case Integer.SIZE: format = "I"; break; case Long.SIZE: format = "L"; break; default: throw new UnsupportedOperationException( String.format("Int type with bitwidth %d is unsupported", type.getBitWidth())); } if (type.getIsSigned()) { format = format.toLowerCase(Locale.ROOT); } return format; } case Interval: { ArrowType.Interval type = (ArrowType.Interval) arrowType; switch (type.getUnit()) { case DAY_TIME: return "tiD"; case YEAR_MONTH: return "tiM"; case MONTH_DAY_NANO: return "tin"; default: throw new UnsupportedOperationException( String.format("Interval type with unit %s is unsupported", type.getUnit())); } } case LargeBinary: return "Z"; case LargeList: return "+L"; case LargeUtf8: return "U"; case List: return "+l"; case Map: return "+m"; case Null: return "n"; case Struct: return "+s"; case Time: { ArrowType.Time type = (ArrowType.Time) arrowType; if (type.getUnit() == TimeUnit.SECOND && type.getBitWidth() == 32) { return "tts"; } else if (type.getUnit() == TimeUnit.MILLISECOND && type.getBitWidth() == 32) { return "ttm"; } else if (type.getUnit() == TimeUnit.MICROSECOND && type.getBitWidth() == 64) { return "ttu"; } else if (type.getUnit() == TimeUnit.NANOSECOND && type.getBitWidth() == 64) { return "ttn"; } else { throw new UnsupportedOperationException( String.format( "Time type with unit %s and bitwidth %d is unsupported", type.getUnit(), type.getBitWidth())); } } case Timestamp: { String format; ArrowType.Timestamp type = (ArrowType.Timestamp) arrowType; switch (type.getUnit()) { case SECOND: format = "tss"; break; case MILLISECOND: format = "tsm"; break; case MICROSECOND: format = "tsu"; break; case NANOSECOND: format = "tsn"; break; default: throw new UnsupportedOperationException( String.format("Timestamp type with unit %s is unsupported", type.getUnit())); } String timezone = type.getTimezone(); return String.format("%s:%s", format, timezone == null ? "" : timezone); } case Union: ArrowType.Union type = (ArrowType.Union) arrowType; String typeIDs = Arrays.stream(type.getTypeIds()) .mapToObj(String::valueOf) .collect(Collectors.joining(",")); switch (type.getMode()) { case Dense: return String.format("+ud:%s", typeIDs); case Sparse: return String.format("+us:%s", typeIDs); default: throw new UnsupportedOperationException( String.format("Union type with mode %s is unsupported", type.getMode())); } case Utf8: return "u"; case Utf8View: return "vu"; case BinaryView: return "vz"; case ListView: return "+vl"; case LargeListView: return "+vL"; case RunEndEncoded: return "+r"; case NONE: throw new IllegalArgumentException("Arrow type ID is NONE"); default: throw new UnsupportedOperationException( String.format("Unknown type id %s", arrowType.getTypeID())); } } static ArrowType asType(String format, long flags) throws NumberFormatException, UnsupportedOperationException, IllegalStateException { switch (format) { case "n": return new ArrowType.Null(); case "b": return new ArrowType.Bool(); case "c": return new ArrowType.Int(8, true); case "C": return new ArrowType.Int(8, false); case "s": return new ArrowType.Int(16, true); case "S": return new ArrowType.Int(16, false); case "i": return new ArrowType.Int(32, true); case "I": return new ArrowType.Int(32, false); case "l": return new ArrowType.Int(64, true); case "L": return new ArrowType.Int(64, false); case "e": return new ArrowType.FloatingPoint(FloatingPointPrecision.HALF); case "f": return new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE); case "g": return new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE); case "z": return new ArrowType.Binary(); case "Z": return new ArrowType.LargeBinary(); case "u": return new ArrowType.Utf8(); case "U": return new ArrowType.LargeUtf8(); case "tdD": return new ArrowType.Date(DateUnit.DAY); case "tdm": return new ArrowType.Date(DateUnit.MILLISECOND); case "tts": return new ArrowType.Time(TimeUnit.SECOND, 32); case "ttm": return new ArrowType.Time(TimeUnit.MILLISECOND, 32); case "ttu": return new ArrowType.Time(TimeUnit.MICROSECOND, 64); case "ttn": return new ArrowType.Time(TimeUnit.NANOSECOND, 64); case "tDs": return new ArrowType.Duration(TimeUnit.SECOND); case "tDm": return new ArrowType.Duration(TimeUnit.MILLISECOND); case "tDu": return new ArrowType.Duration(TimeUnit.MICROSECOND); case "tDn": return new ArrowType.Duration(TimeUnit.NANOSECOND); case "tiM": return new ArrowType.Interval(IntervalUnit.YEAR_MONTH); case "tiD": return new ArrowType.Interval(IntervalUnit.DAY_TIME); case "tin": return new ArrowType.Interval(IntervalUnit.MONTH_DAY_NANO); case "+l": return new ArrowType.List(); case "+L": return new ArrowType.LargeList(); case "+s": return new ArrowType.Struct(); case "+m": boolean keysSorted = (flags & Flags.ARROW_FLAG_MAP_KEYS_SORTED) != 0; return new ArrowType.Map(keysSorted); case "vu": return new ArrowType.Utf8View(); case "vz": return new ArrowType.BinaryView(); case "+vl": return new ArrowType.ListView(); case "+vL": return new ArrowType.LargeListView(); case "+r": return new ArrowType.RunEndEncoded(); default: String[] parts = format.split(":", 2); if (parts.length == 2) { return parseComplexFormat(parts[0], parts[1]); } throw new UnsupportedOperationException( String.format("Format %s is not supported", format)); } } private static ArrowType parseComplexFormat(String format, String payload) throws NumberFormatException, UnsupportedOperationException, IllegalStateException { switch (format) { case "d": { int[] parts = payloadToIntArray(payload); Preconditions.checkState( parts.length == 2 || parts.length == 3, "Format %s:%s is illegal", format, payload); int precision = parts[0]; int scale = parts[1]; Integer bitWidth = (parts.length == 3) ? parts[2] : null; return ArrowType.Decimal.createDecimal(precision, scale, bitWidth); } case "w": return new ArrowType.FixedSizeBinary(Integer.parseInt(payload)); case "+w": return new ArrowType.FixedSizeList(Integer.parseInt(payload)); case "+ud": return new ArrowType.Union(UnionMode.Dense, payloadToIntArray(payload)); case "+us": return new ArrowType.Union(UnionMode.Sparse, payloadToIntArray(payload)); case "tss": return new ArrowType.Timestamp(TimeUnit.SECOND, payloadToTimezone(payload)); case "tsm": return new ArrowType.Timestamp(TimeUnit.MILLISECOND, payloadToTimezone(payload)); case "tsu": return new ArrowType.Timestamp(TimeUnit.MICROSECOND, payloadToTimezone(payload)); case "tsn": return new ArrowType.Timestamp(TimeUnit.NANOSECOND, payloadToTimezone(payload)); default: throw new UnsupportedOperationException( String.format("Format %s:%s is not supported", format, payload)); } } private static int[] payloadToIntArray(String payload) throws NumberFormatException { return Arrays.stream(payload.split(",")).mapToInt(Integer::parseInt).toArray(); } private static String payloadToTimezone(String payload) { if (payload.isEmpty()) { return null; } return payload; } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/Metadata.java000066400000000000000000000064621515451615200257330ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.c.NativeUtil.NULL; import static org.apache.arrow.util.Preconditions.checkState; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.MemoryUtil; /** Encode and decode metadata. */ final class Metadata { private Metadata() {} static ArrowBuf encode(BufferAllocator allocator, Map metadata) { if (metadata == null || metadata.size() == 0) { return null; } List buffers = new ArrayList<>(metadata.size() * 2); int totalSize = 4 + metadata.size() * 8; // number of key/value pairs + buffer length fields for (Map.Entry entry : metadata.entrySet()) { byte[] keyBuffer = entry.getKey().getBytes(StandardCharsets.UTF_8); byte[] valueBuffer = entry.getValue().getBytes(StandardCharsets.UTF_8); totalSize += keyBuffer.length; totalSize += valueBuffer.length; buffers.add(keyBuffer); buffers.add(valueBuffer); } ArrowBuf result = allocator.buffer(totalSize); ByteBuffer writer = MemoryUtil.directBuffer(result.memoryAddress(), totalSize).order(ByteOrder.nativeOrder()); writer.putInt(metadata.size()); for (byte[] buffer : buffers) { writer.putInt(buffer.length); writer.put(buffer); } return result.slice(0, totalSize); } static Map decode(long bufferAddress) { if (bufferAddress == NULL) { return null; } ByteBuffer reader = MemoryUtil.directBuffer(bufferAddress, Integer.MAX_VALUE).order(ByteOrder.nativeOrder()); int size = reader.getInt(); checkState(size >= 0, "Metadata size must not be negative"); if (size == 0) { return null; } Map result = new HashMap<>(size); for (int i = 0; i < size; i++) { String key = readString(reader); String value = readString(reader); result.put(key, value); } return result; } private static String readString(ByteBuffer reader) { int length = reader.getInt(); checkState(length >= 0, "Metadata item length must not be negative"); String result = ""; if (length > 0) { byte[] dst = new byte[length]; reader.get(dst); result = new String(dst, StandardCharsets.UTF_8); } return result; } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/NativeUtil.java000066400000000000000000000101251515451615200262660ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.nio.Buffer; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.nio.charset.StandardCharsets; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.MemoryUtil; /** Utility functions for working with native memory. */ public final class NativeUtil { public static final byte NULL = 0; static final int MAX_STRING_LENGTH = Short.MAX_VALUE; private NativeUtil() {} /** * Convert a pointer to a null terminated string into a Java String. * * @param cstringPtr pointer to C string * @return Converted string */ public static String toJavaString(long cstringPtr) { if (cstringPtr == NULL) { return null; } ByteBuffer reader = MemoryUtil.directBuffer(cstringPtr, MAX_STRING_LENGTH).order(ByteOrder.nativeOrder()); int length = 0; while (reader.get() != NULL) { length++; } byte[] bytes = new byte[length]; // Force use of base class rewind() to avoid breaking change of ByteBuffer.rewind in JDK9+ ((ByteBuffer) ((Buffer) reader).rewind()).get(bytes); return new String(bytes, 0, length, StandardCharsets.UTF_8); } /** * Convert a native array pointer (void**) to Java array of pointers. * * @param arrayPtr Array pointer * @param size Array size * @return Array of pointer values as longs */ public static long[] toJavaArray(long arrayPtr, int size) { if (arrayPtr == NULL) { return null; } if (size < 0) { throw new IllegalArgumentException("Invalid native array size"); } long[] result = new long[size]; ByteBuffer reader = MemoryUtil.directBuffer(arrayPtr, Long.BYTES * size).order(ByteOrder.nativeOrder()); for (int i = 0; i < size; i++) { result[i] = reader.getLong(); } return result; } /** * Convert Java string to a null terminated string. * * @param allocator Buffer allocator for allocating the native string * @param string Input String to convert * @return Buffer with a null terminated string or null if the input is null */ public static ArrowBuf toNativeString(BufferAllocator allocator, String string) { if (string == null) { return null; } byte[] bytes = string.getBytes(StandardCharsets.UTF_8); ArrowBuf buffer = allocator.buffer(bytes.length + 1); buffer.writeBytes(bytes); buffer.writeByte(NULL); return buffer; } /** * Close a buffer if it's not null. * * @param buf Buffer to close */ public static void closeBuffer(ArrowBuf buf) { if (buf != null) { buf.close(); } } /** * Get the address of a buffer or {@value #NULL} if the input buffer is null. * * @param buf Buffer to get the address of * @return Memory address or {@value #NULL} */ public static long addressOrNull(ArrowBuf buf) { if (buf == null) { return NULL; } return buf.memoryAddress(); } /** * Get the address of a C Data Interface struct or {@value #NULL} if the input struct is null. * * @param struct C Data Interface struct to get the address of * @return Memory address or {@value #NULL} */ public static long addressOrNull(BaseStruct struct) { if (struct == null) { return NULL; } return struct.memoryAddress(); } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java000066400000000000000000000053711515451615200314430ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.concurrent.atomic.AtomicInteger; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ForeignAllocation; /** * The owner of an imported C Data Interface array. * *

    There is a fundamental mismatch here between memory allocation schemes: AllocationManager * represents a single allocation (= a single address and length). But an ArrowArray combines * multiple allocations behind a single deallocation callback. This class bridges the two by * tracking a reference count, so that the single callback can be managed by multiple {@link * ForeignAllocation} instances. */ final class ReferenceCountedArrowArray { private final ArrowArray array; private final AtomicInteger refCnt; ReferenceCountedArrowArray(ArrowArray array) { this.array = array; this.refCnt = new AtomicInteger(1); } void retain() { if (refCnt.addAndGet(1) - 1 <= 0) { throw new IllegalStateException("Tried to retain a released ArrowArray"); } } void release() { int refcnt = refCnt.addAndGet(-1); if (refcnt == 0) { array.release(); array.close(); } else if (refcnt < 0) { throw new IllegalStateException("Reference count went negative for imported ArrowArray"); } } /** * Create an ArrowBuf wrapping a buffer from this ArrowArray associated with the given * BufferAllocator. * *

    This method is "unsafe" because there is no validation of the given capacity or address. If * the returned buffer is not freed, a memory leak will occur. */ ArrowBuf unsafeAssociateAllocation( BufferAllocator trackingAllocator, long capacity, long memoryAddress) { retain(); return trackingAllocator.wrapForeignAllocation( new ForeignAllocation(capacity, memoryAddress) { @Override protected void release0() { ReferenceCountedArrowArray.this.release(); } }); } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/SchemaExporter.java000066400000000000000000000111271515451615200271360ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.c.NativeUtil.NULL; import static org.apache.arrow.c.NativeUtil.addressOrNull; import static org.apache.arrow.util.Preconditions.checkNotNull; import java.util.ArrayList; import java.util.List; import org.apache.arrow.c.jni.JniWrapper; import org.apache.arrow.c.jni.PrivateData; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.types.pojo.DictionaryEncoding; import org.apache.arrow.vector.types.pojo.Field; /** Exporter for {@link ArrowSchema}. */ final class SchemaExporter { private final BufferAllocator allocator; public SchemaExporter(BufferAllocator allocator) { this.allocator = allocator; } /** Private data structure for exported schemas. */ static class ExportedSchemaPrivateData implements PrivateData { ArrowBuf format; ArrowBuf name; ArrowBuf metadata; ArrowBuf children_ptrs; ArrowSchema dictionary; List children; @Override public void close() { NativeUtil.closeBuffer(format); NativeUtil.closeBuffer(name); NativeUtil.closeBuffer(metadata); NativeUtil.closeBuffer(children_ptrs); if (dictionary != null) { dictionary.close(); } if (children != null) { for (ArrowSchema child : children) { child.close(); } } } } void export(ArrowSchema schema, Field field, DictionaryProvider dictionaryProvider) { String name = field.getName(); String format = Format.asString(field.getType()); long flags = Flags.forField(field); List children = field.getChildren(); DictionaryEncoding dictionaryEncoding = field.getDictionary(); ExportedSchemaPrivateData data = new ExportedSchemaPrivateData(); try { data.format = NativeUtil.toNativeString(allocator, format); data.name = NativeUtil.toNativeString(allocator, name); data.metadata = Metadata.encode(allocator, field.getMetadata()); if (children != null) { data.children = new ArrayList<>(children.size()); data.children_ptrs = allocator.buffer((long) children.size() * Long.BYTES); for (int i = 0; i < children.size(); i++) { ArrowSchema child = ArrowSchema.allocateNew(allocator); data.children.add(child); data.children_ptrs.writeLong(child.memoryAddress()); } } if (dictionaryEncoding != null) { Dictionary dictionary = dictionaryProvider.lookup(dictionaryEncoding.getId()); checkNotNull(dictionary, "Dictionary lookup failed on export of field with dictionary"); data.dictionary = ArrowSchema.allocateNew(allocator); export(data.dictionary, dictionary.getVector().getField(), dictionaryProvider); } ArrowSchema.Snapshot snapshot = new ArrowSchema.Snapshot(); snapshot.format = data.format.memoryAddress(); snapshot.name = addressOrNull(data.name); snapshot.metadata = addressOrNull(data.metadata); snapshot.flags = flags; snapshot.n_children = (data.children != null) ? data.children.size() : 0; snapshot.children = addressOrNull(data.children_ptrs); snapshot.dictionary = addressOrNull(data.dictionary); snapshot.release = NULL; schema.save(snapshot); // sets release and private data JniWrapper.get().exportSchema(schema.memoryAddress(), data); } catch (Exception e) { data.close(); throw e; } // Export children if (children != null) { for (int i = 0; i < children.size(); i++) { Field childField = children.get(i); ArrowSchema child = data.children.get(i); export(child, childField, dictionaryProvider); } } } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/SchemaImporter.java000066400000000000000000000114471515451615200271340ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.c.NativeUtil.NULL; import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import static org.apache.arrow.util.Preconditions.checkNotNull; import static org.apache.arrow.util.Preconditions.checkState; import java.util.ArrayList; import java.util.List; import java.util.Map; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; import org.apache.arrow.vector.types.pojo.DictionaryEncoding; 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.slf4j.Logger; import org.slf4j.LoggerFactory; /** Importer for {@link ArrowSchema}. */ final class SchemaImporter { private static final Logger logger = LoggerFactory.getLogger(SchemaImporter.class); private static final int MAX_IMPORT_RECURSION_LEVEL = 64; private long nextDictionaryID = 0L; private final BufferAllocator allocator; public SchemaImporter(BufferAllocator allocator) { this.allocator = allocator; } Field importField(ArrowSchema schema, CDataDictionaryProvider provider) { return importField(schema, provider, 0); } private Field importField( ArrowSchema schema, CDataDictionaryProvider provider, int recursionLevel) { checkState( recursionLevel <= MAX_IMPORT_RECURSION_LEVEL, "Recursion level in ArrowSchema struct exceeded"); ArrowSchema.Snapshot snapshot = schema.snapshot(); checkState(snapshot.release != NULL, "Cannot import released ArrowSchema"); String name = NativeUtil.toJavaString(snapshot.name); String format = NativeUtil.toJavaString(snapshot.format); checkNotNull(format, "format field must not be null"); ArrowType arrowType = Format.asType(format, snapshot.flags); boolean nullable = (snapshot.flags & Flags.ARROW_FLAG_NULLABLE) != 0; Map metadata = Metadata.decode(snapshot.metadata); if (metadata != null && metadata.containsKey(ExtensionType.EXTENSION_METADATA_KEY_NAME)) { final String extensionName = metadata.get(ExtensionType.EXTENSION_METADATA_KEY_NAME); final String extensionMetadata = metadata.getOrDefault(ExtensionType.EXTENSION_METADATA_KEY_METADATA, ""); ExtensionType extensionType = ExtensionTypeRegistry.lookup(extensionName); if (extensionType != null) { arrowType = extensionType.deserialize(arrowType, extensionMetadata); } else { // Otherwise, we haven't registered the type logger.info("Unrecognized extension type: {}", extensionName); } } // Handle dictionary encoded vectors DictionaryEncoding dictionaryEncoding = null; if (snapshot.dictionary != NULL && provider != null) { boolean ordered = (snapshot.flags & Flags.ARROW_FLAG_DICTIONARY_ORDERED) != 0; ArrowType.Int indexType = (ArrowType.Int) arrowType; dictionaryEncoding = new DictionaryEncoding(nextDictionaryID++, ordered, indexType); ArrowSchema dictionarySchema = ArrowSchema.wrap(snapshot.dictionary); Field dictionaryField = importField(dictionarySchema, provider, recursionLevel + 1); provider.put(new Dictionary(dictionaryField.createVector(allocator), dictionaryEncoding)); } FieldType fieldType = new FieldType(nullable, arrowType, dictionaryEncoding, metadata); List children = null; long[] childrenIds = NativeUtil.toJavaArray(snapshot.children, checkedCastToInt(snapshot.n_children)); if (childrenIds != null && childrenIds.length > 0) { children = new ArrayList<>(childrenIds.length); for (long childAddress : childrenIds) { ArrowSchema childSchema = ArrowSchema.wrap(childAddress); Field field = importField(childSchema, provider, recursionLevel + 1); children.add(field); } } return new Field(name, fieldType, children); } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/StructVectorLoader.java000066400000000000000000000162121515451615200300030ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.util.Preconditions.checkArgument; import java.util.ArrayList; import java.util.Collections; import java.util.Iterator; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Collections2; import org.apache.arrow.vector.BaseVariableWidthViewVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.TypeLayout; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.compression.CompressionCodec; import org.apache.arrow.vector.compression.CompressionUtil; import org.apache.arrow.vector.compression.NoCompressionCodec; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; /** Loads buffers into {@link StructVector}. */ public class StructVectorLoader { private final Schema schema; private final CompressionCodec.Factory factory; /** * A flag indicating if decompression is needed. This will affect the behavior of releasing * buffers. */ private boolean decompressionNeeded; /** * Construct with a schema. * *

    The schema referred to here can be obtained from the struct vector. The schema here should * be the children of a struct vector, not a schema containing the struct field itself. For * example: * Schema schema = new Schema(structVector.getField().getChildren()); * * * @param schema buffers are added based on schema. */ public StructVectorLoader(Schema schema) { this(schema, NoCompressionCodec.Factory.INSTANCE); } /** * Construct with a schema and a compression codec factory. * * @param schema buffers are added based on schema. * @param factory the factory to create codec. */ public StructVectorLoader(Schema schema, CompressionCodec.Factory factory) { this.schema = schema; this.factory = factory; } /** * Loads the record batch into the struct vector. * *

    This will not close the record batch. * * @param recordBatch the batch to load */ public StructVector load(BufferAllocator allocator, ArrowRecordBatch recordBatch) { StructVector result = StructVector.emptyWithDuplicates("", allocator); result.initializeChildrenFromFields(this.schema.getFields()); Iterator buffers = recordBatch.getBuffers().iterator(); Iterator nodes = recordBatch.getNodes().iterator(); CompressionUtil.CodecType codecType = CompressionUtil.CodecType.fromCompressionType(recordBatch.getBodyCompression().getCodec()); decompressionNeeded = codecType != CompressionUtil.CodecType.NO_COMPRESSION; CompressionCodec codec = decompressionNeeded ? factory.createCodec(codecType) : NoCompressionCodec.INSTANCE; Iterator variadicBufferCounts = Collections.emptyIterator(); if (recordBatch.getVariadicBufferCounts() != null && !recordBatch.getVariadicBufferCounts().isEmpty()) { variadicBufferCounts = recordBatch.getVariadicBufferCounts().iterator(); } for (FieldVector fieldVector : result.getChildrenFromFields()) { loadBuffers(fieldVector, fieldVector.getField(), buffers, nodes, codec, variadicBufferCounts); } result.loadFieldBuffers( new ArrowFieldNode(recordBatch.getLength(), 0), Collections.singletonList(null)); if (nodes.hasNext() || buffers.hasNext() || variadicBufferCounts.hasNext()) { throw new IllegalArgumentException( "not all nodes, buffers and variadicBufferCounts were consumed. nodes: " + Collections2.toString(nodes) + " buffers: " + Collections2.toString(buffers) + " variadicBufferCounts: " + Collections2.toString(variadicBufferCounts)); } return result; } private void loadBuffers( FieldVector vector, Field field, Iterator buffers, Iterator nodes, CompressionCodec codec, Iterator variadicBufferCounts) { checkArgument(nodes.hasNext(), "no more field nodes for field %s and vector %s", field, vector); ArrowFieldNode fieldNode = nodes.next(); // variadicBufferLayoutCount will be 0 for vectors of a type except BaseVariableWidthViewVector long variadicBufferLayoutCount = 0; if (vector instanceof BaseVariableWidthViewVector) { if (variadicBufferCounts.hasNext()) { variadicBufferLayoutCount = variadicBufferCounts.next(); } else { throw new IllegalStateException( "No variadicBufferCounts available for BaseVariableWidthViewVector"); } } int bufferLayoutCount = (int) (variadicBufferLayoutCount + TypeLayout.getTypeBufferCount(field.getType())); List ownBuffers = new ArrayList<>(bufferLayoutCount); for (int j = 0; j < bufferLayoutCount; j++) { ArrowBuf nextBuf = buffers.next(); // for vectors without nulls, the buffer is empty, so there is no need to // decompress it. ArrowBuf bufferToAdd = nextBuf.writerIndex() > 0 ? codec.decompress(vector.getAllocator(), nextBuf) : nextBuf; ownBuffers.add(bufferToAdd); if (decompressionNeeded) { // decompression performed nextBuf.getReferenceManager().retain(); } } try { vector.loadFieldBuffers(fieldNode, ownBuffers); if (decompressionNeeded) { for (ArrowBuf buf : ownBuffers) { buf.close(); } } } catch (RuntimeException e) { throw new IllegalArgumentException( "Could not load buffers for field " + field + ". error message: " + e.getMessage(), e); } List children = field.getChildren(); if (children.size() > 0) { List childrenFromFields = vector.getChildrenFromFields(); checkArgument( children.size() == childrenFromFields.size(), "should have as many children as in the schema: found %s expected %s", childrenFromFields.size(), children.size()); for (int i = 0; i < childrenFromFields.size(); i++) { Field child = children.get(i); FieldVector fieldVector = childrenFromFields.get(i); loadBuffers(fieldVector, child, buffers, nodes, codec, variadicBufferCounts); } } } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/StructVectorUnloader.java000066400000000000000000000121271515451615200303470ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.ArrayList; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.BaseVariableWidthViewVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.TypeLayout; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.compression.CompressionCodec; import org.apache.arrow.vector.compression.CompressionUtil; import org.apache.arrow.vector.compression.NoCompressionCodec; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; /** Helper class that handles converting a {@link StructVector} to a {@link ArrowRecordBatch}. */ public class StructVectorUnloader { private final StructVector root; private final boolean includeNullCount; private final CompressionCodec codec; private final boolean alignBuffers; /** Constructs a new instance of the given struct vector. */ public StructVectorUnloader(StructVector root) { this(root, true, NoCompressionCodec.INSTANCE, true); } /** * Constructs a new instance. * * @param root The struct vector to serialize to an {@link ArrowRecordBatch}. * @param includeNullCount Controls whether null count is copied to the {@link ArrowRecordBatch} * @param alignBuffers Controls if buffers get aligned to 8-byte boundaries. */ public StructVectorUnloader(StructVector root, boolean includeNullCount, boolean alignBuffers) { this(root, includeNullCount, NoCompressionCodec.INSTANCE, alignBuffers); } /** * Constructs a new instance. * * @param root The struct vector to serialize to an {@link ArrowRecordBatch}. * @param includeNullCount Controls whether null count is copied to the {@link ArrowRecordBatch} * @param codec the codec for compressing data. If it is null, then no compression is needed. * @param alignBuffers Controls if buffers get aligned to 8-byte boundaries. */ public StructVectorUnloader( StructVector root, boolean includeNullCount, CompressionCodec codec, boolean alignBuffers) { this.root = root; this.includeNullCount = includeNullCount; this.codec = codec; this.alignBuffers = alignBuffers; } /** * Performs the depth first traversal of the Vectors to create an {@link ArrowRecordBatch} * suitable for serialization. */ public ArrowRecordBatch getRecordBatch() { List nodes = new ArrayList<>(); List buffers = new ArrayList<>(); List variadicBufferCounts = new ArrayList<>(); for (FieldVector vector : root.getChildrenFromFields()) { appendNodes(vector, nodes, buffers, variadicBufferCounts); } return new ArrowRecordBatch( root.getValueCount(), nodes, buffers, CompressionUtil.createBodyCompression(codec), variadicBufferCounts, alignBuffers); } private long getVariadicBufferCount(FieldVector vector) { if (vector instanceof BaseVariableWidthViewVector) { return ((BaseVariableWidthViewVector) vector).getDataBuffers().size(); } return 0L; } private void appendNodes( FieldVector vector, List nodes, List buffers, List variadicBufferCounts) { nodes.add( new ArrowFieldNode(vector.getValueCount(), includeNullCount ? vector.getNullCount() : -1)); List fieldBuffers = vector.getFieldBuffers(); long variadicBufferCount = getVariadicBufferCount(vector); int expectedBufferCount = (int) (TypeLayout.getTypeBufferCount(vector.getField().getType()) + variadicBufferCount); // only update variadicBufferCounts for vectors that have variadic buffers if (vector instanceof BaseVariableWidthViewVector) { variadicBufferCounts.add(variadicBufferCount); } if (fieldBuffers.size() != expectedBufferCount) { throw new IllegalArgumentException( String.format( "wrong number of buffers for field %s in vector %s. found: %s", vector.getField(), vector.getClass().getSimpleName(), fieldBuffers)); } for (ArrowBuf buf : fieldBuffers) { buffers.add(codec.compress(vector.getAllocator(), buf)); } for (FieldVector child : vector.getChildrenFromFields()) { appendNodes(child, nodes, buffers, variadicBufferCounts); } } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/jni/000077500000000000000000000000001515451615200241205ustar00rootroot00000000000000apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/jni/CDataJniException.java000066400000000000000000000024631515451615200302640ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jni; /** An exception raised by the JNI side of the C Data bridge. */ public final class CDataJniException extends Exception { private final int errno; public CDataJniException(int errno, String message) { super(message); this.errno = errno; } /** The original error code returned from C. */ public int getErrno() { return errno; } @Override public String getMessage() { return "CDataJniException{" + "errno=" + errno + ", message=" + super.getMessage() + '}'; } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/jni/JniLoader.java000066400000000000000000000077731515451615200266500ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jni; import java.io.File; import java.io.FileNotFoundException; import java.io.IOException; import java.io.InputStream; import java.nio.file.Files; import java.nio.file.StandardCopyOption; import java.util.ArrayList; import java.util.Collections; import java.util.HashSet; import java.util.List; import java.util.Locale; import java.util.Set; /** The JniLoader for C Data Interface API's native implementation. */ public class JniLoader { private static final JniLoader INSTANCE = new JniLoader(Collections.singletonList("arrow_cdata_jni")); public static JniLoader get() { return INSTANCE; } private final Set librariesToLoad; private JniLoader(List libraryNames) { librariesToLoad = new HashSet<>(libraryNames); } private boolean finished() { return librariesToLoad.isEmpty(); } /** If required JNI libraries are not loaded, then load them. */ public void ensureLoaded() { if (finished()) { return; } loadRemaining(); } private synchronized void loadRemaining() { // The method is protected by a mutex via synchronized, if more than one thread // race to call // loadRemaining, at same time only one will do the actual loading and the // others will wait for // the mutex to be acquired then check on the remaining list: if there are // libraries that were not // successfully loaded then the mutex owner will try to load them again. if (finished()) { return; } List libs = new ArrayList<>(librariesToLoad); for (String lib : libs) { load(lib); librariesToLoad.remove(lib); } } private void load(String 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"))); temp.deleteOnExit(); try (final InputStream is = JniWrapper.class.getClassLoader().getResourceAsStream(libraryToLoad)) { if (is == null) { throw new FileNotFoundException(libraryToLoad); } Files.copy(is, temp.toPath(), StandardCopyOption.REPLACE_EXISTING); System.load(temp.getAbsolutePath()); } } catch (IOException e) { throw new IllegalStateException("error loading native libraries: " + e); } } private String getNormalizedArch() { String arch = System.getProperty("os.arch").toLowerCase(Locale.US); switch (arch) { case "amd64": arch = "x86_64"; break; case "aarch64": arch = "aarch_64"; break; default: break; } return arch; } } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/jni/JniWrapper.java000066400000000000000000000040241515451615200270440ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jni; /** JniWrapper for C Data Interface API implementation. */ public class JniWrapper { private static final JniWrapper INSTANCE = new JniWrapper(); public static JniWrapper get() { JniLoader.get().ensureLoaded(); return INSTANCE; } private JniWrapper() { // A best effort to error on 32-bit systems String dataModel = System.getProperty("sun.arch.data.model"); if (dataModel != null && dataModel.equals("32")) { throw new UnsupportedOperationException( "The Java C Data Interface implementation is currently only supported on 64-bit systems"); } } public native void releaseSchema(long memoryAddress); public native void releaseArray(long memoryAddress); public native void getNextArrayStream(long streamAddress, long arrayAddress) throws CDataJniException; public native void getSchemaArrayStream(long streamAddress, long arrayAddress) throws CDataJniException; public native void releaseArrayStream(long memoryAddress); public native void exportSchema(long memoryAddress, PrivateData privateData); public native void exportArray(long memoryAddress, PrivateData data); public native void exportArrayStream(long memoryAddress, PrivateData data); } apache-arrow-java-19.0.0/c/src/main/java/org/apache/arrow/c/jni/PrivateData.java000066400000000000000000000020651515451615200271720ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jni; import java.io.Closeable; /** * Interface for Java objects stored in C data interface private data. * *

    This interface is used for exported structures. */ public interface PrivateData extends Closeable { @Override void close(); } apache-arrow-java-19.0.0/c/src/test/000077500000000000000000000000001515451615200170665ustar00rootroot00000000000000apache-arrow-java-19.0.0/c/src/test/java/000077500000000000000000000000001515451615200200075ustar00rootroot00000000000000apache-arrow-java-19.0.0/c/src/test/java/org/000077500000000000000000000000001515451615200205765ustar00rootroot00000000000000apache-arrow-java-19.0.0/c/src/test/java/org/apache/000077500000000000000000000000001515451615200220175ustar00rootroot00000000000000apache-arrow-java-19.0.0/c/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200231515ustar00rootroot00000000000000apache-arrow-java-19.0.0/c/src/test/java/org/apache/arrow/c/000077500000000000000000000000001515451615200233735ustar00rootroot00000000000000apache-arrow-java-19.0.0/c/src/test/java/org/apache/arrow/c/ArrowArrayUtilityTest.java000066400000000000000000000153271515451615200305630ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.junit.jupiter.api.Assertions.assertThrows; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReferenceManager; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.memory.util.MemoryUtil; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; 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; class ArrowArrayUtilityTest { BufferAllocator allocator; ArrowArray arrowArray; ReferenceCountedArrowArray dummyHandle; @BeforeEach void beforeEach() { allocator = new RootAllocator(); arrowArray = ArrowArray.allocateNew(allocator); dummyHandle = new ReferenceCountedArrowArray(arrowArray); } @AfterEach void afterEach() { dummyHandle.release(); allocator.close(); } @Test void arraySchemaInit() { ArrowSchema schema = ArrowSchema.allocateNew(allocator); assertThat(schema.snapshot().release).isEqualTo(0); schema.close(); } // ------------------------------------------------------------ // BufferImportTypeVisitor @Test void importBuffer() throws Exception { // Note values are all dummy values here try (BufferImportTypeVisitor notEmptyDataVisitor = new BufferImportTypeVisitor( allocator, dummyHandle, new ArrowFieldNode(/* length= */ 1, 0), new long[] {0})) { // Too few buffers assertThrows( IllegalStateException.class, () -> notEmptyDataVisitor.importBuffer(new ArrowType.Bool(), 1, 1)); // Null where one isn't expected assertThrows( IllegalStateException.class, () -> notEmptyDataVisitor.importBuffer(new ArrowType.Bool(), 0, 1)); // Expected capacity not zero but c array ptr is NULL (zero) assertThrows( IllegalStateException.class, () -> notEmptyDataVisitor.importBuffer(new ArrowType.Bool(), 0, 1)); // Expected capacity is zero and c array ptr is NULL (zero) assertThat(notEmptyDataVisitor.importBuffer(new ArrowType.Bool(), 0, 0)) .isEqualTo(allocator.getEmpty()); } try (BufferImportTypeVisitor emptyDataVisitor = new BufferImportTypeVisitor( allocator, dummyHandle, new ArrowFieldNode(/* length= */ 0, 0), new long[] {0})) { // Too few buffers assertThrows( IllegalStateException.class, () -> emptyDataVisitor.importBuffer(new ArrowType.Bool(), 1, 1)); // Expected capacity not zero but c array ptr is NULL (zero) assertThrows( IllegalStateException.class, () -> emptyDataVisitor.importBuffer(new ArrowType.Bool(), 0, 1)); // Expected capacity is zero and c array ptr is NULL (zero) assertThat(emptyDataVisitor.importBuffer(new ArrowType.Bool(), 0, 0)) .isEqualTo(allocator.getEmpty()); } } @Test void cleanupAfterFailure() throws Exception { // Note values are all dummy values here long address = MemoryUtil.allocateMemory(16); try (BufferImportTypeVisitor visitor = new BufferImportTypeVisitor( allocator, dummyHandle, new ArrowFieldNode(0, 0), new long[] {address})) { // This fails, but only after we've already imported a buffer. assertThrows(IllegalStateException.class, () -> visitor.visit(new ArrowType.Int(32, true))); } finally { MemoryUtil.freeMemory(address); } } @Test void bufferAssociatedWithAllocator() throws Exception { // Note values are all dummy values here final long bufferSize = 16; final long fieldLength = bufferSize / IntVector.TYPE_WIDTH; long address = MemoryUtil.allocateMemory(bufferSize); long baseline = allocator.getAllocatedMemory(); ArrowFieldNode fieldNode = new ArrowFieldNode(fieldLength, 0); try (BufferImportTypeVisitor visitor = new BufferImportTypeVisitor(allocator, dummyHandle, fieldNode, new long[] {0, address})) { List buffers = visitor.visit(new ArrowType.Int(32, true)); assertThat(buffers).hasSize(2); assertThat(buffers.get(0)).isNull(); assertThat(buffers.get(1)) .isNotNull() .extracting(ArrowBuf::getReferenceManager) .extracting(ReferenceManager::getAllocator) .isEqualTo(allocator); assertThat(allocator.getAllocatedMemory()).isEqualTo(baseline + bufferSize); } finally { MemoryUtil.freeMemory(address); } assertThat(allocator.getAllocatedMemory()).isEqualTo(baseline); } // ------------------------------------------------------------ // ReferenceCountedArrowArray @Test void releaseRetain() { ArrowArray array = ArrowArray.allocateNew(allocator); ReferenceCountedArrowArray handle = new ReferenceCountedArrowArray(array); assertThat(array.isClosed()).isFalse(); handle.retain(); assertThat(array.isClosed()).isFalse(); handle.release(); assertThat(array.isClosed()).isFalse(); handle.release(); assertThat(array.isClosed()).isTrue(); assertThrows(IllegalStateException.class, handle::release); assertThrows(IllegalStateException.class, handle::retain); } @Test void associate() { final long bufferSize = 16; final long address = MemoryUtil.allocateMemory(bufferSize); try { ArrowArray array = ArrowArray.allocateNew(allocator); ReferenceCountedArrowArray handle = new ReferenceCountedArrowArray(array); assertThat(array.isClosed()).isFalse(); ArrowBuf buf = handle.unsafeAssociateAllocation(allocator, bufferSize, address); assertThat(array.isClosed()).isFalse(); buf.close(); assertThat(array.isClosed()).isFalse(); handle.release(); assertThat(array.isClosed()).isTrue(); } finally { MemoryUtil.freeMemory(address); } } } apache-arrow-java-19.0.0/c/src/test/java/org/apache/arrow/c/DictionaryTest.java000066400000000000000000000351701515451615200272110ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.channels.Channels; import java.nio.charset.StandardCharsets; import java.util.Collections; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.compare.VectorEqualsVisitor; import org.apache.arrow.vector.complex.StructVector; 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.ipc.ArrowStreamReader; import org.apache.arrow.vector.ipc.ArrowStreamWriter; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.types.Types.MinorType; 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.types.pojo.Schema; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class DictionaryTest { private RootAllocator allocator = null; @BeforeEach public void setUp() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void tearDown() { allocator.close(); } void roundtrip(FieldVector vector, DictionaryProvider provider, Class clazz) { // Consumer allocates empty structures 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 Data.exportVector(allocator, vector, provider, arrowArray, arrowSchema); } // Consumer imports vector try (CDataDictionaryProvider cDictionaryProvider = new CDataDictionaryProvider(); FieldVector imported = Data.importVector( allocator, consumerArrowArray, consumerArrowSchema, cDictionaryProvider); ) { assertTrue( clazz.isInstance(imported), String.format("expected %s but was %s", clazz, imported.getClass())); assertTrue( VectorEqualsVisitor.vectorEquals(vector, imported), "vectors are not equivalent"); for (long id : cDictionaryProvider.getDictionaryIds()) { ValueVector exportedDictionaryVector = provider.lookup(id).getVector(); ValueVector importedDictionaryVector = cDictionaryProvider.lookup(id).getVector(); assertTrue( VectorEqualsVisitor.vectorEquals(exportedDictionaryVector, importedDictionaryVector), String.format("Dictionary vectors for ID %d are not equivalent", id)); } } } } @Test public void testWithDictionary() throws Exception { DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); // create dictionary and provider final VarCharVector dictVector = new VarCharVector("dict", allocator); dictVector.allocateNewSafe(); dictVector.setSafe(0, "aa".getBytes(StandardCharsets.UTF_8)); dictVector.setSafe(1, "bb".getBytes(StandardCharsets.UTF_8)); dictVector.setSafe(2, "cc".getBytes(StandardCharsets.UTF_8)); dictVector.setValueCount(3); Dictionary dictionary = new Dictionary(dictVector, new DictionaryEncoding(0L, false, /* indexType= */ null)); provider.put(dictionary); // create vector and encode it final VarCharVector vector = new VarCharVector("vector", allocator); vector.allocateNewSafe(); vector.setSafe(0, "bb".getBytes(StandardCharsets.UTF_8)); vector.setSafe(1, "bb".getBytes(StandardCharsets.UTF_8)); vector.setSafe(2, "cc".getBytes(StandardCharsets.UTF_8)); vector.setSafe(3, "aa".getBytes(StandardCharsets.UTF_8)); vector.setValueCount(4); // get the encoded vector IntVector encodedVector = (IntVector) DictionaryEncoder.encode(vector, dictionary); // Perform roundtrip using C Data Interface roundtrip(encodedVector, provider, IntVector.class); // Close all AutoCloseables.close((AutoCloseable) vector, encodedVector, dictVector); } @Test public void testRoundtripMultipleBatches() throws IOException { try (ArrowStreamReader reader = createMultiBatchReader(); ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator)) { // Load first batch reader.loadNextBatch(); // Producer fills consumer schema structure Data.exportSchema( allocator, reader.getVectorSchemaRoot().getSchema(), reader, consumerArrowSchema); // Consumer loads it as an empty vector schema root try (CDataDictionaryProvider consumerDictionaryProvider = new CDataDictionaryProvider(); VectorSchemaRoot consumerRoot = Data.importVectorSchemaRoot( allocator, consumerArrowSchema, consumerDictionaryProvider)) { do { try (ArrowArray consumerArray = ArrowArray.allocateNew(allocator)) { // Producer exports next data Data.exportVectorSchemaRoot( allocator, reader.getVectorSchemaRoot(), reader, consumerArray); // Consumer loads next data Data.importIntoVectorSchemaRoot( allocator, consumerArray, consumerRoot, consumerDictionaryProvider); // Roundtrip validation assertTrue( consumerRoot.equals(reader.getVectorSchemaRoot()), "vector schema roots are not equivalent"); for (long id : consumerDictionaryProvider.getDictionaryIds()) { ValueVector exportedDictionaryVector = reader.lookup(id).getVector(); ValueVector importedDictionaryVector = consumerDictionaryProvider.lookup(id).getVector(); assertTrue( VectorEqualsVisitor.vectorEquals( exportedDictionaryVector, importedDictionaryVector), String.format("Dictionary vectors for ID %d are not equivalent", id)); } } } while (reader.loadNextBatch()); } } } private ArrowStreamReader createMultiBatchReader() throws IOException { ByteArrayOutputStream os = new ByteArrayOutputStream(); try (final VarCharVector dictVector = new VarCharVector("dict", allocator); IntVector vector = new IntVector("foo", allocator)) { // create dictionary and provider DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); dictVector.allocateNewSafe(); dictVector.setSafe(0, "aa".getBytes(StandardCharsets.UTF_8)); dictVector.setSafe(1, "bb".getBytes(StandardCharsets.UTF_8)); dictVector.setSafe(2, "cc".getBytes(StandardCharsets.UTF_8)); dictVector.setSafe(3, "dd".getBytes(StandardCharsets.UTF_8)); dictVector.setSafe(4, "ee".getBytes(StandardCharsets.UTF_8)); dictVector.setValueCount(5); Dictionary dictionary = new Dictionary(dictVector, new DictionaryEncoding(0L, false, /* indexType= */ null)); provider.put(dictionary); Schema schema = new Schema(Collections.singletonList(vector.getField())); try (VectorSchemaRoot root = new VectorSchemaRoot( schema, Collections.singletonList(vector), vector.getValueCount()); ArrowStreamWriter writer = new ArrowStreamWriter(root, provider, Channels.newChannel(os)); ) { writer.start(); // Batch 1 vector.setNull(0); vector.setSafe(1, 1); vector.setSafe(2, 2); vector.setNull(3); vector.setSafe(4, 1); vector.setValueCount(5); root.setRowCount(5); writer.writeBatch(); // Batch 2 vector.setNull(0); vector.setSafe(1, 1); vector.setSafe(2, 2); vector.setValueCount(3); root.setRowCount(3); writer.writeBatch(); // Batch 3 vector.setSafe(0, 0); vector.setSafe(1, 1); vector.setSafe(2, 2); vector.setSafe(3, 3); vector.setSafe(4, 4); vector.setValueCount(5); root.setRowCount(5); writer.writeBatch(); writer.end(); } } ByteArrayInputStream in = new ByteArrayInputStream(os.toByteArray()); return new ArrowStreamReader(in, allocator); } private void createStructVector(StructVector vector) { final ViewVarCharVector child1 = vector.addOrGet( "f0", FieldType.nullable(MinorType.VIEWVARCHAR.getType()), ViewVarCharVector.class); final IntVector child2 = vector.addOrGet("f1", FieldType.nullable(MinorType.INT.getType()), IntVector.class); // Write the values to child 1 child1.allocateNew(); child1.set(0, "01234567890".getBytes(StandardCharsets.UTF_8)); child1.set(1, "012345678901234567".getBytes(StandardCharsets.UTF_8)); vector.setIndexDefined(0); // Write the values to child 2 child2.allocateNew(); child2.set(0, 10); child2.set(1, 11); vector.setIndexDefined(1); vector.setValueCount(2); } private void createStructVectorInline(StructVector vector) { final ViewVarCharVector child1 = vector.addOrGet( "f0", FieldType.nullable(MinorType.VIEWVARCHAR.getType()), ViewVarCharVector.class); final IntVector child2 = vector.addOrGet("f1", FieldType.nullable(MinorType.INT.getType()), IntVector.class); // Write the values to child 1 child1.allocateNew(); child1.set(0, "012345678".getBytes(StandardCharsets.UTF_8)); child1.set(1, "01234".getBytes(StandardCharsets.UTF_8)); vector.setIndexDefined(0); // Write the values to child 2 child2.allocateNew(); child2.set(0, 10); child2.set(1, 11); vector.setIndexDefined(1); vector.setValueCount(2); } @Test public void testVectorLoadUnloadOnStructVector() { try (final StructVector structVector1 = StructVector.empty("struct", allocator)) { createStructVector(structVector1); Field field1 = structVector1.getField(); Schema schema = new Schema(field1.getChildren()); StructVectorUnloader vectorUnloader = new StructVectorUnloader(structVector1); try (ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch(); BufferAllocator finalVectorsAllocator = allocator.newChildAllocator("struct", 0, Long.MAX_VALUE); ) { // validating recordBatch contains an output for variadicBufferCounts assertFalse(recordBatch.getVariadicBufferCounts().isEmpty()); assertEquals(1, recordBatch.getVariadicBufferCounts().size()); assertEquals(1, recordBatch.getVariadicBufferCounts().get(0)); StructVectorLoader vectorLoader = new StructVectorLoader(schema); try (StructVector structVector2 = vectorLoader.load(finalVectorsAllocator, recordBatch)) { // Improve this after fixing https://github.com/apache/arrow/issues/41933 // assertTrue(VectorEqualsVisitor.vectorEquals(structVector1, structVector2), "vectors are // not equivalent"); assertTrue( VectorEqualsVisitor.vectorEquals( structVector1.getChild("f0"), structVector2.getChild("f0")), "vectors are not equivalent"); assertTrue( VectorEqualsVisitor.vectorEquals( structVector1.getChild("f1"), structVector2.getChild("f1")), "vectors are not equivalent"); } } } } @Test public void testVectorLoadUnloadOnStructVectorWithInline() { try (final StructVector structVector1 = StructVector.empty("struct", allocator)) { createStructVectorInline(structVector1); Field field1 = structVector1.getField(); Schema schema = new Schema(field1.getChildren()); StructVectorUnloader vectorUnloader = new StructVectorUnloader(structVector1); try (ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch(); BufferAllocator finalVectorsAllocator = allocator.newChildAllocator("struct", 0, Long.MAX_VALUE); ) { // validating recordBatch contains an output for variadicBufferCounts assertFalse(recordBatch.getVariadicBufferCounts().isEmpty()); assertEquals(1, recordBatch.getVariadicBufferCounts().size()); assertEquals(0, recordBatch.getVariadicBufferCounts().get(0)); StructVectorLoader vectorLoader = new StructVectorLoader(schema); try (StructVector structVector2 = vectorLoader.load(finalVectorsAllocator, recordBatch)) { // Improve this after fixing https://github.com/apache/arrow/issues/41933 // assertTrue(VectorEqualsVisitor.vectorEquals(structVector1, structVector2), "vectors are // not equivalent"); assertTrue( VectorEqualsVisitor.vectorEquals( structVector1.getChild("f0"), structVector2.getChild("f0")), "vectors are not equivalent"); assertTrue( VectorEqualsVisitor.vectorEquals( structVector1.getChild("f1"), structVector2.getChild("f1")), "vectors are not equivalent"); } } } } } apache-arrow-java-19.0.0/c/src/test/java/org/apache/arrow/c/ExceptionTest.java000066400000000000000000000122541515451615200270400ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; } } } apache-arrow-java-19.0.0/c/src/test/java/org/apache/arrow/c/FlagsTest.java000066400000000000000000000066131515451615200261400ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.ArrayList; 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.junit.jupiter.api.Test; public class FlagsTest { @Test public void testForFieldNullableOrderedDict() { FieldType fieldType = new FieldType( true, ArrowType.Binary.INSTANCE, new DictionaryEncoding(123L, true, new ArrowType.Int(8, true))); assertEquals( Flags.ARROW_FLAG_DICTIONARY_ORDERED | Flags.ARROW_FLAG_NULLABLE, Flags.forField(new Field("Name", fieldType, new ArrayList<>()))); } @Test public void testForFieldOrderedDict() { FieldType fieldType = new FieldType( false, ArrowType.Binary.INSTANCE, new DictionaryEncoding(123L, true, new ArrowType.Int(8, true))); assertEquals( Flags.ARROW_FLAG_DICTIONARY_ORDERED, Flags.forField(new Field("Name", fieldType, new ArrayList<>()))); } @Test public void testForFieldNullableDict() { FieldType fieldType = new FieldType( true, ArrowType.Binary.INSTANCE, new DictionaryEncoding(123L, false, new ArrowType.Int(8, true))); assertEquals( Flags.ARROW_FLAG_NULLABLE, Flags.forField(new Field("Name", fieldType, new ArrayList<>()))); } @Test public void testForFieldNullable() { FieldType fieldType = new FieldType(true, ArrowType.Binary.INSTANCE, null); assertEquals( Flags.ARROW_FLAG_NULLABLE, Flags.forField(new Field("Name", fieldType, new ArrayList<>()))); } @Test public void testForFieldNullableOrderedSortedMap() { ArrowType.Map type = new ArrowType.Map(true); FieldType fieldType = new FieldType(true, type, new DictionaryEncoding(123L, true, new ArrowType.Int(8, true))); assertEquals( Flags.ARROW_FLAG_DICTIONARY_ORDERED | Flags.ARROW_FLAG_NULLABLE | Flags.ARROW_FLAG_MAP_KEYS_SORTED, Flags.forField(new Field("Name", fieldType, new ArrayList<>()))); } @Test public void testForFieldNullableOrderedMap() { ArrowType.Map type = new ArrowType.Map(false); FieldType fieldType = new FieldType(true, type, new DictionaryEncoding(123L, true, new ArrowType.Int(8, true))); assertEquals( Flags.ARROW_FLAG_DICTIONARY_ORDERED | Flags.ARROW_FLAG_NULLABLE, Flags.forField(new Field("Name", fieldType, new ArrayList<>()))); } } apache-arrow-java-19.0.0/c/src/test/java/org/apache/arrow/c/FormatTest.java000066400000000000000000000232051515451615200263300ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.IntervalUnit; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.UnionMode; import org.apache.arrow.vector.types.pojo.ArrowType; import org.junit.jupiter.api.Test; public class FormatTest { @Test public void testAsString() { assertEquals("z", Format.asString(new ArrowType.Binary())); assertEquals("b", Format.asString(new ArrowType.Bool())); assertEquals("tdD", Format.asString(new ArrowType.Date(DateUnit.DAY))); assertEquals("tdm", Format.asString(new ArrowType.Date(DateUnit.MILLISECOND))); assertEquals("d:1,1", Format.asString(new ArrowType.Decimal(1, 1, 128))); assertEquals("d:1,1,1", Format.asString(new ArrowType.Decimal(1, 1, 1))); assertEquals("d:9,1,1", Format.asString(new ArrowType.Decimal(9, 1, 1))); assertEquals("tDs", Format.asString(new ArrowType.Duration(TimeUnit.SECOND))); assertEquals("tDm", Format.asString(new ArrowType.Duration(TimeUnit.MILLISECOND))); assertEquals("tDu", Format.asString(new ArrowType.Duration(TimeUnit.MICROSECOND))); assertEquals("tDn", Format.asString(new ArrowType.Duration(TimeUnit.NANOSECOND))); assertEquals("w:1", Format.asString(new ArrowType.FixedSizeBinary(1))); assertEquals("+w:3", Format.asString(new ArrowType.FixedSizeList(3))); assertEquals("e", Format.asString(new ArrowType.FloatingPoint(FloatingPointPrecision.HALF))); assertEquals("f", Format.asString(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE))); assertEquals("g", Format.asString(new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE))); assertEquals("c", Format.asString(new ArrowType.Int(Byte.SIZE, true))); assertEquals("C", Format.asString(new ArrowType.Int(Byte.SIZE, false))); assertEquals("s", Format.asString(new ArrowType.Int(Short.SIZE, true))); assertEquals("S", Format.asString(new ArrowType.Int(Short.SIZE, false))); assertEquals("i", Format.asString(new ArrowType.Int(Integer.SIZE, true))); assertEquals("I", Format.asString(new ArrowType.Int(Integer.SIZE, false))); assertEquals("l", Format.asString(new ArrowType.Int(Long.SIZE, true))); assertEquals("L", Format.asString(new ArrowType.Int(Long.SIZE, false))); assertEquals("tiD", Format.asString(new ArrowType.Interval(IntervalUnit.DAY_TIME))); assertEquals("tiM", Format.asString(new ArrowType.Interval(IntervalUnit.YEAR_MONTH))); assertEquals("Z", Format.asString(new ArrowType.LargeBinary())); assertEquals("+L", Format.asString(new ArrowType.LargeList())); assertEquals("U", Format.asString(new ArrowType.LargeUtf8())); assertEquals("+l", Format.asString(new ArrowType.List())); assertEquals("+m", Format.asString(new ArrowType.Map(true))); assertEquals("n", Format.asString(new ArrowType.Null())); assertEquals("+s", Format.asString(new ArrowType.Struct())); assertEquals("tts", Format.asString(new ArrowType.Time(TimeUnit.SECOND, 32))); assertEquals("ttm", Format.asString(new ArrowType.Time(TimeUnit.MILLISECOND, 32))); assertEquals("ttu", Format.asString(new ArrowType.Time(TimeUnit.MICROSECOND, 64))); assertEquals("ttn", Format.asString(new ArrowType.Time(TimeUnit.NANOSECOND, 64))); assertEquals( "tss:Timezone", Format.asString(new ArrowType.Timestamp(TimeUnit.SECOND, "Timezone"))); assertEquals( "tsm:Timezone", Format.asString(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "Timezone"))); assertEquals( "tsu:Timezone", Format.asString(new ArrowType.Timestamp(TimeUnit.MICROSECOND, "Timezone"))); assertEquals( "tsn:Timezone", Format.asString(new ArrowType.Timestamp(TimeUnit.NANOSECOND, "Timezone"))); assertEquals( "+us:1,1,1", Format.asString(new ArrowType.Union(UnionMode.Sparse, new int[] {1, 1, 1}))); assertEquals( "+ud:1,1,1", Format.asString(new ArrowType.Union(UnionMode.Dense, new int[] {1, 1, 1}))); assertEquals("u", Format.asString(new ArrowType.Utf8())); assertThrows( UnsupportedOperationException.class, () -> Format.asString(new ArrowType.Int(1, true))); assertThrows( UnsupportedOperationException.class, () -> Format.asString(new ArrowType.Time(TimeUnit.SECOND, 1))); assertThrows( UnsupportedOperationException.class, () -> Format.asString(new ArrowType.Time(TimeUnit.MILLISECOND, 64))); } @Test public void testAsType() throws IllegalStateException, NumberFormatException, UnsupportedOperationException { assertTrue(Format.asType("n", 0L) instanceof ArrowType.Null); assertTrue(Format.asType("b", 0L) instanceof ArrowType.Bool); assertEquals(new ArrowType.Int(Byte.SIZE, true), Format.asType("c", 0L)); assertEquals(new ArrowType.Int(Byte.SIZE, false), Format.asType("C", 0L)); assertEquals(new ArrowType.Int(Short.SIZE, true), Format.asType("s", 0L)); assertEquals(new ArrowType.Int(Short.SIZE, false), Format.asType("S", 0L)); assertEquals(new ArrowType.Int(Integer.SIZE, true), Format.asType("i", 0L)); assertEquals(new ArrowType.Int(Integer.SIZE, false), Format.asType("I", 0L)); assertEquals(new ArrowType.Int(Long.SIZE, true), Format.asType("l", 0L)); assertEquals(new ArrowType.Int(Long.SIZE, false), Format.asType("L", 0L)); assertEquals(new ArrowType.FloatingPoint(FloatingPointPrecision.HALF), Format.asType("e", 0L)); assertEquals( new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE), Format.asType("f", 0L)); assertEquals( new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE), Format.asType("g", 0L)); assertTrue(Format.asType("z", 0L) instanceof ArrowType.Binary); assertTrue(Format.asType("Z", 0L) instanceof ArrowType.LargeBinary); assertTrue(Format.asType("u", 0L) instanceof ArrowType.Utf8); assertTrue(Format.asType("U", 0L) instanceof ArrowType.LargeUtf8); assertEquals(new ArrowType.Date(DateUnit.DAY), Format.asType("tdD", 0L)); assertEquals(new ArrowType.Date(DateUnit.MILLISECOND), Format.asType("tdm", 0L)); assertEquals(new ArrowType.Time(TimeUnit.SECOND, Integer.SIZE), Format.asType("tts", 0L)); assertEquals(new ArrowType.Time(TimeUnit.MILLISECOND, Integer.SIZE), Format.asType("ttm", 0L)); assertEquals(new ArrowType.Time(TimeUnit.MICROSECOND, Long.SIZE), Format.asType("ttu", 0L)); assertEquals(new ArrowType.Time(TimeUnit.NANOSECOND, Long.SIZE), Format.asType("ttn", 0L)); assertEquals(new ArrowType.Duration(TimeUnit.SECOND), Format.asType("tDs", 0L)); assertEquals(new ArrowType.Duration(TimeUnit.MILLISECOND), Format.asType("tDm", 0L)); assertEquals(new ArrowType.Duration(TimeUnit.MICROSECOND), Format.asType("tDu", 0L)); assertEquals(new ArrowType.Duration(TimeUnit.NANOSECOND), Format.asType("tDn", 0L)); assertEquals(new ArrowType.Interval(IntervalUnit.YEAR_MONTH), Format.asType("tiM", 0L)); assertEquals(new ArrowType.Interval(IntervalUnit.DAY_TIME), Format.asType("tiD", 0L)); assertTrue(Format.asType("+l", 0L) instanceof ArrowType.List); assertTrue(Format.asType("+L", 0L) instanceof ArrowType.LargeList); assertTrue(Format.asType("+s", 0L) instanceof ArrowType.Struct); assertEquals(new ArrowType.Map(false), Format.asType("+m", 0L)); assertEquals(new ArrowType.Map(true), Format.asType("+m", Flags.ARROW_FLAG_MAP_KEYS_SORTED)); assertEquals(new ArrowType.Decimal(1, 1, 128), Format.asType("d:1,1", 0L)); assertEquals(new ArrowType.Decimal(1, 1, 1), Format.asType("d:1,1,1", 0L)); assertEquals(new ArrowType.Decimal(9, 1, 1), Format.asType("d:9,1,1", 0L)); assertEquals(new ArrowType.FixedSizeBinary(1), Format.asType("w:1", 0L)); assertEquals(new ArrowType.FixedSizeList(3), Format.asType("+w:3", 0L)); assertEquals( new ArrowType.Union(UnionMode.Dense, new int[] {1, 1, 1}), Format.asType("+ud:1,1,1", 0L)); assertEquals( new ArrowType.Union(UnionMode.Sparse, new int[] {1, 1, 1}), Format.asType("+us:1,1,1", 0L)); assertEquals( new ArrowType.Timestamp(TimeUnit.SECOND, "Timezone"), Format.asType("tss:Timezone", 0L)); assertEquals( new ArrowType.Timestamp(TimeUnit.MILLISECOND, "Timezone"), Format.asType("tsm:Timezone", 0L)); assertEquals( new ArrowType.Timestamp(TimeUnit.MICROSECOND, "Timezone"), Format.asType("tsu:Timezone", 0L)); assertEquals( new ArrowType.Timestamp(TimeUnit.NANOSECOND, "Timezone"), Format.asType("tsn:Timezone", 0L)); assertThrows(UnsupportedOperationException.class, () -> Format.asType("Format", 0L)); assertThrows(UnsupportedOperationException.class, () -> Format.asType(":", 0L)); assertThrows(NumberFormatException.class, () -> Format.asType("w:1,2,3", 0L)); } } apache-arrow-java-19.0.0/c/src/test/java/org/apache/arrow/c/MetadataTest.java000066400000000000000000000063441515451615200266250ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.*; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.util.HashMap; import java.util.Map; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.memory.util.LargeMemoryUtil; import org.apache.arrow.memory.util.MemoryUtil; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class MetadataTest { private RootAllocator allocator = null; private static Map metadata; private static byte[] encoded; @BeforeAll static void beforeAll() { metadata = new HashMap<>(); metadata.put("key1", ""); metadata.put("key2", "bar"); if (ByteOrder.nativeOrder() == ByteOrder.LITTLE_ENDIAN) { encoded = new byte[] { 2, 0, 0, 0, 4, 0, 0, 0, 'k', 'e', 'y', '1', 0, 0, 0, 0, 4, 0, 0, 0, 'k', 'e', 'y', '2', 3, 0, 0, 0, 'b', 'a', 'r' }; } else { encoded = new byte[] { 0, 0, 0, 2, 0, 0, 0, 4, 'k', 'e', 'y', '1', 0, 0, 0, 0, 0, 0, 0, 4, 'k', 'e', 'y', '2', 0, 0, 0, 3, 'b', 'a', 'r' }; } } @BeforeEach public void setUp() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void tearDown() { allocator.close(); } @Test public void testEncode() { try (ArrowBuf buffer = Metadata.encode(allocator, metadata)) { int totalSize = LargeMemoryUtil.checkedCastToInt(buffer.readableBytes()); ByteBuffer reader = MemoryUtil.directBuffer(buffer.memoryAddress(), totalSize).order(ByteOrder.nativeOrder()); byte[] result = new byte[totalSize]; reader.get(result); assertArrayEquals(encoded, result); } } @Test public void testDecode() { try (ArrowBuf buffer = allocator.buffer(31)) { buffer.setBytes(0, encoded); Map decoded = Metadata.decode(buffer.memoryAddress()); assertNotNull(decoded); assertEquals(metadata, decoded); } } @Test public void testEncodeEmpty() { Map metadata = new HashMap<>(); try (ArrowBuf encoded = Metadata.encode(allocator, metadata)) { assertNull(encoded); } } @Test public void testDecodeEmpty() { Map decoded = Metadata.decode(NativeUtil.NULL); assertNull(decoded); } } apache-arrow-java-19.0.0/c/src/test/java/org/apache/arrow/c/NativeUtilTest.java000066400000000000000000000052511515451615200271650ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import java.nio.ByteBuffer; import java.nio.ByteOrder; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.memory.util.LargeMemoryUtil; import org.apache.arrow.memory.util.MemoryUtil; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class NativeUtilTest { private RootAllocator allocator = null; @BeforeEach public void setUp() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void tearDown() { allocator.close(); } @Test public void testString() { String javaString = "abc"; byte[] nativeString = new byte[] {97, 98, 99, 0}; try (ArrowBuf buffer = NativeUtil.toNativeString(allocator, javaString)) { int totalSize = LargeMemoryUtil.checkedCastToInt(buffer.readableBytes()); ByteBuffer reader = MemoryUtil.directBuffer(buffer.memoryAddress(), totalSize).order(ByteOrder.nativeOrder()); byte[] result = new byte[totalSize]; reader.get(result); assertArrayEquals(nativeString, result); assertEquals(javaString, NativeUtil.toJavaString(buffer.memoryAddress())); } } @Test public void testToJavaArray() { long[] nativeArray = new long[] {1, 2, 3}; try (ArrowBuf buffer = allocator.buffer(Long.BYTES * ((long) nativeArray.length), null)) { for (long value : nativeArray) { buffer.writeLong(value); } long[] actual = NativeUtil.toJavaArray(buffer.memoryAddress(), nativeArray.length); assertArrayEquals(nativeArray, actual); } } @Test public void testToZeroJavaArray() { long[] actual = NativeUtil.toJavaArray(0xDEADBEEF, 0); assertEquals(0, actual.length); } } apache-arrow-java-19.0.0/c/src/test/java/org/apache/arrow/c/RoundtripTest.java000066400000000000000000001175151515451615200270760ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector; import static org.junit.jupiter.api.Assertions.*; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.UUID; import java.util.stream.Collectors; 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.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.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.IntervalDayVector; import org.apache.arrow.vector.IntervalYearVector; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.LargeVarCharVector; 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.UuidVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ViewVarBinaryVector; import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.ZeroVector; import org.apache.arrow.vector.compare.VectorEqualsVisitor; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.LargeListVector; 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.MapVector; 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.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; import org.apache.arrow.vector.table.Table; 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.ExtensionType; 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.TransferPair; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class RoundtripTest { private static final String EMPTY_SCHEMA_PATH = ""; private RootAllocator allocator = null; private BufferAllocator childAllocator = null; @BeforeEach public void setUp() { allocator = new RootAllocator(Long.MAX_VALUE); childAllocator = allocator.newChildAllocator("child", 0, Long.MAX_VALUE); } @AfterEach public void tearDown() { childAllocator.close(); allocator.close(); } FieldVector vectorRoundtrip(FieldVector vector) { // Consumer allocates empty structures 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 Data.exportVector(allocator, vector, null, arrowArray, arrowSchema); } // Consumer imports vector FieldVector imported = Data.importVector(childAllocator, consumerArrowArray, consumerArrowSchema, null); if (!(imported instanceof NullVector)) { assertEquals(childAllocator, imported.getAllocator()); } // Check that transfers work TransferPair pair = imported.getTransferPair(allocator); pair.transfer(); return (FieldVector) pair.getTo(); } } VectorSchemaRoot vectorSchemaRootRoundtrip(VectorSchemaRoot root) { // Consumer allocates empty structures 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 Data.exportVectorSchemaRoot(allocator, root, null, arrowArray, arrowSchema); } // Consumer imports vector return Data.importVectorSchemaRoot(allocator, consumerArrowArray, consumerArrowSchema, null); } } boolean roundtrip(FieldVector vector, Class clazz) { List fieldBuffers = vector.getFieldBuffers(); List orgRefCnts = fieldBuffers.stream().map(buf -> buf.refCnt()).collect(Collectors.toList()); long orgMemorySize = allocator.getAllocatedMemory(); boolean result = false; try (ValueVector imported = vectorRoundtrip(vector)) { assertTrue( 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 IntStream.range(0, orgRefCnts.size()) .forEach( i -> { ArrowBuf buf = fieldBuffers.get(i); assertEquals(buf.refCnt(), orgRefCnts.get(i)); }); assertEquals(orgMemorySize, allocator.getAllocatedMemory()); return result; } @Test public void testBitVector() { BitVector imported; try (final BitVector vector = new BitVector(EMPTY_SCHEMA_PATH, allocator)) { vector.allocateNew(1024); vector.setValueCount(1024); // Put and set a few values vector.set(0, 1); vector.set(1, 0); vector.set(100, 0); vector.set(1022, 1); vector.setValueCount(1024); imported = (BitVector) vectorRoundtrip(vector); assertTrue(VectorEqualsVisitor.vectorEquals(vector, imported)); } assertEquals(1, imported.get(0)); assertEquals(0, imported.get(1)); assertEquals(0, imported.get(100)); assertEquals(1, imported.get(1022)); assertEquals(1020, imported.getNullCount()); imported.close(); } @Test public void testIntVector() { IntVector imported; try (final IntVector vector = new IntVector("v", allocator)) { setVector(vector, 1, 2, 3, null); imported = (IntVector) vectorRoundtrip(vector); assertTrue(VectorEqualsVisitor.vectorEquals(vector, imported)); } assertEquals(1, imported.get(0)); assertEquals(2, imported.get(1)); assertEquals(3, imported.get(2)); assertEquals(4, imported.getValueCount()); assertEquals(1, imported.getNullCount()); imported.close(); } @Test public void testBigIntVector() { BigIntVector imported; try (final BigIntVector vector = new BigIntVector("v", allocator)) { setVector(vector, 1L, 2L, 3L, null); imported = (BigIntVector) vectorRoundtrip(vector); assertTrue(VectorEqualsVisitor.vectorEquals(vector, imported)); } assertEquals(1, imported.get(0)); assertEquals(2, imported.get(1)); assertEquals(3, imported.get(2)); assertEquals(4, imported.getValueCount()); assertEquals(1, imported.getNullCount()); imported.close(); } @Test public void testDateDayVector() { DateDayVector imported; try (final DateDayVector vector = new DateDayVector("v", allocator)) { setVector(vector, 1, 2, 3, null); imported = (DateDayVector) vectorRoundtrip(vector); assertTrue(VectorEqualsVisitor.vectorEquals(vector, imported)); } assertEquals(1, imported.get(0)); assertEquals(2, imported.get(1)); assertEquals(3, imported.get(2)); assertEquals(4, imported.getValueCount()); assertEquals(1, imported.getNullCount()); imported.close(); } @Test public void testDateMilliVector() { DateMilliVector imported; try (final DateMilliVector vector = new DateMilliVector("v", allocator)) { setVector(vector, 1L, 2L, 3L, null); imported = (DateMilliVector) vectorRoundtrip(vector); assertTrue(VectorEqualsVisitor.vectorEquals(vector, imported)); } assertEquals(1, imported.get(0)); assertEquals(2, imported.get(1)); assertEquals(3, imported.get(2)); assertEquals(4, imported.getValueCount()); assertEquals(1, imported.getNullCount()); imported.close(); } @Test public void testDecimalVector() { try (final DecimalVector vector = new DecimalVector("v", allocator, 1, 1)) { setVector(vector, 1L, 2L, 3L, null); assertTrue(roundtrip(vector, DecimalVector.class)); } } @Test public void testDurationVector() { for (TimeUnit unit : TimeUnit.values()) { final FieldType fieldType = FieldType.nullable(new ArrowType.Duration(unit)); try (final DurationVector vector = new DurationVector("v", fieldType, allocator)) { setVector(vector, 1L, 2L, 3L, null); assertTrue(roundtrip(vector, DurationVector.class)); } } } @Test public void testZeroVectorEquals() { try (final ZeroVector vector = new ZeroVector()) { // A ZeroVector is imported as a NullVector assertTrue(roundtrip(vector, NullVector.class)); } } @Test public void testFixedSizeBinaryVector() { try (final FixedSizeBinaryVector vector = new FixedSizeBinaryVector("v", allocator, 2)) { setVector(vector, new byte[] {0b0000, 0b0001}, new byte[] {0b0010, 0b0011}); assertTrue(roundtrip(vector, FixedSizeBinaryVector.class)); } } @Test public void testFloat2Vector() { try (final Float2Vector vector = new Float2Vector("v", allocator)) { setVector(vector, 0.1f, 0.2f, 0.3f, null); assertTrue(roundtrip(vector, Float2Vector.class)); } } @Test public void testFloat4Vector() { try (final Float4Vector vector = new Float4Vector("v", allocator)) { setVector(vector, 0.1f, 0.2f, 0.3f, null); assertTrue(roundtrip(vector, Float4Vector.class)); } } @Test public void testFloat8Vector() { try (final Float8Vector vector = new Float8Vector("v", allocator)) { setVector(vector, 0.1d, 0.2d, 0.3d, null); assertTrue(roundtrip(vector, Float8Vector.class)); } } @Test public void testIntervalDayVector() { try (final IntervalDayVector vector = new IntervalDayVector("v", allocator)) { IntervalDayHolder value = new IntervalDayHolder(); value.days = 5; value.milliseconds = 100; setVector(vector, value, null); assertTrue(roundtrip(vector, IntervalDayVector.class)); } } @Test public void testIntervalYearVector() { try (final IntervalYearVector vector = new IntervalYearVector("v", allocator)) { setVector(vector, 1990, 2000, 2010, 2020, null); assertTrue(roundtrip(vector, IntervalYearVector.class)); } } @Test public void testSmallIntVector() { try (final SmallIntVector vector = new SmallIntVector("v", allocator)) { setVector(vector, (short) 0, (short) 256, null); assertTrue(roundtrip(vector, SmallIntVector.class)); } } @Test public void testTimeMicroVector() { try (final TimeMicroVector vector = new TimeMicroVector("v", allocator)) { setVector(vector, 0L, 1L, 2L, 3L, null); assertTrue(roundtrip(vector, TimeMicroVector.class)); } } @Test public void testTimeMilliVector() { try (final TimeMilliVector vector = new TimeMilliVector("v", allocator)) { setVector(vector, 0, 1, 2, 3, null); assertTrue(roundtrip(vector, TimeMilliVector.class)); } } @Test public void testTimeNanoVector() { try (final TimeNanoVector vector = new TimeNanoVector("v", allocator)) { setVector(vector, 0L, 1L, 2L, 3L, null); assertTrue(roundtrip(vector, TimeNanoVector.class)); } } @Test public void testTimeSecVector() { try (final TimeSecVector vector = new TimeSecVector("v", allocator)) { setVector(vector, 0, 1, 2, 3, null); assertTrue(roundtrip(vector, TimeSecVector.class)); } } @Test public void testTimeStampMicroTZVector() { try (final TimeStampMicroTZVector vector = new TimeStampMicroTZVector("v", allocator, "UTC")) { setVector(vector, 0L, 1L, 2L, 3L, null); assertTrue(roundtrip(vector, TimeStampMicroTZVector.class)); } } @Test public void testTimeStampMicroVector() { try (final TimeStampMicroVector vector = new TimeStampMicroVector("v", allocator)) { setVector(vector, 0L, 1L, 2L, 3L, null); assertTrue(roundtrip(vector, TimeStampMicroVector.class)); } } @Test public void testTimeStampMilliTZVector() { try (final TimeStampMilliTZVector vector = new TimeStampMilliTZVector("v", allocator, "UTC")) { setVector(vector, 0L, 1L, 2L, 3L, null); assertTrue(roundtrip(vector, TimeStampMilliTZVector.class)); } } @Test public void testTimeStampMilliVector() { try (final TimeStampMilliVector vector = new TimeStampMilliVector("v", allocator)) { setVector(vector, 0L, 1L, 2L, 3L, null); assertTrue(roundtrip(vector, TimeStampMilliVector.class)); } } @Test public void testTimeTimeStampNanoTZVector() { try (final TimeStampNanoTZVector vector = new TimeStampNanoTZVector("v", allocator, "UTC")) { setVector(vector, 0L, 1L, 2L, 3L, null); assertTrue(roundtrip(vector, TimeStampNanoTZVector.class)); } } @Test public void testTimeStampNanoVector() { try (final TimeStampNanoVector vector = new TimeStampNanoVector("v", allocator)) { setVector(vector, 0L, 1L, 2L, 3L, null); assertTrue(roundtrip(vector, TimeStampNanoVector.class)); } } @Test public void testTimeStampSecTZVector() { try (final TimeStampSecTZVector vector = new TimeStampSecTZVector("v", allocator, "UTC")) { setVector(vector, 0L, 1L, 2L, 3L, null); assertTrue(roundtrip(vector, TimeStampSecTZVector.class)); } } @Test public void testTimeStampSecVector() { try (final TimeStampSecVector vector = new TimeStampSecVector("v", allocator)) { setVector(vector, 0L, 1L, 2L, 3L, null); assertTrue(roundtrip(vector, TimeStampSecVector.class)); } } @Test public void testTinyIntVector() { try (final TinyIntVector vector = new TinyIntVector("v", allocator)) { setVector(vector, (byte) 0, (byte) 1, null); assertTrue(roundtrip(vector, TinyIntVector.class)); } } @Test public void testUInt1Vector() { try (final UInt1Vector vector = new UInt1Vector("v", allocator)) { setVector(vector, (byte) 0, (byte) 1, null); assertTrue(roundtrip(vector, UInt1Vector.class)); } } @Test public void testUInt2Vector() { try (final UInt2Vector vector = new UInt2Vector("v", allocator)) { setVector(vector, '0', '1', null); assertTrue(roundtrip(vector, UInt2Vector.class)); } } @Test public void testUInt4Vector() { try (final UInt4Vector vector = new UInt4Vector("v", allocator)) { setVector(vector, 0, 1, null); assertTrue(roundtrip(vector, UInt4Vector.class)); } } @Test public void testUInt8Vector() { try (final UInt8Vector vector = new UInt8Vector("v", allocator)) { setVector(vector, 0L, 1L, null); assertTrue(roundtrip(vector, UInt8Vector.class)); } } @Test public void testNullVector() { try (final NullVector vector = new NullVector("v", 1024)) { assertTrue(roundtrip(vector, NullVector.class)); } } @Test public void testVarBinaryVector() { try (final VarBinaryVector vector = new VarBinaryVector("v", allocator)) { setVector( vector, "abc".getBytes(StandardCharsets.UTF_8), "def".getBytes(StandardCharsets.UTF_8), null); assertTrue(roundtrip(vector, VarBinaryVector.class)); } } @Test public void testViewVector() { // ViewVarCharVector with short strings try (final ViewVarCharVector vector = new ViewVarCharVector("v1", allocator)) { setVector( vector, "abc".getBytes(StandardCharsets.UTF_8), "def".getBytes(StandardCharsets.UTF_8), null); assertTrue(roundtrip(vector, ViewVarCharVector.class)); } // ViewVarCharVector with long strings try (final ViewVarCharVector vector = new ViewVarCharVector("v2", allocator)) { setVector( vector, "01234567890123".getBytes(StandardCharsets.UTF_8), "01234567890123567".getBytes(StandardCharsets.UTF_8), null); assertTrue(roundtrip(vector, ViewVarCharVector.class)); } // ViewVarBinaryVector with short values try (final ViewVarBinaryVector vector = new ViewVarBinaryVector("v3", allocator)) { setVector( vector, "abc".getBytes(StandardCharsets.UTF_8), "def".getBytes(StandardCharsets.UTF_8), null); assertTrue(roundtrip(vector, ViewVarBinaryVector.class)); } // ViewVarBinaryVector with long values try (final ViewVarBinaryVector vector = new ViewVarBinaryVector("v4", allocator)) { setVector( vector, "01234567890123".getBytes(StandardCharsets.UTF_8), "01234567890123567".getBytes(StandardCharsets.UTF_8), null); assertTrue(roundtrip(vector, ViewVarBinaryVector.class)); } List byteArrayList = new ArrayList<>(); for (int i = 1; i <= 500; i++) { StringBuilder sb = new StringBuilder(i); for (int j = 0; j < i; j++) { sb.append(j); // or any other character } byte[] bytes = sb.toString().getBytes(StandardCharsets.UTF_8); byteArrayList.add(bytes); } // ViewVarCharVector with short long strings with multiple data buffers try (final ViewVarCharVector vector = new ViewVarCharVector("v5", allocator)) { setVector(vector, byteArrayList.toArray(new byte[0][])); assertTrue(roundtrip(vector, ViewVarCharVector.class)); } // ViewVarBinaryVector with short long strings with multiple data buffers try (final ViewVarBinaryVector vector = new ViewVarBinaryVector("v6", allocator)) { setVector(vector, byteArrayList.toArray(new byte[0][])); assertTrue(roundtrip(vector, ViewVarBinaryVector.class)); } } @Test public void testVarCharVector() { try (final VarCharVector vector = new VarCharVector("v", allocator)) { setVector(vector, "abc", "def", null); assertTrue(roundtrip(vector, VarCharVector.class)); } } @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)) { vector.allocateNew(5, 1); NullableLargeVarBinaryHolder nullHolder = new NullableLargeVarBinaryHolder(); nullHolder.isSet = 0; NullableLargeVarBinaryHolder binHolder = new NullableLargeVarBinaryHolder(); binHolder.isSet = 1; String str = "hello world"; try (ArrowBuf buf = allocator.buffer(16)) { buf.setBytes(0, str.getBytes(StandardCharsets.UTF_8)); binHolder.start = 0; binHolder.end = str.length(); binHolder.buffer = buf; vector.setSafe(0, binHolder); vector.setSafe(1, nullHolder); assertTrue(roundtrip(vector, LargeVarBinaryVector.class)); } } } @Test public void testLargeVarCharVector() { try (final LargeVarCharVector vector = new LargeVarCharVector("v", allocator)) { setVector(vector, "abc", "def", null); assertTrue(roundtrip(vector, LargeVarCharVector.class)); } } @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)) { setVector( vector, Arrays.stream(new int[] {1, 2}).boxed().collect(Collectors.toList()), Arrays.stream(new int[] {3, 4}).boxed().collect(Collectors.toList()), new ArrayList()); assertTrue(roundtrip(vector, ListVector.class)); } } @Test public void testEmptyListVector() { try (final ListVector vector = ListVector.empty("v", allocator)) { setVector(vector, new ArrayList()); assertTrue(roundtrip(vector, ListVector.class)); } } @Test public void testLargeListVector() { try (final LargeListVector vector = LargeListVector.empty("v", allocator)) { setVector( vector, Arrays.stream(new int[] {1, 2}).boxed().collect(Collectors.toList()), Arrays.stream(new int[] {3, 4}).boxed().collect(Collectors.toList()), new ArrayList()); assertTrue(roundtrip(vector, LargeListVector.class)); } } @Test public void testFixedSizeListVector() { try (final FixedSizeListVector vector = FixedSizeListVector.empty("v", 2, allocator)) { setVector( vector, Arrays.stream(new int[] {1, 2}).boxed().collect(Collectors.toList()), Arrays.stream(new int[] {3, 4}).boxed().collect(Collectors.toList())); assertTrue(roundtrip(vector, FixedSizeListVector.class)); } } @Test public void testListViewVector() { try (final ListViewVector vector = ListViewVector.empty("v", allocator)) { setVector( vector, Arrays.stream(new int[] {1, 2}).boxed().collect(Collectors.toList()), Arrays.stream(new int[] {3, 4}).boxed().collect(Collectors.toList()), new ArrayList()); assertTrue(roundtrip(vector, ListViewVector.class)); } } @Test public void testEmptyListViewVector() { try (final ListViewVector vector = ListViewVector.empty("v", allocator)) { setVector(vector, new ArrayList()); assertTrue(roundtrip(vector, ListViewVector.class)); } } @Test public void testLargeListViewVector() { try (final LargeListViewVector vector = LargeListViewVector.empty("v", allocator)) { setVector( vector, Arrays.stream(new int[] {1, 2}).boxed().collect(Collectors.toList()), Arrays.stream(new int[] {3, 4}).boxed().collect(Collectors.toList()), new ArrayList()); assertTrue(roundtrip(vector, LargeListViewVector.class)); } } @Test public void testEmptyLargeListViewVector() { try (final LargeListViewVector vector = LargeListViewVector.empty("v", allocator)) { setVector(vector, new ArrayList()); assertTrue(roundtrip(vector, LargeListViewVector.class)); } } @Test public void testMapVector() { int count = 5; try (final MapVector vector = MapVector.empty("v", allocator, false)) { vector.allocateNew(); UnionMapWriter mapWriter = vector.getWriter(); for (int i = 0; i < count; i++) { mapWriter.startMap(); for (int j = 0; j < i + 1; j++) { mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(j); mapWriter.value().integer().writeInt(j); mapWriter.endEntry(); } mapWriter.endMap(); } mapWriter.setValueCount(count); assertTrue(roundtrip(vector, MapVector.class)); } } @Test public void testUnionVector() { final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); uInt4Holder.value = 100; uInt4Holder.isSet = 1; try (UnionVector vector = UnionVector.empty("v", allocator)) { vector.allocateNew(); // write some data vector.setType(0, MinorType.UINT4); vector.setSafe(0, uInt4Holder); vector.setType(2, MinorType.UINT4); vector.setSafe(2, uInt4Holder); vector.setValueCount(4); assertTrue(roundtrip(vector, UnionVector.class)); } } @Test public void testStructVector() { try (final StructVector vector = StructVector.empty("v", allocator)) { Map> data = new HashMap<>(); data.put("col_1", Arrays.stream(new int[] {1, 2}).boxed().collect(Collectors.toList())); data.put("col_2", Arrays.stream(new int[] {3, 4}).boxed().collect(Collectors.toList())); setVector(vector, data); assertTrue(roundtrip(vector, StructVector.class)); } } @Test public void testRunEndEncodedVector() { try (final RunEndEncodedVector vector = RunEndEncodedVector.empty("v", allocator)) { setVector(vector, List.of(1, 3), List.of(1, 2)); assertTrue(roundtrip(vector, RunEndEncodedVector.class)); } } @Test public void testEmptyRunEndEncodedVector() { try (final RunEndEncodedVector vector = RunEndEncodedVector.empty("v", allocator)) { setVector(vector, List.of(), List.of()); assertTrue(roundtrip(vector, RunEndEncodedVector.class)); } } @Test public void testExtensionTypeVector() { final Schema schema = new Schema(Collections.singletonList(Field.nullable("a", UuidType.INSTANCE))); try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { // Fill with data UUID u1 = UUID.randomUUID(); UUID u2 = UUID.randomUUID(); UuidVector vector = (UuidVector) root.getVector("a"); vector.setValueCount(2); vector.set(0, u1); vector.set(1, u2); root.setRowCount(2); // Roundtrip (export + import) VectorSchemaRoot importedRoot = vectorSchemaRootRoundtrip(root); // Verify correctness assertEquals(root.getSchema(), importedRoot.getSchema()); final Field field = importedRoot.getSchema().getFields().get(0); assertEquals( field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME), UuidType.INSTANCE.extensionName()); assertEquals( field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_METADATA), UuidType.INSTANCE.serialize()); final UuidVector deserialized = (UuidVector) importedRoot.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)); } } importedRoot.close(); } } @Test public void testVectorSchemaRoot() { VectorSchemaRoot imported; // Consumer allocates empty structures try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator); ArrowArray consumerArrowArray = ArrowArray.allocateNew(allocator)) { try (VectorSchemaRoot vsr = createTestVSR()) { // 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 Data.exportVectorSchemaRoot(allocator, vsr, null, arrowArray, arrowSchema); } } // Consumer imports vector imported = Data.importVectorSchemaRoot(allocator, consumerArrowArray, consumerArrowSchema, null); } // Ensure that imported VectorSchemaRoot is valid even after C Data Interface // structures are closed try (VectorSchemaRoot original = createTestVSR()) { assertTrue(imported.equals(original)); } imported.close(); } /** * Tests exporting Table and importing back to VSR. Importing back to Table is not supported at * present. */ @Test public void testTable() { VectorSchemaRoot imported; // Consumer allocates empty structures try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator); ArrowArray consumerArrowArray = ArrowArray.allocateNew(allocator)) { try (VectorSchemaRoot vsr = createTestVSR(); Table table = new Table(vsr)) { // 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 Data.exportTable(allocator, table, null, arrowArray, arrowSchema); } } // Consumer imports vector imported = Data.importVectorSchemaRoot(allocator, consumerArrowArray, consumerArrowSchema, null); } // Ensure that imported VectorSchemaRoot is valid even after C Data Interface // structures are closed try (VectorSchemaRoot original = createTestVSR()) { assertTrue(imported.equals(original)); } imported.close(); } @Test public void testVectorSchemaRootWithDuplicatedFieldNames() { VectorSchemaRoot imported; // Consumer allocates empty structures try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator); ArrowArray consumerArrowArray = ArrowArray.allocateNew(allocator)) { try (VectorSchemaRoot testVSR1 = createTestVSR(); VectorSchemaRoot testVSR2 = createTestVSR()) { // Merge two VSRs to produce duplicated field names final VectorSchemaRoot vsr = new VectorSchemaRoot( Stream.concat( testVSR1.getFieldVectors().stream(), testVSR2.getFieldVectors().stream()) .collect(Collectors.toList())); // 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 Data.exportVectorSchemaRoot(allocator, vsr, null, arrowArray, arrowSchema); } } // Consumer imports vector imported = Data.importVectorSchemaRoot(allocator, consumerArrowArray, consumerArrowSchema, null); } // Ensure that imported VectorSchemaRoot is valid even after C Data Interface // structures are closed try (VectorSchemaRoot testVSR1 = createTestVSR(); VectorSchemaRoot testVSR2 = createTestVSR()) { final VectorSchemaRoot original = new VectorSchemaRoot( Stream.concat( testVSR1.getFieldVectors().stream(), testVSR2.getFieldVectors().stream()) .collect(Collectors.toList())); assertTrue(imported.equals(original)); } imported.close(); } @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); Field itemField = new Field( "col1", FieldType.nullable(new ArrowType.Struct()), 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)); return schema; } @Test public void testImportedBufferAsNioBuffer() { IntVector imported; try (final IntVector vector = new IntVector("v", allocator)) { setVector(vector, 1, 2, 3, null); imported = (IntVector) vectorRoundtrip(vector); } ArrowBuf dataBuffer = imported.getDataBuffer(); ByteBuffer nioBuffer = dataBuffer.nioBuffer().asReadOnlyBuffer(); nioBuffer.order(ByteOrder.nativeOrder()); assertEquals(1, nioBuffer.getInt(0)); assertEquals(2, nioBuffer.getInt(1 << 2)); assertEquals(3, nioBuffer.getInt(2 << 2)); imported.close(); } @Test public void testImportReleasedArray() { // 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); } // Release array structure consumerArrowArray.markReleased(); // Consumer tried to imports vector but fails Exception e = assertThrows( IllegalStateException.class, () -> { Data.importVector(allocator, consumerArrowArray, consumerArrowSchema, null); }); assertEquals("Cannot import released ArrowArray", e.getMessage()); } } @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); Map metadata = new HashMap<>(); metadata.put("key", "value"); FieldType fieldType = new FieldType(true, ArrowType.Utf8.INSTANCE, null, metadata); VarCharVector varCharVector = new VarCharVector("varchar", fieldType, 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); return new VectorSchemaRoot(fields, vectors); } } apache-arrow-java-19.0.0/c/src/test/java/org/apache/arrow/c/StreamTest.java000066400000000000000000000376711515451615200263470ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.junit.jupiter.api.Assertions.assertThrows; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; 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.util.AutoCloseables; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; import org.apache.arrow.vector.ViewVarBinaryVector; 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.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.DictionaryEncoding; 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; final class StreamTest { private RootAllocator allocator = null; @BeforeEach public void setUp() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void tearDown() { allocator.close(); } @Test public void testRoundtripInts() throws Exception { final Schema schema = new Schema(Collections.singletonList(Field.nullable("ints", new ArrowType.Int(32, true)))); final List batches = new ArrayList<>(); try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final IntVector ints = (IntVector) root.getVector(0); VectorUnloader unloader = new VectorUnloader(root); root.allocateNew(); ints.setSafe(0, 1); ints.setSafe(1, 2); ints.setSafe(2, 4); ints.setSafe(3, 8); root.setRowCount(4); batches.add(unloader.getRecordBatch()); root.allocateNew(); ints.setSafe(0, 1); ints.setNull(1); ints.setSafe(2, 4); ints.setNull(3); root.setRowCount(4); batches.add(unloader.getRecordBatch()); roundtrip(schema, batches); } } @Test public void roundtripStrings() throws Exception { final Schema schema = new Schema( Arrays.asList( Field.nullable("ints", new ArrowType.Int(32, true)), Field.nullable("strs", new ArrowType.Utf8()))); final List batches = new ArrayList<>(); try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final IntVector ints = (IntVector) root.getVector(0); final VarCharVector strs = (VarCharVector) root.getVector(1); VectorUnloader unloader = new VectorUnloader(root); root.allocateNew(); ints.setSafe(0, 1); ints.setSafe(1, 2); ints.setSafe(2, 4); ints.setSafe(3, 8); strs.setSafe(0, "".getBytes(StandardCharsets.UTF_8)); strs.setSafe(1, "a".getBytes(StandardCharsets.UTF_8)); strs.setSafe(2, "bc".getBytes(StandardCharsets.UTF_8)); strs.setSafe(3, "defg".getBytes(StandardCharsets.UTF_8)); root.setRowCount(4); batches.add(unloader.getRecordBatch()); root.allocateNew(); ints.setSafe(0, 1); ints.setNull(1); ints.setSafe(2, 4); ints.setNull(3); strs.setSafe(0, "".getBytes(StandardCharsets.UTF_8)); strs.setNull(1); strs.setSafe(2, "bc".getBytes(StandardCharsets.UTF_8)); strs.setNull(3); root.setRowCount(4); batches.add(unloader.getRecordBatch()); roundtrip(schema, batches); } } @Test public void roundtripStringViews() throws Exception { final Schema schema = new Schema( Arrays.asList( Field.nullable("ints", new ArrowType.Int(32, true)), Field.nullable("string_views", new ArrowType.Utf8View()))); final List batches = new ArrayList<>(); try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final IntVector ints = (IntVector) root.getVector(0); final ViewVarCharVector strs = (ViewVarCharVector) root.getVector(1); VectorUnloader unloader = new VectorUnloader(root); root.allocateNew(); ints.setSafe(0, 1); ints.setSafe(1, 2); ints.setSafe(2, 4); ints.setSafe(3, 8); strs.setSafe(0, "".getBytes(StandardCharsets.UTF_8)); strs.setSafe(1, "a".getBytes(StandardCharsets.UTF_8)); strs.setSafe(2, "bc1234567890bc".getBytes(StandardCharsets.UTF_8)); strs.setSafe(3, "defg1234567890defg".getBytes(StandardCharsets.UTF_8)); root.setRowCount(4); batches.add(unloader.getRecordBatch()); root.allocateNew(); ints.setSafe(0, 1); ints.setNull(1); ints.setSafe(2, 4); ints.setNull(3); strs.setSafe(0, "".getBytes(StandardCharsets.UTF_8)); strs.setNull(1); strs.setSafe(2, "bc1234567890bc".getBytes(StandardCharsets.UTF_8)); strs.setNull(3); root.setRowCount(4); batches.add(unloader.getRecordBatch()); roundtrip(schema, batches); } } @Test public void roundtripBinaryViews() throws Exception { final Schema schema = new Schema( Arrays.asList( Field.nullable("ints", new ArrowType.Int(32, true)), Field.nullable("binary_views", new ArrowType.BinaryView()))); final List batches = new ArrayList<>(); try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final IntVector ints = (IntVector) root.getVector(0); final ViewVarBinaryVector strs = (ViewVarBinaryVector) root.getVector(1); VectorUnloader unloader = new VectorUnloader(root); root.allocateNew(); ints.setSafe(0, 1); ints.setSafe(1, 2); ints.setSafe(2, 4); ints.setSafe(3, 8); strs.setSafe(0, new byte[0]); strs.setSafe(1, new byte[] {97}); strs.setSafe(2, new byte[] {98, 99, 49, 50, 51, 52, 53, 54, 55, 56, 57, 48, 98, 99}); strs.setSafe( 3, new byte[] { 100, 101, 102, 103, 49, 50, 51, 52, 53, 54, 55, 56, 57, 48, 100, 101, 102, 103 }); root.setRowCount(4); batches.add(unloader.getRecordBatch()); root.allocateNew(); ints.setSafe(0, 1); ints.setNull(1); ints.setSafe(2, 4); ints.setNull(3); strs.setSafe(0, new byte[0]); strs.setNull(1); strs.setSafe(2, new byte[] {98, 99, 49, 50, 51, 52, 53, 54, 55, 56, 57, 48, 98, 99}); strs.setNull(3); root.setRowCount(4); batches.add(unloader.getRecordBatch()); roundtrip(schema, batches); } } @Test public void roundtripDictionary() throws Exception { final ArrowType.Int indexType = new ArrowType.Int(32, true); final DictionaryEncoding encoding = new DictionaryEncoding(0L, false, indexType); final Schema schema = new Schema( Collections.singletonList( new Field( "dict", new FieldType(/* nullable= */ true, indexType, encoding), Collections.emptyList()))); final List batches = new ArrayList<>(); try (final CDataDictionaryProvider provider = new CDataDictionaryProvider(); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final VarCharVector dictionary = new VarCharVector("values", allocator); dictionary.allocateNew(); dictionary.setSafe(0, "foo".getBytes(StandardCharsets.UTF_8)); dictionary.setSafe(1, "bar".getBytes(StandardCharsets.UTF_8)); dictionary.setNull(2); dictionary.setValueCount(3); provider.put(new Dictionary(dictionary, encoding)); final IntVector encoded = (IntVector) root.getVector(0); VectorUnloader unloader = new VectorUnloader(root); root.allocateNew(); encoded.setSafe(0, 0); encoded.setSafe(1, 1); encoded.setSafe(2, 0); encoded.setSafe(3, 2); root.setRowCount(4); batches.add(unloader.getRecordBatch()); root.allocateNew(); encoded.setSafe(0, 0); encoded.setNull(1); encoded.setSafe(2, 1); encoded.setNull(3); root.setRowCount(4); batches.add(unloader.getRecordBatch()); roundtrip(schema, batches, provider); } } @Test public void importReleasedStream() { try (final ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) { Exception e = assertThrows( IllegalStateException.class, () -> Data.importArrayStream(allocator, stream)); assertThat(e).hasMessageContaining("Cannot import released ArrowArrayStream"); } } @Test public void getNextError() throws Exception { final Schema schema = new Schema(Collections.singletonList(Field.nullable("ints", new ArrowType.Int(32, true)))); final List batches = new ArrayList<>(); try (final ArrowReader source = new InMemoryArrowReader( allocator, schema, batches, new DictionaryProvider.MapDictionaryProvider()) { @Override public boolean loadNextBatch() throws IOException { throw new IOException("Failed to load batch!"); } }; final ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) { Data.exportArrayStream(allocator, source, stream); try (final ArrowReader reader = Data.importArrayStream(allocator, stream)) { assertThat(reader.getVectorSchemaRoot().getSchema()).isEqualTo(schema); final IOException e = assertThrows(IOException.class, reader::loadNextBatch); assertThat(e).hasMessageContaining("Failed to load batch!"); assertThat(e).hasMessageContaining("[errno "); } } } @Test public void getSchemaError() throws Exception { final Schema schema = new Schema(Collections.singletonList(Field.nullable("ints", new ArrowType.Int(32, true)))); final List batches = new ArrayList<>(); try (final ArrowReader source = new InMemoryArrowReader( allocator, schema, batches, new DictionaryProvider.MapDictionaryProvider()) { @Override protected Schema readSchema() { throw new IllegalArgumentException("Failed to read schema!"); } }; final ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) { Data.exportArrayStream(allocator, source, stream); try (final ArrowReader reader = Data.importArrayStream(allocator, stream)) { final IOException e = assertThrows(IOException.class, reader::getVectorSchemaRoot); assertThat(e).hasMessageContaining("Failed to read schema!"); assertThat(e).hasMessageContaining("[errno "); } } } void roundtrip(Schema schema, List batches, DictionaryProvider provider) throws Exception { ArrowReader source = new InMemoryArrowReader(allocator, schema, batches, provider); try (final ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final VectorLoader loader = new VectorLoader(root); Data.exportArrayStream(allocator, source, stream); try (final ArrowReader reader = Data.importArrayStream(allocator, stream)) { assertThat(reader.getVectorSchemaRoot().getSchema()).isEqualTo(schema); for (ArrowRecordBatch batch : batches) { assertThat(reader.loadNextBatch()).isTrue(); loader.load(batch); assertThat(reader.getVectorSchemaRoot().getRowCount()).isEqualTo(root.getRowCount()); for (int i = 0; i < root.getFieldVectors().size(); i++) { final FieldVector expected = root.getVector(i); final FieldVector actual = reader.getVectorSchemaRoot().getVector(i); assertVectorsEqual(expected, actual); } } assertThat(reader.loadNextBatch()).isFalse(); assertThat(reader.getDictionaryIds()).isEqualTo(provider.getDictionaryIds()); for (Map.Entry entry : reader.getDictionaryVectors().entrySet()) { final FieldVector expected = provider.lookup(entry.getKey()).getVector(); final FieldVector actual = entry.getValue().getVector(); assertVectorsEqual(expected, actual); } } } } void roundtrip(Schema schema, List batches) throws Exception { roundtrip(schema, batches, new CDataDictionaryProvider()); } 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()); assertThat(new RangeEqualsVisitor(expected, actual).rangeEquals(range)) .as("Vectors were not equal.\nExpected: %s\nGot: %s", expected, actual) .isTrue(); } /** An ArrowReader backed by a fixed list of batches. */ static class InMemoryArrowReader extends ArrowReader { private final Schema schema; private final List batches; private final DictionaryProvider provider; private int nextBatch; InMemoryArrowReader( BufferAllocator allocator, Schema schema, List batches, DictionaryProvider provider) { super(allocator); this.schema = schema; this.batches = batches; this.provider = provider; 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()) { VectorLoader loader = new VectorLoader(getVectorSchemaRoot()); loader.load(batches.get(nextBatch++)); return true; } return false; } @Override public long bytesRead() { return 0; } @Override protected void closeReadSource() throws IOException { try { AutoCloseables.close(batches); } catch (Exception e) { throw new IOException(e); } } @Override protected Schema readSchema() { return schema; } } } apache-arrow-java-19.0.0/c/src/test/python/000077500000000000000000000000001515451615200204075ustar00rootroot00000000000000apache-arrow-java-19.0.0/c/src/test/python/integration_tests.py000066400000000000000000000442631515451615200245370ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import decimal import gc import os import sys import unittest import xml.etree.ElementTree as ET import jpype import pyarrow as pa import pyarrow.ipc as ipc from pyarrow.cffi import ffi def setup_jvm(): # This test requires Arrow Java to be built in the same source tree try: arrow_dir = os.environ["ARROW_SOURCE_DIR"] except KeyError: arrow_dir = os.path.join(os.path.dirname( __file__), '..', '..', '..', '..', '..') pom_path = os.path.join(arrow_dir, 'java', 'pom.xml') tree = ET.parse(pom_path) version = tree.getroot().find( 'POM:version', namespaces={ 'POM': 'http://maven.apache.org/POM/4.0.0' }).text jar_path = os.path.join( arrow_dir, 'java', 'tools', 'target', 'arrow-tools-{}-jar-with-dependencies.jar'.format(version)) jar_path = os.getenv("ARROW_TOOLS_JAR", jar_path) jar_path += ":{}".format(os.path.join(arrow_dir, "java", "c/target/arrow-c-data-{}.jar".format(version))) kwargs = {} # This will be the default behaviour in jpype 0.8+ kwargs['convertStrings'] = False # For debugging purpose please uncomment the following, and include *jvm_args, before **kwargs # in startJVM function call # jvm_args = [ # "-agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=*:5005" # ] jpype.startJVM(jpype.getDefaultJVMPath(), "-Djava.class.path=" + jar_path, **kwargs) class Bridge: def __init__(self): self.java_allocator = jpype.JPackage( "org").apache.arrow.memory.RootAllocator(sys.maxsize) self.java_c = jpype.JPackage("org").apache.arrow.c def java_to_python_field(self, jfield): c_schema = ffi.new("struct ArrowSchema*") ptr_schema = int(ffi.cast("uintptr_t", c_schema)) self.java_c.Data.exportField(self.java_allocator, jfield, None, self.java_c.ArrowSchema.wrap(ptr_schema)) return pa.Field._import_from_c(ptr_schema) def java_to_python_array(self, vector, dictionary_provider=None): c_schema = ffi.new("struct ArrowSchema*") ptr_schema = int(ffi.cast("uintptr_t", c_schema)) c_array = ffi.new("struct ArrowArray*") ptr_array = int(ffi.cast("uintptr_t", c_array)) self.java_c.Data.exportVector(self.java_allocator, vector, dictionary_provider, self.java_c.ArrowArray.wrap( ptr_array), self.java_c.ArrowSchema.wrap(ptr_schema)) return pa.Array._import_from_c(ptr_array, ptr_schema) def java_to_python_record_batch(self, root): c_schema = ffi.new("struct ArrowSchema*") ptr_schema = int(ffi.cast("uintptr_t", c_schema)) c_array = ffi.new("struct ArrowArray*") ptr_array = int(ffi.cast("uintptr_t", c_array)) self.java_c.Data.exportVectorSchemaRoot(self.java_allocator, root, None, self.java_c.ArrowArray.wrap( ptr_array), self.java_c.ArrowSchema.wrap(ptr_schema)) return pa.RecordBatch._import_from_c(ptr_array, ptr_schema) def java_to_python_reader(self, reader): c_stream = ffi.new("struct ArrowArrayStream*") ptr_stream = int(ffi.cast("uintptr_t", c_stream)) self.java_c.Data.exportArrayStream(self.java_allocator, reader, self.java_c.ArrowArrayStream.wrap(ptr_stream)) return pa.RecordBatchReader._import_from_c(ptr_stream) def python_to_java_field(self, field): c_schema = self.java_c.ArrowSchema.allocateNew(self.java_allocator) field._export_to_c(c_schema.memoryAddress()) return self.java_c.Data.importField(self.java_allocator, c_schema, None) def python_to_java_array(self, array, dictionary_provider=None): c_schema = self.java_c.ArrowSchema.allocateNew(self.java_allocator) c_array = self.java_c.ArrowArray.allocateNew(self.java_allocator) array._export_to_c(c_array.memoryAddress(), c_schema.memoryAddress()) return self.java_c.Data.importVector(self.java_allocator, c_array, c_schema, dictionary_provider) def python_to_java_record_batch(self, record_batch): c_schema = self.java_c.ArrowSchema.allocateNew(self.java_allocator) c_array = self.java_c.ArrowArray.allocateNew(self.java_allocator) record_batch._export_to_c( c_array.memoryAddress(), c_schema.memoryAddress()) return self.java_c.Data.importVectorSchemaRoot(self.java_allocator, c_array, c_schema, None) def python_to_java_reader(self, reader): c_stream = self.java_c.ArrowArrayStream.allocateNew(self.java_allocator) reader._export_to_c(c_stream.memoryAddress()) return self.java_c.Data.importArrayStream(self.java_allocator, c_stream) def close(self): self.java_allocator.close() class TestPythonIntegration(unittest.TestCase): @classmethod def setUpClass(cls) -> None: setup_jvm() def setUp(self): gc.collect() self.old_allocated_python = pa.total_allocated_bytes() self.bridge = Bridge() def tearDown(self): self.bridge.close() gc.collect() diff_python = pa.total_allocated_bytes() - self.old_allocated_python self.assertEqual( pa.total_allocated_bytes(), self.old_allocated_python, f"PyArrow memory was not adequately released: {diff_python} bytes lost") def round_trip_field(self, field_generator): original_field = field_generator() java_field = self.bridge.python_to_java_field(original_field) del original_field new_field = self.bridge.java_to_python_field(java_field) del java_field expected = field_generator() self.assertEqual(expected, new_field) def round_trip_array(self, array_generator, check_metadata=True): original_arr = array_generator() with self.bridge.java_c.CDataDictionaryProvider() as dictionary_provider, \ self.bridge.python_to_java_array(original_arr, dictionary_provider) as vector: del original_arr new_array = self.bridge.java_to_python_array(vector, dictionary_provider) expected = array_generator() self.assertEqual(expected, new_array) if check_metadata: self.assertTrue(new_array.type.equals(expected.type, check_metadata=True)) def round_trip_record_batch(self, rb_generator): original_rb = rb_generator() with self.bridge.python_to_java_record_batch(original_rb) as root: del original_rb new_rb = self.bridge.java_to_python_record_batch(root) expected = rb_generator() self.assertEqual(expected, new_rb) def round_trip_reader(self, schema, batches): reader = pa.RecordBatchReader.from_batches(schema, batches) java_reader = self.bridge.python_to_java_reader(reader) del reader py_reader = self.bridge.java_to_python_reader(java_reader) del java_reader actual = list(py_reader) self.assertEqual(batches, actual) def test_string_array(self): self.round_trip_array(lambda: pa.array([None, "a", "bb", "ccc"])) def test_stringview_array(self): # with nulls short strings self.round_trip_array(lambda: pa.array([None, "a", "bb", "c"], type=pa.string_view())) # with nulls long and strings self.round_trip_array(lambda: pa.array([None, "a", "bb"*10, "c"*13], type=pa.string_view())) # without nulls short strings self.round_trip_array(lambda: pa.array(["a", "bb", "c"], type=pa.string_view())) # without nulls long and strings self.round_trip_array(lambda: pa.array(["a", "bb"*10, "c"*13], type=pa.string_view())) # with multiple data buffers arr1 = pa.array(["a", "bb", "c"], type=pa.string_view()) arr2 = pa.array(["b", "ee" * 10, "f" * 20], type=pa.string_view()) arr3 = pa.array(["c", "abc" * 20, "efg" * 30], type=pa.string_view()) arr4 = pa.array(["d", "abcd" * 100, "efgh" * 200], type=pa.string_view()) self.round_trip_array(lambda: pa.concat_arrays([arr1, arr2, arr3, arr4])) # empty strings self.round_trip_array(lambda: pa.array(["", "bb" * 10, "c", "", "d", ""], type=pa.string_view())) # null value variations self.round_trip_array(lambda: pa.array(["bb" * 10, None, "", "d", None], type=pa.string_view())) # empty array self.round_trip_array(lambda: pa.array([], type=pa.string_view())) # all null array self.round_trip_array(lambda: pa.array([None, None, None], type=pa.string_view())) def test_binaryview_array(self): # with nulls short binary values self.round_trip_array(lambda: pa.array([None, bytes([97]), bytes([98, 98]), bytes([99])], type=pa.binary_view())) # with nulls long binary values self.round_trip_array(lambda: pa.array([None, bytes([97]), bytes([98, 98] * 10), bytes([99] * 13)], type=pa.binary_view())) # without nulls short binary values self.round_trip_array(lambda: pa.array([bytes([97]), bytes([98, 98]), bytes([99])], type=pa.binary_view())) # without nulls long binary values self.round_trip_array(lambda: pa.array([bytes([97]), bytes([98, 98] * 10), bytes([99] * 13)], type=pa.binary_view())) # with multiple data buffers arr1 = pa.array([bytes([97]), bytes([98, 98]), bytes([99])], type=pa.binary_view()) arr2 = pa.array([bytes([98]), bytes([98, 98] * 10), bytes([99] * 13)], type=pa.binary_view()) arr3 = pa.array([bytes([99]), bytes([98, 100] * 100), bytes([99, 100]) * 30], type=pa.binary_view()) arr4 = pa.array([bytes([100]), bytes([98, 100, 101] * 200), bytes([98, 99]) * 300], type=pa.binary_view()) self.round_trip_array(lambda: pa.concat_arrays([arr1, arr2, arr3, arr4])) # empty binary values self.round_trip_array(lambda: pa.array([bytes([]), bytes([97, 97]) * 10, bytes([98]), bytes([]), bytes([97]), bytes([])], type=pa.binary_view())) # null value variations self.round_trip_array(lambda: pa.array([bytes([97, 97]) * 10, None, bytes([]), bytes([99]), None], type=pa.binary_view())) # empty array self.round_trip_array(lambda: pa.array([], type=pa.binary_view())) # all null array self.round_trip_array(lambda: pa.array([None, None, None], type=pa.binary_view())) def test_decimal_array(self): data = [ round(decimal.Decimal(722.82), 2), round(decimal.Decimal(-934.11), 2), None, ] self.round_trip_array(lambda: pa.array(data, pa.decimal128(5, 2))) def test_int_array(self): self.round_trip_array(lambda: pa.array([1, 2, 3], type=pa.int32())) def test_list_array(self): self.round_trip_array(lambda: pa.array( [[], [0], [1, 2], [4, 5, 6]], pa.list_(pa.int64()) # disabled check_metadata since the list internal field name ("item") # is not preserved during round trips (it becomes "$data$"). ), check_metadata=False) def test_empty_list_array(self): """Validates GH-37056 fix. Empty list of int32 produces a vector with empty child data buffer, however with non-zero capacity. Using streaming forces the c-data array which represent the child data buffer to be NULL (pointer is 0). On Java side, an attempt to import such array triggered an exception described in GH-37056. """ with pa.BufferOutputStream() as bos: schema = pa.schema([pa.field("f0", pa.list_(pa.int32()), True)]) with ipc.new_stream(bos, schema) as writer: src = pa.RecordBatch.from_arrays([pa.array([[]])], schema=schema) writer.write(src) data_bytes = bos.getvalue() def recreate_batch(): with pa.input_stream(data_bytes) as ios: with ipc.open_stream(ios) as reader: return reader.read_next_batch() self.round_trip_record_batch(recreate_batch) def test_struct_array(self): fields = [ ("f1", pa.int32()), ("f2", pa.string()), ] data = [ {"f1": 1, "f2": "a"}, None, {"f1": 3, "f2": None}, {"f1": None, "f2": "d"}, {"f1": None, "f2": None}, ] self.round_trip_array(lambda: pa.array(data, type=pa.struct(fields))) def test_dict(self): self.round_trip_array( lambda: pa.array(["a", "b", None, "d"], pa.dictionary(pa.int64(), pa.utf8()))) def test_map(self): offsets = [0, None, 2, 6] pykeys = [b"a", b"b", b"c", b"d", b"e", b"f"] pyitems = [1, 2, 3, None, 4, 5] keys = pa.array(pykeys, type="binary") items = pa.array(pyitems, type="i4") self.round_trip_array( lambda: pa.MapArray.from_arrays(offsets, keys, items)) def test_field(self): self.round_trip_field(lambda: pa.field("aa", pa.bool_())) def test_field_nested(self): self.round_trip_field(lambda: pa.field( "test", pa.list_(pa.int32()), nullable=True)) def test_field_metadata(self): self.round_trip_field(lambda: pa.field("aa", pa.bool_(), {"a": "b"})) def test_record_batch_with_list(self): data = [ pa.array([[1], [2], [3], [4, 5, 6]]), pa.array([1, 2, 3, 4]), pa.array(['foo', 'bar', 'baz', None]), pa.array([True, None, False, True]) ] self.round_trip_record_batch( lambda: pa.RecordBatch.from_arrays(data, ['f0', 'f1', 'f2', 'f3'])) def test_reader_roundtrip(self): schema = pa.schema([("ints", pa.int64()), ("strs", pa.string())]) data = [ pa.record_batch([[1, 2, 3, None], ["a", "bc", None, ""]], schema=schema), pa.record_batch([[None, 4, 5, 6], [None, "", "def", "g"]], schema=schema), ] self.round_trip_reader(schema, data) def test_reader_complex_roundtrip(self): schema = pa.schema([ ("str_dict", pa.dictionary(pa.int8(), pa.string())), ("int_list", pa.list_(pa.int64())), ]) dictionary = pa.array(["a", "bc", None]) data = [ pa.record_batch([pa.DictionaryArray.from_arrays([0, 2], dictionary), [[1, 2, 3], None]], schema=schema), pa.record_batch([pa.DictionaryArray.from_arrays([None, 1], dictionary), [[], [4]]], schema=schema), ] self.round_trip_reader(schema, data) def test_listview_array(self): self.round_trip_array(lambda: pa.array( [[], [0], [1, 2], [4, 5, 6]], pa.list_view(pa.int64()) # disabled check_metadata since in Java API the listview # internal field name ("item") is not preserved # during round trips (it becomes "$data$"). ), check_metadata=False) def test_empty_listview_array(self): with pa.BufferOutputStream() as bos: schema = pa.schema([pa.field("f0", pa.list_view(pa.int32()), True)]) with ipc.new_stream(bos, schema) as writer: src = pa.RecordBatch.from_arrays( [pa.array([[]], pa.list_view(pa.int32()))], schema=schema) writer.write(src) data_bytes = bos.getvalue() def recreate_batch(): with pa.input_stream(data_bytes) as ios: with ipc.open_stream(ios) as reader: return reader.read_next_batch() self.round_trip_record_batch(recreate_batch) def test_largelistview_array(self): self.round_trip_array(lambda: pa.array( [[], [0], [1, 2], [4, 5, 6]], pa.large_list_view(pa.int64()) # disabled check_metadata since in Java API the listview # internal field name ("item") is not preserved # during round trips (it becomes "$data$"). ), check_metadata=False) def test_empty_largelistview_array(self): with pa.BufferOutputStream() as bos: schema = pa.schema([pa.field("f0", pa.large_list_view(pa.int32()), True)]) with ipc.new_stream(bos, schema) as writer: src = pa.RecordBatch.from_arrays( [pa.array([[]], pa.large_list_view(pa.int32()))], schema=schema) writer.write(src) data_bytes = bos.getvalue() def recreate_batch(): with pa.input_stream(data_bytes) as ios: with ipc.open_stream(ios) as reader: return reader.read_next_batch() self.round_trip_record_batch(recreate_batch) def test_runendencoded_array(self): # empty vector self.round_trip_array(lambda: pa.RunEndEncodedArray.from_arrays([], [], pa.run_end_encoded(pa.int64(), pa.int64()))) # constant null vector self.round_trip_array(lambda: pa.RunEndEncodedArray.from_arrays([10], [None])) # constant int vector self.round_trip_array(lambda: pa.RunEndEncodedArray.from_arrays([10], [10])) # run end int vector self.round_trip_array(lambda: pa.RunEndEncodedArray.from_arrays([3, 5, 10, 12, 19], [1, 2, 1, None, 3])) # run end string vector self.round_trip_array(lambda: pa.RunEndEncodedArray.from_arrays([3, 5, 10, 12, 19], ["1", "2", "1", None, "3"])) if __name__ == '__main__': unittest.main(verbosity=2) apache-arrow-java-19.0.0/ci/000077500000000000000000000000001515451615200154715ustar00rootroot00000000000000apache-arrow-java-19.0.0/ci/docker/000077500000000000000000000000001515451615200167405ustar00rootroot00000000000000apache-arrow-java-19.0.0/ci/docker/conda-jni.dockerfile000066400000000000000000000020261515451615200226330ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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=11 ARG maven=3.9.9 RUN micromamba install -y \ -c conda-forge \ cmake \ compilers \ maven=${maven} \ ninja \ openjdk=${jdk} && \ micromamba clean --all apache-arrow-java-19.0.0/ci/docker/vcpkg-jni.dockerfile000066400000000000000000000022361515451615200226640ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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=11 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 apache-arrow-java-19.0.0/ci/scripts/000077500000000000000000000000001515451615200171605ustar00rootroot00000000000000apache-arrow-java-19.0.0/ci/scripts/build.sh000077500000000000000000000063531515451615200206250ustar00rootroot00000000000000#!/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 if [[ "${ARROW_JAVA_BUILD:-ON}" != "ON" ]]; then exit fi source_dir=${1} build_dir=${2} java_jni_dist_dir=${3} mvn="mvn -B -DskipTests -Drat.skip=true -Dorg.slf4j.simpleLogger.log.org.apache.maven.cli.transfer.Slf4jMavenTransferListener=warn" if [ "${ARROW_JAVA_SKIP_GIT_PLUGIN:-OFF}" = "ON" ]; then mvn="${mvn} -Dmaven.gitcommitid.skip=true" fi # We want to do an out-of-source build since (when using docker) we'll pollute # the source directory with files owned by root, but Maven does not really # support this. Instead, copy directories to the build directory. rm -rf "${build_dir}" mkdir -p "${build_dir}/arrow-format" cp -r "${source_dir}/arrow-format" "${build_dir}" cp -r "${source_dir}/dev" "${build_dir}" # Instead of hardcoding the list of directories to copy, find pom.xml and then # crawl back up to the top. GNU realpath has --relative-to but this does not # work on macOS 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}'))") else 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}" done pushd "${build_dir}" # 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:-OFF}" = "ON" ]; then mvn="${mvn} -Darrow.c.jni.dist.dir=${java_jni_dist_dir} -Parrow-c-data" fi 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 [ "${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 apache-arrow-java-19.0.0/ci/scripts/download_cpp.sh000077500000000000000000000034601515451615200221730ustar00rootroot00000000000000#!/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 apache-arrow-java-19.0.0/ci/scripts/jni_build.sh000077500000000000000000000052631515451615200214640ustar00rootroot00000000000000#!/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 apache-arrow-java-19.0.0/ci/scripts/jni_full_build.sh000077500000000000000000000057661515451615200225160ustar00rootroot00000000000000#!/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 apache-arrow-java-19.0.0/ci/scripts/jni_macos_build.sh000077500000000000000000000067171515451615200226530ustar00rootroot00000000000000#!/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 apache-arrow-java-19.0.0/ci/scripts/jni_manylinux_build.sh000077500000000000000000000072231515451615200235660ustar00rootroot00000000000000#!/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 apache-arrow-java-19.0.0/ci/scripts/jni_windows_build.sh000077500000000000000000000102751515451615200232350ustar00rootroot00000000000000#!/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 apache-arrow-java-19.0.0/ci/scripts/test.sh000077500000000000000000000033261515451615200205020ustar00rootroot00000000000000#!/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 -ex if [[ "${ARROW_JAVA_TEST:-ON}" != "ON" ]]; then exit fi 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" pushd "${build_dir}" ${mvn} -Darrow.test.dataRoot="${source_dir}/testing/data" clean test projects=() if [ "${ARROW_JAVA_JNI}" = "ON" ]; then projects+=(adapter/orc) projects+=(dataset) 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}" fi if [ "${ARROW_JAVA_CDATA}" = "ON" ]; then ${mvn} clean test -Parrow-c-data -pl c -Darrow.c.jni.dist.dir="${java_jni_dist_dir}" fi popd apache-arrow-java-19.0.0/ci/scripts/util_log.sh000066400000000000000000000016521515451615200213360ustar00rootroot00000000000000# 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::" } apache-arrow-java-19.0.0/compose.yaml000066400000000000000000000074161515451615200174370ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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: 11, 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: 11, 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: ["git config --global --add safe.directory /arrow-java && \ /arrow-java/ci/scripts/jni_manylinux_build.sh /arrow-java /arrow /build/java /arrow-java/jni"] apache-arrow-java-19.0.0/compression/000077500000000000000000000000001515451615200174375ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/pom.xml000066400000000000000000000044261515451615200207620ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-java-root 19.0.0 arrow-compression Arrow Compression (Experimental/Contrib) A library for working with the compression/decompression of Arrow data. org.apache.arrow arrow-vector ${arrow.vector.classifier} org.apache.arrow arrow-memory-core org.apache.arrow arrow-memory-unsafe test org.immutables value-annotations org.apache.commons commons-compress 1.28.0 com.github.luben zstd-jni 1.5.7-6 apache-arrow-java-19.0.0/compression/src/000077500000000000000000000000001515451615200202265ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/main/000077500000000000000000000000001515451615200211525ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/main/java/000077500000000000000000000000001515451615200220735ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/main/java/module-info.java000066400000000000000000000024021515451615200251520ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ import org.apache.arrow.vector.compression.CompressionCodec; module org.apache.arrow.compression { exports org.apache.arrow.compression; requires com.github.luben.zstd_jni; requires org.apache.arrow.memory.core; requires org.apache.arrow.vector; requires org.apache.commons.compress; // Also defined under META-INF/services to support non-modular applications provides CompressionCodec.Factory with org.apache.arrow.compression.CommonsCompressionFactory; } apache-arrow-java-19.0.0/compression/src/main/java/org/000077500000000000000000000000001515451615200226625ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/main/java/org/apache/000077500000000000000000000000001515451615200241035ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200252355ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/main/java/org/apache/arrow/compression/000077500000000000000000000000001515451615200275765ustar00rootroot00000000000000CommonsCompressionFactory.java000066400000000000000000000037371515451615200355610ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/main/java/org/apache/arrow/compression/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compression; import org.apache.arrow.vector.compression.CompressionCodec; import org.apache.arrow.vector.compression.CompressionUtil; /** * Default implementation of factory supported LZ4 and ZSTD compression. * *

    // TODO(ARROW-12115): Rename this class. */ public class CommonsCompressionFactory implements CompressionCodec.Factory { public static final CommonsCompressionFactory INSTANCE = new CommonsCompressionFactory(); @Override public CompressionCodec createCodec(CompressionUtil.CodecType codecType) { switch (codecType) { case LZ4_FRAME: return new Lz4CompressionCodec(); case ZSTD: return new ZstdCompressionCodec(); default: throw new IllegalArgumentException("Compression type not supported: " + codecType); } } @Override public CompressionCodec createCodec(CompressionUtil.CodecType codecType, int compressionLevel) { switch (codecType) { case LZ4_FRAME: return new Lz4CompressionCodec(); case ZSTD: return new ZstdCompressionCodec(compressionLevel); default: throw new IllegalArgumentException("Compression type not supported: " + codecType); } } } Lz4CompressionCodec.java000066400000000000000000000075641515451615200342270ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/main/java/org/apache/arrow/compression/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compression; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.compression.AbstractCompressionCodec; import org.apache.arrow.vector.compression.CompressionUtil; import org.apache.commons.compress.compressors.lz4.FramedLZ4CompressorInputStream; import org.apache.commons.compress.compressors.lz4.FramedLZ4CompressorOutputStream; import org.apache.commons.compress.utils.IOUtils; /** Compression codec for the LZ4 algorithm. */ public class Lz4CompressionCodec extends AbstractCompressionCodec { @Override protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf uncompressedBuffer) { Preconditions.checkArgument( uncompressedBuffer.writerIndex() <= Integer.MAX_VALUE, "The uncompressed buffer size exceeds the integer limit %s.", Integer.MAX_VALUE); byte[] inBytes = new byte[(int) uncompressedBuffer.writerIndex()]; uncompressedBuffer.getBytes(/* index= */ 0, inBytes); ByteArrayOutputStream baos = new ByteArrayOutputStream(); try (InputStream in = new ByteArrayInputStream(inBytes); OutputStream out = new FramedLZ4CompressorOutputStream(baos)) { IOUtils.copy(in, out); } catch (IOException e) { throw new RuntimeException(e); } byte[] outBytes = baos.toByteArray(); ArrowBuf compressedBuffer = allocator.buffer(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + outBytes.length); compressedBuffer.setBytes(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH, outBytes); compressedBuffer.writerIndex(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + outBytes.length); return compressedBuffer; } @Override protected ArrowBuf doDecompress(BufferAllocator allocator, ArrowBuf compressedBuffer) { Preconditions.checkArgument( compressedBuffer.writerIndex() <= Integer.MAX_VALUE, "The compressed buffer size exceeds the integer limit %s", Integer.MAX_VALUE); long decompressedLength = readUncompressedLength(compressedBuffer); byte[] inBytes = new byte [(int) (compressedBuffer.writerIndex() - CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH)]; compressedBuffer.getBytes(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH, inBytes); ByteArrayOutputStream out = new ByteArrayOutputStream((int) decompressedLength); try (InputStream in = new FramedLZ4CompressorInputStream(new ByteArrayInputStream(inBytes))) { IOUtils.copy(in, out); } catch (IOException e) { throw new RuntimeException(e); } byte[] outBytes = out.toByteArray(); ArrowBuf decompressedBuffer = allocator.buffer(outBytes.length); decompressedBuffer.setBytes(/* index= */ 0, outBytes); decompressedBuffer.writerIndex(decompressedLength); return decompressedBuffer; } @Override public CompressionUtil.CodecType getCodecType() { return CompressionUtil.CodecType.LZ4_FRAME; } } ZstdCompressionCodec.java000066400000000000000000000071471515451615200344770ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/main/java/org/apache/arrow/compression/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compression; import com.github.luben.zstd.Zstd; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.compression.AbstractCompressionCodec; import org.apache.arrow.vector.compression.CompressionUtil; /** Compression codec for the ZSTD algorithm. */ public class ZstdCompressionCodec extends AbstractCompressionCodec { private int compressionLevel; private static final int DEFAULT_COMPRESSION_LEVEL = 3; public ZstdCompressionCodec() { this.compressionLevel = DEFAULT_COMPRESSION_LEVEL; } public ZstdCompressionCodec(int compressionLevel) { this.compressionLevel = compressionLevel; } @Override protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf uncompressedBuffer) { long maxSize = Zstd.compressBound(uncompressedBuffer.writerIndex()); long dstSize = CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + maxSize; ArrowBuf compressedBuffer = allocator.buffer(dstSize); long bytesWritten = Zstd.compressUnsafe( compressedBuffer.memoryAddress() + CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH, dstSize, /*src*/ uncompressedBuffer.memoryAddress(), /* srcSize= */ uncompressedBuffer.writerIndex(), /* level= */ this.compressionLevel); if (Zstd.isError(bytesWritten)) { compressedBuffer.close(); throw new RuntimeException("Error compressing: " + Zstd.getErrorName(bytesWritten)); } compressedBuffer.writerIndex(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + bytesWritten); return compressedBuffer; } @Override protected ArrowBuf doDecompress(BufferAllocator allocator, ArrowBuf compressedBuffer) { long decompressedLength = readUncompressedLength(compressedBuffer); ArrowBuf uncompressedBuffer = allocator.buffer(decompressedLength); long decompressedSize = Zstd.decompressUnsafe( uncompressedBuffer.memoryAddress(), decompressedLength, /* 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(decompressedSize)); } if (decompressedLength != decompressedSize) { uncompressedBuffer.close(); throw new RuntimeException( "Expected != actual decompressed length: " + decompressedLength + " != " + decompressedSize); } uncompressedBuffer.writerIndex(decompressedLength); return uncompressedBuffer; } @Override public CompressionUtil.CodecType getCodecType() { return CompressionUtil.CodecType.ZSTD; } } apache-arrow-java-19.0.0/compression/src/main/resources/000077500000000000000000000000001515451615200231645ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/main/resources/META-INF/000077500000000000000000000000001515451615200243245ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/main/resources/META-INF/services/000077500000000000000000000000001515451615200261475ustar00rootroot00000000000000org.apache.arrow.vector.compression.CompressionCodec$Factory000066400000000000000000000015031515451615200416450ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/main/resources/META-INF/services# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. org.apache.arrow.compression.CommonsCompressionFactory apache-arrow-java-19.0.0/compression/src/test/000077500000000000000000000000001515451615200212055ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/test/java/000077500000000000000000000000001515451615200221265ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/test/java/org/000077500000000000000000000000001515451615200227155ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/test/java/org/apache/000077500000000000000000000000001515451615200241365ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200252705ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/test/java/org/apache/arrow/compression/000077500000000000000000000000001515451615200276315ustar00rootroot00000000000000TestArrowReaderWriterWithCompression.java000066400000000000000000000246751515451615200377630ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/test/java/org/apache/arrow/compression/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compression; 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.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.channels.Channels; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Optional; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.GenerateSampleData; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.compression.CompressionUtil; import org.apache.arrow.vector.compression.NoCompressionCodec; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.ipc.ArrowFileReader; import org.apache.arrow.vector.ipc.ArrowFileWriter; import org.apache.arrow.vector.ipc.ArrowStreamReader; import org.apache.arrow.vector.ipc.ArrowStreamWriter; import org.apache.arrow.vector.ipc.message.IpcOption; 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.types.pojo.Schema; 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; public class TestArrowReaderWriterWithCompression { private BufferAllocator allocator; private ByteArrayOutputStream out; private VectorSchemaRoot root; @BeforeEach public void setup() { if (allocator == null) { allocator = new RootAllocator(Integer.MAX_VALUE); } out = new ByteArrayOutputStream(); root = null; } @AfterEach public void tearDown() { if (root != null) { root.close(); } if (allocator != null) { allocator.close(); } if (out != null) { out.reset(); } } private void createAndWriteArrowFile( DictionaryProvider provider, CompressionUtil.CodecType codecType) throws IOException { List fields = new ArrayList<>(); fields.add(new Field("col", FieldType.notNullable(new ArrowType.Utf8()), new ArrayList<>())); root = VectorSchemaRoot.create(new Schema(fields), allocator); final int rowCount = 10; GenerateSampleData.generateTestData(root.getVector(0), rowCount); root.setRowCount(rowCount); try (final ArrowFileWriter writer = new ArrowFileWriter( root, provider, Channels.newChannel(out), new HashMap<>(), IpcOption.DEFAULT, CommonsCompressionFactory.INSTANCE, codecType, Optional.of(7))) { writer.start(); writer.writeBatch(); writer.end(); } } private void createAndWriteArrowStream( DictionaryProvider provider, CompressionUtil.CodecType codecType) throws IOException { List fields = new ArrayList<>(); fields.add(new Field("col", FieldType.notNullable(new ArrowType.Utf8()), new ArrayList<>())); root = VectorSchemaRoot.create(new Schema(fields), allocator); final int rowCount = 10; GenerateSampleData.generateTestData(root.getVector(0), rowCount); root.setRowCount(rowCount); try (final ArrowStreamWriter writer = new ArrowStreamWriter( root, provider, Channels.newChannel(out), IpcOption.DEFAULT, CommonsCompressionFactory.INSTANCE, codecType, Optional.of(7))) { writer.start(); writer.writeBatch(); writer.end(); } } private Dictionary createDictionary(VarCharVector dictionaryVector) { setVector( dictionaryVector, "foo".getBytes(StandardCharsets.UTF_8), "bar".getBytes(StandardCharsets.UTF_8), "baz".getBytes(StandardCharsets.UTF_8)); return new Dictionary( dictionaryVector, new DictionaryEncoding(/* id= */ 1L, /* ordered= */ false, /* indexType= */ null)); } @Test public void testArrowFileZstdRoundTrip() throws Exception { createAndWriteArrowFile(null, CompressionUtil.CodecType.ZSTD); // with compression try (ArrowFileReader reader = new ArrowFileReader( new ByteArrayReadableSeekableByteChannel(out.toByteArray()), allocator, CommonsCompressionFactory.INSTANCE)) { assertEquals(1, reader.getRecordBlocks().size()); assertTrue(reader.loadNextBatch()); assertTrue(root.equals(reader.getVectorSchemaRoot())); assertFalse(reader.loadNextBatch()); } // without compression try (ArrowFileReader reader = new ArrowFileReader( new ByteArrayReadableSeekableByteChannel(out.toByteArray()), allocator, NoCompressionCodec.Factory.INSTANCE)) { assertEquals(1, reader.getRecordBlocks().size()); Exception exception = assertThrows(IllegalArgumentException.class, reader::loadNextBatch); assertEquals( "Please add arrow-compression module to use CommonsCompressionFactory for ZSTD", exception.getMessage()); } } @Test public void testArrowStreamZstdRoundTrip() throws Exception { createAndWriteArrowStream(null, CompressionUtil.CodecType.ZSTD); // with compression try (ArrowStreamReader reader = new ArrowStreamReader( new ByteArrayReadableSeekableByteChannel(out.toByteArray()), allocator, CommonsCompressionFactory.INSTANCE)) { assertTrue(reader.loadNextBatch()); assertTrue(root.equals(reader.getVectorSchemaRoot())); assertFalse(reader.loadNextBatch()); } // without compression try (ArrowStreamReader reader = new ArrowStreamReader( new ByteArrayReadableSeekableByteChannel(out.toByteArray()), allocator, NoCompressionCodec.Factory.INSTANCE)) { Exception exception = assertThrows(IllegalArgumentException.class, reader::loadNextBatch); assertEquals( "Please add arrow-compression module to use CommonsCompressionFactory for ZSTD", exception.getMessage()); } } @Test public void testArrowFileZstdRoundTripWithDictionary() throws Exception { VarCharVector dictionaryVector = (VarCharVector) FieldType.nullable(new ArrowType.Utf8()) .createNewSingleVector("f1_file", allocator, null); Dictionary dictionary = createDictionary(dictionaryVector); DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); provider.put(dictionary); createAndWriteArrowFile(provider, CompressionUtil.CodecType.ZSTD); // with compression try (ArrowFileReader reader = new ArrowFileReader( new ByteArrayReadableSeekableByteChannel(out.toByteArray()), allocator, CommonsCompressionFactory.INSTANCE)) { assertEquals(1, reader.getRecordBlocks().size()); assertTrue(reader.loadNextBatch()); assertTrue(root.equals(reader.getVectorSchemaRoot())); assertFalse(reader.loadNextBatch()); } // without compression try (ArrowFileReader reader = new ArrowFileReader( new ByteArrayReadableSeekableByteChannel(out.toByteArray()), allocator, NoCompressionCodec.Factory.INSTANCE)) { assertEquals(1, reader.getRecordBlocks().size()); Exception exception = assertThrows(IllegalArgumentException.class, reader::loadNextBatch); assertEquals( "Please add arrow-compression module to use CommonsCompressionFactory for ZSTD", exception.getMessage()); } dictionaryVector.close(); } @Test public void testArrowStreamZstdRoundTripWithDictionary() throws Exception { VarCharVector dictionaryVector = (VarCharVector) FieldType.nullable(new ArrowType.Utf8()) .createNewSingleVector("f1_stream", allocator, null); Dictionary dictionary = createDictionary(dictionaryVector); DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); provider.put(dictionary); createAndWriteArrowStream(provider, CompressionUtil.CodecType.ZSTD); // with compression try (ArrowStreamReader reader = new ArrowStreamReader( new ByteArrayReadableSeekableByteChannel(out.toByteArray()), allocator, CommonsCompressionFactory.INSTANCE)) { assertTrue(reader.loadNextBatch()); assertTrue(root.equals(reader.getVectorSchemaRoot())); assertFalse(reader.loadNextBatch()); } // without compression try (ArrowStreamReader reader = new ArrowStreamReader( new ByteArrayReadableSeekableByteChannel(out.toByteArray()), allocator, NoCompressionCodec.Factory.INSTANCE)) { Exception exception = assertThrows(IllegalArgumentException.class, reader::loadNextBatch); assertEquals( "Please add arrow-compression module to use CommonsCompressionFactory for ZSTD", exception.getMessage()); } dictionaryVector.close(); } public static void setVector(VarCharVector vector, byte[]... values) { final int length = values.length; vector.allocateNewSafe(); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } } TestCompressionCodec.java000066400000000000000000000340611515451615200345200ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/test/java/org/apache/arrow/compression/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compression; 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.assertTrue; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.channels.Channels; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; import java.util.HashMap; import java.util.List; import java.util.Optional; import java.util.function.BiConsumer; 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.util.AutoCloseables; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; import org.apache.arrow.vector.compression.CompressionCodec; import org.apache.arrow.vector.compression.CompressionUtil; import org.apache.arrow.vector.compression.NoCompressionCodec; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.ipc.ArrowFileReader; import org.apache.arrow.vector.ipc.ArrowFileWriter; import org.apache.arrow.vector.ipc.ArrowStreamReader; import org.apache.arrow.vector.ipc.ArrowStreamWriter; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.ipc.message.IpcOption; 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.ByteArrayReadableSeekableByteChannel; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; /** Test cases for {@link CompressionCodec}s. */ class TestCompressionCodec { private BufferAllocator allocator; @BeforeEach void init() { allocator = new RootAllocator(Integer.MAX_VALUE); } @AfterEach void terminate() { allocator.close(); } static Collection codecs() { List params = new ArrayList<>(); int[] lengths = new int[] {10, 100, 1000}; for (int len : lengths) { CompressionCodec dumbCodec = NoCompressionCodec.INSTANCE; params.add(Arguments.arguments(len, dumbCodec)); CompressionCodec lz4Codec = new Lz4CompressionCodec(); params.add(Arguments.arguments(len, lz4Codec)); CompressionCodec zstdCodec = new ZstdCompressionCodec(); params.add(Arguments.arguments(len, zstdCodec)); CompressionCodec zstdCodecAndCompressionLevel = new ZstdCompressionCodec(7); params.add(Arguments.arguments(len, zstdCodecAndCompressionLevel)); } return params; } private List compressBuffers(CompressionCodec codec, List inputBuffers) { List outputBuffers = new ArrayList<>(inputBuffers.size()); for (ArrowBuf buf : inputBuffers) { outputBuffers.add(codec.compress(allocator, buf)); } return outputBuffers; } private List deCompressBuffers(CompressionCodec codec, List inputBuffers) { List outputBuffers = new ArrayList<>(inputBuffers.size()); for (ArrowBuf buf : inputBuffers) { outputBuffers.add(codec.decompress(allocator, buf)); } return outputBuffers; } private void assertWriterIndex(List decompressedBuffers) { for (ArrowBuf decompressedBuf : decompressedBuffers) { assertTrue(decompressedBuf.writerIndex() > 0); } } @ParameterizedTest @MethodSource("codecs") void testCompressFixedWidthBuffers(int vectorLength, CompressionCodec codec) throws Exception { // prepare vector to compress IntVector origVec = new IntVector("vec", allocator); origVec.allocateNew(vectorLength); for (int i = 0; i < vectorLength; i++) { if (i % 10 == 0) { origVec.setNull(i); } else { origVec.set(i, i); } } origVec.setValueCount(vectorLength); int nullCount = origVec.getNullCount(); // compress & decompress List origBuffers = origVec.getFieldBuffers(); List compressedBuffers = compressBuffers(codec, origBuffers); List decompressedBuffers = deCompressBuffers(codec, compressedBuffers); assertEquals(2, decompressedBuffers.size()); assertWriterIndex(decompressedBuffers); // orchestrate new vector IntVector newVec = new IntVector("new vec", allocator); newVec.loadFieldBuffers(new ArrowFieldNode(vectorLength, nullCount), decompressedBuffers); // verify new vector assertEquals(vectorLength, newVec.getValueCount()); for (int i = 0; i < vectorLength; i++) { if (i % 10 == 0) { assertTrue(newVec.isNull(i)); } else { assertEquals(i, newVec.get(i)); } } newVec.close(); AutoCloseables.close(decompressedBuffers); } @ParameterizedTest @MethodSource("codecs") void testCompressVariableWidthBuffers(int vectorLength, CompressionCodec codec) throws Exception { // prepare vector to compress VarCharVector origVec = new VarCharVector("vec", allocator); origVec.allocateNew(); for (int i = 0; i < vectorLength; i++) { if (i % 10 == 0) { origVec.setNull(i); } else { origVec.setSafe(i, String.valueOf(i).getBytes(StandardCharsets.UTF_8)); } } origVec.setValueCount(vectorLength); int nullCount = origVec.getNullCount(); // compress & decompress List origBuffers = origVec.getFieldBuffers(); List compressedBuffers = compressBuffers(codec, origBuffers); List decompressedBuffers = deCompressBuffers(codec, compressedBuffers); assertEquals(3, decompressedBuffers.size()); assertWriterIndex(decompressedBuffers); // orchestrate new vector VarCharVector newVec = new VarCharVector("new vec", allocator); newVec.loadFieldBuffers(new ArrowFieldNode(vectorLength, nullCount), decompressedBuffers); // verify new vector assertEquals(vectorLength, newVec.getValueCount()); for (int i = 0; i < vectorLength; i++) { if (i % 10 == 0) { assertTrue(newVec.isNull(i)); } else { assertArrayEquals(String.valueOf(i).getBytes(StandardCharsets.UTF_8), newVec.get(i)); } } newVec.close(); AutoCloseables.close(decompressedBuffers); } @ParameterizedTest @MethodSource("codecs") void testEmptyBuffer(int vectorLength, CompressionCodec codec) throws Exception { final VarBinaryVector origVec = new VarBinaryVector("vec", allocator); origVec.allocateNew(vectorLength); // Do not set any values (all missing) origVec.setValueCount(vectorLength); final List origBuffers = origVec.getFieldBuffers(); final List compressedBuffers = compressBuffers(codec, origBuffers); final List decompressedBuffers = deCompressBuffers(codec, compressedBuffers); // orchestrate new vector VarBinaryVector newVec = new VarBinaryVector("new vec", allocator); newVec.loadFieldBuffers(new ArrowFieldNode(vectorLength, vectorLength), decompressedBuffers); // verify new vector assertEquals(vectorLength, newVec.getValueCount()); for (int i = 0; i < vectorLength; i++) { assertTrue(newVec.isNull(i)); } newVec.close(); AutoCloseables.close(decompressedBuffers); } private static Stream codecTypes() { return Arrays.stream(CompressionUtil.CodecType.values()); } @ParameterizedTest @MethodSource("codecTypes") void testReadWriteStream(CompressionUtil.CodecType codec) throws Exception { withRoot( codec, (factory, root) -> { ByteArrayOutputStream compressedStream = new ByteArrayOutputStream(); try (final ArrowStreamWriter writer = new ArrowStreamWriter( root, new DictionaryProvider.MapDictionaryProvider(), Channels.newChannel(compressedStream), IpcOption.DEFAULT, factory, codec, Optional.of(7))) { writer.start(); writer.writeBatch(); writer.end(); } catch (IOException e) { throw new RuntimeException(e); } try (ArrowStreamReader reader = new ArrowStreamReader( new ByteArrayReadableSeekableByteChannel(compressedStream.toByteArray()), allocator, factory)) { assertTrue(reader.loadNextBatch()); assertTrue(root.equals(reader.getVectorSchemaRoot())); assertFalse(reader.loadNextBatch()); } catch (IOException e) { throw new RuntimeException(e); } }); } @ParameterizedTest @MethodSource("codecTypes") void testReadWriteFile(CompressionUtil.CodecType codec) throws Exception { withRoot( codec, (factory, root) -> { ByteArrayOutputStream compressedStream = new ByteArrayOutputStream(); try (final ArrowFileWriter writer = new ArrowFileWriter( root, new DictionaryProvider.MapDictionaryProvider(), Channels.newChannel(compressedStream), new HashMap<>(), IpcOption.DEFAULT, factory, codec, Optional.of(7))) { writer.start(); writer.writeBatch(); writer.end(); } catch (IOException e) { throw new RuntimeException(e); } try (ArrowFileReader reader = new ArrowFileReader( new ByteArrayReadableSeekableByteChannel(compressedStream.toByteArray()), allocator, factory)) { assertTrue(reader.loadNextBatch()); assertTrue(root.equals(reader.getVectorSchemaRoot())); assertFalse(reader.loadNextBatch()); } catch (IOException e) { throw new RuntimeException(e); } }); } /** Unloading a vector should not free source buffers. */ @ParameterizedTest @MethodSource("codecTypes") void testUnloadCompressed(CompressionUtil.CodecType codec) { withRoot( codec, (factory, root) -> { root.getFieldVectors() .forEach( (vector) -> { Arrays.stream(vector.getBuffers(/*clear*/ false)) .forEach( (buf) -> { assertNotEquals(0, buf.getReferenceManager().getRefCount()); }); }); final VectorUnloader unloader = new VectorUnloader( root, /*includeNullCount*/ true, factory.createCodec(codec), /*alignBuffers*/ true); unloader.getRecordBatch().close(); root.getFieldVectors() .forEach( (vector) -> { Arrays.stream(vector.getBuffers(/*clear*/ false)) .forEach( (buf) -> { assertNotEquals(0, buf.getReferenceManager().getRefCount()); }); }); }); } void withRoot( CompressionUtil.CodecType codec, BiConsumer testBody) { final Schema schema = new Schema( Arrays.asList( Field.nullable("ints", new ArrowType.Int(32, true)), Field.nullable("strings", ArrowType.Utf8.INSTANCE))); CompressionCodec.Factory factory = codec == CompressionUtil.CodecType.NO_COMPRESSION ? NoCompressionCodec.Factory.INSTANCE : CommonsCompressionFactory.INSTANCE; try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final IntVector ints = (IntVector) root.getVector(0); final VarCharVector strings = (VarCharVector) root.getVector(1); // Doesn't get compressed ints.setSafe(0, 0x4a3e); ints.setSafe(1, 0x8aba); ints.setSafe(2, 0x4362); ints.setSafe(3, 0x383f); // Gets compressed String compressibleString = " "; // 16 bytes compressibleString = compressibleString + compressibleString; compressibleString = compressibleString + compressibleString; compressibleString = compressibleString + compressibleString; compressibleString = compressibleString + compressibleString; compressibleString = compressibleString + compressibleString; // 512 bytes byte[] compressibleData = compressibleString.getBytes(StandardCharsets.UTF_8); strings.setSafe(0, compressibleData); strings.setSafe(1, compressibleData); strings.setSafe(2, compressibleData); strings.setSafe(3, compressibleData); root.setRowCount(4); testBody.accept(factory, root); } } } TestCompressionCodecServiceProvider.java000066400000000000000000000042311515451615200375500ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/test/java/org/apache/arrow/compression/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compression; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertThrows; import org.apache.arrow.vector.compression.CompressionCodec; import org.apache.arrow.vector.compression.CompressionUtil; import org.apache.arrow.vector.compression.NoCompressionCodec; import org.junit.jupiter.api.Test; public class TestCompressionCodecServiceProvider { /** * When arrow-compression is in the classpath/module-path, {@link * CompressionCodec.Factory#INSTANCE} should be able to handle all codec types. */ @Test public void testSupportedCompressionTypes() { assertThrows( // no-compression doesn't support any actual compression types IllegalArgumentException.class, () -> checkAllCodecTypes(NoCompressionCodec.Factory.INSTANCE)); assertThrows( // commons-compression doesn't support the uncompressed type IllegalArgumentException.class, () -> checkAllCodecTypes(CommonsCompressionFactory.INSTANCE)); checkAllCodecTypes( // and the winner is... CompressionCodec.Factory.INSTANCE); // combines the two above to support all types } private void checkAllCodecTypes(CompressionCodec.Factory factory) { for (CompressionUtil.CodecType codecType : CompressionUtil.CodecType.values()) { assertNotNull(factory.createCodec(codecType)); } } } apache-arrow-java-19.0.0/compression/src/test/resources/000077500000000000000000000000001515451615200232175ustar00rootroot00000000000000apache-arrow-java-19.0.0/compression/src/test/resources/logback.xml000066400000000000000000000025221515451615200253440ustar00rootroot00000000000000 %d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n apache-arrow-java-19.0.0/dataset/000077500000000000000000000000001515451615200165235ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/CMakeLists.txt000066400000000000000000000052271515451615200212710ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. find_package(ArrowDataset REQUIRED) find_package(ArrowSubstrait REQUIRED) include_directories(${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_SOURCE_DIR} ${JNI_INCLUDE_DIRS} ${JNI_HEADERS_DIR}) add_jar(arrow_java_jni_dataset_jar src/main/java/org/apache/arrow/dataset/jni/JniLoader.java src/main/java/org/apache/arrow/dataset/jni/JniWrapper.java src/main/java/org/apache/arrow/dataset/file/JniWrapper.java src/main/java/org/apache/arrow/dataset/jni/NativeMemoryPool.java src/main/java/org/apache/arrow/dataset/jni/ReservationListener.java src/main/java/org/apache/arrow/dataset/substrait/JniWrapper.java GENERATE_NATIVE_HEADERS arrow_java_jni_dataset_headers) add_library(arrow_java_jni_dataset SHARED src/main/cpp/jni_wrapper.cc src/main/cpp/jni_util.cc) set_property(TARGET arrow_java_jni_dataset PROPERTY OUTPUT_NAME "arrow_dataset_jni") target_link_libraries(arrow_java_jni_dataset arrow_java_jni_dataset_headers jni ArrowDataset::arrow_dataset_static ArrowSubstrait::arrow_substrait_static) if(BUILD_TESTING) add_executable(arrow-java-jni-dataset-test src/main/cpp/jni_util_test.cc src/main/cpp/jni_util.cc) target_link_libraries(arrow-java-jni-dataset-test arrow_java_test) add_test(NAME arrow-java-jni-dataset-test COMMAND arrow-java-jni-dataset-test) endif() set(ARROW_JAVA_JNI_DATASET_LIBDIR "${CMAKE_INSTALL_PREFIX}/lib/arrow_dataset_jni/${ARROW_JAVA_JNI_ARCH_DIR}") set(ARROW_JAVA_JNI_DATASET_BINDIR "${CMAKE_INSTALL_PREFIX}/bin/arrow_dataset_jni/${ARROW_JAVA_JNI_ARCH_DIR}") install(TARGETS arrow_java_jni_dataset LIBRARY DESTINATION ${ARROW_JAVA_JNI_DATASET_LIBDIR} RUNTIME DESTINATION ${ARROW_JAVA_JNI_DATASET_BINDIR}) apache-arrow-java-19.0.0/dataset/pom.xml000066400000000000000000000157171515451615200200530ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-java-root 19.0.0 arrow-dataset jar Arrow Java Dataset Java implementation of Arrow Dataset API/Framework ../../../cpp/release-build/ 1.17.0 1.12.1 org.apache.arrow arrow-vector ${arrow.vector.classifier} compile org.apache.arrow arrow-memory-core compile org.apache.arrow arrow-c-data compile org.immutables value-annotations org.apache.arrow arrow-memory-netty test org.apache.parquet parquet-avro ${parquet.version} test org.apache.avro avro ${avro.version} test org.apache.parquet parquet-hadoop ${parquet.version} test org.apache.arrow arrow-vector ${project.version} tests test org.apache.hadoop hadoop-common ${dep.hadoop.version} test commons-logging commons-logging javax.servlet servlet-api log4j log4j org.slf4j slf4j-log4j12 com.google.guava guava test com.fasterxml.jackson.core jackson-databind test org.apache.arrow.orc arrow-orc ${project.version} test org.apache.orc orc-core 2.3.0 test log4j log4j org.slf4j slf4j-log4j12 commons-logging commons-logging org.apache.hive hive-storage-api 4.0.1 test commons-io commons-io 2.21.0 test org.hamcrest hamcrest test ${arrow.cpp.build.dir} **/*arrow_dataset_jni.* maven-surefire-plugin --add-reads=org.apache.arrow.dataset=com.fasterxml.jackson.databind --add-opens=java.base/java.nio=org.apache.arrow.dataset,org.apache.arrow.memory.core,ALL-UNNAMED false ${project.basedir}/../testing/data org.xolstice.maven.plugins protobuf-maven-plugin src compile ../../cpp/src/jni/dataset/proto org.apache.maven.plugins maven-compiler-plugin -Werror apache-arrow-java-19.0.0/dataset/src/000077500000000000000000000000001515451615200173125ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/000077500000000000000000000000001515451615200202365ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/cpp/000077500000000000000000000000001515451615200210205ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/cpp/jni_util.cc000066400000000000000000000320351515451615200231470ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #include "jni_util.h" #include #include #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" #include "arrow/util/logging.h" namespace arrow { namespace dataset { namespace jni { jint JNI_VERSION = JNI_VERSION_10; class ReservationListenableMemoryPool::Impl { public: explicit Impl(arrow::MemoryPool* pool, std::shared_ptr listener, int64_t block_size) : pool_(pool), listener_(listener), block_size_(block_size), blocks_reserved_(0) {} arrow::Status Allocate(int64_t size, int64_t alignment, uint8_t** out) { RETURN_NOT_OK(UpdateReservation(size)); arrow::Status error = pool_->Allocate(size, alignment, out); if (!error.ok()) { RETURN_NOT_OK(UpdateReservation(-size)); return error; } return arrow::Status::OK(); } arrow::Status Reallocate(int64_t old_size, int64_t new_size, int64_t alignment, uint8_t** ptr) { bool reserved = false; int64_t diff = new_size - old_size; if (new_size >= old_size) { // new_size >= old_size, pre-reserve bytes from listener before allocating // from underlying pool RETURN_NOT_OK(UpdateReservation(diff)); reserved = true; } arrow::Status error = pool_->Reallocate(old_size, new_size, alignment, ptr); if (!error.ok()) { if (reserved) { // roll back reservations on error RETURN_NOT_OK(UpdateReservation(-diff)); } return error; } if (!reserved) { // otherwise (e.g. new_size < old_size), make updates after calling underlying pool RETURN_NOT_OK(UpdateReservation(diff)); } return arrow::Status::OK(); } void Free(uint8_t* buffer, int64_t size, int64_t alignment) { pool_->Free(buffer, size, alignment); // FIXME: See ARROW-11143, currently method ::Free doesn't allow Status return arrow::Status s = UpdateReservation(-size); if (!s.ok()) { ARROW_LOG(FATAL) << "Failed to update reservation while freeing bytes: " << s.message(); return; } } arrow::Status UpdateReservation(int64_t diff) { int64_t granted = Reserve(diff); if (granted == 0) { return arrow::Status::OK(); } if (granted < 0) { RETURN_NOT_OK(listener_->OnRelease(-granted)); return arrow::Status::OK(); } RETURN_NOT_OK(listener_->OnReservation(granted)); return arrow::Status::OK(); } int64_t Reserve(int64_t diff) { std::lock_guard lock(mutex_); if (diff > 0) { stats_.DidAllocateBytes(diff); } else if (diff < 0) { stats_.DidFreeBytes(-diff); } int64_t new_block_count; int64_t bytes_reserved = stats_.bytes_allocated(); if (bytes_reserved == 0) { new_block_count = 0; } else { // ceil to get the required block number new_block_count = (bytes_reserved - 1) / block_size_ + 1; } int64_t bytes_granted = (new_block_count - blocks_reserved_) * block_size_; blocks_reserved_ = new_block_count; return bytes_granted; } int64_t bytes_allocated() { return stats_.bytes_allocated(); } int64_t max_memory() { return stats_.max_memory(); } int64_t total_bytes_allocated() { return stats_.total_bytes_allocated(); } int64_t num_allocations() { return stats_.num_allocations(); } std::string backend_name() { return pool_->backend_name(); } std::shared_ptr get_listener() { return listener_; } private: arrow::MemoryPool* pool_; std::shared_ptr listener_; int64_t block_size_; int64_t blocks_reserved_; arrow::internal::MemoryPoolStats stats_; std::mutex mutex_; }; ReservationListenableMemoryPool::ReservationListenableMemoryPool( MemoryPool* pool, std::shared_ptr listener, int64_t block_size) { impl_.reset(new Impl(pool, listener, block_size)); } arrow::Status ReservationListenableMemoryPool::Allocate(int64_t size, int64_t alignment, uint8_t** out) { return impl_->Allocate(size, alignment, out); } arrow::Status ReservationListenableMemoryPool::Reallocate(int64_t old_size, int64_t new_size, int64_t alignment, uint8_t** ptr) { return impl_->Reallocate(old_size, new_size, alignment, ptr); } void ReservationListenableMemoryPool::Free(uint8_t* buffer, int64_t size, int64_t alignment) { return impl_->Free(buffer, size, alignment); } int64_t ReservationListenableMemoryPool::bytes_allocated() const { return impl_->bytes_allocated(); } int64_t ReservationListenableMemoryPool::max_memory() const { return impl_->max_memory(); } int64_t ReservationListenableMemoryPool::total_bytes_allocated() const { return impl_->total_bytes_allocated(); } int64_t ReservationListenableMemoryPool::num_allocations() const { return impl_->num_allocations(); } std::string ReservationListenableMemoryPool::backend_name() const { return impl_->backend_name(); } std::shared_ptr ReservationListenableMemoryPool::get_listener() { return impl_->get_listener(); } ReservationListenableMemoryPool::~ReservationListenableMemoryPool() {} std::string Describe(JNIEnv* env, jthrowable t) { jclass describer_class = env->FindClass("org/apache/arrow/dataset/jni/JniExceptionDescriber"); ARROW_DCHECK_NE(describer_class, nullptr); jmethodID describe_method = env->GetStaticMethodID( describer_class, "describe", "(Ljava/lang/Throwable;)Ljava/lang/String;"); std::string description = JStringToCString( env, (jstring)env->CallStaticObjectMethod(describer_class, describe_method, t)); return description; } bool IsErrorInstanceOf(JNIEnv* env, jthrowable t, std::string class_name) { jclass java_class = env->FindClass(class_name.c_str()); ARROW_DCHECK_NE(java_class, nullptr) << "Could not find Java class " << class_name; return env->IsInstanceOf(t, java_class); } arrow::StatusCode MapJavaError(JNIEnv* env, jthrowable t) { StatusCode code; if (IsErrorInstanceOf(env, t, "org/apache/arrow/memory/OutOfMemoryException")) { code = StatusCode::OutOfMemory; } else if (IsErrorInstanceOf(env, t, "java/lang/UnsupportedOperationException")) { code = StatusCode::NotImplemented; } else if (IsErrorInstanceOf(env, t, "java/io/NotSerializableException")) { code = StatusCode::SerializationError; } else if (IsErrorInstanceOf(env, t, "java/io/IOException")) { code = StatusCode::IOError; } else if (IsErrorInstanceOf(env, t, "java/lang/IllegalArgumentException")) { code = StatusCode::Invalid; } else if (IsErrorInstanceOf(env, t, "java/lang/IllegalStateException")) { code = StatusCode::Invalid; } else { code = StatusCode::UnknownError; } return code; } JNIEnv* GetEnvOrAttach(JavaVM* vm) { JNIEnv* env; int getEnvStat = vm->GetEnv(reinterpret_cast(&env), JNI_VERSION); if (getEnvStat == JNI_EDETACHED) { // Reattach current thread to JVM getEnvStat = vm->AttachCurrentThread(reinterpret_cast(&env), nullptr); if (getEnvStat != JNI_OK) { ARROW_LOG(FATAL) << "Failed to attach current thread to JVM"; } } return env; } const char kJavaErrorDetailTypeId[] = "arrow::dataset::jni::JavaErrorDetail"; JavaErrorDetail::JavaErrorDetail(JavaVM* vm, jthrowable cause) : vm_(vm) { JNIEnv* env = GetEnvOrAttach(vm_); if (env == nullptr) { this->cause_ = nullptr; return; } this->cause_ = (jthrowable)env->NewGlobalRef(cause); } JavaErrorDetail::~JavaErrorDetail() { JNIEnv* env = GetEnvOrAttach(vm_); if (env == nullptr || this->cause_ == nullptr) { return; } env->DeleteGlobalRef(cause_); } const char* JavaErrorDetail::type_id() const { return kJavaErrorDetailTypeId; } jthrowable JavaErrorDetail::GetCause() const { JNIEnv* env = GetEnvOrAttach(vm_); if (env == nullptr || this->cause_ == nullptr) { return nullptr; } return (jthrowable)env->NewLocalRef(cause_); } std::string JavaErrorDetail::ToString() const { JNIEnv* env = GetEnvOrAttach(vm_); if (env == nullptr) { return "Java Exception, ID: " + std::to_string(reinterpret_cast(cause_)); } return "Java Exception: " + Describe(env, cause_); } Status CheckException(JNIEnv* env) { if (env->ExceptionCheck()) { jthrowable t = env->ExceptionOccurred(); env->ExceptionClear(); arrow::StatusCode code = MapJavaError(env, t); JavaVM* vm; if (env->GetJavaVM(&vm) != JNI_OK) { return Status::Invalid("Error getting JavaVM object"); } std::shared_ptr detail = std::make_shared(vm, t); return {code, detail->ToString(), detail}; } return Status::OK(); } jclass CreateGlobalClassReference(JNIEnv* env, const char* class_name) { jclass local_class = env->FindClass(class_name); jclass global_class = (jclass)env->NewGlobalRef(local_class); env->DeleteLocalRef(local_class); return global_class; } arrow::Result GetMethodID(JNIEnv* env, jclass this_class, const char* name, const char* sig) { jmethodID ret = env->GetMethodID(this_class, name, sig); if (ret == nullptr) { std::string error_message = "Unable to find method " + std::string(name) + " within signature" + std::string(sig); return arrow::Status::Invalid(error_message); } return ret; } arrow::Result GetStaticMethodID(JNIEnv* env, jclass this_class, const char* name, const char* sig) { jmethodID ret = env->GetStaticMethodID(this_class, name, sig); if (ret == nullptr) { std::string error_message = "Unable to find static method " + std::string(name) + " within signature" + std::string(sig); return arrow::Status::Invalid(error_message); } return ret; } std::string JStringToCString(JNIEnv* env, jstring string) { if (string == nullptr) { return std::string(); } const char* chars = env->GetStringUTFChars(string, nullptr); std::string ret(chars); env->ReleaseStringUTFChars(string, chars); return ret; } std::vector ToStringVector(JNIEnv* env, jobjectArray& str_array) { int length = env->GetArrayLength(str_array); std::vector vector; for (int i = 0; i < length; i++) { auto string = reinterpret_cast(env->GetObjectArrayElement(str_array, i)); vector.push_back(JStringToCString(env, string)); } return vector; } arrow::Result ToSchemaByteArray(JNIEnv* env, std::shared_ptr schema) { ARROW_ASSIGN_OR_RAISE( std::shared_ptr buffer, arrow::ipc::SerializeSchema(*schema, arrow::default_memory_pool())) jbyteArray out = env->NewByteArray(static_cast(buffer->size())); auto src = reinterpret_cast(buffer->data()); env->SetByteArrayRegion(out, 0, static_cast(buffer->size()), src); return out; } arrow::Result> FromSchemaByteArray( JNIEnv* env, jbyteArray schemaBytes) { arrow::ipc::DictionaryMemo in_memo; int schemaBytes_len = env->GetArrayLength(schemaBytes); jbyte* schemaBytes_data = env->GetByteArrayElements(schemaBytes, nullptr); auto serialized_schema = std::make_shared( reinterpret_cast(schemaBytes_data), schemaBytes_len); arrow::io::BufferReader buf_reader(serialized_schema); ARROW_ASSIGN_OR_RAISE(std::shared_ptr schema, arrow::ipc::ReadSchema(&buf_reader, &in_memo)) env->ReleaseByteArrayElements(schemaBytes, schemaBytes_data, JNI_ABORT); return schema; } arrow::Status ExportRecordBatch(JNIEnv* env, const std::shared_ptr& batch, jlong struct_array) { return arrow::ExportRecordBatch(*batch, reinterpret_cast(struct_array)); } arrow::Result> ImportRecordBatch( JNIEnv* env, const std::shared_ptr& schema, jlong struct_array) { return arrow::ImportRecordBatch(reinterpret_cast(struct_array), schema); } } // namespace jni } // namespace dataset } // namespace arrow apache-arrow-java-19.0.0/dataset/src/main/cpp/jni_util.h000066400000000000000000000143331515451615200230120ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #pragma once #include #include "arrow/array.h" #include "arrow/io/api.h" #include "arrow/ipc/api.h" #include "arrow/memory_pool.h" #include "arrow/result.h" #include "arrow/type.h" namespace arrow { namespace dataset { namespace jni { JNIEnv* GetEnvOrAttach(JavaVM* vm); Status CheckException(JNIEnv* env); jclass CreateGlobalClassReference(JNIEnv* env, const char* class_name); arrow::Result GetMethodID(JNIEnv* env, jclass this_class, const char* name, const char* sig); arrow::Result GetStaticMethodID(JNIEnv* env, jclass this_class, const char* name, const char* sig); std::string JStringToCString(JNIEnv* env, jstring string); std::vector ToStringVector(JNIEnv* env, jobjectArray& str_array); arrow::Result ToSchemaByteArray(JNIEnv* env, std::shared_ptr schema); arrow::Result> FromSchemaByteArray(JNIEnv* env, jbyteArray schemaBytes); /// \brief Export arrow::RecordBatch for Java (or other JVM languages) use. /// The exported batch is subject to C data interface specification and can be /// imported from Java side using provided JNI utilities. arrow::Status ExportRecordBatch(JNIEnv* env, const std::shared_ptr& batch, jlong struct_array); /// \brief Import arrow::RecordBatch from JVM language side. The input data should /// ideally be exported from specific JNI utilities from JVM language side and should /// conform to C data interface specification. arrow::Result> ImportRecordBatch( JNIEnv* env, const std::shared_ptr& schema, jlong struct_array); /// \brief Create a new shared_ptr on heap from shared_ptr t to prevent /// the managed object from being garbage-collected. /// /// \return address of the newly created shared pointer template jlong CreateNativeRef(std::shared_ptr t) { std::shared_ptr* retained_ptr = new std::shared_ptr(t); return reinterpret_cast(retained_ptr); } /// \brief Get the shared_ptr that was derived via function CreateNativeRef. /// /// \param[in] ref address of the shared_ptr /// \return the shared_ptr object template std::shared_ptr RetrieveNativeInstance(jlong ref) { std::shared_ptr* retrieved_ptr = reinterpret_cast*>(ref); return *retrieved_ptr; } /// \brief Destroy a shared_ptr using its memory address. /// /// \param[in] ref address of the shared_ptr template void ReleaseNativeRef(jlong ref) { std::shared_ptr* retrieved_ptr = reinterpret_cast*>(ref); delete retrieved_ptr; } // Indicate an exception thrown during calling Java method via JNI. // Not thread safe. class JavaErrorDetail : public StatusDetail { public: JavaErrorDetail(JavaVM* vm, jthrowable cause); virtual ~JavaErrorDetail(); const char* type_id() const override; std::string ToString() const override; jthrowable GetCause() const; private: JavaVM* vm_; jthrowable cause_; }; /// Listener to act on reservations/unreservations from ReservationListenableMemoryPool. /// /// Note the memory pool will call this listener only on block-level memory /// reservation/unreservation is granted. So the invocation parameter "size" is always /// multiple of block size (by default, 512k) specified in memory pool. class ReservationListener { public: virtual ~ReservationListener() = default; virtual arrow::Status OnReservation(int64_t size) = 0; virtual arrow::Status OnRelease(int64_t size) = 0; protected: ReservationListener() = default; }; /// A memory pool implementation for pre-reserving memory blocks from a /// customizable listener. This will typically be used when memory allocations /// have to be subject to another "virtual" resource manager, which just tracks or /// limits number of bytes of application's overall memory usage. The underlying /// memory pool will still be responsible for actual malloc/free operations. class ReservationListenableMemoryPool : public arrow::MemoryPool { public: /// \brief Constructor. /// /// \param[in] pool the underlying memory pool /// \param[in] listener a listener for block-level reservations/releases. /// \param[in] block_size size of each block to reserve from the listener explicit ReservationListenableMemoryPool(MemoryPool* pool, std::shared_ptr listener, int64_t block_size = 512 * 1024); ~ReservationListenableMemoryPool(); using MemoryPool::Allocate; using MemoryPool::Free; using MemoryPool::Reallocate; arrow::Status Allocate(int64_t size, int64_t alignment, uint8_t** out) override; arrow::Status Reallocate(int64_t old_size, int64_t new_size, int64_t alignment, uint8_t** ptr) override; void Free(uint8_t* buffer, int64_t size, int64_t alignment) override; int64_t bytes_allocated() const override; int64_t max_memory() const override; int64_t total_bytes_allocated() const override; int64_t num_allocations() const override; std::string backend_name() const override; std::shared_ptr get_listener(); private: class Impl; std::unique_ptr impl_; }; } // namespace jni } // namespace dataset } // namespace arrow apache-arrow-java-19.0.0/dataset/src/main/cpp/jni_util_test.cc000066400000000000000000000072561515451615200242150ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #include "jni_util.h" #include #include "arrow/memory_pool.h" #include "arrow/testing/gtest_util.h" namespace arrow { namespace dataset { namespace jni { class MyListener : public ReservationListener { public: Status OnReservation(int64_t size) override { bytes_reserved_ += size; reservation_count_++; return arrow::Status::OK(); } Status OnRelease(int64_t size) override { bytes_reserved_ -= size; release_count_++; return arrow::Status::OK(); } int64_t bytes_reserved() { return bytes_reserved_; } int32_t reservation_count() const { return reservation_count_; } int32_t release_count() const { return release_count_; } private: int64_t bytes_reserved_; int32_t reservation_count_; int32_t release_count_; }; TEST(ReservationListenableMemoryPool, Basic) { auto pool = MemoryPool::CreateDefault(); auto listener = std::make_shared(); ReservationListenableMemoryPool rlp(pool.get(), listener); uint8_t* data; ASSERT_OK(rlp.Allocate(100, &data)); uint8_t* data2; ASSERT_OK(rlp.Allocate(100, &data2)); rlp.Free(data, 100); rlp.Free(data2, 100); ASSERT_EQ(200, rlp.max_memory()); ASSERT_EQ(200, pool->max_memory()); } TEST(ReservationListenableMemoryPool, Listener) { auto pool = MemoryPool::CreateDefault(); auto listener = std::make_shared(); ReservationListenableMemoryPool rlp(pool.get(), listener); uint8_t* data; ASSERT_OK(rlp.Allocate(100, &data)); uint8_t* data2; ASSERT_OK(rlp.Allocate(100, &data2)); ASSERT_EQ(200, rlp.bytes_allocated()); ASSERT_EQ(512 * 1024, listener->bytes_reserved()); rlp.Free(data, 100); rlp.Free(data2, 100); ASSERT_EQ(0, rlp.bytes_allocated()); ASSERT_EQ(0, listener->bytes_reserved()); ASSERT_EQ(1, listener->reservation_count()); ASSERT_EQ(1, listener->release_count()); } TEST(ReservationListenableMemoryPool, BlockSize) { auto pool = MemoryPool::CreateDefault(); auto listener = std::make_shared(); ReservationListenableMemoryPool rlp(pool.get(), listener, 100); uint8_t* data; ASSERT_OK(rlp.Allocate(100, &data)); ASSERT_EQ(100, rlp.bytes_allocated()); ASSERT_EQ(100, listener->bytes_reserved()); rlp.Free(data, 100); ASSERT_EQ(0, rlp.bytes_allocated()); ASSERT_EQ(0, listener->bytes_reserved()); } TEST(ReservationListenableMemoryPool, BlockSize2) { auto pool = MemoryPool::CreateDefault(); auto listener = std::make_shared(); ReservationListenableMemoryPool rlp(pool.get(), listener, 99); uint8_t* data; ASSERT_OK(rlp.Allocate(100, &data)); ASSERT_EQ(100, rlp.bytes_allocated()); ASSERT_EQ(198, listener->bytes_reserved()); rlp.Free(data, 100); ASSERT_EQ(0, rlp.bytes_allocated()); ASSERT_EQ(0, listener->bytes_reserved()); ASSERT_EQ(1, listener->reservation_count()); ASSERT_EQ(1, listener->release_count()); } } // namespace jni } // namespace dataset } // namespace arrow apache-arrow-java-19.0.0/dataset/src/main/cpp/jni_wrapper.cc000066400000000000000000001146521515451615200236600ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #include #include #include #include "arrow/array.h" #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 #include "arrow/dataset/file_csv.h" #endif #include "arrow/filesystem/api.h" #include "arrow/filesystem/path_util.h" #include "arrow/engine/substrait/util.h" #include "arrow/engine/substrait/serde.h" #include "arrow/engine/substrait/relation.h" #include "arrow/ipc/api.h" #include "arrow/util/iterator.h" #include "jni_util.h" #include "org_apache_arrow_dataset_file_JniWrapper.h" #include "org_apache_arrow_dataset_jni_JniWrapper.h" #include "org_apache_arrow_dataset_jni_NativeMemoryPool.h" #include "org_apache_arrow_dataset_substrait_JniWrapper.h" namespace { jclass illegal_access_exception_class; jclass illegal_argument_exception_class; jclass runtime_exception_class; jclass java_reservation_listener_class; jmethodID reserve_memory_method; jmethodID unreserve_memory_method; jlong default_memory_pool_id = -1L; jint JNI_VERSION = JNI_VERSION_10; class JniPendingException : public std::runtime_error { public: explicit JniPendingException(const std::string& arg, jthrowable cause) : runtime_error(arg), cause_(cause) {} jthrowable GetCause() const { return cause_; } bool HasCause() const { return cause_ != nullptr; } private: jthrowable cause_; }; void ThrowPendingException(const std::string& message, jthrowable cause = nullptr) { throw JniPendingException(message, cause); } void ThrowIfError(const arrow::Status& status) { const std::shared_ptr& detail = status.detail(); const std::shared_ptr& maybe_java = std::dynamic_pointer_cast(detail); if (maybe_java != nullptr) { ThrowPendingException(status.message(), maybe_java->GetCause()); return; } if (!status.ok()) { ThrowPendingException(status.message()); } } class JNIEnvGuard { public: explicit JNIEnvGuard(JavaVM* vm) : vm_(vm), env_(nullptr), should_detach_(false) { JNIEnv* env; jint code = vm->GetEnv(reinterpret_cast(&env), JNI_VERSION); if (code == JNI_EDETACHED) { JavaVMAttachArgs args; args.version = JNI_VERSION; args.name = NULL; args.group = NULL; code = vm->AttachCurrentThread(reinterpret_cast(&env), &args); should_detach_ = (code == JNI_OK); } if (code != JNI_OK) { ThrowPendingException("Failed to attach the current thread to a Java VM"); } env_ = env; } JNIEnv* env() { return env_; } ~JNIEnvGuard() { if (should_detach_) { vm_->DetachCurrentThread(); should_detach_ = false; } } private: JavaVM* vm_; JNIEnv* env_; bool should_detach_; }; template T JniGetOrThrow(arrow::Result result) { const arrow::Status& status = result.status(); ThrowIfError(status); return std::move(result).ValueOrDie(); } void JniAssertOkOrThrow(arrow::Status status) { ThrowIfError(status); } void JniThrow(std::string message) { ThrowPendingException(message); } arrow::Result> GetFileFormat( jint file_format_id) { switch (file_format_id) { case 0: return std::make_shared(); case 1: return std::make_shared(); #ifdef ARROW_ORC case 2: return std::make_shared(); #endif #ifdef ARROW_CSV case 3: return std::make_shared(); #endif #ifdef ARROW_JSON case 4: return std::make_shared(); #endif default: std::string error_message = "illegal file format id: " + std::to_string(file_format_id); return arrow::Status::Invalid(error_message); } } class ReserveFromJava : public arrow::dataset::jni::ReservationListener { public: ReserveFromJava(JavaVM* vm, jobject java_reservation_listener) : vm_(vm), java_reservation_listener_(java_reservation_listener) {} arrow::Status OnReservation(int64_t size) override { try { JNIEnvGuard guard(vm_); JNIEnv* env = guard.env(); env->CallObjectMethod(java_reservation_listener_, reserve_memory_method, size); RETURN_NOT_OK(arrow::dataset::jni::CheckException(env)); return arrow::Status::OK(); } catch (const JniPendingException& e) { return arrow::Status::Invalid(e.what()); } } arrow::Status OnRelease(int64_t size) override { try { JNIEnvGuard guard(vm_); JNIEnv* env = guard.env(); env->CallObjectMethod(java_reservation_listener_, unreserve_memory_method, size); RETURN_NOT_OK(arrow::dataset::jni::CheckException(env)); return arrow::Status::OK(); } catch (const JniPendingException& e) { return arrow::Status::Invalid(e.what()); } } jobject GetJavaReservationListener() { return java_reservation_listener_; } private: JavaVM* vm_; jobject java_reservation_listener_; }; /// \class DisposableScannerAdaptor /// \brief An adaptor that iterates over a Scanner instance then returns RecordBatches /// directly. /// /// This lessens the complexity of the JNI bridge to make sure it to be easier to /// maintain. On Java-side, NativeScanner can only produces a single NativeScanTask /// instance during its whole lifecycle. Each task stands for a DisposableScannerAdaptor /// instance through JNI bridge. /// class DisposableScannerAdaptor { public: DisposableScannerAdaptor(std::shared_ptr scanner, arrow::dataset::TaggedRecordBatchIterator batch_itr) : scanner_(std::move(scanner)), batch_itr_(std::move(batch_itr)) {} static arrow::Result> Create( std::shared_ptr scanner) { ARROW_ASSIGN_OR_RAISE(auto batch_itr, scanner->ScanBatches()); return std::make_shared(scanner, std::move(batch_itr)); } arrow::Result> Next() { ARROW_ASSIGN_OR_RAISE(std::shared_ptr batch, NextBatch()); return batch; } const std::shared_ptr& GetScanner() const { return scanner_; } private: std::shared_ptr scanner_; arrow::dataset::TaggedRecordBatchIterator batch_itr_; arrow::Result> NextBatch() { ARROW_ASSIGN_OR_RAISE(auto batch, batch_itr_.Next()) return batch.record_batch; } }; arrow::Result> SchemaFromColumnNames( const std::shared_ptr& input, const std::vector& column_names) { std::vector> columns; for (arrow::FieldRef ref : column_names) { auto maybe_field = ref.GetOne(*input); if (maybe_field.ok()) { columns.push_back(std::move(maybe_field).ValueOrDie()); } else { return arrow::Status::Invalid("Partition column '", ref.ToString(), "' is not in dataset schema"); } } return schema(std::move(columns))->WithMetadata(input->metadata()); } } // namespace using arrow::dataset::jni::CreateGlobalClassReference; using arrow::dataset::jni::CreateNativeRef; using arrow::dataset::jni::FromSchemaByteArray; using arrow::dataset::jni::GetMethodID; using arrow::dataset::jni::JStringToCString; using arrow::dataset::jni::ReleaseNativeRef; using arrow::dataset::jni::RetrieveNativeInstance; using arrow::dataset::jni::ToSchemaByteArray; using arrow::dataset::jni::ToStringVector; using arrow::dataset::jni::ReservationListenableMemoryPool; using arrow::dataset::jni::ReservationListener; #define JNI_METHOD_START try { // macro ended #define JNI_METHOD_END(fallback_expr) \ } \ catch (JniPendingException & e) { \ if (e.HasCause()) { \ env->Throw(e.GetCause()); \ return fallback_expr; \ } \ env->ThrowNew(runtime_exception_class, e.what()); \ return fallback_expr; \ } // macro ended jint JNI_OnLoad(JavaVM* vm, void* reserved) { JNIEnv* env; if (vm->GetEnv(reinterpret_cast(&env), JNI_VERSION) != JNI_OK) { return JNI_ERR; } JNI_METHOD_START illegal_access_exception_class = CreateGlobalClassReference(env, "Ljava/lang/IllegalAccessException;"); illegal_argument_exception_class = CreateGlobalClassReference(env, "Ljava/lang/IllegalArgumentException;"); runtime_exception_class = CreateGlobalClassReference(env, "Ljava/lang/RuntimeException;"); java_reservation_listener_class = CreateGlobalClassReference(env, "Lorg/apache/arrow/" "dataset/jni/ReservationListener;"); reserve_memory_method = JniGetOrThrow(GetMethodID(env, java_reservation_listener_class, "reserve", "(J)V")); unreserve_memory_method = JniGetOrThrow( GetMethodID(env, java_reservation_listener_class, "unreserve", "(J)V")); default_memory_pool_id = reinterpret_cast(arrow::default_memory_pool()); return JNI_VERSION; JNI_METHOD_END(JNI_ERR) } void JNI_OnUnload(JavaVM* vm, void* reserved) { JNIEnv* env; vm->GetEnv(reinterpret_cast(&env), JNI_VERSION); env->DeleteGlobalRef(illegal_access_exception_class); env->DeleteGlobalRef(illegal_argument_exception_class); env->DeleteGlobalRef(runtime_exception_class); env->DeleteGlobalRef(java_reservation_listener_class); default_memory_pool_id = -1L; } /// Unpack the named tables passed through JNI. /// /// Named tables are encoded as a string array, where every two elements /// encode (1) the table name and (2) the address of an ArrowArrayStream /// containing the table data. This function will eagerly read all /// tables into Tables. std::unordered_map> LoadNamedTables(JNIEnv* env, const jobjectArray& str_array) { std::unordered_map> map_table_to_record_batch_reader; int length = env->GetArrayLength(str_array); if (length % 2 != 0) { JniThrow("Cannot map odd number of array elements to key/value pairs"); } std::shared_ptr output_table; for (int pos = 0; pos < length; pos++) { auto j_string_key = reinterpret_cast(env->GetObjectArrayElement(str_array, pos)); pos++; auto j_string_value = reinterpret_cast(env->GetObjectArrayElement(str_array, pos)); uintptr_t memory_address = 0; try { memory_address = std::stol(JStringToCString(env, j_string_value)); } catch(const std::exception& ex) { JniThrow("Failed to parse memory address from string value. Error: " + std::string(ex.what())); } catch (...) { JniThrow("Failed to parse memory address from string value."); } auto* arrow_stream_in = reinterpret_cast(memory_address); std::shared_ptr readerIn = JniGetOrThrow(arrow::ImportRecordBatchReader(arrow_stream_in)); output_table = JniGetOrThrow(readerIn->ToTable()); map_table_to_record_batch_reader[JStringToCString(env, j_string_key)] = output_table; } return map_table_to_record_batch_reader; } /// Find the arrow Table associated with a given table name std::shared_ptr GetTableByName(const std::vector& names, const std::unordered_map>& tables) { if (names.size() != 1) { JniThrow("Tables with hierarchical names are not supported"); } const auto& it = tables.find(names[0]); if (it == tables.end()) { JniThrow("Table is referenced, but not provided: " + names[0]); } return it->second; } std::shared_ptr LoadArrowBufferFromByteBuffer(JNIEnv* env, jobject byte_buffer) { const auto *buff = reinterpret_cast(env->GetDirectBufferAddress(byte_buffer)); int length = env->GetDirectBufferCapacity(byte_buffer); std::shared_ptr buffer = JniGetOrThrow(arrow::AllocateBuffer(length)); std::memcpy(buffer->mutable_data(), buff, length); return buffer; } inline bool ParseBool(const std::string& value) { return value == "true" ? true : false; } inline char ParseChar(const std::string& key, const std::string& value) { if (value.size() != 1) { JniThrow("Option " + key + " should be a char, but is " + value); } return value.at(0); } /// \brief Construct FragmentScanOptions from config map #ifdef ARROW_CSV bool SetCsvConvertOptions(arrow::csv::ConvertOptions& options, const std::string& key, const std::string& value) { if (key == "column_types") { int64_t schema_address = std::stol(value); ArrowSchema* c_schema = reinterpret_cast(schema_address); auto schema = JniGetOrThrow(arrow::ImportSchema(c_schema)); auto& column_types = options.column_types; for (auto field : schema->fields()) { column_types[field->name()] = field->type(); } } else if (key == "strings_can_be_null") { options.strings_can_be_null = ParseBool(value); } else if (key == "check_utf8") { options.check_utf8 = ParseBool(value); } else if (key == "null_values") { options.null_values = {value}; } else if (key == "true_values") { options.true_values = {value}; } else if (key == "false_values") { options.false_values = {value}; } else if (key == "quoted_strings_can_be_null") { options.quoted_strings_can_be_null = ParseBool(value); } else if (key == "auto_dict_encode") { options.auto_dict_encode = ParseBool(value); } else if (key == "auto_dict_max_cardinality") { options.auto_dict_max_cardinality = std::stoi(value); } else if (key == "decimal_point") { options.decimal_point = ParseChar(key, value); } else if (key == "include_missing_columns") { options.include_missing_columns = ParseBool(value); } else { return false; } return true; } bool SetCsvParseOptions(arrow::csv::ParseOptions& options, const std::string& key, const std::string& value) { if (key == "delimiter") { options.delimiter = ParseChar(key, value); } else if (key == "quoting") { options.quoting = ParseBool(value); } else if (key == "quote_char") { options.quote_char = ParseChar(key, value); } else if (key == "double_quote") { options.double_quote = ParseBool(value); } else if (key == "escaping") { options.escaping = ParseBool(value); } else if (key == "escape_char") { options.escape_char = ParseChar(key, value); } else if (key == "newlines_in_values") { options.newlines_in_values = ParseBool(value); } else if (key == "ignore_empty_lines") { options.ignore_empty_lines = ParseBool(value); } else { return false; } return true; } bool SetCsvReadOptions(arrow::csv::ReadOptions& options, const std::string& key, const std::string& value) { if (key == "use_threads") { options.use_threads = ParseBool(value); } else if (key == "block_size") { options.block_size = std::stoi(value); } else if (key == "skip_rows") { options.skip_rows = std::stoi(value); } else if (key == "skip_rows_after_names") { options.skip_rows_after_names = std::stoi(value); } else if (key == "autogenerate_column_names") { options.autogenerate_column_names = ParseBool(value); } else { return false; } return true; } std::shared_ptr ToCsvFragmentScanOptions( const std::unordered_map& configs) { std::shared_ptr options = std::make_shared(); for (const auto& [key, value] : configs) { bool setValid = SetCsvParseOptions(options->parse_options, key, value) || SetCsvConvertOptions(options->convert_options, key, value) || SetCsvReadOptions(options->read_options, key, value); if (!setValid) { JniThrow("Config " + key + " is not supported."); } } return options; } #endif arrow::Result> GetFragmentScanOptions(jint file_format_id, const std::unordered_map& configs) { switch (file_format_id) { #ifdef ARROW_CSV case 3: return ToCsvFragmentScanOptions(configs); #endif default: return arrow::Status::Invalid("Illegal file format id: ", file_format_id); } } std::unordered_map ToStringMap(JNIEnv* env, jobjectArray& str_array) { int length = env->GetArrayLength(str_array); std::unordered_map map; map.reserve(length / 2); for (int i = 0; i < length; i += 2) { auto key = reinterpret_cast(env->GetObjectArrayElement(str_array, i)); auto value = reinterpret_cast(env->GetObjectArrayElement(str_array, i + 1)); map[JStringToCString(env, key)] = JStringToCString(env, value); } return map; } /* * Class: org_apache_arrow_dataset_jni_NativeMemoryPool * Method: getDefaultMemoryPool * Signature: ()J */ JNIEXPORT jlong JNICALL Java_org_apache_arrow_dataset_jni_NativeMemoryPool_getDefaultMemoryPool(JNIEnv* env, jclass) { JNI_METHOD_START return default_memory_pool_id; JNI_METHOD_END(-1L) } /* * Class: org_apache_arrow_dataset_jni_NativeMemoryPool * Method: createListenableMemoryPool * Signature: (Lorg/apache/arrow/memory/ReservationListener;)J */ JNIEXPORT jlong JNICALL Java_org_apache_arrow_dataset_jni_NativeMemoryPool_createListenableMemoryPool( JNIEnv* env, jclass, jobject jlistener) { JNI_METHOD_START jobject jlistener_ref = env->NewGlobalRef(jlistener); JavaVM* vm; if (env->GetJavaVM(&vm) != JNI_OK) { JniThrow("Unable to get JavaVM instance"); } std::shared_ptr listener = std::make_shared(vm, jlistener_ref); auto memory_pool = new ReservationListenableMemoryPool(arrow::default_memory_pool(), listener); return reinterpret_cast(memory_pool); JNI_METHOD_END(-1L) } /* * Class: org_apache_arrow_dataset_jni_NativeMemoryPool * Method: releaseMemoryPool * Signature: (J)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_jni_NativeMemoryPool_releaseMemoryPool( JNIEnv* env, jclass, jlong memory_pool_id) { JNI_METHOD_START if (memory_pool_id == default_memory_pool_id) { return; } ReservationListenableMemoryPool* pool = reinterpret_cast(memory_pool_id); if (pool == nullptr) { return; } std::shared_ptr rm = std::dynamic_pointer_cast(pool->get_listener()); if (rm == nullptr) { delete pool; return; } delete pool; env->DeleteGlobalRef(rm->GetJavaReservationListener()); JNI_METHOD_END() } /* * Class: org_apache_arrow_dataset_jni_NativeMemoryPool * Method: bytesAllocated * Signature: (J)J */ JNIEXPORT jlong JNICALL Java_org_apache_arrow_dataset_jni_NativeMemoryPool_bytesAllocated( JNIEnv* env, jclass, jlong memory_pool_id) { JNI_METHOD_START arrow::MemoryPool* pool = reinterpret_cast(memory_pool_id); if (pool == nullptr) { JniThrow("Memory pool instance not found. It may not exist or have been closed"); } return pool->bytes_allocated(); JNI_METHOD_END(-1L) } /* * Class: org_apache_arrow_dataset_jni_JniWrapper * Method: closeDatasetFactory * Signature: (J)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_closeDatasetFactory( JNIEnv* env, jobject, jlong id) { JNI_METHOD_START ReleaseNativeRef(id); JNI_METHOD_END() } /* * Class: org_apache_arrow_dataset_jni_JniWrapper * Method: inspectSchema * Signature: (J)[B */ JNIEXPORT jbyteArray JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_inspectSchema( JNIEnv* env, jobject, jlong dataset_factor_id) { JNI_METHOD_START std::shared_ptr d = RetrieveNativeInstance(dataset_factor_id); std::shared_ptr schema = JniGetOrThrow(d->Inspect()); return JniGetOrThrow(ToSchemaByteArray(env, schema)); JNI_METHOD_END(nullptr) } /* * Class: org_apache_arrow_dataset_jni_JniWrapper * Method: createDataset * Signature: (J[B)J */ JNIEXPORT jlong JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_createDataset( JNIEnv* env, jobject, jlong dataset_factory_id, jbyteArray schema_bytes) { JNI_METHOD_START std::shared_ptr d = RetrieveNativeInstance(dataset_factory_id); std::shared_ptr schema; schema = JniGetOrThrow(FromSchemaByteArray(env, schema_bytes)); std::shared_ptr dataset = JniGetOrThrow(d->Finish(schema)); return CreateNativeRef(dataset); JNI_METHOD_END(-1L) } /* * Class: org_apache_arrow_dataset_jni_JniWrapper * Method: closeDataset * Signature: (J)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_closeDataset( JNIEnv* env, jobject, jlong id) { JNI_METHOD_START ReleaseNativeRef(id); JNI_METHOD_END() } /* * Class: org_apache_arrow_dataset_jni_JniWrapper * Method: createScanner * Signature: * (J[Ljava/lang/String;Ljava/nio/ByteBuffer;Ljava/nio/ByteBuffer;JI;[Ljava/lang/String;J)J */ JNIEXPORT jlong JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_createScanner( JNIEnv* env, jobject, jlong dataset_id, jobjectArray columns, jobject substrait_projection, jobject substrait_filter, jlong batch_size, jint file_format_id, jobjectArray options, jlong memory_pool_id) { JNI_METHOD_START arrow::MemoryPool* pool = reinterpret_cast(memory_pool_id); if (pool == nullptr) { JniThrow("Memory pool does not exist or has been closed"); } std::shared_ptr dataset = RetrieveNativeInstance(dataset_id); std::shared_ptr scanner_builder = JniGetOrThrow(dataset->NewScan()); JniAssertOkOrThrow(scanner_builder->Pool(pool)); if (columns != nullptr) { std::vector column_vector = ToStringVector(env, columns); JniAssertOkOrThrow(scanner_builder->Project(column_vector)); } if (substrait_projection != nullptr) { std::shared_ptr buffer = LoadArrowBufferFromByteBuffer(env, substrait_projection); std::vector project_exprs; std::vector project_names; arrow::engine::BoundExpressions bounded_expression = JniGetOrThrow(arrow::engine::DeserializeExpressions(*buffer)); for(arrow::engine::NamedExpression& named_expression : bounded_expression.named_expressions) { project_exprs.push_back(std::move(named_expression.expression)); project_names.push_back(std::move(named_expression.name)); } JniAssertOkOrThrow(scanner_builder->Project(std::move(project_exprs), std::move(project_names))); } if (substrait_filter != nullptr) { std::shared_ptr buffer = LoadArrowBufferFromByteBuffer(env, substrait_filter); std::optional filter_expr = std::nullopt; arrow::engine::BoundExpressions bounded_expression = JniGetOrThrow(arrow::engine::DeserializeExpressions(*buffer)); for(arrow::engine::NamedExpression& named_expression : bounded_expression.named_expressions) { filter_expr = named_expression.expression; if (named_expression.expression.type()->id() == arrow::Type::BOOL) { filter_expr = named_expression.expression; } else { JniThrow("There is no filter expression in the expression provided"); } } if (filter_expr == std::nullopt) { JniThrow("The filter expression has not been provided"); } JniAssertOkOrThrow(scanner_builder->Filter(*filter_expr)); } if (file_format_id != -1 && options != nullptr) { std::unordered_map option_map = ToStringMap(env, options); std::shared_ptr scan_options = JniGetOrThrow(GetFragmentScanOptions(file_format_id, option_map)); JniAssertOkOrThrow(scanner_builder->FragmentScanOptions(scan_options)); } JniAssertOkOrThrow(scanner_builder->BatchSize(batch_size)); auto scanner = JniGetOrThrow(scanner_builder->Finish()); std::shared_ptr scanner_adaptor = JniGetOrThrow(DisposableScannerAdaptor::Create(scanner)); jlong id = CreateNativeRef(scanner_adaptor); return id; JNI_METHOD_END(-1L) } /* * Class: org_apache_arrow_dataset_jni_JniWrapper * Method: closeScanner * Signature: (J)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_closeScanner( JNIEnv* env, jobject, jlong scanner_id) { JNI_METHOD_START ReleaseNativeRef(scanner_id); JNI_METHOD_END() } /* * Class: org_apache_arrow_dataset_jni_JniWrapper * Method: getSchemaFromScanner * Signature: (J)[B */ JNIEXPORT jbyteArray JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_getSchemaFromScanner(JNIEnv* env, jobject, jlong scanner_id) { JNI_METHOD_START std::shared_ptr schema = RetrieveNativeInstance(scanner_id) ->GetScanner() ->options() ->projected_schema; return JniGetOrThrow(ToSchemaByteArray(env, schema)); JNI_METHOD_END(nullptr) } /* * Class: org_apache_arrow_dataset_jni_JniWrapper * Method: nextRecordBatch * Signature: (JJ)Z */ JNIEXPORT jboolean JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_nextRecordBatch( JNIEnv* env, jobject, jlong scanner_id, jlong struct_array) { JNI_METHOD_START std::shared_ptr scanner_adaptor = RetrieveNativeInstance(scanner_id); std::shared_ptr record_batch = JniGetOrThrow(scanner_adaptor->Next()); if (record_batch == nullptr) { return false; // stream ended } std::vector> offset_zeroed_arrays; for (int i = 0; i < record_batch->num_columns(); ++i) { // TODO: If the array has an offset then we need to de-offset the array // in order for it to be properly consumed on the Java end. // This forces a copy, it would be nice to avoid this if Java // could consume offset-arrays. Perhaps at some point in the future // using the C data interface. See ARROW-15275 // // Generally a non-zero offset will occur whenever the scanner batch // size is smaller than the batch size of the underlying files. std::shared_ptr array = record_batch->column(i); if (array->offset() == 0) { offset_zeroed_arrays.push_back(array); continue; } std::shared_ptr offset_zeroed = JniGetOrThrow(arrow::Concatenate({array})); offset_zeroed_arrays.push_back(offset_zeroed); } std::shared_ptr offset_zeroed_batch = arrow::RecordBatch::Make( record_batch->schema(), record_batch->num_rows(), offset_zeroed_arrays); JniAssertOkOrThrow( arrow::dataset::jni::ExportRecordBatch(env, offset_zeroed_batch, struct_array)); return true; JNI_METHOD_END(false) } /* * Class: org_apache_arrow_dataset_jni_JniWrapper * Method: releaseBuffer * Signature: (J)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_releaseBuffer( JNIEnv* env, jobject, jlong id) { JNI_METHOD_START ReleaseNativeRef(id); JNI_METHOD_END() } /* * Class: org_apache_arrow_dataset_jni_JniWrapper * Method: ensureS3Finalized * Signature: (J)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_ensureS3Finalized( JNIEnv* env, jobject) { JNI_METHOD_START #ifdef ARROW_S3 JniAssertOkOrThrow(arrow::fs::EnsureS3Finalized()); #endif 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 * Signature: (Ljava/lang/String;II;Ljava/lang/String;Ljava/lang/String)J */ JNIEXPORT jlong JNICALL Java_org_apache_arrow_dataset_file_JniWrapper_makeFileSystemDatasetFactory( JNIEnv* env, jobject, jstring uri, jint file_format_id, jobjectArray options) { JNI_METHOD_START std::shared_ptr file_format = JniGetOrThrow(GetFileFormat(file_format_id)); if (options != nullptr) { std::unordered_map option_map = ToStringMap(env, options); std::shared_ptr scan_options = JniGetOrThrow(GetFragmentScanOptions(file_format_id, option_map)); file_format->default_fragment_scan_options = scan_options; #ifdef ARROW_CSV if (file_format_id == 3) { std::shared_ptr csv_file_format = std::dynamic_pointer_cast(file_format); csv_file_format->parse_options = std::dynamic_pointer_cast(scan_options) ->parse_options; } #endif } arrow::dataset::FileSystemFactoryOptions options; std::shared_ptr d = JniGetOrThrow(arrow::dataset::FileSystemDatasetFactory::Make( JStringToCString(env, uri), file_format, options)); return CreateNativeRef(d); JNI_METHOD_END(-1L) } /* * Class: org_apache_arrow_dataset_file_JniWrapper * Method: makeFileSystemDatasetFactoryWithFiles * Signature: ([Ljava/lang/String;II;[Ljava/lang/String)J */ JNIEXPORT jlong JNICALL Java_org_apache_arrow_dataset_file_JniWrapper_makeFileSystemDatasetFactoryWithFiles( JNIEnv* env, jobject, jobjectArray uris, jint file_format_id, jobjectArray options) { JNI_METHOD_START std::shared_ptr file_format = JniGetOrThrow(GetFileFormat(file_format_id)); if (options != nullptr) { std::unordered_map option_map = ToStringMap(env, options); std::shared_ptr scan_options = JniGetOrThrow(GetFragmentScanOptions(file_format_id, option_map)); file_format->default_fragment_scan_options = scan_options; #ifdef ARROW_CSV if (file_format_id == 3) { std::shared_ptr csv_file_format = std::dynamic_pointer_cast(file_format); csv_file_format->parse_options = std::dynamic_pointer_cast(scan_options) ->parse_options; } #endif } arrow::dataset::FileSystemFactoryOptions options; std::vector uri_vec = ToStringVector(env, uris); if (uri_vec.size() == 0) { JniThrow("No URIs provided."); } // If not all URIs, throw exception if (auto elem = std::find_if_not(uri_vec.begin(), uri_vec.end(), arrow::fs::internal::IsLikelyUri); elem != uri_vec.end()) { JniThrow("Unrecognized file type in URI: " + *elem); } std::vector output_paths; std::string first_path; // We know that uri_vec isn't empty, from the conditional above auto fs = JniGetOrThrow(arrow::fs::FileSystemFromUri(uri_vec[0], &first_path)); output_paths.push_back(first_path); std::transform(uri_vec.begin() + 1, uri_vec.end(), std::back_inserter(output_paths), [&](const auto& s) -> std::string { auto result = JniGetOrThrow(fs->PathFromUri(s)); return std::move(result); }); std::shared_ptr d = JniGetOrThrow(arrow::dataset::FileSystemDatasetFactory::Make( std::move(fs), std::move(output_paths), file_format, options)); return CreateNativeRef(d); JNI_METHOD_END(-1L) } /* * Class: org_apache_arrow_dataset_file_JniWrapper * Method: writeFromScannerToFile * Signature: * (JJJLjava/lang/String;[Ljava/lang/String;ILjava/lang/String;)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_file_JniWrapper_writeFromScannerToFile( JNIEnv* env, jobject, jlong c_arrow_array_stream_address, jlong file_format_id, jstring uri, jobjectArray partition_columns, jint max_partitions, jstring base_name_template) { JNI_METHOD_START JavaVM* vm; if (env->GetJavaVM(&vm) != JNI_OK) { JniThrow("Unable to get JavaVM instance"); } auto* arrow_stream = reinterpret_cast(c_arrow_array_stream_address); std::shared_ptr reader = JniGetOrThrow(arrow::ImportRecordBatchReader(arrow_stream)); std::shared_ptr scanner_builder = arrow::dataset::ScannerBuilder::FromRecordBatchReader(reader); JniAssertOkOrThrow(scanner_builder->Pool(arrow::default_memory_pool())); auto scanner = JniGetOrThrow(scanner_builder->Finish()); std::shared_ptr schema = reader->schema(); std::shared_ptr file_format = JniGetOrThrow(GetFileFormat(file_format_id)); arrow::dataset::FileSystemDatasetWriteOptions options; std::string output_path; auto filesystem = JniGetOrThrow( arrow::fs::FileSystemFromUri(JStringToCString(env, uri), &output_path)); std::vector partition_column_vector = ToStringVector(env, partition_columns); options.file_write_options = file_format->DefaultWriteOptions(); options.filesystem = filesystem; options.base_dir = output_path; options.basename_template = JStringToCString(env, base_name_template); options.partitioning = std::make_shared( SchemaFromColumnNames(schema, partition_column_vector).ValueOrDie()); options.max_partitions = max_partitions; JniAssertOkOrThrow(arrow::dataset::FileSystemDataset::Write(options, scanner)); JNI_METHOD_END() } /* * Class: org_apache_arrow_dataset_substrait_JniWrapper * Method: executeSerializedPlan * Signature: (Ljava/lang/String;[Ljava/lang/String;J)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_substrait_JniWrapper_executeSerializedPlan__Ljava_lang_String_2_3Ljava_lang_String_2J ( JNIEnv* env, jobject, jstring plan, jobjectArray table_to_memory_address_input, jlong memory_address_output) { JNI_METHOD_START // get mapping of table name to memory address std::unordered_map> map_table_to_reader = LoadNamedTables(env, table_to_memory_address_input); // create table provider arrow::engine::NamedTableProvider table_provider = [&map_table_to_reader](const std::vector& names, const arrow::Schema&) { std::shared_ptr output_table = GetTableByName(names, map_table_to_reader); std::shared_ptr options = std::make_shared(std::move(output_table)); return arrow::acero::Declaration("table_source", {}, options, "java_source"); }; arrow::engine::ConversionOptions conversion_options; conversion_options.named_table_provider = std::move(table_provider); // execute plan std::shared_ptr buffer = JniGetOrThrow(arrow::engine::SerializeJsonPlan( JStringToCString(env, plan))); std::shared_ptr reader_out = JniGetOrThrow(arrow::engine::ExecuteSerializedPlan(*buffer, nullptr, nullptr, conversion_options)); auto* arrow_stream_out = reinterpret_cast(memory_address_output); JniAssertOkOrThrow(arrow::ExportRecordBatchReader(reader_out, arrow_stream_out)); JNI_METHOD_END() } /* * Class: org_apache_arrow_dataset_substrait_JniWrapper * Method: executeSerializedPlan * Signature: (Ljava/nio/ByteBuffer;[Ljava/lang/String;J)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_substrait_JniWrapper_executeSerializedPlan__Ljava_nio_ByteBuffer_2_3Ljava_lang_String_2J ( JNIEnv* env, jobject, jobject plan, jobjectArray table_to_memory_address_input, jlong memory_address_output) { JNI_METHOD_START // get mapping of table name to memory address std::unordered_map> map_table_to_reader = LoadNamedTables(env, table_to_memory_address_input); // create table provider arrow::engine::NamedTableProvider table_provider = [&map_table_to_reader](const std::vector& names, const arrow::Schema&) { std::shared_ptr output_table = GetTableByName(names, map_table_to_reader); std::shared_ptr options = std::make_shared(std::move(output_table)); return arrow::acero::Declaration("table_source", {}, options, "java_source"); }; arrow::engine::ConversionOptions conversion_options; conversion_options.named_table_provider = std::move(table_provider); // mapping arrow::Buffer std::shared_ptr buffer = LoadArrowBufferFromByteBuffer(env, plan); // execute plan std::shared_ptr reader_out = JniGetOrThrow(arrow::engine::ExecuteSerializedPlan(*buffer, nullptr, nullptr, conversion_options)); auto* arrow_stream_out = reinterpret_cast(memory_address_output); JniAssertOkOrThrow(arrow::ExportRecordBatchReader(reader_out, arrow_stream_out)); JNI_METHOD_END() } apache-arrow-java-19.0.0/dataset/src/main/java/000077500000000000000000000000001515451615200211575ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/module-info.java000066400000000000000000000022001515451615200242320ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ open module org.apache.arrow.dataset { exports org.apache.arrow.dataset.file; exports org.apache.arrow.dataset.source; exports org.apache.arrow.dataset.jni; exports org.apache.arrow.dataset.substrait; exports org.apache.arrow.dataset.scanner; requires org.apache.arrow.c; requires org.apache.arrow.memory.core; requires org.apache.arrow.vector; } apache-arrow-java-19.0.0/dataset/src/main/java/org/000077500000000000000000000000001515451615200217465ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/000077500000000000000000000000001515451615200231675ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200243215ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/000077500000000000000000000000001515451615200257465ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/file/000077500000000000000000000000001515451615200266655ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/file/DatasetFileWriter.java000066400000000000000000000052441515451615200331170ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.file; import org.apache.arrow.c.ArrowArrayStream; import org.apache.arrow.c.Data; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.ipc.ArrowReader; /** * JNI-based utility to write datasets into files. It internally depends on C++ static method * FileSystemDataset::Write. */ public class DatasetFileWriter { /** * Write the contents of an ArrowReader as a dataset. * * @param reader the datasource for writing * @param format target file format * @param uri target file uri * @param maxPartitions maximum partitions to be included in written files * @param partitionColumns columns used to partition output files. Empty to disable partitioning * @param baseNameTemplate file name template used to make partitions. E.g. "dat_{i}", i is * current partition ID around all written files. */ public static void write( BufferAllocator allocator, ArrowReader reader, FileFormat format, String uri, String[] partitionColumns, int maxPartitions, String baseNameTemplate) { try (final ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) { Data.exportArrayStream(allocator, reader, stream); JniWrapper.get() .writeFromScannerToFile( stream.memoryAddress(), format.id(), uri, partitionColumns, maxPartitions, baseNameTemplate); } } /** * Write the contents of an ArrowReader as a dataset, with default partitioning settings. * * @param reader the datasource for writing * @param format target file format * @param uri target file uri */ public static void write( BufferAllocator allocator, ArrowReader reader, FileFormat format, String uri) { write(allocator, reader, format, uri, new String[0], 1024, "data_{i}"); } } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/file/FileFormat.java000066400000000000000000000020741515451615200315630ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.file; /** File format definitions. */ public enum FileFormat { PARQUET(0), ARROW_IPC(1), ORC(2), CSV(3), JSON(4), NONE(-1); private final int id; FileFormat(int id) { this.id = id; } public int id() { return id; } } FileSystemDatasetFactory.java000066400000000000000000000054711515451615200344020ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/file/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.file; import java.util.Optional; import org.apache.arrow.dataset.jni.NativeDatasetFactory; import org.apache.arrow.dataset.jni.NativeMemoryPool; import org.apache.arrow.dataset.scanner.FragmentScanOptions; import org.apache.arrow.memory.BufferAllocator; /** Java binding of the C++ FileSystemDatasetFactory. */ public class FileSystemDatasetFactory extends NativeDatasetFactory { public FileSystemDatasetFactory( BufferAllocator allocator, NativeMemoryPool memoryPool, FileFormat format, String uri) { super(allocator, memoryPool, createNative(format, uri, Optional.empty())); } public FileSystemDatasetFactory( BufferAllocator allocator, NativeMemoryPool memoryPool, FileFormat format, String uri, Optional fragmentScanOptions) { super(allocator, memoryPool, createNative(format, uri, fragmentScanOptions)); } public FileSystemDatasetFactory( BufferAllocator allocator, NativeMemoryPool memoryPool, FileFormat format, String[] uris) { super(allocator, memoryPool, createNative(format, uris, Optional.empty())); } public FileSystemDatasetFactory( BufferAllocator allocator, NativeMemoryPool memoryPool, FileFormat format, String[] uris, Optional fragmentScanOptions) { super(allocator, memoryPool, createNative(format, uris, fragmentScanOptions)); } private static long createNative( FileFormat format, String uri, Optional fragmentScanOptions) { return JniWrapper.get() .makeFileSystemDatasetFactory( uri, format.id(), fragmentScanOptions.map(FragmentScanOptions::serialize).orElse(null)); } private static long createNative( FileFormat format, String[] uris, Optional fragmentScanOptions) { return JniWrapper.get() .makeFileSystemDatasetFactoryWithFiles( uris, format.id(), fragmentScanOptions.map(FragmentScanOptions::serialize).orElse(null)); } } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/file/JniWrapper.java000066400000000000000000000062211515451615200316120ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.file; import org.apache.arrow.dataset.jni.JniLoader; /** * JniWrapper for filesystem based {@link org.apache.arrow.dataset.source.Dataset} implementations. */ public class JniWrapper { private static final JniWrapper INSTANCE = new JniWrapper(); public static JniWrapper get() { JniLoader.get().ensureLoaded(); return INSTANCE; } private JniWrapper() {} /** * Create FileSystemDatasetFactory and return its native pointer. The pointer is pointing to a * intermediate shared_ptr of the factory instance. * * @param uri file uri to read, either a file or a directory * @param fileFormat file format ID. * @param serializedFragmentScanOptions serialized FragmentScanOptions. * @return the native pointer of the arrow::dataset::FileSystemDatasetFactory instance. * @see FileFormat */ public native long makeFileSystemDatasetFactory( String uri, int fileFormat, String[] serializedFragmentScanOptions); /** * Create FileSystemDatasetFactory and return its native pointer. The pointer is pointing to a * intermediate shared_ptr of the factory instance. * * @param uris List of file uris to read, each path pointing to an individual file * @param fileFormat file format ID. * @param serializedFragmentScanOptions serialized FragmentScanOptions. * @return the native pointer of the arrow::dataset::FileSystemDatasetFactory instance. * @see FileFormat */ public native long makeFileSystemDatasetFactoryWithFiles( String[] uris, int fileFormat, String[] serializedFragmentScanOptions); /** * Write the content in a {@link org.apache.arrow.c.ArrowArrayStream} into files. This internally * depends on C++ write API: FileSystemDataset::Write. * * @param streamAddress the ArrowArrayStream address * @param fileFormat target file format (ID) * @param uri target file uri * @param partitionColumns columns used to partition output files * @param maxPartitions maximum partitions to be included in written files * @param baseNameTemplate file name template used to make partitions. E.g. "dat_{i}", i is * current partition ID around all written files. */ public native void writeFromScannerToFile( long streamAddress, long fileFormat, String uri, String[] partitionColumns, int maxPartitions, String baseNameTemplate); } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/jni/000077500000000000000000000000001515451615200265265ustar00rootroot00000000000000DirectReservationListener.java000066400000000000000000000110511515451615200344520ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/jni/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.jni; import java.lang.reflect.Field; import java.lang.reflect.Method; import java.util.concurrent.atomic.AtomicLong; import org.apache.arrow.util.VisibleForTesting; /** * Reserving Java direct memory bytes from java.nio.Bits. Used by Java Dataset API's C++ memory pool * implementation. This makes memory allocated by the pool to be controlled by JVM option * "-XX:MaxDirectMemorySize". */ public class DirectReservationListener implements ReservationListener { private final Method methodReserve; private final Method methodUnreserve; private DirectReservationListener() { try { final Class classBits = Class.forName("java.nio.Bits"); methodReserve = this.getDeclaredMethodBaseOnJDKVersion(classBits, "reserveMemory"); methodReserve.setAccessible(true); methodUnreserve = this.getDeclaredMethodBaseOnJDKVersion(classBits, "unreserveMemory"); methodUnreserve.setAccessible(true); } catch (Exception e) { final RuntimeException failure = new RuntimeException( "Failed to initialize DirectReservationListener. When starting Java you must include " + "`--add-opens=java.base/java.nio=org.apache.arrow.dataset,org.apache.arrow.memory.core,ALL-UNNAMED` " + "(See https://arrow.apache.org/docs/java/install.html)", e); failure.printStackTrace(); throw failure; } } private static final DirectReservationListener INSTANCE = new DirectReservationListener(); public static DirectReservationListener instance() { return INSTANCE; } /** Reserve bytes by invoking java.nio.java.Bitjava.nio.Bitss#reserveMemory. */ @Override public void reserve(long size) { try { if (size > Integer.MAX_VALUE) { throw new IllegalArgumentException( "reserve size should not be larger than Integer.MAX_VALUE (0x7fffffff)"); } methodReserve.invoke(null, (int) size, (int) size); } catch (Exception e) { throw new RuntimeException(e); } } /** Unreserve bytes by invoking java.nio.java.Bitjava.nio.Bitss#unreserveMemory. */ @Override public void unreserve(long size) { try { if (size > Integer.MAX_VALUE) { throw new IllegalArgumentException( "unreserve size should not be larger than Integer.MAX_VALUE (0x7fffffff)"); } methodUnreserve.invoke(null, (int) size, (int) size); } catch (Exception e) { throw new RuntimeException(e); } } /** Get current reservation of jVM direct memory. Visible for testing. */ @VisibleForTesting public long getCurrentDirectMemReservation() { try { final Class classBits = Class.forName("java.nio.Bits"); Field f; try { f = classBits.getDeclaredField("reservedMemory"); } catch (NoSuchFieldException e) { try { f = classBits.getDeclaredField("RESERVED_MEMORY"); } catch (NoSuchFieldException ex) { throw new AssertionError(ex); } } f.setAccessible(true); return ((AtomicLong) f.get(null)).get(); } catch (Exception e) { throw new RuntimeException(e); } } /** * Get the given method via reflection, searching for different signatures based on the Java * version. * * @param classBits The java.nio.Bits class. * @param name The method being requested. * @return The method object. */ private Method getDeclaredMethodBaseOnJDKVersion(Class classBits, String name) { try { return classBits.getDeclaredMethod(name, long.class, int.class); } catch (NoSuchMethodException e) { try { return classBits.getDeclaredMethod(name, long.class, long.class); } catch (NoSuchMethodException ex) { throw new AssertionError(ex); } } } } JniExceptionDescriber.java000066400000000000000000000027631515451615200335440ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/jni/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.jni; import java.io.PrintWriter; import java.io.StringWriter; /** For native code to invoke to convert a java/lang/Throwable to jstring. */ class JniExceptionDescriber { private JniExceptionDescriber() {} /** * Convert a java/lang/Throwable to jstring. See codes in arrow::dataset::jni::CheckException for * more details. * * @param throwable the exception instance. * @return a String including error message and stack trace of the exception. */ static String describe(Throwable throwable) { StringWriter sw = new StringWriter(); PrintWriter pw = new PrintWriter(sw, true); throwable.printStackTrace(pw); return sw.getBuffer().toString(); } } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/jni/JniLoader.java000066400000000000000000000073671515451615200312550ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.jni; import java.io.File; import java.io.FileNotFoundException; import java.io.IOException; import java.io.InputStream; import java.nio.file.Files; import java.nio.file.StandardCopyOption; import java.util.ArrayList; import java.util.Collections; import java.util.HashSet; import java.util.List; import java.util.Locale; import java.util.Set; /** The JniLoader for Dataset API's native implementation. */ public final class JniLoader { private static final JniLoader INSTANCE = new JniLoader(Collections.singletonList("arrow_dataset_jni")); public static JniLoader get() { return INSTANCE; } private final Set librariesToLoad; private JniLoader(List libraryNames) { librariesToLoad = new HashSet<>(libraryNames); } private boolean finished() { return librariesToLoad.isEmpty(); } /** If required JNI libraries are not loaded, then load them. */ public void ensureLoaded() { if (finished()) { return; } loadRemaining(); ensureS3FinalizedOnShutdown(); JniWrapper.get().initialize(); } private synchronized void loadRemaining() { // The method is protected by a mutex via synchronized, if more than one thread race to call // loadRemaining, at same time only one will do the actual loading and the others will wait for // the mutex to be acquired then check on the remaining list: if there are libraries that were // not // successfully loaded then the mutex owner will try to load them again. if (finished()) { return; } List libs = new ArrayList<>(librariesToLoad); for (String lib : libs) { load(lib); librariesToLoad.remove(lib); } } private void load(String name) { final String libraryToLoad = name + "/" + getNormalizedArch() + "/" + System.mapLibraryName(name); try { File temp = File.createTempFile("jnilib-", ".tmp", new File(System.getProperty("java.io.tmpdir"))); temp.deleteOnExit(); try (final InputStream is = JniWrapper.class.getClassLoader().getResourceAsStream(libraryToLoad)) { if (is == null) { throw new FileNotFoundException(libraryToLoad); } Files.copy(is, temp.toPath(), StandardCopyOption.REPLACE_EXISTING); System.load(temp.getAbsolutePath()); } } catch (IOException e) { throw new IllegalStateException("error loading native libraries: " + e); } } private String getNormalizedArch() { String arch = System.getProperty("os.arch").toLowerCase(Locale.US); switch (arch) { case "amd64": arch = "x86_64"; break; case "aarch64": arch = "aarch_64"; break; default: break; } return arch; } private void ensureS3FinalizedOnShutdown() { Runtime.getRuntime() .addShutdownHook( new Thread( () -> { JniWrapper.get().ensureS3Finalized(); })); } } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/jni/JniWrapper.java000066400000000000000000000115411515451615200314540ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.jni; import java.nio.ByteBuffer; /** JNI wrapper for Dataset API's native implementation. */ public class JniWrapper { private static final JniWrapper INSTANCE = new JniWrapper(); public static JniWrapper get() { JniLoader.get().ensureLoaded(); return INSTANCE; } private JniWrapper() {} /** * Release the DatasetFactory by destroying its reference held by JNI wrapper. * * @param datasetFactoryId the native pointer of the arrow::dataset::DatasetFactory instance. */ public native void closeDatasetFactory(long datasetFactoryId); /** * Get a serialized schema from native instance of a DatasetFactory. * * @param datasetFactoryId the native pointer of the arrow::dataset::DatasetFactory instance. * @return the serialized schema * @see org.apache.arrow.vector.types.pojo.Schema */ public native byte[] inspectSchema(long datasetFactoryId); /** * Create Dataset from a DatasetFactory and get the native pointer of the Dataset. * * @param datasetFactoryId the native pointer of the arrow::dataset::DatasetFactory instance. * @param schema the predefined schema of the resulting Dataset. * @return the native pointer of the arrow::dataset::Dataset instance. */ public native long createDataset(long datasetFactoryId, byte[] schema); /** * Release the Dataset by destroying its reference held by JNI wrapper. * * @param datasetId the native pointer of the arrow::dataset::Dataset instance. */ public native void closeDataset(long datasetId); /** * Create Scanner from a Dataset and get the native pointer of the Dataset. * * @param datasetId the native pointer of the arrow::dataset::Dataset instance. * @param columns desired column names. Columns not in this list will not be emitted when * performing scan operation. Null equals to "all columns". * @param substraitProjection substrait extended expression to evaluate for project new columns * @param substraitFilter substrait extended expression to evaluate for apply filter * @param batchSize batch size of scanned record batches. * @param fileFormat file format ID. * @param serializedFragmentScanOptions serialized FragmentScanOptions. * @param memoryPool identifier of memory pool used in the native scanner. * @return the native pointer of the arrow::dataset::Scanner instance. */ public native long createScanner( long datasetId, String[] columns, ByteBuffer substraitProjection, ByteBuffer substraitFilter, long batchSize, int fileFormat, String[] serializedFragmentScanOptions, long memoryPool); /** * Get a serialized schema from native instance of a Scanner. * * @param scannerId the native pointer of the arrow::dataset::Scanner instance. * @return the serialized schema * @see org.apache.arrow.vector.types.pojo.Schema */ public native byte[] getSchemaFromScanner(long scannerId); /** * Release the Scanner by destroying its reference held by JNI wrapper. * * @param scannerId the native pointer of the arrow::dataset::Scanner instance. */ public native void closeScanner(long scannerId); /** * Read next record batch from the specified scanner. * * @param scannerId the native pointer of the arrow::dataset::Scanner instance. * @param arrowArray pointer to an empty {@link org.apache.arrow.c.ArrowArray} struct to store C++ * side record batch that conforms to C data interface. * @return true if valid record batch is returned; false if stream ended. */ public native boolean nextRecordBatch(long scannerId, long arrowArray); /** * Release the Buffer by destroying its reference held by JNI wrapper. * * @param bufferId the native pointer of the arrow::Buffer instance. */ public native void releaseBuffer(long bufferId); /** * Ensure the S3 APIs are shutdown, but only if not already done. If the S3 APIs are * uninitialized, then this is a noop. */ public native void ensureS3Finalized(); /** Initialize Arrow Compute. */ public native void initialize(); } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/jni/NativeContext.java000066400000000000000000000030361515451615200321660ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.jni; import org.apache.arrow.memory.BufferAllocator; /** Context for relevant classes of NativeDataset. */ public class NativeContext { private final BufferAllocator allocator; private final NativeMemoryPool memoryPool; /** * Constructor. * * @param allocator The allocator in use. * @param memoryPool Native memory pool. */ public NativeContext(BufferAllocator allocator, NativeMemoryPool memoryPool) { this.allocator = allocator; this.memoryPool = memoryPool; } /** Returns the allocator which is in use. */ public BufferAllocator getAllocator() { return allocator; } /** Returns the native memory pool. */ public NativeMemoryPool getMemoryPool() { return memoryPool; } } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/jni/NativeDataset.java000066400000000000000000000046361515451615200321360ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.jni; import org.apache.arrow.dataset.scanner.FragmentScanOptions; import org.apache.arrow.dataset.scanner.ScanOptions; import org.apache.arrow.dataset.source.Dataset; /** Native implementation of {@link Dataset}. */ public class NativeDataset implements Dataset { private final NativeContext context; private final long datasetId; private boolean closed = false; public NativeDataset(NativeContext context, long datasetId) { this.context = context; this.datasetId = datasetId; } @Override public synchronized NativeScanner newScan(ScanOptions options) { if (closed) { throw new NativeInstanceReleasedException(); } int fileFormatId = -1; String[] serialized = null; if (options.getFragmentScanOptions().isPresent()) { FragmentScanOptions fragmentScanOptions = options.getFragmentScanOptions().get(); fileFormatId = fragmentScanOptions.fileFormat().id(); serialized = fragmentScanOptions.serialize(); } long scannerId = JniWrapper.get() .createScanner( datasetId, options.getColumns().orElse(null), options.getSubstraitProjection().orElse(null), options.getSubstraitFilter().orElse(null), options.getBatchSize(), fileFormatId, serialized, context.getMemoryPool().getNativeInstanceId()); return new NativeScanner(context, scannerId); } @Override public synchronized void close() { if (closed) { return; } closed = true; JniWrapper.get().closeDataset(datasetId); } } NativeDatasetFactory.java000066400000000000000000000067671515451615200334160ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/jni/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.jni; import java.io.IOException; import org.apache.arrow.dataset.source.DatasetFactory; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.SchemaUtility; /** Native implementation of {@link DatasetFactory}. */ public class NativeDatasetFactory implements DatasetFactory { private final long datasetFactoryId; private final NativeMemoryPool memoryPool; private final BufferAllocator allocator; private boolean closed = false; /** * Constructor. * * @param allocator a context allocator associated with this factory. Any buffer that will be * created natively will be then bound to this allocator. * @param memoryPool the native memory pool associated with this factory. Any buffer created * natively should request for memory spaces from this memory pool. This is a mapped instance * of c++ arrow::MemoryPool. * @param datasetFactoryId an ID, at the same time the native pointer of the underlying native * instance of this factory. Make sure in c++ side the pointer is pointing to the shared * pointer wrapping the actual instance so we could successfully decrease the reference count * once {@link #close} is called. * @see #close() */ public NativeDatasetFactory( BufferAllocator allocator, NativeMemoryPool memoryPool, long datasetFactoryId) { this.allocator = allocator; this.memoryPool = memoryPool; this.datasetFactoryId = datasetFactoryId; } @Override public Schema inspect() { final byte[] buffer; synchronized (this) { if (closed) { throw new NativeInstanceReleasedException(); } buffer = JniWrapper.get().inspectSchema(datasetFactoryId); } try { return SchemaUtility.deserialize(buffer, allocator); } catch (IOException e) { throw new RuntimeException(e); } } @Override public NativeDataset finish() { return finish(inspect()); } @Override public NativeDataset finish(Schema schema) { try { byte[] serialized = SchemaUtility.serialize(schema); synchronized (this) { if (closed) { throw new NativeInstanceReleasedException(); } return new NativeDataset( new NativeContext(allocator, memoryPool), JniWrapper.get().createDataset(datasetFactoryId, serialized)); } } catch (IOException e) { throw new RuntimeException(e); } } /** Close this factory by release the pointer of the native instance. */ @Override public synchronized void close() { if (closed) { return; } closed = true; JniWrapper.get().closeDatasetFactory(datasetFactoryId); } } NativeInstanceReleasedException.java000066400000000000000000000022311515451615200355470ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/jni/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.jni; /** Thrown if trying to operate on a native instance that is already released. */ public class NativeInstanceReleasedException extends RuntimeException { public NativeInstanceReleasedException() { super("Native instance has been released"); } public NativeInstanceReleasedException(String message) { super(message); } } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/jni/NativeMemoryPool.java000066400000000000000000000045211515451615200326440ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.jni; /** C++ memory pool(arrow::MemoryPool)'s Java mapped instance. */ public class NativeMemoryPool implements AutoCloseable { private final long nativeInstanceId; private NativeMemoryPool(long nativeInstanceId) { this.nativeInstanceId = nativeInstanceId; } /** Get the default memory pool. This will return arrow::default_memory_pool() directly. */ public static NativeMemoryPool getDefault() { JniLoader.get().ensureLoaded(); return new NativeMemoryPool(getDefaultMemoryPool()); } /** * Create a listenable memory pool (see also: arrow::ReservationListenableMemoryPool) with a * specific listener. All buffers created from the memory pool should take enough reservation from * the listener in advance. */ public static NativeMemoryPool createListenable(ReservationListener listener) { JniLoader.get().ensureLoaded(); return new NativeMemoryPool(createListenableMemoryPool(listener)); } /** Return native instance ID of this memory pool. */ public long getNativeInstanceId() { return nativeInstanceId; } /** Get current allocated bytes. */ public long getBytesAllocated() { return bytesAllocated(nativeInstanceId); } @Override public void close() throws Exception { releaseMemoryPool(nativeInstanceId); } private static native long getDefaultMemoryPool(); private static native long createListenableMemoryPool(ReservationListener listener); private static native void releaseMemoryPool(long id); private static native long bytesAllocated(long id); } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/jni/NativeScanTask.java000066400000000000000000000031141515451615200322460ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.jni; import org.apache.arrow.dataset.scanner.ScanTask; import org.apache.arrow.vector.ipc.ArrowReader; /** * Native implementation of {@link ScanTask}. Currently RecordBatches are iterated directly by the * scanner id via {@link JniWrapper}, thus we allow only one-time execution of method {@link * #execute()}. If a re-scan operation is expected, call {@link NativeDataset#newScan} to create a * new scanner instance. */ @Deprecated public class NativeScanTask implements ScanTask { private final NativeScanner scanner; /** Constructor. */ public NativeScanTask(NativeScanner scanner) { this.scanner = scanner; } @Override public ArrowReader execute() { return scanner.execute(); } @Override public void close() { scanner.close(); } } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/jni/NativeScanner.java000066400000000000000000000123271515451615200321360ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.jni; import java.io.IOException; import java.util.Collections; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReadWriteLock; import java.util.concurrent.locks.ReentrantReadWriteLock; import org.apache.arrow.c.ArrowArray; import org.apache.arrow.c.Data; import org.apache.arrow.dataset.scanner.Scanner; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.SchemaUtility; /** * Native implementation of {@link Scanner}. Note that it currently emits only a single scan task of * type {@link NativeScanTask}, which is internally a combination of all scan task instances * returned by the native scanner. */ public class NativeScanner implements Scanner { private final AtomicBoolean executed = new AtomicBoolean(false); private final NativeContext context; private final long scannerId; private final ReadWriteLock lock = new ReentrantReadWriteLock(); private final Lock writeLock = lock.writeLock(); private final Lock readLock = lock.readLock(); private boolean closed = false; public NativeScanner(NativeContext context, long scannerId) { this.context = context; this.scannerId = scannerId; } ArrowReader execute() { if (closed) { throw new NativeInstanceReleasedException(); } if (!executed.compareAndSet(false, true)) { throw new UnsupportedOperationException( "NativeScanner cannot be executed more than once. Consider creating " + "new scanner instead"); } return new NativeReader(context.getAllocator()); } @Override public ArrowReader scanBatches() { if (closed) { throw new NativeInstanceReleasedException(); } if (!executed.compareAndSet(false, true)) { throw new UnsupportedOperationException( "NativeScanner can only be executed once. Create a " + "new scanner instead"); } return new NativeReader(context.getAllocator()); } @Override @Deprecated public Iterable scan() { if (closed) { throw new NativeInstanceReleasedException(); } return Collections.singletonList(new NativeScanTask(this)); } @Override public Schema schema() { readLock.lock(); try { if (closed) { throw new NativeInstanceReleasedException(); } return SchemaUtility.deserialize( JniWrapper.get().getSchemaFromScanner(scannerId), context.getAllocator()); } catch (IOException e) { throw new RuntimeException(e); } finally { readLock.unlock(); } } @Override public void close() { writeLock.lock(); try { if (closed) { return; } closed = true; JniWrapper.get().closeScanner(scannerId); } finally { writeLock.unlock(); } } /** {@link ArrowReader} implementation for NativeDataset. */ public class NativeReader extends ArrowReader { private NativeReader(BufferAllocator allocator) { super(allocator); } @Override protected void loadRecordBatch(ArrowRecordBatch batch) { throw new UnsupportedOperationException(); } @Override protected void loadDictionary(ArrowDictionaryBatch dictionaryBatch) { throw new UnsupportedOperationException(); } @Override public boolean loadNextBatch() throws IOException { readLock.lock(); try { if (closed) { throw new NativeInstanceReleasedException(); } try (ArrowArray arrowArray = ArrowArray.allocateNew(context.getAllocator())) { if (!JniWrapper.get().nextRecordBatch(scannerId, arrowArray.memoryAddress())) { return false; } final VectorSchemaRoot vsr = getVectorSchemaRoot(); Data.importIntoVectorSchemaRoot(context.getAllocator(), arrowArray, vsr, this); } } finally { readLock.unlock(); } return true; } @Override public long bytesRead() { return 0L; } @Override protected void closeReadSource() throws IOException { // no-op } @Override protected Schema readSchema() throws IOException { return schema(); } } } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/jni/ReservationListener.java000066400000000000000000000021611515451615200334000ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.jni; /** Listener of buffer memory reservation. Used by native datasets. */ public interface ReservationListener { /** * Reserve bytes. * * @throws RuntimeException if request size cannot be granted */ void reserve(long size); /** Unreserve bytes. */ void unreserve(long size); } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/scanner/000077500000000000000000000000001515451615200273775ustar00rootroot00000000000000ArrowScannerReader.java000066400000000000000000000070171515451615200337170ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/scanner/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.scanner; import java.io.IOException; import java.util.Iterator; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorUnloader; import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.types.pojo.Schema; /** An implementation of {@link ArrowReader} that reads the dataset from {@link Scanner}. */ public class ArrowScannerReader extends ArrowReader { private final Scanner scanner; private Iterator taskIterator; private ScanTask currentTask = null; private ArrowReader currentReader = null; /** * Constructs a scanner reader using a Scanner. * * @param scanner scanning data over dataset * @param allocator to allocate new buffers */ public ArrowScannerReader(Scanner scanner, BufferAllocator allocator) { super(allocator); this.scanner = scanner; this.taskIterator = scanner.scan().iterator(); if (taskIterator.hasNext()) { currentTask = taskIterator.next(); currentReader = currentTask.execute(); } } @Override protected void loadRecordBatch(ArrowRecordBatch batch) { throw new UnsupportedOperationException(); } @Override protected void loadDictionary(ArrowDictionaryBatch dictionaryBatch) { throw new UnsupportedOperationException(); } @Override public boolean loadNextBatch() throws IOException { if (currentReader == null) { return false; } boolean result = currentReader.loadNextBatch(); if (!result) { try { currentTask.close(); currentReader.close(); } catch (Exception e) { throw new IOException(e); } while (!result) { if (!taskIterator.hasNext()) { return false; } else { currentTask = taskIterator.next(); currentReader = currentTask.execute(); result = currentReader.loadNextBatch(); } } } VectorLoader loader = new VectorLoader(this.getVectorSchemaRoot()); VectorUnloader unloader = new VectorUnloader(currentReader.getVectorSchemaRoot()); try (ArrowRecordBatch recordBatch = unloader.getRecordBatch()) { loader.load(recordBatch); } return true; } @Override public long bytesRead() { return 0L; } @Override protected void closeReadSource() throws IOException { try { currentTask.close(); currentReader.close(); scanner.close(); } catch (Exception e) { throw new IOException(e); } } @Override protected Schema readSchema() throws IOException { return scanner.schema(); } } FragmentScanOptions.java000066400000000000000000000020601515451615200341050ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/scanner/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.scanner; import org.apache.arrow.dataset.file.FileFormat; /** The file fragment scan options interface. It is used to transfer to JNI call. */ public interface FragmentScanOptions { FileFormat fileFormat(); String[] serialize(); } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/scanner/ScanOptions.java000066400000000000000000000131261515451615200325050ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.scanner; import java.nio.ByteBuffer; import java.util.Optional; import org.apache.arrow.util.Preconditions; /** Options used during scanning. */ public class ScanOptions { private final long batchSize; private final Optional columns; private final Optional substraitProjection; private final Optional substraitFilter; private final Optional fragmentScanOptions; /** * Constructor. * * @param columns Projected columns. Empty for scanning all columns. * @param batchSize Maximum row number of each returned {@link * org.apache.arrow.vector.ipc.message.ArrowRecordBatch} * @deprecated Deprecated. Use {@link #ScanOptions(long, Optional)} instead. */ @Deprecated public ScanOptions(String[] columns, long batchSize) { this( batchSize, Optional.of(columns) .map( present -> { if (present.length == 0) { // Backwards compatibility: See ARROW-13257, in the new constructor, we now use // null to scan for all columns. return null; } return present; })); } /** * Constructor. * * @param batchSize Maximum row number of each returned {@link * org.apache.arrow.vector.ipc.message.ArrowRecordBatch} * @param columns (Optional) Projected columns. {@link Optional#empty()} for scanning all columns. * Otherwise, Only columns present in the Array will be scanned. */ public ScanOptions(long batchSize, Optional columns) { Preconditions.checkNotNull(columns); this.batchSize = batchSize; this.columns = columns; this.substraitProjection = Optional.empty(); this.substraitFilter = Optional.empty(); this.fragmentScanOptions = Optional.empty(); } public ScanOptions(long batchSize) { this(batchSize, Optional.empty()); } public Optional getColumns() { return columns; } public long getBatchSize() { return batchSize; } public Optional getSubstraitProjection() { return substraitProjection; } public Optional getSubstraitFilter() { return substraitFilter; } public Optional getFragmentScanOptions() { return fragmentScanOptions; } /** Builder for Options used during scanning. */ public static class Builder { private final long batchSize; private Optional columns; private ByteBuffer substraitProjection; private ByteBuffer substraitFilter; private FragmentScanOptions fragmentScanOptions; /** * Constructor. * * @param batchSize Maximum row number of each returned {@link * org.apache.arrow.vector.ipc.message.ArrowRecordBatch} */ public Builder(long batchSize) { this.batchSize = batchSize; } /** * Set the Projected columns. Empty for scanning all columns. * * @param columns Projected columns. Empty for scanning all columns. * @return the ScanOptions configured. */ public Builder columns(Optional columns) { Preconditions.checkNotNull(columns); this.columns = columns; return this; } /** * Set the Substrait extended expression for Projection new columns. * * @param substraitProjection Expressions to evaluate for project new columns. * @return the ScanOptions configured. */ public Builder substraitProjection(ByteBuffer substraitProjection) { Preconditions.checkNotNull(substraitProjection); this.substraitProjection = substraitProjection; return this; } /** * Set the Substrait extended expression for Filter. * * @param substraitFilter Expressions to evaluate for apply Filter. * @return the ScanOptions configured. */ public Builder substraitFilter(ByteBuffer substraitFilter) { Preconditions.checkNotNull(substraitFilter); this.substraitFilter = substraitFilter; return this; } /** * Set the FragmentScanOptions. * * @param fragmentScanOptions fragment scan options * @return the ScanOptions configured. */ public Builder fragmentScanOptions(FragmentScanOptions fragmentScanOptions) { Preconditions.checkNotNull(fragmentScanOptions); this.fragmentScanOptions = fragmentScanOptions; return this; } public ScanOptions build() { return new ScanOptions(this); } } private ScanOptions(Builder builder) { batchSize = builder.batchSize; columns = builder.columns; substraitProjection = Optional.ofNullable(builder.substraitProjection); substraitFilter = Optional.ofNullable(builder.substraitFilter); fragmentScanOptions = Optional.ofNullable(builder.fragmentScanOptions); } } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/scanner/ScanTask.java000066400000000000000000000023571515451615200317600ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.scanner; import java.io.Reader; import org.apache.arrow.vector.ipc.ArrowReader; /** * Read record batches from a range of a single data fragment. A ScanTask is meant to be a unit of * work to be dispatched. The implementation must be thread and concurrent safe. */ @Deprecated public interface ScanTask extends AutoCloseable { /** Execute this ScanTask and return a {@link Reader} instance. */ ArrowReader execute(); } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/scanner/Scanner.java000066400000000000000000000031371515451615200316370ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.scanner; import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.types.pojo.Schema; /** A high level interface for scanning data over dataset. */ public interface Scanner extends AutoCloseable { /** * Read the dataset as a stream of record batches. * * @return a {@link ArrowReader}. */ ArrowReader scanBatches(); /** * Perform the scan operation. * * @return a iterable set of {@link ScanTask}s. Each task is considered independent and it is * allowed to execute the tasks concurrently to gain better performance. * @deprecated use {@link #scanBatches()} instead. */ @Deprecated Iterable scan(); /** * Get the schema of this Scanner. * * @return the schema instance */ Schema schema(); } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/scanner/csv/000077500000000000000000000000001515451615200301725ustar00rootroot00000000000000CsvConvertOptions.java000066400000000000000000000027241515451615200344330ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/scanner/csv/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.scanner.csv; import java.util.Map; import java.util.Optional; import org.apache.arrow.c.ArrowSchema; public class CsvConvertOptions { private final Map configs; private Optional cSchema = Optional.empty(); public CsvConvertOptions(Map configs) { this.configs = configs; } public Optional getArrowSchema() { return cSchema; } public Map getConfigs() { return configs; } public void set(String key, String value) { configs.put(key, value); } public void setArrowSchema(ArrowSchema cSchema) { this.cSchema = Optional.of(cSchema); } } CsvFragmentScanOptions.java000066400000000000000000000067551515451615200353730ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/scanner/csv/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.scanner.csv; import java.util.Map; import java.util.stream.Collectors; import java.util.stream.Stream; import org.apache.arrow.dataset.file.FileFormat; import org.apache.arrow.dataset.scanner.FragmentScanOptions; import org.apache.arrow.dataset.utils.MapUtil; public class CsvFragmentScanOptions implements FragmentScanOptions { private final CsvConvertOptions convertOptions; private final Map readOptions; private final Map parseOptions; /** * CSV scan options, map to CPP struct CsvFragmentScanOptions. The key in config map is the field * name of mapping cpp struct * *

    Currently, multi-valued options (which are std::vector values in C++) only support having a * single value set. For example, for the null_values option, only one string can be set as the * null value. * * @param convertOptions similar to CsvFragmentScanOptions#convert_options in CPP, the ArrowSchema * represents column_types, convert data option such as null value recognition. * @param readOptions similar to CsvFragmentScanOptions#read_options in CPP, specify how to read * the file such as block_size * @param parseOptions similar to CsvFragmentScanOptions#parse_options in CPP, parse file option * such as delimiter */ public CsvFragmentScanOptions( CsvConvertOptions convertOptions, Map readOptions, Map parseOptions) { this.convertOptions = convertOptions; this.readOptions = readOptions; this.parseOptions = parseOptions; } /** * File format. * * @return file format. */ @Override public FileFormat fileFormat() { return FileFormat.CSV; } /** * This is an internal function to invoke by serializer. Serialize this class to string array and * then called by JNI call. * * @return string array as Map JNI bridge format. */ @Override public String[] serialize() { Map options = Stream.concat( Stream.concat(readOptions.entrySet().stream(), parseOptions.entrySet().stream()), convertOptions.getConfigs().entrySet().stream()) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); if (convertOptions.getArrowSchema().isPresent()) { options.put( "column_types", Long.toString(convertOptions.getArrowSchema().get().memoryAddress())); } return MapUtil.convertMapToStringArray(options); } public CsvConvertOptions getConvertOptions() { return convertOptions; } public Map getReadOptions() { return readOptions; } public Map getParseOptions() { return parseOptions; } } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/source/000077500000000000000000000000001515451615200272465ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/source/Dataset.java000066400000000000000000000024101515451615200314730ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.source; import org.apache.arrow.dataset.scanner.ScanOptions; import org.apache.arrow.dataset.scanner.Scanner; /** A container of Fragments which are the internal iterable unit of read data. */ public interface Dataset extends AutoCloseable { /** * Create a new Scanner using the provided scan options. * * @param options options used during creating Scanner * @return the Scanner instance */ Scanner newScan(ScanOptions options); } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/source/DatasetFactory.java000066400000000000000000000032651515451615200330340ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.source; import org.apache.arrow.vector.types.pojo.Schema; /** * DatasetFactory provides a way to inspect a Dataset potential schema before materializing it. * Thus, the user can peek the schema for data sources and decide on a unified schema. */ public interface DatasetFactory extends AutoCloseable { /** * Get unified schema for the resulting Dataset. * * @return the schema object inspected */ Schema inspect(); /** * Create a Dataset with auto-inferred schema. Which means, the schema of the resulting Dataset * will be the same with calling {@link #inspect()} manually. * * @return the Dataset instance */ Dataset finish(); /** * Create a Dataset with predefined schema. Schema inference will not be performed. * * @param schema a predefined schema * @return the Dataset instance */ Dataset finish(Schema schema); } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/substrait/000077500000000000000000000000001515451615200277665ustar00rootroot00000000000000AceroSubstraitConsumer.java000066400000000000000000000126731515451615200352310ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/substrait/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.substrait; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Map; import org.apache.arrow.c.ArrowArrayStream; import org.apache.arrow.c.Data; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.ipc.ArrowReader; /** * Class to expose Java Substrait API for end users, currently operations supported are only to * Consume Substrait Plan in Plan format (JSON) or Binary format (ByteBuffer). */ public final class AceroSubstraitConsumer { private final BufferAllocator allocator; public AceroSubstraitConsumer(BufferAllocator allocator) { this.allocator = allocator; } /** * Run Substrait plan. * * @param plan The JSON Substrait plan. * @return the ArrowReader to iterate for record batches. */ public ArrowReader runQuery(String plan) throws Exception { return runQuery(plan, Collections.emptyMap()); } /** * Run Substrait plan. * * @param plan The JSON Substrait plan. * @param namedTables A mapping of named tables referenced by the plan to an ArrowReader providing * the data for the table. Contains the Table Name to Query as a Key and ArrowReader as a * Value. * 

    {@code ArrowReader nationReader = scanner.scanBatches();
   * Map namedTables = new HashMap<>();
   * namedTables.put("NATION", nationReader);}
    * * @return the ArrowReader to iterate for record batches. */ public ArrowReader runQuery(String plan, Map namedTables) throws Exception { return execute(plan, namedTables); } /** * Run Substrait plan. * * @param plan the binary Substrait plan. * @return the ArrowReader to iterate for record batches. */ public ArrowReader runQuery(ByteBuffer plan) throws Exception { return runQuery(plan, Collections.emptyMap()); } /** * Read binary Substrait plan, execute and return an ArrowReader to read Schema and * ArrowRecordBatches. * * @param plan the binary Substrait plan. * @param namedTables A mapping of named tables referenced by the plan to an ArrowReader providing * the data for the table. Contains the Table Name to Query as a Key and ArrowReader as a * Value. * 
    {@code ArrowReader nationReader = scanner.scanBatches();
   * Map namedTables = new HashMap<>();
   * namedTables.put("NATION", nationReader);}
    * * @return the ArrowReader to iterate for record batches. */ public ArrowReader runQuery(ByteBuffer plan, Map namedTables) throws Exception { return execute(plan, namedTables); } private ArrowReader execute(String plan, Map namedTables) throws Exception { List arrowArrayStream = new ArrayList<>(); try (ArrowArrayStream streamOutput = ArrowArrayStream.allocateNew(this.allocator)) { String[] mapTableToMemoryAddress = getMapTableToMemoryAddress(namedTables, arrowArrayStream); JniWrapper.get() .executeSerializedPlan(plan, mapTableToMemoryAddress, streamOutput.memoryAddress()); return Data.importArrayStream(this.allocator, streamOutput); } finally { AutoCloseables.close(arrowArrayStream); } } private ArrowReader execute(ByteBuffer plan, Map namedTables) throws Exception { List arrowArrayStream = new ArrayList<>(); try (ArrowArrayStream streamOutput = ArrowArrayStream.allocateNew(this.allocator)) { String[] mapTableToMemoryAddress = getMapTableToMemoryAddress(namedTables, arrowArrayStream); JniWrapper.get() .executeSerializedPlan(plan, mapTableToMemoryAddress, streamOutput.memoryAddress()); return Data.importArrayStream(this.allocator, streamOutput); } finally { AutoCloseables.close(arrowArrayStream); } } private String[] getMapTableToMemoryAddress( Map mapTableToArrowReader, List listStreamInput) { String[] mapTableToMemoryAddress = new String[mapTableToArrowReader.size() * 2]; ArrowArrayStream streamInput; int pos = 0; for (Map.Entry entries : mapTableToArrowReader.entrySet()) { streamInput = ArrowArrayStream.allocateNew(this.allocator); listStreamInput.add(streamInput); Data.exportArrayStream(this.allocator, entries.getValue(), streamInput); mapTableToMemoryAddress[pos] = entries.getKey(); mapTableToMemoryAddress[pos + 1] = String.valueOf(streamInput.memoryAddress()); pos += 2; } return mapTableToMemoryAddress; } } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/substrait/JniWrapper.java000066400000000000000000000057751515451615200327300ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.substrait; import java.nio.ByteBuffer; import org.apache.arrow.dataset.jni.JniLoader; /** * Class that contains Native methods to call Acero C++ Substrait API. It internally depends on C++ * function arrow::engine::ExecuteSerializedPlan. Currently supported input parameters supported * are: * * 
     * - arrow::Buffer: Substrait Plan (JSON or Binary format).
 * - arrow::engine::ConversionOptions: Mapping for arrow::engine::NamedTableProvider.
 *
    */ final class JniWrapper { private static final JniWrapper INSTANCE = new JniWrapper(); private JniWrapper() {} public static JniWrapper get() { JniLoader.get().ensureLoaded(); return INSTANCE; } /** * Consume the JSON Substrait Plan that contains Named Tables and export the RecordBatchReader * into C-Data Interface ArrowArrayStream. * * @param planInput the JSON Substrait plan. * @param mapTableToMemoryAddressInput the mapping name of Tables Name on position `i` and theirs * Memory Address representation on `i+1` position linearly. * 
    {@code String[] mapTableToMemoryAddress = new String[2];
   * mapTableToMemoryAddress[0]="NATION";
   * mapTableToMemoryAddress[1]="140650250895360";}
    * * @param memoryAddressOutput the memory address where RecordBatchReader is exported. */ public native void executeSerializedPlan( String planInput, String[] mapTableToMemoryAddressInput, long memoryAddressOutput); /** * Consume the binary Substrait Plan that contains Named Tables and export the RecordBatchReader * into C-Data Interface ArrowArrayStream. * * @param planInput the binary Substrait plan. * @param mapTableToMemoryAddressInput the mapping name of Tables Name on position `i` and theirs * Memory Address representation on `i+1` position linearly. * 
    {@code String[] mapTableToMemoryAddress = new String[2];
   * mapTableToMemoryAddress[0]="NATION";
   * mapTableToMemoryAddress[1]="140650250895360";}
    * * @param memoryAddressOutput the memory address where RecordBatchReader is exported. */ public native void executeSerializedPlan( ByteBuffer planInput, String[] mapTableToMemoryAddressInput, long memoryAddressOutput); } apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/utils/000077500000000000000000000000001515451615200271065ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/main/java/org/apache/arrow/dataset/utils/MapUtil.java000066400000000000000000000026661515451615200313360ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.utils; import java.util.Map; /** The utility class for Map. */ public class MapUtil { private MapUtil() {} /** * Convert the map to string array as JNI bridge. * * @param config config map * @return string array for serialization */ public static String[] convertMapToStringArray(Map config) { if (config.isEmpty()) { return null; } String[] configs = new String[config.size() * 2]; int i = 0; for (Map.Entry entry : config.entrySet()) { configs[i++] = entry.getKey(); configs[i++] = entry.getValue(); } return configs; } } apache-arrow-java-19.0.0/dataset/src/test/000077500000000000000000000000001515451615200202715ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/test/java/000077500000000000000000000000001515451615200212125ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/test/java/org/000077500000000000000000000000001515451615200220015ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/000077500000000000000000000000001515451615200232225ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200243545ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/arrow/dataset/000077500000000000000000000000001515451615200260015ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/arrow/dataset/OrcWriteSupport.java000066400000000000000000000033541515451615200320040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset; import java.io.IOException; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.Path; import org.apache.hadoop.hive.ql.exec.vector.LongColumnVector; import org.apache.hadoop.hive.ql.exec.vector.VectorizedRowBatch; import org.apache.orc.OrcFile; import org.apache.orc.TypeDescription; import org.apache.orc.Writer; public class OrcWriteSupport { public static void writeTempFile(TypeDescription orcSchema, Path path, Integer[] values) throws IOException { Writer writer = OrcFile.createWriter(path, OrcFile.writerOptions(new Configuration()).setSchema(orcSchema)); VectorizedRowBatch batch = orcSchema.createRowBatch(); LongColumnVector longColumnVector = (LongColumnVector) batch.cols[0]; for (int idx = 0; idx < values.length; idx++) { longColumnVector.vector[idx] = values[idx]; } batch.size = values.length; writer.addRowBatch(batch); writer.close(); } } apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/arrow/dataset/ParquetWriteSupport.java000066400000000000000000000116631515451615200327040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset; import java.io.File; import java.io.InputStream; import java.lang.reflect.Method; import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Random; import org.apache.arrow.util.Preconditions; import org.apache.avro.Schema; import org.apache.avro.generic.GenericData; import org.apache.avro.generic.GenericRecord; import org.apache.parquet.avro.AvroParquetWriter; import org.apache.parquet.hadoop.ParquetWriter; /** Utility class for writing Parquet files using Avro based tools. */ public class ParquetWriteSupport implements AutoCloseable { private final String path; private final String uri; private final ParquetWriter writer; private final Schema avroSchema; private final List writtenRecords = new ArrayList<>(); private final GenericRecordListBuilder recordListBuilder = new GenericRecordListBuilder(); private final Random random = new Random(); public ParquetWriteSupport(String schemaName, File outputFolder) throws Exception { avroSchema = getSchema(schemaName); path = outputFolder.getPath() + "/" + "generated-" + random.nextLong() + ".parquet"; uri = "file://" + path; writer = AvroParquetWriter.builder(new org.apache.hadoop.fs.Path(path)) .withSchema(avroSchema) .build(); } public static Schema getSchema(String schemaName) throws Exception { try { // Attempt to use JDK 9 behavior of getting the module then the resource stream from the // module. // Note that this code is caller-sensitive. Method getModuleMethod = Class.class.getMethod("getModule"); Object module = getModuleMethod.invoke(ParquetWriteSupport.class); Method getResourceAsStreamFromModule = module.getClass().getMethod("getResourceAsStream", String.class); try (InputStream is = (InputStream) getResourceAsStreamFromModule.invoke(module, "/avroschema/" + schemaName)) { return new Schema.Parser().parse(is); } } catch (NoSuchMethodException ex) { // Use JDK8 behavior. try (InputStream is = ParquetWriteSupport.class.getResourceAsStream("/avroschema/" + schemaName)) { return new Schema.Parser().parse(is); } } } public static ParquetWriteSupport writeTempFile( String schemaName, File outputFolder, Object... values) throws Exception { try (final ParquetWriteSupport writeSupport = new ParquetWriteSupport(schemaName, outputFolder)) { writeSupport.writeRecords(values); return writeSupport; } } public void writeRecords(Object... values) throws Exception { final List valueList = getRecordListBuilder().createRecordList(values); writeRecords(valueList); } public void writeRecords(List records) throws Exception { for (GenericRecord record : records) { writeRecord(record); } } public void writeRecord(GenericRecord record) throws Exception { writtenRecords.add(record); writer.write(record); } public String getOutputURI() { return uri; } public Schema getAvroSchema() { return avroSchema; } public GenericRecordListBuilder getRecordListBuilder() { return recordListBuilder; } public List getWrittenRecords() { return Collections.unmodifiableList(writtenRecords); } @Override public void close() throws Exception { writer.close(); } public class GenericRecordListBuilder { public final List createRecordList(Object... values) { final int fieldCount = avroSchema.getFields().size(); Preconditions.checkArgument( values.length % fieldCount == 0, "arg count of values should be divide by field number"); final List recordList = new ArrayList<>(); for (int i = 0; i < values.length / fieldCount; i++) { final GenericRecord record = new GenericData.Record(avroSchema); for (int j = 0; j < fieldCount; j++) { record.put(j, values[i * fieldCount + j]); } recordList.add(record); } return Collections.unmodifiableList(recordList); } } } apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/arrow/dataset/TestAllTypes.java000066400000000000000000000330601515451615200312430ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.channels.Channels; import java.nio.channels.SeekableByteChannel; import java.nio.channels.WritableByteChannel; import java.nio.file.Files; import java.nio.file.Path; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.Objects; import org.apache.arrow.dataset.file.DatasetFileWriter; import org.apache.arrow.dataset.file.FileFormat; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.Float16; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.DateMilliVector; import org.apache.arrow.vector.Decimal256Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; 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.LargeVarBinaryVector; import org.apache.arrow.vector.LargeVarCharVector; import org.apache.arrow.vector.SmallIntVector; import org.apache.arrow.vector.TimeMilliVector; import org.apache.arrow.vector.TimeNanoVector; 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.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.LargeListVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.StructVector; 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.ipc.ArrowStreamReader; import org.apache.arrow.vector.ipc.ArrowStreamWriter; import org.apache.arrow.vector.test.util.ArrowTestDataUtil; 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.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel; import org.apache.arrow.vector.util.Text; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; public class TestAllTypes extends TestDataset { @TempDir public Path TMP; private VectorSchemaRoot generateAllTypesVector(BufferAllocator allocator) { // Notes: // - IntervalMonthDayNano is not supported by Parquet. // - Map (GH-38250) and SparseUnion are resulting in serialization errors when writing with the // Dataset API. // "Unhandled type for Arrow to Parquet schema conversion" errors: IntervalDay, IntervalYear, // DenseUnion List childFields = new ArrayList<>(); childFields.add( new Field("int-child", FieldType.notNullable(new ArrowType.Int(32, true)), null)); Field structField = new Field("struct", FieldType.nullable(ArrowType.Struct.INSTANCE), childFields); Field[] fields = new Field[] { Field.nullablePrimitive("null", ArrowType.Null.INSTANCE), Field.nullablePrimitive("bool", ArrowType.Bool.INSTANCE), Field.nullablePrimitive("int8", new ArrowType.Int(8, true)), Field.nullablePrimitive("int16", new ArrowType.Int(16, true)), Field.nullablePrimitive("int32", new ArrowType.Int(32, true)), Field.nullablePrimitive("int64", new ArrowType.Int(64, true)), Field.nullablePrimitive("uint8", new ArrowType.Int(8, false)), Field.nullablePrimitive("uint16", new ArrowType.Int(16, false)), Field.nullablePrimitive("uint32", new ArrowType.Int(32, false)), Field.nullablePrimitive("uint64", new ArrowType.Int(64, false)), Field.nullablePrimitive( "float16", new ArrowType.FloatingPoint(FloatingPointPrecision.HALF)), Field.nullablePrimitive( "float32", new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), Field.nullablePrimitive( "float64", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), Field.nullablePrimitive("utf8", ArrowType.Utf8.INSTANCE), Field.nullablePrimitive("binary", ArrowType.Binary.INSTANCE), Field.nullablePrimitive("largeutf8", ArrowType.LargeUtf8.INSTANCE), Field.nullablePrimitive("largebinary", ArrowType.LargeBinary.INSTANCE), Field.nullablePrimitive("fixed_size_binary", new ArrowType.FixedSizeBinary(1)), Field.nullablePrimitive("date_ms", new ArrowType.Date(DateUnit.MILLISECOND)), Field.nullablePrimitive("time_ms", new ArrowType.Time(TimeUnit.MILLISECOND, 32)), Field.nullablePrimitive( "timestamp_ms", new ArrowType.Timestamp(TimeUnit.MILLISECOND, null)), Field.nullablePrimitive( "timestamptz_ms", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), Field.nullablePrimitive("time_ns", new ArrowType.Time(TimeUnit.NANOSECOND, 64)), Field.nullablePrimitive( "timestamp_ns", new ArrowType.Timestamp(TimeUnit.NANOSECOND, null)), Field.nullablePrimitive( "timestamptz_ns", new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC")), Field.nullablePrimitive("duration", new ArrowType.Duration(TimeUnit.MILLISECOND)), Field.nullablePrimitive("decimal128", new ArrowType.Decimal(10, 2, 128)), Field.nullablePrimitive("decimal256", new ArrowType.Decimal(10, 2, 256)), new Field( "list", FieldType.nullable(ArrowType.List.INSTANCE), Collections.singletonList(Field.nullable("items", new ArrowType.Int(32, true)))), new Field( "largelist", FieldType.nullable(ArrowType.LargeList.INSTANCE), Collections.singletonList(Field.nullable("items", new ArrowType.Int(32, true)))), new Field( "fixedsizelist", FieldType.nullable(new ArrowType.FixedSizeList(2)), Collections.singletonList(Field.nullable("items", new ArrowType.Int(32, true)))), structField }; VectorSchemaRoot root = VectorSchemaRoot.create(new Schema(Arrays.asList(fields)), allocator); root.allocateNew(); root.setRowCount(2); root.getVector("null").setNull(0); root.getVector("bool").setNull(0); root.getVector("int8").setNull(0); root.getVector("int16").setNull(0); root.getVector("int32").setNull(0); root.getVector("int64").setNull(0); root.getVector("uint8").setNull(0); root.getVector("uint16").setNull(0); root.getVector("uint32").setNull(0); root.getVector("uint64").setNull(0); root.getVector("float16").setNull(0); root.getVector("float32").setNull(0); root.getVector("float64").setNull(0); root.getVector("utf8").setNull(0); root.getVector("binary").setNull(0); root.getVector("largeutf8").setNull(0); root.getVector("largebinary").setNull(0); root.getVector("fixed_size_binary").setNull(0); root.getVector("date_ms").setNull(0); root.getVector("time_ms").setNull(0); root.getVector("time_ns").setNull(0); root.getVector("timestamp_ms").setNull(0); root.getVector("timestamp_ns").setNull(0); root.getVector("timestamptz_ms").setNull(0); root.getVector("timestamptz_ns").setNull(0); root.getVector("duration").setNull(0); root.getVector("decimal128").setNull(0); root.getVector("decimal256").setNull(0); root.getVector("fixedsizelist").setNull(0); root.getVector("list").setNull(0); root.getVector("largelist").setNull(0); root.getVector("struct").setNull(0); root.getVector("null").setNull(1); ((BitVector) root.getVector("bool")).set(1, 1); ((TinyIntVector) root.getVector("int8")).set(1, 1); ((SmallIntVector) root.getVector("int16")).set(1, 1); ((IntVector) root.getVector("int32")).set(1, 1); ((BigIntVector) root.getVector("int64")).set(1, 1); ((UInt1Vector) root.getVector("uint8")).set(1, 1); ((UInt2Vector) root.getVector("uint16")).set(1, 1); ((UInt4Vector) root.getVector("uint32")).set(1, 1); ((UInt8Vector) root.getVector("uint64")).set(1, 1); ((Float2Vector) root.getVector("float16")).set(1, Float16.toFloat16(+32.875f)); ((Float4Vector) root.getVector("float32")).set(1, 1.0f); ((Float8Vector) root.getVector("float64")).set(1, 1.0); ((VarCharVector) root.getVector("utf8")).set(1, new Text("a")); ((VarBinaryVector) root.getVector("binary")).set(1, new byte[] {0x01}); ((LargeVarCharVector) root.getVector("largeutf8")).set(1, new Text("a")); ((LargeVarBinaryVector) root.getVector("largebinary")).set(1, new byte[] {0x01}); ((FixedSizeBinaryVector) root.getVector("fixed_size_binary")).set(1, new byte[] {0x01}); ((DateMilliVector) root.getVector("date_ms")).set(1, 0); ((TimeMilliVector) root.getVector("time_ms")).set(1, 0); ((TimeNanoVector) root.getVector("time_ns")).set(1, 0); ((TimeStampMilliVector) root.getVector("timestamp_ms")).set(1, 0); ((TimeStampNanoVector) root.getVector("timestamp_ns")).set(1, 0); ((TimeStampMilliTZVector) root.getVector("timestamptz_ms")).set(1, 0); ((TimeStampNanoTZVector) root.getVector("timestamptz_ns")).set(1, 0); ((DurationVector) root.getVector("duration")).set(1, 0); ((DecimalVector) root.getVector("decimal128")).set(1, 0); ((Decimal256Vector) root.getVector("decimal256")).set(1, 0); UnionFixedSizeListWriter fixedListWriter = ((FixedSizeListVector) root.getVector("fixedsizelist")).getWriter(); fixedListWriter.allocate(); fixedListWriter.setPosition(1); fixedListWriter.startList(); fixedListWriter.integer().writeInt(1); fixedListWriter.endList(); UnionListWriter listWriter = ((ListVector) root.getVector("list")).getWriter(); listWriter.allocate(); listWriter.setPosition(1); listWriter.startList(); listWriter.integer().writeInt(1); listWriter.endList(); UnionLargeListWriter largeListWriter = ((LargeListVector) root.getVector("largelist")).getWriter(); largeListWriter.allocate(); largeListWriter.setPosition(1); largeListWriter.startList(); largeListWriter.integer().writeInt(1); largeListWriter.endList(); 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; } private byte[] serializeFile(VectorSchemaRoot root) { try (ByteArrayOutputStream out = new ByteArrayOutputStream(); WritableByteChannel channel = Channels.newChannel(out); ArrowStreamWriter writer = new ArrowStreamWriter(root, null, channel)) { writer.start(); writer.writeBatch(); writer.end(); return out.toByteArray(); } catch (IOException e) { throw new IllegalArgumentException("Failed to serialize arrow file", e); } } @Test public void testAllTypesParquet() throws Exception { try (VectorSchemaRoot root = generateAllTypesVector(rootAllocator())) { byte[] featherData = serializeFile(root); try (SeekableByteChannel channel = new ByteArrayReadableSeekableByteChannel(featherData)) { try (ArrowStreamReader reader = new ArrowStreamReader(channel, rootAllocator())) { final Path writtenFolder = Files.createTempDirectory(TMP, "writtenFolder"); final String writtenParquet = writtenFolder.toUri().toString(); DatasetFileWriter.write(rootAllocator(), reader, FileFormat.PARQUET, writtenParquet); // Load the reference file from the test resources and write to a temporary file on the // OS. String referenceFile = ArrowTestDataUtil.getTestDataRoot() .resolve("parquet") .resolve("alltypes-java.parquet") .toUri() .toString(); assertParquetFileEquals( referenceFile, Objects.requireNonNull(writtenFolder.toFile().listFiles())[0].toURI().toString()); } } } } } apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/arrow/dataset/TestDataset.java000066400000000000000000000142061515451615200310740ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.IOException; import java.util.ArrayList; import java.util.Iterator; import java.util.List; import java.util.Spliterator; import java.util.Spliterators; import java.util.stream.Collectors; import java.util.stream.Stream; import java.util.stream.StreamSupport; 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.util.AutoCloseables; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; import org.apache.arrow.vector.compare.VectorEqualsVisitor; import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; public abstract class TestDataset { private BufferAllocator allocator = null; @BeforeEach public void setUp() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void tearDown() { allocator.close(); } protected BufferAllocator rootAllocator() { return allocator; } protected List collectResultFromFactory( DatasetFactory factory, ScanOptions options) { final Dataset dataset = factory.finish(); final Scanner scanner = dataset.newScan(options); try { final List ret = collectTaskData(scanner); AutoCloseables.close(scanner, dataset); return ret; } catch (RuntimeException e) { throw e; } catch (Exception e) { throw new RuntimeException(e); } } protected List collectTaskData(Scanner scan) { try (ArrowReader reader = scan.scanBatches()) { List batches = new ArrayList<>(); while (reader.loadNextBatch()) { VectorSchemaRoot root = reader.getVectorSchemaRoot(); final VectorUnloader unloader = new VectorUnloader(root); batches.add(unloader.getRecordBatch()); } return batches; } catch (IOException e) { throw new RuntimeException(e); } } protected Schema inferResultSchemaFromFactory(DatasetFactory factory, ScanOptions options) { final Dataset dataset = factory.finish(); final Scanner scanner = dataset.newScan(options); final Schema schema = scanner.schema(); try { AutoCloseables.close(scanner, dataset); } catch (Exception e) { throw new RuntimeException(e); } return schema; } protected void assertParquetFileEquals(String expectedURI, String actualURI) throws Exception { final FileSystemDatasetFactory expectedFactory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, expectedURI); final FileSystemDatasetFactory actualFactory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, actualURI); List expectedBatches = collectResultFromFactory(expectedFactory, new ScanOptions(new String[0], 100)); List actualBatches = collectResultFromFactory(actualFactory, new ScanOptions(new String[0], 100)); try (VectorSchemaRoot expectVsr = VectorSchemaRoot.create(expectedFactory.inspect(), rootAllocator()); VectorSchemaRoot actualVsr = VectorSchemaRoot.create(actualFactory.inspect(), rootAllocator())) { // compare ArrowRecordBatches assertEquals(expectedBatches.size(), actualBatches.size()); VectorLoader expectLoader = new VectorLoader(expectVsr); VectorLoader actualLoader = new VectorLoader(actualVsr); for (int i = 0; i < expectedBatches.size(); i++) { expectLoader.load(expectedBatches.get(i)); actualLoader.load(actualBatches.get(i)); for (int j = 0; j < expectVsr.getFieldVectors().size(); j++) { FieldVector vector = expectVsr.getFieldVectors().get(i); FieldVector otherVector = actualVsr.getFieldVectors().get(i); // TODO: ARROW-18140 Use VectorSchemaRoot#equals() method to compare assertTrue(VectorEqualsVisitor.vectorEquals(vector, otherVector)); } } } finally { AutoCloseables.close(expectedBatches, actualBatches); } } protected Stream stream(Iterable iterable) { return StreamSupport.stream(iterable.spliterator(), false); } protected List collect(Iterable iterable) { return stream(iterable).collect(Collectors.toList()); } protected Stream stream(Iterator iterator) { return StreamSupport.stream( Spliterators.spliteratorUnknownSize(iterator, Spliterator.ORDERED), false); } protected List collect(Iterator iterator) { return stream(iterator).collect(Collectors.toList()); } } apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/arrow/dataset/TestFragmentScanOptions.java000066400000000000000000000346201515451615200334350ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertThrows; import com.google.common.collect.ImmutableMap; import java.util.Arrays; import java.util.Collections; import java.util.Map; import java.util.Optional; import org.apache.arrow.c.ArrowSchema; import org.apache.arrow.c.CDataDictionaryProvider; import org.apache.arrow.c.Data; 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.scanner.csv.CsvConvertOptions; import org.apache.arrow.dataset.scanner.csv.CsvFragmentScanOptions; 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.ValueIterableVector; import org.apache.arrow.vector.ipc.ArrowReader; 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.Text; import org.hamcrest.collection.IsIterableContainingInOrder; import org.junit.jupiter.api.Test; public class TestFragmentScanOptions { private CsvFragmentScanOptions create( ArrowSchema cSchema, Map convertOptionsMap, Map readOptions, Map parseOptions) { CsvConvertOptions convertOptions = new CsvConvertOptions(convertOptionsMap); convertOptions.setArrowSchema(cSchema); return new CsvFragmentScanOptions(convertOptions, readOptions, parseOptions); } @Test public void testCsvConvertOptions() throws Exception { final Schema schema = new Schema( Arrays.asList( Field.nullable("Id", new ArrowType.Int(32, true)), Field.nullable("Name", new ArrowType.Utf8()), Field.nullable("Language", new ArrowType.Utf8())), null); String path = "file://" + getClass().getResource("/").getPath() + "/data/student.csv"; BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); try (ArrowSchema cSchema = ArrowSchema.allocateNew(allocator); ArrowSchema cSchema2 = ArrowSchema.allocateNew(allocator); CDataDictionaryProvider provider = new CDataDictionaryProvider()) { Data.exportSchema(allocator, schema, provider, cSchema); Data.exportSchema(allocator, schema, provider, cSchema2); CsvFragmentScanOptions fragmentScanOptions1 = create(cSchema, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of("delimiter", ";")); CsvFragmentScanOptions fragmentScanOptions2 = create(cSchema2, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of("delimiter", ";")); ScanOptions options = new ScanOptions.Builder(/*batchSize*/ 32768) .columns(Optional.empty()) .fragmentScanOptions(fragmentScanOptions1) .build(); try (DatasetFactory datasetFactory = new FileSystemDatasetFactory( allocator, NativeMemoryPool.getDefault(), FileFormat.CSV, path, Optional.of(fragmentScanOptions2)); Dataset dataset = datasetFactory.finish(); Scanner scanner = dataset.newScan(options); ArrowReader reader = scanner.scanBatches()) { assertEquals(schema.getFields(), reader.getVectorSchemaRoot().getSchema().getFields()); int rowCount = 0; while (reader.loadNextBatch()) { final ValueIterableVector idVector = (ValueIterableVector) reader.getVectorSchemaRoot().getVector("Id"); assertThat(idVector.getValueIterable(), IsIterableContainingInOrder.contains(1, 2, 3)); rowCount += reader.getVectorSchemaRoot().getRowCount(); } assertEquals(3, rowCount); } } } @Test public void testCsvConvertOptionsDelimiterNotSet() throws Exception { final Schema schema = new Schema( Arrays.asList( Field.nullable("Id", new ArrowType.Int(32, true)), Field.nullable("Name", new ArrowType.Utf8()), Field.nullable("Language", new ArrowType.Utf8())), null); String path = "file://" + getClass().getResource("/").getPath() + "/data/student.csv"; BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); try (ArrowSchema cSchema = ArrowSchema.allocateNew(allocator); ArrowSchema cSchema2 = ArrowSchema.allocateNew(allocator); CDataDictionaryProvider provider = new CDataDictionaryProvider()) { Data.exportSchema(allocator, schema, provider, cSchema); Data.exportSchema(allocator, schema, provider, cSchema2); CsvFragmentScanOptions fragmentScanOptions1 = create(cSchema, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of()); CsvFragmentScanOptions fragmentScanOptions2 = create(cSchema2, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of()); ScanOptions options = new ScanOptions.Builder(/*batchSize*/ 32768) .columns(Optional.empty()) .fragmentScanOptions(fragmentScanOptions1) .build(); try (DatasetFactory datasetFactory = new FileSystemDatasetFactory( allocator, NativeMemoryPool.getDefault(), FileFormat.CSV, path, Optional.of(fragmentScanOptions2)); Dataset dataset = datasetFactory.finish(); Scanner scanner = dataset.newScan(options); ArrowReader reader = scanner.scanBatches()) { int rowCount = 0; while (reader.loadNextBatch()) { final ValueIterableVector idVector = (ValueIterableVector) reader.getVectorSchemaRoot().getVector("Id;Name;Language"); assertThat( idVector.getValueIterable(), IsIterableContainingInOrder.contains( new Text("1;Juno;Java"), new Text("2;Peter;Python"), new Text("3;Celin;C++"))); rowCount += reader.getVectorSchemaRoot().getRowCount(); } assertEquals(3, rowCount); } } } @Test public void testCsvConvertOptionsNoOption() throws Exception { final Schema schema = new Schema( Collections.singletonList(Field.nullable("Id;Name;Language", new ArrowType.Utf8())), null); String path = "file://" + getClass().getResource("/").getPath() + "/data/student.csv"; BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); ScanOptions options = new ScanOptions.Builder(/*batchSize*/ 32768).columns(Optional.empty()).build(); try (DatasetFactory datasetFactory = new FileSystemDatasetFactory( allocator, NativeMemoryPool.getDefault(), FileFormat.CSV, path); Dataset dataset = datasetFactory.finish(); Scanner scanner = dataset.newScan(options); ArrowReader reader = scanner.scanBatches()) { assertEquals(schema.getFields(), reader.getVectorSchemaRoot().getSchema().getFields()); int rowCount = 0; while (reader.loadNextBatch()) { final ValueIterableVector idVector = (ValueIterableVector) reader.getVectorSchemaRoot().getVector("Id;Name;Language"); assertThat( idVector.getValueIterable(), IsIterableContainingInOrder.contains( new Text("1;Juno;Java"), new Text("2;Peter;Python"), new Text("3;Celin;C++"))); rowCount += reader.getVectorSchemaRoot().getRowCount(); } assertEquals(3, rowCount); } } @Test public void testCsvReadParseAndReadOptions() throws Exception { final Schema schema = new Schema( Arrays.asList( Field.nullable("Id", new ArrowType.Int(64, true)), Field.nullable("Name", new ArrowType.Utf8()), Field.nullable("Language", new ArrowType.Utf8())), null); String path = "file://" + getClass().getResource("/").getPath() + "/data/student.csv"; BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); CsvFragmentScanOptions fragmentScanOptions = new CsvFragmentScanOptions( new CsvConvertOptions(ImmutableMap.of()), ImmutableMap.of("skip_rows_after_names", "1"), ImmutableMap.of("delimiter", ";")); ScanOptions options = new ScanOptions.Builder(/*batchSize*/ 32768) .columns(Optional.empty()) .fragmentScanOptions(fragmentScanOptions) .build(); try (DatasetFactory datasetFactory = new FileSystemDatasetFactory( allocator, NativeMemoryPool.getDefault(), FileFormat.CSV, path, Optional.of(fragmentScanOptions)); Dataset dataset = datasetFactory.finish(); Scanner scanner = dataset.newScan(options); ArrowReader reader = scanner.scanBatches()) { assertEquals(schema.getFields(), reader.getVectorSchemaRoot().getSchema().getFields()); int rowCount = 0; while (reader.loadNextBatch()) { final ValueIterableVector idVector = (ValueIterableVector) reader.getVectorSchemaRoot().getVector("Id"); assertThat(idVector.getValueIterable(), IsIterableContainingInOrder.contains(2L, 3L)); rowCount += reader.getVectorSchemaRoot().getRowCount(); } assertEquals(2, rowCount); } } @Test public void testCsvReadOtherOptions() throws Exception { String path = "file://" + getClass().getResource("/").getPath() + "/data/student.csv"; BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); Map convertOption = ImmutableMap.of( "check_utf8", "true", "null_values", "NULL", "true_values", "True", "false_values", "False", "quoted_strings_can_be_null", "true", "auto_dict_encode", "false", "auto_dict_max_cardinality", "3456", "decimal_point", ".", "include_missing_columns", "false"); Map readOption = ImmutableMap.of( "use_threads", "true", "block_size", "1024", "skip_rows", "12", "skip_rows_after_names", "12", "autogenerate_column_names", "false"); Map parseOption = ImmutableMap.of( "delimiter", ".", "quoting", "true", "quote_char", "'", "double_quote", "False", "escaping", "true", "escape_char", "v", "newlines_in_values", "false", "ignore_empty_lines", "true"); CsvFragmentScanOptions fragmentScanOptions = new CsvFragmentScanOptions(new CsvConvertOptions(convertOption), readOption, parseOption); ScanOptions options = new ScanOptions.Builder(/*batchSize*/ 32768) .columns(Optional.empty()) .fragmentScanOptions(fragmentScanOptions) .build(); try (DatasetFactory datasetFactory = new FileSystemDatasetFactory( allocator, NativeMemoryPool.getDefault(), FileFormat.CSV, path); Dataset dataset = datasetFactory.finish(); Scanner scanner = dataset.newScan(options)) { assertNotNull(scanner); } } @Test public void testCsvInvalidOption() throws Exception { String path = "file://" + getClass().getResource("/").getPath() + "/data/student.csv"; BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); Map convertOption = ImmutableMap.of("not_exists_key_check_utf8", "true"); CsvFragmentScanOptions fragmentScanOptions = new CsvFragmentScanOptions( new CsvConvertOptions(convertOption), ImmutableMap.of(), ImmutableMap.of()); ScanOptions options = new ScanOptions.Builder(/*batchSize*/ 32768) .columns(Optional.empty()) .fragmentScanOptions(fragmentScanOptions) .build(); try (DatasetFactory datasetFactory = new FileSystemDatasetFactory( allocator, NativeMemoryPool.getDefault(), FileFormat.CSV, path); Dataset dataset = datasetFactory.finish()) { assertThrows(RuntimeException.class, () -> dataset.newScan(options)); } CsvFragmentScanOptions fragmentScanOptionsFaultValue = new CsvFragmentScanOptions( new CsvConvertOptions(ImmutableMap.of()), ImmutableMap.of("", ""), ImmutableMap.of("escape_char", "vbvb")); ScanOptions optionsFault = new ScanOptions.Builder(/*batchSize*/ 32768) .columns(Optional.empty()) .fragmentScanOptions(fragmentScanOptionsFaultValue) .build(); try (DatasetFactory datasetFactory = new FileSystemDatasetFactory( allocator, NativeMemoryPool.getDefault(), FileFormat.CSV, path); Dataset dataset = datasetFactory.finish()) { assertThrows(RuntimeException.class, () -> dataset.newScan(optionsFault)); } } } apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/arrow/dataset/TextBasedWriteSupport.java000066400000000000000000000042361515451615200331440ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset; import java.io.File; import java.io.IOException; import java.io.Writer; import java.net.URI; import java.net.URISyntaxException; import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.StandardOpenOption; import java.util.Random; public class TextBasedWriteSupport { private final URI uri; private final Random random = new Random(); public TextBasedWriteSupport(File outputFolder, String fileExtension) throws URISyntaxException { uri = new URI( "file", outputFolder.getPath() + File.separator + "generated-" + random.nextLong() + fileExtension, null); } public static TextBasedWriteSupport writeTempFile( File outputFolder, String fileExtension, String... values) throws URISyntaxException, IOException { TextBasedWriteSupport writer = new TextBasedWriteSupport(outputFolder, fileExtension); try (Writer addValues = Files.newBufferedWriter( new File(writer.uri).toPath(), StandardCharsets.UTF_8, StandardOpenOption.CREATE, StandardOpenOption.APPEND)) { for (Object value : values) { addValues.write(value + "\n"); } } return writer; } public String getOutputURI() { return uri.toString(); } } apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/arrow/dataset/file/000077500000000000000000000000001515451615200267205ustar00rootroot00000000000000TestDatasetFileWriter.java000066400000000000000000000106371515451615200337350ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/arrow/dataset/file/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.file; import static org.junit.jupiter.api.Assertions.assertEquals; import java.io.File; import java.util.Arrays; import java.util.HashSet; import java.util.Objects; import java.util.Set; import java.util.stream.Collectors; import org.apache.arrow.dataset.ParquetWriteSupport; import org.apache.arrow.dataset.TestDataset; import org.apache.arrow.dataset.jni.NativeMemoryPool; import org.apache.arrow.dataset.scanner.ArrowScannerReader; import org.apache.arrow.dataset.scanner.ScanOptions; import org.apache.arrow.dataset.scanner.Scanner; import org.apache.arrow.dataset.source.Dataset; import org.apache.commons.io.FileUtils; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; public class TestDatasetFileWriter extends TestDataset { @TempDir public File TMP; public static final String AVRO_SCHEMA_USER = "user.avsc"; @Test public void testParquetWriteSimple() throws Exception { ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a", 2, "b", 3, "c", 2, "d"); String sampleParquet = writeSupport.getOutputURI(); ScanOptions options = new ScanOptions(new String[0], 100); final File writtenFolder = new File(TMP, "writtenFolder"); writtenFolder.mkdirs(); final String writtenParquet = writtenFolder.toURI().toString(); try (FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, sampleParquet); final Dataset dataset = factory.finish(); final Scanner scanner = dataset.newScan(options); final ArrowScannerReader reader = new ArrowScannerReader(scanner, rootAllocator()); ) { DatasetFileWriter.write(rootAllocator(), reader, FileFormat.PARQUET, writtenParquet); assertParquetFileEquals( sampleParquet, Objects.requireNonNull(writtenFolder.listFiles())[0].toURI().toString()); } } @Test public void testParquetWriteWithPartitions() throws Exception { ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a", 2, "b", 3, "c", 2, "d"); String sampleParquet = writeSupport.getOutputURI(); ScanOptions options = new ScanOptions(new String[0], 100); final File writtenFolder = new File(TMP, "writtenFolder"); writtenFolder.mkdirs(); final String writtenParquet = writtenFolder.toURI().toString(); try (FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, sampleParquet); final Dataset dataset = factory.finish(); final Scanner scanner = dataset.newScan(options); final ArrowScannerReader reader = new ArrowScannerReader(scanner, rootAllocator()); ) { DatasetFileWriter.write( rootAllocator(), reader, FileFormat.PARQUET, writtenParquet, new String[] {"id", "name"}, 100, "data_{i}"); final Set expectedOutputFiles = new HashSet<>( Arrays.asList( "id=1/name=a/data_0", "id=2/name=b/data_0", "id=3/name=c/data_0", "id=2/name=d/data_0")); final Set outputFiles = FileUtils.listFiles(writtenFolder, null, true).stream() .map( file -> { return writtenFolder.toURI().relativize(file.toURI()).toString(); }) .collect(Collectors.toSet()); assertEquals(expectedOutputFiles, outputFiles); } } } TestFileSystemDataset.java000066400000000000000000000566071515451615200337540ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/arrow/dataset/file/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.file; 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 com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.primitives.Primitives; import java.io.File; import java.io.FileOutputStream; import java.io.IOException; import java.util.ArrayList; import java.util.Collections; import java.util.HashSet; import java.util.List; import java.util.Objects; import java.util.Optional; import java.util.Set; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.stream.Collectors; import org.apache.arrow.dataset.OrcWriteSupport; import org.apache.arrow.dataset.ParquetWriteSupport; import org.apache.arrow.dataset.TextBasedWriteSupport; import org.apache.arrow.dataset.jni.NativeDataset; import org.apache.arrow.dataset.jni.NativeInstanceReleasedException; import org.apache.arrow.dataset.jni.NativeMemoryPool; import org.apache.arrow.dataset.jni.NativeScanner; import org.apache.arrow.dataset.jni.TestNativeDataset; import org.apache.arrow.dataset.scanner.ScanOptions; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.ArrowFileWriter; import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; 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.Schema; import org.apache.avro.generic.GenericRecord; import org.apache.avro.generic.GenericRecordBuilder; import org.apache.hadoop.fs.Path; import org.apache.orc.TypeDescription; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; public class TestFileSystemDataset extends TestNativeDataset { @TempDir public File TMP; public static final String AVRO_SCHEMA_USER = "user.avsc"; @Test public void testBaseParquetRead() throws Exception { ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a"); FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); ScanOptions options = new ScanOptions(100); Schema schema = inferResultSchemaFromFactory(factory, options); List datum = collectResultFromFactory(factory, options); assertScanBatchesProduced(factory, options); assertEquals(1, datum.size()); assertEquals(2, schema.getFields().size()); assertEquals("id", schema.getFields().get(0).getName()); assertEquals("name", schema.getFields().get(1).getName()); assertEquals(Types.MinorType.INT.getType(), schema.getFields().get(0).getType()); assertEquals(Types.MinorType.VARCHAR.getType(), schema.getFields().get(1).getType()); checkParquetReadResult(schema, writeSupport.getWrittenRecords(), datum); AutoCloseables.close(datum); AutoCloseables.close(factory); } @Test public void testMultipleParquetReadFromUris() throws Exception { ParquetWriteSupport writeSupport1 = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a"); ParquetWriteSupport writeSupport2 = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 2, "b"); String expectedJsonUnordered = "[[1,\"a\"],[2,\"b\"]]"; ScanOptions options = new ScanOptions(1); FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, new String[] {writeSupport1.getOutputURI(), writeSupport2.getOutputURI()}); Schema schema = inferResultSchemaFromFactory(factory, options); List datum = collectResultFromFactory(factory, options); assertScanBatchesProduced(factory, options); assertEquals(2, datum.size()); datum.forEach(batch -> assertEquals(1, batch.getLength())); checkParquetReadResult(schema, expectedJsonUnordered, datum); AutoCloseables.close(datum); AutoCloseables.close(factory); } @Test public void testMultipleParquetInvalidUri() throws Exception { RuntimeException exc = assertThrows( RuntimeException.class, () -> new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, new String[] {"https://example.com", "file:///test/location"})); assertEquals("Unrecognized filesystem type in URI: https://example.com", exc.getMessage()); } @Test public void testMultipleParquetMultipleFilesystemTypes() throws Exception { RuntimeException exc = assertThrows( RuntimeException.class, () -> new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, new String[] {"file:///test/location", "s3:///test/bucket/file"})); assertTrue( exc.getMessage() .startsWith( "The filesystem expected a URI with one of the schemes (file) but received s3")); } @Test public void testParquetProjectSingleColumn() throws Exception { ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a"); FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); ScanOptions options = new ScanOptions(100, Optional.of(new String[] {"id"})); Schema schema = inferResultSchemaFromFactory(factory, options); List datum = collectResultFromFactory(factory, options); org.apache.avro.Schema expectedSchema = truncateAvroSchema(writeSupport.getAvroSchema(), 0, 1); assertScanBatchesProduced(factory, options); assertEquals(1, schema.getFields().size()); assertEquals("id", schema.getFields().get(0).getName()); assertEquals(Types.MinorType.INT.getType(), schema.getFields().get(0).getType()); assertEquals(1, datum.size()); checkParquetReadResult( schema, Collections.singletonList(new GenericRecordBuilder(expectedSchema).set("id", 1).build()), datum); AutoCloseables.close(datum); AutoCloseables.close(factory); } @Test public void testParquetBatchSize() throws Exception { ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a", 2, "b", 3, "c"); ScanOptions options = new ScanOptions(1); FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); Schema schema = inferResultSchemaFromFactory(factory, options); List datum = collectResultFromFactory(factory, options); assertScanBatchesProduced(factory, options); assertEquals(3, datum.size()); datum.forEach(batch -> assertEquals(1, batch.getLength())); checkParquetReadResult(schema, writeSupport.getWrittenRecords(), datum); AutoCloseables.close(datum); AutoCloseables.close(factory); } @Test public void testParquetDirectoryRead() throws Exception { final File outputFolder = TMP; ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, outputFolder, 1, "a", 2, "b", 3, "c"); ParquetWriteSupport.writeTempFile( AVRO_SCHEMA_USER, outputFolder, 4, "e", 5, "f", 6, "g", 7, "h"); String expectedJsonUnordered = "[[1,\"a\"],[2,\"b\"],[3,\"c\"],[4,\"e\"],[5,\"f\"],[6,\"g\"],[7,\"h\"]]"; ScanOptions options = new ScanOptions(new String[0], 1); FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, outputFolder.toURI().toString()); Schema schema = inferResultSchemaFromFactory(factory, options); List datum = collectResultFromFactory(factory, options); assertScanBatchesProduced(factory, options); assertEquals(7, datum.size()); datum.forEach(batch -> assertEquals(1, batch.getLength())); checkParquetReadResult(schema, expectedJsonUnordered, datum); AutoCloseables.close(datum); } @Test public void testEmptyProjectSelectsZeroColumns() throws Exception { ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a"); FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); ScanOptions options = new ScanOptions(100, Optional.of(new String[0])); Schema schema = inferResultSchemaFromFactory(factory, options); List datum = collectResultFromFactory(factory, options); org.apache.avro.Schema expectedSchema = org.apache.avro.Schema.createRecord(Collections.emptyList()); assertScanBatchesProduced(factory, options); assertEquals(0, schema.getFields().size()); assertEquals(1, datum.size()); checkParquetReadResult( schema, Collections.singletonList(new GenericRecordBuilder(expectedSchema).build()), datum); AutoCloseables.close(datum); } @Test public void testNullProjectSelectsAllColumns() throws Exception { ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a"); FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); ScanOptions options = new ScanOptions(100, Optional.empty()); Schema schema = inferResultSchemaFromFactory(factory, options); List datum = collectResultFromFactory(factory, options); assertScanBatchesProduced(factory, options); assertEquals(1, datum.size()); assertEquals(2, schema.getFields().size()); assertEquals("id", schema.getFields().get(0).getName()); assertEquals("name", schema.getFields().get(1).getName()); assertEquals(Types.MinorType.INT.getType(), schema.getFields().get(0).getType()); assertEquals(Types.MinorType.VARCHAR.getType(), schema.getFields().get(1).getType()); checkParquetReadResult(schema, writeSupport.getWrittenRecords(), datum); AutoCloseables.close(datum); } @Test public void testNoErrorWhenCloseAgain() throws Exception { ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a"); FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); assertDoesNotThrow( () -> { NativeDataset dataset = factory.finish(); dataset.close(); dataset.close(); }); AutoCloseables.close(factory); } @Test public void testErrorThrownWhenScanBatchesAgain() throws Exception { ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a"); FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); NativeDataset dataset = factory.finish(); ScanOptions options = new ScanOptions(100); NativeScanner scanner = dataset.newScan(options); List datum = collectTaskData(scanner); AutoCloseables.close(datum); UnsupportedOperationException uoe = assertThrows(UnsupportedOperationException.class, scanner::scanBatches); assertEquals( "NativeScanner can only be executed once. Create a new scanner instead", uoe.getMessage()); AutoCloseables.close(scanner, dataset, factory); } @Test public void testScanBatchesInOtherThread() throws Exception { ExecutorService executor = Executors.newSingleThreadExecutor(); ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a"); FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); NativeDataset dataset = factory.finish(); ScanOptions options = new ScanOptions(100); NativeScanner scanner = dataset.newScan(options); List datum = executor.submit(() -> collectTaskData(scanner)).get(); AutoCloseables.close(datum); AutoCloseables.close(scanner, dataset, factory); } @Test public void testErrorThrownWhenScanBatchesAfterScannerClose() throws Exception { ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a"); FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); NativeDataset dataset = factory.finish(); ScanOptions options = new ScanOptions(100); NativeScanner scanner = dataset.newScan(options); scanner.close(); assertThrows(NativeInstanceReleasedException.class, scanner::scanBatches); AutoCloseables.close(factory); } @Test public void testErrorThrownWhenReadAfterNativeReaderClose() throws Exception { ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a"); FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); NativeDataset dataset = factory.finish(); ScanOptions options = new ScanOptions(100); NativeScanner scanner = dataset.newScan(options); ArrowReader reader = scanner.scanBatches(); scanner.close(); assertThrows(NativeInstanceReleasedException.class, reader::loadNextBatch); AutoCloseables.close(factory); } @Test public void testBaseArrowIpcRead() throws Exception { File dataFile = new File(TMP, "datafile"); Schema sourceSchema = new Schema(Collections.singletonList(Field.nullable("ints", new ArrowType.Int(32, true)))); try (VectorSchemaRoot root = VectorSchemaRoot.create(sourceSchema, rootAllocator()); FileOutputStream sink = new FileOutputStream(dataFile); ArrowFileWriter writer = new ArrowFileWriter(root, /* provider= */ null, sink.getChannel())) { IntVector ints = (IntVector) root.getVector(0); ints.setSafe(0, 0); ints.setSafe(1, 1024); ints.setSafe(2, Integer.MAX_VALUE); root.setRowCount(3); writer.start(); writer.writeBatch(); writer.end(); } String arrowDataURI = dataFile.toURI().toString(); FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.ARROW_IPC, arrowDataURI); ScanOptions options = new ScanOptions(100); Schema schema = inferResultSchemaFromFactory(factory, options); List datum = collectResultFromFactory(factory, options); assertScanBatchesProduced(factory, options); assertEquals(1, datum.size()); assertEquals(1, schema.getFields().size()); assertEquals("ints", schema.getFields().get(0).getName()); String expectedJsonUnordered = String.format("[[0],[1024],[%d]]", Integer.MAX_VALUE); checkParquetReadResult(schema, expectedJsonUnordered, datum); AutoCloseables.close(datum); AutoCloseables.close(factory); } @Test public void testBaseOrcRead() throws Exception { String dataName = "test-orc"; String basePath = TMP.getAbsolutePath(); TypeDescription orcSchema = TypeDescription.fromString("struct"); Path path = new Path(basePath, dataName); OrcWriteSupport.writeTempFile( orcSchema, path, new Integer[] {Integer.MIN_VALUE, Integer.MAX_VALUE}); String orcDatasetUri = new File(basePath, dataName).toURI().toString(); FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.ORC, orcDatasetUri); ScanOptions options = new ScanOptions(100); Schema schema = inferResultSchemaFromFactory(factory, options); List datum = collectResultFromFactory(factory, options); assertScanBatchesProduced(factory, options); assertEquals(1, datum.size()); assertEquals(1, schema.getFields().size()); assertEquals("ints", schema.getFields().get(0).getName()); String expectedJsonUnordered = "[[2147483647], [-2147483648]]"; checkParquetReadResult(schema, expectedJsonUnordered, datum); AutoCloseables.close(datum); AutoCloseables.close(factory); } @Test public void testBaseCsvRead() throws Exception { TextBasedWriteSupport writeSupport = TextBasedWriteSupport.writeTempFile( TMP, ".csv", "Name,Language", "Juno,Java", "Peter,Python", "Celin,C++"); String expectedJsonUnordered = "[[\"Juno\", \"Java\"], [\"Peter\", \"Python\"], [\"Celin\", \"C++\"]]"; ScanOptions options = new ScanOptions(100); try (FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.CSV, writeSupport.getOutputURI())) { List datum = collectResultFromFactory(factory, options); Schema schema = inferResultSchemaFromFactory(factory, options); assertScanBatchesProduced(factory, options); assertEquals(1, datum.size()); assertEquals(2, schema.getFields().size()); assertEquals("Name", schema.getFields().get(0).getName()); checkParquetReadResult(schema, expectedJsonUnordered, datum); AutoCloseables.close(datum); } } @Test public void testBaseJsonRead() throws Exception { TextBasedWriteSupport writeSupport = TextBasedWriteSupport.writeTempFile( TMP, ".json", "{\"Type\": \"Compiled\", \"Language\": \"Java\"}", "{\"Type\": \"Interpreted\", \"Language\": \"Python\"}"); String expectedJsonUnordered = "[[\"Compiled\", \"Java\"], " + "[\"Interpreted\", \"Python\"]]"; ScanOptions options = new ScanOptions(100); try (FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.JSON, writeSupport.getOutputURI())) { List datum = collectResultFromFactory(factory, options); Schema schema = inferResultSchemaFromFactory(factory, options); assertScanBatchesProduced(factory, options); assertEquals(1, datum.size()); assertEquals(2, schema.getFields().size()); assertEquals("Type", schema.getFields().get(0).getName()); assertEquals("Language", schema.getFields().get(1).getName()); checkParquetReadResult(schema, expectedJsonUnordered, datum); AutoCloseables.close(datum); } } private void checkParquetReadResult( Schema schema, String expectedJson, List actual) throws IOException { final ObjectMapper json = new ObjectMapper(); final Set expectedSet = json.readValue(expectedJson, Set.class); final Set> actualSet = new HashSet<>(); final int fieldCount = schema.getFields().size(); try (VectorSchemaRoot vsr = VectorSchemaRoot.create(schema, rootAllocator())) { VectorLoader loader = new VectorLoader(vsr); for (ArrowRecordBatch batch : actual) { loader.load(batch); int batchRowCount = vsr.getRowCount(); for (int i = 0; i < batchRowCount; i++) { List row = new ArrayList<>(); for (int j = 0; j < fieldCount; j++) { Object object = vsr.getVector(j).getObject(i); if (Primitives.isWrapperType(object.getClass())) { row.add(object); } else { row.add(object.toString()); } } actualSet.add(row); } } } assertEquals( expectedSet, actualSet, "Mismatched data read from Parquet, actual: " + json.writeValueAsString(actualSet) + ";"); } private void checkParquetReadResult( Schema schema, List expected, List actual) { assertEquals(expected.size(), actual.stream().mapToInt(ArrowRecordBatch::getLength).sum()); final int fieldCount = schema.getFields().size(); ArrayList expectedRemovable = new ArrayList<>(expected); try (VectorSchemaRoot vsr = VectorSchemaRoot.create(schema, rootAllocator())) { VectorLoader loader = new VectorLoader(vsr); for (ArrowRecordBatch batch : actual) { assertEquals(fieldCount, batch.getNodes().size()); loader.load(batch); int batchRowCount = vsr.getRowCount(); for (int i = 0; i < fieldCount; i++) { FieldVector vector = vsr.getVector(i); for (int j = 0; j < batchRowCount; j++) { Object object = vector.getObject(j); Object expectedObject = expectedRemovable.get(j).get(i); assertEquals(Objects.toString(expectedObject), Objects.toString(object)); } } for (int i = 0; i < batchRowCount; i++) { expectedRemovable.remove(0); } } assertTrue(expectedRemovable.isEmpty()); } } private org.apache.avro.Schema truncateAvroSchema( org.apache.avro.Schema schema, int from, int to) { List fields = schema.getFields().subList(from, to); return org.apache.avro.Schema.createRecord( fields.stream() .map( f -> new org.apache.avro.Schema.Field( f.name(), f.schema(), f.doc(), f.defaultVal(), f.order())) .collect(Collectors.toList())); } } TestFileSystemDatasetFactory.java000066400000000000000000000040611515451615200352670ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/arrow/dataset/file/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.file; 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 org.apache.arrow.dataset.jni.NativeMemoryPool; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.Test; public class TestFileSystemDatasetFactory { @Test public void testErrorHandling() { RuntimeException e = assertThrows( RuntimeException.class, () -> { new FileSystemDatasetFactory( new RootAllocator(Long.MAX_VALUE), NativeMemoryPool.getDefault(), FileFormat.NONE, "file:///NON_EXIST_FILE"); }); assertEquals("illegal file format id: -1", e.getMessage()); } @Test public void testCloseAgain() { assertDoesNotThrow( () -> { FileSystemDatasetFactory factory = new FileSystemDatasetFactory( new RootAllocator(Long.MAX_VALUE), NativeMemoryPool.getDefault(), FileFormat.PARQUET, "file:///NON_EXIST_FILE"); factory.close(); factory.close(); }); } } apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/arrow/dataset/jni/000077500000000000000000000000001515451615200265615ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/arrow/dataset/jni/TestNativeDataset.java000066400000000000000000000026561515451615200330310ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.jni; import static org.junit.jupiter.api.Assertions.assertNotNull; import org.apache.arrow.dataset.TestDataset; 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; public abstract class TestNativeDataset extends TestDataset { protected void assertScanBatchesProduced(DatasetFactory factory, ScanOptions options) { final Dataset dataset = factory.finish(); final Scanner scanner = dataset.newScan(options); assertNotNull(scanner.scanBatches()); } } TestReservationListener.java000066400000000000000000000123111515451615200342120ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/arrow/dataset/jni/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.jni; 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.io.File; import java.util.List; import java.util.concurrent.atomic.AtomicLong; import org.apache.arrow.dataset.ParquetWriteSupport; import org.apache.arrow.dataset.TestDataset; import org.apache.arrow.dataset.file.FileFormat; import org.apache.arrow.dataset.file.FileSystemDatasetFactory; import org.apache.arrow.dataset.scanner.ScanOptions; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; public class TestReservationListener extends TestDataset { @TempDir public File TMP; public static final String AVRO_SCHEMA_USER = "user.avsc"; @Test public void testDirectReservationListener() throws Exception { ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a"); NativeMemoryPool pool = NativeMemoryPool.createListenable(DirectReservationListener.instance()); FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), pool, FileFormat.PARQUET, writeSupport.getOutputURI()); ScanOptions options = new ScanOptions(100); long initReservation = DirectReservationListener.instance().getCurrentDirectMemReservation(); List datum = collectResultFromFactory(factory, options); long reservation = DirectReservationListener.instance().getCurrentDirectMemReservation(); AutoCloseables.close(datum); AutoCloseables.close(pool); long finalReservation = DirectReservationListener.instance().getCurrentDirectMemReservation(); assertTrue(reservation >= initReservation); assertEquals(initReservation, finalReservation); } @Test public void testCustomReservationListener() throws Exception { ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a"); final 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( rootAllocator(), pool, FileFormat.PARQUET, writeSupport.getOutputURI()); ScanOptions options = new ScanOptions(100); long initReservation = reserved.get(); List datum = collectResultFromFactory(factory, options); long reservation = reserved.get(); AutoCloseables.close(datum); AutoCloseables.close(pool); long finalReservation = reserved.get(); assertTrue(reservation >= initReservation); assertEquals(initReservation, finalReservation); } @Test @SuppressWarnings("UnnecessaryAsync") public void testErrorThrownFromReservationListener() throws Exception { final String errorMessage = "ERROR_MESSAGE"; ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a"); final AtomicLong reserved = new AtomicLong(0L); ReservationListener listener = new ReservationListener() { @Override public void reserve(long size) { throw new IllegalArgumentException(errorMessage); } @Override public void unreserve(long size) { // no-op } }; NativeMemoryPool pool = NativeMemoryPool.createListenable(listener); FileSystemDatasetFactory factory = new FileSystemDatasetFactory( rootAllocator(), pool, FileFormat.PARQUET, writeSupport.getOutputURI()); ScanOptions options = new ScanOptions(100); long initReservation = reserved.get(); assertThrows( IllegalArgumentException.class, () -> { collectResultFromFactory(factory, options); }, errorMessage); long reservation = reserved.get(); AutoCloseables.close(pool); long finalReservation = reserved.get(); assertEquals(initReservation, reservation); assertEquals(initReservation, finalReservation); } } apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/arrow/dataset/substrait/000077500000000000000000000000001515451615200300215ustar00rootroot00000000000000TestAceroSubstraitConsumer.java000066400000000000000000000641571515451615200361300ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/test/java/org/apache/arrow/dataset/substrait/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dataset.substrait; import static org.hamcrest.MatcherAssert.assertThat; 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.io.File; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Paths; import java.util.Arrays; import java.util.Base64; import java.util.Collections; import java.util.HashMap; import java.util.Map; import java.util.Optional; import org.apache.arrow.dataset.ParquetWriteSupport; import org.apache.arrow.dataset.TestDataset; 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.vector.ValueIterableVector; import org.apache.arrow.vector.ipc.ArrowReader; 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.hamcrest.collection.IsIterableContainingInOrder; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; public class TestAceroSubstraitConsumer extends TestDataset { @TempDir public File TMP; public static final String AVRO_SCHEMA_USER = "user.avsc"; @Test public void testRunQueryLocalFiles() throws Exception { // Query: // SELECT id, name FROM Users // Isthmus: // ./isthmus-macOS-0.7.0 -c "CREATE TABLE USERS ( id INT NOT NULL, name VARCHAR(150));" "SELECT // id, name FROM Users" // VARCHAR(150) -> is mapping to -> {ARROW:extension:name=varchar, // ARROW:extension:metadata=varchar{length:150}} Map metadataName = new HashMap<>(); metadataName.put("ARROW:extension:name", "varchar"); metadataName.put("ARROW:extension:metadata", "varchar{length:150}"); final Schema schema = new Schema( Arrays.asList( Field.nullable("ID", new ArrowType.Int(32, true)), new Field( "NAME", new FieldType(true, new ArrowType.Utf8(), null, metadataName), null)), Collections.emptyMap()); ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a", 11, "b", 21, "c"); try (ArrowReader arrowReader = new AceroSubstraitConsumer(rootAllocator()) .runQuery( new String( Files.readAllBytes( Paths.get( TestAceroSubstraitConsumer.class .getClassLoader() .getResource("substrait/local_files_users.json") .toURI())), StandardCharsets.UTF_8) .replace("FILENAME_PLACEHOLDER", writeSupport.getOutputURI()))) { assertEquals(schema, arrowReader.getVectorSchemaRoot().getSchema()); int rowcount = 0; while (arrowReader.loadNextBatch()) { rowcount += arrowReader.getVectorSchemaRoot().getRowCount(); } assertEquals(3, rowcount); } } @Test public void testRunQueryNamedTable() throws Exception { // Query: // SELECT id, name FROM Users // Isthmus: // ./isthmus-macOS-0.7.0 -c "CREATE TABLE USERS ( id INT NOT NULL, name VARCHAR(150));" "SELECT // id, name FROM Users" final Schema schema = new Schema( Arrays.asList( Field.nullable("ID", new ArrowType.Int(32, true)), Field.nullable("NAME", new ArrowType.Utf8())), Collections.emptyMap()); ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a", 11, "b", 21, "c"); ScanOptions options = new ScanOptions(/*batchSize*/ 32768); try (DatasetFactory datasetFactory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); Dataset dataset = datasetFactory.finish(); Scanner scanner = dataset.newScan(options); ArrowReader reader = scanner.scanBatches()) { Map mapTableToArrowReader = new HashMap<>(); mapTableToArrowReader.put("USERS", reader); try (ArrowReader arrowReader = new AceroSubstraitConsumer(rootAllocator()) .runQuery( new String( Files.readAllBytes( Paths.get( TestAceroSubstraitConsumer.class .getClassLoader() .getResource("substrait/named_table_users.json") .toURI())), StandardCharsets.UTF_8), mapTableToArrowReader)) { assertEquals(schema, arrowReader.getVectorSchemaRoot().getSchema()); assertEquals(arrowReader.getVectorSchemaRoot().getSchema(), schema); int rowcount = 0; while (arrowReader.loadNextBatch()) { rowcount += arrowReader.getVectorSchemaRoot().getRowCount(); } assertEquals(3, rowcount); } } } @Test public void testRunQueryNamedTableWithException() throws Exception { // Query: // SELECT id, name FROM Users // Isthmus: // ./isthmus-macOS-0.7.0 -c "CREATE TABLE USERS ( id INT NOT NULL, name VARCHAR(150));" "SELECT // id, name FROM Users" final Schema schema = new Schema( Arrays.asList( Field.nullable("ID", new ArrowType.Int(32, true)), Field.nullable("NAME", new ArrowType.Utf8())), Collections.emptyMap()); ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a", 11, "b", 21, "c"); ScanOptions options = new ScanOptions(/*batchSize*/ 32768); try (DatasetFactory datasetFactory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); Dataset dataset = datasetFactory.finish(); Scanner scanner = dataset.newScan(options); ArrowReader reader = scanner.scanBatches()) { Map mapTableToArrowReader = new HashMap<>(); mapTableToArrowReader.put("USERS_INVALID_MAP", reader); assertThrows( RuntimeException.class, () -> { try (ArrowReader arrowReader = new AceroSubstraitConsumer(rootAllocator()) .runQuery( new String( Files.readAllBytes( Paths.get( TestAceroSubstraitConsumer.class .getClassLoader() .getResource("substrait/named_table_users.json") .toURI())), StandardCharsets.UTF_8), mapTableToArrowReader)) { assertEquals(schema, arrowReader.getVectorSchemaRoot().getSchema()); int rowcount = 0; while (arrowReader.loadNextBatch()) { rowcount += arrowReader.getVectorSchemaRoot().getRowCount(); } assertEquals(3, rowcount); } }); } } @Test public void testRunBinaryQueryNamedTable() throws Exception { // Query: // SELECT id, name FROM Users // Isthmus: // ./isthmus-macOS-0.7.0 -c "CREATE TABLE USERS ( id INT NOT NULL, name VARCHAR(150));" "SELECT // id, name FROM Users" final Schema schema = new Schema( Arrays.asList( Field.nullable("ID", new ArrowType.Int(32, true)), Field.nullable("NAME", new ArrowType.Utf8())), Collections.emptyMap()); // Base64.getEncoder().encodeToString(plan.toByteArray()); String binaryPlan = "Gl8SXQpROk8KBhIECgICAxIvCi0KAgoAEh4KAklECgROQU1FEhIKBCoCEAEKC" + "LIBBQiWARgBGAI6BwoFVVNFUlMaCBIGCgISACIAGgoSCAoEEgIIASIAEgJJRBIETkFNRQ=="; ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a", 11, "b", 21, "c"); ScanOptions options = new ScanOptions(/*batchSize*/ 32768); try (DatasetFactory datasetFactory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); Dataset dataset = datasetFactory.finish(); Scanner scanner = dataset.newScan(options); ArrowReader reader = scanner.scanBatches()) { // map table to reader Map mapTableToArrowReader = new HashMap<>(); mapTableToArrowReader.put("USERS", reader); // get binary plan ByteBuffer substraitPlan = getByteBuffer(binaryPlan); // run query try (ArrowReader arrowReader = new AceroSubstraitConsumer(rootAllocator()) .runQuery(substraitPlan, mapTableToArrowReader)) { assertEquals(schema, arrowReader.getVectorSchemaRoot().getSchema()); int rowcount = 0; while (arrowReader.loadNextBatch()) { rowcount += arrowReader.getVectorSchemaRoot().getRowCount(); } assertEquals(3, rowcount); } } } @Test public void testRunExtendedExpressionsFilter() throws Exception { final Schema schema = new Schema( Arrays.asList( Field.nullable("id", new ArrowType.Int(32, true)), Field.nullable("name", new ArrowType.Utf8())), null); // Substrait Extended Expression: Filter: // Expression 01: WHERE ID < 20 String base64EncodedSubstraitFilter = "Ch4IARIaL2Z1bmN0aW9uc19jb21wYXJpc29uLnlhbWwSEhoQCAIQAhoKbHQ6YW55X2F" + "ueRo3ChwaGggCGgQKAhABIggaBhIECgISACIGGgQKAigUGhdmaWx0ZXJfaWRfbG93ZXJfdGhhbl8yMCIaCgJJRAoETkFNRRIOCgQqAhA" + "BCgRiAhABGAI="; ByteBuffer substraitExpressionFilter = getByteBuffer(base64EncodedSubstraitFilter); ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile( AVRO_SCHEMA_USER, TMP, 19, "value_19", 1, "value_1", 11, "value_11", 21, "value_21", 45, "value_45"); ScanOptions options = new ScanOptions.Builder(/*batchSize*/ 32768) .columns(Optional.empty()) .substraitFilter(substraitExpressionFilter) .build(); try (DatasetFactory datasetFactory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); Dataset dataset = datasetFactory.finish(); Scanner scanner = dataset.newScan(options); ArrowReader reader = scanner.scanBatches()) { assertEquals(schema.getFields(), reader.getVectorSchemaRoot().getSchema().getFields()); int rowcount = 0; while (reader.loadNextBatch()) { rowcount += reader.getVectorSchemaRoot().getRowCount(); final ValueIterableVector idVector = (ValueIterableVector) reader.getVectorSchemaRoot().getVector("id"); assertThat(idVector.getValueIterable(), IsIterableContainingInOrder.contains(19, 1, 11)); final ValueIterableVector nameVector = (ValueIterableVector) reader.getVectorSchemaRoot().getVector("name"); assertThat( nameVector.getValueIterable(), IsIterableContainingInOrder.contains( new Text("value_19"), new Text("value_1"), new Text("value_11"))); } assertEquals(3, rowcount); } } @Test public void testRunExtendedExpressionsFilterWithProjectionsInsteadOfFilterException() throws Exception { // Substrait Extended Expression: Project New Column: // Expression ADD: id + 2 // Expression CONCAT: name + '-' + name String base64EncodedSubstraitFilter = "Ch4IARIaL2Z1bmN0aW9uc19hcml0aG1ldGljLnlhbWwSERoPCAEaC2FkZDppM" + "zJfaTMyEhQaEggCEAEaDGNvbmNhdDp2Y2hhchoxChoaGBoEKgIQASIIGgYSBAoCEgAiBhoECgIoAhoTYWRkX3R3b190b19jb2x1" + "bW5fYRpGCi0aKwgBGgRiAhABIgoaCBIGCgQSAggBIgkaBwoFYgMgLSAiChoIEgYKBBICCAEaFWNvbmNhdF9jb2x1bW5fYV9hbmR" + "fYiIaCgJJRAoETkFNRRIOCgQqAhABCgRiAhABGAI="; ByteBuffer substraitExpressionFilter = getByteBuffer(base64EncodedSubstraitFilter); ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile( AVRO_SCHEMA_USER, TMP, 19, "value_19", 1, "value_1", 11, "value_11", 21, "value_21", 45, "value_45"); ScanOptions options = new ScanOptions.Builder(/*batchSize*/ 32768) .columns(Optional.empty()) .substraitFilter(substraitExpressionFilter) .build(); try (DatasetFactory datasetFactory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); Dataset dataset = datasetFactory.finish()) { Exception e = assertThrows(RuntimeException.class, () -> dataset.newScan(options)); assertTrue( e.getMessage().startsWith("There is no filter expression in the expression provided")); } } @Test public void testRunExtendedExpressionsFilterWithEmptyFilterException() throws Exception { String base64EncodedSubstraitFilter = ""; ByteBuffer substraitExpressionFilter = getByteBuffer(base64EncodedSubstraitFilter); ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile( AVRO_SCHEMA_USER, TMP, 19, "value_19", 1, "value_1", 11, "value_11", 21, "value_21", 45, "value_45"); ScanOptions options = new ScanOptions.Builder(/*batchSize*/ 32768) .columns(Optional.empty()) .substraitFilter(substraitExpressionFilter) .build(); try (DatasetFactory datasetFactory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); Dataset dataset = datasetFactory.finish()) { Exception e = assertThrows(RuntimeException.class, () -> dataset.newScan(options)); assertTrue( e.getMessage() .contains("no anonymous struct type was provided to which names could be attached.")); } } @Test public void testRunExtendedExpressionsProjection() throws Exception { final Schema schema = new Schema( Arrays.asList( Field.nullable("add_two_to_column_a", new ArrowType.Int(32, true)), Field.nullable("concat_column_a_and_b", new ArrowType.Utf8())), null); // Substrait Extended Expression: Project New Column: // Expression ADD: id + 2 // Expression CONCAT: name + '-' + name String binarySubstraitExpressionProject = "Ch4IARIaL2Z1bmN0aW9uc19hcml0aG1ldGljLnlhbWwSERoPCAEaC2FkZDppM" + "zJfaTMyEhQaEggCEAEaDGNvbmNhdDp2Y2hhchoxChoaGBoEKgIQASIIGgYSBAoCEgAiBhoECgIoAhoTYWRkX3R3b190b19jb2x1" + "bW5fYRpGCi0aKwgBGgRiAhABIgoaCBIGCgQSAggBIgkaBwoFYgMgLSAiChoIEgYKBBICCAEaFWNvbmNhdF9jb2x1bW5fYV9hbmR" + "fYiIaCgJJRAoETkFNRRIOCgQqAhABCgRiAhABGAI="; ByteBuffer substraitExpressionProject = getByteBuffer(binarySubstraitExpressionProject); ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile( AVRO_SCHEMA_USER, TMP, 19, "value_19", 1, "value_1", 11, "value_11", 21, "value_21", 45, "value_45"); ScanOptions options = new ScanOptions.Builder(/*batchSize*/ 32768) .columns(Optional.empty()) .substraitProjection(substraitExpressionProject) .build(); try (DatasetFactory datasetFactory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); Dataset dataset = datasetFactory.finish(); Scanner scanner = dataset.newScan(options); ArrowReader reader = scanner.scanBatches()) { assertEquals(schema.getFields(), reader.getVectorSchemaRoot().getSchema().getFields()); int rowcount = 0; while (reader.loadNextBatch()) { final ValueIterableVector sumVector = (ValueIterableVector) reader.getVectorSchemaRoot().getVector("add_two_to_column_a"); assertThat( sumVector.getValueIterable(), IsIterableContainingInOrder.contains(21, 3, 13, 23, 47)); final ValueIterableVector nameVector = (ValueIterableVector) reader.getVectorSchemaRoot().getVector("concat_column_a_and_b"); assertThat( nameVector.getValueIterable(), IsIterableContainingInOrder.contains( new Text("value_19 - value_19"), new Text("value_1 - value_1"), new Text("value_11 - value_11"), new Text("value_21 - value_21"), new Text("value_45 - value_45"))); rowcount += reader.getVectorSchemaRoot().getRowCount(); } assertEquals(5, rowcount); } } @Test public void testRunExtendedExpressionsProjectionWithFilterInsteadOfProjectionException() throws Exception { final Schema schema = new Schema( Arrays.asList(Field.nullable("filter_id_lower_than_20", new ArrowType.Bool())), null); // Substrait Extended Expression: Filter: // Expression 01: WHERE ID < 20 String binarySubstraitExpressionFilter = "Ch4IARIaL2Z1bmN0aW9uc19jb21wYXJpc29uLnlhbWwSEhoQCAIQAhoKbHQ6YW55X2F" + "ueRo3ChwaGggCGgQKAhABIggaBhIECgISACIGGgQKAigUGhdmaWx0ZXJfaWRfbG93ZXJfdGhhbl8yMCIaCgJJRAoETkFNRRIOCgQqAhA" + "BCgRiAhABGAI="; ByteBuffer substraitExpressionFilter = getByteBuffer(binarySubstraitExpressionFilter); ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile( AVRO_SCHEMA_USER, TMP, 19, "value_19", 1, "value_1", 11, "value_11", 21, "value_21", 45, "value_45"); ScanOptions options = new ScanOptions.Builder(/*batchSize*/ 32768) .columns(Optional.empty()) .substraitProjection(substraitExpressionFilter) .build(); try (DatasetFactory datasetFactory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); Dataset dataset = datasetFactory.finish(); Scanner scanner = dataset.newScan(options); ArrowReader reader = scanner.scanBatches()) { assertEquals(schema.getFields(), reader.getVectorSchemaRoot().getSchema().getFields()); int rowcount = 0; while (reader.loadNextBatch()) { final ValueIterableVector booleanVector = (ValueIterableVector) reader.getVectorSchemaRoot().getVector("filter_id_lower_than_20"); assertThat( booleanVector.getValueIterable(), IsIterableContainingInOrder.contains(true, true, true, false, false)); rowcount += reader.getVectorSchemaRoot().getRowCount(); } assertEquals(5, rowcount); } } @Test public void testRunExtendedExpressionsProjectionWithEmptyProjectionException() throws Exception { String base64EncodedSubstraitFilter = ""; ByteBuffer substraitExpressionProjection = getByteBuffer(base64EncodedSubstraitFilter); ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile( AVRO_SCHEMA_USER, TMP, 19, "value_19", 1, "value_1", 11, "value_11", 21, "value_21", 45, "value_45"); ScanOptions options = new ScanOptions.Builder(/*batchSize*/ 32768) .columns(Optional.empty()) .substraitProjection(substraitExpressionProjection) .build(); try (DatasetFactory datasetFactory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); Dataset dataset = datasetFactory.finish()) { Exception e = assertThrows(RuntimeException.class, () -> dataset.newScan(options)); assertTrue( e.getMessage() .contains("no anonymous struct type was provided to which names could be attached.")); } } @Test public void testRunExtendedExpressionsProjectAndFilter() throws Exception { final Schema schema = new Schema( Arrays.asList( Field.nullable("add_two_to_column_a", new ArrowType.Int(32, true)), Field.nullable("concat_column_a_and_b", new ArrowType.Utf8())), null); // Substrait Extended Expression: Project New Column: // Expression ADD: id + 2 // Expression CONCAT: name + '-' + name String binarySubstraitExpressionProject = "Ch4IARIaL2Z1bmN0aW9uc19hcml0aG1ldGljLnlhbWwSERoPCAEaC2FkZDppM" + "zJfaTMyEhQaEggCEAEaDGNvbmNhdDp2Y2hhchoxChoaGBoEKgIQASIIGgYSBAoCEgAiBhoECgIoAhoTYWRkX3R3b190b19jb2x1" + "bW5fYRpGCi0aKwgBGgRiAhABIgoaCBIGCgQSAggBIgkaBwoFYgMgLSAiChoIEgYKBBICCAEaFWNvbmNhdF9jb2x1bW5fYV9hbmR" + "fYiIaCgJJRAoETkFNRRIOCgQqAhABCgRiAhABGAI="; ByteBuffer substraitExpressionProject = getByteBuffer(binarySubstraitExpressionProject); // Substrait Extended Expression: Filter: // Expression 01: WHERE ID < 20 String base64EncodedSubstraitFilter = "Ch4IARIaL2Z1bmN0aW9uc19jb21wYXJpc29uLnlhbWwSEhoQCAIQAhoKbHQ6YW55X2F" + "ueRo3ChwaGggCGgQKAhABIggaBhIECgISACIGGgQKAigUGhdmaWx0ZXJfaWRfbG93ZXJfdGhhbl8yMCIaCgJJRAoETkFNRRIOCgQqAhA" + "BCgRiAhABGAI="; ByteBuffer substraitExpressionFilter = getByteBuffer(base64EncodedSubstraitFilter); ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile( AVRO_SCHEMA_USER, TMP, 19, "value_19", 1, "value_1", 11, "value_11", 21, "value_21", 45, "value_45"); ScanOptions options = new ScanOptions.Builder(/*batchSize*/ 32768) .columns(Optional.empty()) .substraitProjection(substraitExpressionProject) .substraitFilter(substraitExpressionFilter) .build(); try (DatasetFactory datasetFactory = new FileSystemDatasetFactory( rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.PARQUET, writeSupport.getOutputURI()); Dataset dataset = datasetFactory.finish(); Scanner scanner = dataset.newScan(options); ArrowReader reader = scanner.scanBatches()) { assertEquals(schema.getFields(), reader.getVectorSchemaRoot().getSchema().getFields()); int rowcount = 0; while (reader.loadNextBatch()) { final ValueIterableVector sumVector = (ValueIterableVector) reader.getVectorSchemaRoot().getVector("add_two_to_column_a"); assertThat(sumVector.getValueIterable(), IsIterableContainingInOrder.contains(21, 3, 13)); final ValueIterableVector nameVector = (ValueIterableVector) reader.getVectorSchemaRoot().getVector("concat_column_a_and_b"); assertThat( nameVector.getValueIterable(), IsIterableContainingInOrder.contains( new Text("value_19 - value_19"), new Text("value_1 - value_1"), new Text("value_11 - value_11"))); rowcount += reader.getVectorSchemaRoot().getRowCount(); } assertEquals(3, rowcount); } } private static ByteBuffer getByteBuffer(String base64EncodedSubstrait) { byte[] decodedSubstrait = Base64.getDecoder().decode(base64EncodedSubstrait); ByteBuffer substraitExpression = ByteBuffer.allocateDirect(decodedSubstrait.length); substraitExpression.put(decodedSubstrait); return substraitExpression; } } apache-arrow-java-19.0.0/dataset/src/test/resources/000077500000000000000000000000001515451615200223035ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/test/resources/avroschema/000077500000000000000000000000001515451615200244335ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/test/resources/avroschema/user.avsc000066400000000000000000000017451515451615200262760ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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.dataset", "type": "record", "name": "Users", "fields": [ {"name": "id", "type": ["int", "null"]}, {"name": "name", "type": ["string", "null"]} ] } apache-arrow-java-19.0.0/dataset/src/test/resources/data/000077500000000000000000000000001515451615200232145ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/test/resources/data/student.csv000066400000000000000000000000701515451615200254140ustar00rootroot00000000000000Id;Name;Language 1;Juno;Java 2;Peter;Python 3;Celin;C++ apache-arrow-java-19.0.0/dataset/src/test/resources/substrait/000077500000000000000000000000001515451615200243235ustar00rootroot00000000000000apache-arrow-java-19.0.0/dataset/src/test/resources/substrait/local_files_users.json000066400000000000000000000035161515451615200307200ustar00rootroot00000000000000{ "extensionUris": [], "extensions": [], "relations": [{ "root": { "input": { "project": { "common": { "emit": { "outputMapping": [2, 3] } }, "input": { "read": { "common": { "direct": { } }, "baseSchema": { "names": ["ID", "NAME"], "struct": { "types": [{ "i32": { "typeVariationReference": 0, "nullability": "NULLABILITY_REQUIRED" } }, { "varchar": { "length": 150, "typeVariationReference": 0, "nullability": "NULLABILITY_NULLABLE" } }], "typeVariationReference": 0, "nullability": "NULLABILITY_REQUIRED" } }, "local_files": { "items": [ { "uri_file": "FILENAME_PLACEHOLDER", "parquet": {} } ] } } }, "expressions": [{ "selection": { "directReference": { "structField": { "field": 0 } }, "rootReference": { } } }, { "selection": { "directReference": { "structField": { "field": 1 } }, "rootReference": { } } }] } }, "names": ["ID", "NAME"] } }], "expectedTypeUrls": [] } apache-arrow-java-19.0.0/dataset/src/test/resources/substrait/named_table_users.json000066400000000000000000000033011515451615200306670ustar00rootroot00000000000000{ "extensionUris": [], "extensions": [], "relations": [{ "root": { "input": { "project": { "common": { "emit": { "outputMapping": [2, 3] } }, "input": { "read": { "common": { "direct": { } }, "baseSchema": { "names": ["ID", "NAME"], "struct": { "types": [{ "i32": { "typeVariationReference": 0, "nullability": "NULLABILITY_REQUIRED" } }, { "varchar": { "length": 150, "typeVariationReference": 0, "nullability": "NULLABILITY_NULLABLE" } }], "typeVariationReference": 0, "nullability": "NULLABILITY_REQUIRED" } }, "namedTable": { "names": ["USERS"] } } }, "expressions": [{ "selection": { "directReference": { "structField": { "field": 0 } }, "rootReference": { } } }, { "selection": { "directReference": { "structField": { "field": 1 } }, "rootReference": { } } }] } }, "names": ["ID", "NAME"] } }], "expectedTypeUrls": [] } apache-arrow-java-19.0.0/dev/000077500000000000000000000000001515451615200156545ustar00rootroot00000000000000apache-arrow-java-19.0.0/dev/checkstyle/000077500000000000000000000000001515451615200200125ustar00rootroot00000000000000apache-arrow-java-19.0.0/dev/checkstyle/checkstyle.xml000066400000000000000000000251431515451615200226770ustar00rootroot00000000000000 apache-arrow-java-19.0.0/dev/checkstyle/suppressions.xml000066400000000000000000000045241515451615200233160ustar00rootroot00000000000000 apache-arrow-java-19.0.0/dev/license/000077500000000000000000000000001515451615200172765ustar00rootroot00000000000000apache-arrow-java-19.0.0/dev/license/asf-java.license000066400000000000000000000014411515451615200223320ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ apache-arrow-java-19.0.0/dev/license/asf-xml.license000066400000000000000000000014451515451615200222150ustar00rootroot00000000000000 apache-arrow-java-19.0.0/dev/release/000077500000000000000000000000001515451615200172745ustar00rootroot00000000000000apache-arrow-java-19.0.0/dev/release/.env.example000066400000000000000000000023341515451615200215210ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # The GitHub token to upload artifacts to GitHub Release. # # You must set this. #GH_TOKEN=secret # 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 apache-arrow-java-19.0.0/dev/release/README.md000066400000000000000000000216241515451615200205600ustar00rootroot00000000000000 # 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. apache-arrow-java-19.0.0/dev/release/bump_version.sh000077500000000000000000000037721515451615200223540ustar00rootroot00000000000000#!/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" 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 apache-arrow-java-19.0.0/dev/release/check_rat_report.py000077500000000000000000000034471515451615200231770ustar00rootroot00000000000000#!/usr/bin/env python3 # # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import fnmatch import re import sys import xml.etree.ElementTree as ET if len(sys.argv) != 3: sys.stderr.write("Usage: %s exclude_globs.lst rat_report.xml\n" % sys.argv[0]) sys.exit(1) exclude_globs_filename = sys.argv[1] xml_filename = sys.argv[2] globs = [line.strip() for line in open(exclude_globs_filename, "r")] tree = ET.parse(xml_filename) root = tree.getroot() resources = root.findall('resource') all_ok = True for r in resources: approvals = r.findall('license-approval') if not approvals or approvals[0].attrib['name'] == 'true': continue clean_name = re.sub('^[^/]+/', '', r.attrib['name']) excluded = False for g in globs: if fnmatch.fnmatch(clean_name, g): excluded = True break if not excluded: sys.stdout.write("NOT APPROVED: %s (%s): %s\n" % ( clean_name, r.attrib['name'], approvals[0].attrib['name'])) all_ok = False if not all_ok: sys.exit(1) print('OK') sys.exit(0) apache-arrow-java-19.0.0/dev/release/rat_exclude_files.txt000066400000000000000000000015711515451615200235220ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. .gitmodules .github/pull_request_template.md dataset/src/test/resources/data/student.csv docs/Makefile apache-arrow-java-19.0.0/dev/release/release.sh000077500000000000000000000063201515451615200212540ustar00rootroot00000000000000#!/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" 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" apache-arrow-java-19.0.0/dev/release/release_rc.sh000077500000000000000000000163711515451615200217470ustar00rootroot00000000000000#!/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" 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 < \ "${FILTERED_RAT_TXT}"; then echo "No unapproved licenses" else cat "${FILTERED_RAT_TXT}" N_UNAPPROVED=$(grep -c "NOT APPROVED" "${FILTERED_RAT_TXT}") echo "${N_UNAPPROVED} unapproved licenses. Check Rat report: ${RAT_XML}" exit 1 fi apache-arrow-java-19.0.0/dev/release/verify_rc.sh000077500000000000000000000076141515451615200216330ustar00rootroot00000000000000#!/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)" TOP_SOURCE_DIR="$(dirname "$(dirname "${SOURCE_DIR}")")" if [ "$#" -ne 2 ]; then echo "Usage: $0 " 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!" apache-arrow-java-19.0.0/docs/000077500000000000000000000000001515451615200160265ustar00rootroot00000000000000apache-arrow-java-19.0.0/docs/Makefile000066400000000000000000000012011515451615200174600ustar00rootroot00000000000000# 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) apache-arrow-java-19.0.0/docs/README.md000066400000000000000000000015751515451615200173150ustar00rootroot00000000000000 # Documentation Build with Sphinx. ```bash cd docs pip install -r requirements.txt make html ``` apache-arrow-java-19.0.0/docs/requirements.txt000066400000000000000000000020561515451615200213150ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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 apache-arrow-java-19.0.0/docs/source/000077500000000000000000000000001515451615200173265ustar00rootroot00000000000000apache-arrow-java-19.0.0/docs/source/_static/000077500000000000000000000000001515451615200207545ustar00rootroot00000000000000apache-arrow-java-19.0.0/docs/source/algorithm.rst000066400000000000000000000121611515451615200220470ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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. apache-arrow-java-19.0.0/docs/source/cdata.rst000066400000000000000000000414661515451615200211470ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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 8 8 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 8 8 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 apache-arrow-java-19.0.0/docs/source/conf.py000066400000000000000000000043741515451615200206350ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT 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'] apache-arrow-java-19.0.0/docs/source/dataset.rst000066400000000000000000000256101515451615200215110ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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. apache-arrow-java-19.0.0/docs/source/developers/000077500000000000000000000000001515451615200214765ustar00rootroot00000000000000apache-arrow-java-19.0.0/docs/source/developers/building.rst000066400000000000000000000561661515451615200240430ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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 11+ * 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. * For JDK 11, due to an `IntelliJ bug `__, you must go into Settings > Build, Execution, Deployment > Compiler > Java Compiler and disable "Use '--release' option for cross-compilation (Java 9 and later)". Otherwise you will get an error like "package sun.misc does not exist". * You may 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 ``jdk11+`` 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 8 8 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} ... apache-arrow-java-19.0.0/docs/source/developers/development.rst000066400000000000000000000143271515451615200245610ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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. apache-arrow-java-19.0.0/docs/source/developers/img/000077500000000000000000000000001515451615200222525ustar00rootroot00000000000000apache-arrow-java-19.0.0/docs/source/developers/img/conbench_benchmark.png000066400000000000000000004234231515451615200265610ustar00rootroot00000000000000‰PNG  IHDR‘Ãᄼ€IDATx^ì½÷»U×}¯«?áþpϽ79ÉIìØNìãzÜç8n)NìãÄNìÛr‘­^è*H¨ ‚êê ¢WÑ;D½÷Þ6mÜõŽ­ïÒd²ÖÞ  ÷ó<ï³YsŽ>Æœ‹õ™ß5ÖyI)¥”RJ)¥”RJ)¥”ª£óÊ”RJ)¥”RJ)¥”RJ©P“&òÑ£GӡÇӆƒiÿiï¾ý""""""""""ò¢¦‰|äèÑlïÞ»/íܽGDDDDDDDDDD>¢g"7<”vïÙ{\BùèQ5‘ÙºbßþÇ%‘.ÙDÆ@fo‹òIùh“Md~4¯|BDDDDDDDDDDä¼ Çl »öìM{÷HûÈy]N#""""""""""g7ç•4Çî½û²q¼aÓæ4oÁÂ4}Öìôö‚Eù5Ç9_Îs2 Ü“Y6¦÷ž}ûO¨Ì÷“çdÓÜ8ÐFÎ7eê7—&D¡œætQlKùÜé䃎C¹ñúƒ”Y‹zóJ]M­›r'’–²ËÇ¡V;Êy›«§%i‚òŸl(·¹²[’¦ÞÏ·´Ïõhi;š«§¹¶r®©2â|Seij=Õ£%ýhŠmã©¢¹±äÜɾ_ˆˆˆˆˆˆˆ´”2‘ù»vý†ôlŸS«vÓù\˜~öó_¥_]pQjݾSzþ…þ•óß×]òìo8˜öîß̇ù8¾nÃÆ\7ÿŽòI‡áP‹¦>Œ‡Ñ´|媴~ã¦\fs¾ iûÎ]iÙŠ•iÓ–­9O9M™Æ¶7ä¿M•&H™bž8Æ@-3›c;víN«Ö¬M7o©i¨‘†~fóÖmÇ6¤ßºmGzkî¼ü``Ñ;KÒâ%K+çt@[¶n¯¶åÃhPï’eËÓ‚E‹ßWr?¶½×怿ÓgÎJsßž\ú÷C˜Ü¬M[·7÷óæ/H+V­n²ýaRq­­Y·¾YsŽ²XGóÞ^×TºYƒä/¯±òZŽzO4M¹-Ô7û­¹y]óŠ×SPo âš*÷™ôÑê ]­>qغ}GõºÚ½ïØrb\[Ò•ç(úãϹ¦ú\‹b;Xsµî…ű¥®X;Å4ñ:ßK6/‰1¡ŽX+µÊhî~T„s¬Ñ·ç/ÌùŠåÄÜÓÇ\RÿºýhªüZDhc­99b Ší¬·£í‘¦¹±$Íò«Òü…‹Ž‘ÓE‹MdÌÝEK–¦öº¤_üúw©ç}¤£Æ¦é³Þªü“W>ÞáškÓâ¥ËŽ3YšâÀÁCÙÐmݦ}úÅ/ »0j00/¾äòô¿þfú«o|3]~E«lä’gF¥îÿó£§úÁÒ~øÏ~”~ø£©û—ôô³Ïg3¥\ÇfÍž“þý?ž¾ü•¯WÊý›Ôóî{ó¹zæyƼ9>ý蟒¾ü寥oçoÓ“O?sÜýrŸhçU­Ú¦_þê7iéò•5ÓslðÐaé¦î·¤›n¾%u¿ùÖ*K–6Ž¦fÏÃ<–Úw¸:óÈcg< Òa0ÜÖãŽÔ¶]‡Ôéš.é•W_Ëç¢_Œ'Æ%u‘¦s×kÓ°#1=H3{ÎÜtÑÅ—fÓôö;îJ=ï¹·æXžjèÓ¢ÅKÒe—_™¦NŸqBkêdB;zøÑtË­=ꮑ¦ Ý3fÍÎcºpñ;Ùd¼µÇíé‹.Éå’æý”[,3š¹êбSº¦s×ôÜó}³)ÅÜbð^Ý©sêÓ÷źóHý}½Ÿ|:µïxujסcêõàCUS«œ>êðúÀ\gq-6"ÝxÓ͹¬ãá#G¥{ó0–˜—<øpêСS®ëž{ïÏ×yq-¯Y¿>ן×û»iÞY²¬É¶ÄSý\º|Eºó®»¹¶nèvS¾.ÊãÐxËçŸyîù´cçî<&°mÇÎôò+ò9útëm·§IS¦VÛÂ_ÆéÑÇžH]¯»>uéz]º÷¾^ùL\óü0irn}ºöºòØEý”Áéǟx2'igž×Z÷ZPÆêµëÒ#>žº^ÛØŽûî -~gé1íÀŒ¼¯×¹ŽWwjÛ¥ïm¾î*m¿ýλª÷’þ/¿’×QŒ †ëó}^ÈuÐÞÛn¿#Í™÷vµ惇&7ßr[jÛ¾± Æ•5¶koíµNž~/½œÇ˜± =Ì )hË“O=“ç‚ús?*ïŒs»ös?ž©ÜóYó-½–(ƒû_—k¯ËýdNX»åt-²XoÌYÛïÌëS¦MÏë&îé¬CÆä–Ûzä5Dÿhó³ÏõI»\›ç„똇>Åy§O¬e檥ëADDDDDDädÒ"™ó?í;uM—^Ñ&M™>3àçCòÁÃG#ˆ+l'Mž.¼ôÊÔ±óµií†u ŸrÙ˜LßüÖwÒÇ?þ©ôÙÏ}1€nŒ>Þ~øþ%ýÍ·¾›ú¾Ø/=Þû©ô¹Ï1ÿ«ßäÒ˜ÍwÜÕ3ÝVùÐÞãŽ;ó‡ìN»¦ÿïþ(›$¹ÅºÈƒÉòÝïþ}úÎ÷þ.½2àµüþüÉÇÓƒ?r\z Ÿo/X˜¾øů¤ŸüëÏÒƒ†dÓ„<䯑LŸ† ‘þ÷7¿>öñO¦Ï}áKÙà ´hrp¬×e“å…ûg^ì÷RZ³v}Õ0 ?˜©Ã†¬”=<]zÙÙˆ£¬˜ÌŒÌ* Ãß_xqn/FE¯¥ËÓU­ÚdlòÔiéÑǟȦÛø‰“ªmãïëeÓ#‡q½ûÞûòØpjÍ+å×:Qv‹óQW½üEà ùÊ«Z§é3f¥†C‡kæòsŽÍ_«Mås{ökÖ”ÛKZ |Œ²˜»0¸¢ÜÈ”çã’ùá51~÷û‹ÒÄISªéšÊßÔyÚ†iÚ¦mû<¯'OIýú¿œç#9"è1¨ú¾Ðï˜1,ŽÉž½û³yÉ¥—§7I#FN­Z·MwÝ}OÕ<Œ~G~ÖÄkßÈk™ºX›¬©'ž|*}s|·‹/¹,½9~B>‰H;¯¸²UÎ;rÔ˜<.ö¤Ád‹4´¥˜¦Ö˜Ó/Ld® Æ#?dš9;÷ó‘ë‰ëŠëbÒä©ÇqüC#‘yÁà>sŽq¤\Æoô˜7³QȵÈ}Šº16»ßrk6Kœ¯QÆ⛬_Úr¯óµG›˜#LË} s„á{é¥W¤—^~5;.u»±{jÕ¦m~uÌ\Õ¸Vh+Æ.&%¦*×=í`Ì0z#ž±½þ†s_ב£Ç¤«¯éœï?´•4üÅ„¾¦K×|ox¡_ÿ<.˜Þ±æïýdºüŠ«òüŒ›01ßw(ó”4ÜïZ·i— TÆìåW¤K/¿"=ôÈ£Çô¥¼ž0«i/Zˆ¸¥W^Ù:?páa}Œ¾rïfÝcþÓ6Ò`Tç¶Þ}‡4ÃGŒÊsÀœÐÖÌI1m½2Ú³k÷Þ<ŒóÌœ1<4êó‹y þñpä·ü¾zM`ú³Ö¾18£ÿ<8áÁ!uS>÷bæuÀkk¾çˆˆˆˆˆˆˆœjZd"³×1[Xi& 0 æÍ'>ù9²õH%=æÓU­ÛfÃw]Á *æ¹±û-éó_ø_Ù´>\i¦ãOöŸéŸü¯Ùè(~äç«Ñ˜âíÚ_|øÑô¹J^¢z3ÒóÐFL5Ìú‡)O›è[˜d˜U$DÍ5>œÏcÖq ³†×ƒÍæÅ‚…‹«Fáƒ=’Í!Œ#ƇH=LŒô0̉Ž£þhí¹«ç=9ú˜4D6cŽ`Þaøa:-_µê8£†z1­F›V®Z“ÏSæÕ´Êš!?QÐC‡HoÍ™—óÇ€HHL-ò3w”ÉqLdŒï)S§çÈFÆ€¿Ñ^Ì7LYL¢69Ï_úÏ|ñzÂÄɹ1·abr²ÆXoÙª€y…¹rõšÜ^úŒIˆÑ&2e³n1ØHÇz`Œ‰,Ä€š9û­ãŒWÌxÌ*ÆžNæŠñbsžôÔGT:ærDbÆ8Õ;O[0ý1ôØ ®Š`†‡¡‰‘J¿“±ãƳ•Æ+Æ(fëŸ /¼$÷‡>R߆›Ó¸ñó<2g˜¦a"sþî{îKOô~*·! GÖà=œ¯æÔA{™'L4¢:z4+ièg9Mÿ—^©^Ì ‘žŒ Q°Œ} /ä£DT3Ô¼Škó6òs•:\Œ7ÇÉ‹AÌC€‡*c{°qkîWŒ'er b¸2N´•c¬ùß]xQ^ëÔùؽóuÛçÜpãMÙ¤ç5y¹ž¹®c¹f¸^ϸwÅËý€ñÇpd>cYÇ´ƒ-=¢ŒÆ蜹oç1᯹f¸ÿóxÏ}÷Wï%ÃFŽÌm¡ž|/jhH=Þ;µ¯ÜÏ6mÞšûφû7öŸ4\“¬iÚDÝD´3FÜ xË­¬;(0&´—ñ¥¬7ÇMÈçˆäÅlÆDeüésΚäú ™±áZ£ò±^iÏeW\™¯É¸Dù¬êëœòY—Dýò:?­Œ1ÑäDÇö9Å2˜›b±†XÜyØÅ}œ¿Ì mÈ×è»?> ´“ù¤±}é‰t±d-0–\#‘ñe,ÊÊ"""""""§‹fMd>,óºUÛŽy ‹ˆÄâïÛ ¥k»uOó.Î&r>^ù{ÇÝ÷¥6:¥ |u¹ða» çø€>lĨl4¼üêké3Ÿý|6Cø`N$.&H|E™ÙÍÿó³_È^Ũ°Fcpuú__þZ6ø0^®£à÷^’¾ÿ?¨ÖO½¯¿18ý÷?ú“¡XŽòÂ`»ŒŸÿòWÕ2ù‹™M;Â.ö ó‡(·FcãålŒ“Ž>a0e‰áƇ˜Ta˜¨l'&)i–#!—½»%„©6hÈÐ\Æ/QˆÑꄸàwV£ 1Éغƒs¤ã/‘„az3”ÈÄ×^# Q—˜„Dò³²ãÕ×TMzò` Ð>¾RŽ¹Äù…‹ÞÉù1?0ŽØ ‚ü€AEù‘Ÿ9¦ mÚ´ÏeaŠÓ>Ld̨›o½-·‹ÈJò¿úÚë9/f*f†i|œ6>û|Ÿœ–cDÆÒØ«–h]ê#²3‡ú0b©#«uëvéºë»eóšhG RÌ%LdÒ`3lÙ÷c]¿±É\R>&'ë‡üD¯c0cþ±‰IÄ8=õô³9 Û£ÐNÚLþ;{Þýž‘¼kO~øQ>O8iÆõ{¥ÒFÖçµ×߯Žóïënè–ë¤ æ„¬Uæ‰ÈW¢(1`Ãࣟô‰^Òp ŸÑúÄkÆ0LdÊ‹½x¹–0¶G W 1äç:3Œ|Ìæ-mo2MÛ«×ÚÓÏ<—ÛÌšb.¹î/¿òªª‰Œ1ÉŒOæcÖ[sr™¬ÕxÀ_æxÛιOŒ†,fõ¨Ñcs~ŒÍØØ×÷”A¹Ì¥ËVT¯WÈ÷Ãw·ö`ý`ºó`Þ¾`èFûø˺f­ÅÚ%Jœ¹cì¹.#•sÅ=u£\gùÞ±beNƒYÎ:fž¸?ðà†µCZÚF^ ÿë»ÝXú{ ¦-&&ë‰}‹™ó¸GÍ7?GQ‡ñ‰9]¢ î0¯Dä2?´#þ°ž˜Gîaܧè mæþ¦4ãµmûÎ\ßLamòP#µûVú:~¤¼öX?¬yî“Ü#¨›o0®ù=­r¯"J˜÷Æ›z`„ùϵÅŠ2.¾ø²ôê€×«û]ÇuÂ5M¤ú´3óÚ9ÖÐFÖb~€TY›”O?Ó0­é÷ÒLžÒ8¯˜øDdSGña˜ˆˆˆˆˆˆÈé¤Y™Üó,Lçÿö¼÷qDó“xæì9û\bŽìn4e†Ž•Î¿à¢l.75œ1ú½ôJÕDæÃs˜!‘æëùוô­+åûAšülkñù/|©º§r¹.ÊCà Ž=ý'úgi`åƒzÑD¦Þ›·æ¨b ¬È“M¢ƒÒýñŸæYʨק>/ô«šÈaªaèLŒBŒ¢Øë˜ÈIÌÌ,ÌÒ`ˆÆ×ÍÀLÀt!òöðµq ©h ã‡AY|•C“20-# ‰öļÂhÎÙ¼·Ó%—]žÿò³2o‹1xH6<0œ(Óš1ÆLÆ´%º6Ì4ÌIö¥G˜D¤Ç4¢¿ìcJü£ÈgòÒGÚˆÁƒ9ƒ©†ƒÁ¸jíÚÜ Ý؇—×D_véz]68Fô4æùè±oæ~AHi sÀ–DV†!ó\Ÿ¾¹<ÚÓ¸7w»Ôí¦îÙŒ #S–­Sh#Æí¡.æ˜^ÜN´*‘´˜hÌuD™bnqžHV C^‡Ot(¯ùJ=‘±Œ=&7fWœÇ+ŸÇ`¢}±Þ³ õÖÜtÉ%—g³‹×¡‘˜[´y"j—yü¦¢=1{™0©È‡¹Ç¾Í¤y¾ï ¹˜–”Áß+®h•÷”-®MÖíK¯¼šó1®Ì7&hŽ <5›tÌ ë æã’uJ”)i0¶ËiX礡îØb²™+úxÑE—æ5C Cúܱã5ùZÁ˜ã:Ž=‚i+?Ȇ±Îöäa›†¢‰Ì8–þa\ÞÿÀƒù:íýÔÓUÓ>®w® Ž³Ö[·m—ç­hÇ¿ hLMæœ5ûÓÎC¬ÅˆŒÅ¼& ûçrÏ]µzm¾Î0bcÿkÚÁ½†v°½F+ÆhÔMÚLeÛ¶íÏ3.\#‘æÆî7ç½wcp/aÍ1–´—:81cÅÃ)Ö íäüSÏ<›ç+æx`õ«_ÿ6GⲞX{¤aÛÖk“õEÛ⺋q#O6‘+sÏk¶òáuݤe¼¸oPãÇýùæás…¡Û‡{FÌ5é9ÆšåÞAä1ý _qo¡ æ'ʈ¶q/ámáX=9·û¡‡«QÈ1~ű¤Ü;Hc°{Ͼ| ÑŸr™"""""""§‹fMdLãé³f§Ÿý×ùizåÃþ¾ý&2°§Î˜•~üÓÿ<îø´™³rú³çäÈå2kÉP6‘ãÿÆœøñO~š·À$.žçÃ9¤_øâ—S÷[n«…åüÃ÷£‘‹&2Û|òÏ?}ÜÇÙdÚ°±Rç·r”cÑD&¢ôþðsä]1OÒMä0 0€"JSSSÓ³Š([ŒÛ¼}E¥/˜-͘ a6`œa`Q?ÛNMd~ ƒ£›ô;]“#ïŠ&òÔi3²aGd!uÙûqr“ ³†ôŒ3}"*s‹¶:DüÀcñGa”M$*çˆVŽkóŒ±W1ë‘u‰M[WÚyWMÓ÷Ø4”ÏC æxë¶yŒhÆlDóÆ:À fýsÍ1ߘ}·ÜÒ#×ÁúÆ¥ì˜ÌdÆ=ŒÉ¸0.y«‹Ê˜S&†k¥h´²–ض„þ2<*ÿÐuÆC¶¸ˆëƒ’=ʹ–Xk±†1y1âi{%3\›ÔÇ#Dâb°ÇÎqœ5˜Ûqsc;x`E©Ÿ¨Z®7ŽA?ÄÇX‡1ÛzÄ}‚cÜ¿xpù÷Þ#Ê"?Q¸±'9s@»¾P}‰ý²)›þ1±ž8ÆåÇ ã^õMdÒÞ{¯c¶áæˆ:xx¸0VÔK”4cÎ7Øî‚û+ÑßD Ó拇˜î˜ÐÜ x`Ã8_uU›¼Æ)ƒm_^˜p p^soa;"Ö¯yS~ãà÷–Û úÀzgmPW¼_lÙº=#m¬÷^#""""""rªiÖDæ<ÛV‰<|ä{‘Èfq£¹ïß~úé_~òoÕñäe;~Œ¿Å²ù÷ïþ]þQ¨0§ÉCTæŸýÙ§òÞ¢eC Ø®²‰å†A²cçîl8Pí##‡´Ã¨àǘ0G"Jøš5_oÆ¥]ý}Ïó3cfãW³ç6îã‰!È6 ÑGÚCŸ1RˆŠ¥½D b¦d“ªRQ½z‘‡rùš5‘·˜¸ìåùû‹.Îf!å†'fÕò«òž¿L1§ä'â0Ld¶ïÀtÉ?VøŠ8m!}ñ‡õ"?Q¨˜1UÙDîÜ¥uK"t1æc €w–,ˆF噉©‹ÁF{ù7ç‹&2Æpq bc¾aîc:ÅV±FˆLÇüÂÂðaoטSÆ“q-ö¿¸ÅQÏŒ&uŒ!õŠ Idh­ó}´ƒ6Pëƒã­ñ£\´?Ld̾âÜUÉ`ØÅ~µ´‰²(“ã¦<¨ ’ó€ Xý,î‰Ì±øK_iía_a @Ž1g¬¬qúF”;}aËÚJÞbà°½ #‘£´=|1Êã:¢\T°–™oú€ñº~æj›ªQã]’¯®Iæ­eØ»—¾3æ´—ë­(¿Ü'Xa¸s-Å5i8׆plËã…á‹QIdi˜ü±æùËÚŬÏ[Ü x-ÏY쥌ÁÈusk“5e”ÛAä7cÂÃ^³]׆kôgÅÊÕÕ}yã¡×}qKî%¬l® îM/šËoX÷DîßqGÏ|M±†x(À悹%"=ö9Žë'LdŒËb$òÃ>–óðšh|ÖæçhW˜¬˜îùÚ¹±{ÞS·XÆC&Í÷Ò<þÄ“ùšà!óΘP×?’‘ÈÃ'‚=ÚBš`"3÷ü›í0uYpÙeWæõÊšcíaÌczb:G9´Ck”±ež¹†c«…¸?ñ0uÃ8ѺЌÕ0ÔGÒ2v”Ær¾?ìÛ—Çœô ‡åHÔªüîµq1ÞÌ3´Êíà8 [悽…éo˜ëÀš/þ¸ÛvÄ7 âšÇ çš‹mh;?ªëƒ4ÑoÀÄåíd>Z½p?£Œ.]¯Ëã‘È´…5ÊÃ’º&òöy<ù¡¼xèA>êg« ¶3Á¨g­Ñ'LiÊ <Ö }ââ‘é±å m¬41?X!™1f\yeë<å2˜'úÈ}ŒûA¬!òÒGà_Ì=}ÇÌg¬»ò{õ2–y›—Cï%iÙj#;Ô󊈈ˆˆˆˆœ.š5‘èãç_èŸ~ñëߥÉÓfT£DË?¬Çñ‰S¦¦ŸWÒõí÷RÎW.«e™Ï?Ö-}òSŸÎæÔŠUü€Ø’lÔ…É@¾nìž¾ø¥¯T£(£Ìüáüµ9’“†ô§LIûø'Ó­=îÈÈç-X÷Y¾àwå¼äáëÎ=ï¹/ç¡üW+e°ÿñ㽟Êç‰^û‹O6u½ö†ªI@dy6mi4þ¢OE™>Ñî~ý_ÎFb66nʦ&_ÇøÀ$Á˜Åè6¼ñÇ—øÚ9æ­ä§]üPÇÈO¾JO ;êÆ8%ZÓ/L¾¾ýûß_œ† ™±w.{.cä`ôq™EÛöDƘcLøŠ8Æ æXDÖƞȘ{˜±w,ýÁp¡Ìz&2yˆdŬÆìf ( Ó#CðdšÈk×mhüA­¾/¦­;vdc£.ÚSÏD&J37G³Ž•Ç“ž×˜i´Ÿ± Ûu`È9È6!Å=TË&2s‡ùÉƘٛ¶nÍæã‡Æ¢¾ü5øÂy¶„à<õåõpGÏÜnö:ÆødX ? ÅØ™y`ÛÌg>ÐöÔ¦ Œ‘ÔñC‚ÌF)eb¦‘†m (³3xÜø‰¹ ¢+y¯×w·ü 2“20Y§D¶2–´ñÀŒ­+ž{¾o^»¬;êaNX‡˜ñå4¬sÆœm˜"Yi/÷ Lɉ¼´q^Ƈö2&Œë€HR"G—0«¹ø˞Ì+ ØN€zIðÞWúĬLxê™9ë­¥Ž »~Ó¦|]Í`ÛÒ°®rk—®iAå~‰)Íñ—qaݳæÙ¦ù¡½D"S.k1®OúÃõAzæ—(Úx0ăÒs-0´càƒs;â‡>c‹®yÎ3vlŸ‘¿-PY«´•4Œ#×å2ÑŒA<ä"Ú—4DJ3.\»ÌmÜâÁÄg]¬^³.·“‡¬WÖ-k/¶‚ ½£Ç¼™×yì‰\ËDfNx.×%sÉ%j=¶;áþÆ79¸W1nÌ[÷0÷ÜâsËÚ¡.æŸ29æ>×!ë—4”A^Êà&ýâ~ȺeýðoæŒk‚{sK9€9Îü—£9Ç}žµÍƒƉûBŒ%}e«#îÃaЗß;EDDDDDDN-2‘ù »vÃÆÔ¡óµéâË®J§NËnù NT`6R+¯'Lžš~É©S×ë³áx"QS˜/ô{)}âSŸÎú)Ãë¿ý?þäO?‘¾òÕ¿ÌF1|æ~>›3‡ŽÉi?õ矩|Å2i×ù¿úm.3"i÷]wß›>þgŸJßýÞߧÏ}áKù/ çÉóËó“kÌêˆëtM×l>ÿíß}?G-³?3&AŒÃ/ÎÿunyÂ|¤OÏõy!ýù_|&×Ïk .Ì ™0=‰neïML_Œ0 ¾º^ŒÖ$ †+¦ÿŽ=~cŽ0)#ÓÓ3" ÊcRŒ Ì_¢·1Êèf†vì‹KßØ–HS¢‘Ûµë˜Í&¢°1 #zy¢,L0úpe«ÖÙ¥L¢§1±0XÂÆ„¤œ0Q1ìÈ‹YLýô Ó‹sŒ'¦ Pä§ßTD c*åýB_~µ:žì³Û¹kãžÈ¼&fæ?ebŽ1Œ#ãÉ6 ÅHdÚAZ4‘ùA,ÒŸ`€Ñ/Ú‚9É1æ˜×äg[Æœ¯·wèЩº½DôŸ9fn¢Êcîbï_Ì6̼êùgkœ¯¬}¢‰Êüeeý1NŒ&90¦ÔA;˜3æ°S§.Ù$b“=‹1ïba¶’öR>cÂrŽ¶cìñð‚xʸúêÎÙŽí,HÇ<³†þ2—{ÌãDµ2ÿä£þù÷ Ò`–Óð &Ö20¬Öý¦?¤£.Æ™1gkŽqMÐæ‰5EÑï€ë“yÃ$ ót¬Ö>†#ãȸð`ª¸4mã8k‘6‡A¿}GãöüX$sDæ…qfŒ¸cÿq W®£Øf…5ʼócŠÔA9ñ#‹”Á: ïÅ-&øA¹ƘÒ1¶<üà8ëŸtü›¨mÎEŒè¸—´jÓ6ÏAq\o¾A@;úC½ìUÍ8nÙ¾=¼QéÃò7Ö &8×}fnxMßË&2‘ÔŒ×{ÌÛÐa#ª÷#Æòßš7¯º^1–ùæ ?JÙô“#”MLxʤÝñ—í>⾞ËX³.¯_ê‰2Øz…2âþÅÍ.ú{"}æá‘ïñ~÷l´nÝ.]pÁ…Ç%mHÃ<aν+òŠˆˆˆˆˆˆœnZd"ÃÞýûÓâ¥ËRÇÎ×æˆä;ï¾? 1:oi1tøÈtGÏ{s2ò;¥¾k |XÆÀ#‚‹è|ˆÇ1jLåÃ÷ØlžÔôɃR뫾˜ F'OÍj#ݘ7ÇgCá™çú¤å+·’¨—‡ôü›a"{1»×½ûuðbö<Õ>­lìSìcÊyÌAÚ]5I*c‹¡ÊþÀDüb ‘®ØúJ¤+æľ¯Å¾òC–2øÚ4Q·å4˜IDÓazaZfãÁÖ ‘'úE$"f‘˜ÙÌK±\òcŽÓO¢l—­^·qcWÆ#•:0ŒˆdÄp,®ÊŽöG›IO?ù:;åDbä©wžsVÌ;Ññã{ùu¥}ñÕ}ú™#F+ãLÙ³—1-¯1ú‰ñýúƒòƒò<3ÑO jŒTÖsѤ¤¬aÖ Ñ±¬±b9ü›±Ä¤Æ‹²c¼k¥¡Žbš¢ˆyˆAÿ¨‡Hå¸ÖHÏÓŒSî¬Ýz÷&òÐ'ú]ADÊ•ó£Új™Å4å6ËÊuå–ëŒ7"ÞØ ¢ØêÊjªÜzm‹úŠý©u¬Xþ‰æ/·‰cÌ×K_ƒbßòù=å•Ç?Æ¡\f+¶?Úƒ!E„$ÛÇ´Üžr›Êýiî|¹î"å1z¿óX«ű+¦)·§<–-©§%iÊmy?ãZ¦<ÏÁíûó©WFù\­4µúÓT;N$M±åzÊ}¨•®%c_LS\Ÿåtõ楹±hnÞš;ßÒ4A­5ùëå+–]„oÖð-"èã› å¼"""""""§‹óNôƒ)é1Ä0Èæ/ZœfΞ“ÿòšã'Zž|ø0gDZe]ŽJ”÷cJäå}÷?£ ½.DäDà>L„5Û‘Ä·*ÊiDDDDDDDNç9r$íÞ{âNsÄUDíkŒ,§‘³‡ˆ ÓìRðæÁ›H‘Óñ†Â›×‘#GÒÁƒ;w&@» ÚY>'""""""r6Ÿiù»cÇŽ4zôè4sæÌêgq?óŠ¼4‘å#o¼i,]º4 6,=÷Üs©OŸ>iĈiÉ’%ÙH>Uo&Ô»qãÆ4eÊ”´nݺ3ÎH¦=ëׯOS§NM›6m:eã """"""rªà³íž={ò¿ÑÛo¿®¼òÊtË-·¤½{÷¦£Gæ¿û÷ï?«?÷žÍm—³Md9çá»oß¾4xðàÔ±cÇtÕUW¥N:¥«¯¾:¿¡p¬oß¾i÷îݧäfL2í%—\’ÆŽ›#~Ëi>LhíºôÒKÓ´iÓθö‰ˆˆˆˆˆˆ4òæ͛ӣ>šÆæΛzè¡Ô®]»üùÿõ×_O³fÍJ÷Ýw_ö0’Ëeœ `‚oß¾={åsg±…È©ðVäÃGYÎy¸yacwïÞ=¥kÖ¬I«W¯Îæî­·ÞšÍä~ýúÕ¼ÑEsùx-j¥ÅD¦Nê7n\õ†ZNW‹HW«]ï'iÊõÒž7ß|3ÏôéÓO¨}"""""""6|Ž]°`AêСCjÛ¶m6ù[ü7[µj•îºë®SDv*áÔóæÍË‘Õ|ËútlÍy"Ä·°g̘‘vîÜyÖ¯4&²œÓpSå+,¼Yp£Å<æ+,þͱk¯½6§Y¹reõFÌ 7"Þ\x¢É±òµ0[£¼mÛ¶UÓF9Ey„ ùk5”¹eË–ºOé¢î]»vå-&xÊXŽC˜ãÜ 7lØË-—uvëÖ­¹q,ú&2{E¡Zý9Sá³ëO<‘ ãöíÛçÏù>ø`ºû¹ |îg«I>§—óŸéðÙƒ¶uëÖÙ0/{A|ÞÊ~Cñûlµä™={¶QÉ""""""rÆÂçd°ø|Þ¦M›ü-äÅ‹çϲeñ{Hle‰ÁüÚk¯åcgÛç\>»ã3`„/\¸°¦‰ý}çw²áL:‚ÍÂøå<ûF36œÇ{`ÜV¬XÑl~ü‚ð¾+ËË—/ÏÈÊáOtéÒ%Mž<9¦• l9»ÑD–snVÜì0¯»îº|S«÷&Á 1ž°ño¾‚‰Š±Š‘Š¡‹QK41û*a$c³=uçÎóSÎÞ½{gÓš'œ<}ãæŠ0‘)‹}˜áá‡N?þx΃¹<`À€ª¡Ë ˜ºyrúÈ#ä2»víš_Oš4)—·hÑ¢\{;åÑFÊ£nØ”E™ßãOe<äÇüF˜È¼™rœ±¢Nö’¢NÞxc)?mù0á3<¦çË/¿œƒ¡0Š1‘1iÂ`?äøÜK&i½Ï¹ÅHÜ"øõ|…SMs&2íÚ±cGþÌïpÓM7å¿ô›È`úÊ·“ñ|¸ýöÛsÊc¿år~ØðSÊã\¯^½ª¦2{4ô6hР4lØ°ª×ÀßÑ£G×F9»ÑD–sn <%À妆ñÛÜS0ÎÅËMsxÈ!¹n–ÜL1k1’G•ßˆðÅææJt0é€(éˆ2™t²D/óÆC:žö‘—7ºØÓó–›.éB”QÄl]±dÉ’êòÆGyS˜7Ì'Ÿ|2¿^¶lYNÃSX mêD´ûŠ+®È7nê˜á´£Øt”Iä4ã€ÉìÍ_DDDDDDÎ$ø Od-[VðœÏÒ|&fËŠâ7‘ù‚tüå©eÄ4~X¯ÑË¿y²GZöJf› ÞØ ƒ›<7߀cÔÅ›7LdnîaúFycdË nê”IÙ”[üš}áÍ”ú!~X1ˆö…‰ŒqÍ–šÈ""""""r&Âç]Ä6¾D$};}úôüyó”ÏÐÃ|›• ä€ãñ9™oöRù(3úÃ2G›2‘ñØr ²¢?Á7– ®£ßø|c:Îãyà}x‡Ÿ@~<‡r~Æl¾õÜ”‰Lûð/ðVêm"g/šÈrNÓGnhÜdùѼZ7{ÌÑW_}5]zé¥9"—hcŒWÞ|j™Èl Á›GÑD.š¯ü­g"s³.š±a"cÔöíÛ7×Í›{ÅÌWQ^cæòfÁÖ¼’'ÚXËDÆdæ ”MñË&27{¾^Âñ0‘yƒÕD‘³‰#’–Ïÿ|®Žâ€cÀçärþ2E#™ÏÞ’}˜2Md¼ÄçyàŸïéŸ>}²O@{cìòf¼ ‚Í"/iØÊ9ò‡ÏùùÖ4Û„ò éÞ½{çHîJÃ#!h/¾…&²‘Èç&šÈrNÃM‹­0pùÁ;n‚Åhä¸éñ†À˜7nvìბËM4ÒsÅÀå „¯{pÓd;‹–šÈl+Á>òÅOÊä)!íc?&¶° n¶£@œ§<Ä6ìÅÌ›¯Ï4g"So”]þQöx¢ö~ÆLfdMd9ÁÅXå¼|™í,ˆª%Šø|¿?D$-Aaõ"‘ÎS.ß<Æ€ý°MÑ¢‰%LŸ^zé¥ü›Kl+Ág}Æ„->ðCèË‹/¾˜£ŠÙÛoaøðáÙ`ÇÆ[`«‹r~Æ‘üxœÇPÆ“ >| <ƒØâ3¢žÙ⳩-EåìDYÎiÂXåɦ.Oä¶mÛ–o¾Ü@1m‰8Æhåkäá&Ç×^8ÆÓ4~é•›%_ áÍ&~h…ùÚ™+ohìC{ñFÄ7ñxSâÍŽ›8Qѱí[qôìÙ3g¯§–D"S7vÞ$‰žæGü(‹~=:÷/¢­D‘³>?³=eüè=ßæs.&*Ÿ£ù,Œ¹IàX­­.ëÁgäIª¯£3— 4Àøås>p˜å¼&XŽtìwŒáAtxg\ˆLæ÷ ˆFæ<~ ÁfÅüD)“O/S™ßrÂgàžFxæÑ.ò÷aïrrÑD–sÞ(ØÊ‚¯e`œrƒã¦ÈãqÓä-?jÇéyÓá†ÉN=rä1iyÝ»wïüD(æ+Û`”÷DÆ|å×^Ù†a‰Œ‰Œ©ËM¼W¯^ùf1ŒIÆ2f/¦7ÇÙŠƒºyMÝl»A½ìqŒ¹[ü¡¾0‘¹aSv<5åfΛ 7yúÌMž¯ð`6sƒG¼¹ÒòžÈ<¼ì²ËršÈ""""""r&Ãgàø\ñÉgj¶¯àXlQq&Âï‡Øb“èà2|kšÀ7úÍ79Æç{LõØÆ’ü˜éøxä¡<üŒç»zùiçWåð 0)?%¢ŽÙ9Úc$ò¹…&²|$Àå¦ÆW7ø *?ÇW6ø¥Ñ;vToŠÀ¿¹áõïß??‘ÃH&º— 湡†AKÄ1Æ,7ÖÈÏ_~øŽôlwAÝÜ@1ƒ1“ùšod”ÉÍš¯¡P^Ü\ãÍŽ}Š0|y2H4‘ÉDEó¦ÇMŸ¼Ã—Ÿ­„^‹0ÇËiÂàß|›™ 5Lb¶³  Œí1ð.ŠÞB­üAÀEa>Çy^Û#çšÈò‘!ntﳕi|Ó3¹[·nùÏî_,-AY>rpcäæY6Lëq*ÞtZZfKÓµ”“]žˆˆˆˆˆˆˆœàƒ¬ÆññÍcì8¦O -AYDDDDDDDDä#†qlmÙÒà:ÐD‘ºh"‹ˆˆˆˆˆˆˆˆˆH]4‘EDDDDDDDDD¤.šÈ""""""""""RMd©‹&²ˆˆˆˆˆˆˆˆˆˆÔEYDDDDDDDDDDꢉ,""""""""""uÑD‘ºh"‹ˆˆˆˆˆˆˆˆˆH]4‘EDDDDDDDDD¤.šÈ""""""""""RMd©‹&²ˆˆˆˆˆˆˆˆˆˆÔEYDDDDDDDDDDêrÞ¦M›’ˆˆˆˆˆˆˆˆˆˆH-ÎkhhH"""""""""""µ8/)¥”RJ)¥”RJ)¥”Ru¤‰¬”RJ)¥”RJ)¥”Rª®4‘•RJ)¥”RJ)¥”RJÕ•&²RJ)¥”RJ)¥”RJ©ºÒDVJ)¥”RJ)¥”RJ)UWšÈJ)¥”RJ)¥”RJ)¥êJY)¥”RJ)¥”RJ)¥T]i"+¥”RJ)¥”RJ)¥”ª+Md¥”RJ)¥”RJ)¥”Ru¥‰¬”RJ)¥”RJ)¥”Rª®4‘•RJ)¥”RJ)¥”RJÕ•&²RJ)¥”RJ)¥”RJ©ºÒDVJ)¥”RJ)¥”RJ)UWšÈJ)¥”RJ)¥”RJ)¥êJY)¥”RJ)¥”RJ)¥T]i"+¥”RJ)¥”RJ)¥”ª+Md¥”RJ)¥”RJ)¥”Ru¥‰¬”RJ)¥”RJ)¥”Rª®4‘•RJ)¥”RJ)¥”RJÕ•&²RJ)¥”RJ)¥”RJ©ºÒDVJ)¥”RJ)¥”RJ)UWšÈJ)¥”RJ)¥”RJ)¥êJY)¥”RJ)¥”RJ)¥T]i"+¥”RJ)¥”RJ)¥”ª+Md¥”RJ)¥”RJ)¥”Ru¥‰¬”RJ)¥”RJ)¥”Rª®4‘•RJ)¥”RJ)¥”RJÕ•&²RJ)¥”RJ)¥”RJ©ºÒDVJ)¥”RJ)¥”RJ)UWšÈJ)¥”RJ)¥”RJ)¥êJY)¥”RJ)¥”RJ)¥T]i"+¥”RJ)¥”RJ)¥”ª+Md¥”RJ)¥”RJ)¥”Ru¥‰¬”RJ)¥”RJ)¥”Rª®4‘•RJ)¥”RJ)¥”RJÕ•&²RJ)¥”RJ)¥”RJ©ºÒDVJ)¥”RJ)¥”RJ)UWšÈJ)¥”RJ)¥”RJ)¥êJY)¥”RJ)¥”RJ)¥T]i"+¥”RJ)¥”RJ)¥”ª+Md¥”RJ)¥”RJ)¥”Ru¥‰¬”RJ)¥”RJ)¥”Rª®4‘•RJ)¥”RJ)¥”RJÕ•&²RJ)¥”RJ)¥”RJ©ºÒDVJ)¥”RJ)¥”RJ)UWšÈJ)¥”RJ)¥”RJ)¥êJY)¥”RJ)¥”RJ)¥T]i"+¥”RJ)¥”RJ)¥”ª+Md¥”RJ)¥”RJ)¥”Ru¥‰¬”RJ)¥”RJ)¥”Rª®4‘•RJ)¥”RJ)¥”RJÕ•&²RJ)¥”RJ)¥”RJ©ºÒDVJ)¥”RJ)¥”RJ)UWšÈJ)¥”RJ)¥”RJ)¥êJY)¥”RJ)¥”RJ)¥T]i"+¥”RJ)¥”RJ)¥”ª+Md¥”RJ)¥”RJ)¥”Ru¥‰¬”RJ)¥”RJ)¥”Rª®4‘•RJ)¥”RJ)¥”RJÕ•&²úHèСCéÀiÿþý""""""""rgƒw£Î\i"«sZGÍ7¢½{÷ŠˆˆˆˆˆˆˆÈLCCCörÔ™'MduN‹§Xå’ˆˆˆˆˆˆˆˆœ™‘|fJYÓ2 YDDDDDDDäì/GyÒDVç´ØW§|3‘3¼uæIYÓÒD9{ÐD>3¥‰¬Îii"‹ˆˆˆˆˆˆˆœ=h"Ÿ™ÒDVç´4‘EDDDDDDDÎ4‘ÏLi"«sZšÈ""""""""gšÈg¦4‘Õ9-Md‘³Mä3SšÈ꜖&²ˆˆˆˆˆˆˆÈÙƒ&ò™)MduN냚Èû÷í«”±¯ñï»ÿÞ·ïøt§“}•v”µ„}{¹úò./§=Y0VÔQ>þ~¡ï-í?s_¦¥yO'ÒÖSÁɬÿDË"íN(Ïûáƒ^ßµx?뤹ôÍ/¦kiÚ2Må­wüdðAÊnIÞ¦Ò4u®%ç[J½±­uìdñAË®×f‘³Mä3SšÈêœÖ5™6lÙ™ÖlÚq [¶ïJ û÷¥ûOÿô={ö¤7¦;ww®)0Š·ïÜ]iÿö´¶ÐúÇùƒZÞ—Ã û[Ôwêd¬¢Ž“Á¦M›ÒöíÛ;^f÷îÝyœ6lØP…¼»víÊæf9ý©€ú¶mÛö¡9›7oÎcu2ê§/;vì8îx-¸æ¨{ÅŠ'¼NËPóÅ|²ö‹Æ4e3Çœ+ç HKò’Ž7uO }v믕¿©öA¤gÝE[ÊuœgŒ©Ÿ²âXK„pŒ6çˆc ¹,ê¶ÆùZm?U}+e“·VÙQuÇ5[l[Su_7Õï–® ˆûN¹¬“òúknÜš¢8g-m[-b=¯¿Zk ¢m1n´n‘“ÿUgž4‘Õ9­ò!£ôüG&¦¯t–¾vãð*_¿ixº°÷”´põ–™©' >äcšMŸ>ý„ÍÉà ÒàÙ«Ò—®š¾^è üã]cÒÈ9«³1^ÎWSxç®Ýéå©ËÓÜ››íûÁûÓ}C¤ß<:)íÞ½'}núºuëÖÜw ¦úŽù±råÊtûí·§o¼1uëÖ-ÿo¹å–4bĈãòœ zõê• Ymß믿ž»0ìÞzë­tÿý÷ç¼wÝuWš4iRÝ1džgÒ²Æy=a„œïî»ïÎÜqÇyN1‹íÆ,}üñÇóø„iI; ”zôè‘Ë|ôÑGÓ’%Kªæ ãöÜsÏå¶Ã /¼ë s‘qzú駫}{å•Wò5çgÏžî»ï¾\6c6yòäšóP&Æö’—²_{íµ<ö´;ÚN}1î¬÷×bÝ'N¬¢ŒÅ°aÃòXqþá‡N‹/®¶›õÓ¿ÿêš`~›ZC”wÏ=÷T¯%Žñ—ör3§1”±~ýúê¸Ru1nõæ½Xe8°:g½{÷N«W¯®Ù¶zD9ô›ûãM˜ËâZ gÊç=…ÓfêæZ9ѺEDDDDNü?TyÒDVç´>È`LÒoÝ62ýý£Ó°9«ÓÀ™+Ó¾0}ºó ôÓ^ãÓÎ]{RÃ}éÐýÇl A^LTþ9‘ËDñ9زHÞ2˜‹-JóçÏ?î\s9x õŸº<ý_W¼šµ( ™½*÷çÙñKÒwzŒJ_¾aXZ¹q{µ]ô‰v¡†Æ~ØßÍŒýĘũҙjùžET3ùn|uNúþcò룇Ï7gVׂ¾/]º4Í™3§¥YL“eË–¥k®¹&›93gÎÌæFÓƒ>˜Úµk—æÍ›W5)c¨h.Æñâù¢yÄ9(®7Œ˜>}ú¤£GVχ‘T,›:‹uSF­òÊuÛÊ¿Ã<Œr1Í1…ÂD¦¬8WîC=H‹ñ8wîÜfǨë±ÇË>sQãÑ'ê/÷©<¾¼# “ëºë®ËÆ?m‡²FÚ©S§æטv¾å1¡ž~ýú¥öíÛg#yøðá¹,Ìõˆ`%]Œ?ePçõ×_Ÿ†šÛ‚ñŠN^ƲS§NÙØ£lLMæÆŒ“ \Öu>|8_§W_}u6óÆ—žxâ‰Ü–3fT×Q±ï¬ ÉŽ;fS“va´bêaÓLAÖr±ŸäÇà¼üòËóØÄZ õÑnúÓ½{÷lºc’ÊãAß^~ùåÜVŒCÚļÝyçé†nH#GŽÌõRf.õqýÐÎGy$÷v2ûµúVœWêg}ЖQ£Fec¾mÛ¶yÜh7cËuÃõK½ÂÔE”Ã}€sÌã›o¾™3î\çŒ;ý¦<枇ÌsÊ:¡]¬Ï.]ºä1aܘ_樸¾‹k‚¨\B1F”ϼ3·ŒEëÖ­óÜÇ5FZî1Q~¬YîCÅë.Ê/^üe,è+ùwêeÂ`u]ïþëü—]vY6ÈyM»œËçaÁ‹/¾˜ÛDûIC9¼f `bs½1fåºEDDDDNšÈg¦4‘Õ9­Z²[ Æè·o•~õØäJI³) éHCºkðÛé/:JËÖoMk7ïH“®Ë†2F2·‹VoI3—nȾç­Øœæ¯ÚœVmÜžúOYžúM^–VnØvBfj˜j%Dgžè‡zLä—¦-Oÿ÷UÒ¼U[R:| ±/GÒàY«Òk5 [°6›á´ëí•›SßIËÒÓã–¤qóצ½{ö¦­;v¥¡ïF3wî7+Í^¶1íÙÓq¼lÝÖôê´éù‰KÓ¬¥«¦r÷sÓîy3-]»5½P)¯ï¤¥éµ[N¨ï€éGäãªU«šS  ÌÎ;gƒçÈ‘#ÙHÁ"ºê7Þ¨š0D=ÙHÔ'&=epœèA̲ø;~üøôöÛoWëa0ž,X êiÓ¦UÍ?ÎcþaØu‰ÑÆy¢ )›¹¤,¢{ÉusŽô”‡I‘¡QyH ˜iaÌ-_¾<÷ƒŽqb+ L:Œ¹0W; ̈è,[-wŒÁ–D$Ò>Ì5 Âx ÷‡üŒ=ãN¤*}Z»ví1†9&5}"?yh3ýÀä¢~æ0Â0º":S—~“–1£Nú&/!f"y1ii†#ãH»H‡IãMÛ0ß1Ñh3†ç"?%!&+í¥¬wÞy'ŸÇpÄhä<óŠ)‡qËœ°ö¢f‰^m(‡µÇú íÔEŸXG¤ÁÅLDÔO=E™³~h3‘¯˜Èô‹1ÄÆ Œ¶3ŽÁ´—õ¹ÉxÄyÌCÆ+î3Œ=mh;i0&¹vûgžy&¯/ÒFß01P£]œ#?× uÇ8P6ãFc܈æ&?å°–™c õPÕ¬¢–¹¶#Êãú¥ïœ§ß×^{m6x)—òéoÌ#eóoÌåXSÔÃXp3vqÿà8ãÃ:`1ã'Ìaæˆ62ž”ב÷”U,Ÿq%-÷mÖ[”O{ Ö õržré_äeRi£m¤›2eJ¦|]’†k)¾u&rœ ó™qà'â™ù¤LúÂúŽuiX_Å5'""""r:ÐD>3¥‰¬Îi5g|5EÙDΑ¹¯‡¤k_šÍÔÍÛveãô/o–Ídò`Ävé7;ýÛýãÒáÊ¿Ûö™‘¾[)çGw¿™¾}ëÈôñŽo¤ïT^¿³vk‹#’1Æ01ËŠQ”-¥h"ÏY¹9%"øÝê{†ÌOŸë:$-]·µrü@z~ÂÒÊëÁy«‹oUÚû±SÏÁo§ù•|wû¨ô±ŽÓg:J—ôž’Ë2kUŽdnÜöcXúd§Aéž¡óÓáJY·½>/}îÚ!¹¿ß¼edúók¥¯VÒÍZ¶±ÅF2sH!hK¢âŠ&òèÑ£³™D˜(”¡ƒ©…™‚шáIZLÎaü’ó„ãclÏ@Ô&Ž±ó@Ä$y8Ï_Ì&Ì1ÚÀ×é1òˆ&ä8yÙÚ€>`LéÇWòÉKêÂðŒèV ¾Fù†¹QŠ˜wäÅÐÄ¢–|UòÚ´i“#41±1‘fæ‘£˜FÍaŒ;¦*uÄ å4EHCä0í¦m{d䧽[Œÿf3†ôÝò Ky˜aDˆF{1¾ˆ\¥|ÊÅDæßÔå3>±¥†4¦\Ô룔èVÊãayÊ ,Ú€qKDé†É‡I[1;i_D±ÆVÌÆãOݬ+"™i ç)ƒ¶‘‡>ñ€€õsÊúâ/&4ë/ÚˆyN›1¼ãžF™ÕËy C¢¤cmÒÒóh/}cN1`é툾aÂbrS.Æ6mms+Ò³ÆÉK0cQìÑ­¬1ƃ¾qÝĵÁ<‘‡>a`cÒÇÒ3¦ÔsL¿¹F¸N1ry°ýf\Y±c\Y?Ì/í¢ý§}äe`úsœ6Ŷ@{8Æbž2nŒm€¸¶#ÂœõÀ½˜-2˜ËˆDÌnæ3 cÚÀµGýê¤Á„f½sO ?×)åÑ&êe]“.æ—´l‰¹K±=mcÍžº€ùd 3W\G9^¼NyMúÄ=Ž>Q?ÚÜ7)'âЗâ·6DDDDDNšÈg¦4‘Õ9­j"ïöÑ™>—æ¨Ü'ß|']ÓoVúl×ÁéÎAóÒÑà é™ñKÒWo8ÖD¾úÅYéÇ÷6šÈíûÎLØöõÔgÂÒ­<æíµéWJ^œ£uËõÖSèÓ÷Ó'Lä—§-ÏÇÿþÀøôû'¦¤ ›œþéÎ1éOÚL_ r¶«øF÷©íóÓÓŽÊ¿öíKí*íÇüÞ´uWZ¼fKúêÃÒÝCæ§m•ó«7mO_í6<ï½zÓŽlªwé?;}¬ÃÙ”î9x~®óá ÓŽ]»ÓÌ¥Ó§;NÝÌÉm*·³˜r©[>W&Ld L9¾¶¡…A†q‚YE:ÌŒ7ÌLŽ-Œq‚ÉB#æ‘£˜S˜G˜°˜U`DXb uHD)ubìãŠô»ŒÌ#&íÃøÅ,£l X¢©›ùÆ¢|ÚaÅ9Œ0ΑÓ”ˆPŒ/ æV­Ze#ˆô˜Oô‹(HÆsS »%c`€Å~²åseh&!Æ7í  –Œ &cÀ8-KÛbÿWÌ[ 4ÆŒ±ÁHÃd # ø7Q—˜gÌ%sÌ\-Ëü2?Ô[`€a¸1—\3avR?í¡^ê$"–¶…ñIÛGL8Lá0àˆ0Æœc=`Šcè1ÔѺa¤2Ǥáß´7"f©›ùÆP¤.ŒAþÍú`,ˆbeNY3Ì/cAù…< ,úƒQ)<8Gi;&r~aÓNLZÖýæX¿Ìý¾•·£ˆ±Çˆež0–Y <´ =Ìwô:˜;ÚÎxûÆ_Æšh^ÒS&09é3Æ9ë6"d)›ë‡s\'Œ Ѹq]ÍQÊÂ`gÇ`0ǹFÈÏuÆ5Êbä2¿ÔÇõ{ÕUWåzYcÜ#èevâ±Öæ~È°c€‰E»PlA{) “£—cŒ+Q´˜ôÅñç+ö‹òh3ó‚¡IŸùº}l‹@ŸÃ4Äøb,1aË&2Æå°5F5cCÿÂpÅDd\1%#‚vÑŽ!G^ú‡iŠÁH[)ƒ(U \æ ¢NcýÐGÌ[úÄyæsC³’¾Pcÿ_LAòsž¿œ§¬7úÇÚe¬0[y‚bSNqM2æŒGü`ãA_9ÏZ#ò7ú†!ÑÇœ®Ö-ùÆ%ö{æz¢oŽŒéx@ÂuÅšŒ¾&yDº} ú3“qá:¥^ÌuL[ò3Ƭi®5Æõ_ÏDŽmD(›cD3WŒ'×cAcoêؘ¶Æ7 ¸¾1Ó™Ÿ0ê)“5ÁÓ÷¸>h[-™q©g"Ó_ÚÉuÁü±îé'm¤Ÿú´…¹Äˆ§m´•s ësxØž(ôŒ?–èglg×'}âay⺈sa ÌšeNŒB‘ þ/­Î0*éF2}a}R.æ*Ûp¯ˆìf0:£Ú†1JÝ\kñctŒéÉsɵ€a̓Ÿ˜sÒœ6rÝsý³Î˜ÿX÷äaŽiå1Ï:d.ȃ)ËuF:0Vô‹õÒœ‰Ìü3^ñÀ‹¶ðƒqgíÄ=¯Ö óȸÑ>Œh®! }Lmêf|™ fò•¯%ÚÀkòÑòqœr1¯c $Ö2Ç€s,úÃkÆ‘qÆ|Ç„oéýBDDDDäd£‰|fJYÓ:&òùNJépCŽ ÌW ׿½}tŽ.~nÂÒô¥ë†¦u[v¤ƒ û³1üÛÇ&§}wOäj"czuFt\|è?QŽÛùÐ{{"ÓÏ¿ê>"]ôä”´`Õæ¼g3?¸iÛÎÜïGG-Jÿûæióö÷Lä§Þ$ùPzmÆŠôÇÕHäÆHã…«·äñÁœ&ù%ù?OÀDÆHÄü,î‰ÚEs“Šca¦aêa¾`Ra¼ùˆ‚i†i‰é‚ù‡ùƒùŠ!X6‘1‰ø :Z˜E.Ô‡ÙÄœµÄDƸkÎDæáåb>žд#³‰þ‡!Dß0^10”0ž"R³ÒnÆ úEô›ëŽµÁ¿¹FÂDŽ(\æCkŒë„20=1Ñ©'Æ£˜5ýçkŒ9c|9ÎœÄuÎ߈DÆ€>“–qiÎD&-×'ù÷0È£LÖQø¬gÖk±â÷•hó@>Àl§nŽS&ídŒè;ù9Ïý„õ€é\\¬Öë!ÌuÈÅv´/®[ƇHkÊgüâ~w"×¾ˆˆˆˆÈÉ„ÿ‹ª3OšÈêœÖù&ò/”Ž Oö1f¯_ÌT¶wúÖªôÿ¶ž¿$mÛ±; ˜¾"}¼ãÀôïLøÀ&2íÇÐÀŒŠh²rš–P4‘_Ÿ±2›Ål_Á–l«ñßÛ¾žúM^–&,\—þ°Íké¥)ËÓÞ={Ó´%Ò7o™þê¦iÕÆmi{¥ìÜíÕ9iã¶]9™Êûõ#“ÒÊ Û+Çvæ-2ø½5›v| ™¾Çêa˜´´ïEŽˆ:ŒŒ¢ #²“„¨LÌ\¢©sÓã‘ôõLd $L>ö+ÅP$o˜QDo’þd™ÈlD"ÒL $Œ)Ú‰)IÚz&2ÆÂ|Ãp¥¿¤/YÚŠ9Ń‹–üa‘¢‰Œ †Y†¡È6 ïrñc„¼&"“ŒõÍÒW^chÅžÀŒ3i¹#—>ĞȘétÌ ãÄœ0®¤!2—qà<†)Æk˜¼¬ ÖF;ídNxÀÀžÈlñÀø3ÆDR6mˆ­˜ÎÓ>Öec¬Ò×è;k²˜oòcÜRëÓ•È]^cbòÑ_Öí‹õE~þ26¬GL`ú˽˜àß´=öû&?&1ë‰ñâ ¦yŒgq?èè×óƒ®9ÒÒöòس^ém#:–²ùK_è3}gQV¬!úNzÚÀ\3oDýÒoLSÈĞƼÆÈsÖ"kŠþ°ö£,òóš‡ ´Ó˜ºC"¤™SÊŠ¶ÇC êÁt¦/”ÉÊC%Æ0¶3a°è;×Ö•W^YÝ9ÖyÙDæí¥~ÊäßÔOd=pMrÀPf2nœgŽ¸WÐ6Ì_ø°^©›9å>A]¤g-GÛ8Æ°0¨yëq!÷¢¸6™7Öé#:?®wÖ3÷.Ö2õDzÊ×¹ˆˆˆˆÈ©FùÌ”&²:§u2Lä_mßµ;ýü¡‰éO; L_»qxúû;F§ŸöŸÁD¾òÙiéß<ÖDþ›[[f"c<•‰éóAú‚‰ü &r«ù‡í>ÛuHæ3]§/^74Ý<€m öæèãÿx`|Þ÷ù¯o‘÷B¾ô©©¹oí_˜™·»¸øÉ©éÚ½ž~YÊ4seÞúK×͆ò®’úOYžMb¶³øÇ»Ž5‘ÿ㖙Ș2˜DÞˆ‘Y4`1ü0¤s2~X£„t=Dða ðµqÎc@aœP/¯Ë&2e‘Ó& Qòb>ñãx˜SÌÆç‹&2†PqOdŒÐ0‘I_3>Pæ"ó9EŸbËŒ5Œ;L5 èØG8Läø ~DIb‚a¶aìñºsH¹Œç©•vÑ7 @ÖE”ϸҖØ£’óä¶c„’—¾b‚³Æâ¦u˜Ü˜‚¬%L;òS.†dDÐÒGÖ"ùicô-¢£151R9Ïœsœ¾².ÈOý¯1¯ËwÌZÆŽu†ËØðcL4*mŠ¾±Ö0‡É™Xkì1;#bŸ¶ûÆúí3h?óH´ó´ ƒ“µ?*H¿iõ0Ÿä®ÆãÔKû‰Fæu`xǸP×0e³ni})¶:0›1¸ùËz€(›‡,Œ åcˆÇ9ÊÀ¸uU6‘I&2ÇÉQõG\ñæXSÌel¥SŒö¦ý±N#<Þ hoŒcDy±^ÇëEÛø‘îÅ-`8ÎC)¾P_q¾iKÜÿÊ׺ˆˆˆˆÈ©DùÌ”&²:§u¢&X ßù«6§Ek·æm,ªÇ÷a¶îÊQº±ò–»Ò˜ykÒ°·V§U·§u•ã ÖlÉi—®ß–½ûï(cÞÊÍ9]ñX-0*‰ <‘HÜZPQÂ3–nLÓ3Ò´ ¼&š˜}ŒIC?ÙûyÌÛkÒÙ«Óâ5[sþqó×å¨åÜ×í»r>ú€ñÌ–«7mO£æ®ICg¯JËÖmË?¢·¿RÞÊJÙƒyO´£c±´’¦8¦µÀl"îDL´uD×mZ„ò(7Ö1¥0J0kˆŒ Ã”¢Œâž¹G+þH¯1f1¼I† ó†Ym§¢¢éùãßQ6ii_¹¦Ô…9I=RD¾rœ-0´É{ÜÍåSf4Ñ—Mçû÷€vë`òpžüa“7êai;†ul}£˜ñœ§ý±guñ<}¢íô¥Öú"k“ò¹Ô3ô¯ò¸ÓúNù´³86”ßr oä²›ûrßâx±lÖ;çi;ãº+Ží¢}´³Ø/Ò1äeÝÏG±f¸füfþhK´·ØvÆ–¼±¶0¥ËeǸq=q>®5Ú‚ñZ¼þöÞÃ]Š*ÝÛþS¾ëúÞ÷;gÎ83grÎÑ grÎùL2" ˆ‚¢b¢’sÎ9çœs†½a“ÖW÷Ú¬MS;¶tÑsß×ü†Ý]U+uÛµÖ¯žzŠþЇTnéñ”ϸRFº`Uºïcžßο¼Æ,¦nÚ‘?–öÐnÆŽÏ/}Oòâ»Zúûÿ”Í÷‹¶§ýx?Eæç?oþ½Ú»”RJ)¥”j1ϕ⡉,Mc ì+æ*Ê¿AZk”^0Mªkó mŒÑŠÒq¥ו{ªvŸ|¹ ©Ôø¹¥6ç•oGi_’¹Ì~©l/}(£´ÿ¥ï×ëû…ñÉ·¯!]kß“Y”W2±Ú7E—š.¥¯Kß˿΋R}¥u•î“ߟ}K÷Ï×϶TOê¯ù7lCõ§òòû5$öÍÓ•ª¡¶ð:S~[cc˜ŽiH¥Çåo¨=Õ¶—Ö]Ú†ÒíßPÛ“òm»ÜöÔÖƶç/Õåƽ±±i¨o—{t¹¶]n{é¶|»KÛÖÐöÆÊNï7¤†>Ó†ÊNÛÓqéï†öË—›/??®ùí •›¯»¡cSßÚ^º_~{:6¿ï•Ž›RJ)¥”R7JÌ]¥xh"KE“_D+¥”RJ)¥”RJ©âJ¹˜h"KE£‰¬”RJ)¥”RJ)uóH¹˜h"KE£‰¬”RJ)¥”RJ)uóH¹˜h"KE£‰¬”RJ)¥”RJ)uóH¹˜h"KE£‰¬”RJ)¥”RJ)uóH¹˜h"KE£‰¬”RJ)¥”RJ)uóH¹˜h"KE£‰¬”RJ)¥”RJ)uóH¹˜h"KE£‰¬”RJ)¥”RJ)uóH¹˜h"KESSSSïÇH)¥”RJ)¥”RJSx9R<4‘¥¢9wî\¨®®®÷ƒ¤”RJ)¥”RJ)¥Š%¢ñr¤xh"KÅÃÏ©S§4“•RJ)¥”RJ)¥ (<¼ ä⢉,ÿ6œ?>þ)¥”RJ)¥”RJ©âÏFŠ&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ˆˆˆˆˆˆˆˆˆˆ4Š&²ü[q^)¥”RªŒ**ùv*¥”RJÝHIñÑD–ŠçdÍÙ0`Ú¦ð÷çç…vŸ~ ”RJ)U&ý©ïœðäصáЉšü”决ÿè©ÐkÔšðǬMùv*¥”RJÝ(ý°Ç´ðçç…36‡SgÎæ§,R4‘¥¢©®9Z¾²(ü×=#ÂûZŽ h5J)¥”Rªlz6yÏÝ#Âožœ;•ŸºÜ0v®?ï5#¼'›#½¿•s$¥”RJ•Wx6ïÍæ%÷½¶8œ:­‘\D4‘¥¢ymæ–ðÞ#‡Z¾O)¥”Rª¼í:leYnß<ŸUÚnÐÒð_÷ ¯×.¥”RJ©r ïæ–#Â[s¶å§/R4‘¥¢¹³ÿÂx5+ÿä”RJ)UNqó“ǧ‡š3çòÓ—wîÔúîcSÃß;ª^»”RJ)¥Ê©[ZŽ ­.ÊO_¤h"KEóË'fÆ[4ó?JJ)¥”RåÔ[ _|hB8~òL~úò®s¤êtø܃ãÇZk"+¥”RªX"ðwOÏÎO_¤h"KE£‰¬”RJ©"JY)¥”Rª¾4‘‹‹&²T4šÈJ)¥”*¢4‘•RJ)¥êK¹¸h"KE£‰¬”RJ©"JY)¥”Rª¾4‘‹‹&²T4šÈJ)¥”*¢4‘•RJ)¥êK¹¸h"KE£‰¬”RJ©"JY)¥”Rª¾4‘‹‹&²T4šÈJ)¥”*¢4‘•RJ)¥êK¹¸h"KE£‰¬”RJ©"JY)¥”Rª¾4‘‹‹&²T4šÈJ)¥”*¢4‘•RJ)¥êK¹¸h"KE£‰¬”RJ©"JY)¥”Rª¾4‘‹‹&²T4Ía"àÞQ±Œ¤dúï{oÜ¢‹ºnd}y•»~¥”RªUI&2ó„÷¶ÅÂïC­ëïS.Ñ–´U¨6)¥”Rªqi"Md©hšÃDþA÷©á½gÔég½f„[;OŽ ¯wÛ\ýP¦ïtþ'S9?ÔùÍΓÂ÷›Zo›RJ)¥®]•`"3O`¡÷í.SB‹W…Ö¯- yfNøÄcãE÷üþ7ZÌ£>Ón\øiÏéáSíÆÆ×´ëý-Ëß6¥”RJ5,Mä⢉,ÍõšÈ˜ÄK¶ 'kΆS§/êXõé0s;h°¾›F2eÏ^·?L_½7|0û›ÅO~ŸwSÔ?né®°4ƒOf B¼ù}”RJ)uõºÙMd dæ%ÝG¬Žó¢RVl?~ÓgÖuÍÁšCÔÿ¿ÏÍ gÏwõ_#¥ï{uIxhð²ð‘6õ÷WJ)¥Tù¥‰\\4‘¥¢iyó¾aýîc¡å+‹B›×–„öƒ–…á vÄò‡.Ø$ìËBê–ìÇ.ÝÎYj.ßÒ¢6 ?†·4°±­ôVÐTÿ²­G¢MÃ{ï©Ý†JûD¼.=ž…iþÖÒ†êJmIQÎÜîIy¥ÇÎZ³?lØ}<|ªm­‰Ì¶ÿºçÒ6(¥”Rêêt³›ÈÌ!ž·.–7gýð÷ççÅ;§:½½"šÊ{œ ßì<¹n¾Ã¼!Í-˜k¤ ãÌAxôai^’¶—ÎmR9W³?ú~·iaðÜmáç½f„ÿ¸kx˜»nœÛÝèôdJ)¥”º2i"Md©hšÃDÞ²ÿD˜ºzoøovÑ$¾gDØ’-@Vl;>Ú¦6oòÇÚŒ wô[º[Ú¾±4|óÑÉá}YÝ¿L¸ëå…ág=gÄh˜îÃW…‡ÞZ¾Þéb3ÿþù™9¡ó•áÑ¡+Ã_Ÿv˜ºK¶ó7¿{rVè–Û%ÛNJ˜ß/;ö¶¬N¢}þÐwNè:tUhÿæ²ðùlaHŒGÞY÷ÿS¶ÅjŒÊDùÔÕkÔêÐúÕÅá³íÇÅò~Òsz¸í¥ùáëL m^_þ–-§¬ÜÖí:Mdkßé2%<8xYøcßÙ.¾”RJ©kÔÍl"sþgžAÛ™£0×aŽÀ\£– îÐ)›‡¼çn ß‘á·Ù<ƒ™yÎ?^˜WWÎ:Œ­.ßï65´|eqœGÝžÍmØþçgæƹÏÃo/5sº&÷¶þþ_zxBÜÎ܆y{Ÿ dsµ¯tœè]VJ)¥TÁ¤‰\\4‘¥¢i.yúš}Ñ8潶£Y¨ #íŒQ¹Ÿh;6ŒZ¼+ÖyèxMçC4™Üczøt»qaß‘“aßÑS¡ºæL8ZU÷[¹ýHøÂCããä3ã×Ç÷H•Qsælüû¹‰b”̼lqV]s6TŸ:}@„Ï/zÏŒÇMÃ"îxö^¬;cÁƃaÛªp8kãé3çâmœ,œþóîá1jèüùó1EÇc§âþDVcŽ¿0yC¨:u&¶Þž·-Œ_¶;šÈÉú΂mÛþª¬¾ÓáÏ}çÄEa~Ì”RJ)uyÝÌ&2s.R£™ ¥m\àþh›1á'O_~xbœ«ôµ&ÎEj²9 óx'›cªëç=gÄy ó$æ '.ŒÇÔ•{âßi®ÂÜäóƇŸõœÞøþ§Jöß{"|öÁqñ‚8t¼,È æ\yf®ù‘•RJ©‚I¹¸h"KEÓ&òÚlÁqêô¹hÖnÝ_¶¬Š‹ £Ù"ˆµüÇÝÃC×l!u.[±â8N,bˆâýl»qa硪hüÌ^ž³6¶ï×}fF38[ …Ñ‹wÆ(™¯všf¬Þ[DùÌZ»?îKä0†t‹‹ââéå©›Â{²EÛº¬}˜Ì˜Äl8cKÜÿÕì_êúmŸYÑ0º`G|Ð Fó„å{â¾`L˜œµ‘ì—žPw[*}%2šÅ×ôUûªíGbÑÖý'ÂŽU±·´¼¸`TJ)¥ÔÕéf6‘™ôËæ!Ì_~Ø}ZŒ@.ÝŽ‘œRm1;sö||Æí'…/e}fŽ­^Y/PŸË Z±ípŒæ!}Ì7€;¹>‘ÍUÒüä_/Î?Ìög´œý³}¹Cj[#ûñÌÝX@Z2Lh.´oÞ{"Öõñûk¶—ïŸRJ)¥Ê'Mä⢉,Ms˜Èë÷‹Ñ.Ãî£ï #íˆÑÁgΞ ã—îŽi,l:sÿ=8xyh;hY|:9‹ÌÛ_ôšv®c—x­C*`±óØ°Uqö£Ç§ÅÈÚûµNÃmÙBéSmÇ……›j;<ØŽS³ëPuŒ‚&²˜Héi«÷Å¿Ñ=ŲØ£¶<¢™‰J³dwüûݦ…_öžn{iA¼¥tÇÁ글ýZ§Iᩱµ‹. äÿÌÚJ[¦®Ü;#¯á÷OÍŽ·¦æÇJ)¥”RW®›ÝDî?mS,ëG˜È ¤·b¾òÞ#Cß w[q1:=S!õ{èüíñÂ4êëÿcî‘Íu&¯Øç'”Á1ÜýÁ?è6µÞþSòû?S»?s¢?^8–‹ðÌ“¸P¿öÂVÜ]–o·RJ)¥Ê+Mä⢉,Ms˜È˜´SVí ÿÏ¿†ÄE ÌØ×.DÑŸxÙ¶Ã1mÄáã5áh¶0;– £wõŽ£á½fÆ¿ßš»-šºü –Fż>sk4›¿õ蔺EXzõÇœÈÄ[>ӓЉ†.5‘'d‹'M”Ož@øq¶(£ï,1‘G/ÙËì?uS4„içÊmGbªŠcLä§Ç­ §Ïž‹Q@DQÿ¤¬lØ´÷xŒÀ¾pG|ßÈ¥”RêÚu3›ÈÌ9zZËb.Ãü#߯ {ŽÅœÈ¤ÊbnAnâô|†O¶£‡É§Ì…m"‹Û½¹,–ËyÓ¾ñ*æ.äFŒàh"Ÿ¡íµû3?Šûｸ?i*àî—/5‘içÌ5ûâ܇<ÎéÁÂJ)¥”*Ž4‘‹‹&²T4Íe"—>XýŸ;††Žo¯ˆuüê‰Y1ªek¶ØùF§Iñ¡v<¨îî—ÆÛÝúȤ°÷HÃ&ò=Ù><ȳüž{†Ç(_CÎdnñ$¢yA‰‰L»ò&òÄ+0‘‡-Øþpᇸߔá &„[²c^˜¸!œ<}.>\†h!ò òP@ mú?mÕÞh‚åá‰aÐì­ñøû^]rÉ‚Q)¥”RW§›ÙD晘¿Ü•5=›1?!Ís¢ƒÉ? m-­›ç›˜`Ìx¸0ÐßÌæ?î>-ޑŃKMdRk5e"·t©‰Ì][Wb"ÏZ³/¬Ýy4÷Á{ë÷M)¥”Rå•&rqÑD–Š¦9LäMûŽ‡ÕÙbƒÅy‡oï7?[-ûŠ),ȯ‡© /LÚ¾ÛuJøÛsóÂÁã5aùÖÃáKÙö½GO6h"³ !¯2:sÖíÿ¢÷Œ¸pbGžâ…›5‹‰Ì`M6wúMŸYá}Ù<„¹H›×–Ä‹ë\4g¾C °§Æ®Û~šÍW˜?•¦³Àd.MO‰Ì܆9F5ûó Šd"Çý÷\j"³?&2Æ53¯ú?w ½G¯‰é3HÑÅEöÏ÷M)¥”Rå“&rqÑD–ŠæzMdô­G'‡ïu›¾ß}ZxMÚŠ´ð`!ˆXØ´Ì8˜Ë, ‰Òá!/ßé:%|µãÄڜƙ0d)ã3"`ØïÖ¬ž:Ž¸Í“÷0hyÿÖΓ/iÓ·)¯åŠù‹¿þȤº²YÐQ6 ©dð²ÏסÌQ1‘?wõ_ŸhŽþ¬½D"Uýݬ#Oà…c¿‘÷­.“c?Ò{D±ßG³có㥔RJ©Ëëf7‘Q|FBëÚy‘¾¤®àzŸn7î’ùó%æA}vn|ppzöÇ2_aÞRÛžÚý™×°¿‹ó¦vµsæMä2¾šý¹˜ÏßÌóTß²ù\š‡åû¥”RJ©òI¹¸h"KEÓ&2 $:y±)Ý/ÞÚ™íËuSZFéëZãøbdqÚ‡'Œs<—¾_ú:_^~{]Ù¹ýÓ>´»´öçu,3Û–OQQ¯üûjû™o“RJ)¥®\•`"£ÒùÏ--FÔÎ(7ÎsZÕŸçÔ¦Æ`>ÒøܧtÞtÕûçæ\ñõ½µs9 d¥”RªxÒD..šÈRÑ4‡‰¬”RJ)ÕܪY)¥”Rª9¥‰\\4‘¥¢ÑDVJ)¥T¥‰¬”RJ)U_šÈÅEY*Md¥”RJQšÈJ)¥”Rõ¥‰\\4‘¥¢ÑDVJ)¥T¥‰¬”RJ)U_šÈÅEY*Md¥”RJQšÈJ)¥”Rõ¥‰\\4‘¥¢ÑDVJ)¥T¥‰¬”RJ)U_šÈÅEY*Md¥”RJQšÈJ)¥”Rõ¥‰\\4‘¥¢ÑDVJ)¥T¥‰¬”RJ)U_šÈÅEY*Md¥”RJQšÈJ)¥”Rõ¥‰\\4‘¥¢ÑDVJ)¥T¥‰¬”RJ)U_šÈÅEY*Md¥”RJQšÈJ)¥”Rõ¥‰\\4‘¥¢ÑDVJ)¥T¥‰¬”RJ)U_šÈÅEY*šd"(û!RJ)¥”*ŠŠb"çÛ¥”RJ)UNi"Md©h~Õgføï{G…¶£”RJ)Uióå‡Ëg"¡Ãø𑬠ùv)¥”RJ•Sh5*ü¾¯&rÑD–ŠfÉšÍaîÃaûÁ*¥”RJ©Âh˾caþŠ áìÙsùéË»Îé3gü¬î­û×k—RJ)¥T9µ~ç¡°lí–üôE €&²T4;wìgO×äß)/çÏ…Û·…sçn¼‰|öìÙXwÖˆü&‘²r¦æTصsgþm)šÈRÑlß¾=œ:u*ÿ¶ˆˆˆHYÁÈݺukÙLärÕ-"""Ò'Ož ;vìÈ¿-@Y*Md)"å4rËY·ˆˆˆHSh"Md©h4‘EDD¤ˆ”ÓÈ-gÝ""""M¡‰\\4‘¥¢ÑD‘"RN#·œu‹ˆˆˆ4…&rqÑD–ŠFYDDDŠH9ÜrÖ-"""ÒšÈÅEY*Md)"å4rËY·ˆˆˆHSh"Md©h4‘EDD¤ˆ”ÓÈ-gÝ""""M¡‰\\4‘¥¢ÑD‘"RN#·œu‹ˆˆˆ4…&rqÑD–ŠFYDDDŠH9ÜrÖ-"""ÒšÈÅEY*Md)"å4rËY·ˆˆˆHSh"Md©h4‘EDD¤ˆ”ÓÈ-gÝ""""M¡‰\\4‘¥¢ÑD‘"RN#·œu‹ˆˆˆ4…&rqÑD–ŠæzMdY ‡.ÛBëôéÓaÿþýáÈ‘#ùM7œªªªØþ‘k§œFn9ë¾ÑÔÔÔ„„£Gæ7‰ˆˆHÑD..šÈRÑ\¯‰Ì×Ã?{ì±²™¸«W¯÷Ýw_èÛ·o\•“aÆ…Ö­[‡Q£Få7‰ˆˆÈUPN#·œußh–,YçQ½{÷çÏŸfòÔ©SìY³â8ˆˆˆH±ÐD..šÈRÑ\¯‰ÌëÞ{ï >ø`Ù"X–/_n»í¶ðøãǨär2hРð¯ý+¼óÎ;ùM"""r”ÓÈ-gÝ7š7†^x¡nî²yóæp×]w…¶mÛFSYDDDŠ…&rqÑD–Š¦9Ld¢WˆFÞ²eKX°`A˜>}z|¿akÖ¬‰ÛæÏŸS>$öìÙ-Zöîݢ6mZ˜={v½Èfö›;wnܾtéÒºv¯X±"ÜqÇ1‚f÷îÝñXö¡o‰µk×Æc(sÙ²e±ëÖ­‹‹#êä5e;v¬î©ÍS¦L .¬Û^]]ËÙ°aC4Ï©“}ß~ûípûí·‡!C†ÄýØF„ýË÷GDDD§œFnsÖ½sçÎ0gÎœ87aÎ’¿à¶3ßÈo§ Ì#HÆ\†2¸xNû˜óÌœ93F ïÛ·/îϱlç.-R1çš1cFœÑÞ'ʘ2ÓÝ[ÌOx™L9cÆŒ‰íÚµ‹óºC‡ÕµGDDDÊ&rqÑD–Š¦9Lä6mÚ„ûï¿?tèÐ!F®Œ±ÌÂ0\_~ùåp÷ÝwG³=ôÐCaÕªUq;©0^}ôÑXdz‘ÅÇû°ˆiß¾}ÜmwÞygxâ‰'¢©K9Ô‹‘ØÎ~DÐM¤Û þîÝ»ÇÙNûõëÿMåbDÓ^V¯¾újÜ—²[´hÿíÑ£G8qâDìwË–-cŸ)û¯ýk,‹t”5bĈØvÊ#2™Ž‘+£9Ü«¥¹êÆäen“æÌ%ž}öÙºù Æqévæ7Ï<óLÝÝ]ýû÷ï3GjÕªUÜïž{î /¾øb¼ ,×±cÇh3×`^ƒºtéR7'b¾òüóÏÇc™gQÏk¯½ëX¼xq|Íœeüøñ±˜Èˆc¹P."""ÅA¹¸h"KEÓ\&2‹ŒR"\ <òH4d‰haÂöõëׇ‘#GÆýYÜP7ÛÙŸEÔäÉ“cÄÌ<ß#"‡HJƒ9K„M·nÝÂ?þñøšÈÌ^ŒÞÁƒǨߴp"20)ÃCšå!ÚCduRÑÅDã`³(£)Uó¶dx‹ˆˆH1ÐD..šÈRÑ4—‰ŒH7DÏÌ‚„H`D,^0t‰fa¡„ñÊÃvܸqqñCTM‚(cÌXnï$  ª>}úÔ-æh7å`Z‰Ì¾:uª‹öÅæ_LdRVeR'ftºm”v² ²3gÎÄþ`³¸ÂèN‹2ÒZ$™EÛ,ú؇‡ëñ/‹E#EDD®žæ0r¯•æ¨›¨^æ\TNlÚ´)Î_øc9?ÿáÂ8¦2ó(æ'Ìc˜Ï¤ö²`dŽÁÝV)Mf2w=1Ÿââ=óæ(Û¶m‹Û¹KŠ2¸+ ؇㙻íÚµ+Ή˜³p·0wÃDƸöÁz"""ÅC¹¸h"KEÓ&2‹ ܩ‚ƒ… "`ˆþÅlÅ8Nêܹs4X1iYd±¸!‚9Á­œÃäù#2¹ÔΓr"c§¾Lš4©ž‰Œ‘RhÏ8ÓÆÀmœ˜È´#ü©§žŠÑÑô>ÐnúD¤3ãÆû¤Ø(ÍuœLd€ô$"""WGs¹×JsÔ=tèÐhîòoC4´gF`î2ç`^Å܈}¸c ˆnÆ &J¸ªª*¾ÇÜ#o"c§çOyÌö·Þz+¾æâ6å7f"s7V2‘SÞd)šÈÅEY*šæ0‘Y„`¨ò€:ày˜«¨ü°%s–(`êbñB4/·r²@"Ê·1™ÈeÒW©CŽä´˜!B™Û7Y|¥&2·ŠBS&2&1ÇÿóŸÿŒF2&1Ñ@)š¹ÔD&R(o"Ó–—^z)Œ;6ÖÉXp˪ˆˆˆ\9Íaä^+ÍQ7s æO>ùdLš“Üá€IDATM‘ÞcþBú.æAi{ª‡üĤ«àb{ŠDÆf_(5‘ÓN™Èéa{ÉD&åh"‹ˆˆÜÜh"Md©hšËDN‹¢†1^ÓƒñX|L™2%¾&'yöÈŒÁ\š3¹1™(^ŒföÅÀe"—1f9“yÍš5Ín"§èhÞäNO*gaG¤rS&2Ç‘«™h¯^½êneÍ?]DDD§9ŒÜk¥9ê&˜Ú1ÿ4hP=zt4x™ ‹ùF-s7Þx#ÅÜÙÅvæ0`À€n"Óoî¨bžC ^‹ˆˆHqÐD..šÈRÑ4‡‰LÊ¢j0j1~YˆAC$ `$2$.zØÎb‰§Œó°`¡„ ÌÃ阮,º0¥aݺu11ïQ ÊdqGŠ ÞÇðM}!‡2“1ÁŒœÒYÝL9ÞÉDÆPNé*ˆ¦¦¼æ8ÌbLpþ&‡azX‹½R™å’0 ÓCuHÍ!"""WFs¹×JsÕÍÅkæÌS˜ÿ0gbÞ“Êe;NóæJÌ%Òy¹Ž©œÒY`úRæs2‘ÙŸ}¸Ž‰ŒAÌÅöd"󜶗¦³À„Æl¦<.¬3Oáâ?í¢î¾}ûÆ‹çG€ˆˆˆMä⢉,ÍõšÈ°‡ŠOó¦â‚9ÜÐÃäØm<°…EN‚¿8pÉÓ¿)ïàÁƒ—´¿i/aÿ¾ìKãt»(?ª¥e²}R40Æ6ÛKIû¤:éæ5f8§vò7‹KúsøðáK˜lcŸ”§0™S]"""re4—‘{-4gÝ̘1§`î'mgŽÄ<¤æ1Ì!’©\:ÿHó—4÷ Þc;JÅcN’æ/߇9õ–^ç=.¨sѼtÎ&"""åG¹¸h"KEs½&²ˆˆˆÈ»As¹WK9ëi Mä⢉,&²ˆˆˆ‘r¹å¬[DDD¤)4‘‹‹&²T4šÈ"""RDÊiä–³n‘¦ÐD..šÈRÑh"‹ˆˆH)§‘[κEDDDšB¹¸h"KE£‰,"""E¤œFn9ëi Mä⢉,&²ˆˆˆ‘r¹å¬[DDD¤)4‘‹‹&²T4šÈ"""RDÊiä–³n‘¦ÐD..šÈRÑh"‹ˆˆH)§‘[κEDDDšB¹¸h"KE£‰,""ru:t(9r$ÿv³QSSöïßߨ¨ÿüùóùÃêÁ>ÇÏoº)(§‘[κEDÊÍÉÓgÃÌ5ûÂøe»ÃÊíGÂÉš³aסê‹:Xö9Neû])ÕYs§¯w…%›…ñKw…9ëöÇþå9{î|Ø}èdlgžªì½=‡OÆ}®•#U§ÃÁ㕳îf,“šÆòF@Ýkw 5g®ÿ<Íwò\ö¥bu¹ï*u7ô}i ö=uúúÛ šÈÅEY*Md©6lØ,XpEëåX¶lYX¾|yþí°oß¾ððÇ.]º¼kæìªU«BûöíÃ>XOíÚµ O<ñD8}útþ°z0mÛ¶ O>ùd8sæL~sá)§‘[κEDÊÍŽƒÕá³íÆ…t›úMÞ&.Û¾ÐaBøÊÃÃW:NŒÿ~ùá á›'…×fn‰¦ÛåµhgøÉãÓñêËŸ¿Þ úŒ^¾ÛuJøjÖö݇ªó›£9þµŽ“Âð…õ¹qKw…o=:9lÙw"¿éŠé5jMøûóóòo7È´U{ÃÄå{òoŠ=‡«ãwaòÊßΑ‹v„Ï?8>ÜúȤ°vçÑü櫦ë°UáGݧ…Ÿ÷šõ‹Þ3Â/3µ¸$›ÔÿnW:?O¾_ÏÚðÏæ‡5M´cǪÐò•EáÙ/·>29<üÖŠpè:/(h"Md©h4‘ED¤R:thèÖ­[þíkâ¹çž /½ôRþí0a„кuëвeË0{öìüæfáĉÑFÊ÷Ýw_èß¿Ø´iSX¿~}1#ÿvƒÜ=`Qøßçææß.\ xzܺ°~ϱü¦w¿=?/üé™9aßÑS͉6µÁï6‘Çî;'~÷ù š½5þ·ôÀK.éÜšÈÅEY*Md¹Ù!òsã÷‘G .¬‹æ·dÉ’0uêÔh¤–Fä>|8Ì;7Lš4)ƒy +V¬?þxèÕ«W<6'‰þ}ôÑGÃàÁƒ£ÁܧOŸ:“qçÎ1r¹Ôt$åÅÒ¥KCUUU4ƒy½k×®0oÞ¼hRÒ–•+W†iÓ¦…ùóçÇöäaŸûï¿? >ü’÷÷îÝæÌ™¦OŸûžçèÑ£1jy̘1¡GáÕW_­ÛF©7µ'ÕÍ{ôuݺuuQØÇŽ ‹-ŠãÇ{i,H«AßRD4)8+"µ˜°k×®½æÈðr¹å¬[D¤Ü`¨~ù¡ aé–Cñ5&2¦ÚÖý—Fâ8v*|â±áù‰ÏC»W‡¡ó·G³lýî‹ãyÛÃ7™t‰Ñ¶tËáðÊôÍ1Ò÷ð‰šº÷ToÌÚºI+ö\’:cë¾±ü~S7†)«ö†Ó9#qI#å.Üt0FÏ6f"¬Í˜0xN}yÄÂÚþoÚ{v>™µkwÕƒ¡~úì¥m ²™q{yÚ¦°jû‘ÐeÈÊðë>3붓ŽaH6N/fý–•ŸR],Øx0üöÉY1vÊÊ=áXuí¼³”¾˜¾)Ì]¿¿^ŸÃsÞ†õ¢bWï8ßO¬Úq$F¾ïèɺ÷o>ǶZ¼3ÈÚt¤ª& ÍúA;y×@{eãY:¾;VG£tàŒÍq¬ÓùŸÔó³úwÒ£ðÙpA‚ úŸê]RO)Ô9+ûtŸ£IM’Æ}cöÙðÙѯÛ.¦;~òL¼x@ô0ß‹qËv×mkŒYû~üøôøçÁ<'ùÑìóLðßÉß;*šÑy–o;·‘&Ñ?ëãgÛ‹ãu­h"Md©h4‘EDäf£oÀ€1êöÞ{ï ;wŽ (F/f*F*½/¿ür4f1A1„IAz¶ñ/f*Çð=öØcuæîêÕ«cùÛ¶m‹Æj«V­ê¢‚‰JnÑ¢Å%Q¸¤¥8xð`èСCxöÙgcz ê$Wñ‹/¾£„»wï÷{衇êÂœ£1‘‡ V÷u?ðÀ1­í£o¼ñF4>³fÍŠãÁ=Íþ)3 êKí¡ bÆ©_¿~áž{î‰m£/ŒûÐFÚúÔSOÅ21äéoj/¯ï¸ãŽý |&Œùk¯½Vצ«¥œFn9ë)7ÉD&ê’‰Ìmù¥`Ž~ Õ¨ðöÜZãuîúá›'Ç4¤®à˜Á¶%³³÷èµásíLJ_=13¦% @J1lÁöð©¶cÃÿt׶Ú Z·aNÕùý¬ÌVŒßû_[Îœ=£I3Ê¥¾÷˜¶î«m7íõ˜È˜è÷û}+ëçwºL ¿è5#é­. 'NÖ¾¦_ë4)|=ÛéßW;Nì[{ÓªGÃ7œ½71üî©ÙáÓíƆ?<=;gÆ„HÞÿ}vnÜNY˜¶¤†HõÐÏC¹ \€¸íÅùáãÙxóc$ŸÈÆà­ì3#ýÊ÷›š}ÞÓcÝ}ƮǬȾ£¤¾h1`QøÈ}£ÃŸ³>6Ų­‡³ãÇÆ”" Á…„gåL¼Ðo¨5–§„ç&Ô¿¨©Ný¥†1mú@«‘—ûW‹&rqÑD–ŠFYDD*Îeo¿ýv4;‰BÆü4hPÌ]œÌÓ7Fc˜¨c"1K“é»yóæЩS§9ËÄœ¨æ^x!TWW×EÒ 80£Àû˜®ï¼óN|±ŠÙ:jÔ¨ø£ºk×®±M’ÇÐ3f̈QÂÌ¤Ä ê™h^LeÊ&Gi¾ã¼‰Œ¡M=´~L™wß}wŒ(úÞ»wïhQËÀDÔztŒOe}?kæ'Ÿ9&2Qâ|ößè4)üëÅùÑøåsbÌÞßjd¼ @~ѻ֌&%Û¹0Á˜­ÉÚDJêY¾õÂEûGc‰¸.…o§ûýS³Ãƒƒ—Å¿1uIñp6¦Œ}"%&-}^·ûXì7ˆ /~ÎC»þ’}®ÿx¡ñÜÕî”Í…‚ŸùO{ÎÝF¬.Ù³–W§o‰¹¿hˆoÞw"~FדOZ¹¸h"KE£‰,""•)zöìÿƈMÑ¿Ƥ~àߎ;ÖåÆD%mæ*¯1h“iH$.QË LQ"{y3š|ÅD,cºb(Q·˜À,–ˆVfÊÅÄPM)%â‚°{÷ØŽRˆ0¦M¤ÆHäMä´OéÂ6cX¿òÊ+ñ5æ'56mŧßËÔM™´­4J˜þ`"Á ôØ4¥P&}¦NLu¢·éãLÿ¨‡”¤ìÀd¾ž‡–ÓÈ-gÝ""妞‰¼pg4ùÈýJ$*ÂäÃì9²6çþÒ­‡Â‡ïž³&F¿“ Ãùý-G†‘‹vF³£®ãÛËc.éØô½_è0>ì?z* £—úÙS™ÈÞãÕg¢9ùß­FEó“ãÖï:Ž^H}@ä+Fei¹-b¹b´è¢M‡.k"ÕšgdÎD&=†d)äé%JvK¶ãRš¾0†ýDm:‹ƒÇk®ÃÕaí®£1uCëWÇñ]¸ñ`ÜN”1e¦èOŸ Þ¡ .ô+û÷{M i ª–¨eƹïøÚ¨Øg²‰þ%J›¼¿ŸÏÆ£œr'ÆŒÈg"ž©ãÚ˜¹l#‚˜Débc8“¶ä3íÇ…©«÷ÆÏéC­GÅí¥üî©Y±ßÔý“žÓ£žØy¨*ÚäÓÞ°çX¬‡ï©00“S= AŽáG†¬ˆ¿5w[ÖŽK#}«OQÖ}Ƭ9›ÿã®á ^ÈC[ø욊^¾íH¸¥ÅˆKLd¾Ó˜È ¿Ø¿ãF„uéƒô6gß!LäIšÈ‰&²T4šÈ""R)`´b"cü‘ߘHaŒäá‹q‹é™ò YK´.i$0fÉœ \LäR“7Eûbè’ñ7)ˆlÒ]`¼’g˜DåbFµËþ¼)Š·4E¤HéÒ‡àåMdÚÁ>)sâùçŸÏ<óLü{ĈõÚÊ1´mË–-Ñ`ÇDž8qbÝñŒû³˜°O>’xÊ”)ñ}úµxñâÍ"†\Ô˜Õ˜ò¤ö ¼Ô„¿Êiä–³n‘r“7‘S:‹µ;ÅHaDhïÑkÂGï#+1à’ÑŒYŠñø›>³ÂûÎ ÓVï‹F2æ&ÃØýh›1Ñ4d?Dt)F2Ñ»oÎÞM@Ž¿³ÿ‚ðû§ggÇNŽæ+Œ[º;üã…ùÑ0&•ÄÝƨÞ{²ã1‚Sý©\RM`â™Ú˜‰Lþ[LMrêæ!]†é¶UaýžãÑð$¯p)÷½¶$ü¬×Œ˜ªC•|Á¥C7™Èäè%ÍÆO{Nfø]//ŒÑÂÉDæ…Dõ9‡¿ûØ”ðÕ¬¯i¼~•éOÙ¸v~§ÖHÍÓ9«ëgYÙD}sÌ“c×Å÷;½½"|ô¾ÑY;fÕÏoŸšS`ð…ûÐ[µi/ØÜ|¤‰ ¼ »ÅHeÆ„Ï–TŸ{p|Ø–Ë™M¿ØŸï ÑâD'øŽ‘‚bê…´änõûÛêa¬Öíjø¡}ôˆwxnÒ†˜äxõ¥Ƥ Á '·6¦?Ÿýå¸ÿõ%1šžhðÆ ²™‹ô9A¬÷ä˜Ú¥ðÝùRöߦ{sþ-G†Y"æ¯Mä⢉,&²ˆˆT ­Á€ H®_R/}ËkÄÃßxp‘Âü DóP9Œdr &2y–#s´oß¾1¿1Qºˆ4DÞ’ú‚:0‡Iÿ0dÈõKÊ à<‹ÙŠ›Êc{2}DÓ†Ò‡ÓåMdr1œÚÔK;0nùãœ(ãÔN„!Š‘Ìxnc¹4Ê™6& ìŸÒU$ˆ>&Ec‰‘Íó‘ИÓ^ó.cbQ·DÓÊCÙzŽZó1SG*— ¤o ¨“¶1™¶’ƒ¹ÍkKò›¢ùúåì8Ò%ð°@¢«‰¬.3–¨[̵]/ò•(_R\Ðg UÌñ”ZaËþªðñl숔†–‡Û_ZÿÆ€%Ä#LÓ›ýóÖ71»dË¡ÕMT5Q°©@z‘Ú¼ÔgêÆ‡Ï ƒŸè] ÿ®C/>0£=Fë’“˜‰8æ)Æ*&+¡Œ!MT1†-ǤÃÈCÔ/¦.yƒ:1_1d1™!o"†1õOž<9¦Ñ Š—èdŒe ÌDÞD®ªªŠF5ÑÕáäR¦Íô‹Hf"‡KxW ‘ÁÔÁCiOS&2æ)&:uc\3~´™zÆ_wûÜ~ûíu†û‚ âxs™\Ë×C9ÜrÖ-"Rn2‘É‹YV 9z1ð†ÌßMCLR"1îØÖcÄêqŒQIJ MLÂå[ÄÈÛ¾1÷îɘÖëýí¹yѤ%‡í€é›b Š%[‡ïwŸÓ0ì:X<÷õNc~_ŒO G" 1Bcú…¬\R8)‹1M*‡??37æÓm*'2tº2įduc“Vcƒ°Ý ZCC–\¿¤óX¼ù`Üçå©›¢QˆY1{WÿÑ\°¡vû€i›c¹<Ô #ú[]¦„û^]ÍNŒù»ú/ ·´&\HAT3¦+ã}&×þYùDI“³“(XŒöRS¶Œj>‹Ï¶s7_HQë"—2‘ÊôoSV>Ÿuq µ+5‘ÍÚ>œ-õÑn>Ç[;OŠÇc¦~:+kbÖ¦šÓçÂÏ{ψÑÕäbæ³'g6c‚ĨÔ€‰ü™¶ãÂœõâ…Æxjªgï‰hvcÞ7fw2‘_"Ò1ééåòYaÒrAbMöÝâ;x9™‡ÆÏ`yý³×î ÞZ¶8¿GD¶“ë›´|‰¬ç¿í<ùڌ͡Ϙ51êžy"¥‰¦§m|¦äç5;®Mä⢉,&²ˆˆT ˜¤Z ¹ˆ™`%Ë{¤¬ÀÜ%o/ѺDÑ1Ë6¢x1P‰NÆ/†)é0lIkÁñù€aMôpŠ:Æ„å8|—,Äyó¶Ô´%Ó“–²ù—Ùað–±ô³:A¤æÀô&꘶§‡æAÁ\jD'0˜1xif{:ȉL1͘ÛiŒh#ÿi\:o ÷1fx*‹¨gÚ…©œú­”ÓÈ-gÝ""å&o"cÐÙŠéX Ùw»N í/¬¤øuŸYÑ`ãAx</Eäò¶o?:9FQ²­˜ºD´’ö3ú7–Æ4˜Î¤r —Hè^_SJüæÉYñõ­'ÇÈXÒ``8Ãà9[/)—|ÊÉü¾œ‰LÛ0 1¢i?ûÅzÏ€EÑ4LIÚ†‘™ê!j–ü»gÏÖž÷¿ÿ}vn¬Ÿí?è6í‚™]›Ã˜ˆ`Rb¢âÖ¬Œ¦4Û#•Âè%»â±¶D«Öæ‘^ `êÄ°lýê’pü䙺¶çyeúæ𞻇ÇüÔ¥ ƒ–þ}6;Ú•rÓ>¢ÄD,ßûêâ8Ö<ŽïyŽ·í¯ ûžŒ¯1‘Þ‡òÙÓgR8ßâÃ{Ï /MÞXWöÎlŒ¾ÔaB˜¹v_Œ8'…Gi=˜Ûä—nŒñÎ%f7Ù¤¡øâCãc½\dµ¸ö¢øÊ ̧ɃÏ犙œç¹‰ëÃ{ïQ—¢3™Ès"©é+ÿ­àA…ÛéÁ‰äÏþVöÝÿR6æ|_ó˙ڗC¹¸h"KE£‰,""•ѯ¤œ 9A´.¦2q©±‰Q¸{÷î…ËD¼ôö#úS”kü–CÌÖÒÈñ:=l(‹÷˜ðçáýÍ›7Ƕa,çIÇ––á¾ùº2»C”þ°>•¶‡º96o>s é5h#ÇæIÇÑžDc}½ZÊiä–³n‘r“Lä•Ûj.L¾ÝGNÆÔ y0W÷¹ø›Oôå†ÝÇcJ…ýŒW¨ÊÊ B•¨ÔÄ¡51‡ðêG/ÉCKô-ɘ›É¼ÅN‘žÔi¹pÓÁh<ç/Z––˃éK·nÒDNÐÿ9ëDóèÛR(3š’Yù´gÑæC—ô?qúì¹Øfú€9}âÔ™KÆ#“~Ç…ÂLwöÝ™µ³4ZuWöšÜ¾›÷]þ¡µ·ñø³õÏc´“Ï—ÈßcLþí¡¾ÙX1žÆ³öáz'/ùܨ“èßÛx0ÞÅy=¢?¥¦7Ç3n©»æ œFî»U7s.æ/‡Îo¹f8oUWWÇârÛkjjâöÆÀ´¼Zß •ÙÔ¼áĉ¶©!.×N¶>}esoÚÕÔþiÌ®t.Àø0NAYW;†¡‰\\4‘¥¢¹^™ã;tèºtéŽ9’ß|Ep"˜9sf˜0aB8~üx~sƒ¼ñÆáÞ{ïLjˆŒÅ?<5;ô¿>¿éšX°ñ`¶`GhNÿCòLß{lj8xìÚÏWÊ™³çÃ_Ÿ>Ýnl˜·á@~s=¶î« ßî29,Þ|(¿éªaýÂøÍY¿?¿éŠyjܺð©¶cמibü÷óƇ_=13¬Ûu,¿{³3{ݾð­G'‡Çjò›n–í‰õ?Ùðb«¹y·ŒÜ+áݪ{ÅŠá½{÷Žâ/s«I“&…éÓ§_Õ"[DDþ}˜7o^èÚµkx衇BçÎÃŒ3êÎ_¥³fÍŠï³u;û'ÎzóÍ7ÃÃ?×õO?ýtضm[Ýö]»v…gŸ}6Ëö×^{í²ë~ ÎwÞy§®Ì>}ú„-[¶äw‹õ<úè£açÎùMõ ƒ ª+³oß¾—˜¨›7oO=õTl':th“fóš5kBÏž=ã¾;v cÆŒ©3“3λiÌÛ… æJ¸ãñÊ+¯ÔÕýüóχ}ûöÕmß½{wCÚ º¾¹­&rqÑD–ŠæzMdX­[·Ž?†GÍo¾"¸úøàƒ†Ûn»-Fç\ /¿ür¸ýöÛèQ£ò›DD Ã{L \“ûšè5jMøÉãÓóo_¯LßMÑÜ7:¼9{k~s³³dó¡ðñûDŽ϶îci~s=¶ì« Ÿ}p\˜·á`~Ó5Áøuº"ÿöÓuتðå‡&„ù„e[‡¥[‡ w„¯tœþúìÜw="yÚê½á3ÙØ•ËDµxW¼p£¢ëß-#÷Jx·ê^ºti¸çž{â"¾±È§RXtßu×]Ñxn*’KDDþ=áâ$çÌÝåË—Ç9o,^¼¸nûÝwß^}õÕxbÝ¢E‹°qãƸ}àÀá¾û)F)Æ)"¸ -Z´(L›6-´iÓ&ôïß¿ÉsÒˆ#b|Qæc=:uê`À4ž?~èÖ­[ Óh Îů© ,<òH4 ¹3ùàÁƒÑOÀ¬^²dI¬·U«Vaøðáù¢"ÜIÍyõÉ'ŸŒû³c4eÊ”¸qâõ믿ÿ~饗BË–-£Q‡¶1¦x"˜÷ô´å‰'žˆçyæ0«é?íf ©»_¿~MŽáåÐD..šÈRÑ4‡‰Ì9W7mÚæÌ™&Ožÿ.…Û6ùA8qbܾråÊøƒ‹Ìmûöíã3?ø” l_½zu˜:ujÔ† êÊã$rÇw„Ñ£G‡õë×ÇrgÏž}Å‘Ì"Rlm:úOÛ^˜¼1L_½7œ¹`Îí˶_Œ¦m­‰LÓ³þ2>ô™ïI‚Ï“zýîcáµ™[ÂÙ8ó7Ð÷‰Ëw‡ït™þùâ¼ø , åÈÛó¶‡Ùç¶jûÅÈšmªâwˆ1»xË¡Ðsäêð?ÙgÀ¸—B„òç³þ?y1²tmÖÊ|{Þ¶˜ö"Áw„¶ò¥hëþuÛaçÁêðNÖ¦—&×~ÒçÂwä &„µ»Ž…!Ù÷ÿ>æ–DVÓ_Ìmê4{küî3ƧNŸ c—ìŠã:kíþK-Ûö×îËwšÿ~¨ýn‘:ƒ±>–}Þü76iÅž0|áŽx }oO>Ç•q‚݇«Ãà¬N¾ßÎ¥ãq-¼[Fî•Ð\us< x¦,T™ 1ßaAžLd¢¼ˆ?~|\\¦…=‹Ü‘#GÆvÛ¶mÃܹscú0`þÄ‚ž9 ô´0‘/^xá…h¦¦óç¯îÝ»Çõ2`€rá2ÝÍ‚1Œ¡Iô1´œ“¸ã%Á9èÎ;ïŒkìµk×Æ¿Óy ˆØeíÞX42ç'ÊçÎáD1e9omÃH½ÿþû£.g"s>Ä”N&/оÔ6Ρ˜Û¥iÆ ‹fnCíäÜ™ßÿ¹çž‹¦6ãʦ¿£¿cðàÁñ5óÎéø(´ó4e&ØFÙƒC91dÈ8F×ãÃh"Md©hšÃDæq2áª'Âü²ØêàG’y~@¹RÊ*WüøÑçÄ…8†}¸âÇŠ«¥ìÇ{”É>œ 8Ar»åp•›íˆÛZ¸Ú'"7/ÏNX>ùÀ˜5úë'f†µž¿.n{}ö–ðG&…¿=7/¼¿ÕÈÐ}Äêpgÿ…á/}ç„ÿé2%|ºÝ¸hvº2|"+ã½g„ïu>×~\¹¨öV¹_fïÝýòÅ[Òþú윉›Œ¾1KvÅèÜRãˆ:%}õÿ1«Ù_öžYgÜöŸº)FÙþ¨Ç´ðÓ¬í·fûQNŸ1kãö©+÷fíÌß?5;î{G¿…ü¶ìûílíýeÖïO<06´¸(æ¡í2de4)9æ§=§‡•ÛŽÄ÷[ X¾ôЄð»'gÅYÚuøDM4·¿Öibh™mÿdVîϳ>³)¤Çøh›1aõŽ£aÚª½á#mFGSúOÝ>•ÕO¾äÄÃo/ýe\êÊÎö!•õÝ•}ôá7}f…¯g}Gþ‰CÇkÂ;LˆãÄþ|>§o®ÛNûZ½R[&õ0N¤²à3]˜µ#³›qa|~œñG³6clf8õµãÄìsŸ¾Ûujì;Ÿ;fæž®ó/<4>üýùyÑDÇtþfçÉá‡Ý¦Æ:9ƒ0ªIÝð—gæ„÷·^˜´!<9v]üŒªO;fìáªÓáoYyÿ|a~ü,S–Ïš2•M¤2¦<´}ciøvVî{LÏÆjføLÖ?Ú@®eÀä%U}ç³æ»ÙæÕ%±ÿìôÿûY™|Ç0´?xï¨0hVíÅ×'F¯ _ÉŽeÛo²ïß™_guÜöâ‚ð³^3²òfÄý_ŸY{K))EHÇA?Æú}g‡ÓgÏÅ &|ß[½²8Ž‘ÖÃlŸËÊä"&3ŸùW²}ø.aŠÓ>·?eßCÊâß×a$7—‘{-4GÝ”ADóæ(,€¹s‹¹QY”M¤ hæ:̃˜G±c˜‹ì¼æ}ÄÅód$÷èÑ#ƒxŸÅîž={òM‘ ‡ ,RC&H•Ä9µ4çά¯K! –èXî$f]ÍÅÌësÎ=˜´DÞ¾õÖ[—¤… ˆ s¶±‹—œ‹8ßa¤&0a‰æÂ(PÁ_–±–¿œ‰L@í,=ÏÑ>ÚI¹;™òK} ΣœKј،Q)˜À¤¢ÀXfÌ’aœ E©¨€HlÊ&eõówiïsÞ¦ øÔEû1~s¢”1ªD®L4‘¥¢i.™NNFëÖ­‹')~Ðù!ç„A>$ÌàôÃIÛ1–¹2Ç1ü`s²a1Å+·©°à"?ѲeËbd·ˆPmN·ép·Â‰Ì ˆ2Js8‰ÈÍŽ”ÏO¼8á{à¥Ñø…Ás·…ÿ{ç°h–‰ñˆy÷Þ#Â3ãׇýÇNEø––#Â[Ù¾D¶bbÝ÷êâh–a°¾4eS4ÚxcóË'D“Œ_è0xyøË3sùÜÄ®Ó;+¢©L1"€1øˆ8ݲïx,£ÇˆÕ1:“ÈWLÇÿ¼{xxrl­‰üÏæ…Ûû-ˆæ%>âcÙ1G²6­Ø~$|8+‹}h#Q²“–ï h52F.S^«ÂÏ{͈‘¡´mÀÔMÑäÛr!Š•(fÌHê#ñ YŸÇ/Û_çé0xY4[6cú>:te|½-+“—HV Ï˜çÀD îÁqá‹YÙ²²fŸÙ‹“6Ʊ˜¹¶6j–èWŒr õd^¿#¦Dþæ?tŠ²%b÷¿ï#^é?ãpgÿq –n>£u‰‚¥N`âÍçÑ [ÞsÏððNV&(mnóú’h²N;~‘Ÿãɬ¿¼Æ¤mÿ沬°XDŒ2ÇÄf,‰"æ»ö÷ççÆÈuƈHdÚƒiAúÛ>³¢1þé¶ã¢! v‹ý|gní¹c¹ý›KãEÚÕá­åñÂf1Í\ôøH›1uŸ3ñ_Ÿ›¿|WgmÂ,_¾ípŒ:¦þÎCVÄö`æb`cÑ·´Æ/ÝÇpκýñuëìûOŸkNŸ‹eÿã…Ú\ˆ¼±$üù™9цqÙq|'k¢Š p¢ÜùœùoäóY¹¸p÷Ë‹¢¡¾zgm*«'ol=º.ªšñàóO&®…æ0r¯•æ¨;¥®À$&JŠÅ2æ1sþeNÄ­´éN,^¿ýöÛqþóÌ3ÏÄ9wx±åB= TÝäUÄ<æÂ:ïq+2F2Ñf×Ó^¹ùá|¹‡»`®"è*™· ÖëÜ“‡ux¯^½bdsCwøQÜ®]»º(ç† Uëtî(N`¤>þøã1ú¹ Ø+1‘óÐNÊ#šs'Q¾œK9Ïçð”þ²´ aúY eÐ7Ϋôìر—lÇË ÂȧŒqB úPjèbzc´s75Ð.| ÎåÜYļ ¡ÑWƒ&rqÑD–Š¦¹Ld~Óí ü¨cþròb;Wùƒ˜š_|±.‡ÛKÒ-5¼—®.’ПQé-+,Æ0“1žY0qRH'"Nia¶jÕªºcDäæ"ÜΆ‘÷?]§†_õ©5;1†ÿëžaõŽ‹Æ†ïoŸœ•|Àðà›ËbÔe©—ÂmþbSVî ›ö/Ýz8Öñ³ž3Â=Å㨟‡Ì½2móŃ/€‹±†)Fº¢™“‰LjÌ·Òô ˜”_é81ô]›Î‚è[úµdË¡˜>áÏÏÎæàã§bŠ¢fç]ˆVMýyÛKµ·¿=1fm4]“·M„ì/²×ogc‚ÙËØ}Jô5©'02ËÇL[0d[ \œµçplÓß_˜dÌBø׋ bÄ.ÌZ»/šÒ”˱Dì–¦éÀ0¾{À¥!õýc|€öb¤’‚‚úº _Mc¢°¡ÕÀEufh‚4|V 6ˆH‘±}UŒœæBCêt²2Fà–šÿºñs_Q{naüº«]`ð"ªó£–1$G4mzcÖÖ0({M´ûæ½QD¾ó9q±à‰ìsE¿µ¼.’S—(䵩²¦LÊî’ÕqL ¾g· ÌÞïgß¹§Ç¯ ³ºhÓìµST`BoÜ}<–Mj‹OgõoÞ{1ýi(~ØcZü›ïcœ"[0죉Ë/Fî×™‹0ìr¸ª&椫š}‡o{i~¼x€iÎXó¦] LäOeß+þ{û»†Ç´ "­é7Û^›±%~§¹pA6Íbä^+ÍQ7Ñÿõ¯Å%H½Å\%¥³ ’‹()æHDÆöíÄ™9maašîã–b.ºS‹gn ¾ÞõˆˆÈÍ ¿ÿ¬‘YSs~`.ÀÅȦLäÒHXŒQöå|“OMÉù‡ó ç.~¦¨¤fàcdÇz¨f.t6d"—¦¸êɛȜ)+•Ye ”K°Ádœ#”Oûè^é'87°–§)™ú1‘ÃRJMäRx˛Ȍ&2ÞmãœN›‰â¦_lã!€M=øïrh"Md©hšÃD&B¸4Q~ÊÃÄ"ˆÛLXa*³â“mÃ??°\åL&2O.n!áx® 6DiNdà˜ŽIW EäæƒHL wkçIáOÏŽ·Ò“º ™ÈoÎÙM±#UÓÖüáé91Z9qGÿá_/ί{ ŸDå¦ÛþI…€9öÜÄ á¾×–ļ¾¿î3+¦¬øjÇ ñÖü<ä‡ýa÷iáç=gÄüÄÔƒ™Š18`ÚæñZzë>¿Kõ`´ã×傉ÌLÏõ˜OƱ¤lòR“«˜!¥i@¾*¦ÉÈ3gý˜&‚(é—¦nŠcK”tú\ø ùάÝy,´Ë>Ó»û_4܉6þQö½ê;~}X´éP3Ìï†À˜gìRÞoÀDþA÷Z™ïȯŸ¼ÝÍE>+È—è2teŒ¢Ìe.špá¨oLdÜGînŽeìȃ Êÿ}Y߸3€q º4×6¹«[dÿÍ¡ŒÍÅ…|¾ç«¡9ŒÜk¥9êf!‹‰LTX‚ Ý)'2·³HeÑʉ¹ï3âVWà¶âd"Ó&î¶b±Ê<Š9QMéÉñ̃JoI‘X3sžÀ.]?‰L*…|0éR: ^ÌZR)ñ~ʽŸ`½Î0œwHQš:’H[ÎW˜¶<ÔŽÜÊ…±&Ç\M`–Ò6ê)¥!™¼É”•Êä ^ kÚI”1e¥PçT”ƒ©ÜÐ]Êìƒ'Qj¢ÓÆc–±ám)¤ŸàBoòFsÞ.,f éÆ4wa—>ÈðGð2Š’¾R4‘‹‹&²T4Ía"óãÌ'] ÄfAÄ"‰m® ¦¨bNRËlkÈDN?„$¶çNÀ‰Û4Y`q5‘tšÈ"•ÑŒD‚bêrë?<5v]Œ,¢<¿Öqb݃½€Ð/¯{IF~ät‹>,Þt(æÿ%ÛsïÏ}çÄ21̾ÕyrŒÅ`K&\£Œ<¶D'££ “•H[¼öá֣ÆÝM`~ÇvÒK7™¹D&S’è^Œ6öÃDÆdMÆ,ЂŸõšZ^0V{^Ûœ¸íÅùáÎ~µQÊiLb¢t‰@ÅÔ&%F"[÷Ô¬ð—gçÆ~ïÎöE;V‡ïu›þñü¼Xi+HaAZò÷òÐ5 Bƒzò…è^öåÊ+…híÜ7&F/“št"ö««rÈLJÆ´ÓÛË£™šxô£wÅÖÃ1Ú›h^¢˜Sîá{²Ïƒ(ràsçâùŠD2M;sMmd/$º¨½[cûSmÇÅèt |6ô“|À<Ø‘Ïœ4 L䆬Çwï׈…;¢ýÙìsÇlOí‹qÏ÷àþ¬ÏäN$ùéqë¢áÊ÷ ¥ì c¾ôãy_&2\ò&riÄ0ãDô>H$Ÿ1)-ˆr¾‡˜È\ÌHÇκ0v#‘ÛŽeˆÇ,ç `€/Þ\» cImB$z‘×¾@j#÷ZiŽº‰üºí¶Ûââ;-T1”¹Í—*wp1‡baÎB’:ÉÉHW2‘‰ˆJá™C1gJ#bÞÅ{D'©EÊŒô°>ù÷€ó†%FqCÏâÄz<‡RÎdÖÖ€±Ë9¥4‡qŸ#— •–”›Ü•ÌÎ#FÔíÃNÖúÜSJC&2äËL]Î¥æt‚ÄDJ—ž³IGA42‘Àyðð+8&R#™>c`§;‚€s+cÐPúÆØéyPHHÿ8?“¢ô)åÇõøšÈÅEY*šæ2‘S~?Œb>¼üðÌIŠIr÷ñ£ÏÃïxÍ‚ˆn®¦ôÜÃm0ìG™œÐø‘çöŽáÇžà”΢ÔDæjâõþ‹Hy!÷*ÆÔ¸¥µy|1øˆDýA·iÑh$=ÀW¾ÔDþýÓ³Cû3ëC­GÇWÌ@‡yˆÑ†q¤P mÆÆ&¿!ä˽¥ÅˆhšÂ¶'ÂÃY¹p˜Èä¾%·,¦/é1þõÒ‚¹ûÿ³÷îU\çÚ÷ò]ïû½ß)o’“ÞcÇINr┓vÒ““œØqì¸70`z1`ªé½ØÛ€mzBôÞ»@4QÕ@¢<ßü±ÄhKl!g†ñý»¸/´§­5eïµÖ=ÏËðÈ+‘Iäw¤/³eáÅz ø‰tâq`Ìi…epJtŸùÿnê+„âÞ‚þ+ÁWd1®ÇœäÚŽÞˆTIš´3Œ£¹™ÉgÆÛ,Cd/F2Ę’¬0Ny1æ/ùùi³Â<ŒÓ|í cwÚ7‚¿ˆ¦ý¢­ é/˜¬ÜHÍ5‘s¡n¡‘ "·žçì7ûÄK™ÏYÚÛÅœ í;û…ñŽÁ±`û!ú˜¨dÚbŽÇ C:3À 'Uó(ßÚúP>¦=矃º²m–///÷㈱¯‘M}'ÈDN/2‘E¦¹[™Š;küH2àÁøå™A2À3˜é4R,G2ÓÂc—<šBÃÂ5o’‡á.&ÛdÀÅÁ`³ ÓBÂ{P~¼ÙþÝ<"„HÌX¢Z1·¾÷Ê2ÏWŒ‘IDã¨%Ül$b8n"c|õ|¯qê^ÒÆ ú0üˆÜäqýø ¾0…1y±ZˆX&Ï-å„hbr÷~¦í×7U‚‰|êBo£;€þØØúhs^FÉw2x)#)A0»‰îÇüfúºý·ÎÛ¢m'ý8…ïf>‘Ùœ?"¦1Ç¿]‡|?8.¤»à¦CKi #·¥´VÙô(Ò—aIt[žžb€Î€ƒ˜>ý "¾èÑb`J ƒ€A67è¸UÅ#ǬG¿‰írƒ>÷d!„Ù†ö€@.Ú R1¤…hgBêÚ LL]Æàü¹ÊøS4<ííã|‚ÀBJ¥ød"§™È"ÓÈDBˆä!rùññ<‡³h“—öT ";$iä&Y¶B!D>d"§™È"ÓÈDBˆäá‰@^ô&Z©î6}‚HI¹I–-„B‘™ÈéE&²È42‘…B‘F’4r“,[!„"2‘Ó‹Ld‘id" !„"$iä&Y¶B!D>d"§™È"ÓÈDB!DIÒÈM²l!„Bˆ|ÈDN/2‘E¦‘‰,„Bˆ4’¤‘›dÙB!„ù‰œ^d"‹L#Y!„i$I#7ɲ…B!ò!9½ÈD™F&²B!ÒH’Fn’e !„BäC&rz‘‰,2Ld!„B¤‘$Ü$ËB!„ȇLäô"Yd™ÈB!„H#I¹I–-„B‘™ÈéE&²È42‘…B‘F’4r“,[!„"2‘Ó‹Ld‘id" !„"$iä&Y¶B!D>d"§™È"Ó`"ótãÆ I’$I’¤ÔèêÕ«‰¹ÁDæÿÜzI’$I’$%©šš™È)E&²È4üð9rÄJKK%I’$I’R#LÜ’’’ÄLdʦ¹õ’$I’$IJRx8eee¹Ý‘d"‹LC$ruuµGûH’$I’$¥E¤ÛJ:¹¶¶¶I½$I’$I’’TUU•"‘SŠLd‘i”Y!„i$ɼÄI–-„B‘åDN/2‘E¦‘‰,„Bˆ4’¤‘›dÙB!„ù‰œ^d"‹L#Y!„i$I#7ɲ…B!ò!9½ÈD™F&²B!ÒH’Fn’e !„BäC&rz‘‰,2Ld!„B¤‘$Ü$ËB!„ȇLäô"Yd™ÈB!„H#I¹I–-„B‘™ÈéE&²È42‘…B‘F’4r“,[!„"2‘Ó‹Ld‘id" !„"$iä&Y¶B!D>d"§™È"ÓÈDB!DIÒÈM²l!„Bˆ|ÈDN/2‘E¦‘‰,„Bˆ4’¤‘›dÙB!„ù‰œ^d"‹L#Y!„i$I#7ɲ…B!ò!9½ÈD™F&²B!ÒH’Fn’e !„BäC&rz‘‰,2Ld!„B¤‘$Ü$ËB!DSjkk½}2‘ÓŒLd‘id" !îjj¯Ùñ³ÕvíúÜYM¸\wÍ.T×6Z¶öêu_¿îê›BUW®ZEM]îä åWÔ\m4íÔ…Ëv¡ª¶Ñ´|°¿ñ:\½v#Òõ&ºÛG–©‹¦å+–½ê—ŽÕ¥ê¦û¥îºïå4GÕåÛ·ªÛ‡¡::/¨9*ó”ͱßqô‚Ÿ×æ8}ñ²¿ƒsÃqäøp<òq…ãØL8þ¬ýFóÇ€z²?·ƒu9ÖÍQŽ?¸nàß©s5®!Î ×È•f®ê?‡èÆmö¡5HÒÈM²ì|P¯Ó§OÛ™3gRW7!„÷&çγ .äN¶ªª*+++³êêêÜYÞ•——[MMM‚z\ºt)w²-]ºÔúôéc'NœÈuGœ={ÖvîÜiÛ¶ms/äNû1¶nÝ:7p7oÞlGŽÉ]ä‚Läô"Yd™ÈBˆ{…{NÛƒ¯ÚÙŠü¿YSV–ؼ²Ì¾Ýc©ýfð*[w Ü§ï9~ѾM/>Y‘³Æ-ÎV^±ÎÓvØôZfßê¾Ô¿ÁæYþvt™¾Ýz¼·³á3ÆÜC#ÖÚ°…ûcKÝ̽¿Œ\g½ÞßåŸ1uÿ:zý¼‘ýrà ׯ­´ŸEŸ‡ßÜæ–Ãç|Ú¯"ñÿÏYçw¶7˜µkö—Û¶ïô(ˆö¯À×»\ÛÔ(¼-ÿÂ[¬Í”- ÓØÆ›ÑqýQŸB?®?ë¿ÂælºÕq-ÎÉËoo³ïÝmŒŽã¿÷,°ô.´I˹™Ÿ ë3“6yÝàwŽãw¢ãø§áklstÞ'¢}i?µþ8~7Úþ‹Ñy(;kpvðT…_“¬Ëþvxk›àÂ]§ì7ëöèŒÉE‡okøÃúƒgíÁhNÝÜÇê·¯­òm?}WÆ,-ök/0xÞ^¿æ8‡èVDÇãTÃüÖ&I#7ɲ㔔”Ø‚ lÇŽþ™{×®]­W¯^Ͳ…Bˆ;vnôèÑ6~üøFÓO:e¯¾úªõïßßo\æ‚ÁÜ»wo[¹reî¬;bÆ ¶gÏž†ÏÔcܸq±%ÌJKKíå—_¶U«V5š~'°Ÿ‹-²Ž;Z§NüÿvíÚÙäÉ“}_> Îo¾ù¦›ÙÓ§O·:ØáÇsû‡ 9½ÈD™F&²â^¡`ç)ûFçÅv®òö£3Öµ/¿¼ÐF,>`kö•ÛC#×Ù7º,öÈÍÝÇ.Ú×;/²'š7^DÀ”Ã$]°µÌfo:nÿÙo¹†µW›š­¹‹aýêì=öÉ6ó¬Ó´í ó0K1ÝÌÝ[ãö ]°Ïþ÷Ó³¬óÍm°þÌÇì‡Ý̺ªÄ†DËü[TÎÄå‡|þÇÔ¤Þs6·YÑò+÷žñï±sÕöÍ®Kì…76ÛªhÚû£uçÛEM;¾o¯>bÿôìl7‚‹wœ´OEe œ·ÏVï;ã†ñ;,°íGê£VÚ¾¹ÅËEÛNx=Ôw¹ýψ5 æê¸eÝÔŒ×mU´ÛEâÆá¦Æ-eþëósÜØ °>p({Þ–2ûÏW){mCÙÝgì°ß¼¶Ê·s»(ôGÇnð›ûË~wˆê²6:×ðk;ϯ»\¦DçŠãøÜäM ÓÏßg_î¸ÐÞZuÄÏÍ#cÖ¹Y\~é²GÿeäZûAtœ–DÇ{áÖ~Ó㱨nœ¢—1áÚ¿ÈÍâéQ™˜ã˜Î°­ä¼×­ã;Û½nCî·O·oï­/m(?–ãûq:*ÿÈ™*ÿ~azs Œ®U®¯ikë÷úù‰èšä8‡óx´¼itRk‘¤‘›dÙq/^l>ú¨r¹}ûöÖ¹sg†B!îÚ¹¡C‡ÚÈ‘#¦ƒòKõí æ|··"³s!òzñö“Ñ9/¶ ÑqÝQzë‘ÊùÛÊìÁÞËÜPýtÛy6®à–‰LÚ‡¯uZd#h˜63ºv>ß~GG}ýçk<:w鎓 i<ö–]r3žs¸ÿD}to—ØùâFÅìMÇlbá!Ûtèœß`€þ³÷؃} ¥i÷ÖVûå • êG/øqüktŸžXo"S׺-µ±ë“(bêOÄð¾—ì xÄ} h÷iûR‡…žŠƒhá/FsM¸ðã¾Ë½~ý¢ºý°oa£Èá¶S·z4~¸FH‘±xû ½´Øj¾_ѤòàÆÄ¿÷(h”.„ã"©O^¸ìÑÜÅ'no\:UéÛ¹ä€Gï$}Ks$iä¶VÙô±xd–Çp1„C?x|˜Ïûöíkxœ–>QxL–Gd'NœhÏ=÷œ²·lÙâu"«[·nÞ×Z»v­o—mÄ©««óþ}+úf‡ÕßlÎ ÐÓÃ`€.„âãíÜ°aÃlÔ¨QþÓ¹_¿~"<èmé&Ö¬YãË{;˜šø´IEEE ãzüÚº `ÓÑ>ÍóÌ3öôÓOû ˆ†BññÃŽÈL¢VI½@dåÓ7z$pÑîSöÅö |é0¸0Õz¼»ÓS`.c³ Ì1 B ÙÀÎÒ‹nvþ°Ïr$ÿ¿‡­±¯v\ä‘¢ÕW®¹y†Q™ŒÃgŒÅ_\i57µ7W•ØW¢u~Mû¯¨_ˆêCd,ËcŒQ÷ø¸ Öþ­úQ ~ŠÊð£þDÊ~¯W£ÛÜ2¢ÏU^±ŸXas7—YÏh¿ÚÝŒ2Íô[Ý–Ú™›&íµ7ìÃV[Û)[hŒf¾ ¤Ï â™ô'¬³;ú„ëïÑ5‰±}2Ï1â&È—:,ð—ë#e¦®:bÿòüû}ô½z):¦˜ß÷EÇ2œ7R¥<=icCD69³ÙŸ£å.Ï`s´†‘ÛRZ£lú6a€K¤J †”20fN_ˆÁ{H2}òÉ'ý¢¥èWÑ7Âæñb¾˜Élƒ<žË œ>oÎíœ9s|€N¿‹>f5ƒ`új”K4õb –ß}÷ÝDÉB‘´s´#´˜ÂÜÜ|ýõ×›¼pŽi´5¤ž ­m7"‰¼Å(¦­âÉÀàd›ä f;ø ˆuhóhÙFmm­ 2ÄÞ~ûmÿ0£I¤Œ7^6¶‹ö™(iÚ5Ú>L`ÌZ¶ƒ‰Û\Žcʧ½£ýøfY>Þþ± –ãÚhnÀ†<ÌÂ!b0¡igçÍ›çuÇxfŸÙGæÑöR?žâØq¬‚ï¼4û2ËLäô"YdšÖ2‘ù1/})S¦| ‰Ì‡F#4Pà¸r. vœô(î¿Ûàù|‰oÿvýË ?Œ‰¼fÿ7]I¥ØUzÁ÷ #0⃉Œþåèx¹çïã7xŠŽõúY¿"ß¿?Xë‘ÆD'“BH9‹Û¼¹ÅÞZ]âF3QÖDÏs¾~1pEÃqäº"7òÄ¢CV][o®?XîçCž¼Ô/ž™ÜÒDOó‚FŒ[Ž×>uÇ€&‡s`ÉŽž;ÜÁØ%☺Q¯·¢ë–k–é?ˆ®¯x>f`=¢ã¹6ˆ¤Æô3bÑd¾ÝË,ûÍÝë×7l0¤CR„ÔG%oòë/žòâNi #·¥´FÙ Bè†-ÑÃôƒB$2‘O͙ȘÇq™§µâ&2ƒ_ d¶Cði¶_YYéƒbÊ$ê8À —íÒ×¢FtfÀÝædBqoC;‡9J$-7"·cäÒŽ`ÖHgOçLä}L›„L„-þfnüex©´=D(‹60žR)n"c S©%˜Î8l#˜ÈÜ”çô¿‰L9<?`ÄR6mhØW"‰I¿AâmäíLdnβ.Ûà]¹„€9Új¼Úo¢¡ÙÚì–"9½ÈD™¦µLd‰?ŽÁ£! zø˜Gn$~\e" !šƒˆcLä÷7ÜzAQ¬Ddb¶’VàsíæÛæØËìBê 7ÌH¢Hçn¾Õ¡"ª’Gî131‘1„ã/&£[‡Y‹ ½åð9OM7½0m1ÙÈË‹Ù0F©g€œÀ¤yà~qš3‘)ç©›/Y `L z4Ú&‘©OŽß葵DQŠ9Ùa궆—‘ï™:ðBÀ8èÔãâÍ(f ò—ý]½÷Œ½òþ.7jãpxiéúÏÞí³Ee³f ‘ÊìGÙù7"9Æq0gÉ}Ëy!¯nlŽ®ULäç[iÓFnKi²‰¬âý<ÆÊ`ñ¨0}!Ld»üÍ`5D@1H§/Ldæ3Ó!D"7g"cXñÄ —ùa`'Nt¹°°ÐÑ Âé—ÉDBˆ7!'rÈC DÿbÆ’Ã7DÌÒ…—ïA0‘I÷ÀXŸô¦lôM<Ã|>óB:ŒÓx»ÁDŽGüb®âàQÐfb‡tÜœåF,fí˜ÈÀÐܹi:XŸT¼è¶ó8¾ßLd¶À&W3í4ûK¤1G€6šÈc¢“wîÜÙð¢Aö…cF;ÍSJ-E&rz‘‰,2Mk˜È4ü€ÆMdÞHÊ ‡eró9Ò8üè™ÁL¸ËäWÂX–‰,ÄÇ“kQ'“ÔzDæblõµÛ#M‰l%ïñ'^œk¿~m•§m r’ùDn+©7–1ëxÄž<­ÇÏU»1H[ ʽe—üE}DØn:tÖ£/É+L~_^ÄÇãøßëµÌž¿ÁMI¶ùÊ:;^´sµöÝK݈;|ª2ªS•¿œÈ`^̳c:r"2*krÑaOw€1Nß6olñ<ņ!7'2QÉDžæFŒ’2€ô¤+ ÞD[cð²¯DNÏÙ|Ü÷#’t[Ÿ÷´D>ç¦â¢YÉ« Zƒ}å¾3þB7¢N©)$0ÉÑ˱ç8cxv–ÇÜÝó’3õQÉ]¦íh¨ÛC£ÖùqÃDí7{›åD3sÎÇûyÆ  ùu¿‰3‘OEÇ‘H]^®Ç9aû”Ã͈qÞ!:Ž¤ßΦ'Âd/-¯?Çÿç™Ùž¸f¸ÖŸõó„NÞb^Lˆ©ûɨnä-&j™ŸŽÃ´c90`î^r´ŸºÍ#ãIûéŽÉË5À1%g3×>ç sxY…§ /òÛÕ—ä áÚ>íߣãÖûÍL{Ì|ê6±ð°Ï''1ç(Dd“Jƒý`Ø/¾O¤Ãàûq®âŠÍˆ–cþ»ëK=ÈƃçÈ[´?Ú×ù¶˜èñ¨¾¼\ñ íz´:åc6ó}$ê™›¤Óøl´<wKi #·¥´FÙܧ¯Ã ›” ’é…` þ0sà vx´–ÏD'™÷D0¸%’™í1Èåå<ô¯â&2ƒW%æ\…—± ƒ[RYP&둃eÊT$²BÚ9"Œ¯Ç!r—¶$…a"Ç—Ád%¯0éHÛ€)J~`<LhòcÚbžâ`À–––ºh_Y'˜Èl7¢ÁD&•j܈¥½c¶IhCiÏ‚·ÀDæ&ks&2õd]<<"Ž¹©J41ûÁËóð-Ø&ËRÊÄ°Æ/aßð=hÏ1¹QÌ>‰ŒùÌ1äæ,Ûdÿ9n´µÔ“cǾ°-Ž[x‰mµ™ÈéE&²È4wk"Ó0Há.aÌ?œ¼mœÁ"døáeðB£ÄtîèÈ‘„ÑÌ·âã ¦æåý]x[›\Áã%xnOMØèÿ#æÇ#CyYùf‰€dýô.lˆ<Æ4$¢–èÚæa>=‰<Ü5ûêsÚ2í¿×2›‹î$ýÀOû¹ø@÷%žKwûÍ‹Åyá-R`"?Õ Óù—ƒVxÙ¼ÈãÜ4nÑú¹i+Ïßç9š‰ZÍc——ÞQoòàòbÁðÒ³ºk×mð¼½þR9Že?ù’É[Ìö0ï1éɃ̵® äy_Q]ç¹¢‰òæ`ß¹@6 _tÀ£ÏùNÍ (Œ¾ZÛìŠCk¹-¥5ʦD”1ƒHúC<ŽË€šÏ½À “¾ý±D‚1Ÿœ”À@s˜e0’ 3e Ê€„òTtêM1HÓ×b@ÌcÊáÝ ¸1”©›Ld!„øxr"ÇŸ0>§M"’–—ÜÅ—Á ¥]#ª3˜ôKä;Æ0&˜(dn|†4LóX.D>³^x -Æ1¦1OÏ„y<}Ãz´}Üp¥íÂg8|ø°·‹” =¤-§Ù PÇ‚‚oC=i i—IGAÙ´•L£,ê‰øÌûØ&u ©:0†‰¶&]`–S/êÖ é1};HßAÎ\Œ¿æŒ;R~ÑJ´p<÷s.¤Ï`rá¸qÜs#¡uò7¯[Y…§îˆG‡yœoÎ+õ¢îÍ_Žas:ÓHWÂ:£8Ô%¼¬0Àö‰6'‚›|Ó˜á¹uÆÜç<u› éˆnn^¶›{Ù?¢³9VÍå×Ç‘¨îÜãìËh>û Ü8§CÝ*š¯‘ÂÌçêÂw"~-…ìQÚ\ŸÍDÈs ÕßÜöɯíÛ¾ù½:q¡Æ#¶›nåÃÑFnKi­²”28e ÞNïë0¨epÉ`˜eCÿ(ÞâoÒXpƒÓ—> ä²}1ÖÃ2ôÃØ~|{aù»ÉÅ(„";Ð~ä¶!@Æì(¾ mFrxÑ°,m'm^HAÁ|ÚÄü .‚í„ùüÊ ÐöqS6¬ñ?7Xƒ!^h#óµÝ”K»Lš Œß°>ë0/·ž(,C[Êþ…òÙÇÜö•y,“›ƒº}MÙÍEJß)2‘Ó‹Ld‘iîÖDB!„ø(h-#·%$Y¶B!D>d"§™È"ÓÈDB!DIÒÈM²l!„Bˆ|ÈDN/2‘E¦‘‰,„Bˆ4’¤‘›dÙB!„ù‰œ^d"‹L#Y!„i$I#7ɲ…B!ò!9½ÈD™F&²B!ÒH’Fn’e !„BäC&rz‘‰,2Ld!„B¤‘$Ü$ËB!„ȇLäô"Yd™ÈB!„H#I¹I–-„B‘™ÈéE&²È42‘…B‘F’4r“,[!„"2‘Ó‹Ld‘id" !„"$iä&Y¶B!D>d"§™È"ÓÈDB!DIÒÈM²l!„Bˆ|ÈDN/2‘E¦‘‰,„Bˆ4’¤‘›dÙB!„ù‰œ^d"‹L#Y!„i$I#7ɲ…B!ò!9½ÈD™F&²B!ÒH’Fn’e !„BäC&rz‘‰,2 &2?@ ’$I’$I’Ò¢«W¯&fä™ÿsë%I’$I’”¤jjjd"§™È"ÓðÃsäÈ(I’$I’$¥I%%%>XúGƒyLÙ¹õ‘$I’$IJZx8eee¹Ý‘d"‹LC$ruuµ–$I’$I’Ò¢ÚÚZ(%e"Sv]]]“zI’$I’$%©ªª*E"§™È"Ó('²B!Òƒ¤¤Mä$ÊB!„ȇr"§™È"ÓÈDB!DIÒÈM²l!„Bˆ|ÈDN/2‘E¦‘‰,„Bˆ4’¤‘›dÙB!„ù‰œ^d"‹L#Y!„i$I#7ɲ…B!ò!9½ÈD™F&²B!ÒH’Fn’e !„BäC&rz‘‰,2Ld!„B¤‘$Ü$ËB!„ȇLäô"Yd™ÈB!„H#I¹I–-„B‘™ÈéE&²È42‘…B‘F’4r“,[!„"2‘Ó‹Ld‘id" !„"$iä&Y¶B!D>d"§™È"ÓÈDB!DIÒÈM²l!„Bˆ|ÈDN/2‘E¦‘‰,„Bˆ4’¤‘›dÙB!„ù‰œ^d"‹L#Y!„i$I#7ɲ…B!ò!9½ÈD™F&²B!ÒH’Fn’e !„BäC&rz‘‰,2Ld!„B¤‘$Ü$Ëþ¨©ªª²²²2»xñbî,!„"u”_ºb§/^Î|Çœ¼Pc+öœ¶ë7rçÜ{ÈDN/2‘E¦‘‰,„Bˆ4’¤‘›dÙ5K—.µgžyÆ&MšäŸOž”Y^^îõ 6555±-!„"­$iä¶FÙôAÖ­[g%%%vâÄ ïÏð7ý»råJïÃlÛ¶Íjkk}:ƒBúUÛ·o÷”………VPPà}ž8Ôk×®]>íÙ³§Q]Ù}£ÐÚ¼y³o(;Ô‹~àÔ©SíÅ_´W^yŧ‡4üÏg¶±jÕ*_/pæÌïW9rÄëItóŠ+|y>JKKmÍš5vàÀ†iB!îM.T×Ù¬MÇmø¢6qù!ÛWV?¾ç!–5ûËmïñhüð¬)8hÓÖµã罞©²7W–ØÐEûmÚÚ£v*'EÅΣlRÑ!­¿éй†é/¿µÍ·Áö»h£ Šm|á!;x²"¶¦Ù¹ÊZ{Ci´n±­;Pn×rrV¼»¾Ô~=h¥]©«o+Du™«ËéKõu¹TSg…»NYEôÿÚh;s·”Yiy•ï_å嫶`ë ÛP|Ö—=tªÒÞ^}Äë»fÿ™†2·¹`+÷žñãåWlîæãþ?„㵫´åi¥d"§™È"ÓÜ­‰Ìà¡M›6Ö¹sgëÝ»·=òÈ#6aÂ;uê”æqÉçŸÞž~úiëß¿ÃdõêÕÖ®];{öÙg­C‡¾ Ûa@3fÌði,óÒK/ùvÜ`$/X°ÀÚ¶mëë³ÿ6¬aà3pà@{î¹ç¬oß¾>0b}Dä C!„B¤›Ö0r[Jk”yJÿ‡þP—.]ìá‡ö›ß˜¿:uò¾ }þ?~¼ß˜Çx¦ßC‡uBŠi¹€AL†yaý^xÁÞ|óM«««óºOž<ÙçµoßÞ·Å6ÆŽëó1|üq›2eŠÍ™3Ç×gû¡Ü}ûöùÀ´OŸ>Þwb۬߱cG7¦º0eèþýï÷åØå©1ÆŒãó¸™/„âÞå\Å7a¿Õm©ýmÌzûÏW—Û}]Ûîãõãï? _c÷w]bÿÑ«ÀµÖ¾ÿÊ2ûNÛVrÞço9|.ú¼Ô~Ôg¹=6vƒ/ûÇa«Ý¬…IËÙW;.²ß¾¶Ê~9p¥}®Ý|µ¤þd×iÛí¾Î‹£u í/#×ÚÝ–Ø·»/µâ›Fò¡Ó•öÓ~E‘–ÛÃ#×Ù×:-²ïkõmøµk7¼~ã—ôÏÔ¬ÿã¾7ëíÇŸ†­±šÚk¶ëØE߯§&l´Ï·_`ÏNÚdÓ×õýùëèuöÙ—æÛ”‡mÉö“ö¨N¿Pdÿ3b­}¥ãB{ùímV{õºMŽæ³Mr0Ã[«Jì=5ÓÞ^}Ô?c&+Ú‡·×ñÏ-A&rz‘‰,2ÍÝšÈ °08tëÖÍÞ{ï= 5j”$4…B” ƒ 5 ~† â¢VмÿþûöÔSOÙĉ­¢¢Â·Å6‰Ä¹pá‚OïÑ£‡Gâ9ƒáÌ`Ó™f dÖÇ|†×^{ÍO nvîÜióæÍóu^~ùe  !„"Ý´†‘ÛRZ£lú(ÜȦ?ÃôÙ³gûÓQ={öô>ý¥Ã‡ÛСC=QÇçÏŸ÷¾ ëÐ÷aþܹsý3¦2}$ \Ì]Ìi iú9½zõòi<¹E¿ C˜þÙþýûíìÙ³6räH_žˆhÊ¡OFNäÊÊJ†u$Š˜§¶èWÑ{ã7¼ßÅÍûÐ×cø¾‘ƒ:²]ÊÄT¦Üا¯Æ4ÊBqï2oK™}¯W-¯ê"=Ä7».±©«êÐ?Xk_ï¼Ø£k‰´åex¿²Ê§¡ÛmÆ7[¯ÔÕt­ÞwÆMÚ}Ç/Yiyµ}éå…6tÁ~7~ë®^·~³÷Ø÷zØÅê:_—eY‡mŒQ ö‡ƒ¶ fØõÔ@F!„H?­a䶔Ö(;­D“þˆäÅxÅôż¥3zôhïóÐwá‰-úO˜¯§OŸöuHñE¿ˆõH}Aÿ‡ç¬ o„1Ltq¸O¿ˆr¦OŸî7íéÿL$2ËbR©.è£qƒHO†ÉM½é³Ç#<åEÊ0 pêƒanγmnâ³/l“e英zaŸBÜÛ`þž¾xÙM_R3ô|o§}æ¥ùžžþ0lµ¿¸.μ-Ç혥g«Ýl%2cvƺ£öÔÄMö…ö lïñ‹6{Óqûf·%¾ýÀ•«×Ý°Æ€îøö6ûýú§q€¬¿ºÚF.>`ªjÝÐm7u«½¹ªÄST¼±ò°Ým¯×û»|ùW¢ÿ_¸iàBEÍU;« &4uaßœ¬°~nŽÍÚx¬aùwÖµ/w\ØBƒT^÷›é<D;SmÇ£wTîåÚköëÁ+í•™»ígýWx:s÷ØãÖ5Z÷N‘‰œ^d"‹LÓ&2 a AÞ;¢~IEq‹Ä©Âà‚ÁƬY³Ühf¢…˜0€ @2ˆš6mšGë„Ç#‘3sæLüõ¯Q7D+“{™:Dä*rýÉDB!î ZÃÈm)­Q6&2}›Aƒ5l‡üÅL£DŸ…>ý úH¤£‚‰Üµk׆÷80 F3QÈô•ø›ñRgÐï1b„¦î<ÆöCÚ/úbÓ!9˜È˜½,C ¸iOŸŠuyZ,@¹˜×,OÙ”Ç:qƒ˜ý£nì Û§ÞQ!„âÞæÐÉJûÓ°Õž6â™I›¬ã;Û=%¦-`ê’+9κârûòË íÀ‰ +ÚsÚ~ԷУ“ÛLÙbíßÚæé+È£<¹è§ª¨º|µÑúѲ½Õ–`,¹ì\G%“ƒtDcÐ>1a£½½êˆUT×yÚ Ê,ÙqÒ~Õå÷CV××eê6ûZTLäýQ]IY±ýH½·oEÛùaT¿zcîæ2Â>~®>²9@¤2F2ŒYZl¿¼Ò6:g?ï_äf9)@ÖŸµ?_ko¬8ÜhÝ;E&rz‘‰,2Mk™È 4ˆö¢mˆ^a`Ìã2 ý'n²‡—ѱ7ÕéÃìݻןîâ‰,ÒcHçÒz‘¾‚~Txa7àyGë`B3/n"ÓWÂDÆÈƦ<"¡é«Ñg‹×ºÓwc`ZÇeÓ$·2û@@î‹•…BÜ{¼ôæV7zIc¤pÀœ &2ù¿1¶†ÙëE‡í[Ý—Ú¡“ö‹+¬ÃÛÛÒYϘ”¼œTÊ%§+Ö%‡2f5i1H?Ñœ‰ù¡É“üøøzO‚\ÒDLÿ$ÚVÈçÜRd"§™È"Ó´†‰ ÁDæG›—»0˜aàÕdðBT ƒ%&¬Ç ‰&DÉ΂He f>óFpò"“¦"DàP: x(‡¼Ë|f~QQ‘×Ó™õã&28 ~d" !„é§5ŒÜ–Òec"Ó!ME0‘‰.æé«o˜È]ú8O<ñ„›»<Õ…‹aÌÍoROϘå1‚1…IážøâÉ.n°‡—â7é‰N&ò™: >¼Ñ1¢ˆ9MùÃ[!ĽÏÓ“6y$qÉ™*O3ñêì=öÏÏϱ! öy.â‡F­³Oµ™ç/È#B—üżˆãÊËW=’#úÔ…Ën°b ³ü²§Ü˜æ¥{¤•8x²ÒÓFð‚½ÿ¸Âó#cÄ>:¦±‰ü»!«=‡2ô¹ÛË"Í‘ÉSV”ا^œk…»NÙßÇmð¨à9–ì½ÌÚOÝêé3ÈጹK]Vì9í)*¾Ø¡±‰LÞg ó`"WEÿÿaèjûYÿ"O‰AncRfð2À Ïú2˜Êÿ=lýÓs³mâòúÀ¹qÅž[™bÂß 2‘Ó‹Ld‘iîÖD&J˜<Ç NxJ€¨G HD0¸ QÆðî»ïº‘Ì ˆˆ'DÞ0Ðbp€Ó—ù ‚XÃ8<ÚÉ›Ãt1?ääã¥/D!Û£Ìð¨'/ayÌir6 !„"Ý´†‘ÛRZ£lndÓáFz0‘Aæ+ý›Ú‚]ô¸‘Nˆ¾O\…~‘ËÄÕ«W7D'#úLñ´Ütç)1Ö/¦ïÅþpÃiáÝÜh'¢˜zPÑÎUUU~ã?¼<™ý >ÌÀ`L#ÅEî1¢ïE(Ã[!ĽæèúºYûÝ^Öþ­­ž¾\Æ;Ž^°¿Ž^oÿ=lµýjà >æÅqϧ/Õç9ž²â°Ýße‰›Ålã÷¿®´ß¼¶Êjë®ÛÚånÊ>­KäÆà…n3vÚ3“nE9c"cZYRoc “7™äGæÿqËŠmϱ‹öã¾…vôLã´DHßçu)ðº¼:k·ýbÀ ûÓˆµ¶ñÐÙ¨þK<Š90}íQûÕ • &2Rn¡^ñuè/± é(ây 8räˆO‹æöþ$4)5BÊŽæê-„âÞäRMç7>q¾þ·¿¦öš§ ¸\wÍ#sG->àÓHQÁ2¹ö‘¿¬®²¾m ùØYÆÿõóÙ‘ΈHÞæ-ÑÌqÈORTÎ\ºìe—ßL¹Áö·—ÜÊmçTN]X‡j¢ªyYFu€Õ¹ûÃ2eÑ~ŠŽAsùœ}[ÕµvýæŠü>ÚöÕk-ï[d"§™È"ÓÜ­‰,„BñQÐFnKIªln°ÅË“\ñ'¼î0³©;i00‘Iy!„"û`¦ò¹×æíÍ%>"d"§™È"ÓÈDB!DIÊÈ…¤Ê&º—<ÇsçÎm6‚9ÍcÎœ9^^Žœû™BˆlB~ßžïíj’ÎA|tÈDN/2‘E¦‘‰,„Bˆ4’”‘ I–-„BÜkªA·ÿqÈDN/2‘E¦‘‰,„Bˆ4’¤‘›dÙB!„ù‰œ^d"‹L#Y!„i$I#7ɲ…B!ò!9½ÈD™F&²B!ÒH’Fn’e !„BäC&rz‘‰,2Ld!„B¤‘$Ü$ËB!„ȇLäô"Yd™ÈB!„H#I¹I–-„B‘™ÈéE&²È42‘…B‘F’4r“,[!„"2‘Ó‹Ld‘id" !„"$iä&Y¶B!D>d"§™È"ÓÈDB!DIÒÈM²l!„Bˆ|ÈDN/2‘E¦‘‰,„Bˆ4’¤‘›dÙB!„ù‰œ^d"‹Lsüø1«­•‰,„Bˆtqýú5+-MÆÈÅD¦ì7þñe !„BäãÊ•ËVV&9ÈD™æ`I©¯¨¶ê+W%I’$I’R£ŠšZÛ°$ùêµk^veM]“zI’$I’$%©s—ªìБc¹Ý‘d"‹LóÈèµv×Åö’$I’$I©Ñ·º/±÷-´ªËWs»/9—jêì‡}–Ù·£:äÖK’$I’$)IÝ×e±ý}üúÜî‹H2‘E¦ùÝÕöÉ6sí³íæK’$I’$¥FŸ~iž}³ÛR«LÀD¾X]gßè²Ä>Õ!·^’$I’$IIê/ε?X›Û})@&²È4¿ºÚþ­í<û|û’$I’$I©ƒ¤ouOÎD¾¿ëû\T‡ÜzI’$I’$%©Oµ™g)9ÈD™F&²$I’$Ii”LdI’$I’¤¦’‰œ^d"‹L#Y’$I’¤4J&²$I’$IRSÉDN/2‘E¦‘‰,I’$IR%Y’$I’$©©d"§™È"ÓÈD–$I’$)’‰,I’$I’ÔT2‘Ó‹Ld‘id"K’$I’”FÉD–$I’$Ij*™ÈéE&²È42‘%I’$IJ£d"K’$I’$5•Läô"Yd™È’$I’$¥Q2‘%I’$I’šJ&rz‘‰,2LdI’$I’Ò(™È’$I’$IM%9½ÈD™F&²$I’$Ii”LdI’$I’¤¦’‰œ^d"‹L#Y’$I’¤4J&²$I’$IRSÉDN/2‘E¦‘‰,Inʘt€IDAT’$IR%Y’$I’$©©d"§™È"Ó´†‰ÌúüˆÅõé¶ó£WÓe[[”Aù 4sçI’$I’tï* &2}¢Ï¾ÔtŸ‰¦Ñù\3ë´TêI’$IÒÇC2‘Ó‹Ld‘iîÖDÆÄýã°ÕöظõöèØz=6nƒýöµUö΋ïjÛ$v”ñ×Ñëì;=–úÀ‰Avî²’$I’$Ý[º—MdÌá/w\d³Þ~ЧÐû' ûýýã¾…ö×1ëì+Ñ2w{ÓÝËzy¡=4jýG¯e2’%I’$)ã’‰œ^d"‹Ls·&2ƒ¢=eíÚõVwõz½®]·k×nXéÙj{zÒ¦»Ú~>ùçˆ5^VÇw¶Ùÿ}a®u¾Ã†/Ü Ê¶jt$I’$IÿXÝË&2uÿNï š¿×þ5꣄yü=¡ð ÏûÕ •w}ó;”UõÁúÍÙÓ¨,I’$I’²'™ÈéE&²È4­a"—œ©²ÝÇ.ÚC£ÖÚßÇmðHä>3wÛ…ªZ;}ñ²ý{Ï_ŽØ'^œëÂðš˜Ï4~ ÃÿaËÇëÈ`‰eÐC78»LÛnÿë©™V¸û”?_m_h_ÿXgn}%I’$Iº7”/Ø×ÄDž¼üÏûÍàUÞbyú;Ÿ¼)þÊ!Uýú>D23/,ÿ¯ÏϱoFu½~ㆠœ¿×þϳ³}¹°>Û¦ÌOÇúEŸ|qž¯Ïß¡¦±Í°\胡OPv»[e7ן“$I’$飗Läô"YdšÖ0‘”WYÑÞÓöž™Ýùÿ‹0SV”x¿ym•4H=A¤ð„ÂC6xþ>ûå •¾üg^šgö.´~söÚ›+Klä’öèØ >(!¢¸û»;í‘1ëíÓ/Õ²¾Û³ÀúÎÞc?ìSh¾Æì†Y§w¶Û‹S¶ØÁSv¾²Ö†/Þo?~uy³y%I’$IJ¿²b"Z°Ïþåù9n¾"þž3‘?õ›0éï0}R$úK_í´ÈûYÏNÞdOMÚä}¶á‹ØŸG¬µ/tX`¦Õ÷©F-)¶'Æoô§Áèý.ZŽˆd¶Éúô›ÎÛëëÃ÷…76[çiÛ½>¤Àè;k·½^tØÆ/;hmßÜj_ì°ÐÕ)Z†§¾þ•9jq±GN)šÞùfÙôç~}³?—{ $I’$Iúh$9½ÈD™¦µMd”‰ô•—Úê}åv¥îºý°Ïr­Ù_n7®ß°#gªìrí5+¯¸bÿ=lM4H[lûÊ.ÙÕk×mïñKVQSg7n܈8[ìkùçY›Ž¹Íö½ÎëÞöÍ-ö‡¨þ˜ÈßÙn«ö±ëÑö£¾ÎßƬóºäÖY’$I’¤ô++&2Æ/Ó¾Ñe±‹¿§¬,ñy¿¼ÊÜuÊýóñ³Õvâ|ÿ=jÉï÷¬æ¯ªµš+×|z¯÷vÙÂm'üï“.Û¹ŠZ7£®“õ|—u™¾Ãç½õ£þßgfYï÷wûçÂ]§¼ßv_T‡Êš«¶pk™Ýße‰íú`päL¥•_ºâss£º,ªË‘Ó•V[wÍÓ‡µ‰¶ÉS_õËWYMÔŸ»ÕíÏÃ×ÜUR’$I’¤/™ÈéE&²È4­a"Ÿªô,DO_{Ôö»äUŽˆNà<6vƒGà|¯g?Wc+÷žö0±ð?BÉ€q_Y…M‹¶ƒ}æÒe{w}iƒ‰Ì‹ô€(š`"×GÓÌ·¢=§ìhy•}ÿ•eMê*I’$IÒ½£,˜È˜¬ÜPg{—nŠ¿É_Ìôÿ¸Ò~=h•]®»f½gîöôÜxçÉ*n°õ‹ù Ç"–aÌÒb¡Þ×;/¶¹›Ë|‘ÈÃlŸˆæzv¶ÍXWêóN_¼âOxùS\Ï¿¾Ùž‹D‰v¦FÔ2Ç›íy=NVø²o¬(ñrÏßëfõãã7xÿŽŠúÛŽœwÓYï£$I’¤^2‘Ó‹Ld‘iZÃDÞâR4ºæ)¸¥åÕV}åš`ˆ‚ù×ççÚö#<ºeèÂý6$Uî8zÁ®DËY|ìlµ›Îs·”Y÷vúçKÑ@ç›Ñ@èÚÈÿòÜ[ºó¤•œ©Œ:‹|—[_I’$I’î eÁD®½vÝ6>gã Ùä¢Ã.þÞUzÑŸÀ"™÷8ü¨ïrûÛ˜õn$¿³æˆ÷¡öŸ¨p“xùîÓVuåª};:¤ ›¾î¨ÐÒSÐûÏW—Ó²þs÷¸¼z¹úe¤´ V°ó”ïÓ†®ñÈèšh{À_é¸È~új‘›Â}gï¶Ù›Ž»I¼hû 7KNWzÔŸj3·¾.Ñùº¨¾?×î^7¯^¿aßï½Ìû…¹ÇB’$I’¤Ö•Läô"YdšÖ0‘IgÁ`…üÃD­0 ùù€þb½Ñ ‰Ç4÷”]òÌeçj\¢ÁѦCçüÅ{?èShï¬=êÓSz`40¹/À`"EóA&217MdÕ”‰,I’$I÷®²`"C¿¹{ìÿyr¦¿Üñ÷Ø‚bŸÇû!èѪ©½õ“ªlÙÎSnüŸ¬7‘Wì=mG£þÓ—;.ô'¶æn>îi»èoÑÇ¢¬ºÕ¿XoÀ¼½þ¢aÞ3ýÔÄMv)Zö¥)[ýeǃæíóô_Ëwò>Õ¯¯r³ºúÊU;|ªÒ–î8éý·Å7MdÊÅ8¦Gù\»~£¡/Ç“hôç6Ÿu#[&²$I’$}ô’‰œ^d"‹LÓZ&òÊh@BÄqxi ƒœ{NÛ¹ÊZt)ŠþÆ@þV·¥öù›ž6S¶X¿Ù{ìÃê_Ãà„ÁÚ/®ôeÄü¬_‘Œ(³£Ñ7<–É _€Ç0sMäe»NÚáÓ•ö¥hÐõé—Z¾_’$I’$%«¬˜Èƒìsó8ÌãoRMÀOú-·a‹öûßÏNÚäýªOFó1eIñÅ låÞ3õ†rÔO"ÊxRÑa7ŒyÉÝ¿¾0ǧ=2z½ocÀܽöOÏͱõ-tc˜”'/ÔØ×;-öèb ã+µ×¬Ë´ö¿Ÿže3ÖõmÕ¿tož§Ï ßµhÛ  (=[mó·–Õ×+êÛÍÛR›ZŠf_˜NJ ž4û2¦v3ÇB’$I’¤Ö•Läô"YdšÖ0‘ž­²UûÏx¿0Ú½³úˆG«<ػО‰FðÞ†cžª¢Ã[Ûü1¼,æÉ }ÞûëK}@ôßÑü’ӤŨòHä}Ñ€çÔÅËöÜäÍöÔÄ/9¢ é,¢A/‰!ß`·é;ìß{,U4²$I’$ݣʊ‰üZ3&ò䢛&ò«ËmÌͨän3vxz rQ¾ö^f+öžñɘÈôÙþ0lµ§Â *øѱ쩨EÚ ÀD¦?D9ëœõi«÷Õ÷ÑH%Ï_)ô›ï³6óTôÓ~?dµ›Ê°åðyûaŸB¿©¿p[™G"2êÛ=õÅ`Φãþräooóô{Ž_ò”ä I’$I’ZW2‘Ó‹Ld‘iZÃDfà°dçÉF&2?j}Þßí9üpðyäâ>p!UoøÞsü¢ýbÀ ß/äc:/w!í\©Ì§ÝÔmõ‚!Í@ã¹¼âŠ==q£ýnÈ*Ÿ†)ýÏѲí§nõÇ6ÙÎã6x¹¹u–$I’$)ýº×Mäow/°Ê¨D:‹x‰¿1Šé#‘þëgýWØ¡Ó•n ³¯%gªüi®ºè3OjÍÙ|Üß#QÿÄÖî>c‡¯¬õ”ô›èSÑÿé=«þå|C™2úÎÚcÿüÜ¿QÏöI—A=è#‘‡ùÔ…Ëþòcæí*½àý,>÷‰¶…Q=sC©›È¬Ãñ²`_CîjÔŸ#ÚóûÓJe!I’$IÿÉDN/2‘E¦¹[}³Û»Vl“ÇÁ_ÈOÌà‚ǵÎ~ûÚ*Ïå‡òù±üØØ öçák|›a0Â2ßëU tÖÙúú£t[êƒ)ÒVP†ç¼YöÑço÷¨ŸŸ[WI’$I’î ÝË&2úBÔ_¡þ_ëÔ8Í“o˜yôcèç|#ê+ýψµ~sœå‰êåÉ-úMÌ£_ß6}7ê÷—hnȳú?¤ `ûÞ‹úFñ>Ê­/æûÎ;ٯ®ôí0ÿ¢²Ù>¢üxÝ)ûîK|ßÝìÏÉ@–$I’¤œd"§™È"Ó´†‰Ì ¯¹ÇhäÎã3å1ØÈ]‡A”Ï‹4¹i(ˆ€a^Üt øßásî4I’$I’î-Ýë&r؇Üþ }™ørñþQ}ÄñÍþLβ¹ë4êÅʺ]?,·>Í—=¯ÙmÞnÜù’$I’$}t’‰œ^d"‹LÓ&²$I’$IRk+ &²$I’$IRkK&rz‘‰,2LdI’$I’Ò(™È’$I’$IM%9½ÈD™F&²$I’$Ii”LdI’$I’¤¦’‰œ^d"‹L#Y’$I’¤4J&²$I’$IRSÉDN/2‘E¦‘‰,I’$IR%Y’$I’$©©d"§™È"ÓÈD–$I’$)’‰,I’$I’ÔT2‘Ó‹Ld‘id"K’$I’”FÉD–$I’$Ij*™ÈéE&²È42‘%I’$IJ£d"K’$I’$5•Läô"Yd™È’$I’$¥Q2‘%I’$I’šJ&rz‘‰,2LdI’$I’Ò(™È’$I’$IM%9½ÈD™F&²$I’$Ii”LdI’$I’¤¦’‰œ^d"‹L#Y’$I’¤4J&²$I’$IRSÉDN/2‘E¦ùÝÕöÉ6s} &I’$I’”}ú¥yö@·äLäûº,±ÏDuÈ­—$I’$IR’úÄ‹síO#d"§™È"Ó<4j­}­ó"ûf·%’$I’$I©Ñ}]Ûû,³ªLäKÕuöýW ìþ¨¹õ’$I’$IJR_í´È·>·û"R€Ld‘i>j媬¢¦N’$I’$)5ºPyÙö—صë×s»/9W¯]ó²/V]iR/I’$I’¤$uæ|¥—”æv_D ‰,2ÍñãǬ®¶6w²B!D¢Ü¸~ÝJKÚõLäk×®yÙvãFî,!„BˆD©½rÅÊÊŽçN)@&²È4ÇŽ³+ÑB!DšÀÈ=z499©²…B!òqùòe;~\&r‘‰,2Ld!„B¤‘$Ü$ËB!„ȇLäô"Yd™ÈB!„H#I¹I–-„B‘™ÈéE&²È42‘…B‘F’4r“,[!„"2‘Ó‹Ld‘id" !„"$iä&Y¶B!D>d"§™È"ÓÈDB!DIÒÈM²l!„Bˆ|ÈDN/2‘E¦‘‰,„Bˆ4’¤‘›dÙB!„ù‰œ^d"‹L#Y!„i$I#7ɲ…B!ò!9½ÈD™F&²B!ÒH’Fn’e !„BäC&rz‘‰,2Ld!„B¤‘$Ü$ËB!„ȇLäô"Yd™ÈB!„H#I¹I–-„B‘™ÈéE&²È42‘…B‘F’4r“,[!„"2‘Ó‹Ld‘i> ™í9rÄNœ8a7nÜÈ-„B´ˆºº:»xñ¢UUUåÎr˜á«®®ÎåÔÔÔøüÚÚÚÜYNee¥oñÃB'žm6×–ÒVTTØ¥K—îh›¡žÍm3@=)ûÃ@=Î;×l8çÏŸ¿í1Ë…cOÝ®^½Ú0íSWêÃÿq5Wæ‡%I#7ɲóÁ5^ZZêý·ÛÕ­¼¼ÜJJJüzB!Ä-è«ÐÕô]è3Ýn~>èK±mÚêÛAÿ‚íß® ÿ°ÈDN/2‘E¦i yË–-6mÚ4Û¾}»>pà€=÷ÜsÖ§OŸÛô?JøA.((°3fØ©S§rg !„¸¡¡]éÒ¥‹uîÜÙ¦NÚÈ$Û¼y³õêÕËç¡éÓ§7´oüÿþûïûº¨gÏž¶fÍš†uÏž=kãÆkXwàÀV\\Ü0¿90QçÍ›g]»võuºwïn+W®l˜úôi9r¤Ï{ùå—mðàÁ~ƒ5lsîܹ ûØ£G[µjUîbn ¿òÊ+¶qãÆÜYͲgϯƒ–8k×®õc?f·3¦1ÃßxㆺõíÛ×· ´µ{ŽE·nÝ\l—e×­[—³¥O’Fn’e烸S§N~>¹9ѯ¿þº=þøã¶bÅ ÿÌ9xûí·íàÁƒ9K¶œb›B!ÄG7¼—-[æ}”Ðï™3gN#׶-̧³hÑ¢GŸaÖ¬YÖ±cG_—>ÚŽ;rsX®wïÞ·mÃ?,2‘Ó‹Ld‘iZÃDfñØcÙÌ™3ý3ëK–,¹m´×G ƒÛ'Ÿ|òM!„駬¬ÌÚ·oo“&M²C‡¹üâ‹/Ú{ï½×0¿]»v6aÂ7µŠŠŠìù矷¥K—úüÅ‹ûò˜j‡¶)S¦XÛ¶m}Y£GöAÃÎ;mß¾}nøb‚滺aÿaºpáB¯Ó;ï¼ãep#FŒá¦7Fë¶mÛ¬_¿~nNÇ#xsY¿~½×›AK|›¡-ã öáµ×^óå¨Ãí`¿öïßïûŽáØ¡C‡F&2éK/½äfƒ)¦S6f?Ð zÃ]@Ê#JŒ†uØvøa'bŒ6?ì x©Ïüùó=ºŠGR˜GYÔ?ÜdßÙ•R>óLÝÍñBÑ~k1<ã‘Ç<ƒI‹ÉÇ|LÒx‡œhÌþýû»ÁI”,¦f€i,Ì:´q3Œöˆòng€R&&*í^€ß~ÚÚ¢KØ>íD€í38¹]Š¶É cøðá Ó ÝË6ýÄ&B™må3‘i1ÝY–m0˜‰·£´¡ïñÇ+§“'Oz[M›N=(ARC›õ·nÝÚ0-Àþs3·9CúNHÒÈm­²y´ó6ôW8¦D¾C8Æ Pé#Ñ? 7 ¸1B߃iüO?8\[\kÜaÛåZ ý“½{÷úvÙ&×È«¯¾ê7<0“駅帾‰¦§ÏCÿ&7õÛçœSÎ%ßÖ%‚ó¶I\[ô…˜ŽÇ€úVDzÑã{ÀþÓ/ 7y¸YÃÀ; pm}x:zq¬Ø.ÇmÆ¿c¡ïŽ?ËÓ¯Ž)ljúñ}?Êç{A]Øwú³ñë€mÆëËyB‘,áw›¶+€ÉÊ4ú%Ë—/÷°¡àwˆàÇiôí´w´‘ôÄ Û£M¦= ÐΰÝÐÞu =¦Ý$RY&rv‘‰,2ÍÝšÈC‡õÁ7ƒm"Žˆ´ ÑYü8òãF§ý駟öA,ƒÐgŸ}Öç0À§a@?óÌ3~÷.ÉÜýc Ìt¶Ïÿ˜ᇒÈ)¶Á@ 1?<~Ìz”C^xá7hT(²)ñ7Ó¨#QbO<ñ„7,<†Â¼P&K¨#ÿ3H‚°=u¡ŽÔ‡ÏcÇŽmQ%!„̓)Æ Âø#…“'Oöß}<CôpÜÃ8â·žN;ˆøJÚ P:üXeqƒsŠvävé'¸aIû…é‡QÍ´ƒüO¨mÃ!Cn‰Ì6iOˆ|‰CÄ'ƒöm…\}˜ÂÁùÀ¸c»´q´mñ}Ä(cD=‡ú²L0Ó1Éh1'ÙÇ(nî1H£Å€‹C=©3ýƒ\SòNi-#·%´FÙ ,‰§oÀõF…¿é+a.c¸Ò?""‰cÿðÃûuNô:ç&Þ¢B‡sHD8ÛâA¼Âö€›f8þô߸æ1ÎùLyÌ ý8Îu Ç–éÁ ¸€ý×—ï퇉TBñÑÁï:}úgô»hgñ#¸áO_ˆwð[OûƲôøL?†Ï!-7Ç1€óßîÆúC´1´ÝôX?$Þ o„'‡Ay´ß2‘³‹Ld‘iîÖDæÇAhñtØéH3à`pÀ‘tÀé”óã.tÎ1˜…N>?º °t±w ‰!¿1-FÜÙ£q  -4DNÑágÐÁ|îò1Ÿpò63b Îgaüøò37D Ѩ0¨ D…QÇ7ß|ÓëÈ …F!Ô‘²hòÏ@å1íšË›O›C{ÂoüÿÏÞ8YRTíÿk_ådß7pQ@dSYdUAYdœ "†}‡™ˆ™‰ ~ñ)~§Í)nßéé¾3Uoñ|#2ºï­%³–›çœ'Oe­÷bvÓömf.L?A´kvŽèÇæ±-ëaŸÖ“ݶjK+Ⱦ7""ìÓPD†¬Kv;ØC¶¸AeϤ9Dxmbÿ†/äcOÙÒáÓE›aBîfYEÝΟónð@ ZþKM{"ÃÕÿνudÉò7œW÷‹ù±=íd.ê#Ž8¢GÝŸöiÙøD]ÙS®£ß€@ÕúÁ%H»ÿ}o{÷ŸÉ6ü0A¬ú –r탰êúº'ÅbËù>öɳOƒ%öI(ö™gŸ²þëwcð†ï&pvüü*Çkß²»üÌê5 ãØìϱz—ù§®‹ãª¤Ç¡?п#÷:ßÏÿü8÷©éÖ´Áq9ê¯D¾Ÿ'×ZßOößÏg¿kþ©ö:&ƒ=Úëœj¯ó·^BáÐài{ì '’Ù\}?ØG>•XÒw›zŒ½¯÷ ä%<³yì°Í&^=º¶cÙì ¤¦ö䓲¹ì%½""ò¼‰ˆfÍVEdbÊ¡†Àh(" 4tÚÐIëÜ9úÀ ”jN¹QA±9ýLwA<ÖÑ3[,ÑÙ~,È rEnýz™q°_¢² F@\ÐÁ ½ŽAW£:õ¸ðóÏ?߯#pf¬Ëµs j«ãt„BX-œoB­¾Ûcæ•Õ[Àm¦ò"™(Ä>áôûCˆ²D$ëD@l"‚é×%„W‚”}ø¿e‘ˆ\s8³SÙ@öH6ŠP…ЬìÞP˜^$"äDDf¿†"r}džÚQO&chÏãÁ’ã ò!ZlG|S¶"¾«r7Ë*ê®À”XÌ÷pÿ”ïAt%Øò—ø2umdJ•À[/9t½ œ|~“ûD©).ø)|~‹:ÝSßz±ž@ÖçòWÜc•ÍwヹçË«¹¹Ýwu/hÑUûྱÏÊ\oDDv¼µMeôqý¾´ÅýWO“9ÿëQ"r‰å0â\:ÇÚé\ø;ÿœ‡zrÎ ‡{Y{ˆö…à_›ë·UóW ØÍ%^í­¬l‚!„qa«ª_æë|1~q™Ý¤QX‡/ÇVøŸ~¡¯gûØ ¶vøÔÙ¶œ.ÁŽÑUØ'>(;Ë6òøMÕ–‘‡ƒðJDäé9ÌšUˆÈèý‰È–× tÒ‚N~@Á‹}ƒuü–Ñóø$_.0•E#˜ô¨Ãºõ8b=þX"reBÉPÑY3UxT6P‰Èê—©<‘+sÌ: D+"[.)<c?‚«B«ÃlG|˜ÑKÐNgA@,”“nîX6€à<Ì@¶N LlYûB="°e„9bÛC¬bCj½Â4Äl™œla+ƲwÄ1Ù¦²Z쨨0%"/š"£æèk÷µUÙ¾ ¦²­mö¤Ðöœµ°{Úé /Ê 1Èë®êek«r7Ë*ê&dòaœ3O*¹ŽõΆ‘+˵ꑽl9·½ŸýoФæD&fÖ5qËï’ »HD.Á·Dd÷mebÕ´]¶Õ÷7ñ—¿ã·ÒbßÕ&÷Ïz"r îÈ*æÿ¹Çëžu•8ìw"@wÏØN;øuü,bîzÓÉ Dd>—`î[>ÿQοýj—￱¡ˆÌß³?Ç&+ͱžá:Y_»êܯ×Þe¢w!„Cƒ¾ŸÐ÷ëÃÁï1­…ø½õ§Ø'}9MdQr@ME¶ ƒìFûVÍ“ì©j>»A aùŠlљߴY""O—ˆÈaÖ¬BD–µ"Hðˆ#¶""„•ÕaDPP#2žÞÿ‚Zb@eŽTJD.W@U"5ÃA FhçVDdAX Ú!Ð!„°Ø θ>v‘ÀÉ>°íKâ¾DR}:¡‹ÍY$ºZÎF™ñ¨£€¤‚‘–Ùø`_o=UVûô=dù»ö~•ý/@ŠÈî¯Ê`ç#Öù¨— »÷Ý}ýÖ‹VD®)ÄZ™ÿXç_o}Ë%lFD&¶;íõ¹žd°OçÀ½Ÿ÷a„¸è“Í[/¦g#ôù|чSŸ‰ãõëúw6TRBaq¿Áüex¹*»À¦²Ù"öIÖ3;È/õW;ØB6}+ƒ‘§KDä0kV!"—H, ÑË¢á”s¸un¾ãä×Ûåk®"ÁFHDaÁH½žÃoŽ½ †`+ 3‚«šCÏ£)êò¹2bˆÈ3F¡FµQÀL|üh£ÀEv›`@ý5ß1êE/­ˆl2uJD¶Ç!@d£žû{‹k!„cª}«þ˜Øeî}™å|³e–™‰lì RŸ.ˆÐ¿Ë*&Õöì ±É¶lU»ÜßšN`õ¨¢Çä‰Å«zF0ÂÎ Ó2*e³‘—9ü‚ û$4×>Ù™áK»ìŸ¬ ÌýA¬yb Ûj¯9]Ùâš¿Wö©À‰x&c#ëE„lt{Ž”6k›hè\¯jnØU¹›euóøü¾?ÁõTøõb·ö%„þÊRâoðÜ„Qû!¬º†î#¢r‰È^~XóÛ^=|¦šæÁg>’kã>·zi° V;‰¿ÖÑ.ׯã³'ÀøPîŸkŠ/ÛÕ>mÏ·sÿð«üÖÔÉ?ò;¬é,,×övÇd?Ú#^€ïØÕ»ÞK(áºÀ—VD6ØTSKÔù÷;tþõuþÝßõ2ÀzA§s§^õ·"²câOú]T{ùÚ[sC×¼š!„ƃÍ$вIõbqõßbuüñ5øsõ4a˜ÁFð–Qv›­–•ÌþøÜ·ðɬëéê­yºDD³f"²àØc—:lY_2bÏFîì[¬CælC†’,eŽ~Í9Yve[ Jý/há¼ë৕ݢƒV—m4‘ºŽ…€`ÿ„ QÁ¨¤ XPÎìŸa &Ø#S£ È1U¶‰@¶‚êVD.£Á(©sUñ†BøoÖ¬@@_«èkÙÂl MÄ-‚1̺÷8ØxßÛ¦¶U¬CH%L³!í¾Óf,›ã”Ðæ±DûQ§"0©ì]B+;¥ Ñ©²-×cÑ> „©§iˆ‡“UÚfj;ožÒ¶:õòZ°£l¨`‡p_ëµçH)»D=6·Í`Þ «r7Ë*ê6À×(߃ßÀÿàÓ)ëÍðîãV€tö|¾…kã¾à{¸>îg¢k î+ë2µ—èëÚ¹†à¿¸_ÔÅÿ²Žc#^ó}øY®¥)¹êÚi{ |hƒ{‘¯Víä‡Õ>Ýç¾—éì·TÓZÈèÒ.zݳÚÞ>¾ë˜ê>¬9,Õ»‘löš‚¦_™ÀüQÇ¿õÛ'¬;÷Ös¼Šcâ«V†~7…Og}ËáwL°¯ù¦µWâC½ ÏºÚ;|ie!„C DŸÞ¾¯ˆŸW}¾þ›-Õ—ßÃN•¾À3ÐjlžuøuñØmÚmÙ3vq½'·ødt‰á‹Ÿ”ˆÈÓ%"r˜5«‘¡ÓÌÈ4ádëÔj¿–qì+3I€b™u*X̶ëö'Øhß‚! è<äÙŸíêñBu“=^¢Õ.u3þ¯7¨ÖúÚØfµëÈÒ‘a#ˆ·žË<ÊB«§úc}l[†öË:'[› ÷ý¢íõÿe†Ë}ÞHàÀNÜõÿöM cw†ömËö ¶³µOû£ìò"qWö&[:œã¸¶© ÍõÎaÛ×ã@Žs¬BÈÝ,«ª›ÏÐúåK¸7꼶¾G‹û˜ÔšuíÛë9ÜOëß®kÜfŽÛŽø©}‹¦Žpìü¨õ–/Ú§õµ¹Úç¯ö´¾ß¢sjÿêñ{©åö%QÁuh‹ï¬çxÛ{pQC߯~ïíùoû‘á>«ÿö5‹ÚBa<ªÿ⻶g³ø=뉸låú~°-âü¡-ªRë±lÍþ|>ëiÏ"Ÿì@PoDäi9ÌšU‰È_Ed»x\t£sR†Baã¬JÈÝ cÖ¾À ½ìøzeßU–u!„p0!N{érÍiÜßPƒˆÈÓ%"r˜5‘7ÇcÌ%¹•·ª†Ba1c ¹cÖ¾@––ó¿¨l5ƒ+„BØ(ëÙ£1mQDäé9ÌšˆÈ!„B˜"c ¹cÖB!„°ŒˆÈÓ%"r˜5‘C!„0EÆrǬ;„Ba‘§KDä0k""‡BaŠŒ)äŽYw!„Â2""O—ˆÈaÖDD!„ÂSȳîB!„eDDž.‘쉈B!„)2¦;fÝ!„Bˈˆ<]""‡Y9„BSdL!w̺C!„–yºDD³&"r!„¦È˜Bî˜u‡B!,#"òt‰ˆfMDäB!L‘1…Ü1ë!„BXFDäé9ÌšˆÈ!„B˜"c ¹cÖB!„°ŒˆÈÓ%"r˜5‘C!„0EÆrǬ;„Ba‘§KDä0k""‡BaŠŒ)äŽYw!„Â2""O—ˆÈaÖ‘wíÚÕK))))))))S){öìMÈU¿ºµaØ®””””””””1ËÎ;#"O”ˆÈaÖèx¶oßÞJ))))))))S*Û¶mMDV÷°=))))))))c>ÊŽ;†îK˜‘ì©LdZJJJJJJJÊTÊ2‘÷îÝû¥v¥¤¤¤¤¤¤¤ŒY’‰<]""‡Y“9‘C!„0EJÈ,jƬ;„Ba™yºDD³&"r!„¦È˜Bî˜u‡B!,#"òt‰ˆfMDäB!L‘1…Ü1ë!„BXFDäé9ÌšˆÈ!„B˜"c ¹cÖB!„°ŒˆÈÓ%"r˜5‘C!„0EÆrǬ;„Ba‘§KDä0k""‡BaŠŒ)äŽYw!„Â2""O—ˆÈaÖDD!„ÂSȳîB!„eDDž.‘쉈B!„)2¦;fÝ!„Bˈˆ<]""‡Y9„BSdL!w̺C!„–yºDD³&"r!„¦È˜Bî˜u‡B!,#"òt‰ˆfMDäB!L‘1…Ü1ë!„BXFDäé9ÌšˆÈ!„B˜"c ¹cÖB!„°ŒˆÈÓ%"r˜5‘C!„0EÆrǬ;„Ba‘§KDä0k†ˆ¼gÏž>ðzýõ×û ,„B‹Ï?ÿ¼{ꩧº?ýéOÝŸÿüçnïÞ½Ý;ï¼Ó}üñÇÃU»Ý»wwï¿ÿ~¿ÍO?ý´_fÂzœø]»v}©°…!òÃ?ì>øàƒ¾þýaŸêúì³Ïöùž½Ö£´m‚¶ªk؆ávU,Ó®á÷Š¶Ôù°¿áò²óÚ ˜±¯§Ÿ~º»ûî»ûómûÍ2¦;fÝ!„B˜ü,¾ÙG}´ß""òt‰ˆfÍÁ‘ßzë­îÌ3ÏìN;í´îÝwß.þ:À;ï¼³»é¦›úN5„BXÄÀþð‡Ýw¾óî†nèmÞ/ùËîþûï®ÚÝsÏ=Ý~ðƒ}Ä×;wv¿ùÍoºÓO?½·köóøã÷Ëìë‚ .èÎ8ãŒÞî)gŸ}v¿ÞwÜѯóÆot?úѺSO=µ;å”Sº‹/¾¸Û¶mÛÚþ‡<üðÃý>Ô§Üzë­½@ë8ìÇwU×YgÕ×}ýõ×÷ÛÆõ«_­mû½ï}¯ûç?ÿÙ/{õÕWûm†mõ÷ùçŸïþö·¿u'Ÿ|òÚ²Z~î¹çö8œÛ;–ùÿÉ'Ÿì—½ùæ›}}„åx »ð »“N:©Ã7˘Bî˜u‡ƒ‹ßÓïÿûî–[nÙÒýB!l„×^{­»è¢‹zÿ?øãÿ¸÷›¶BDäé9Ìšƒ!"ËðÒA ý¿?èáo}ë[Z?„BØ(Ä@Â0ѨøÅ/~ÑÝ{ï½ýÿ%™³¿ûÝïºN8¡y[ù®»îêN<ñÄîïÿ{÷Ê+¯tW_}uoßÞ~ûí^`$ÒþãÿèËO<ÑŠ{ì±½Ðlùe—]ÖwÞyÝK/½Ô½ð ýþ ¬mÅ‹/¾Øo{Ûm·uÛ·oï…î£>º{ðÁûåÄ^í¨º,?æ˜cú,k“ Õ2¯,?ýéOûÏt?ù䓵v*Äߟÿüçý±ìر£·¿ÚÜË÷¿ÿý¾½üçé»ßýnwóÍ7÷Ç¡-Ï=÷Üš'".³ép¼DfY7›eL!w̺ÃÁÅýìwá·æÞ!„×%ðEÿýïwÏ>ûlï/ñÑ=ù¶Q""O—ˆÈaÖ¬JD€Ë–l?óÌ3}“Œ¦…eby´U¶±@ý‘GéRYSøÃzÑùøãï—ýë_ÿZ¸Í£>º¥ÇbC!|õX$"ÿä'?éí <#SYÆ-[ÔŠÈ=$¤N ¢Q™­b}Ù¿ýíoצº ÒV]`ÿØ»áT¸ñÆûv´SA]~ùåÝÏ~ö³f­ÿ"£š\SQØ/;\vµõ±Çk¶úï2A QzDfË+@!D…×ÝØlËKD~ùå—#"ÿ ¸ÿø2¦ø0Ýœ3ƒÄxA¥å!LVø^†|]ír@¨ ñeÔô"·ß~ûš/ÕNRË \È̵¼üBË ,ÄàÏ=ôÐCýq¸·ƒc±O÷÷²~_öËoTÜkíyn—k—úMøÝÈtÿË_þ²vîeçÑà‡¶عï¾ûú6ÉìçGhÑfÛû‘ý6üV x„BIü3ü›|Î9çlI‡‰ˆ<]""‡Y³ Y@!SꨣŽê;Hâ± Ù£2”dCùÿÛßþvÿx-§Ýºwˬo;`²B$2°m#Ëj]!„¯‹Dd‘³ Z)Ä¥óÏ?MDfoLñÞ{ïõÿjÿøÇ?övÉwCLaÑfãÚFæò%—\Ò‹pÕ+¯¼²±[¡¸ ¦ Û{üñšk®iÖúâ!¬Ú!˜ 6¶Gþç?ÿémkMiQ çÚ¶‚1L4,dûŽ¨G¸65ˆsSèuuî""Qø¸ãŽëý>ΑGÙ_S0qßg×Ñu²žÏÎs ¨ËT:ì°ÃzÑîAõü.™ñËp]'ÙêÖ÷—/uÅWô÷¨c4EI-×>Ë=f+sݵ”Eï{S•|ó›ßìÛÊwt ÖmÛ,;}h¿cr>Ô믺ýë_¯ †øÔòZǽª=¦LÓõi# „xšM›¬ãø´ÍSnüKÿk£û–Yç»üNây!„p00àé©4 l?;Æ^]{íµÃUˆˆÈÓ%"r˜5[‘u\cNºŒY5„^_pk¤Ã^NºŒËÚ‚Kmð?G^ФC”}bÛmÔËŽ;­ ûóˆ#ŽèFçÒ”l¿,ÛEDDþbD癟$›Ö½àš~øá}v¯ï\ßò¡ÜCæ´&zxpšzÅy¯ûŠÈʲ_Å.CÖßɽbÀú2Ý ®–\°ÜüÖ®£ãU¯ú –u]gmt ®»AŽ«®ºªÿ418" Ø}-k>¥ jõ[× ‡Lkþ¢í}v^J6w¸6»‡}GtçWj‹c*Ùy´½sëœ ÔCOèpncõðS+kÞ~Ô_ƒ>!„ÂÁ€ÆQ¶”ý–D·ìý!"òt‰ˆfÍVEdŽ{•5§ ˆ\ÓYpè& AÉ¥—^Úw¢‚sÞ=Îá;#sX´Á6y|3„BÀªDd#Ю»îº>hK„)/ ˆ¶Ó.o YD0SxyqŒÍ[öâYuÉš”åITâQ}mXϵEöh;5*»”X¹™ÌÂ[‹ý8þ:7ÄLB ¬êEBkDä/^ÔÈGj3ÉùEø;DQ"®{¤|(ËùO¦V!˜ÊÀݬˆl‚u;Ø àT¿'Åܳ–Ë–.\cm"Ȫ«2Žkן€«~ÙÃTd;»W µ„òõpŒîkÇ#“¾H›ÒÃý\Ëý^ 瑈ì·ç·±YÖt›½]OµiŸ¬0ÛÛ?8„B8X°åü.¾¨A|öÜ€.ûZvl3DDž.‘ì٪ˆlî<ξL™B`ìBA ¥¹pÎ=¦Èq·Ù#°%"×c¸¡nÃù·ÍzAs!„0d+"2±É|©5‡-dgzAŠlÌ‚XL05S 1—¸ÖNÁîún(Òž¼©¹™Ûé" mSÿ¢©(´×TQl'qÑ“;¦@P¿zA f¿‡âï"1Õ9uÜ‹üˆˆÈ_ÌyHÌf®»†î#ƒ ¦j ÄOן$8,¹DUA£õˆÈDÞVL!WýŽÏ£´DäVЕ-åþ0á¾ä•+žßVþ™ ØïÅ`‚Gu%,‘³õˆÒí LµÉ}dNqYÎír·mˆÃ²ˆ‡"²§Jø.Ù9s~á8œ3~$ÁÚ9,y8ÈB!¬6Ø k;hÉÖ²užÚ,‘§KDä0k¶*"{üQ$ã©æeh@<6(+¹²XŒ¼AÛf"—ˆ,‹¤6A‰ýÖ Úm""‡BØ([‘‰°ÕVdƒýµ"®GâÃt‘‡ÓFU‰WµOZe¡šþ ¦ X4g2ØIu•@ÖB¬$@ÛÑR/³]ô«áÃý/ t(†Êf&h‡DDþbÞjA¢sZ÷”) Üî5>Îð)+S¸¾ÎÿH3A´2ùUî‚èþDd÷‘më^u Çêç“Éð%2{Acáé/Y¿2‚ù‡¦³ÐF™Âð]Í“\Â2á[Ƽý-º¯ZJØ®éÊ |xy$aZ0md(¼›^ÅwDyçÉñ×ÜÌð=?ó@Ddçw#ç0„BØ Bë]Qÿ‚Ö±Þ”`!"òt‰ˆfÍVEd‘X$kʼDz;|™ÅœzνT†±õtš\ªŸˆ,‹D6‹ì,ÛXgÑ6mB!,ã@Dd"¬?ÛÊõ ¡˜€G°b¯j`²w ÃÇkîbóÌ«XöÇ>z,Ÿ€æIõø†ê"$‚„µ6H0‚Úácør‰b„]uÉr±½RíR‡'{dƒ.G‰mö±è\=Ùa6™m'²õëeTGDþbÂýÄrÝKÀuÞ\#SGc}.Ê <W&9Üo>›ŽÁKðìÏúÉDv¿0P¡˜˜l[#2z݃–ûÎrÓR¸þÖ';ü2÷{‰È0¯1¿O&2±×=¥2‘ (´¸¯ëeÊÄaçÅùqÜDsÇd¹b9ÑÙý^YÃ~›¦ó¨ÌcçÔ2ŸK¬7Btö´Ü×~Wêôô»SŸmœÏE2!„Â*à3±Sü=Oë°sì[&¡n³DDž.‘ì٪ˆ û WfŒ EСc@ ¼+Dàì`™Kúš¾ÂKYjŽ= š÷·M!„°?DD&XKDqVV&b›¼d® }ê N , ªŠLçz|‘hÆ®ËÔÉ–úLàR¢^s×ÚȾÊþl!ÎÕ¶íö%JCûd–VèöÕ6­@^ÔKoµÇ>âw{®Z""I¦-‰cÐÜ@ˆ—\þÍW\Ñ §|%ü»ÿLñ@ø­A÷¨óOÝî ÷¦mÕÏ?k_ìè:ck9a¸®¿kKvÍÛl)ß»öÚ ÍŽÍo¤žHÛO>ùdÿ;PŸzÝ÷^ªW˜ÊÂ9©å’Ú6Ëàö[¬iÎøî[‚4›ðl‘²£Ý÷~®)ü†mk½õ~!„Â* ‡´ä±©ìrk÷6CDäé9ÌšUˆÈl™ Y°ØÎmWÃÊ‘õRó3Öò68ÓZ^íÙÈ6!„Â2JDn_¶Ö]$³el‘—ž-²™ìÓzb*ì—}ô$M›©i¿öWßùQi÷mÝEmTÇp»*–m}CØ×ýÙY¶™½o_ ¸SsDDþû0¨îezíýSs"™]"ì¢)»,#ð×¹ôÙ~ˆ£²š†…¸Ì¿+̱Øw½ìxˆåîMˇ÷–ÏëÝîÇT¾ˆßê_Ô¦6 @;Ôç^jïÏv¹eÚ½è~5ÃG¬Á”öwRþb»ßEÇ1ô;C!„ƒÛĦ±™‹ìÞyºDD³fU"r!„0EJD¾ôÒKûÇÛ Raµ˜úÀÔ õ‚ÝͲJ!÷@9u›'ñßøFn& 4]…i¿LéЙ¸²k7³Ï­"³YýÃ6É öÔY;×v!„6GDäé9ÌšˆÈ!„æ ÐËÂ<êNÈŠÍ;øÜ}÷Ýýô ÄD™Ë›åP¹ëq(êü™kÚ´ ›WëÅŒÃ2&öL¡M!„Âœˆˆ<]""‡Y9„ÂWY‡ŽU ‡‡BÈ]1ë!„BXFDäé9ÌšˆÈ!„B˜"c ¹cÖB!„°ŒˆÈÓ%"r˜5‘C!„0EÆrǬ;„Ba‘§KDä0k""‡BaŠŒ)äŽYw!„Â2""O—ˆÈaÖDD!„ÂSȳîB!„eDDž.‘쉈B!„)2¦;fÝ!„Bˈˆ<]""‡Y9„BSdL!w̺C!„–yºDD³&"r!„¦È˜Bî˜u‡B!,#"òt‰ˆfMDäB!L‘1…Ü1ë!„BXFDäé9ÌšˆÈ!„B˜"c ¹cÖB!„°ŒˆÈÓ%"r˜5‘C!„0EÆrǬ;„Ba‘§KDä0kˆÈ;wî샥””””””””©”Ý»w&äª_Ý{öìùR»RRRRRRRRÆ,4œˆÈÓ$"r˜5:žíÛ·÷RJJJJJJJʔʶmÛF‘Õ=lOJJJJJJJÊØ…²cÇŽ¡û&@Dä0kd"{B€–’’’’’’’2•" X äÿC YÝ{÷îýR»RRRRRRRRÆ,»víJ&òD‰ˆfMæD!„Â)!W°t¨³îB!„edNäé9ÌšˆÈ!„B˜"c ¹cÖB!„°ŒˆÈÓ%"r˜5‘C!„0EÆrǬ;„Ba‘§KDä0k""‡BaŠŒ)äŽYw!„Â2""O—ˆÈaÖDD!„ÂSȳîB!„eDDž.‘쉈B!„)2¦;fÝ!„Bˈˆ<]""‡Y9„BSdL!w̺C!„–yºDD³&"r!„¦È˜Bî˜u‡B!,#"òt‰ˆfMDäB!L‘1…Ü1ë!„BXFDäé9ÌšˆÈ!„B˜"c ¹cÖB!„°ŒˆÈÓ%"r˜5‘C!„0EÆrǬ;„Ba‘§KDä0k""‡BaŠŒ)äŽYw!„Â2""O—ˆÈaÖDD!„ÂSȳîB!„eDDž.‘ì٪ˆ¼gÏžnÛ¶mÝöíÛ»½{÷‡Ba?ìܹ³{÷Ýw»Ï>ûl¸èK|þùçÝG}Ô½÷Þ{ÝîÝ»‡‹{ñóÃ?ì>øàƒ…v™Í·íÇÜïkÈþÚòÉ'ŸôÛ/ò®ê~ÿý÷{ÿ`«Œ)äŽY÷ªqMøjþ†éâ^ã—¿öÚkÝ®]»†‹Cáÿ|“ò3–ÁÏ°ÞßóCø#Cöçñ{ø ì;>¤Úf[ñ¸Cöçí¯mJDäé9Ìš­ŠÈo¼ñFwÊ)§t'žxbßá†Bac`î¿ÿþîŒ3ÎèN?ýô¾üñ\7¸xüüç?ïíîi§ÖwÞyÝ‹/¾¸¶üÍ7ßì~ô£õËO=õÔî?øA/„Ï=÷\wî¹çöÛZ~õÕW¯yÚrß}÷íÓ–;ï¼sMDX]wÝuýv¶ÿÎw¾Ó=õÔSkûTýä'?YkÛ\Ðýûßÿ^[¾ÆrǬ{U¸¦=ôPMŽ<òÈîöÛo®&Aàûßÿ~wüñÇw/½ôÒpq!Lž§Ÿ~ºûÞ÷¾×ûlÏ5×\Ó ¶Cè?þñ{Ÿ¥`wÿð‡?ô>„íù#ü’{í‡ßR~æùçŸ_Û^RÛÅ_Ü/Wì›_Tð—øMåýâ¿Ø§m<ðÀšd6³Ä`{7Þxã>>Ðßÿþ÷µmßzë­¾¾ªû¢‹.êo ÿû®–;vÛl…ˆÈÓ%"r˜5[‘ÆÝsÏ=ÝÝwß½p´0„B‹yá…ºcŽ9¦»í¶Ûúà‡=ýö·¿Ý=þøãÃUû Šˆ{òÉ'÷âíþóŸ>`9묳ºO?ý´;%Ò+¯¼Ò ¹^xa˜É~ø\~ùåý¶O<ñDw 't7ÝtS¿m9úè£û I[ØuŸÿñôË ÊÇ{l÷è£öÁÐUW]Õ·åí·ßî—W`÷Ì3Ït¯¾új÷Ãþ°ûîw¿»n6ÏFSȳîU!ø=óÌ3ûë&ømÂô œþù}Ÿ9„ð Ée'tR÷«_ýª÷#ˆ¬>ÿæ7¿®ÚÝu×]Ý·¾õ­Þo)ø>| þ{Å7â‡üë_ÿê—ÿö·¿íמ|òÉþ‰+®¸¢| °Ë0æïèCù?ü ³Ï>»»òÊ+{ÿÉ€9{Ȳ-?Šsýõ×÷ûÖçªëæ›oîë.ì±Çë—k[ú׿þµ÷~ýë_÷>—Ÿ`Í—*ÿË  ~]ñY±Ì:Úb›õ’6BDäé9Ìš­ŠÈ„ã¿üå/}©GïtÒ2«tô2`« …ÀÒã­þ󟻇~¸}“qeÔoÑã$!„Âœa'eÒT¦ [H|•m<´‹m<ìf!ÐĆ< äÿGydm9»+«‘­Ìù¿Í|¹ãŽ;úLðEd”=Ô¶E&3±˜í&· Ç5Õ' bÿŠl!HSŸ6l–1…ÜUÖ-SŠ8ïzóy\Kìر£L£þ󟽠/ Kð]øžåÂõáCÝ{ï½ý5Xƒö庺÷»ßõÇcPÀurh‹v¹Ö‚Q×Rýûßÿ¾h‘iþ§?ý©?÷ØË/¿Ü·Ë_mzðÁûýÖ± 굑ÏWA2ÿPŽ³=&¸‡íÿõ×_ï·q?:7ÃÇ¡eÓ´C)ÜêjÏ­:!`Y¶—ã–éæœ:?ŽÝ¹rÍì—«.û7XSh—ú-ÓVm¯Ç²‰ŽO[ñÎ;ïôËí£Ú·ˆ‘ýn òØÿ 7ÜП×ztúÙgŸíÏKíOñ™0Bc¢/姴ý¶þ±õsÀnð7º—]vÙÚ÷?ýéO{¿£à›èù4ƒæü–‚ %ìÊ~f;ˆÀÕï 7?Å`7;{ÜqÇí3˜ª¿'æV–±v¶Sqi›6j»,â_þò—kËØ8ç쎾߾ÕW±‰Ðg6]ÛˆÇ;Ê6Þ7KDäé9Ìš­ŠÈ:>ÆB­å_rÉ%ýc“:uŽ°õˆ#Žè;^Ë}ö½€ôë_ÿzÿØÉ*‚´Báÿ ì{)›¥E€$¸Î),@bo‰Hû-ˆ!$ Ê,o B¢ï‡Ö‘¥ÜîW6å÷¼ôÒKûL磎@m=Æ… J; •Õ¶6{Rpã;BØfY¥{ ¬ªnç‡?$˜dòd'€‰>—Å7ªÏ5M‰ ú°Ã[;÷‚]A¯€t™pè~ð(±ý*Ú Ø Ët=çœsºÃ?¼¯‹ÿ&XV·6uÔQ}p[1Ô} kÌ1XÇv|»[o½umÁ2÷$h£Ì,ç’+ ˾µG]îÝzYF™eŽÍrÿ[G™í Ý~Ún¿ŽÇrÇ×ükOeñûXÇq,‘]Ç£¨Ûþ«m²Úì£ÚRƒ;v k=þÖñ¡ƒxí\9ÿ¦tùÙÏ~Ö}ãßè³õ‡¿ë–ÊVsüÚR×B=ð¥}v®µÍ¹øæ7¿Ù ë!„0&|€ÖO€>K¿V¶Ô`{¾7pŽÊ$6ˆ×¢ßäŸèLJ"1ÛÃ1fàM_Ì^ìÛkPÔ€";ÑÎ7OÔÕß²!l­¾ºE[ôëDh¶Ì`cÁ°›ü&ƒæúc‚qÁob òÁ´­úSÚ¶•§N""O—ˆÈaÖlUDxÖœGqÙgƨ›GÐÊ (‘}æ¤MÔ¡·s…B_ˆ°)™Ÿ-2yHCÛLŒ"ܵó º<®YÙ¬–·Ù©²4Ùg™ÄLZ› $€²\$Po²Ùl¾‚õjj ”nšZÞfÐj¿ Ëñl–U ¹›au;¯T×E I`uΈ~ÄuAo ¸ NùT×^{mï3 ¬ Žæl$Ž&!6o¿Ëæœvëº:Ñî÷êÔ&÷„v¹†êíEP–µ^B&¿í–[né—›>Å>ñ|?0aW›ÈöYÙÅ‚iûpÿ¸O ÄOuªƒ¿ÈtO „ ´Ž›(L€v>œíxk?áÔgÙÙ®uK÷·éVøŸêÿÒgÆ0«¿E6™}j‹ãT·kD¬à§ãí_ÝŽQ†Á÷¶G‘ëE•Î‡u¬¯>ó€:?Öó=ñb/6,Ùv·c$²h›óh¹{„oíÚ»G´ÓyŽ/B˜„bý—Ì‚!Øê' T—ˆÌ>é'=ÅÑbº }½¬ã¡èÊF²iÑØ þP;Ç1ᙽ³ºô¯íôì”åúTÇÃt¾•9”­Ç¾éû[Ø6š/e¹6ÚÁ†È8&¬ëÃÛ m‰xên³“”ˆÈÓ%"r˜5«‘(õRA@›y¤|´"2§·Œ>„Bø*Q"rÍI\nDD¤Ô”Q•UZÈø)™%€jEd“å„3"òPÐ&ıقëµs5 Ķ YÞfáh?ÿ@ôfY…»YVQwù?æ‡,¢Î»Ì&™Èü!Ap ²¨\G‚¥kIÝŒˆ AyeÊÊJwLÉ%涺ﳄKBª¬WT5'6ñXðÞЦq àîODÖVuW†±s¡ÂqÕwÖ5mJ °O¬Lb÷Ÿ{Ñy ²ÎUK} ßÎ[e6{Ìy„]ívž*c™Ø -•ý[ƒ1Š¶¸Vêô{ ^ÔMô.AÀ~‰Ç°¡ ‘kÇh{ÇÕœQŸë6Ìü !„1!–šŠGÿçi•ò;ØOÁÔ€ûaОâX$"v‰Èú^}zÍŒ¡ˆÌjëè5Ç|‰È­Ä÷)™Í>¦-2 ³~[<)£ÿ÷ä›Ù ºDdý6!]=¸‰ÈÚ¶Û°‘§KDä0kV-" 4tðm'«óŠÈ‡öm©!„ÂW ‚žG4[.âÓð±w"™ Ef!€ ljI–×t @ù®¦³`µûõ½åÄ3e(bq¹Ë4?l! Öt–×  º¾ÊÓY‡S  ù]®Ãßþö·^04_¤ ²%X–ˆ\s]»æ‚шÈ|®‘]C×Læ/q¸DJ×Éuä—ùëz[ÇËŽøt‚þÊ^n¯¯ë¾LD–9¬öé>ÔŸÝÛm÷‘gM¡á<ð ¨öYçÒ–E{.]/íulľ){/}.™`]0®—ºüQçÑ5á»:6×€ˆëøˆïÄç°õ:Æ‘7òÈr͉ì\—PB€÷ÛjEd êr¬íé!„06„Y}/A¸~K?mÀÒ2ƒÏŠ>Rߦ¯å‡°5Ãu¶¯dவ]0À§6ÈH{ÐG·ÓÕ”5Åðe¿žª)#;š(Ü¢-„g6ŒMi} z«¦³Ð6Ç^Ô{*j: 6¤jƒ1̬>P""O—ˆÈaÖ¬RDæÜzDÅフ½²ï""‡BÿE6§ ¦ì0!LP%³…@¦´ó»ZARÀBÈj_¬W™œ ²·‚ ALý_˜Æ@'£¹ÚâODLÁŸ@­™Õ§nâ¦@†Ð%󶩯 "”U¹›euËŒ2°.0®óJàtÞdœ{"«¦÷*PPk9QS°Ë‡"®–XèúÕËò6+"«³æ•t|„\ÙÉ•Ñêz*eÙj‘”WbÎ`ÛwÄ]×ß1´×[°-¨vOiñ˜8ÚŠ×qØ;Nu¬'"¤Ðm´×@ˆAçªî»ª³²¨K^ÆŠÈ >®ù¤+s¿ÚÚŠÈ~[mq¾ çí‹ù±?YvŸóçœØ§vlåþ !„UÂΰžiçûc0[†2á•fÀÇ #kõoìUýš>‘mý¬Á;ý~Q/¬óÄI ÆíÓ'lšþ²^¬§o%øö¥o­ëéÇ«ÝÚài«š—Ÿ®}E ®·/ÖkŸ a ÛëµO˜€¯¦m[ÑC""O—ˆÈaÖ¬BDæXsÀ+Ô™sê›FëÅ%"3œß­tš!„Âÿudv²Ä(¡èXsû+[ÓceAWe¶L8l*qJÀcN?‘ýz<“ð+ð±ÑÄD=š0K#~ rˆ_Ú"ÓÔçz²ˆ€è3ј¨'k‡¨V™5æ€ ’iAõ·Y?Ê*„ÜͲŠºÅDA×T J®)dF¹¾þw\'×Âõ$ÎÖu!ÒúìÚ ¼ÔÖ·ŸˆÈü3שDdb¦ëØŠû܇‚yõÄÐfÂ/ñÛ5÷YÀ\óY–oGwŽêÅ|Ž—€`_–;ËÚ–•îîU¢1›(lŸ„Ùʸ.ÁÀq,q{¬X;dn»ßíJ½`Ð=N ®i2öw~@DvßÖô! {^‰u"²ï‰Ì~K„kç˱V¶¶âë;ÏÎkèÜ-»Ÿι}µ"²cq\~—Îëï;÷”9ŸÕ¯ÍËöBöªž~ÒÒönìLUQØNÏÿÐß±GìUõ‡Äf6_ÂÑóyôã•éÌï`ÔϱúFöغ"ùOö¡%PÃ|öú^b²º ¢²Y5—Oˉ¿ô‡í̓lÿlM„/Uþ?ˆÎ¢øŸýf—œ6Ç6|ÍyºDD³f«"²‘=£ˆœZŽ.ÌÓfRÇj™Q»z±ž @FŽå¶ !„¾ª†‰PÄ2vQ€C¬­ ‚ð&(!fÂ|zD#‚’uÙXOAÐ%ÊY®–jzªˆ|ê\Õà.´åÁì—U[dxV[ˆz‚-ûµ\0Ô¾%€G̪¶ æÚŒŸÍ° !w³¬ªn±K¬ä/Õ;!LgAtºü&%¶²VùVWÛ Î-3° cª}‘á"ìƒ(éšÛ¹Úy%!Ø5ÕõX¯^d;ÓH¸g,7xá>#ŒVf¼ÌsÛÛÖzÄeŸù|¶¯,´šŠƒÀ*˽2°—ïjE6˜m‰Îuÿ¼PoKÇ6œ£Òï£^Ð×Î}¹„ít ª-î{ÇAD@G׎ïêwæ”°¬MΗã–}^bBeHP‡ßM Þ‹ Ø›’Ãõ)œ8RƒA®‡sâºþÖó{S—í¶z¯†ÂV0Щ_Òßé£ýž>qüv™>›Xkå‹ðKêüˆúZë°m­-ãï”°¬oÔ/·SHð—ôõ–Û{Qó«£æ.V¯í ì—!ص}o¹~¹µ?jËÿ²6¢^¬^¬[m³®m¶BDäé9Ìš­ŠÈ:\#Jûök·N[æDÙ&:ýZŸ£Ü®B!|Ua‡‰|ÃÇ?ÙÐEöR EðZ$’YWPD|Z$*Ùƶë=Z/(YÔ–Bf±åÃ9›¡n¢œº·’]S¬JÈÝ «¬»® ¿¨¦µ@͉,Ðu]œ×öíí…¶/ë…Au_öeé ñ¾Ú>ôÑÊ[t\–i£v,Z®m‚^ËjNävz÷Ž§Ì*û¼ênq¯·ÇR8þa»†mÇz粦¼ *È3Sëóu“á9R´uXOߺ‡}?\ÇÿŽÃoÉwÚmÇáÛ¶Õú¾#Û¢šíkx}†u×ÿm©zCaLªýÔ"ô±‹ú.þ;´žFÁªÑ!úNö…hmGQ>P  YÏ+ÊöñP_ù@‹lhµÍ:‹Úv DDž.‘ì٪ˆ¼8ò_ûÚ×ú ¥B!„°J!÷@9u=M› kv3¢»iN®¿þú>û«-²z‰»›ÙçF1íÉÿüÏÿ¬M»16¦˜øú׿ÞgGË4«ì2â…,jçdxž”ʶ>T˜®Ã´&ÃvH´ðˆs!„°""O—ˆÈaÖ YvˆÈ´½†B!,ãP¹ëq(ê&`ÞÿýýœÄ«ÈJ:”˜–ÁcÇ™wøP ›WÒ‚i9jnãBa®DDž.‘ì9"r!„Âr(„Üõ³îB!„eDDž.‘쉈B!„)2¦;fÝ!„Bˈˆ<]""‡Y9„BSdL!w̺C!„–yºDD³&"r!„¦È˜Bî˜u‡B!,#"òt‰ˆfMDäB!L‘1…Ü1ë!„BXFDäé9ÌšˆÈ!„B˜"c ¹cÖB!„°ŒˆÈÓ%"r˜5‘C!„0EÆrǬ;„Ba‘§KDä0k""‡BaŠŒ)äŽYw!„Â2""O—ˆÈaÖDD!„ÂSȳîB!„eDDž.‘쉈B!„)2¦;fÝ!„Bˈˆ<]""‡Y9„BSdL!w̺C!„–yºDD³†ˆ¼sçÎnïÞ½))))))))“)»wïMÈ-yÏž=_jWJJJJJJJʘå³Ï>‹ˆÿüó””””””””É™6¥±Dduû;lWJJJJJJJʘe×®]ÉDž(‘ìɜÈ!„B˜"cÎKùä“îóÏ?.Zƒ/õꫯö¾kL>þøãî•W^évìØ1\´e›ãä˺î!„æ Û÷öÛow}ôÑpQϧŸ~Ú/÷·`'wîÜÙ7,CÅ~m¿ÈbkÞyçîý÷ß_‰VyºDD³f«"òo¼Ñ|òÉÝ 'œÐ:Ç;»þúëûàå`ò÷¿ÿ½ûõ¯Ý=ýôÓýgÅI'Ôxâ‰Ý{ï½7X;„B¾ÀYgÕ=óÌ3ÃE̘Bî˜u¯Šûï¿¿;òÈ#»_þò—û|/ðüþ÷¿ß}÷»ßÝ'PBd>õÔS»ÓO?½ŠÇä‘GéŽ>úèîÊ+¯\ù5yóÍ7{ÿñì³Ï^z>–ñ—¿ü¥ûÕ¯~Õ‹Ñ!„¦Çßþö·Þ?9í´Óú>ÿÚk¯]ëóÙü?þñÝgœÑÛ<ëüáèíÍgŸ}Öýà?X³‡Šõ|þÝï~×o¿gϞnê÷k[ö䩧žZ«Û è\ÐoC;¹ä’K¶,GDž.‘ì٪ˆ,À¸÷Þ{»?ýéO}G¦èœ8∃püæ7¿é;ì°^´# -ŠÑÂB!¶oßÞÝsÏ=}D싈<>÷Ýw_÷­o}«»úê«ûÏ|£ÿûßÝUW]ÕsÌ1}Ì[ë vÄÛ§Ûýë_»oûÛÝW\1\´eˆÈø¯‹²Ç6ÂÏ~ö³îk_ûZ÷Øc …Bš‡¤·›o¾¹{ýõ×»‡~¸ÿ|÷Ýw÷Ë Ì|ƒ¯„Yß³9ÿøÇ?z?àÅ_ìEaÉkü›;ï¼³_ßvøýïßrÊ)Ýã?ÞïŸMðùÃ?ìõ‘óÎ;¯»ôÒKû'^ìëœsÎéE—= ´?""O—ˆÈaÖlUD`<ôÐC}†¬d‚®¶ãŽ;®»ñÆ»çž{nm½?ÿùÏýw·ß~{IJ6ŒüùîÙgŸíÅaûò˜‡Žù·¿ýmwÝu×õÛU†‡ïeÑ‚Œä©_æ³:´G§ ûЙ3öÃ`Ô2|ðÁ~Á¯Î_ÝÖ©GL ÃöÚàøL—B!„ÿb0ùÜsÏíÎ?ÿüîøãˆüÿ òmn¹å–î†nè…]S3@ð'ýç?ÿÙû–óuZ?CÀ*JP þ ¿…`¿¿'¾†"²Œ©o~ó›ÝQGÕ_£3ÏoV><Æ«Ý|-Å ûHñpûÖg+YÐý /ôm”õ¥-|Im·½ó´‘¤‚‘¿óïtÿú׿z_Ïö®‰s"»þ®»îê}ÆšÖƒïìü:bÃ÷¾÷½^PøùÏÞ·î'Ÿ|²?çÞÿ­`à :×ÖÓ{Ž…¿šé4Baõ°ž¾i /¿üòî§?ýiÿ¿J"o¡¯¾øâ‹û§ž‡°—]vYß·[eOÕQ°óìdž°¹mf2;CHÞŠyºDD³f«"rMgA86½„¬¤‘=9Í:QB¯ Å÷‚Îú£>Ú,¨á¼Ûæÿ÷û ‰h+ø©G 9çFôŒà "Õö§\tÑE½ãoyMgá±pª÷Øcí×÷ÿ~ô£>kYð%0ò½ÇRÔí±Oëèôa4R½ŽQû´]{Ì1B!„/ ˜ÀØV>at«¬JÈÝ «ªÛµo¾Dù²µ'Y«üœšŒÿb9?ƒP ®§®² VíÏã±ËŠÈÄì[o½µpÕ7‘ šö[~Éz‘‰Û ¶½ãå[9–ÿøǽË_Àû®æo Úùk˶w̦³àËñ×øzþ·Ü¹“Iαíë\Û¾Dûõ "Ûó^~¤¶i áX&™e>ó5á\:O¶#”?jÛþð‡½oí<×u¯{€X óTó+}ï¼ØÖµ*Q"„Âê  °éúW6Å${R¶Õ ž áÂ:úöš®¢Å@¥þ»)½7€-xþùçûÁQƒœ+)c‹èÚaQ2±""O—ˆÈaÖlUDæx—C/À$èT9÷„V#uærŒM$/‹˜³L6Ç,Žµm®¹æš~¹N\6“V‚8†€Àú:vsÏùlÿ:q¨@G`bÔ¨m¹ì»,b¶¶ÈüT”8-ûǹyb£‹:þ¿ªS@vá…öAE!„ö…mˆüžž" . Žó3ø2Îù²~ù>Ö!’  ó;‘Ï_üâýa•ÝÊ)ѱ}ªkC¹Lm§³¨¤€j ŸŽODä܈ˆ,ƒ˜¨Ê¯âÿñ»ø_êwìÚ¯>B±sÊgÓuj,]Ûó!‰ã¶—Ùn{÷ŸN[lãÙÚ|DˉÇs¾(ŸÔöÚ+`w®eÉ/£Ddׂ¯j[ÙÅê3Å…s亸·Ýv[¿ ‘À¾ù‡~ò“ŸôËxàþ$JÄv]]+×Í>%C8Ç®3_—Híx¬ëè±_bBs…!–ÿ³1‹^°§fÿØYýuK½ƒI^xòŠ ÑÇÓضԓ)¥³€ÔÿÛ§ARë÷} DDž.‘ìY¥ˆÌ' µÓIrša5í„@ˆàë±=ësžu¨Fì8Ú:]S¡S÷ØGÞ£$‚R¾Ìeèܲ| *Y†zJ0.lœ|»GKJð®·}k †qq^JgtF#3«x£j!„07""ÿþ ŸÃ wÁW2PN8•‰Ì¯ñžÊ>å›Ô<Äþ'¯ZDæŸ Eäj‹Çw«-|©Î‰LT=üðÃ÷ñ¹øˆ2ªíÎ#_KàMx%˜ò­øƒÖ#Þ®·=‘ÕÓkü3í+lCÈ•\ ý|M~¶ðל×ò/—]ǚ΂?«.ÔyÖF¾vØ?Ñ—J ·¬¦1Ï´¶ÔœÈ|Ъ[;ù“ücë¸/\»z:Ï÷Äuôþ²¦C!l1lê Oãxj¹FÈ“-b~¶n0„}agÚi1¬§¿§M”F``”ÖQ6L?ÏøÞ £9’=•´•>?"òt‰ˆfͪEdÙ­ˆlßmB­ŽÒ´ó™‡H¶Š ’,BÀa=#u:x™ÄhÎv²‡—‰Èz· •BGnSZTÖ´  ‚#¶¢Ô¥cf.ÚkßöçØ,!„BØ—ˆÈÿ…XJ4¬ìUØŸàS ëÉ+~í“Uü ~•à°DäÊzâ›Èö]µˆlþa>Nûâ:mq-7""Ê«ýΡc• pvLÚMˆ¸ Òë8ù€ëmÏk_¬W"·©9*ÓXf7ÿM[ùš’d}[ß#ÃËænçD®9˜]ŸJ‚à;jƒ ßzêõ×þëþ(ÑŠÈŽ¯bíQ状B¶5_¶üfëÛ§„ŒB«Ã``M}TÐØB6„e};›\ f-töw8ÅQšˆÜN1¥·®þžýfûÚ}Ò@lSSWm†ˆÈÓ%"r˜5SD®ì ±À¤Þ^*˜ÐI­_™:­ˆÌ çŒÛ¶ëÓ¡ëlKDŽj~¹VD6â'cÙ~9ò…€Îw2›e’è܉È*´"²}Ll#ÈÈJ„^„B_E""ÿÂ'ÿ‡€X¾?ÂÀºÁhÓ2øß” Ïa®]~‰tA(q–ßb¾b˜šKÐ+Ð]¥ˆ, ­wD”`\mÙˆˆÌ§ÓNb*‘×ùãÙ'a¸ÄðšZÂñò ÉöÍ·³\@»½óã<ÉÜZODæGÖ£ÂjžJ‚®}™Ñ"ø²Î‡ ßœ–píkº þ% ð­ço½DI{<¶¬-•¸`ÎLŸe­¦õ°A¿s+C¹^ HH¨éDê%Ò!„V{Ò>É6š’|Vö®}9Þzbyøn$6Õ¶ìTÁ®±e¿J+)Ø'ûªyö7CDäé9ÌšUˆÈ‚EŽ¶ŽÑ¾Ì/Ĺ–¬sHê$&œisÔ jUb/±ó£Gù ¿uˆ¼‚ Ûéœcs” +¸‘É!˜ò?'œ¬m>k‹ D ¦3·?Y&@)YG^sÔ BdŒxÑŸà@`Póëy\%„BûBD®9€·Ê*„ÜͲŠº ´DM~…§¯±|ª,TóóAƒ×| ‚2ߦ2ˆ >óµ €×@½ýìïÅz¦ ã³ðZø<ü!>R‰Èü±z÷Dµ¥^àWï½X¡•Ÿ§m|&/¬Ì`ËJ '{i±wOXN 'ŒÛmOˆn·7­‡§Â´Íy,ù–[né¾ñoô»û¬uøÎe›¼•‰¬nÇìØù°ö%¡O.aÙõÓ.S’Þéá:Ùž8Ì¿¶Žu»kçõ8V¢‚¶k#?—¸\ 5¥F!„Õ@[`SèlûË®²Wì=û¡ß7 ÷ÜsÏõ¢²¹ŽÛìaƒ ìaÛ÷ÃölÆ÷1H30PK ®iªØ5Ëè ž–¡™d:‹y9Ìš­ŠÈaÁ†Gü*‹Æcæ‘ÓÉr¤á1s½ùZ'Z‰9àtâ‚BÖo̓¬Ó%8ëœuƵž¶›³H‡md‘ˆ¬C&ë°!ÀU¶‹åØZÇ|yÚZÆ@ÇÎä8/:}Æ¥Ú! ÐŽE“ð‡B_u¿e8ðÍ·á“ñ‰\¯‚¸Ìÿ)Ÿ‡f׳|¼eðÃZŸ‹U™¶%æËùÞþ ÅÖ•,P÷ʲíesh ²%"KB°.Ž‹Pï˜kz >æþ®a¾¤„†šfÃy·¬¥^¸,³ºÚ¢€óäúñM!aÁúöçÚ:þ:V´ltm¬Œgçº2¡C!¬v@B˜>W¬_öÄ ;¨ß¦ °­þ*Ökçé§AxyÑËO Ȳkl€íدöi!}¿>^½ Ñz6|DDž.‘ì٪ˆÌæ+­3­sµ/Êó¿leu»®NݲáôÖ<Ù¦– ׳­Œaèzm±Ž:@í¶‹Öo¿k1J(¸ªŒB!,fh‡7˪„ÜͰʺ þA²õ‹jNd¢°úø*‹ü m°m `ûìó³ ~‹ÀÕ£²öi½EþÕ"_§–iGeGÕ>øZÃzÚºêz¯çsÁ±´¯}—ÿמãõ¶¯c~7ÜõÔüÆÚæZjë°íŠóØîÛ5à{¶ûµ>±š8àš-Ê´¯kRõ¢ŽÇþoëjëzËC!¬ûÒÎ=ßÚD6¥-­ÍñÿЮ a§•¡]‚m-cWá#EDž.‘ì٪ˆB!„p0X¥{ Šºe ›Ò@ÆÒþÓE˜&Âû!d'·Å;&Lýµ™}®‡º†õT]¦ÜXe]«†`Ê m¶ßù[ô¥!2œ;ì°~Š Ê‹2ÑB!„CEDäé9ÌšˆÈ!„B˜"‡BÈ]CQ7aöÁì³ZW‘•¦Gó(´y4Û)@B!„1ˆˆ<]""‡Y9„BSäP¹ë1fÝ!„Bˈˆ<]""‡Y9„BSdL!w̺C!„–yºDD³&"r!„¦È˜Bî˜u‡B!,#"òt‰ˆfMDäB!L‘1…Ü1ë!„BXFDäé9ÌšˆÈ!„B˜"c ¹cÖB!„°ŒˆÈÓ%"r˜5‘C!„0EÆrǬ;„Ba‘§KDä0k""‡BaŠŒ)äŽYw!„Â2""O—ˆÈaÖDD!„ÂSȳîB!„eDDž.‘쉈B!„)2¦;fÝ!„Bˈˆ<]""‡Y9„BSdL!w̺C!„–yºDD³&"r!„¦È˜Bî˜u‡B!,#"òt‰ˆf yçÎÝÞ½{SRRRRRRR&S r%ä–ˆ¼{÷î/µ+%%%%%%%eÌòÙgŸEDž(‘ìÑñlß¾½{ýõ×SRRRRRRR&Sˆ¸Û¶mMDV·6 Û•’’’’’’’2f¡áìرc辄 9̙ȅøüóÏSRRRRRRR&SdÚŒ‰ìï°]))))))))c–]»v%y¢DD³&s"‡BaŠŒ9/ñ˜u‡B!,#s"O—ˆÈaÖDD!„ÂSȳîB!„eDDž.‘쉈B!„)2¦;fÝ!„Bˈˆ<]""‡Y9„BSdL!w̺C!„–yºDD³&"r!„¦È˜Bî˜u‡B!,#"òt‰ˆfMDäB!L‘1…Ü1ë!„BXFDäé9ÌšˆÈ!„B˜"c ¹cÖB!„°ŒˆÈÓ%"r˜5‘C!„0EÆrǬ;„Ba‘§KDä0k""‡BaŠŒ)äŽYw!„Â2""O—ˆÈaÖDD!„ÂSȳîB!„eDDž.‘쉈B!„)2¦;fÝ!„Bˈˆ<]""‡Y9„BSdL!w̺C!„–yºDD³&"r!„¦È˜Bî˜u‡B!,#"òt‰ˆfMDäBac¼óÎ;ÝgŸ}6üºÛµkW÷Ö[ouï¿ÿþ>¢ãîÝ»»?þøKå“O>Ùg=ß½ñÆÝG}´öÝz7ß{ï½îí·ßÞÇ~ïÝ»÷KõT]–;wîìÛúÁtŸþù>ß·U|o=û.Sª Žg¸Î¢su Œ)äŽY÷ªp?þë_ÿêï­!u½ëpüûßÿî^}õÕ}î—1ð;x饗Jpìº:FÇê÷Bá«¿äÍ7ßü’Tðc,g ë}úé§_òBn»žýÚ~‘ÿÃ_â½ûî»Ýž={†‹˜ˆÈÓ%"r˜5‘C!„ýC;묳ºÇ{lŸïŸ~úéî;ßùNwÚi§u'žxb÷“Ÿü¤_?üpwüñÇ÷˪œqÆÝÙgŸÝÊ?üáÝ©§žÚ/;ùä“»[n¹eÝàB`ò£ý¨;å”SúmÎ;ï¼î…^è—ÆN?ýôµ}U]>?ùä“ý:ûÛßúº­§®_þò—}À„k®¹¦;餓ֶ­}]qÅ}ðtî¹ç~ißÖ¿ûî»ûí‰~¾;óÌ3ûbûË.»lÝcÙc ¹cÖ½*î»ï¾î[ßúVwõÕWïó½A .¸ ¿o][nÝ«îaƒ%cò׿þµûö·¿Ýß{«>ÿîw÷±ßÐfƒï?ÿùÏý9}å•W†‹B!L”G}´÷OÊ_ûÕ¯~µæ±ù¿ÿýï÷ñî¸ãŽþ{‚ðÅ_ÜÛÖGòùæ›oî·7(yã7®ùQêyâ‰'Öê~ùå—{ŸÍ~ù_?øÁzf+DDž.‘쉈B!¬Ï믿ÞÝsÏ=Ý÷¿ÿýîè£îüñµe²Iëk‰ZþÏþ³{æ™gºgŸ}¶/^xawÉ%—ô‡ÏÇsLwï½÷öÀ<Ð{ì±½Ø;Dv(Qö{ßû^ŸE¹}ûöîÇ?þq/ Ë*%ÚŸº”矾»êª«úöÉ|!ˆwÜq}ÛÔ%¸!þîw¿ë÷¿mÛ¶}Úª –µÕýâ‹/®íû¹çžëîºë®~¹uá¾ûÝïö™¯Ä5“}.ÊôÙ(c ¹cÖ½*†"²@ØuùÅ/~ÑßwDÿ  ù‚÷ß_ÜOcÒŠÈ«¦‘eho„ŸýìgÝ×¾öµ…¿ÓBÓƒæq 'ô>ÿ Ê>óeà3»øàƒö¶áOúSo‡þþ÷¿÷~€ò¡Ôú„·ß~{oWþñôûggùƒ~øa/öòÝøp|7û20¿ÕÒˆÈÓ%"r˜5[‘$Ç¿üå/}”ŽY€*3K@+€¹îºëº;ï¼sŸÇt´¾»öÚkûå:êZ.x0ògµîÿøǾ³®Œ™BáPÀ¾É!" 8Zù©§žê©ö±Ç_ÿú×½H¼(0°­ìOÅ ØɆ©é®•ý2Ä6MAÔ¦E‘ ™ÓeKe hÚG/ûÛßö©‹¦/¸á†z‘z‘`²hØæâúë¯ï~ó›ß4km1…ÜUÕíšçe+9G­¿Ã“‰î>xZ~ë­·v¯½öÚÚö&ÜÚ×JË'ªŒ÷õŠÈÚðÍo~³;òÈ#û™RCÙ@†Låý¡²òÝC®ûC=´Ïv¦[ÑÆòó ÄȾ/œW÷f¾dýŽJD¾òÊ+û`ݽ¨Èhqÿ;_êwN=><¤Dd™bλó¯~O ¸æ~Ü×Î]Ýûîw×ߪˆ?ýéOûþÁ5Vˆ ŽG[yä‘}¦Ëp­]GýƒåèE¿»B«Ç$çœsÎ>vêòË/ïûñúŸÈ[èÓe볇° ü¼„g» ʲ[ÃîYÆ®°¹lNAtÖž­L«yºDD³f«"²‘:A© Öè›,§£Ž:ª”ãüóÏï?s´127 êôˆõ¸þÐV=öj$ðˆ#Žè÷g¤Nçþ¿ÿû¿}¦×V‹ !„Û%`¯Úé,|'«ÄrÙ›ÿùÏúÌ`ÂÔ¢{ÙÄBû$€ÉVaÿ|fû†­×ÈV6ñð‘HA‰à‡Ý,ÍDì63˜]¾ôÒK¿$”.Ùõõ‚B!­‚â›ýȪ1-ñضQo«r7êê(ÿç”?dP‚ì^ò™Ïd™¿ü!‰§àvØak÷˜øË·Z4xÐ2‘e§kÁ~õµ"²{Ù½¤îýMgÁsozeÎÛN»à~;vìè¾+?Ï'è.?Žøëê¼X×€¡™Àj¿„Þò-ë¼øÁoÁ÷íöêß߼⎷eö{´OÇ¡}!\ßû®Dkâ°å¾—=fÛº^ŽÉ½B v<Úä|X_&šã• á|h«ímçºà·’©BacèÛ=1¦Ïe§ôïl“AKÈõ$UaOW•PÜÂ~YVOíðåØv›pl`“Í.{ç/ÿHa_ùlì]=µ¶Y""O—ˆÈaÖlUD–eäãl!‡Úg£k:U3‡Z 8ía™T%Ȇ̼m:Ú6Ó+„B8”xf("ùþð‡½í#ÌU¦q »Æ–µY'JYšNKXdl$CQ0D˜"رZˆÑö·¨Ž¨eÝA¡«lôvž½n3¢m#»“`H¤$2Ó´m+Y6«r7Ã*ê6íHù?„Hâª9­]oçß Çù$œ .ŸH™ÞOB¤õ+Ð%" ^mcÿËŠÈ…¬(÷G;«æ{´|®=MV½6˜¶Dâ!Wv±,\÷‚ _8nõ ²ùy¶q¸gÝ–„pœoƒîMÛ8G®ƒ'‹Œg/–ÔVí·/çØ.ËZ¿Œ‘ýVe˜ÙÖoÀuÒ>mq ´å¶Ûnë·èÛ7?×½nÞsË+k›ðà:àWۇ맽·ý»¦–,2ðäwîÚnå÷BáÀ`?Ø6Œ¯³hà‘6°³ì[ ÑÖ@bë;±ïì›h0Ó½mÙ¥zú‹`×ØöÂú5(ºY""O—ˆÈaÖ¬BDT·/b@r®+xà0 t 2—d:<‚‘<ƒUg[¯`F–SWíÈ`!„p¨Y&"=ZO”Lmm+ÛÆ–Dmñx¥ïÙMs ËÞõyÑt…ºOð°±-D)mXô&„DEv×C‘ÔñY¶Þ|±+ûoŸ ò?±¬Í`%4²û‹²ª7Ê*„ÜͲŠº_þP{- ¦;'®ƒsMÌ$8VFª×µ%°ÖµZµˆ,3v+"2QõðÃ_›KÄaÇUo¼ççi3?p^bºº ½­ŸYb¶wüDh¾lú¢¦â°­ó&ç[ìÑk]"özSÉ5…ã'æÂú|ÑÁµÃþ ã~_üQí/1ø«-5'2qÙ¶Öwn´‡`ìµOÆ8?¹„wYܦY$^„B8xð‘ôñü1ƒþÞÑÄvÙdöª~¢ÐdzÁíÔ`;Ù}í˶=Ä.Ú†ÀÌÆÈzf/؉­ØˆÈÓ%"r˜5«‘u´t~-õˆ†Ž“Ó,˜dÔ‹‡@^0äíðöÁA/Y°®såÄs¼‡óà…B‡’E"2:òØ8ö®Í01ßk½`¥h°2VZˆNıEŠ„>Ù½öE \´ŽàHý2c†¼ð ½p¨^¢÷bš I΢ÇìÕO€#z·ÈœfOËØÔΡp~ ¬BÈÝ,«¨›?Dl$¤ökš¾²æþ-œ·šÆApX"2a‚F÷ǘ"2QÜqµS³¸þŽË€‚€Úµç¿Éø²¾Ï•‰\祽ÚóÒ¾X¯Î¿sX"²A⻶ öBµ a¿Ÿaf~K‰ÈõÂÉBƘóì÷m{Á½ö›âÂßvÚ—¡ˆìq皢­Úc~ÍšÖhC%H8>O/ä]!„pðñtô0«Ø@¾þ½ô¶Åg~Ù¢'¹Øý<»ÐB§ "·SL±glµ'VØoûmç­k~Ùf‰ˆ<]""‡Ys¨Ed»Ì)¦ ºö!hEdŽv;§{+#u!„ÂVX$"›ûx(øz‰Z;O1AŠ¨$(i…ÙúÞ@j Aª2}­Oœó×úÄ+âÖ¢à¦`ƒÙÔáËÑ CekóqdCÉÀ…û‡¿ä;A¨¬Uâ¦lWhX@:–ˆì¾×¦Ÿÿüçýùq®ÔáX øßòºjžaEÝÎ ?Pöpù 7ÝtS^´Ùol=Yvw=ìž«9)Ýo¦ŸØßýVÓYؾ„yçÞùÖþJZ pªýü%6Ào±Dd‚,sžÛ©eìǵv\S>q=Uç¯s±è·Baõðµ.ºè¢}¾3`YÓm–ï¶là[¿Î6 á ŽÚ¶õ}WS™¢©ÕJ@ж¯­÷; öºje‰üÂz!sî7VËêü»¿ñoô¶ßV%$8¯üMçͱ Ö—áxcù™~»uˆÚ5øRb½v …u/tMœâ°ûßàŽc4XdþsûvNd¿¹Î¦ÿÐF¿¡H0•‰B8ø×'ó¡Ø¡vk: O›HZãÿ”ÅGjEZöÌ6×ÏÚÞ€/[Ád›Ù?}>„ý°FËÊGR_ÙàÍyºDD³f«"²NÔ¨žŽP¦ d‘$*k†ƒ,€ð¢=Ùê´`\g«ÓðZ.èt+S™Î98menÅBa³jÙ.H ˆ2XªHÚƒ,P°ÞÌw÷´AI‘MYY«^"V/êS¯åö%SYa7Õ-ËÄ5ö¶ìoAˆckÙÑÚVQ/A¬²'MC`ÿ뉈ÄmÑ¢Lf¶ZÝÎv²åË2¦7Â*„ÜͲªº»õÔÁ²²e‰¥ÄDßR ®7Q´”Î%‰ßd™@ÔuÒ¾eÈ´%tzº«Åõ%âMk@€àÍ—“5-èÝ2h ¸55„k_¢¶½æ”Ôn™ÉÖu¿U¦¯{–àZçE½õÛò[q£ëü_íËÀ ø“îiÇGìuoû½ìïz9¯êò[&«_û ^VÙâ˜d’Ÿ…f®•u Â8žºN®£óë÷÷ŸÖïÕ±ªÏ¹ßÊcÌ!„6Û` &ý0…1àG§`ÿØ€Ö?R¬×Îýop—´h*1û`ëØ$ûføE…vý~ùŠB‡6ç@‰ˆ<]""‡Y³U¹±­õhcëÈ×:•í$ÐÄT†R»Ái»?Ø×¢ïC!„CEkdžßû†º(û¶h»B@bû¡=.{Úþ¿¨ §ÉXT×p›¶m}‹°l™H§^Y5%€nõ­BÈÝ «®›¿c°½àkNdA©z*˜bÑ´2–êžr¿n‰šÃ"Ó¹®ï¢c^ûö;ûî¯ÝoÝ_þj3ÑyX‡ãúy¾®·è¼Ôñµë.ò-á^SOûÛáÛJ:¶ÝwÚÛîGæ±úÛý:§2Îú®Œ´!öaÛ¡˜àx\«öÅ“…6j«ó5<÷!„>l «ò«Õ!†¥µ ÃÏ‹ðT²È³­e‹ìüfˆˆ<]""‡Y³U9„Bá` `[¥{ ŠºÍÇhŠY»›M!›Vvr[LÑ k|3û„ý÷Yû•å¾Ùý î^dd ´aûeµ™aë!óøë_ÿz?Å…©:†Bq!„06‘§KDä0k""‡BaŠ !w=EݦºðR6Sy-ÊZ ã`žp/W2·å¢§ B!„±‰ˆ<]""‡Y9„BSäP¹ë1fÝ!„Bˈˆ<]""‡Y9„BSdL!w̺C!„–yºDD³&"r!„¦È˜Bî˜u‡B!,#"òt‰ˆfMDäB!L‘1…Ü1ë!„BXFDäé9ÌšˆÈ!„B˜"c ¹cÖB!„°ŒˆÈÓ%"r˜5‘C!„0EÆrǬ;„Ba‘§KDä0k""‡BaŠŒ)äŽYw!„Â2""O—ˆÈaÖDD!„ÂSȳîB!„eDDž.‘쉈B!„)2¦;fÝ!„Bˈˆ<]""‡Y9„BSdL!w̺C!„–yºDD³&"r!„¦È˜BnÕýù矅B!ŒJDäé9Ì"ògŸ}ÖíÙ³'%%%%%%%e2Å ÷Ø"òîÝ»¿Ô®””””””””1˧Ÿ~y¢DD³Fdz}ûöîõ×_OIIIIIII™L!ânÛ¶m4YÝÃ6¥¤¤¤¤¤¤¤Œ]h8;vìº/aDD³F&²G!<®™’’’’’’’2•²wïÞÑ3‘ý¶+%%%%%%%e̲k×®d"O”ˆÈaÖdNäB!L‘)̉ûì³î7ÞèÞ}÷Ý^|,Ø×?üpŸòÑGuüñ>â¤ï|ðÁÚwë±wïÞîwÞéÞ|óÍ>€h¿·ïa]Êž={ÖÖ«¶¾÷Þ{Ý矾Ï÷ö*¾·ÞpßUª Ž{¸Î'Ÿ|²¶ÿÍ0¦;fÝ«Â5~á…º;v¬}çZúþµ×^ëï£ö çú•W^é^~ùåƒâ“ºßŸþùþ¸C!Ì—]»võþÑÛo¿ÝûECøC|9v¡õ¡ ~ ›éïz°%|œõ`Óøt»wï.êñ½:Zq3DDž.‘쉈B!ìAÙgžÙ=öØcû|ÿä“OvguVwÚi§u'œpBwÅW¬‰U?üpwüñÇw§žzj_¬súé§÷ë¿ÿþû½@yÇwt§œrJÿýI'ÔÝ|óÍë‚žK/½´_ßþÎ=÷Ü^믾Úï¿–)öyòÉ'wO<ñD¿Î#<ÒC}ÕUWõ‚6®¹æšîÄOܧ­öuùå—÷Âù9çœó¥}[ÿ®»îê·饗úïÎ8㌵:´µ°”1…Ü1ë^÷Ýw_wÄGtW_}uÿÙ€Àõ×_ß_·cŽ9¦ÿëº×=0ßø¥«æ†nè?üðîÞ{ï.Ú~[Γ{} ¢{!„/cД_Äžðqø í ª~ß±}ìÍyç×mß¾½_Fp¾ýöÛ×ü4Ëo»í¶/ Ñ;wîì.¸à‚îž{îÙçûËÔ¿žÀ{ã7vgŸ}öR¡z#DDž.‘쉈B!¬ãÎ;ï샆£>º{üñÇ×– Hª×^{mïÈ?ýôÓ}àA´1ùÙŒÙ'á€IDATgŸí…Þ*^xᚸúÌ3ÏôbÞý÷ßß Ä=ôPwì±Çv>úèZ…õm'è‘Eúúë¯÷‚5AZPC lëyñÅ{‘XûdäD;î¸î¦›nêëzꩧzÑ›h ÇùÿØ»gˉkýûí¯® &ØØ`lÀ˜œ³Ém²ÉÁ0dL² ¶1`3U68«€*ö[ÝZçíÑÕÞçÌ9›‘Ï·J5s¶¤îV«Õ½ÖÓK­¶¬®“^N‘øÕW_íöqÔî¿ÿþn¿ß@ #4‹&U>›¸—þ&ò_z饳ìïjcъȭÿ.<|pWÛá¹çžëÎ?÷Üsû»B!LóÄß+®¸¢³ÇØ"'œpBg'ÿLÎ_xá…‹3Î8£{óÅøÎf9çœsºý&ÙÙc=öXÉl¼8üðÃ7üöàƒ.Î;ï¼Åa‡¶t>óÌ3ãàUJEÄT¤ G†è+:¹1ƒóüóÏoüfÜU&ÿö!2˜k|5¦rjÚ1ßøëÚú‘60®r¾†¢¢½.zöÙgwâzáøÏ×ŘBîºòv¾v¢®M6¸µl‰{ä^‹×®ì¿õÖ[o¾ùæÆùî7›§–Ònî¹çžM—g(™e¢A{ІL@žœYŽôV"Æ9ÇÊ¢ûWù õ¥¬ÚÀÕW_ÝEZi{mý±9]¯ó ×u%"ÛL¬°ùØÊ߶WÇù­êéµ×^ÛØ·Š‘ÕÙ“O>ÙOP¯å<Ø°þn¯‡«¬ÊBp7#rÌo&`2Æþ²sý]¨ïGy¤Ë˵Ø/¿BëGŸm³¿”g3Ùˆ2ŠÌRNv™rŠDצˆºUv6œv©Œµ”Šv"±cœ/¬:pMµœK-“b¿²×„ŠçCùë|iËÃ$Ïf‘¥Çf”‡ë¶6LÐçð«cAP?®¡ê¦ì[åò¼#Ô½²*GÕ‡ûâ™V„t¿žEÃÉ»™ðBaßa‰&n'¾MVêý&°M2[è¶Ûnë&¾ Ï…±ÊWÍÁ¹Ž-ôítö• á>&O½ÆvèG"³%.¿üò PiGDž/‘ì٩ˆ,2ŠA͈á!=Xþf,ïÞ½»{-—ñ­3e¬›¹³F›Sà•[†¸c8Ù'F7§‡èÌàæ”KS”ˆNža.=¿ß8~š !„v A´/" y¯º§Im$ 8„*Lëàp*½ÆÄM²n&èA±‘ÈØÿ8 1šXVoðôQÂ5‘Œ8ØÂq2~s††/ñ®ˆvÎØM<æt9†˜¼c]BîvXGÞ„FB){…Øê~ÚÝgQªìârE¸›D/ÁS´“¶rå•Wví£&®¢ÙÓF,Ñ.g×R“ÄN©ýì·U 9¼Ú‹r‹¨Õ>”“Sìþ³ýÊA]¹E`kÏ`všˆ_×i¢EzÎa– ëú=?œŸvÚiݽ0q!=¿s¼k96èf¤”¦¼]˜Ó­<ÒSGze±_{ußÜõã¹7y£,žq»DÇè¼iGì–‡çQùÜsÏ€{l‚‰p@Àš !„°^hú`6 Œ&üØ7Æ­K.¹¤ë÷Iýe·ØtÞtf—Õ[@EMÒ÷EdoÞ°èd× ¼Þx1¦?h‘çMDä0kÖ!"WÔH½ÞgQz†9Œk*GÊñd¶u‰D}è¤+:¥œ$b²ÙD<.Ž‚óuä résN ép„ !„ºY%"¹ˆgÖNå$›Ú±•Cak£±k×®îw¯Ä‹Ôä”øÛ8ÚwjŠZ΀Hl<í;!NR…ÜB$Lr®DæôER¢²}}!¼0¶÷#Wk}ç²@d6nouÉ!Ö!än—uäíClì•ö^°‰ˆ”‰\âdÝo¿×›Vꓼ.®Kû©èY6ØÐê[Dg¹—ÚUÝw爼¯% ÉÚqØ„I›µ¹Šâ­Hzí˜Í®ó;Ð1~ƒkÓ]«g¯D[Ž?›ÏŽw ý%Ñú´ËYê€ÝéYp-Äj¶©öJtp¼ç~×R¯:³e+8B9<Äcç»oö³q+2ÛÿÙ³êkÙsBaç/L^/L>ÖĽq@ŸÜNœëã…¾Qè§M2êˇƖ!Ù¸m|¬Éõ‘ké2ã$-¤ÆìŠD6Á¸""O—ˆÈaÖ¬KDfH—ó(ƒÓRQL"Dt”ú²2Ç9,ÙÄ:œ§}lGD–¶·¼ä©Hsè>÷±$L{mëÃ3Á –® LjÄצÙv^N¼çDyëµâ¹Ò(Ù5³KáùqŽe%´;ºvm oQÏÒMf½âöä!JD6iS¸võP Ä„Zúƒ€Àõ6Jì'ƒ-ªlÚ·²ØìSuí9wmòó;›V]8§>DBa½³&[-4™íÄ»ñFß\cˆÉEý´%)Jî3$"› ¤m?LÌ—L šèg#úÍøEÓ°/w´,ÅyºDD³f‹Èfút¤ö×ë%K^p²D«Ø_¯,‚!n0Ð9GD!„0C"²qÌ’í8Êq!@•FŒ$Ú˜ZÁÎï"P-íÔb<«ˆOã£G|4l£fúƒ•SÄj GÃXm©‰þò…%¦8PË–†"»6K´pÀL ·ËgàŒû}Gn_X‡»]Ö‘7Û†ƒÊá¬ûQË4p@E®—ðZ6”úRo~#Ίteó”pÏâœÚ¶*"kcÚP-ë }3¥¯]ÛV‰ÉäWÂ6ˆ°Ú ܇]‡g¡ÖÁvµœ;ŽHìx¯ñÂ9ÒÓž]û2™ýçù’(²ZÚsáõ`õÒNÎ áºk©˜úˆ¤hiå#rC¾µ^²k1‘RÏ ±ß=sOô& +o l]ýÂõH×=V¯ŽW~B¶· B!¬ý,Û‰llë#P­µÍP“þƵÚ_¶Ë‘kt ËÙˆÉÆLea/I›íe\ô/{ÐøaÚ‰yºDD³f"r}¨¤]΂¡\_igX×Á‰Äƒr€8Gfáj­:Ñ9:W†}uТDïίù°žó Æz?Ú'„BXÄ1ã]û*¤×+º„£`|%*2±Ä*aª½ ÑÄ¢!-@t%úëÚ×è}é[DŒýÄ-Nˆe8?Öü#:Ö+‘ÏŠ$î‹Ÿå¦sœoRV”L ˆ!y,‹–¬å6ú¯`Šì!zˉƒÎ'Ší»/¬CÈÝ.ëÈ»–a ~Vä¬öRQ®Ú‹ú¶áM¢‹”j'Å “þõÊ!­µ‰¥±Ù‡õL°—ˆÈì5v—sÙlÒqÏäç^±¿–¡.JÌVN²gA¹Dš‹4VékçÚZ-5Q_¶g:Þó!-y—8®p¨•«|Ï“ºb?jŸœpõä8õ¤KOÔðfËqXnÂ=píγ„…´¤_hîs-·Ææ´r!o¿¹Fϧ½"—-íÂÞU/x`gÃÖ„“ýòö&qÃõ¶k‰‡BXúlý·‰<}vÙH– cgØ$l*vŠýÆ!|&%ÙmúéÖ¾b#µ¢3ˆÈÒ±\Õ2ŒÆÀeË‚ÉßضjÜÝ ‘§KDä0kv*"…uÆœ¤ú°‰Î[§\‘ÆœKF7CÛñà`:Oä‰Q_ëꉰòšÜ>N“×BD£ Òo?ú£#×/‹ž !„v‚HRãV<Ž‡‚8f#صkÓˆtœ‘>¬Wö‰z6âT‰bœ"i‘åk”–ñÔÆá‘÷ž={6Ò#VÖø[ȈyŽ'vÕùÒó{E\*cE†¡l®È[Lë•×uˆš^ö:èVY‡»]Ö•·:`Ç°[¯DÏjC&$8½l(Ë ÔÇ EïVä2J]{ÙEö™DwŸÊnZ†Hh¶‘ v“αµm@›Ð®û‘ë}£R÷V9ÇæROêH›Svû,EfÝbé{ìgk²Õ´uBpµu¿ËÛ„z¨hvö©²jSÕÞ°Ò–‡ºP/eW®¢ÞpàZ«üý5ÊÙ¯µÔGû (Ÿ‰"Ï‹k4yDHð ³ß\ßÕs=»‚”+?ǹþþ²/!„vŽ IãK;¾ù›ÆP‘ɵ¬¦ñÀd|¢_ß,=Ôž[6RÿzúpÂõVÌò1ž-³L˯ÞÜÙ.‘§KDä0kv*"ƒsÐ:Y"FüݾÙ?ösŽlý×(«Sogè*¡ô‡~ !„ÖɲqÆïƲ!!®Æ§UpJ¼ŠÙ_ã¸=·þ?´µ¬*㲭ج¬ËÒná ÕÃvX—»Ö·((vM»”H­‰ÌÉ•›gè5\u^ŽnýíxiZ{h#tÖýuüªö£L&úiÔÖ:ºlFËžôÛ*×^c¿}ÁyÎïG»ì7ËW´ö¡ùdñÅ_dË–-[¶lÙ²Mfã -"›hï—+[¶lÙ²eË–mÌíã?Žˆêÿ¼øøã;CþÿøÇâË/¿Üøÿ¯ýkñïÿ{¯í?ÿùÏ^â¤cÞyçîßÍøâ‹/ï½÷ÞbÏž=‹O?ýt¯ß¥=”ß矾qœò+ëûï¿¿W\ÃйŽ7î*s¿¿?ûì³î|¢kS›±{'Œ)䎙÷ºøàƒ:Û‰­WhŸÚño¼ÑíŸÊô—¿üeñÚk¯}-6©gå•W^™Ìõ†BX ƒ2Ä?ÿùÏn|vLÁÎéÛ(µcŒ-CûË[fµv »@™Œ©Ž_†1§l£å`¿9FDäé9ÌšoŠˆÌø?øàƒ]tQן~ú鋃:¨û=„Bøº!@wÜq‹ßýîw{ýþüóÏ/~úÓŸ.Ž9æ˜Å‘GÙS%V=ñÄ‹üà‹ÿøÇÝö“Ÿü¤;N:œÉwܱ8úè£Ç{ì⨣ŽZÜtÓMKÇeŽÉyç·‘Ö‰'ž¸1¾ùæ›ÝoÒjóúÑ~´xá…ºc”GÞòòûå—_¾øïÿÛí»úê«»ü«¬Î•–ëá,tÒIiÛ¤ñÃþpqï½÷vçëÚ*óÏ?¿··Ë˜Bî˜y¯‹ßþö·‹ï~÷»‹«®ºªû[»¼ä’KGqÄâðÃïÚÀ­·Þº´½í/´¯j»­à½.n¼ñÆÎf|ðÁû»¶„gë¿øE×ÖM¨„Bøú0þþüç?_\wÝu{ýn<þÍo~ÓÙZl c†~]¿,¸¬µaŒ)Žqì[o½µxæ™gþÏ~û0lœplßbsÉ— |Ùe—mìsÞÓO?½Wù àM>¯¿þú^¿£/ùþìg?ëʼ""O—ˆÈaÖ|SDd,GÈ€b`á´sŽ–ÍP†Bë€Cp÷Ýw/N=õÔÅ÷¾÷½Å³Ï>»±ÏÄ¡¸á†ºèN‚.ç€3B±ñëÏþs·ýéOZœy晋 .¸ W_zé¥Å÷¿ÿýÅ£>ºøûßÿÞ‰¼¾!ÇÄñœ™ÓN;­+ÇáÒK/í-¢’+›ÈNÂgG9ß~ûí.m¢¡¼^~ùåÎѹùæ›»ôEB+_ÿÜsÏu"1‡Íä­ÈÕÚO0¾ÿþû;g¨œ ŽÜÉ'ŸÜ nÊÇi“æND·1…Ü1ó^ì¤C9¤› p=þõ÷9çœÓµQb¿¿µ»1!"k§ž%%ëæW¿úUˆ >¶ƒg zî¹çöw…BX„Ú‡zhqñÅ/;ì° [ªØµkWgw<õÔScþÐCíìoMµ6 ;E:&ÛEK»ÝocOé׉±ÆžÖ^3ùΞ»í¶Û:;æÚk¯íÆ)é²÷®¿þúÎF*×ïÊCf+²™ vÄ/ùËn2Þqò’ï 'œ0±¼U""O—ˆÈaÖìTDfø‹Ìxä‘G:‡õ׿þu7kHä5ãÆÉÔirBuÞ…NÓ¬C6@üþ÷¿ßËÑtî}÷Ý×+ýǼ$JDæðÜu×]Ý«Ãu¼AÇñœfN2'€è1ƒƒò]sÍ5]¿ÉÎa!„¯Ÿ'Ÿ|²~ oœ—VD6v‰ti—¸0®UäJâ0‡¦Æ®‡~¸rk ã §åÎ;ïlOëð $Á¹¢ŠÁIQ¦ÖY) ÇüF¤²ñ˜H׎ùdoö´KpÆJQ«íRÁZH³àPëÖɘBîºòv¾{ÆÖa£ÜsÏ=Ë–HÿöÛoïDJQRö‹Do#˜ìkm÷Š½£x¥w%"KWuÿµ—r:¥Cí;êËÐΔEzþå¸ÊåÙÐD>—m§ “ œlç³ë:KD¶ýáèlF‚÷<°W{e뙬`é'‚ÀV(ÙdÐc=ÖåÏ.dzæØ°þv? 6¢{£¼žoü] ê±®Çýhy;êWYÙ¥®©žûBÿã[«Þ¸2A_°EÎ8㌮Ï-Œ/ú•XË2±=]€ˆË¶Âuá…vcñ“mÖ¾ÍâÍケÆö³Î:«ÛLØ·vѹt‡óGí²dûJDäé9ÌšŠÈ:E¯z˜‰ó/gØÌ¡E$é0u¤œSœ ΂×c‰ÂûÍØ1îÁð¥åœöµÎ3YŽO­‹lÀñ·ˆ,ùJc_³Žœ‡J§Ê(ß,‡Ba+»Œ+ír~3ŽÖzzœN‹IÐ>ˈ )8L¢A‰M ¦þ… h&PÛõù*Z¸ïDד½â¯ýëâÅ_ÜkÂÖ¸l¼í ¥œ"c¥1~ΕžZª‚#Çéãp]yå•Ý¿D³¡õ£÷…u ¹Ûay«kb+›„³ÉNbÛˆFr¿µ%¶º®W\ýíÿ"Åáõ×ï|ç;]Û€¨%¶N?Òiˆ6™ÎâX³›Àl*ùm%Ùˆ¶f ¶;J9ÙQ®Ó„„¼´qQûul9ØÖfv}Îó¯}uïšj‰•Z‚ÅõIO;Á\›V^çWrÂ7ƒˆ\åv¤/íSN9¥³M˜pÀ‚ÇÜdòiç%¾»OžoQöpeUŽº.Ï®çL}Òý&÷[~þŸu™CasØ'&å “~únvOMÐ 4²3ØcúÛ?ú˜4,ûL€êókbјàØö hßg2®ÛZØMÆ¿þd°òƒMNº.éíľ@Däé9ÌšŠÈœbF3#šÓìõ k êT‰¹:_Ñ5%3ÖÍâ|Enx½„C# Ç _Æ='ÇoRÆx‰ÈDcyŠ~æ„;–£Å1â(1ú•æ˜r68ï:wéqD8Â!„Âf ‰ÈCÞ¸dÜ!ðõÅW¢‡Gôn»F0¡’€e¢ÓhlôúåVÖ&dÉ˺Æm´'8'$ãçN˜ÉÙþõG‰^"=‡0ærÚˆh‘4„4N§®Ö,¬(ží²!w»¬#oN$ÛCÝpÝ÷—ý!zJ*Õ±U„²z׎D3±gˆòÚ…ˆx\Õ³s–EX%"‹„m©dù°Ÿ6‹h椋ÂUVo‘B‰Æl+íÏý7¡ÌDTdåáz\‡¶ÏŽ«j6£ôüÎ6ôlI]Voù›=ç^HGz~7‘"bY°7û %¹]í×óyöÙgwé© v¢gATšgÙ}S6uãy T+‹‰»m÷VôµuŽ'‡¿Ù¼ž÷Õ9„ÏÃfÂ!„ÿ+"{“EŸ¯Ÿ¦5°©Œƈömgˆ 66Ðú°ÇL°;oÈ>1žéûW½¡cŒò—ñ«ÿ¼!Ù[gl2ã¨1Ó¸nŒq }ûm_ˆˆ<]""‡Y³¹>TÑœG†yEtèÀušŒmÇ‚ëゃÅÑ‘ÂQh£O`F°ÖDnEäz-„CB|æ\xíÐL¥ÎÙyœ›Î^g‹²û ö‡BC¬‘ ¢Rˆ‚¢0 T5Þ œŠ6 „DN‡Ç|-œ–­%L`6†w9Ymd2ŒéDà6 ¹…ócÒUyˆŒý|Œ¥ö-{ÅSº¬Vèöãq+æ‰z6î»®í²!w»¬#oBißÖ 0j+êJ[bË´kíúÉ.õD^·ˆÌö³‰¶Ê3©¿ê:ÙQî¥û^KŸhßÙr`9ÞÖøöV™Èç 0 {ÕX{õ[EIG a@Õë“# Fk‡®Õ³G´e×ÕÇ–êC„[ù®¢–³¨@³Ç6äÌ»6¢òj¯áZþ„bûåðlTp„rÔÒ/êÒ}+[·¢ÊÙ»"ÒݯþóBáÿÒ‘&>õÙeøMK¤-Œ1Æš¡eÁÀ2¶Ô˜ÚÇÄŸñ®]Þ¨0N·Wä!;iHD¶4“±˜W˜ì7®Ô˜·""O—ˆÈaÖ¬KDfH—ˆ¬ÓÖQV”¯ÙA³†:d+cš¡mPT†íŠ+®è6Ž•áÝ®µØ~X¯‘E1ëœÃaWfm†Ý¥l–r ¸xE,W™}íÛε•ˆìšÙ¥àX;‡=§¾¼Ý¦¬Ê /íÔ³#ª«uÖ‡(y×®]¿©OuÏ®Á¸–þ‡¼kÍçûå‹.´o¿¹Vé(O½e§ND0ûݵ© ÂB–Q !„Íé‹ÈÞvn—H‚qÈ"íqõ†Ï2q–=Öê-Ò“/¡!lì5–êûËÖ2‘ÙAÎk?[k*Ó*¶KDäé9Ìšý-"‹F!òÚ_:l³·ÜrK—F½fÙ¾+¸>¬×_ÎÂìç€ÓRéqÌ 4‘C!¬ƒ!Ù8f³G;í:ůù›(m…Y¿—úQ¢D·Šô%ØÝœgìAl|r~ oñ(gÿ‹ßÆ{¢—Ò,F½ÕÃÑi£eZD»¶¾mœ%òµËxpàŒûÈí²!w»¬#ooWÝ3g­Ì>1QnDή¨©ZkQûQo&œÃieÕÛYêU®m_DdË(( AÓ}ÒžÊí2Ëp¼üZ‘—S¬-XžÃR®C™ëÕ^×YËYøM›t|} Ï$G­Ÿ¬=/‘E"{¾8ö®¥Ö@ö|(¿zze¹…ˆ\6dM¼T½²=áyUeT'„áˆýlÎZŽ¦½+ZM¹Éuî^²+ò™Pn¢€Ý¹lÎBÿ?}Ù8b¬j#Œýfr²Ø6(¨lèƒÁÆs}»IÈ!êM öz0v+ŒqÞÀYňl<3®´ƒ5~³Äm—ˆÈÓ%"r˜5ë‘EDµËYpP:è µ„ˆÈ:]8ÃZçËPçpžÖµV'Ýk&*kE)s¨ãoƒ‡œC ãõ±Ž€h»¨y9ß+ˆŽa¸{½¸®•‘ï#*œˆBa3ŒOÆ;‚RALå,”HgÌ1Þyã¦NŽñnèÕJÑÄÆ*´8.Æ/Ë[€`M¸2.–B”#̊ƱvK¸ã qŽ†>–ÇRVåwŽs‰˜ÆÈ·K l¿ ÞBøt-ý¢šë66‹´ôV!¼ÄÓí²!w»¬#o÷¥ìœâZY4«{àïZ·Öxl—ôÒüm€=ÄaÖØE›­¯k2€=f’ÝdÁWZÒ`O)‡²µË< ¡.Lº³µL°Ý8ÚÎÑ˖Զ؂¢®µÑZj¢ì,m\ÞÚI-w!½ŠèªÞ‰Z†ç¡„i÷@½¿'Måw~kÛ-Cð‚²ª7癌‘–:,G_µ$…g°>díÙoŽõæz+r¹–Žqÿ<ðÀÎvß]ŸýêͲ&õq½v-ñBï¼ÙÑ¢/5Ö°ôÝÆ6cM&ùŒ¹úø!Ø*&þÛ·ƒZØYÆ•¾ D¿0^*Ok±ùúë1‘#í[jÊåzh&ŠgÆ eíOøï ‘§KDä0kv*"ëTÔŒøêDE†˜),'ØkˆœZ³‚iÃQåds–Ú«¶äÔr¸ì¥Âùõ/G€¡ÏáÈx=“±îoƒ‚ã  ;'€h­Ó®¯“ƒcÎ1IDH!„­ÀÐ7Þ1þ[8„^ãžMäLû¡2oàS ·}j}cã™qʸeì¬qÙæw"2!ËqÖn5žÚŒ‹Ò®`éÉú¢ q“Ðåø:×&="WE\*£ô–}hpfŒŠð!ŠëÕ2sîÚ/™o‡u¹Ûe]y«öá–âMªjC}ΦöÃ9µ_ý±yʱäôªK6ŽÍ>b²ûTŽó2¤ï^X¢¬´¿ÉK[òfÖв(}ØZRçŠÂÕ–œ«žÔ‘6§ìÒå‹ÄÕÆj½emÚ›cgçW[g—¹V6[²Ö˜dºFmªÚ›ö)my¨ böª¨ü¢Ê¦<ÃU~‘â-®Ç¤ŽôÛg@ùØŸõœ¸ÕAÃoõQIEõì:FÙÙ¸òsœëzvB!ì7Hú~º~—̇קOÛè^o}˜DßÕ,]Ôb<Ö÷/{CÊ2ÆÅvÙ °oÊÞjm(ý~2WÒèG,{cF [Àyþ¿ìãÇ[%"òt‰ˆfÍNEdˆ`ê¯Øÿ­ÿwQká áxt97mCéÕÚ|ÅfÇýB!,cÙ˜áwãÏÐ’[kKÎóúãØÐ6t|Ÿþ9Cç÷ÿî³j_Alª‡í°.!w;¬;oB$›¦M¯ÖD¶Ì‚ºí¾,*©ÖÙ-êxtEű•×Ð=uLü— i¿IyçõÓ²µ¶Ñ×¹CmÕ½WκFéöëOžÎïÛžCå\ö›ztíq|¨ì~ãü·éÔùíÚÌcËdn±Ü~+jíï¶Lu=ËÚ¼ºS¾6¿B«êÿ ãd;´,;§XµYžõûÐ6IJ߱jœßW""O—ˆÈaÖ¬CD!„BX7ër÷…ý‘·ïAXaè#>[AÔ®µ‡û›(`Q¾[M“j©çõÓ²4E?*kjX[r¨ì"Ù¼ö¼Ê¡‡(ËžyY„óîÝ»û‡„B“""òt‰ˆfMDäB!L‘ý!ä.cäm©K*ˆ˜ÝLè _"“kËZN#„B˜2‘§KDä0k""‡BaŠì!wcæB!„°ŠˆÈÓ%"r˜5‘C!„0EÆrÇÌ;„Ba‘§KDä0k""‡BaŠŒ)䎙w!„Â*""O—ˆÈaÖDD!„ÂSÈ3ïB!„UDDž.‘쉈B!„)2¦;fÞ!„B«ˆˆ<]""‡Y9„BSdL!w̼C!„VyºDD³&"r!„¦È˜Bî˜y‡B!¬""òt‰ˆfMDäB!L‘1…Ü1ó!„BXEDäé9ÌšˆÈ!„B˜"c ¹cæB!„°ŠˆÈÓ%"r˜5‘C!„0EÆrÇÌ;„Ba‘§KDä0kˆÈŸþyÿçB!„Q!àŽ%ä–ˆüÕW_õw…B!ŒŠ@ÀˆÈÓ$"r˜5Dä?þ¸’³eË–-[¶lÙ¦²‰²[Dæ¤õË•-[¶lÙ²eË6æöÑGEDž(‘ìÑñ¼ýöÛ‹wÞy'[¶lÙ²eË–mRÛîÝ»G‘åÝ/O¶lÙ²eË–-ÛØ gÏž=}ó%L€ˆÈaÖdMäB!L‘1×%3ïB!„UdMäé9ÌšˆÈ!„B˜"c ¹cæB!„°ŠˆÈÓ%"r˜5‘C!„0EÆrÇÌ;„Ba‘§KDä0k""‡BaŠŒ)䎙w!„Â*""O—ˆÈaÖDD!„ÂSÈ3ïB!„UDDž.‘쉈B!„)2¦;fÞ!„B«ˆˆ<]""‡Y9„BSdL!w̼C!„VyºDD³&"r!„¦È˜Bî˜y‡B!¬""òt‰ˆfMDäB!L‘1…Ü1ó!„BXEDäé9ÌšˆÈ!„B˜"c ¹cæB!„°ŠˆÈÓ%"r˜5‘C!„0EÆrÇÌ;„Ba‘§KDä0k""‡BaŠŒ)䎙w!„Â*""O—ˆÈaÖDD!„ÂSÈ3ïB!„UDDž.‘쉈B!„)2¦;fÞ!„B«ˆˆ<]""‡Y3†ˆüÁ,^zé¥.ïBá›Âž={~øaÿçî·wÞygñÞ{ï-¾øâ‹ß?ûì³nÌûç?ÿ¹×ö¯ýk/qÒ1»wïîþÝŒÏ?ÿ¼+‡ü>ùä“ßå+í¡üœSü÷¿ÿíÎýûßÿÞ ¥…k:×ï_}õUWæþ~úé§ÝùÒªcjûÏþÓ»]ÆrÇÌû›†6ðꫯ.þò—¿l©¾´_v 6”ÑýéOÚ«]úÇ×^{mñå—_nüöÊ+¯,~ÿûß/Þÿýîÿ/¿üòÆöç?ÿyã™á»=ðÀ]›áØ_|qqçw.î½÷ÞÅÿøǯåzú<û쳋ßþö·×á_ýDýý‡?ü¡Gôõÿ!Œ«ú ûÿüóÏ÷ÙÀ8¨žÚqÚ8Ìç$¶ݶ²q¶Ø,Í>ú¨?¤1T×ÿþ÷¿»ýîs;F¾ûî»]?ÚŽÍð·ßëþ« í²ïÃë‹Õíßþö·½~o‘–:ºçž{wÝu×â…^Øë•W»lÛž¼Ü‡*«6ëåïóñÇw÷©lƒe8îw¿ûÝâöÛoïÚrÿYX†²ªûeu» öÉàêw/ìÂïO<ñÄâŽ;îX<øàƒ‹7ß|s¯û£ ž©¡óÿñlùº¶CDäé9Ìš1DdÅÿûÿoqå•Wöw…B“„³tÌ1Ç,žy晽~î¹çÇw\·ï‡?üáâ‚ .Øp¬8GqÄâÇ?þñÆö“Ÿüdqì±ÇvGä¶Ûn[}ôÑÝoGyäâW¿úÕÒqÙ˜}Î9çl¤s 'tÎ586~“V›×QGµát?öØc]>Êû£ýhq饗nˆW_}u—{®c.¼ðÂÎÙ;ñÄ÷JÛõþà?èD pJëÚª>Î;ï¼½Dõ}eL!w̼¿i¨'mÿøã_Úv[ˆx`'j…0evíÚÕõ{CbYB‘¶M°Ý Æ’Ÿþô§àc‚>èD¬úÞ÷¾·øÙÏ~Öõ½6ýªçLòÒgßtÓM½T÷†ÀtÖYguãÀ©§žº8ýôÓ»~^¿>$z­‹·ß~»+ó¹çž»q Æc†k®¿ æ7ß|óâ ƒZ\sÍ5ýdº~øâ‹/^|ðÁÿ¨¯‘qn×ôýïã^O9å”nùä®=H礓NÚž ÔÊ^×®MºWl™n¸¡Ó¥£ü„è>êü°Ã[¼þúëý]¦Õ!䡇êêD[Qþª«!L@(OÕýÙgŸÝ ý›AÐ>óÌ3;û¬ þþú׿îîk‘Q½ó•³êB])³ºS‡%’k‹žKö^×UöÞ×ADäé9Ìš1Dd=CÀ¬]!„0eÞzë­Åo~ó›ÎyàÔ‹®*8Â/gθÆÙàhA´‡“SeóŽÌE]Ô9b¢¯¤É!”ÖSO=Õ9jþíãxÎ(G¦"œ.»ì²Î!çprh*Ûo¼Ñ9:˜Št–6'” ÂÉäå7‰)_㌿ý®LÆÏC=´»–~4rõ-®Ùƒ+®¸¢ž‡0N£ˆœ%tºNõ¨?7þ+—E°Š0& —8é<"61Õ8«ÌÒ¼üòË»¨ai:^šÆ oúh/q÷ÇXªþ•®Ý˜¯ïÓ¦L$½~øán¿ûK$‘•AÁR>%"»VOãóêšMà˜|® ^u|Úi§uåÔfØ-ìc|Ý'e%»OÎu±»&‘[Ô¿k\&"k5 ¾µMõ,íÆ=ªg¢EÛpÍ&O\7[ƒÍÁžÚlœ&Î+;‘º`w)§ºt}®×ß%À+HÐ[=ÿ&oØ€îÙ­·ÞÚýÆvÓFŸ|òÉ´ i²Ñ""ûˆˆfÍNEdüÝwß݉ÂIŽ³S+ºÉ€c@0(VªÓ×™×,±s9è~wœÎZtÊÐ+Ã!„Âþ„ Ë¡äÜM[™Ä¡äh×^{íâüóÏtj8ÁrHyä‘Îù¨W9ßœ¤!qÃÊqåTp•iÈi$s^jÂVDŒ´Û1Ÿ#ÎyzÕÒuˆŠ$æàrœË©Çõ×_¿!H¯‹1…Üuå-Ñêê‡ïÞÖ²%„vŽ(1I;³Ÿ Bh(ìmÈ¡…{Å1%&l¶ü‰v᛬xúé§7"‰#Ò»ñÆ»ß8Óm "ì#9·mW¢‡ýÊO!L´"²6îZ]“2´mt«"²´Ø˜ÊbCÙÙ”Dºörm\½j~ùË_ve«ýžWeT§ÎSfù²QM”8Þ5ö”W~u϶*"†yQ"²gM›Òù+úSþ F[²O;3/ø?õÌ´þ‘g©žµ‘õ§Ä4ÿ¯hFç÷Çð¥Dyj×Æ ãÒ²çHÙM-m`ì!|–¯ÅŸÓ¯hïú…öyP~}„1è–[né¢3A Ô¯y†ôõŠ¿º ÛFïB?Hœì‹^òSÇDIBs‹k#8ËÔ9ôgþ&Øöq\a]¹‰€m”°~À8:4f‚_Ü ‰P^¿¹nçKÓØPÈGšúT×N¬¬ÉU¨#b>Œ‘l€¶?%Hk3Ð7‘µ‡¢íˆ,?÷N„­öRãÄ&”sHì$üd]—zÕŸ»ç-Ú©{¡múÿvDdç±9´M}*;Âõ°ÙMêÊdùPT/Ý@{n£¯µc×½jY÷Mž&í…ÉöoÏ¥ò°}j"D-Ëá¹w®ûè°ã†&ÿÙHÕVk ó ÙhþÝÉòž‘§KDä0kv*"3(TŽÍl±ÎÔ€à_¯}rÈ!‹K.¹¤ë uøßùÎw:C¢»D&ià_x!„06å@qüÛßDÐ9㜧ÖÙm1^#‰çÖñœ^Ñ,+s6úHŸ“ÒF茿ý×?ë"â*"„Îw;®—9B}¡TºÆïÖYk!ØI¿aâ1ƒ#Ë!æ rv:¼.!w;¬#ouÍId±…+웚HЖ&lÚ{È=u¯Àñe3•CO\ÐF8éË„‚K¬•/ûJ¤O¬±V6›ß*}åb—)‘B¹ìg³±ál„&û¥'mE™‰ö©ý_Ú®IY¥S [‘ heG*‹k–:,1… ¢|Ž!8NÙ*R‘àøZŠ…©\®ÑõKßñþ5ñϨë©ë³ŸÀÒÎÂü)Yd‰kD+l/íšUË8–ÿC,ÔwhŸÚ–vèxÿ×êJD®hLi—àY"r‘êÙRžVÕÏj÷Cý”ò觗ùR™ìº¤áYU6×£íó×`"GÙ<Óú“§Æ"×"“0ê¡–Öp]C"¯qÅsØÍÔ‹å#œO@¬k!yN]cáMHWž³[L€yví#;ÎØ[è›<ÏËÖUÖôEf"!ñÒä”(\i¶c£qBšDN÷ȘZ‘Êlý½1úsõÛ¢~cÌTEœÕ×Fk#ú½‘Ý/é‹ÄÕ·õëµO-­ Ž†¨ö ­»Æ¾Ø¬nÔƒ6¸ÝHd ®UÛ`3¨ý¼cÕ¡¶e\Цú¨Om¨àõóC¢{áZågEß¼Yÿ¬Œ+ò(Å¿ÒsÏ乕·šØbêx¨ ¢‰k­mí Dd“3Ƶþ„¿:¡?ôæ—öibÞ©|êų`’¡&eÕO{ßÜãu9 mÚ=3FÈÓ%"r˜5ë‘u’²z­Rço0­(3q9½c†DdƒM½:Åùf4”ÑB!ŒÍ*™ÓÍñ%¬· íëöÆ5ÎD… Î#Q‚#Ãñª¿9á}¼ ,:ÞqĺþÚÆtNK…ÜÂ9"d:ŸÙÏGí[¶Î&'œß:¢þOi… uaÜo_ëÝWÖ!än—uä͹eïpL N§¶Â6Ò–Ø?µ*8›e,ˆ¹;‘‰¢Ÿà>ºõª¹ûMè0yQëržÛûÆùöá+×CPRžö#X®Gy”‹GÔò7Œˆk™ yj^BɾˆÈlÌzM›ó]Ž¹ë©5P•MÛô{ ÖöË_^D:°{•`Q¯)C#|”ðOìqDu䜡å]Â|éG"k{í‡ówÚ HŸ¡ï/Q«Úªç€ÏcÓ¦<“ž;éÔ0ž…jÓ Îi§Ú´ÝZÊ…Peâ¤úmѶCQ¨ú-bmíØ5ÉÛx Œ-žSý€çÁ ƾŠ,vÍ®K¿Ò¿.ªI—Zž¡e•ˆ\“Bê¯>çÙU6ý’7ª^Õ31P>}JD®79œ£?¾­?QlŒPªïóoERëkõòÔïôi& àÞ,K³|Xý¶ôƒú"×¥?s=îc‹tMÀ¹DbíÉä.!¼6»|l(£>IùôþUÞ~ôw‹ë#œ¶¶I¡½höƒä¥L"Ãmú@õ]¿ql;"²Iãò›8(Ø4&µ¿·âm×Ⱦ©o<,ë“åáÞU4¿º«땈Ü µeåcÿH¿ý2Âù"ÉÕwgêK›S¿í=³?J/Œë&*ÕçNì DDž.‘ìY—ˆlÀ,Ùhƶf è¢ÍDäzƒƒkˆBa* ‰Èª¾HÊ‘4ÞµëPÞŒ“­ƒWëïõJ§chl6žrž‘ýul å‘Wp¬­Ò(g°…s^ÂC_\ᓳäµÐN\ß1%Ð9vhmí²!w»¬#oöqÒ¤@aÀ}äÄrP „…B½¹¿Ú ç°D䊜SÏ„§­ŠÈ챊|Ó¦8ÕÚG­»ÜŠÈÊDð µi{µÞup‰ÊÓFŠNˆ(D0eÚ˜òqÆ¥ëú”(Ä9ßYZ­HFL.ÁZ]ŠPfCª³­•¿/"×N‰ÈµŒJD&rüÕ<•Ý³@,s8ýCÏe˜/}YÛÑÇ„²e"r­ÉkÅó¢šD0jÇ„ÁVDæ+• ¶lMd}¯g­&ùLégÛ5z â¬4†Öt•·r{ˆbýèXBYúõÑÔê -®K”sÿº,ÇAøÞW¹Ö ÖgÕÒ1ê]´¶:é‹Èļvr®è‹ÈP×"` ºÊçº7Ê¡V÷þm׶ÏØ¥Ìúa¢_+(¶iÖ$qªÎÛ7iÕWù¹ê¹„ÌÂø@ÐT¿ê°ú¶mÄ1írò¬É5ý5ÑSúËÆ,åÕ·µkD|oûLJ(»ÉŠþúÔ-%ÂÖÇåZ,%¢,CË4¨WקM¹¿ì)õ*omHûö¼,†]¯çDÿÞ®[ÝGý÷ÙRÆõâ6œvãyÕFkâ ·ûcœu}¥Uô!r«í”@ì™lÛG¡=·Öó,äÙ`}›i;DDž.‘쉈B!lΈ\Ž~ë ˆ.ã`—a_ÂB-PøÈÖ‘95õÚ¿cŒÑÎ#ˆÉDéô_l!ô)gßA±$çÚ=$ØË1ZöŠe][ÿã="8\í2Ætã~+¸ì+ër·Ë:òæTZJ3Ëé®×ÄEì9œmm¨D!Ωz«HY‘Ll¦ê•oÛVEd(6‘µ1‚i½º]çR•ÁõgØlîw}ÌH$Á©–!ÓFý]BŽë i{ûºœÅ2Yþžƒ¬«.ˆ>òî‹È%¯‘Eæ×G’ZÇ_»&ô°_—=;ažì«ˆ¬í×ú·ÚÒÐÒ?úxmª]΂`ç™/ѸƖþ²~Úy –(1¸|°â¥g»ý*<ƒžéz£DTg]YQœòq]ú*ÏSa¢Èuµõ Ü®Ëñ„Å}]΢DdâµëV>ý‰tÕO+"ë“ô©µ´CK‰È„L}.ñ°—\×*ÿÒ9ÊÓö­„t׫ÜÞüÑŸ¶ë,×GîÜ+}¡ñ·}CˆØÙ¾…á~Õ}5&+b¼>¬×®…Zfª‘¥Ó¾1ä>[úÄ+Ü7u£O¬åP ÷L¿§þõ뛉Èé.»¯#ã”­¿Ä œg2ŽPë_ý¸7ªô¯ì$ã‰ô¡öìùSÊßϳ kß6u¤¾7Ú Ñ¸¢§Õ/›¬ý°°ß<µž¾sÏí‡n¡=Š\ï‹Èý5¤¡/0–¨ ×n‚ų´*b|_ˆˆ<]""‡Y³S٠ʸo—³0Àyí±^3ÒqÖG`cP:ðÀ»A:pƒ^ Ø:dƒœlÙ B!ìO8qÆ»~4G‚SIH%Èï8å˜8Ï1C‚™ñÈÈÉ! p’•%àÉ‹³ÌQ$Ç8h„F9A›ˆPÂC9vCËãq„‰iÎq¾ôEd•³ÏYæ4 E1Cù\K?²CÌù2)LX špú *ýc÷…u¹ÛeysJ½ ÌÆáÀŠøsÙLÒ&"©ogº>¶CÄ­öBLò7›ˆƒÍ –Qb³50 ì±Zµ¢Ù´±Š ×¤_Ñqe(}ç™èÐFëUiB“6éšüë×C¬Òöµ{iú›˜ÂÖÓˆ´U¯ pÀ¯f/C»®ß™ç@}ioÚQÁõ¨;é ”·¢:=›ò*_:%Ü—hFDsŒúò,‘ˆæìS"¾ëb·îÚµ«Š¾%¸çåhÿÚNåªßç³ð‡´Ï‹g^¨!y¦=ÄD}€q„ð¤Öõ‘Æ’6¢Øó_Âd‹±@û.aš8æÙz+¥ò÷|ðÉ«úgýˆ>¤Ò¨HZo ŒK%6¢–ÚhÓõ\x¾]—è}cŽë$@ê׆>¬g¼’O=Ë… "ýj[ÏcE‰ªϳ¾%r}Ï)÷Èóï&PóŒÏêÞ¤Ñ2\›zp_\›<ô¹êLÙôµŒE¥©®*pŠ §¿rŸÕ¥cÔ…ÈSåáwWtº¶£Ïj?ôGÌuüˆ¬õV ¡¸ÿÑ@é×G—M0Ö:ù®QžÆmÃo%ÒkëÊXK-ÃùÚ¢´Ô· åÓÞ¼!ÒŸD(ìÆ—ž íÄõjGžýñ6áY”6±\ÞeiÓî›:rýËÖSV×ûÂy®›m¤ìÚ»ç¢Þ"“žçÙýÔöˆÉòó›6¦þk9 u>ôÖ• ˜š4P÷Æõ]6¬ÿ<ì ‘§KDä0kv*"… :åšqÔÉ·ë é8ký Çè8#5ƒÌX0W'jÔ`k«(—BaL8Æ;ŽN Çœ@h\3ö ÚHaÑ?"lÚh´‚ãodN®s+Úª^e唈Z‘†|G#ÞÙjýår$9‘³þ:„PŽ™sŒÅu¾ô”·¢’”QzË"9Rœß!ÁÂn¬W®ƒ¨7M´/¬CÈÝ.ëÊ[\Ô q€Z¯ár 9­êœØc¿vÀ¯Hr·º$šÚ8Âì&çô#úúpxÝ‹´+÷O(G™ó¯håôkkÊ©<ÄBl+VpZ D²çØÚž6¨Îˆ+Ú*1FûšìÐæý¶™(ë™c J[„#äíÚå/å×–åãYÈ@ˆ j±? ÕÊXÏÎùÊZulÒÆ1µäAC:êŽP¦~ýè½0ˆUžþˆö«¯o# °Ú(ßž9íE_èYЖˆTÚ’6/­øD2‚ÓVж XÄÔvýýe8ž8ª=¯z.<«ûÒÞK(o‘¾¨Î¡5‹×Èåz«a«a÷庠.V雥iÿ²ó«¿ï÷YëbÕ=.ôïÚÄN'Çä¥]™j»è;E§®[àòÆ‹ˆW“O!|Û‰ˆ<]""‡Y9„BSd¹ËØy‹À².b».uáÛ‹Wâû¯ÔA<²l@-£1ELYzbxCahíÙ¾DDž.‘쉈B!„)²?„ÜeŒ™w!„Â*""O—ˆÈaÖDD!„ÂSÈ3ïB!„UDDž.‘쉈B!„)2¦;fÞ!„B«ˆˆ<]""‡Y9„BSdL!w̼C!„VyºDD³&"r!„¦È˜Bî˜y‡B!¬""òt‰ˆfMDäB!L‘1…Ü1ó!„BXEDäé9ÌšˆÈ!„B˜"c ¹cæB!„°ŠˆÈÓ%"r˜5‘C!„0EÆrÇÌ;„Ba‘§KDä0k""‡BaŠŒ)䎙w!„Â*""O—ˆÈaÖDD!„ÂSÈ3ïB!„UDDž.‘ì!"þùçýŸC!„F…€;–["òW_}ÕßB!„0*#"O“ˆÈaÖ‘?þøãNHΖ-[¶lÙ²e›ÊöÙgŸ."‹ôé—+[¶lÙ²eË–mÌí£>Šˆÿüóß?ýôÓÅßÿþ÷Åûï¿¿×öÁt"eñücñ·¿ý­;ö«¯¾Úø}c¶r¼ýöÛ‹>úhãwùJ[ZýüÚqþ?ÿùOWÖwß}wñå—_nüî†Îuûì³ûDÛToºé¦îÿÄ%Îñœ†×_½Élþæ™gvŽq•,Í'Ÿ|²‹â%P#ýÝÇñgŸ}vw¾±[DéW\Ñ9"D^c-!¹òRnâ¡ýÿüç?»È—ïÿû‹;ï¼³ËëOúS' ÿêW¿êÒE]e•Î /¼Ð9;DG‚ªóÛ´ jDæ×^{­;ÿXœrÊ)]”³ò¼E;'y<´íÑ}w- ú|ðÁ÷},”ÍÄ„6¦½}¸üòËÿó?ÿÓM…BÅ«¯¾º¸ùæ›?ýéO;»®eO\uÕUd¼c/š 7‘ÊvÛ.‘§KDä0kÖ!"‹¤zì±Çº(¦k®¹¦›U»ÿþû¿ùÍoºÎ#*êåÐCíDdQ'‘C!|xúé§çž{ngð¼Z™,¸}ÿÚk¯]œþùƒâãã?Þ‰ÑD\<òÈ#[ËJ\Eßu×]íiYbà‹/¾¸ñgD™ˆ¿}ˆ¾œb1î»ï¾.ív¹ QЧvÚ^ËZ¿üå/;ÑlhŸë=ýôÓ=ôÐÆo„ó_ÿú×ÍQ;gL!w]y«?“ÚÑÿŽ;îè–ÁÝ„6¦mØ/r¸îì#øÖ\ÄW÷òÖ[oÝhGCxUÖý0É¡¸—ì³{ï½wñÝï~·ËGÚW^yew¯+ºVÙ´ å1A`"¢Ú<§XYê:LdÊeòƒ-è­3öàsÏ=·Ñ~Ô#ñU^6i,å•®ük2¤íwMlЪg÷ò·¿ýmw]‘íUãÛo¿½K‡üÇ?þ±»6uT÷B9A8ßdÊ 'œÐ=ƒ^xa7i²Ó¶Ba—½}Å^ ЊÈ&éu5éo£³""Ï“ˆÈaÖìTDæðr€;ì°®Ãñ¤“$ûWô#Ÿ£Éàæl3À#"‡Bø&aí_K5´Qˆ~#¤“Dÿ·¢tûp"ÆíùDJã"çC:Deç-ùdýcbd+XKË䬴[8Äá6Ú”`I «È`ÿZ΂ Öä+J¹·>È3Î8cc© "ljxN@ô‚:Ù ër·Ã:ò.éðÃßx…õCé¢D"©gìöÙØPì)6T-SrÙe—-¾óïtÂ&8œC‘N}L:œ|òÉi›@Ð.‰ÈòЖ Ìî³2ØOt5!ÁN«}"o ­Dß ¨rU˵<qZ:¢åµuÇ”ÍTÊ­üö+—<6‘ÝéK¯®E=¡kqm²ö+Ÿkªòi§"ÀüF,ÞµX‚Å3å:è .íºvéIËúã]tQWfuà>z›@ !„ ¶œq¢Ü¯IÓW^y¥³ D%³÷†&é÷…ˆÈÓ%"r˜5;‘9X v†µˆc‘/^ûå$”CRp8öED!„ðM‚ÐÔ‘ †¼¥&ˆdD§¾QOäÑé˜Öi B‰Pµ¼€±”Gœj£…—AxSÎH_è|ôÑG;q¬ƒ[äKL$ÖzÉ…ë$¾‰t‚ÀÍAúÃþ°ñ›¨NÂÑÒ¤±´‰p¢r–­ï¼Ö!än—uä-êÈ»ºÑjR(Lœä@²›ìgCÙPê];b/i¢jµ‹Z℈¼•^mÎz×%Ø*‹k±œE‰£î¥eGDÇXE‹`w¼ri‹î³|ßeaBîu×]וɄ„I“ ê*·×{?üðÃ.?¢­7×üí_íH^«®„_mU{w¬z”ŸëÐÖ]›ý®•}éº=WÊç¹cçšôpMåíb±rió"¥' u@8öøM½¢Ý;פì;Yë;„Â<1Û‘~øánü4`Lò–‹ñÐÿw"$GDž.‘ì٩ˆÌùdP3æ 'ç•qÞ:°œ!F|DäBß4V‰ÈÄ9‚£H_Ë<ëZñ”‘Ï©è”àJ !JÛµkW'ª‰â]¶–°1›è<c™\û‰_Ë>œ&j™8(\?Ÿ§žzªK[¤ìĹþþO$m¿6þÒK/uNR»4þ²!w»¬#ïv鈂K¸UWÚ›ˆØ^7룈l(NævDd8–Óê^—¸«íHOº…ÿ›ÈÐ> ¨Úªü« p†ë7mÎñçœsÎƤ‰ûo"ÁR+–ˆ°áܵ‘¯-(w‰¯ž­^‡å]”¹¬¯z$H×~yD`eö,¨ë­ˆÈÒ°FáZÕ‹WŽÛ¿ûÏq!„€!™gRÓäqa\1F½y¶U""O—ˆÈaÖìTDæ1ºÖ…‘‘C!Ì…!™àV¯Ç„dDëøNaàUÞέ˜2áohlö ¾¥"”ƒØ;ô:=‘l™xKäSKXÙúHÏ>"s_\†å)Œí„Èã;±Åߎ}â‰'öú}_X‡»]Ö‘·%ˆŽÞÐ*ÜWv’öDŒ$ĶñÕí>üKDv¿ÁiÔ>¶"¾J«DäºßCÖ#œöEdoMÔ„€2Y®ä’K.ÙX3X:YB8ûNT³‰mH/aWô²óEªW›UµÜÆf×A˜V>ëU–”±>s¿WŠ]·jÏÈnÑÊõ\XnMdX+"KC„ua}çˆÈ!„¶ÊˆlØXg_aü2.¶ß¹ØW""O—ˆÈaÖìTDyÄb¤×V8Ø¢=úËYDD!„ðMeHD&Èù­ý0Š% ÚuŠ‰f,T+ÌúðÖNÂÂÇÐN=õÔ.bÓ1œçù›ØhI ¢Ý2j„±þ2Q¢Y_náÔ(» Q×VQ­[ ]„Ar»¬CÈÝ.ëÈ[[!¤ŠÜÖ~¤Et‘dbÁÛ\õ=‰²¡, Þˆ¿Ú—ˆLHÕØOÄÑ­,Q"²{f-Fç\QÃÖÖv´õ²ë´o& ˆÄ®£Ê b¬rû—`K€®e5`=ç­^]=zŽˆÆî9ªGÑüÖ›l÷ÃÌ~óŒ±sM¾X²Cô=œSu¿U™pîšä94BáÛ͈Ìf2‰ùꫯnüf? ¥ýØÞ¾yºDD³f§"2gÁ+ y%œaŽhHDæ3â7sB!„)Ah"„µQˆƒZV‚&Ú„(å#\Åé<…WîûøX¡‹@Hˆ%ðq*Œ ^7 Œ"›MÂÿ`öÙ•µ¤‘R$±­I,"T9ˆ~®Á¹Ê+’ºŽõÿö•ÿ>>èzûQÇ– "K­íK‘JðÞÉ—Ç×!än—uä­žJ¼$pjIvû«DdÑV;ªÆ5aI [M"xÓKì¬UÖûNûѦÀÚ›t<ðÀ½–Ø –pÀ]›ÐΤO.¹¿ÙzŽWéX˜«]؈°^ÿwÝ¢•µ/Ë«8ßóá:ý«\Úäfס}kKÒ;묳ºöíYénI‹Z“¼êùòË/ïÒU¯µvw‰Î¢¼Õc-·¡¾¥/j\´âºõÆýVQàµÎ²~@ì¤m„B˜&X¡lº¢ÖÔ7®²ßLÈói(ý ÿ}!"òt‰ˆfÍNEdpnìHÎ G•³Òÿ¨‡šÑ-™ÑB!|S †ÛHø›PH4%°Ë,ûPío?DW­ƒlY çC­WâñŠ0kmZ‘¤ÆOb¡ÌæoNH½"YÑÊ¢0[J¬ój“61±¢*•Qz^ÿ‚ IXú {‚“¤N9å”NpºÖ!än—uå­˜êÅ[[\pÁ†HOÌ'dFÝûÕeÊ6#Ћfcq:í#̺Oê|38«Ž•¿¶eâ@& ÚŸ¥X8¤Ú˜{ÙÞgퟭµY RODbiK—“,ú׮ˇ÷Ô¡aQöòqÚ£¶´Õë%/­ªGBrÁeÕs•Á„FMˆ¦¶‡s•ÁD%7l®Àî÷öÙQ×~«hzQÚúu`ID#‡Bhaó×úK†YÎŒb¯ÉXãÖNˆˆ<]""‡Y³¹ÉÁ Ö!‘ûk"‡BsEÔ©7n¶‘Íùs¥85%‚ï”u ¹ÛaÝy³µúA¬5‘ “ f.³ÉˆÉýˆ%¿~ûÑSd|[v÷d×¢^8ÃC ʧLý(øÇô£Ó]³Iåî_‡èø6¯¡zlÙl¿²ï´}ÖR3!„¾`üÛ®­Ø'"òt‰ˆfÍ:Eä¢>fâ5Âe_w!„BXź…Ü}aä-Òü ƒꢖ¶#JŠ®µŠ¨ßv»õÖ[»HÛí¤9„]ky]‡èçuÛ©!„Â7ˆÈÓ%"r˜5_‡ˆ¬C=âµØu§B!„oûCÈ]ÆþÈÛÛZÖE¶vâªèÝB!„–ˆÈÓ%"r˜5_‡ˆB!„°Sö‡»Œ1ó!„BXEDäé9ÌšˆÈ!„B˜"c ¹cæB!„°ŠˆÈÓ%"r˜5‘C!„0EÆrÇÌ;„Ba‘§KDä0k""‡BaŠŒ)䎙w!„Â*""O—ˆÈaÖDD!„ÂSÈ3ïB!„UDDž.‘쉈B!„)2¦;fÞ!„B«ˆˆ<]""‡Y9„BSdL!w̼C!„VyºDD³&"r!„¦È˜Bî˜y‡B!¬""òt‰ˆfMDäB!L‘1…Ü1ó!„BXEDäé9ÌšˆÈ!„B˜"c ¹cæB!„°ŠˆÈÓ%"r˜5DäÏ?ÿ¼ÿs!„¨pÇrKDþꫯú»B!„FE `Däi9Ì"òG}ÔuBÙ²eË–-[¶lSÙ>ûì³ÑEdeè—+[¶lÙ²eË–mÌíÃ?Œˆû¬+Ç[o½µøðÃ7~7®J»Ÿ—£Qüë_ÿêÊjüÿâ‹/6~w í¹µùÝ®Ìý}¶?þ¸;ÿË/¿Üë×òÁlz=«SÈ3ïB!Ì‹>ú¨³Ýh!Ÿ~úi÷>yºDD³fˆÈ:Ék¯½vqÄG,=ôÐÅ÷¿ÿýÅE]Ô9›!„Â7ÂèüãÅ3Ï<³×ïO?ýôâ'?ùÉâØc]üð‡?\œsÎ9ݱxüñÇ»±ïè£î6çÛOÌ%¼ÞtÓM‹£Ž:jqÜqÇuç/[Ñ·…¨yÆgtisÌ1‹Ÿþô§‹^x¡Û÷æ›ov¿ÿèG?Ú+?i>ÿüóÝ1?üp—·óäi,þç?ÿÙíûÅ/~ÑÛž{ä‘G.Î?ÿüN(>ñÄ÷J[:Æó{;ÿOúS÷»zp-ÊwÞyçí%Tï+c ¹cæ=7L(\qŋ믿¾kó!„· t?ûÙÏ:Û‰}uꩧ.vïÞÝ?lŸˆˆ<]""‡Y³?Däûï¿ñÝï~wqÊ)§,î¸ãŽÅ™gž¹8äC:Gy'J!„Â× qö–[nYœp ‹ï}ï{‹gŸ}vcŸ(_ªý¢nÿò—¿t"êu×]×íÿïÿÛÿ׿þµ‹þõcà%—\Ò‰i/¾øb—&!š¨,m¢óO<±‘G!Òø¬³Îê6Nƒ‰Ø«®ºªËÈ+BYò²‰v! Œ+ú˜è{÷Ýww¿öÚkè{ã7vé¾Û²þþ÷¿ïDä{ï½·TEδi?øàƒíšñÀtNáuÏž=]ÕÏNÆù1…Ü1óžÚ‹vnr!„Bø6ÁFcCš¸gkÑ_ØK&éw2ÑyºDD³f"²Îo×®]‹+¯¼²sX9žœÎÛn»­ë4u"9Ëà¤~øዳÏ>;ËZ„B˜4"/¸à‚Îø'ª¶"²ñNd‰e+ ¤Æ½!ññÑG]œ|òɈ‹Gy¤b+:“à*Jå®»îjOëÍI`–gaùäâàƒÞkb„hü?ÿó?8.ß9ì> &*8ËÒ#”[J¥®Ûµ™Øh'BB!„9cÂ×8[`‹ëm®íyºDD³f§"2‹cÀÙäpr-YÁàLˆ )|4ˆSÁ9"b%„Bø&@¼í‹ÈCžpj|#˜õ#Kˆq–ž°^rÙëÿ"C‰t„8K?u·2.s<vW_}õÿ:½dœ%ð !_‘­ÅzK¨0Ž¶½M4ÑѤ1q² ¦+ áš8i+Â%aÙ2ÛeBîvYGÞDö×XMÕEéj'¢~Ýgí†#ém.mùkýiö»Ê„‚¶!‘E€šE‹¯Âñ%R‹íþ°Û¬­­Ý³¥%*݇‰ÖÚ9QXÞ®Áù&ª.ØyÚ¬ÉkjkkÚ„ècW°€e[´}Ñð®A»òí¤NC!„oÆ]cxû1äz»ÌZÉÛ%"òt‰ˆfÍNEdΧ—RX뇀ÌlXŽ©è&Ž“N“³Ùv¦!„ÂTY%"ÅŒ§"}EèîZñÔ>"Zÿ\¢›ßE [‡Xä¦hÍŠ‚Ó`¹ çÛ¥4j?gÅ[BC(Ë…^ØMþ­½,ò”ÈØN·½]__ ÷šf;¦[¶°¸“ ãu¹Ûey‹Ô%¶h m€(ÛŠÈ„fKCüi‰öîÇvDdQòõÑÆŠ`w=œWûØeÚ4{­f·$‹üµ‘Ä›‰ÈÎ7‰PQVÎqM&M m—>´3‘¿)"òt‰ˆfÍNEdι}¥’pìuE…èëßqfËù‰ã•INI»6P!„0U†Dd_+þÁß3ÂYa=b“«­ÈJL#ÆŠ$n±T€1rhl&؉‚–²8Eµoÿüç?÷wuâÁ (2¶áÔ2 &†‡Dlã·óEµ¶(k»ÌD';–0¾]Ö!än—uäMì'¶Z¿¸ðZk?ÙÄz{¿D‰;ïöÛoï중K€5ù°™çÜ~{b§Ù8´Ú¥­]ºÄ2(òì±Ç6DdòÕ&´ÑVD6!Áî«ö-ºs\"²ç¤DäZ×9„Bø6`3öÉÛÂ[GÆÖ¡oZl•ˆÈÓ%"r˜5;‘½N[ëìU‹hc™ã`æÃÄ!$s°9œ\³p!„ÂÔ‘‰hƺ6êØXG0«ñhféËY´øÝ:³ýumMºzSÇGɈ†Ò&dú›øK˜k£Yúˆ%2s.Z8,Æ^‚ö2j}Ûöƒy-õ!iµ˜06Ö·ËxXÂÊ8/"y»¬CÈÝ.ëÈ›˜*"W»aû¨7m¡‰ÌFªåCˆïÚ… z<‘Lèå€jÄcË]ô' úˆR'îºgøúßüëï“N:©‹D®å4|8& N9å”®L&$Ü?‚°%6<Úc}0y«"r Ꭽ<4AB!Ì‘ŒÛob<ðÀ°¼ì­¯­yºDD³f§"²ÎËš‡"b8$¾*^1‡‡sÁaålp 8:LÇû‚7§8„B˜:4k¿Š.|¬ÎrÆ>KYˆÜ$Šä­ˆKçíD¥ö±\±Îr^i´üGƒsQ«'Ÿ|r·L”IY‚ä7ÞØ•ÁR6ÇAðIl¹Š¾PçCwÊ¡ŒÎq.‘’¸YÇ M/{½R²±œÐÙâÙÄo":GIÔ©1ÿÛ¼&2Ç${ê>»~6R?Ùf’ଳÎêŽ59Op-!XÄ]6—ócåÍ°>5Ìd‡õ˜Ùfl°Š&w¯¾î«2ÉÃ~‘Çʯí±ë”ÝR-6euŒ6]Ò/Ùdƒò9VÛdàù˜lèGï‡BsÅXjLeZ ´– ³tTßVÛ""O—ˆÈaÖìTD†W 92¢T¼šë5Gd}ÔGÇi½GŽ§RŽW5ûë8†BS…Jî…à*2“ØKD%Ö¶‘¢Lìo?DW*‰ºÄ6â1ÔxY¬Þì±O„7ã(‘ðh«¿ }p¡Ðr-„PKuë\›ñXäs ¥¢T•U4ê&…E¿­k+òš`ê:Ôƒ«-Kg«¬CÈÝ.ëÈÛý`#‰ò%¸ØEyׇôJD& k7êž-EL6iPøÀžvÀ†ÒM2¸—C}Ü+ùˇ€ëž?ôÐCmŒKô•·ôÝ;krûÈ_í×FO;í´n¿û/_m•#L86qá#ze×yKM~Ú€úS—DkíÕs2ô,„Bs…&büd#ÙTÀv‰ˆ<]""‡Y³¹à333íÅAÐ099ÙüøãÍßÿÝ^ÜìÝ»·ù駟šmÛ¶5»wïn/nþúë¯ôÛíÛ·ÏYÿGY¦¦¦¨ÿþûo«?üðCó矴 ¦§§›Ÿþ9mÿŸþi/^4%EnɼUhË}øá‡é_ù_|ñEóÄO¤´jÕªôïo¼1VC‰\/Jd4GZ"¬>ðÀÍC=4g0<ûì³Í²eËš§Ÿ~º½HDDdü÷ßÍš5kšsÏ=·¹ì²ËÒ¿?þø¼õ2‚÷¦›njÎ?ÿüæÒK/m.¾øâ$íÄñ5×\Ó\xá…iÙW\Ñlݺuvù{ï½—–ñ[¶qË-·4¿ÿþ{ZFY¨‹£,çw^ r8–#…W¬X‘¾g9¿Ÿ˜˜˜Ý6âð†nh.¸à‚æ’K.Ié“O>™]> %EnɼUžþùæÄOLÿö÷Þ{oóØc¥ëSDDäX„z‘vÙÚµk›—^z©¹õÖ[›ë¯¿~¬ºM‰\/Jd4GZ"ÿöÛo) %¨ÜµkW{ñ,ÜXO9å”ÀŠˆˆ|úé§Í駟޼óÎ;IÒòDÎgœÑlØ°¡½j’™+W®Lb˜‘ÈtÒ>úè£ÍE]”~K ‚`¾ù曓lFßyçÍUW]•D0õqˆa~‹pFô>òÈ#iû߹ʲqãÆ´œúqüí·ß¦ÐÏ=÷\sÎ9çÌ>At×]w5W_}uáÌ(e:ÙþÎ;g÷a±”¹%ó>V!8>餓š_|±½h$ÕÎ5xùå—·‰ˆˆ3à:î¸ãŽÙÏï¾ûnsã7*‘ŠYM™ kõêÕ)€¼ûî»SïÓÄÞsÏ=)h%ð$ üꫯÒo^ dYþú믧ÿŸzꩳ™àmË–-i3£PHâ±ÝuëÖ¥‘Z_~ùålÖ¯_Ÿ¾c»"""ÇÔ}ˆß| êQêƶ¼¤#–ºô­·Þšý.:j¾ñêÎà»ï¾Kß!ã6mÚ”êã|š ¥DBˆPt^êõûî»/µ}Œ8¨ç)Ò™ÿ³m‚¢€À†¼Ç\Räö•7ý…^H甶O>ùdê( xäðßÓŽ¡ÍƒŒ§·þøã´>ßQ&@ä3:—‘å]ÐyÀ9#xeû?üp:Ÿä¹ÿþ´SŸpî¹>¸v8ç´ÕØï>ú(ý†²S®£>ó[®Y¶É:'Ÿ|r*+×ûEþ\_@ñæðÒÙÁ:äA' Ç…rsíÒ BG£çéì ¬óÁ±|íµ×R;’Ñ÷kÊEòkž2SÎh³>óÌ3©½ üŽãÊ~Ó®ŒQÐüæ›o¾™m“¾ýöÛij—€¶'ûÅöØv^N¦x¡ËïضÇt1""rü@=|ûí·Ï~¦í¤D.Jd4ãJdF3]yå•ÍYg•nŒ×^{mš–‚Æò×_=ûx.Á%£§Þÿý” ]BË:Œ6 ‰ŒˆfDë± Öå‘‚èW^y¥Y²dI „v Ðx4—ß;Ÿˆˆ ™Nù”óòË/§½1D€€ã©êÖ€ú›Ç#ß|óÍ$Yþ믿Î.GžñS^ îØn.¿>ûì³´œix´ñ–C]|Ýu×%ÑÇz¹´$ðaú „!#£YÎ\Ée£GÆJ_"wúÈQ{Ûm·¥'­8´y¢§vZŽ#Ž!í³Ï>;-ç8²>¿c9‘å´sâš #vÐý÷ßßÊñ`öíÛ—¶wæ™g¦|i‡ÑÞ¢M…¥ Å9å3×m¹N8!>GÈ’/ëÓÎãÿtD>眑æÑ–ã÷´õh‹qÝ ¿Ùm¼âüqÚ7gòFÞ2R>¶ÁïYÎßÄ|pH’ëu>8V´AÉ?¦{ᘲ½8NÈ^¶Åñ_¾|yjWòÊÃùFö"Áù=ç‹uãØÇ1lŸß €2±yqœXÎ1çïÊ”¿#ÚÉ,g]=ˆˆ?(‘/”È2hƕȳÑøÇkõÁh‚†Æ4¼%EÐ@àÂo|ðÁ$…¹ù1‡#rFëTÐÀæ;–±~CÞ †ïh°¼°Œ€†í1Âcœ±ˆˆÈÑ‚ú ‰ÕžÆ‰§mgíº™—Õ!§rù„TD®Q7òÒ–ç£.IL°‚œfÔs^O2ò’匜Dfæ#9£­ÀzHÀ¹Šè¤ü3–#ºƒ½œœ^,}ˆÜQé#oFˆ#!¥´]hóÐN{'Fx#3Y'ÚDä‰ì§]Ã(nFÁ"€9!‘ ‹˜¤mÔ9:Ùq»gÏžt®»^Î;£nù?‰¶ËãZâœ2—QÅtà³®!:"X—ve Ü\#H\ö ‰Ì”&ÈX®ÈìB›ÎŽ+²˜m0ʘm0ú˜õ)3!ôØwö¥ë\p, ÈÉŸÑÑü02žýb›;27 @ç íIö•òò·Ç¹`}Î",ã;m\¶Í±g„6ó·Ä牉‰$µãÉ<>Jäã %² šq%2£“b £E( Bã†F áKƒž Œ@•@„`! ¥Qδ©!¦ ²HÚÌ%DðB> <ž¿¥ñ/""r,#yÛ/”E6uIä|´/bŒuWt¨† ÔÁŒ,E¶_âBG0Ë•3ˆÆF/‡e½ü%~1’šRŸÇvÊOÝ?NÝ܇È•>òŽ9‚óŽŽ!R1—È´‰ò)º6ù²“Nƒ <ŠDæ só'µ lŸÑ¾LaØeûHP@¶ÒÆbŠ‰ƒh;ÎUn°Ý…$2£ÞÉ‹ëÉ£vQË Úzì;‚–kp¡@;$2Ç Lƒ£’cJ4®Qþ^8vü=°}¤:íY®UŽ+ӯǡOc ÄttpBB3²œïhŸr\8Vƒ§žz*­Ïwe„9mטšDDD†ùøB‰,ƒf\‰ L'Áãª4–ù—GI™õ"h¦1¦B‡Df^>¶AÐ@#ŸÆÕ“ßÌõ’ÀÃ¥‘;*}äLEªÒ~ BÐæIɈßyÌï:é<‰× ‹‘ÈHμDyØ>B”Ž{Ä+×bì+£eéÌϧËàZC¶’xòŒr#„ÍÄþæ™k.$r¼¬‰L‚‰Ë>\k´96¹D^Hº†DæXF™hw²ï”‘Îì/eæo†i×ÞüŸ”Kdæ-¦#á©8ÊKdŽó7s¸îù=mUCŒ°F“?åþüóÏÓ6cº ~?Î4/""rl¡D>¾P"Ë W"ôÐ ‘  €‰`’–Æu¶`'‚Ø,&©|‡Dæñʘ£Ž åDУ QC`À´Èä… ‘š ¨@pE]G}LÊhE`4f<ÆТƒ–9Vdõ Æ#ߌ˜ï–åÔņc€4à3LjŽþÙG¦˜‘ã%òñ…Y͸6iLhø<–I›†2Â!ov'À$  1M£›G ŒxܽÁÍ–à˜mÒøfDO£’™‘ ‹‘"""njإŽD€Q§åó#‚©\@À@}IÝÈúß|šD1#C©oYŽÎ_´Gï`ˆ®|z ä0us”…Žß!Ê4 =,GNçSg ܨ·#oä)£;Ç¡‘;*}äÍ1ã,r¬iÇ Ù©àÈHa$<íÄ&ÓƒåóKs°m&„.²“k&^7œ:ò‘¨<éÅ9eJòfÊ@Ìò=ùçûŠ€¦ÃáK;Œü¸6ó§¾ÅÌ÷Ñ–CSΤŒh&hf9í¸‰‰‰´«W¯žÍ‹N ®%Ö‰MÒòBÒR>®©|º”6¬ËõNyâØñ÷€8çšçÿŒ<¦œäÅöx‡/\óLMAç ‚™QÄÀùaÄ5eçZÎ剗P"ÙÌwgŽç‡s ¨`€ùÅržˆå""2|”ÈÇJd4}Hä€Ñ2ñR–¹hO7`îZø 늈ˆ •ÅÔÅÔ¹ümCPÒ%©øm—í* ¿ëZNÙºò^ }ˆÜQé#o$;RIÉü¾ŒŒEü32•àli™øFüÎç4¦Éazd.#ÓóÄwˆQFÆ2¥´¥æÛþ|pºÚj m“}ŒéZº ®kšr°? Bhï+ß-æ½Ǿ®Ñ¶;Ÿào†ãÐn ‹ˆÈðA"Ó!0°N‰<\”È2hú”È""""}чÈ•>òFT2Ò•)Å1%£qãeÃŒPezF¸Ž£Õ Niœ'F3r—)Mé#ieÉÍàö¾ò]ÌY,""RÔÓ<ÄÔ$žÚ™ë%Ê‹A‰\/Jd4Jd©‘>Dî¨ô•7"™QÆLqÀ4Èb?`ÛÈO¦‹¹¥û„í#’y^×hb9rРÓ)›HÌ‘ÏTHã(ׯ¿þÚþ:ABÄË‹cùÝwߥ²q>sb¿YÎ>´Éù\ù²¯?üðC*×?ÿüsÈ2꨷ß~;9ƒ={ö4ï¾ûnóÎ;ï4ÿý÷_Zgjj*ývÇŽ‡ÄÎäͺü&êÓü8pà õçƒ}üøãÓöF…²SÆHÔñ£lã@ù)Ó|°_Ÿ~úéœ}¡D®%² š>$òÄÄDsÛm·5ï½÷^úL¶téÒæŠ+®˜³qP;ß|óM*ÿUW]•nÎýõW³råÊæÁL•ˆˆˆyú¹£ÒGÞŸ|òI³lÙ²æꫯž ¶Û<öØc͉'žØ¬_¿¾½HDdl×]w]sñÅ7—\rIú|øá‡ÍÙgŸÖ!†¤àµ×^Û\tÑEÍ…^˜Ò¦M›f—öÙgÍe—]Ö\zé¥ib­È“{íË/¿Üœþù)†<÷Üs›+V¤8¬ bŽí<ðÀ³ß!vo¹å–”'ûÈïŸtà±ì‚ .8(‘ÿi§–)¼ÿþûé;¶{Þyç¥ãÁþ´AªÞpà ÍM7ÝÔ^”@ ŸsÎ9ÍÚµkÛ‹äsÊ)§$q¸|ñÅ©lœ?Ê~óÍ7ÏJrd(ŸYÎqeyÄä@>ì¿e¿8¶Äºòúë¯OËX‡ôõ×_§e»'žx"]÷Ýw_:?üp:¶ÄÅé5k֤ߒÇYg•bæˆÿ©wß}w:§<òHúüæ›o¦Ï—_~ùì9n‹ë¹@ÊŸqÆ©SvTž}öÙ´/q r­¼òÊ+íÕäùçŸOÛÁ£Ìrºý·Ö7JäzQ"Ë éC"?ùä“Í’%Kš_|1}F´Ò ¡Â›/hª*Ä(?üòË/©1@å¬D9:ô!rG¥¼ ¸iOðdÖ|ÛYµjUsÒI'56lh/‹ýû÷'AÈ`FrOBæ! CÜ=ýôÓiq!æC"4É‹„pDØ!^yÒ“XêÑGM)xë­·67Þxcº?"Ï<óÌæ­·ÞJ±#=‘|O=õÔAù°.ÛE³/€ìd»W^ye’ªüñG’›ÈJò¡lÜÏ#ÿ²ß”rR6ä+1-Ç…u8&Ð6ˆÈÓO?½¹ýöÛÛ‹’Ôã{–Ï%'±‹Xd?CÔ.Dpö‡QÆHj>s\â˜SV¶MÙŸyæ™$uùçiüÐC¥sÏ:sÎÇ‘ò._¾¼¹ãŽ;R¼Ë:ˆx¶Çrç祗^JeÁ þ£Ó“NÄ1HS:N‘«QŸ!“Ù_êA ì¬`g}:ØÎÅBpžˆÇGíèà¹óÎ;ÓñáØ°=®¹FfwÁyãør-w \öëkË–-íE½¡D®%² šq%òk¯½–*l*zl_xá…TAP°ŒÊ …Š›(•Ð]wÝÕ<þøã)o¾»÷Þ{›{î¹'=æ’?Ãc7¬Ç Ÿ^Юdžò#FöD-•‰Ç¨èXþùçŸ'ñMåKo*• óºuëRž”—žR*9zWÓs,""r<@€Æ£›ŒV¢^%¸Ÿk×bèCäŽJyÜÑž`äUl‡BÛ‚ïiœ|òÉJdé¤!Â2]¬Äm1ERÑG¬…¤DL1µ÷(îåÄtùh`ä#†‰Ï¸ÏÇÈPdÛB,ü‘KLÆèWc.½(2‰‹$%¦Ìï¹ÄÄ—9L€'~äþ IJÈO–Èdöx² ±*#QcÀ±R:åÊñhçÍþ0 ™{8ùµa„*â™ãÖ‰Ì(܈i…Kçb€TgÿYÎo‰Oc: âižfÉK>uÇ4íJÜ‹|Ÿ˜˜Hçšý̧ýàÜ’?ÇŒ),8?ù9á7ÄÈiÊÊùA:Ës¼‰‡ϯ¿þzŠ¹)o@;Íqà÷ÄÍWÎâ™}A¬æO*Ÿ_sÍ5³>‚óñ?Ç$âyÎÛâšåš$_:aY?Àlܸ1y…v;„<¹>©ƒç«×É+¶M‡m€Ž-矎:bòãÆ5Ìy£Ü÷ß žKdþ®Þxãôw…´_½zõAS®Œ‚¹^”È2hÆ•ÈÜœ©¨è¤BeAEFEL…ÇÍí©§žšD,ëPÐC‘ð/§@!z2Il‡eTZ =ÂBÍcT<&ļOÀÍ• ’D%M%HyøL@øL>HfüžÇ”(OòF8о@D\# hKÐF¢=AûƒöŸ•È"Ò7ÄyˆÔ<~""¾B#]‰¸W#ψÄIˆCâ7ÖáÿH1 Nb¤-¿a¤-ÓIpC:r¯ä÷ÄbåŸþ9É8„ËÈ—õòyŽ‘„MFÌ":ói¸3j–ßìbŽué ‹rq Ž|e=ÊÆïÛóS¾¢ùwÄ°@QN„v>Íåç;Ä32mä '‰‘l+žÐ 8–»»!‘)ÛbÛü›O‡Á¶ILüI'-ŸCÔG‡@ÄÏÀS³lÁŠTå¸ð/R—ÑÝÔG1º›ý@Š²¿äGÒ²>ÛAtrþ‘ÆHUŽõŸéeÄ*ë3ª—ß²>ÇïÙ&uÒ¹Ë ®öq£#ò²>BWÀ9çúb4×eg?¢Îäp<);ß…Hæ;ŽòÉ`œ_ÊGìÏ9¦ !ÂcNfÎçkÑŽø¦,1e×ûÃqÅ-o\5¿c?87–ó·熶Ç‚¿3öuÙÏùÞ—p8(‘ëE‰,ƒf\‰LÏéÊM“!–›&=‘ÜÜøž‰6•‚˜ €uè%¥b 7™@‹ å¡‚£C¥FŠŠåÜø»^²@E@eÁ¶cär<:DŠ^UÊÃÍž ‰^XOü†Jˆßñ¶67{Ê?ÎÍ^DDdˆPO¼Åˆ'‚K‚Îqœ>Dî¨ô‘7mŽ íÚ'Œ@¢-ÃÈ+‚N$ …ï”È"Ò7Üg¸Ç )‰`Á=<= ÅU‹Ð$&˧Z`9"ÎÄSùH5: ó§IÙ6q\tµc'>3R“uÚ†XÙ– p¾Cì…\EÊ…DΡYÇŒ”mßÃÉ—åL߃Ìä¸P&;ûÅ(V¶A^ä ÈÊ\†r<Ms#s‰LlŒü¦>äœw†D&E¶çóëò{b^Ž玲Ä_D2Ç…'x‰‰‰W¢tØrþ”ìG.É)?¿”… Ž¼‘øÄä”óLLÏþ³ŸäÍ9eûÈa .¦4Ÿçˆµ¹ÎðÄ×1ý#ÿr½P#™YžƒdE<#qYÎõà(×ûGYÑqN9î\›ìåe¢¬#Ýé@ ³ÿŒ"ÿè£Òߟã<Ñ^‰©UÊJ~±-òbŸØ_ÎI´oü7œBs_³œ't.û9—lc:*JäzQ"Ë W"7ln’qæÚ–È,'h**R*1¤,P9PÙ³-nÌTîôF6uTè]£šW"S¹Fy©thüP‘Pö\"sl¢¢ˆ^R9Fù0ꇺ› —>…ãЇÈ•>òÎ%2'^Ø”?ZL°îœÈ"r¤@\1*’ûBŽ˜‡—Ü£y†„ ˆÝmKüF F ĈO †âÏHPd&‰xíçññ÷9b?‡ÈCº!yÛïÏ!ãžÉ6‘}@¬†¸EêÅú”e.‰ŒXäÞL%ï|ª‡ˆz9䉠DˆSn¦Š  Ôiˆwä0œb¶™ºÑSkÇ-d7Ò1±”/y°]ÊS)Pd*õŸC²Ç‚¼(y#Ù‰qÙ&ç—rño{žad5—ŒÐf{|~î¹çR|ͯ€Ñ×äOü˱æÜ#æ!¦„h?•‹$Efr­QÊG9ãår!Où>¿Ö€éø 瀼د¸¶¸Î(+#™)/î =­$ûɹáXq<ö\+\Äîl›i#8œÛ¼NG$3œú™ãŒk›ù=çŽAf¬ò}åï":"¸ù»Ê;KøàŽé,p ”¶í‘ò‹A‰\/Jd4}Jä¨,ç“ÈTT€8Ž@*æÊ%2?•7W*:TÎ4r¨(º&ÀÆB8$rÜèÛ9OÜð©pç“Èñ{**9‚!‚~êl‚1dÃBÓPuчÈ•>òÎ%2mÚ“|GF~-[¶L‰,"G "‰EìE£fs‰ŒhcÊ$3‚móæÍÍ Ì½x IGÜG"†ãþÅý’ѱ¹ãÿä1"1Ÿ‰±æzAeB`"óNÈQÈCÐÒQÉgÊÂ(T:1™.!Ó¸w³_¹0&â½P>îÏù\½”/æ“æ˜ù’b:"D%2!/âã‰Ä%åC¢S'2"›QÔ,g[ÄÀSŽçqÊg8qžØ#‘™Ü¯ì+åEÈñ-S\°_lãËÓAˆ`F ç£Ä9ÿìë³ÛÉIq,ˈŽa—DFæ£o!¦j@žSïå™k‰íåõ*ï âzãúá8 Œãú¢“x|(7~ Ý¦`45ÇŽ˜=¦ß 3€k‰ý ‰Lybžé :XÈ›|)[;oÎç…cÆ9á2rŸë”ãŒX&ßö4\‹ä ÇOpÝÆ5Á‹ùûb•ÈÃD‰,ƒ¦‰LãILÅÀMr‰ÌM˜|ÌC"8^vÇzyE؆J‰ž@*?zRùLƒ†C¼‰x±™cDƒH7 "<朿žà‘`ŽoTú¹£ÒGÞ¹D¦=DJûƒà”¶€cNg!"G•GÈÆ ä æ¯'"r‰ÌœÄÄ@¹$þâ>Ì¡K —?‚$Cˆ1Ê•ø+´Ü㈡&&&R|‡€DD·G#H‰ çº2ÒñÆh^„&ÿRfä!²0æ~&æ (#åŽ'a9H<Ø„ÈDv†”$+lj—رOyÞÈD„ e%–¥ ±œ„D#á­Êþ“‰û?‰Ž3ËÚOÚæ©Ù~ñ2ûó3ÕB@y©wˆgÙ‡öË 9OļÈXÖay¼G¦|Ç1]h$2eDŽçÄè^bhŽs.‘9?Ä÷yG*R–:’sÌh`åpÎÈÀoó6åã:`d/ò<^ÄtHrŽ;ò7Ÿ7Á±Eä2ýG{_Øo|+x…8‡”k•k‹cË~ñ9Ÿ’„¿9:ø—ýÅ…äs~#¦9‡´ FE‰\/Jd4}Hd‚F0-7qz iDÄ‘íÒ¥KS¥T¶ÜйÑG/3ÛX²dÉìüM12==wÜà©“¨óÁ£¡Ü¨i´PP1Ó°à_$2y°-&Í9*H~C†Š4^„CùiPP*zVóÊIDDDþ7‚Œº“ŽcFf!¨çÛ½.†>Dî¨ô‘7ò¦Sš`ž@˜íÅˤhÛpÒNBÐV™O”ˆˆŒC¼LqHƒtÌë'¾B÷)îÛHeÖ'¾C‚q?à'!Úbîa1š’mñ=gÈêÉÿÿr:êb-â)þ%³!ávù²Xž$Î!žctiÀ>!ç(WL-‘ÏÅ̾1-bL­ÑIN¬HÙI]S(" ‰çƒrÄt!m±är›8œ}§Ì”sE¬É瘚u‰­©[#¦Ž˜˜cÏ9 Ìüžz%Ÿž‚)-˜Òƒí±ñm¼“Xõ6ƒ¿âú þ¦þŠr0º›˜:Ä'õ#‚–:-ânÎ3ñ?ľ2H+à83Eå!?ÖÏã{êBÎ=Ûãü³}À ˆ9ïœÛí×È_ÊÀ¶Éƒeü&^ŽG¹9Ž,gÛùqF ÇrÊÆò¹ž€Fì³}F?eŽ ¿å!õÙw¦Eá8±­ø›a=þ®XkeT”Èõ¢D–AÓ‡DŽÚ=»}À#&ù œ  L…–'¾Ëo¤ünœào.úÞžˆˆÈ °ã҇ȕ#™w#‘£²l.!6¬›Ë5þß¾§wme!6ܳÇÉq:ß~IÚùÎ5pª«Þà·sý&`û]ËñãÔw]e› ÎQûšÊYÈOÌ—_ûüó¹}>º>*[çù~K9òŽõè\%r½(‘eÐô)‘ô2ÒãJïsžøn®—4ˆˆˆÈ±Ç‘¹ Q2o‘.”Èõ¢D–As,Jd>%EnɼEDDDºP"׋YYDDDj¤¤È-™·ˆˆˆHJäzQ"Ë Q"‹ˆˆH”¹%óéB‰\/Jd4Jd©‘’"·dÞ""""](‘ëE‰,ƒF‰,"""5RRä–Ì[DDD¤ %r½(‘eÐ(‘EDD¤FJŠÜ’y‹ˆˆˆt¡D®%² %²ˆˆˆÔHI‘[2o‘.”Èõ¢D–A£D‘))rKæ-"""Ò…¹^”È2h”È"""R#%EnɼEDDDºP"׋YYDDDj¤¤È-™·ˆˆˆHJäzQ"Ë Q"‹ˆˆH”¹%óéB‰\/Jd4Hä´¿) ·”È ‰<33Ó^$"""R*‘ëD‰,ƒ‰¼{÷îfÿþý&“Éd2™LÕ¤}ûö5Û·o/&‘É{zzúr™L&“Éd2•L»víR"WŠY 7ž;v¤M&“Éd2™jIttONN“ÈäÝ.“Éd2™L&Sé„Ùššj7_¤”È2hœÎBDDDjÄé,DDDDÅé,êE‰,ƒÆ뉈ˆH”|¹]ɼEDDDºðÅzõ¢D–A£D‘))rKæ-"""Ò…¹^”È2h”È"""R#%EnɼEDDDºP"׋YYDDDj¤¤È-™·ˆˆˆHJäzQ"Ë Q"‹ˆˆH”¹%óéB‰\/Jd4Jd©‘’"·dÞ""""](‘ëE‰,ƒF‰,"""5RRä–Ì[DDD¤ %r½(‘eÐ(‘EDD¤FJŠÜ’y‹ˆˆˆt¡D®%² %²ˆˆˆÔHI‘[2o‘.”Èõ¢D–A£D‘))rKæ-"""Ò…¹^”È2h”È"""R#%EnɼEDDDºP"׋YYDDDj¤¤È-™·ˆˆˆHJäzQ"Ë Q"‹ˆˆH”¹%óéB‰\/Jd4}Iä™™™ÙÿïÞ½»ùàƒšÏ?ÿ<a"""2l~þùçæ±Çk~øáæ‘GI驧žjöíÛ×^õ°))rûÈ{ï޽͇~Ø|üñÇUµ‡8'”ë“O>©ª\"""G’ß~û­Ùºukóã?& ›ƒÏÀ|ýõ×ÍöíÛ›ÿþûï åm¨?'''›¯¾úªÙ±cÇA>d.þþûïæÛo¿m¾ÿþûäKÆE‰\/Jd4}Häõë×7Ë—/o6nܘ>÷ÝwÍÒ¥K›+®¸"P"""2l’gŸ}vóÒK/5¯¿þz’ÈçŸ~ šF¥‘;*}äMzÖYg5_|ñØm­>!è¥\—\rIUå9P—¿ùæ›Í\Åyç—üÅ/¿ü’–#Œéø>çœsÒòsÏ=·¹ÿþû›éééÖ–þÇRÇ9ížXÿ¹çž›·cö³Ï>KmêÝ‹.º¨¹òÊ+›~ø¡½Ú¢P"׋YMùÉ'ŸLÒøÅ_LŸÇ}ôQóå—_.؃'"""Ç>›7on.½ôÒÙzÛ¶méó?ÿüÓZóðéCäŽJyó{ÕË/¿|ì¶V°ÐH§Ã¹M¹|û*WÐGùDDDú„6 §¯½öZÌÓSÔÍ<ð@Zþ駟6gœqFóÎ;ï¤å|fýuëÖµ¶ô?^yå•$¤µÌúêSN9¥ù駟ګ¦åäuß}÷5ýõW }ÓM75·ÜrK’Ñ£¢D®%² šq%27P‚n²×^{mêc„ b™e“[¶lI·`"šW¬XÑ<úè£)¸B6ß}÷ÝÍ]wÝÕlØ°á Þ;÷`$Óí·ßžzúòÞ:nÀkÖ¬IËn»í¶´=*9ú„DŽÇCy\S‰üùßÿM#šh'=øàƒÍ¦M›Ò:|GðI;ˆôꫯΎ|úõ×_Sûç7ÞHS„ÝyçÍûï¿ŸŽ5/ÓQ0òûŽ;îHÁ)O‚¦Ã÷ÜsOjw½ð ÍÎ;Ó²ÅHd¦'‹¼Ö®]›Ú]+W®<(/Ú’«W¯žm“Ñv‹6yñ™§Ö&&&Ò>PVö‡ žcAùŸþùƒž^#È~æ™gÒ2Fƒq¼Æ9""rü‚<¾ì²ËšÂ‚:Aç'.áÆo<¨#”ºêÖ[oýìß¿¿¹æškRÝP?!žçjóà4Î<ó̃3ë2Š™i0FE‰\/Jd4ãJd‚ÉàÑÏ Ð`î¿ÓN;­¹úê«ÓÍíÙgŸmN=õÔô˜ëðxë駟žþá…¦àŠž?A ° n¬l—›;Ëø½}ÜÜ ´Ø&Á ù±ü÷ßo•PDDDŽ4JäC ‰Œ¬ÝµkWj#tÒI©Ý‚Ð¥ýuï½÷¦ö ¶r¼ø?â—NxD-'íÚCK–,IÒ˜vÛá{~C;Œví)¶K;é‰'žHß±œïÙîÍ7ßœD-AëáJdäð|yíÙ³'µ»xD—XÛˆdÚl¬KY”ióѤÍGÛÏlmÔ³oüžý/¦l¬G²ò»ˆˆÈbA¶"sêI:5q ÀÿéÔÌaðuP»žœššJõºÑÙI]9Ÿ‡xï½÷R˜OïEy¨×ò2-%r½(‘eÐŒ+‘ ®¸iÒøg 7dFŸô\ýõéæÆ÷/Œ &  ˜`FÉdÐÓå¹îºë’Dþâ‹/Ò6Þzë­!BÐÁï Z ˆŃÐfÄ‹ˆˆˆ]”ȇÂïC¾ÒÎ!`dTm¼l”¶Ë²üñGj+ÑÞa¤.£–Tiï  ‘ÊÏèœçqX~C>ÈYž # e=D.ËÿüóÏô(-búä“ONí.‚ÎÕÈO?ýô¬€&/ÚYHcò"Ÿo¾ù&µÍB“û˙ڌö™ö!£«¯ºêªt,üñ$×!FyÆqG€SVžhcÄûÒy¾ ]DDäpàIüu+s2™º'‡÷;0(®ý‚`êA:?©ëzè¡ô”#“©S©'Û¼ýöÛ©£8™¥Ô“ %r½(‘eÐŒ+‘Iè ‚bNä¹$r>‘#OâFKAÀÀ¶"¸!``D =ƒB! ô"¦ :è!ä•q‚= %ò¡ð{æLD¢Ò^¡mÄ”SÑ–¡C[‡Ä1cä/m'ÞO§;ò”¶SÀÈ'ÖaT2PF`ÚI¼%žGlÙ.-SY°]þO;‹©%3™¼^~ùåô™ãˆcòŠ§ÇÙLžòn¸!µïØgä/m3öA irÍë±>A82mð™¶e§-È6‘Ù D¦Ý¤¾¤nfº&`õ >!‡§_¨—Ú/×cÔ±«V­šý?»``[êýx")`43uãÖ­[³5‡¹^”È2hú”ÈôÂÁ|™yqŒD&€‰%¹D&!XàæNï#Œ¹IÈ81ò˜D‚–ÅtäÉÛáEDDäè¢D>”È15ÿÒþá‰* mC ÊËu•K›†vr™wFp ‘ÈŒ`bNå $rÌÇÈÔÈb¶…D¦M…p%ÏØ.Û$mܸ1À‹•È“WL#† &?„/ÛcD5²™‘Æ1*:$òòåËÓ±$ñ. Ê×–È”—‘ZüËgd3eg(åŸë¥E""" A½u#s¨—è´Ìáik|FþÎ&ˆ©(òNÍèÌ¥¾jC¾´¢îÞåÄ6òw>-%r½(‘eÐô!‘iÔ#‰¹i²-‚žQ%27k!°!pbлÇ2+¡ÌÜü ¢¸a#”óé0DDDäè¢D>”˜™ãÀ´¨´U±ËÈ'q¥íC‡y¼Ì‡¹y1²”„xF"çopï’ÈLF"8eÄp´ñèx'¤/í¨q%2Û'Ÿ˜Ÿ™÷ý´ñF‘ÈŒøB–3uÇiݺui9×ré/9\x‘Ë™s¯`Z¨h³ÄôšÔG@Ì“2´ ™L}•{~ÇTO¸ `*¦xY,mêÌù L×ÄÖ´ FE‰\/Jd4}Hdn‚A̹ÇKñ¸QÆ‹õÖ¬YÓ,]ºtöF‹8æÑLŒè‘c¼0&#!p b=F¢pS'HáÆÌM™GùL@†\ŽíЈˆˆÈÑE‰|(­´M,ÙζmÛR9‚”¹‰é4§=s3"9^†G`É¥xI].‘é¼?á„R< vE»hË–-é3í#¶KÌcºˆ\òfŸxâ‹rѶZ¨ H‹¼"0fÛÈiò¢ÍE@NÑÇH^ö‡6`ŒTfÊÁt!‘ÞË–-k6mÚ”¶É¼Êü†—õÅ»5x,˜v%š} ?äs>"[DDäp .£¤.Ã3PwÑ9™ÈX–SwRÒ™KýÓMlØ°!Õ™Œ †õëק:Šé<iLL]äxŠ†:áL½ÉKt‘Øï¾ûnÊŸ:˜2Eò((‘ëE‰,ƒ¦‰Ì›F™ˆž›'£\x‘ }‚ ‚n–,ãJ`> nÌÜL#¸äÎ:Ñ;HðÆ gx¬‘€ebb"Ý„‘ÉŒ|&Xa9ëÑ»7N°'"""£¡D>‚Iž”¢Íò“–6 &í/d1$óÿòí($/ð/óÌcœËS4Á)+PVFQ±]Ž;0-O€1 K›+–цâ3Ûm?¦Û†À™¼ ¬ÏvÉ )NÇ>å§=Fbà qÍ ƒ¸ ÊÑ™ ›åÈcà8Ónä˜D;²®\¹2Ž ÛŒQ]"""‹wÄ»èÈ$Q1…T/Š¥Ó“eÔ{!y™2W;ê2ž&¢~b}Þ×D]ðtu̧ÌS4t³>ÆÔÙÔŸã D®%² š>$ò‘†ì|A’zÜ°ˆˆˆŒGHä€Æã]"/ŽÛ‘h±!öç‚'ÂL€œ':æóÂòˆsß'G‹ˆÈÑb1uuìbê>Öï«]¡D®%² šcA"‹ˆˆHÝ y<“i©™Êh¦$ài¥Q9Ú"7§dÞGö‘ù†™V,O|Ç9çQ[92(‘ëE‰,ƒF‰,"""ãB ƒ|\»vmš®€GG™o°kìB”¹%óéB‰\/Jd4Jd©‘’"·dÞ""""](‘ëE‰,ƒF‰,"""5RRä–Ì[DDD¤ %r½(‘eÐ(‘EDD¤FJŠÜ’y‹ˆˆˆt¡D®%² %²ˆˆˆÔHI‘[2o‘.”Èõ¢D–A£D‘))rKæ-"""Ò…¹^”È2h”È"""R#%EnɼEDDDºP"׋YYDDDj¤¤È-™·ˆˆˆHJäzQ"Ë Q"‹ˆˆH”¹%óéB‰\/Jd4Jd©‘’"·dÞ""""](‘ëE‰,ƒF‰,"""5RRä–Ì[DDD¤ %r½(‘eÐ ‘8ÐþZDDD¤(ÜR"7$òÌÌL{‘ˆˆˆHQ¨D®%² $ò®]»šééi“Éd2™L¦jÒÞ½{›íÛ·“Èä½oß¾CÊe2™L&“ÉT2íܹS‰\)Jd4ÜxvìØ‘þ5™L&“Édª)MNN“ÈäÝ.Éd2™L&Sé„Ùššj7_¤”È2hœÎBDDDjÄé,DDDDÅé,êE‰,ƒÆ뉈ˆH”|¹]ɼEDDDºðÅzõòÿ‹~çÞ¨ÕØÈIEND®B`‚apache-arrow-java-19.0.0/docs/source/developers/img/conbench_runs.png000066400000000000000000001723311515451615200256150ustar00rootroot00000000000000‰PNG  IHDR¥ûÜó«€IDATx^콉—U×}gï¢{u~t¦î¤3Ú±c'NbG:ù%NìNÛ±%Ûš'$bЄf!„¡’ ÌbžçQÌÄ(„˜Š™ÓoŸÒ·|¹¼ªWxª’÷^k¯ªz÷ÞsÏ9wxë}î©ó¾DDDDDDDDDDDêÄÊ/ˆˆˆˆˆˆˆˆˆˆˆ\(.x(}êÔ©tüĉtäè±ÔpøH:ÔpXUUUUUUUUUUI½`¡ôÉS§r}àPCÚwà ªªªªªªªªªªê… ¥;–òTUUUUUUUUUUÃóJ;BZUUUUUUUUUUky^Biæ.Ü™wúÐá#©áÈ‘üs¿óP«ªªªªªªªªª~®=çPúä©Sgý¥†5ä‘Õ;víNËW®Nó.NËW­ÎP³¼¼ÍùrÏgÙ„èŸU™Ÿf›óm­~ Ž,oé!A­ux½-:#4bý?Ú›/]–û¦¼}½d¿ËW®ÊǨ¥:°lÝúiÕ{kZ\iëž½ç¶ó3úé ÚÌÕ–ʉókׇ{Ëøbâ<ýðóƒÊuµdÙò|lZÚXí㜎×ã8Ðßû6þ]>&Qç¢åòk¹bÕêªÇKUUUUUUUÏÝs ¥ùróùÀN±zíúÔ©K÷ôï?½4õy¼š8yJZ°hIšPùÙ»ïéß/¹4uîÚ#­Y¿áŒ°³%;žƒŸ7vL?úñÅ9Lˆ°ê½µëÒåW^¾þßÈ^s]‡¼.Û0Jûþã?¥¿ûûHÿ½Ì~ïþWúŸ_:,Âå}ñÚÂJðo?L_ùÓ¯¥¯ÿå7Ò#}ÍËÊ¡Oq›ÉS¦¦øÇN_þÊ×Ò__ôôìó/œÜ”ÛD€}]‡ÓòÓ\çjëóÚ[ïŒM=ï¾7Ý}Ï}§¹¶ÒÊ(=9àétó-]²ž”á8†”³bõêtÿƒ¥ŽnI]ºvK¯½>*‡ÔÑ.úsÞ‚…y_7u¼9uë~[7~âiçë,^²,]qÅÕ9&ð/×ùB»ÿСØuº¹szëí±90+¯Òî'Ÿz:ÝwÿƒÍ?¤'½›n»ãÎÜ?Ý{ÜžFŽz#?h Í„ŠW_s]îŸæÎ]öÅz}í›û®Çmw¤AÏ<Ûö•×¯&çe<ÒûÑ4`à ü€'¿›®»þ†ü@¡VYq]<ð`¯Ôë‘Þyʘ=w^ºýŽ»N;‡î½ïÜ7 /Éëp—Ï5þ~oͺÓεƒçöÑÎ>õÍu.Ÿ¿¬ËÃÇŸèŸî¾÷¾@Æ:üäüiê«Û+}5øÙT×!<ìÝ籦þüìóùØS6í\º|EnC±¾wÞuw>–ͧ¢ìƒcüÚÈQy[Ž/¯q®||à@¾öîêyOêÚ­G^>vü„Ó¶]¶be¾V®­ÜwƼõö÷ÖáZ¤‹×ZsFŸ |fpÓµüP¯GÒ²å+Ïèß(ÿ1o¦;îì™F¾þFÓþy€¸ïãýËéÜ%=V¹ǽ³i­-¯C}8hC>Þ=nK>öøiÇ›~^³n}®'×dç.]óúÅcY–smü„I¹ßXýÐ7“*ïô3ýÍyGŸÒoQßq&æãQ<Þ÷Ü{¾Öº.}ð€ƒzrÿ+/UUUUUUU=wÏ)”>~âÄ6'aÀÖí;r }Íõ7¥y å ‚àá؉“ù'Ïž;?]qÍõé–®·5Kå²Ê²AÐ_}ã¢ô[¿ý;éÿèKyZMÐ|Ñ·¾“^ñJzfðsyùO.þY0WîÝ'‡%=Ü;‡Ë]»÷HÿéW~5=ÑïÉ ÷Å60ßúöwÓw¾û7iÔè19Lúµ_ÿÍø–×GÚE¸öÅ/}%ýë÷ÿOIoé|kú/¿þ[iä¨ÑUGLŒ¼=v\úË¿úëܦ?þâ—sÈ!Z1¸¢N÷ëŸníÚ=½úÚë•v¾šñÊkMaë=öøéšë®Oã'NJcÇMHW_{]ê×ÿ©¦Q–¬KÐC(9söœôâð—Óe—_™Þ|ë¦ÑÚk×mH×w¸1‡™sæÍÏÁãW^fÌšÝT7~¾1úÍjïع»id:ÇHÆr ëDûXëa9¼ŠuÂâ¹Âß„¨„Ä+W¿—·-¯}Ȳ§ž˜Û”ëXeü=uúŒtéeW¤~O>•ÃÛ!Ã^L?¿ôòôÊ«#óqçø ϯœÛGŸ8c?Ô`Œ°‹0m¤ÉiÔcrÊÛVlw± Â7Îõ7Ü”CŽ}<éÝ)é†;æ 1®',–ýF9?Ò'ßk®½>?üˆ°›ãþò+¯f9Ÿú>Ñ/ýôg—æ0–vò0…ó·x® ¯\[Œ§ÞŒj½ÿƆ&NÎae—[»¾S¿á/¿’.»ìÊtùWåó.úžÐœ07úŠ`•¿yh£q¹~oêts å_¯\OƒþOÈ£hé‡w§LKW]}m:ì¥|MPß—†HK–.OëR>7¢Žü¤oìõpºòªkrŠF_ct>ÿy°0uÚŒôô gr;F¿ùV^‡ó€öuêÔ9mØø‹ûÅãKÙŽó¯J7w.Ò¿Üۄ¦ïN–ûˆcO]ËýKÐÏ1þÙÏ/KýŸ÷Ç:Œ4æAçÒ›o¿“û˜{[¸nÙ'? uËëôèqG^F9äwÝ}OÓñæþrÛín8ß©ÃÆÍïçí8š ò9–>ôði#û‹Ç÷ήîMq,9\¯ô7 èoî•1= ˸½:²t/ÜÞ8¢¼ÜÿX¼Ÿò÷Òe+òyÄ…ò}GUUUUUUUÏÝs ¥Ïfêæ~aØðôï—\–æ-X”Ã->ü/[±*ÝÒµGZ±ª18äõYsçåÓ/¾üjÍQj„ „pŒ<&0ð™ôGü'¹¼ã'Oå„ðwÊÔééT¥Îøä€é÷~ÿrXIùì;–ÛÕ÷‰þéKò§yd_9`=‚íÿö;¿—æWÚq²²>á #´ÿê›U Ò)÷Ξ÷äPzýÆÍéD¥”û¯ßÿ·ôOÿû_ó:ÅP„ Œ€ïßüVÿú?õt¥W¯YÛ 0r³^Œ’õ! £mò®Þ=q"ûθñùµU«×äõ º&#T%Ü#à!x"zþ…¡¹Ÿ ‰è Ê'$„Œ@úx2:–:þQwöCxÄñ"¤Z´dinG´Ÿ¯Ë+çÛ^ο ’ –匀ŽõÙ–2h¡8![ô=uçA¯ó7#Š “ß~g\ÔiûŠPšìp—ó&2ÿ$´d_Œ®½ïóúœ[ô3û`D=õYýÞÚ J3:—ý°?¦4ˆŒ ¿_ú2ð¯ëpCS¨KùÔƒ°‘Qè„Šô-}Éà öA üÜóCrøW ¥o¼©SjgÍ™›ûšm©[ñü¢BABÌa/Ï#Q LÙ/ûg9uå'ç.çm§¬C úTåªv®±œcÍhÙçYgÁ¢Å9Ðç? (;ŽýÄþ_2,‡õñpƒò8®­ô}>ùÝ©9 '°æ5FmÓq>²/F4¢ržñm½ã®ž¹mqÜ"˜Í×[å'mãz~gìø4köܦÿ @æ0òvÊ´é9H¥Ÿã¾D€Ëƒ,öýƨl®–s~6ç=³æ4…´kBt®G®î)ÅPšýÒ‚àø/ ¶åa ×-–ò}«Ò–mÛwæºQNô/åÐ.êΨz®Ó¥©+÷D‚ê¹ó4Ý׬]Ÿ#ª¹Oq–×¡ï ¹ù/ú“úÌOŸ1«é|àÞʃF-ó ò½¦ro¡ Öáú&è'h¦Î´™syúÌYù8°ú(BiÚÿtåþÎÃ?~ewö¼;¦{#¸6ªŸq¿dªêË9÷¡X×÷¹â’¨ªªªªªªêùóœBi‚ærÕä‹ ;ÜtKž²£JßÑóÞ´ê½µ9,àu~öêóXºé–®98,jeYF<~âätôøñôêÈQéþð‹¹\ ÂÂÖ%¤xjÀ #.Ö!d!HŠp / Ȳ#•6$\EØ׌™³ÓÏ/»¼i´ñíwÞ• h#ëÐv‚£M›·ä}qÌ9zÌ[¹\MGF‹²L1“u˜[;ÂÌÞ>–GNtDH²-¡8Ó`B²ÚÆ(I´%J8Y v{=Ü;O#À1#f Â3¦y 0ãÓ—´ßÔ‡ð²ó­]s¨‰VlOßÓW“œ7„‰ÓfÌl „5Ž§]Ôƒ6ÐÞ[»uoÕÇ‚õc>]öK]YNXFÛa ;Ù¡:uf$-eÐ>êD°ÇˆfF G(M`Ë6œ 1­Û^¯­<ÍM¥mŒÒ&4§/"”.ž·Ô…ó‰å„éìÑo<Y±ru>לs¨÷É®5®»8‡¨+<8~ñ ‡eŒe.Soð€@½JÇ90sÎœÆÑö‡çÎsúrÊýI¹ŒþgŠŽ9¯½ôòˆ|.pÜ” ì¹E=8_YN; éΩh#ÄÙû¤ŽJsì¸ØåÑÎvìÈSS0 åsÎqm`÷ÛnÏç/ÓapNs=ò“ós—:pŒi?£òù¯®Âf$ÀçþEPÌ=ˆ>+WÎ…xH}O̾¹Öy膥ÕL›‹£×ÙŽúPÖã¿$xPUáN™ö¥xáA u?rôxš9Þôå<ûÜ ùzà˜°ä¸q½0}ÇèwÂsÂkú…¶s®Ñöøîõ¬Gý)—ÑÏ„ÇŒÞçoúiJ¸7Ò_å{!õåáõçXr¼éƒÁÏ=ÚŽ÷#iUUUUUUÕ ã9…Òå«9ùrÃå+Wç/5dé†Oæ÷% ÌY°xÉiá3Ã;ã'æ/?\Y«›3‚Þ—G¼ÖJÇ軫øÉô_ýÚ×sðT.—`ä¾J_üÒ—s0X^žÛR)³0†Ìœ=7ýúoüvzó­±ùïbÝhã7¾yQê~Ûy±¯Q£ßL¿úk¿‘GíR·æÚ4쥗›BiÊfd%aÊ¡/æu¡¤:ÝÒúFæ=÷¦‹u …Š#@ù –eÔ‡0'F7G[-]šËZ°pq™’ƒrc~$’1šr2 ²M} ºéopœÙ7ÿZß½Çí§…ÒŒH%'x#,½ä§?Ïsïþð£\W¦ ¤b]àâHLÂN­èKúœàiÌ›oç>dz Ö§®{>ú8!!a#…YÎSF„fzÜîŠÓ‡´ŸQÂ7ÜÐ1‡ÅüÍrÖ¥1}Ó¿p¬>Ü»7V&d#0Žü¤žÍô Ç€ 9ÂvFÙR>!*!£ÄÙW´ö³.u/†Ò_ò³¨"nßµ+,Í# #¡™*„ ußþƒù¡Cs¡t]+çýPIʹGßq®!Çž`1¦aˆs—úóàà‘GÍ £~c,ç¡ýE¨É9M@Yœ9O®¿¡qº˜xøÀƒ('΂fú“›¾[²|y.‡eL­A0K?#Aý=÷Ýß46£ÏyÀHi‚oêÍúŒ¢þFî Jǽ…ÑØÃÒ>Ü;ÿäïx8BUΔ#Èò „ •p”kŠ\cœûl³~Cã~ZYιOÙΦ–  ÎÎs^§Ý¬K?Ci¤LÀ1bšãI¿Å}sû…Š÷©+÷ ®‘MïoÉûå8s Ä=ŸëŸm_ymäi÷rUUUUUUU=Ö%”æƒÿ¼…‹Ó~ø“ü¥†:ñúÜù Ó?ÿÿV}ý?¼8-ü$¬+—YM‚Žr(ËøàâýÓ¿¤oþõ·ÏáøpŒi2¦«’Žõþæoÿ>qb¬C}™ª€)=^ÿ$(Šõ J=ÉÜЄUÅPš0äWþ¿ÿ’G„6~°^1”Ž`“p—†ò·~°=ä_ó Ãd™Vâò˯ÊÓPPOF*F€Á!á jŒ2$”&°ŒÀ™},ñïû„ÄnŒVuøIXJpEÈƾK¨CSÁ-s^³ŒúòZkBiÂÎNý‰AI½#Ä$À%$`¤ül™2„`œÑ¼¼Fÿ(ÓŸQÂd@~'¼%à¢_ù;‡f3f涳/F|24!Çš¾ 4#0cD%û ®Œâ¥ýÑVÚƾ mã³.#f B Zoérkžgœ× ׯ½¾Cžú£X_"#wãØT ¥©côY1˜d¤3ýÊŠL¹ÀT6„Àr›ßßzÆõÂ:9$¬ôeÅ9C8úhßÇsÿq®Å4„¤q^ÐÎÂTÂPÂAê˜ šm9~œ§Ls,iK¾n¶ïÈ#‘ésúŠ7ÜÔ1'`,^c|i!ëðà‚s”é¨/Ë…9· sé+Žûz¢ÿ“Mç^\ÿŒúžW¹÷0ªŸ9¨ã˜å°¸J³ ç)j8—9̈wÎQîC ósn=ðP¯\í'0ežã8¾È½úÓÌY͹Ï(_Î1ò0/57âËúÙå•×äóŠe\Çô×]¸#Ê%\¦\®Fpsıäzã:áþ@{ÙžH„ËÄδ+¯óüéëP6Ëy`ÁÈyú”cM(M?1'Œb§NƒnÊep.Ðœ‹ÌYÎ:œ+L…Òp¸q:× £Ücú^ãXpì©/ÓtÐïÜ_â¡ç#Óã¾:3ï5¦èOΠʦ/¸wr¿b;FƒsÝò Š}ò‚étø»µïqªªªªªªªzvžS(ÍèrÕÌ#¥W­Î#Ÿ'Lz7ÅHéV/X”Ãê…‹>óúÛã&TÖ¿ò´i=jÙ\(Íö„'?¹ägéO¿öçMS2”·% dÞç˜â¢\>R7¾¬ðŸÿåûMÛ2ï+_ŽH`U,›ð‡ä¢o}70ÌåÛ }ñ¥ôÛÿõwsøe•­JG¹P 1b“°†zGɈVFeBñe|¨„·„`ó%q„½dQ=ÑvêÄèfgFGR&¡[ÌKË:Ô‡¶ÇÔ ¼NȨåî˜ó•Q¬ÑF^¯JN–Æl×ÊøÅѳçÌË!¡+£—™ƒï§cÞ~;‡“„R1z’ÞÅc#F)Ÿp)Qâ_øyãBÈÔ%üNàËÈmê—C³ãÇsßEhÆ|¼Ôeö¼y§{„´‹×Ø.BÐØ7íb?„dH»é÷â±æ<(ž“,+‡Ò“¦LišÚ"ÎBABDæ¦È y(CØÎyöÜÆ9¶ãœâ'uæš(Ž’¦8×(+Î5êB@ÉúÅ×øIóu>aÞ£vi £\ ª ¹Þ8–ô#õ!Œf{æ/'èdÚÇ(f²ÄüÅq~û“}2§4ëЗ¼Î9"ìdoÌL{^¹>8¿EÎ1 › ¥yré[Î[\^ãa×\~À±«qŸM¡tCã]Ò&®™¸¨#!*Ç…õy¨Á¹F]8¾\'ü$ÔeÞòØŽ/ïãLjsˆy¹¹>© ºò¼Óã'äãÌuDxËT)œ?öÅ}‘ë‚eC0PÎÙ8Ïò:ÏWY§WcˆÎt"<”`„3ëó$¸ß½X¹wÅ9À5Â5J;9ΈwîQœ lÇwFµÇƒ;ú“QÍ\SJçcØÐø3ú›0;¦¿a;Ž'çlœ‹œ»œ÷BBBéJ3×mŒLŽ@‘P®µ¡tô˨AA$ý@[ßÁýA_PwêÛR6Ó¹ð;¡y±? ¼)ƒÐŒ • ‘0ŸÑ¡Ïš—EPLØÏëJ ¼Æ¨mÚKðEù¬ÏèWFA³è?êA(FØÉ1Œ)b$-e°¶ÏÙr(#¥ Ñùó—Ò]w}îÏø‚:$¥Þì‹6²aŽesП1Jš}RÎyÚÌ^s‘Q̻쓑½Œp/FŒó”󂑺¬Ïqd¿Åzp¼Ù'ëP6!èýûä~‹c–ï#;åѶ¼ÎÔ„ˆÑŸ¼ÆuBYŒFgßyîì9óš®-Öñêk9”æücÔ2Çë¾d\-”æØp-±BÙ}_ÆȃF^ÓŽ¥YÎõAOX7kªi¤t¥|Ž#ÓÁ®~¸goÕ¼}Ç®|/£ÎÔ•éašú·r½Ó_ñß”ÏBFÿ"uEÚƹD ¿hñÒÜ>êƾÙa-!4åp¼ËëÄhoî´QÒôK䓇G?9~Ÿ<¨àEèËqŠûç:ýÎÔô3Ç:Þ° ëp âú"Hça#øãXÒßïÜǸFY)xâ %ã~KøM(M¹üÎÃ,öOÿò_ JSï~O>•ëÁ¶Å{±ªªªªªªªž?Ï)”>úɈ¶ÖÈ—"þJú÷K.M³ç-È¡Aù‹y}ú¬9•õ~ž^~edS ÖË¡4Û: 2µá&û"ð‰©,ØŽ/ðã Ya"RÇ×^#š„A„!ŒBüÍßúoùßáÙ~銕é«öõtÙåW5¨|‘ßÃ}Ëû |¦õøÏ¿öià3ÏæåÌÕú»ÿýÒíwô¬”Ù8 –žônÜ&Âr(M›XN[(ƒ²—Øb ‘ùòB‚FRçÕIpÁ!Á0¿S.S^° #0™o˜“rù7ü4UÖaˆË.»2A¼F@D s3/_±*ÿM0Aa9”æ'#] Þ˜‹6Fv–2ʺV(MpÆT 3ÂvúŠi8à­8æ´‹à;BV¶%ešÚËt„]¬O½cNi6ú¾f4%#C©e0åóÓf2‚AÂ8FRG¶a´/á6íç¸Pú”Q¤ì‡à—Q¥™‡mÛwæ©Š#{{WÎB?BOêJhÇr‚âõøY-”¾¤rÝîqÒqì| à¨3£Ž©Çšß9çèç)S§ŸöevœŒœ/Ž’ŽP•p™P`“¾!D¥„¿;¾XŽ „ƒôõã¡íŽ@žë/êÁ¨ZdúŽë&Eú…cÆ2FS¿a/ ÏeÇt$ì“}s±Ÿ¶í¨\cƒóvìƒ~á‹MÿÒú•¿éÎFs®¬Ò'ìƒ`·¹Pš× @ÙŽ@™k…ý³O€pÄô5JSÇ—©l:ÜxS>X‡s‚ý5Í)Í~:vÊ£§écÎP¿s-Shþæ¸Òo\OLÇB›sÿ~ÒÇü¤^Œ*¦Ž„ÛÇ‚úñ3æ#FÎ-ÊŒþc¿\;œûôç>1çV|s³/Ž!ëð;_šyà`ãô/ïN›–÷ÛS£Šã~I}‹ÇšvÒfz®¯8§Ü¿| d×ÆãÄÌqÿÉëŒi~ʉ̼Nÿw¾µk^ŸvÇ›RssßÊ×zçÆ6ñ°&Ê!@fžoÎ=ÚE9\£1=íˆiHègÊa×<ÐAΡ³ã^È®Êà 7©kÌÁÍ(Êã}€{+e.\´$OS<÷UUUUUUUõüyN¡4ðe^åB›“À`íú9˜þÑ%—¦^}ú¦wÆOLsæ-HoŸ|¸OúÑÅ?ÏôúÒ¶FÂF2åá_‰Ó“§LËaRÈh¼„Ù†9^ JË!8ÞÆ—«ñe3:³_ÿyÄjñ‹â~±Íü¦mXŸßaLð2âÕ‘iû'SEÄ6«×¬Ë£Ë_âF¨mÊÁpe9_àG½£¾”ÃèXF˾1úÍü…m±}”E[ ]‘HÀµvCcà[l+?LŒ€$´-¯“ÿå}þ‚<ê‘«ƒ‚õ}¼_Y܆Ðбط¬O`E¨G?RwÖ¡~H*32—íã£]Y‡0Œu¨£NéM¾Ð¬ŒÇþØŽQ¼LËAhHÙ¼NÙì›À‹s€P‹°–cZl¿³ çmgtq|‘e°ÿeËWæy²W­^“÷C_GÇyÀ ¶'|Œ9Ì›ÖihÈõà $9î”U<>ñû¦Í[r9Q>í%D§ú…‡"ÅíŠò:£néCÁègÎ7^+Ö»h¾f*}Cý9O˜þ ŽI”Ë:„¸<,¡¢ÕÊã5ú•Óôg±¯(—cξ(‹2ã˜Å:ùz¨¬Ãèsf”ÿÛßyÀU/®CØŽãÏ:ü^n¯qß"Xf4tsýr è¿â:”ÃëÜ—ÍN?ïQßâ:œs~O`®]êÊyÇ:ñà ÖAúû×hŨ²<Âñ“kͬÓö#ǯ-·°ªªªªªªªªªª¶-/H(N=v<:|8ÔNí¡ªªªªªªªªªúËí ¥‹PŸÛ8p ÿË—/O×_}º÷Þ{ÓÁƒÓ©S§òòÇ·ÛϽ‘m”_W­·†Òª­›6o@£GN7ß|sêСCêÖ­[êÚµkþ׆ –ßœ.Ä#çÌ™“®ºêª4eÊ”<:»¼Îg)õy÷ÝwÓÕW_Ÿ ·µú©ªªªªªª¶$Aí®]»Ò“O>™ -Z´(õë×/uêÔ)uéÒ%5*ÞíÓ§O3fL¦Ëe´É6öìÙ“ÎXÖŒ)S.D¶¢mKCiÕV:tèÐü„ô¾ûîK .LÛ·oÏ£ƒ,X|ðÁ¼løðáùÆY¾yžÍ“HÖ+¯K(=oÞ¼¼iÓ¦5Ý ËëU3Ö+שlk׫V?ê3uêÔÐÓü]m=UUUUUUÕ¶(ŸcW¯^:wîœ:vì˜Ãh~åµo¼1õîÝû‚ J»2õè²eËÒ=÷Ü“Ö­[W—©HÏÆø/ñ¹sç¦}ûöµ»þÕ³ÓPZµ†¼1qÓæ èÈa4ÿ²sJó;SjÜ~ûíù jÓ¦MM7vn¨ü¾ÿþ<­O"ËOü"ÄŽ÷Ã?ÌëòÔ5Ê)†Ò3fÌÈÿFÄ š§¸QÇòÍ:öÍXÜÔ:tÆæØo¬GÈN]Ëåñ{¼¶{÷î\ÇØoüŒPš§ÉôIµv¨ªªªªªª¶Uùì:xðàüùŸÿˆæsþ€Rß¾}sX}Ë-·äÏý„¦|N/oßÖå³;Énºé¦´jÕª32‚02ÚXÎÊËé³Èª-/oYJ±¼(ƒ,LáÖ[oÍ9F{ìcm½†Òª5äæÈ›+s&s“,¯ÃMvâĉùß|x²7YjÞÐx É›ÙC=”^ýõsÃe;n¸½zõÊo Ï=÷\ºãŽ;²<òHÓ—)D(}à 7¤7Þx#Ú¾ûî»›Ö›4iÒisZ±ïÍ›7§Aƒåõn»í¶¼Y³fV·G}4?>ÿëQÏž=S=òHpæÉ&@Àœ²™žƒ}±Ï;ï¼3=þøã¹îÑBiž³m±÷òYœ€ÏéÍ}ΑÂÕln› m­Pšv~ôÑGé駟Οí æùÉ<Úüg5ëó_Ñýû÷ÏŸów×]w¥ûï¿?=öØcyylOÛ3 .þëœì~Œšûá‡Îa4ƒæ(7r”먟 ¥U[ò–-[òâÖÎ#‹5MËM•›mL¹Á¨k‚éxƒ›>}z¥yCÛºuk~ Þyç¼#”!FJS—+Vä›9R?FbsÃgN(ê̾ù›'’e³oÔcíÚµùÍ7Ó%K–ä²`ñâÅùMöÙgŸÍSf¼™0MI¬G¸;í‹é;x3b=`Ý·Þz+¯ÇrßLTUUUUUµ-ÉgxF=ªò™6BéÙ³gŸöŸÒ|ž8p`ÓÔ×]w]þOáòçÜZù<Îó»(ÿIÌ.×Ê.„µBiþ~õÕWs°ƒôÐÆ_Óv>û¿üòË90fú–3šþ ÿè¯W^yå´íÙ9a4|9”¦ï ¥ù"I0p‘ÒŒÄnM£íWCiÕä&É¿•Üòd›es7ĘfÌMš ›7£xªÊ2Â`nЄ¹Üð ¥ m o!Öã ŒzôèÑùuBiÞô^zé¥ÓÞù§‡”AÍ¿Í°ï'žx"ÿNMpÍüÎ<ù¥LBhFSóÆAHN=hõdÎhFU?õÔS¹„êì—)B¢~ñ&K˜b„¿ £Åv,]º4¿±óÔÓù TUUUUUµ-ÊH^QFKJ3Ú—”ϱ|îÞ°aCþïaBi>oóÝS{÷îm6#às:ëòœòÏÖ £¬æþ ûBÚR(M;ÀÆ@³_|1·™:á;Á4õæ?À ®c9e00ŽÑÔ„È ¸‹Ü¢¼=£¬Ül( ÷¬Ëñ(‡æúùÒPZµ#¤e¾eþ-eÏž=Uß8xÑÁŒ&æ‰S`0…E1@æ†Ë›oBÜè¹a3O¡-Oã†ËOÞà(ƒ9¤!¾è7Éâ™ß9ÍÝðáÃó o0é<ä cÈkì‹QÚ¼™JÇ›M±½¼ÑJS6o&ÔƒéHŠÿbD["„g½)M=£~,#§nÌ[m(­ªªªªªªmÑP™®’™ÏËÓŒ˜fa+¯ñy›ÑÎÐ\ Íë|.æó>ŸÃ˜F@Í÷L1hì³ [[ ¥i?ƒÙ¤6a„Óò‰øòAÚMhOàËɘæ“P:¹‘9”·ç?µÀÇe·JS?Bi2–Ïjš­†Òª5d*¾è€7Âßjo¼6räÈtÕUWå‘ϸ¼‰qc.‡Ò̯DÀýàƒæ›o„ÒÅ0—ŸÍ…ÒLçQ wY7FUJ³oÞ`¸©¿ýöÛy¤uÈhl^c´6Órð†Zy]¥y6lXnKµPšÅáK‹Ówð&e(­ªªªªªªíI>3o2ƒÈøLM#œC^Ã!C†ä¬ \FÑb0M(MðYÒX ¥ùŒ|žG–Åwj‘ðù6ÿMØÌ’3RšÁm±-ëðŸå#¥cð[lOAÐÌ@¿Áƒç©>b€÷ÙG„ÒŽ”þük(­ZCn‚Ü@ „yJÊkÅp–åÜXã©)o2Ü<™Ì›N|Ë,ërSf5o1=FÌ)ÝšPšõá©m¼9D™ãÆËõã_b•»uë–ÿ}XÎúÀ—&2ëL× ¥Ù.Êf;ˆõÔ™‚„ý0•G|a£¡´ªªªªªª¶7ùÜÎçd>OLóÙ;¾øïYbà*¿óY¸5ó³œrù<Ï<ÒŸuÈZ ¥ ËùOoê…Îä ¸£ÍLÕI™;šø™gžÉÛ3g4}Àgÿ]»v弄À¾oß¾9[ˆ9©«mO&Àöä%ÌUÍt!Lõ™Ĩj²„Öô±¶_ ¥UkH˼G¼)q£dJþæ…eÌÛÌDý·,cnš¼Yñ7t\nÞ3C|Îå32Ÿçù¬Ë´LsQmjÏæŒÏÑå×ë-õ ô%/ C(N÷Éç~¦å &@æï{ï½7ÄÎ1r™ ™Àží=Í”$LÝÁèg–3ò™Œ¤¸=Óu°=™á7ß…EÎ@Ùä-LßA¶À”K°]œD?J«¶BÂ_Þxø7‚Xn <å¦Ë¼Í„ÏüNðË ë3bš‘Ò¼q#æFËÓÁ¶]Ì ™7¹«¯¾úŒ9¥ù‚Àk¯½6‡ÚÀröÍS[Âä^½zå?oÍLÍûfD4û$Ô¦ì›:ÃažJò÷СCÏ¥yCâ‹"øf›xÓ¢„ÛüÍípÚQžSš/)¸æškòŒ¡´ªªªªªª¶ec? RXÆT™¼Ÿ‘ÛBÀüi¤þLÁÉwL•eyíã y›pÕÑ~òŒÁn„Ô j8p`Sß”·g ^ô? õYÆ©Ô‡ÀŸ<%FE3ø²©#¥?¿J«¶ÒmüÚk¯¥Þ½{çy¡ ©û÷ïŸG óo.Å7&nÆÜ@ c ° r¹IO™2%ß ãÍŒù  CbþŽmõÌú|©B„ÈL¿À0ŸoŒ”ÉH‚àxƒˆí¹‰óo3„ˬÇ4³fÍjj a2O?©O„ÈñÅüQ„Üü²êÃüÚñ¦_<€Œì¦Ź§ùÉÛªÇ몪ªªªªªíAU>×_o¯FQÍÈ4b²~FÖ@6Á[3è©I —ùRD¦!o(f Õ¶ÙËeü]ÌSâ¡@{ ÿµuJ«ž…ܹ ó¦Ä¸ñZù&[\Ÿe†Ò"50”VUUUUUUUm?J·} ¥Ej`(­ªªªªªªªÚ~4”nûJ‹Ôà\Ci¶/ÛÐÐpÆzõô\ÛžM[X¯¼nk·¯¶íÙZÞž¿Ï¶Êeœ[Û¶æÖ©µ}KË·-•ßÜ2^?räH³ýÿiMsû+¯W~­¥×[³¼Ö¾k-/Zíš8›íUUUUUUÕPº=`(-Rƒ³ ÇÊnÛ¶-½ÿþû§¹sçÎtðàÁs.ûÓxàÀ´}ûö´oß¾3–­|ðAÚ½{w‹m¤­{öìÉû,†nÔƒ¾hn{^'¼Ü¿Þžºówy½Öøá‡æí£Ü>ú(mݺµéµ–d}Ú@[ùÉküdÛ²±¼¹ubY< øøãs]¢­å}‡Å}ÇöÑŽ½{÷6•Ë)‹ý5×o-Õ­5²o¶¡_‹õÂØwôyqßÔ‘:¿÷Þ{ùü){þ¦Î\7år«Y­›»®ÊÇ©[œŸœgüuŠcÒ\?ÆrÎk–SN¹­¬Ãñ©vŒÊ²=ç$íˆíËm+Çò1lMŸ©ªªªªª~Þå³”´m ¥EjÐRˆTKB¤'žx"uïÞ=Ý~ûíMÞvÛm©oß¾iݺuçTþÙJ}vìØ‘æÏŸßr•×i­„_½{÷No¾ùf:vìØËñèÑ£iåÊ•éá‡NwÞygºûî»Óرcó¶ì»ÿþéÅ_LÇ?cÛ‚ÇŸî½÷Þ¼ýƒ>˜fÏžÝêzàÿüóÍžo^³É“'ç~á\:thÞ–säÉ'ŸL6lÈm£=ì‹:Ç1¼çž{ÒðáÃÏ([UUUUUõ—M>wIÛÆPZ¤Õ¤ÖJÐDØ÷è£æ‘¡’HèFPݧOŸ< ’©ìFð¿#ua=,‡n­‘ú¬^½:­ZµªÅ°¯¼ õˆýF8¾þúëéÔ©Sg,§~›6mÊ<(áßË/¿œ:tèæÌ™“×%œ|öÙg›Ýžà¸cÇŽ9¼$ $„%Ì,†·Å>Áb»ýg̘‘CE¶+†Ò„¶=zôÈÁ߉'òkÑßEymÙ²ey{‚OBE^ã8Ž5*‡ÕH8O8‰„©¬GHHHÈ8qâÄ4a„<–í ã»uë–CÐq ‹ h#á3çJçÎó¶¼Æ:¬K I¿¼óÎ;©K—.é…^ÈíXºtiºùæ›sXK¿4(ÿ½dÉ’ÜVÎÅ®]»¦7Þx#aqÔ{ñ˜û•Ñ¸lÏv<`¡½lG¿²Í¢E‹ÒM7Ý”ÃTBSŽoìæÍ›—)u$„-CÎú‹¾#ì®VŸâ9øôÓOçãHð=}úô\Î+®«â¶œ3ô÷­·ÞšÃýè'ŽësÏ=—,XtêÔ)Ÿw„î”G]9†,~äA ýH=iëÈ‘#s?sžS>ç=ûça õ!¼§¿Ø?e”Ï/êB͈kÚ=nܸü7!6mãüˆ@2)›í^}õÕ|ÞÄ1äõŒ²iñ~š{†ªªªªªj{ÔPºíc(-RƒóJ3Z’°-Âe~g„e„d„I„Rlƒª„¡$2ò’  z6u‹‘˜ì‡ÑÒ­ ¥Y‡Ñ ‹/Î*Áu  ‚„½r„Ü,'fjê„Єf™Ñ—„±„ò„d„ò‚¯„j¼Nùl˺„Ð?þxþ›~" e]^c?ìýÈxØE¿0²”  a‹¡4#¡*ýÊïÔùòåMíæ'ë¸2– „‚1Ò•}FJèʈVÊ °$´ŽÐX‡×ÙŽ°•°~̘1¹]@Ýo¼ñÆhR>uà|`5£‘YŸú³?Î ÂOÂÓ(›à”ãÁ6„«œw#öa'Ç ØŽºrlÙiKœ{´à•~& %\Ž©W¨;¡/Ç:p|‹¡4m¢®œk´ãA?à3]ÇŸ@—ö° ç#Á-£èé7öËÈãM¸8F¬ÇúÔ—~ ,B[ö…‚fêLÙlKøÏêM]8XFˆOýYNyì‹óí±ÇËç,Çœþçõy@ÐÎöüäØ°œ>䚤Oxè¨wÔó›ºQïˆ9_âÚâØs\ø›þã¼¥ÿ)8Ï ÉéGöI=ã˜R>ëQ¯8wé?ZЯÜ3¸ÎÎæž¡ªªªªªÚ^å³´m ¥Ejp.!N„ÒƒÎeG,2Ý£jcÔ.a[žŒ4%ão¶'¬bz$ˆ#´«6MAsRÖ'`&p«JG˜ÇhMÂsÚAI= þ¨uáoBH‚[–3…£RYNHúÚk¯åŽPŒ0”“:S> Û˜RýAYÔ‘p`‘p4Ög{ÚNGHÈ~½L¨! !bŒ6f_ô/aa 7õâðë³O‚ÆMMÈíãDÿîÑçåéøãF1A5A7myûí·s`L0H˜H™ŽìŸlÏñ1bD>ÆQ>íg:ÁãÂ… s?EˆK›èg΂e¶a-õ¤|Ž¡(åÐþ&”ÀÇŒ´f„56džðœõ1˜Ð/WFÓ/q®±.}LÛ¨s1”fÄ0ÛRg^ã!û昲ÚA?Ó—Ô‡cB[`Ùǃcʾc¤4}CØËrú™cEßqp>ðP$dð:û£¾ ^çä\£.ì?¦Aá5¦ûˆó+ÖevP/ú‡`Ÿv³ës®qÌxÂ1^³fM./FCsœhçW\ÓœœË´ë¿éÌ:´‡´›öPåsÅT#gF¾S§˜Ÿm8? ü)‹QêqNrNðЇ~#8§ß¸N¨kk着ªªªªíU>IÛÆPZ¤Åòl%è"H#dô#Á$A#§ ‹ø›À‹ŠÐ¨J@šñ7#: ¯MK¹„plÏ(J­ò~«É¨YÂÉ7¶ªMtŒÞdd&£SÙž€Œ¶0‚™¿#¨&äoBoB3–T2 •°™à’u Üø£­i&íŠò ÒØÁmLk@ÀÊkÔ…à•~e}F¿²‚Hþf„**a^Œ’¥ HË¡4mºþúëóÈ]öEèÇ1µá-Á"á)Ç…°•cT¥iGÊ1²—טr„©J‹˜ƒRBDFòës¯ |æ–¦žÔŸzrŸa9ç0Ë9?ègÎ öO›i{KsÀ«ªªªªª~^4”nûJ‹Ô 5nsÒöñ“à(F&~Ö ¥™&€³qc¥\BBFG¶&`b}ÂT¾Ö~Á!û!°d¿¬õ"Ü"P$à$&°$+•„aŒ %l%0åwGFiÒlC½ E %#|#ˆct4ᣫù›}(ÞQúŽ°8B6öû§Þ„uÅéIØí®JÇ<Ó±ÿdéWÊ þÝ@pÎ1¢ÿXŸ²©e8ó°!Žûf$4 bþ_ÂAÖã¸F»Ø'a+s;s^vÒVOÂ{¶‡¥ Ji+á(Ëé¶%ç|¢/cHÛ½ òÙ7Ób°_FøòÀcÂzÔ•>æoööºló]¿x¶®CiÊæ\¦®+ú„Ÿ1Åpì8?mMÌùÁs®ʦo!Îq ¯ yÙÛ3ýýJ{(ƒsš¶s¾ü3 Ç,F¼s< •iD(M}éGÎ1öÏ6L5Ãôcú’ub.qÚÎrÊb9Óϼ”Ç~é;®þŽQâ´Pžÿ6 8¦Mô}Äqa?ô)ý†Œ¤çüæøÄ9BŸPÇ™ÀBiÂÈf#œ%øŒ/£cô''AË¡öÎW(͈QF[2º¸ضd™“Ô‘}³/³˜ò‚’uxõi# £rÙa&]£¬Çß¼N™Œ,%àcÎÒ®Œ"§„„¼+á#ÛÐÑgLÂr¶'p¥¯búÊm.”Ž/‹Ù¬OÈÉèUBYÚ_zG¨G[ éâ<×LßЯ_¿¼¯è[EêAXJŸ°OÖ§Ž#¡!ugäx‹„’L)Á(nÖ 9OMM[ù#~©7û NŒ„et3u"~c9S\ÄHlŽa~öM˜Êq¥ßé‚a^/åŒ@»¥Pšõ A9·Ù.ŽA.Ç’í"”&4'„%¤¥ÍÑO”K=(‡ýQïØ7ç eÄ—S²_êÉC‚ø2@ú”ãÄ(gBäb?Rú‘cH?².çA¼”Áèæ8žܬG?R?úšPœšš@ÑVŽ#ç ep-óð€p›íÑöÔQÎàô£ši32ª™×Ø.Î^‹{ý#˹ž9Ž<èb¿Ôÿ2à<§à9Ÿâ Q9nÅyÏË×»ªªªªªêçIC鶡´H ÎW(] ‚˜³ôÛ™K–Q•„u„RŒÌÍ|®¡4ÛÇ–¶6”f_|¹£o™Ó–-¾XŽÐ‘r è ب7û ,¥¾´‘Ð’Œp–À7¾p¶ÓïÌ¥LPF8G°JI°GIèÆhTÆâöŒ¦ÝôI,' $L¦™º8‚”`‘þ,OßÁœÒì‹à vQ~|q!+Ra*ûŠ/ž£©¯1ªºx|YF0IÀM]X ™°—cË1'0$(åØP&Ç’å„RǾ㋠©?uç8Ð„ÿÍ©Ô…”@šp“’}ó“íiõ#Ü&d§Wú}³Ž!DŽЗà•úÐŽ }×CµP ”Ù7,ûŽ –:ÞÒï1ÅeÐwÅô/çû/YäÁ#îù›ÑÁlC ×é;YkBfFÓñ!æ¿.ö#m¡hÇ‹ó‰s€cC?ržržð3ÂxÎê̾9o#øGô ǃ¶s\¸ŽÙ?×ç(õ¢Þ´)®i¦À¡Ÿ¸7p\Y‡e1ç4}J¿ó¡;õ£)Ÿë„}R7B{êFû9®\“„×12ž²ãüâÚ ¿Ï垦ªªªªªÚ4”nûJ‹Ôà\B13BÚr(Í”\L„Å„S¢Œf$\bz-öÏŒ0.†Ò¬ály¿E w Ñ^϶-„yLu@¸FøÈOÂϘ³™¶‚1Ê–zøÔúE°JHIÛ3‚“ÐŒº°.";j̽£C ã­L¹„¸üd ÂRú0¦(àuÊgd)#µ‹Á;åDYþ¢@ÖgÔ6áûç'Áb„Î!ÇŽ¹¥ cºÚGH„˜12YNÝh_±nˆ1CL3ÁÈ×ø²BêÎöÅ}GMÝ(3êÆHîⶵ”ÉúÍœ/Å~#$縱œ€”ãÇ•íqNÙô-A(ÇŒmc:ÎMBð8)‡¾¢ ñXmßüM¿1š—𙾠(¥_9÷ h© }Ì1á¼ \ú™ ?ŽQñ¡.ô mg?´ƒú2R8®•b?rPûf9¡6£¸ik”MxË5}AÝb9õf.hÊc{‚dFíÛÊùózÇtl‡Œ`#ñ #æ°FÊâ'!9£«y˜@¹Åó“ãmâ˜QnœôKŒüç\ˆréW®[®%ÚÈñ/ÞKTUUUUU?oJ·} ¥Ejp¶AnYB:‚·bPÊïb„·Œše1‡,¯1Ò’ õÙ>~/–[œã·9 AÙû«µnYÖ§n„xLë_p}B¨'?™n€ŽvÄòد³œzD ÛÓÂGÂy‚F½Ì(‡ýÅö„²Q§XÎþXΈÚâ(æbh;Ë(+^‹z³Œù…i_ÌQ\­¨¡nL!å2‚•Ÿåm¢nôe—ëÆ~8ÎÛÔºðZµcT­þ¬Ë6°”_˜û¦ž ô?#h}ÓÎ7¶/¶;ÚE™,‹s—u‹Ó¿Dû¤¼oŽ)#Á‹ûf¿Å¶×~(öu´•:ÒÔ#jD[h;õeY\Oå>l®‘‡ì›ó3^/Ö6°ïxS<†qŒÙ>ÎÏârú~db½«-§~´?¦I‰¶qO(ŸŸ,§,êÄ2®˜Ê#êÎ2ŽËh3e±¯j竪ªªªªêçI>IÛÆPZ¤Í\­µ’†*C>¥&‹áVñ÷ZåV³µë5'û¢^å}ÆßÍ-Çòòr¹Ñ¶O³}k–שöZ컥í›+«Õ¹¼MseŸ}Ï™jÛ´Ô¯Íí»¸,¶kn½òkµö]m›jë—שÕÖæ–•×+—[,»Ú²rÝÎvy¹]ååQ÷b»ËýÕÜqªU~yY”U^OUUUUUõó&Ÿƒ¤mc(-RƒjAªªªªªªªª¶M ¥Û>†Ò"50”VUUUUUUUm?J·} ¥Ej`(­ªªªªªªªÚ~4”nûJ‹ÔÀPZUUUUUUUµýh(Ýö1”©¡´ªªªªªªªjûÑPºíc(-RCiUUUUUUUÕö£¡tÛÇPZ¤†ÒªªªªªªªªíGC鶡´H ¥UUUUUUUUÛ†ÒmCi‘;v쌛›ªªªªªªªª¶MÉr¤mc(-RƒS§N9ZZUUUUUUUµH†C–#mCi‘VpòäÉtäÈ‘ÔÐÐpÆÍNUUUUUUUU?[ÉlÈn ¤Û†Ò"g76UUUUUUUUm{JûÁPZDDDDDDDDDDꆡ´ˆˆˆˆˆˆˆˆˆˆÔ Ci©†Ò""""""""""R7 ¥EDDDDDDDDD¤nJ‹ˆˆˆˆˆˆˆˆˆHÝ0”‘ºa(-""""""""""uÃPZDDDDDDDDDDꆡ´ˆˆˆˆˆˆˆˆˆˆÔ Ci©†Ò""""""""""R7 ¥EDDDDDDDDD¤nJ‹ˆˆˆˆˆˆˆˆˆHÝ0”‘ºa(-""""""""""uÃPZDDDDDDDDDDꆡ´ˆˆˆˆˆˆˆˆˆˆÔ Ci©†Ò""""""""""R7 ¥EDDDDDDDDD¤nJ‹ˆˆˆˆˆˆˆˆˆHÝ0”‘ºa(-""""""""""uÃPZDDDDDDDDDDꆡ´ˆˆˆˆˆˆˆˆˆˆÔ Ci©†Ò""""""""""R7 ¥EDDDDDDDDD¤nJ‹ˆˆˆˆˆˆˆˆˆHÝ0”‘ºa(-""""""""""uÃPZDDDDDDDDDDꆡ´ˆˆˆˆˆˆˆ\¶oßž.\˜Ý²eK:uêTy•ÏŒC‡¥;w~ª:5=‘&,ÝžFÍÛ’mü(ØW^|ÖÜ3ryúÞý“Ó?WÎOä<áç Ï.H'ª<ˆà¾ïõåé;w7^ ?í?+-ÙÜüuσœkÍMßî9¡rþML·[œví;\^M>gJ‹ˆˆˆˆˆˆÈyƒ‘ÇC‡M:uJ“'OÎ!4Áï¼yór@=`À€tâĹÞ6ž^³m_—lú(Ý1biú½ŽcÒÔU;K%œÉȹ[ÒßÜ31í;ôÙ„Ò0o݇é;wOLÛ«„Òñ_¹õí4böûåEiô‚­9]¿ã@yQ«!ÄüQßå—«Òñ……éʧç”_nS0Jøù)ÒÆ]Ÿ>¨ÿ´\Q雫ÎÍu¨5š¹5ÌxoW6}c>sSvè´é¯nŸ:UŽCµQнF¯L_®œ+/ÏÚœ¦­Ú•~X9®wÿ䘾Áï®OË £G_Ÿ»%ýí½“šBi²¾Ùkv§§&¬ÍõÎO9ÊÈÓA“ק¾cßË¡pqZ ‚qʼ²ì­Eœ6…åÎ*”[‘º`ýžCé¯v}'½R%”f”7ígÄ+A(#™7UÚÉëOVö3qÙŽ>Ù´ó`zqƦÔüš´pÞtßÈåéÇÿ"”¦Ÿ?û¾ó^zq榴coc=§¯n 9ÐgZnS>À¡#ÇóßOVÊc ’òþàäÉSiêÊyZ‹"‹+Ç‘ð4 >'­K#fmÎÇ:àx0œºÑ»öI»+«´ƒzŽ˜½9ÿ GŽL3*ueZ’€þy©²î€Êù@;"0f$ý´U;ó”ó+Ç€cóR¥Í» Û2µÉÐéûy…ý–L©1qÙöôMI—?='Mªô¯ÁÊ-§ç¦lÈíbÁ¾†cy=Î;΋Qó¯Í–8q]ú·G§ŸÖ¶€rþÿû&å‘ïÁŠÊ¾¿Ôù­|^”áX|±²lòŠM¯q}L3|~1”‘óƨQ£òtû÷ï//Ê2¯>|8=óÌ3©sçÎé±ÇK÷ÜsOêÖ­[…@úÿ±÷ÞAzç½î_·êVݪêœ{Ž}lËö=NG–dK¶œ$9Kr’%Y²%Y²²%Š  ‚$HŒ`HHDÎ9çœÃ.ˆœsÆ}çéE/ƒo» ¾vöyªÞÚý&ô„îéyû7ï¼CÊÇ{,<ùä“qþ}÷ÝS€ôìÙ3<ýôÓQü~ã7âòˆÐ:u ?þxxꩧⲬó /Ä2˜vÛm·…#FÄåçÌ™î¿ÿþ1ýÄOÄm!˜ÉÍ~öë×/ÜyçQ¼.’DiÄLH¢ôö‚ˆF$õG;ŽŽ¹šaÆš½QôEÐ#:”èâþ3kã¼$J“ú‘ò‘á«ÃgºM ßë;?ü]ÏiÑ6îª?§çn Ÿ¼g\øò“3Ã÷^œÿ¿ó­¥qÂ)©CØÆ^Z~?›wû‹ÃÙ:r_ˆÉŸé6áb¹SÃßgånÚUáŒXùADiÄÚM»Å(Ù¿|pRøB¯áÛÏÏÇù“~ £ ²’í#ÇôíççÅãæ÷·_˜ç#öÿŃ“Ãç™#~™÷o}fEÁ÷þA+ⶈLÿÎ óbÞæÇ΄o=7'¦«øÞ‹óbš‘›³í=yyþkyÎÅ—ž˜Ï Úrº*þ~nü†ð©ìü°/1Ƀ—"â™öÝl›ŸËö—t%ˆ¹_ÌÊâ<þø•…1zžßìx ‘„ØñËv†Ogå’ö‚tÐe\è”Õ ؇¯ôž×ýb¶Mæÿnv®9&ÄuÒ©üý£YÈŒ4¤‰ùÒ“3®¦i;·¿¾8n÷OîŸn}}Qëiclm|ý™9á÷²6ñ؈5Ù aÕ¶Ãñ|Þôê Ïùn \ü~Ví¬›÷dõßil<7 r«S×/N¼2=×ûšOƒˆý±;G‡ù÷ç–”²¡(-"""""""× DbRa\ Rg *#>#ñL¤2B1ÿ¯Zµ*Ü~ûíaìر1ÝÇÁƒ£ÐŒÀŒ°Í:ÌL#ú™òXž¨gæݱcÇзoßpúôé8}ëÞ½{Ü~Ê)}ìر¸=Ê3fLÌaÆ 1º©Hé$JZ´= à ]ž›¸!<;~Cè1luø\÷ÉáG// ‡Oœ‰‚áç{N OŒ^ÛPâ21‚ÑΈ³ˆDú’³zÅňÞÓgë¢ØŠÈxöÜùðÍçæDÑ:ñâ¤Qþæ¼  O_]Ÿn…ˆùŽý—ÄóMT3i1>ÒqtxmZM¬Ï‘‹¶…wºT‘òÿúÔ¬x÷;þ(«[¶Ç²ˆìì7Ť\"š‰2küýÕ>õ).8Ïœ³ñËwÆèb"œ´9þU·C´6lºR´ElÿVv²2þσ nRjð›²óòÑ;GÇõ7dó³ã¨ðN‹â{Sf¤>¾ùÜÜØ# Ê 7_Ú7ê Aüñ\ôtâí¬MüU)áà±K×)O¨£|ô´”Ei¹n üöèÑ# m•@äEdF€&Ú9é9ÈE½~ýú°fÍšpï½÷†½{ëÓ* Ñœ"aÁ‚QÄæʤ !4–£[·nQ|N :³oÿÐ!eaM‘æp…(½xG ɉüݾó£!"¬ÅˆÔPŸ"âÇÄßD®"Ð’j‘AzܲQ”FÀ{`ÈÊðNŠ¢ê›3jb „c>œGÊ"Ž/ç¬ÛËúrï™á³O ÇNž‹&Bé·_˜×Ch%Rº Z?>—/÷æ×ÅÈbÝ%5¯*J© J.DJÿv§1WDÒþè•…™-ˆ‚1bìä\$-tè¿$|ýbNiR‘ð¡@DRDpöQ›(] *œÈ^8u¶.F:5Ž‹ôD/çbäuÒ‹üí#Scjx6ûËú´×ç'lˆç˜4oΨçöîw–ÅÈ`Ro JãÙ9 emØy4>T "¼ïÄQc×…»†_»mdÅXáu»¤öÊ·+³cüÍ£.ÌiÓÿÒ{Vxlä•¢4mÈ󃹨ïÚì\síLQ”.5ŠÒ"""""""rÝ }i0ˆl.BôòC=s9Áüúë¯_6ŸTD;“3Qš(èýûëÅ-¢¡¥GÝ°üüùóDé)S¦Ät)g4ЈÒ,“hL”fDé–|è°RúŽ”ƒè_¢R×ø ââdzÿÿ­Âõ¼(¢’zã¦W…¹Ù<„mDir#Ð’b!-‡|ï»+bèצmŸ{xJøçç…Ζ†fóþªûÔA”\ÆDø’2î¾tülø‡Ç¦…¿é1õ²ãbþ£.‰¼yx0@ú ¢–‰P&ÒÎöáÔ&©R t~{Yx߉ðÝç…‡®¼¬,"‘‰JþÒ3cŠ Òspì¤áœÌX»7L[½'|æ¾ 1ÕHžÛÞX…x¢¿Æ_º©a>µ)„`@g;_¼¸êU‘;(ͱùŸ‰ÀN‘ó‰o>;'<0xeXŸ•KDvjËMqWvÞ‰H?}1B½wÕ?tÈG9óô#¤F)B^sÒ‘äÓw¬Ù~$¦½!ǵ”Ei¹nöayܸqÅY1Ê™Ü͈Ïo½õVLó‘Òeb4)7ˆ˜F¸nmQš)6‡ÆDéšÝEé‹ Fÿ̵֣{¢ ‰à9oã%¡|¿D²Æô‹êÓwUúäèµá›ê.æ=Ò8Q½i÷ÑðñNcÂÓkb„-³šHi„eRO¤4ˆ‡Vì é0*FÂ>=n]ø§Ç§7¬—Êe}¢n›JßÁ¾]‹H\ä¾A+âz¤~ ÑߤôHœÏ ÿÚ3³ãºäZFpç#‡y¾Ów^̳}æìù(¸"€¦ô5{êÏÝÒšúóݱÿÒps¿ú e§œÐ‰)+wÇ•àüwûíY[â9'ºˆ²æ<Ïå¢&Báôäéº(ö>:üR´/¹¯§]_ëê.ÄrÉ+ÝåÝe1ßs}¤ôž°vÇ‘˜Ò‚,&¨ç~bFè–;>Àø¥'g†W¦\)J“›ô IœFÀ&úýO˜î°¬aùtHv^”&ýGQ”Îç”FÐNé;(›²¿0cÆŒ0`À€XN‘¢(=|ÑöðG]'D:Ïîç⇇.Ø…Y"y‰FÌ$%‘¯D2oÝ"Œ\¼#F?#~®Û~$–×kÔš(f“K™tDÿ² ¢H_ž²1¦ò@x%W5i'ÞÏÊD Fp\Rs FGuLd2éAˆê%Ú™â­ûNÄÈèÏgå¦Û5%JÓ£×ÅÔ D€³ÿÛ÷ŸŒ)3ˆŒ}øb1‘·ä+&:—ˆm>þø„ácwŽ “WîŽâ÷]o/9Ÿi™O¤ò‡;ŒŽšË+ž‡zò›m€t%lŸs¹)kß\’íÛø¸ˆÕ¿Ûùê¢4çâÃwŒºì†‰é«w‡ÎYnÙ[ßþïywY<ö÷ÇvvS¿…á³óš¢¡‰gßI#B7‘îœc¸°ðy¼¥¼(J‹ˆˆˆˆˆˆÈuåÔ©SaØ°a1G×®]£8Ìÿï½÷^œ—˜9sf•1æ?÷ÜsQ †•+WF™¨j@l~æ™g.K±±hÑ¢ùŒ(ˆÌÇ “(ØŒÈÍ2 rF#>ù¨‰Ôf9@°&Š»wïÞQ'’Q¼R’¢(=véÎ(Fòº<|àŽX„b@Èýþ‹óßÜ7)¦,@|½¤^hEx$Š™u¡óóL-$Ÿ0‚6Û%ª–(XD?ä/ôšž³®>röe1͹?ÝmbøËlŸYIÏA¤3°=Ê%ê•rù€"i7àj¢4eô¾:|&+›r9> xïÅô°n瑘¾ãǯ,ˆ¹‚Y†0¾:usÜw Zú¶×G!ãûÊS³b>î$Ž#23õþ¦Ç”ÐsØêð/½gÆ™ˆÛ¬GÊ„L¢¥‰.gŸ˜ÎºþM}´oàÜ÷£Ð;zqýùOnƒˆp„}Îá¿eûµ~g}”1û×{ôº†eOœ® ÷^·‡HÌöI•Âù#Rš6‘ÒX rÿ€º¿¸"°T’«šè÷9¥yHA RºÐ†È ê”}ÙÑH=ýèÕ…—åo^¾åP}i#rÿ¤(Ž§ý"UÓÖn¿I] >̈°¼>q n‰„æá ¿šóÀùãš; ç(åÉæ/Ç”ÎÇžßÿƒÙo¶WܧJ°^qýûÃŒÜà ¶YIèæs¼ì_*ègŽ;¿>õÂCÒR¤‡ Ží&*­_i;•¨´ŸDÓfx0’ßN<´Ý<ãìûU„ò8Îåʨoƒ'cjrwçaÿŠmT-©MæYÊ‹¢´ˆˆˆˆˆˆˆH @”& ùugËé>úŒ]£pók)DÞ~¢óØ°²‘ð‰ˆ€¢´ˆˆˆˆˆˆˆH øÖ×ÅÃ/MÚXœÝ&A`ÿVv<|`ˆãk…üØ|„|Å""¡(-"""""""ÒBHc@ʃJé Ú"éx*¥+i)e9'"R=¥EDDDDDDDDD¤ÕP”‘VCQZDDDDDDDDDDZ Eii5¥EDDDDDDDDD¤ÕP”‘VCQZDDDDDDDDDDZ Eii5¥EDDDDDDDDD¤ÕP”‘VCQZDDDDDDDDDDZ Eii5¥EDDDDDDDDD¤ÕP”‘VCQZDDDDDDDDDDZ Eii5¥EDDDDDDDDD¤ÕP”‘VCQZDDDDDDDDDDZ Eii5¥EDDDDDDDDD¤ÕP”‘VCQZDDDDDDDDDDZ Eii5¥EDDDDDDDDD¤ÕP”‘VCQZDDDDDDDDDDZ Eii5¥EDDDDDDDDD¤ÕP”‘VCQZ¤sáÂ…pîÜ9­‚?¾xºDDDDDD¤uuuWŒ©4íẔxû@QZ¤sâĉ°~ýú°|ùr­‚­^½:ìß¿¿xÚDDDDDDä"gÎœ ›7o+V¬¸bL¥i×b+W® »ví*65)ŠÒ"픳gφ5kÖ„-[¶„ƒjlçÎѱ:räHñô‰ˆˆˆˆˆ´{ˆhÝ´iSظqc8pàÀc*M»Û³gO¦÷íÛWlrR"¥EÚ)tôDŸ:u*¾ëVÚåÆÿÚÚÚh""""""r9Œ' ä9vì˜ãJíºAt‰­[·®Øä¤D(J‹´SHKAONii^âÉ¿ˆˆˆˆˆˆ\)!‰hE¹žHÇÛÝR^¥EÚ)ˆÒä“–¦Ù½{·¢´ˆˆˆˆˆH’(í‡éäzsèÐ!Eé’£(-ÒNQ”nŠÒ""""""•Q”–j¡(]~¥EÚ)ŠÒÍCQZDDDDD¤2ŠÒR-¥Ë¢´H+ÁMúôéÓ7L®-Eéæ¡(-"""""×›m|x­(JKµP”.?ŠÒ"U'cúôé¡OŸ>áÁ =öX6lX8|øpœÏ—e‡–/_^X³º(J7jˆÒ8l|´/U7Û|ûí·ÃÉ“'‹³DDDDD¤ræÌ™0iҤлwï8>ìÕ«W9rd8~üxqѪÃXãÝwß 6l(Îj6ŠÒR-¥Ë¢´H•A€¾ãŽ; AƒÂÌ™3È#Â]wÝž}öÙ(J"J?ôÐCaìرÅU«Š¢tó¨†(ãöÈ#„¹sçg5°téÒØ.Ž=Zœ%"""""m„[Æ…wÞygxï½÷âø¥:„W_}5œ={¶¸JUa¬Ñ¥K—0gÎœâ¬f£(-ÕBQºü(J‹T‘cÇŽ…ûî»/ 8ð²éóçÏ7ÝtS;/\¸ʉ'ƨj"h+=%g9ÊÃq@ÈNðÊWÞyÁ(Fà—ë%JSîöíÛãþåa_kjjâñ´e®&Jã„9rä²:¡ˆ€¢¸™¦óϼ„ûï¿?L:5.ǹ£Ž€r¨~S×ùóJû`[•¢§)‡ípÞ‹u!"""""?} 4~üø˦ófí­·ÞÇOyÒø¯’à˘€òX&‘W°N~lÀtÆ Œ_”ÍxuáÂ…qM~DDDDDDZ|ü‡z芄fÆs¤íÀ×' …qVk`¤é#è… Þ¼EÀ&à)4¨Â8ƒ±QÐÀX4 Å·V"~øT¥›(JË€¢tùQ”©"”8 äÎ?ß³gO|­Š\‚cc¢4,ä©8¹¿Ò‡S±eË–Åß8é#ˆé)8¹¥ùM5"v‘j‰Ò8Cßüæ7cù öŸ|È7Ó¤ÙøÎw¾_a«D%Qa™üÏÉù¢p0q‹¯&Jó° m³)Qšeq*y….Ÿ’cÈ!1:!š²òë’ï:EY‹ˆˆˆˆÈ߀!Õ©ó$ˆË¼õÈ›´¥0Žãƒò d|~ÆoŒïÈûœ`üÁ[¸Œ“åEiÆ|û†€˜ôÝÆi,ÃÛ»Œ5¥ÄYŸ±oä6†¢´T Eéò£(-Reˆ¾åI7‘ˤa@ÈÄÉ :Á“›7F>•Žë¤‘ç…ÈjÄfæQ'P¿øñÜ \"N±œO ’§Z¢4ûþõ¯=Ì™3ç²ßyGçF‹66F%Q¨œ9"ÞIMBò?N p ó²šíÒRi'D4Ðòu—Ú û“Ê¼Ó ¹±@ð%h… Þže,Pü†oD¸Ä’¤HLc Æ]Œû˜ÎX$¥¤<Æ€”Å›¹y˜Î·ŠX‡¿IüfŒAÙ%[0–a¼Ê˜“ù)òºŠÒR-¥Ë¢´H;åz‰Òy±µ­Ñ”(-"""""ÒžQ”–j¡(]~¥EÚ)×K”.;ŠÒ""""""•Q”–j¡(]~¥EÚ)ŠÒÍCQZDDDDD¤2ŠÒR-¥Ë¢´H;EQºy(J‹ˆˆˆˆˆTFQZª…¢tùQ”i§(J7Ei‘Ê(JKµP”.?ŠÒ"íEéæ¡(-"""""REi©ŠÒåGQZ¤¢(Ý<¥EDDDDD*£(-ÕBQºü(J‹´S¥›‡¢´ˆˆˆˆˆHe¥¥Z(J—Ei‘v üÚµk‹“¥À¶mÛBmmmq²ˆˆˆˆˆH»çÔ©SQ”>{ölq–Èbß¾}Ò•Ei‘vʹsç¢(½}ûöpæÌ™èDh—ÛÁƒ£ƒuìرâéi÷\¸p!lÞ¼9ÔÔÔDº8¦Ò´k±£G†U«V…›œ”Ei‘v N©)^µ+ Ñžˆr© "â–-[®Oiڵ؊+ÂêÕ«c¤´”Ei‰ù¿´+§þ""""""ruŠã)M»«««s,ÞNP”‘VCQZDDDDDDDDDDZ Eii5¥EDDDDDDDDD¤ÕP”‘VCQZDDDDDDDDDDZ Eii5¥EDDDDDDDDD¤ÕP”‘VCQZDDDDDDDDDDZ Eii5¥EDDDDDDDDD¤ÕP”‘VCQZDDDDDDDDDDZ Eii5¥EDDDDDDDDD¤ÕP”‘VCQZDDDDDDDDDDZ EiiuNŸ;Öï:¦­ÝæmÚv>U\DJŠ¢´ˆˆˆˆˆˆˆˆˆ´*c–ï×{FøÕÎcÂ/Þ9*üRf»w|¸õí%áý}Ç‹‹KÉP”‘VáÄésá®AËÃ/uþûm#2þë-ÃÂÿ{ëðð³wÔÿþäƒÃèe;‹«J‰P”‘V¡×¸uQ€þù#c”ôMý‡¾Ó6‡£Ö„Ïô˜~.›þ³·é:.¬Ýy¤¸º”Ei‘*³fÇ‘°póþ°`óh k„u;†Óg늋V…Í{Ž…»'·[÷Ÿëwþô¶/"""""""7‹k†_¿{l¤±ã¨ðöÜ-—Íßqðdø³G§†Ÿ»cdŒ˜þæËóÂés­£ŸHë¢(-Reþö©áÿ»kTøÄý¢}ü¾ ÙïÑáûÌ ïï?Q\üºsÛ€%ᛯÌ/Nn5îº2|åù¹ÅÉ"""""""Òθkàòðßn£¡?Õcr8|âLBX¾õP8Äÿ÷½·*üÌí#Â/\Ì3½¨æ@¡)ŠÒ"UæO¾ÚwnØ}øTØyèdØvàD½|GøŸF‡NƒV¿îÜüö’ð—ç'·÷ Y¾ôÜœâdiGœ9w>üýS3cjŽu~£ËØð͗燯÷>ûø´°æbªŽ»-?sÛˆøñCì×gÕ^^”Ei‘*ƒ(ý~ ‹“Ã?==+ükßKħÎÖ…ó¶„ÛÞYžš°!Ôì­ÿÒlÝù aÀÜ÷ÃÊ­‡Ãð%Ûãü‡G­ [ö_úí… !Œ_±3t²"txwY¹lG\ny{IøúKó”5»ãº/ó7_zÊøÞ¢íÙïýaÂÊ]áölÝ^ãÖ‡ƒÇÏļM÷_ןÍcp>+wäÒq;?ykqx~ê¦pèÄÙ8'›Ão¯Û<;yc˜¼zwx`øêðåœ(½~×ÑÐ{üú˜Æ¦dËt˜••íç+Ó7‡£'ëË‘òÀxÿ³OÿJ“ÔIt”IDATãŽzÁùCG…ÿrË°ðçM ûŽŽË?}. Ôÿ㎑qRx<=qC¡$)ŠÒ"UQú/ÏDzŽõØ©sáð‰³aî¦ý1",Ðéþ{¶Ì'˜n{{Iø›'¦ÇT«·‰bð÷œþðáIáûÌ ?~sqøµ»Ç„¿ê5=¹(àÞ?lUø•N£Ã·²2¾Öwnø…#C·l:ò]—…_Îæ‘“‰u?Ócrø{ƆUÛÇu¿øÌìl_Ƈ¿ë=#ÜòÖ’ð›]Æ…¿~rFøë¬üŸô_¾ð̬x£˜¶vO\¾ûÈÕáW;×o‹åÙ.)Bà¥ìx>‘×_>6-ü¯lˆé¸$J#ZÿnvŒß|i~Øwôt¼p[\ÿ;ýde,8ønö?OOEDDDDDD¤<œ«;ôˆ”Fpƨ¿õÒ¼0{ý¾¨;Ü>`iÌ'æ)ýÖÜ÷‹EI P”©2¼¿Òytø³žSŸ>2%üþƒãWfn÷­È«(ˆÔ;¬_éB_{qnøÞk cÄ3¢ôçzMoˆHž¸zwÖI«v+·ê02¼–{¥÷„õáSOŽÏ-ñê‹"ôö'ÂGï^žV/ˆé™Ùá3Ù~í=r*þ~sΖðÞ44ŒXº#þF0ÿd¶Ï½Æ®Ën Âßöž^QS¿¡ŒGF­ œ­_—Ý\X÷ÿºé½˜#Šm#¨w²2|­ï¼ý{Y9·¿³4œ8Sÿ‘DrÄò“gÎÅßÓÖ퉢4늈ˆˆˆˆˆH¹ ¨¡9‰Î¿Ðqdøû'ÄÀ·/=;;¦õøPnlèR>¥EªÌ_<>-üõÓÃ{‹·…÷m C2ë3iC¡ST0âìï<01t{oeè2teèšýEüýÝl¹¨ÿäÑ©¡×¸u enÙw<|ôÞqaAÍ(ÿv·ñáà±KB.B6¿…o}{IøÊ —hèíTÞ?õ™î¸¬aÞ„Õ»ÃoÜ=6ì³B—!—>¸8aÕ®ð±®ãb^èQÚŒ^óIw|wYøÄýãòczlìºðŸo.ŠÑΈүά Ÿî19=ui_HßÁM†m ¼ßÔqC~jXòþÁðÀðUQ¸ÿå»F‡/<=+ ñ"""""""R>xÓû¿ß6"ŠÏ¤ò ç‰3ç¢`M*S¦ÿ·Û†‡?xhRª¥œ(J‹T™Æ>tøÆìÚð³wŒˆQϷ䈎I /2uÍžðÊôšpêL]øãžSà Dé…µãG y½e]Öy'øhá?ôžSrÜñÎÒ«ŠÒD<'¥?ZA”~<[ž ?ßaDüða‚¾ÁQ=/wº2挂A ·†ŸÉn,ï-Þ¿<}s˜½qòãVì ?Ÿ“A ¶5L‘òp¶î|xzÒ†ø½+Äç_Ïþò†ö§»On¤ÿæÉ1 ¨”Ei‘*C1é)HU1tѶhOŒ_~ãî1á[¯ÌçÎ_ˆù‘H™A¾å%[…á‹wÄß. .\Ÿê19Ö­^ô¾}À’˜ÚcÂÊÝac¶>Çø‹GÅã‘ò²dËÁðý×Æô¦¤ûDœ&Ð Íáð ¿5Uv¥EªÌ·^Y>þÀ„ðɇ&Fñøw3£“í6teØ“KS1iõî‘ŒÀû±nãC‡w–…C'È }!üó³³cîèÄÖýÇãSÄ5ãï¥YGþ…gf…gëbÿñÊüðþþú¼Î÷ [n°¤a]øòósb”2|¿ß‚Ðst}¾h˜¾~oüc>…Æ¿õž´1ŠÊŽXó?‘ŠãËÏΉB4ÇÕmØʘ7ûËÏ͉¯Ý$³.üàõE ¿kö)>zdÛÜwôtøa6ï·³Ðïdçèº×‹ïõ"""""""R~Ð|C/8y¦®8[JŠ¢´H•!Z™üÎy;WWYtEôÝsäT8xüò'‚”Á¼9™ù@b^»eÚþc§3»|]Þsu— _ÿçE`òPSv~‹Ûgÿöf7 Rv¤ßDGó›eóPöÓ²ã?uöÒæ@¶Ï; 'N{ó);ŠÒ""""""""""Òj(J‹ˆˆˆˆˆˆˆˆˆH«¡(-""""""""""­†¢´ˆˆˆˆˆˆˆˆˆˆ´ŠÒ""""""""""Òj(J‹ˆˆˆˆˆˆˆˆˆH«¡(-""""""""""­†¢´ˆˆˆˆˆˆˆˆˆˆ´ŠÒ""""""""""Òj(J‹ˆˆˆˆˆˆˆˆˆH«¡(-""""""""""­†¢´H©«« gÏž çÎÓ4MÓ4MÓ4MÓ4MÓZhh+R>¥EªÈöíÛÖ-[ÂÖ­[5­tVSSÞÿý+¦k7¶Yoš¦iå7úyúûâtMk+f.§Y¯ÚµšÊŽ;Šr‹”Ei‘*²mÛ¶pêÔ©páÂM+á œ9sæŠéÚkçÏŸõÆÅyš¦iZyŒ~ñ‡~¿8OÓÚ‚ác⳧kmÛ3F.N×´¦ìäÉ“1àOʇ¢´Há†{úôéâd‘RÄMi;àÔQo¼'""å…~Qš~_¤-‚‰Ï"傱1cd‘–Àà Eér¢(-RE¥¥Ì(J·=¥EDÚŠÒÒÖQ”.'ŠÒr-(J—Ei‘*¢(-eFQºí¡(-"Ò>P”–¶Ž¢t9Q”–kAQº¼(J‹TEi)3ŠÒmEi‘ö¢´´u¥Ë‰¢´\ ŠÒåEQZ¤Š(JK™Q”n{(J‹ˆ´¥¥­£(]N¥åZP”./ŠÒ"UDQZÊŒ¢tÛCQZD¤} (-mEér¢(-ׂ¢tyQ”©"ŠÒRf¥ÛŠÒ""íEiië(J—Ei¹¥Ë‹¢´HQ”–2£(ÝöP”i(JK[GQºœ(J˵ (]^¥Eªˆ¢´”E鶇¢´ˆHû@QZÚ:ŠÒåDQºýp=ï?ŠÒåEQZ¤Š(JËÄÑ£Gþ}ûÂÂùóç/›WWW<oøÍEQºu8räH¬7ê§XoˆÔgsûEi‘úçÇÇþž¿ÅÁ<÷[úûæÞw¥¥µÁG9tèPlÃø.Ŷ‡¯Bn®¢(]N¥Û\¿«V­ . µµµ ÓOœ8–-[§ïܹóÒ WAQº¼(J‹TEi¹‘` €s€`Á"Í;sæLn­ÆQ”nöïßÏõ®]»®¨·½{÷¶¨Þ¥EDnLèŸ÷ìÙûhúö¼ GŸ½{÷î8ešÓ‡+JKkƒ_’Ú)~e¾íá/âǤö÷gCQºœ(J·¨ç3f„1cÆ„•+W6LçÕ¤I“ÂèÑ£ÃæÍ›sk4¢tyQ”©"Í¥qÊè×®]Ö­[ÿnܸ1:s ~r=!Ò–6™7@ z2}šß³9¢4ƒbÚ3íšÈ/i9L¾ ÓƯֆ›+JÓNñW0|—¢‹´fÍš8_Që§OsEiÚHÑMF}ÒÆš‚vDz 'N +V¬'Ož,.&U‚zž5kV7n\Œ˜N;v,Ì;7LŸ>½Yí ¡(]^¥EªíÕDi©Gy$tìØ1ÚwÜ:tè:uêž~úé°iÓ¦â*‚óüú믇‘#G^á‰T¥i7Üä“ñ;?˜hŒæˆÒ< ðÁÃí·ß_Ó’–SI”æ¼ë »Z½5W”¦­Üwß}áÎ;ïŒýõ‡ÑG=ðÀaéÒ¥ÅU*•þè7ÞˆâHsÀQ½é¦›Â‹/¾øS&DDZ‹ÆDiúãJý=𚢹¢4÷„Î;Çþ”¾ž¿üÆ?mI$Û7~îºë®xA´øi2vìØð“Ÿü$ôïß¿8Kš (Jí?¦ègbW Vh®(̓ üÚ,y à:xöÙgÚök¯½–[S~4W”5jT¸å–[bÝ%£Žé'n»í¶Û´-üHÖáZƨÿǼYÛ–Nc¢4×5cÆ£ôÍ vQ”./ŠÒ"U¤%¢4ŸþùðØcÅ7Ünݺ5;ß‘-?þñã>zÕèiT¥ù‹Ãž¢n™G›mNûi‰(S¨(}mT¥S½!\¤zãAÔÕê­¥¢4ƒ;ú¤¡C‡†·ß~;ŠôKôQWëÛôk *òiS J3x`0!"Ò^hL”¦_ç^›Ä>–á÷Õúû–ŠÒ=¯¾újìïß|óÍpï½÷†[o½5öÅW+ãz€¯K@÷™Ÿv4#¢ôÍ7߬(ÝBŠ¢t¾ 㣤héô6èÕÚpKDé$V"<8°aí*ÿÐEQú§OsEéùóç‡gžy&ú¡Ø /¼ºtéë_têÔ©ÅUxï½÷â5üðÃÇDåöèÑ#ú—”Ù}Z{§1Qšº›0aB?~|³ÇŠÒåEQZ¤Š´D”ææ:sæÌ8 §Ž×•… %x iðàÁ1"ú­·ÞŠ=NÛ–-[Â+¯¼ÅmÖ2dHCÇó‡ƒÍ cØ°a¡6÷±i?T¥›3¯1®—(Íö/^Þyç8Xàï’%K¢dq$É;F'*‚¶Oôe'èŒ1"¶q–ÇÁ¡­sµu*‰ÒÍ™W‰–ŠÒô'¼î˜@0¦O¢NÓkÝ,‹Sɹ'"šzH"9õ…°¸ýÒK/…yóæÅu¨[êüÝwßëá°¦¨¿$J³<Û¦ÿÃòùèDDÊFc¢ôÕæ5FKDiÄ`„»¼P4|øð(ê<ñÄ ÷cDüO„¿ü[3¼NÄ}—¾ž¾>¾iMz„oîܧÓ6“(ýÐCE_?–2fÏžïoì>/þ,eâ #n?~<Þ;دAƒ]öfN~Ÿñ1ð¡ âüÊI¯‘³¿ì;÷¤¼(͹g¹üþÊ•4&J_m^c\‹(ßÒ½{÷eòäÉQ¨®$Jãw ŒQÏ´IÚfÞVe:~gjké >-í¶8eʔ؞Y–c”¦i®(]„sž‚¶¨K/DÛBÄ& ‚~$A…HÍ›È):—úœ6mZô3Ó¸:ÿ¶qšÏ¼ızÛÓ·„ô íñ àó¦ 톿í1PS¢4猾œqB~œÑŠÒåEQZ¤Š\«(à&Æ •ƒ îôú7Õ»ï¾;Îg¹ Dg‡Œ¿¯Ã2àü3¨§,¶ß³gÏxîY>õ5Ôë0>)õK L$¤WÄ™Îúôƒl—Èú)Älú8Ú¿©Ûšššâ®Šˆ”‚¦„ç¦æ5FKEilÎœ9qåË—Gß“>Q~Ÿ{î¹ØÓ7'4‰0Üoy[7üRd*÷vî ì3 0§(GÊH};"4bÿ32¸o0ŸeÙ/î½zõŠ÷îñì/óØÞä7@Ù&Â6÷ΑÞÜcð}ñ'Ò}‡û ÷Îô›ýüþ÷¿ŧ$Jó•·ƒ86¦Ÿv÷LSÂsSó£¥¢4íß:M²ñ[ðehSüM¢4û‚ßAc=Úí€ñ,Z´(¶'Úí†eh[´Àåzjïl“ù½{÷¾êد½s-¢4×+uK}à[r}7"3õA½ñP1Eþ›9´?Ò¥>þ&õy±}©_Ò|Œù}ûö))BØ'ÚýÇþçÆô <Ü ]äÇÞü¦¶'¥™Î}¦¶¶6Zñþ¡(]^¥EªÈ¥—-[odÜè¸ùâá¥åø‹ÓüÔSOÅ›'+éÕz8ãÜ8)ƒA!YÇ^Ǻ}Ñ”ðÜԼƸ¢4×&ƒR Ê)ŸË‘öL›%êg„ˆÇ®]»F‡’è%–'ò&‘-÷ß\‡¨Ý¶NSÂsSó*ÑRQ:Eq.“8À Á \P¿ Îh D?$1›ˆzê–߈ÍôSD¨ñ›úÃI¥ÍQw8úÔ-e'¡‚hjŽÁùMåiË4%<75¯1Z"J§ô>ùþž~: )< ä^KD(¦DD#ô².û†8(Òϳ÷b¦Ñ—ãs¦û‘‡¬Cä!Ûéׯ_CšîïôõüæÞ’‰ðw)éøˆÆITæ «´MŽƒà îeÜ?Xí!ŒâsðКèi¦±Mö _ÑœóAÛE°Bd¤<¦^íÞÙÞiJxnj^c´D”¦ñ7©;"ùÙ9xPÎtÚ%Ûæ!9c({ðU ü¡ ‚~ð-‰zÆ¿Ieða“@M;à: 'ó‰>Í ‡é~…8øƒÐÒwpϧsb›ÜðØ'übh@Ô`)3Á}ˆýH>~þ4â#¾pŠŠMÂx:öü‡ÅñGò)8NÓ2\¦„ç¦æ5ƵˆÒ´ üþ§Î©kÒ²ð?b"b2â%â3ËÑ®iψÖéƒóŒ§h£´ ®ÖOØ{ñ€„rhg¤!Úg—•!u\5i‰(Í8!™ëß3}c]ÆÊÃôéC©´1|O‚]x8•ÞÒ`\BäauH?“`D{¢F.ß—ÁôAô%iŒC™)žþ™ý¤Ÿ¡MáëRËðÀ.½5ÒhL”æ¼¥ôÌã!$¾ÿÕÆ1ŠÒåEQZ¤Š´T”Né 8üÜøÐQ#´á!éÁ ç‡'²8{„XP‘üŸ^½Ë?Ô)«(]¬·j‹Òô7ù\ΠĨOÎ7¹y2}t‡Ÿ·9øݘ(ÍÀ°˜Z…}§­Ðò:LÇ‹P¢(-"e'/<çi QšÏü€Ÿû@JÓà’Ò7á³âObéžÍ›1 ôÓ@ " “hÍ=€²ðQÓ>ñ (ÿ¡Ã$*±]¶™¥ñs“0•Di¢¹_`øÜCØ&QŠÜ›R,"ÛÇÇH¢4ÛL9f!‰Òì÷8þG„,Þƒår¨Ï¼ðœ§5Déô…àê5µßÔ>“(¿‘Þ£íоk1>¢M×ÖÖF±Œ2™Çz)Í\¥iÛiìÆ–ä¯*J7MsEi|CD]®?àwá!~)×7Â4ã\Êf@@CêKèéhô!ŒI’èœOiɾÑ>ø›påç#žâ3åè;hk)z›öJ»a;ű7ìŠ}{™áúM¢t@úã6æuÎïj}‚¢tyQ”©"-¥¹áÒi#2ñ1Rà sãK¯%‘(}4†›+‘Ò8{8AÜhq˜pÄqšx²Ÿ2àæÌvp¸éÜ¥üàP×8B´I~8´Íä´†(Mtí›Æ6qRhÓ´í…h‰ØQšvŽðÌò\ìÊåƵӖEiÎ-ƒœ$Dplôù>¥5Di¢æè¨K"Ò‡a™AQDÀ¾á 3xH(‚õG9Dgê‹vÉ>ñš-}}mŠÒ"Òž ï¤¿§/§OG¼ãw>ÿikˆÒôÛ©¿§¯F„£?ç¾J?M?œÒ(Q.ý:÷ö ¿µ1QšHCöŸòˆ:Noêáß2 !ß÷zŠÒÌãRžßt?IBT^”ÆOÁ/J¤ôIøÏ)ÅGŠ†”+¡ ã£à«Tjí!J#L²´Å$d"T2-EJ# rm¤T0©Ni_Do¦Ô1”EÛJ‚$þ¿¥?8Í¥iø„Ô#用fú™4†àú*Úí$¾pÊõÌ[y||0ùš´%úê@.ú4ú‚äÃÒ~y‘ê°8Ÿ2R„<ýÛ§¿¢HWe,ƒÿÌ´4^gm’7¯æ—ê„q×)MHý„oOKßN=q¯£¯HÆoÆ4M=øS”./ŠÒ"U¤¹¢túð Žy¤|~8Ïéc 8VL#‘›7=–à ¢ƒ'j‘uqÌpÐùÍÍÇ×ÓˆùÑ~ ͱ¤íCÛ¡-r#OVTS”¦íåÛw2¢ª‰BaPHúœ<9Ú,m”ýãÁJ~ŠÓ’¢qq<ÙgœBÖIí§G¶-‹Ò8n•êãOõVMQ:},ˆú¡ù͹¥ÏáAç>‰ôA ˆá¼³\ª¯4ðc}òêÑ&¨S¦áüÓDZÄö ñ‚>Šè»äœ2¨à2ü)ˆ•ú{úY¨¦(Í2܃ӽ“þž¿ôËÜW‰òÃŧ¤ŸfY„=ÒsÿEã~K?“@Hbâ5ûA¤rú°-ÿsá7ó,¹¯ãÛ&Q:}ø6=Ä$ò}Bp¢$ñxø™DiîCÜCˆ„äavºÿ 4ó"ÇÉ6ˆØ&~Çœ¥&Ó‡ ˆKÖa9êA®$ù‘Å6Ì=ª)JóÐ߃‡ ´EÚtK1öüê”·LˆWÖ¡m0®âÁþ oß!6Ò¦X_ß…1¿Ùc7Ú/m;½¥Çv“¿š¤2Í¥¹'Ÿ2õ=i ‘Þ^¨m‹ëº`9|T|Jêšþ#}W©¶¶¶á Rz ú'¶Oôs~>}åÑ?ò-EÜç#¥!=aÿi´%ú¤$n—®uÆŒ èg1òwcÔ;ý;QïüÎÏçZJ±*¡(]^¥EªHsDi:Xn”ë³ÓèÏR4RÊH;H¢4ÑK<ˆÀq)ôÍ©?¯ô²š¢4ý1}4÷ä|ŸO¿Œ¸’Þb¢Eñ9é·1à$’{—>=ÿAZîLK9 ñ=™–îùÜWx8M_Ïqq_ça%¢p_à~‚ø„?€ˆÿ›D?þâðp3‰Òø¸ÜCX…ßé^CÙ¼…ÅöÓ6ñSÒ6€ÊHoçp/f™´ž\ ç±9߆ñ-Sû«¦(M;§-‘Ú±‘¶øL[I¾#þuJ°?|ƒú§-âcàw²Ï€¨Æu@»¡ÞÓ˜‹ë‚ýGä¦]¥?3ž£ÒS„¬T¦9¢4×#"pcãú…|¸"´ÞÄ H ÿ”10õˆš°D‚¿›üX¶™>™Ï˜#ͧñð•éSè?©‚,ûÈvypG»£m¶8ŒÇˆ’N¢3çëžû@>§tzƒ›·W¥Û'ŠÒ"U¤9¢4ÐA'g:YcÐ!3PÉ¿ÒÉòyç.¥kÈ—ÃÍG+ "íÚKœ{EQ’6Å<^µjŽ¸ Í¥¡Rû.¶sÚ%í3m;?¿RÏÏça *yç8F $ˆh©”Ǻ-Á1s^R½ë†óšû0òš#JC±®Šu_Ž}Ì?è*Ö/ÿÓ'Û õNå—­Tß•¦‰ˆ” ú8DRúhúöb¿È=€y*Íé ›+JC±¯oª¿e:b vµ~ºÒ4À7Îû²‰â½£¸~ñ7WZ†}å^“¦å·SÜ&T*#Mkν¶½Âù¡}ÒNóUÒ< 3¯¹Vš+JC¥zlª-%¨OÚFŠÎÏ“ÚiJw˜/£Ryiš4MsDiHç³1+Öe%è©CúÏÆÒVR}Dc[MͯÔî´_¶ÝžÄè<7*<¨DÜOç‰~Ÿëšèò¼!87Õ¿*J—Ei‘*Ò\QZ¤5À©J¹¼Šà8 æ¦ülÍ¡¹¢tµA%†oD»•€ M”LqÀÛ¡>¨·JB2‘$-©7Ú@sEii]ÓŸW˜óþž¿Ícš+J‹\/hwø,EŸ$ù © 7‡–ˆÒÒvh®(-m®uîC•îi-EQº¼(J‹TEi)37Š( D"ðšékxÝ–7ʾÝH(J‹ˆ´¥¥­£(]N¥åZP”./ŠÒ"UDQZÊÌ$JKóP”i(JK[GQºœ(J˵ (]^¥Eªˆ¢´”E鶇¢´ˆHû@QZÚ:ŠÒåDQZ®Eéò¢(-RE¥¥Ì(J·=¥EDÚŠÒÒÖQ”.'ŠÒr-(J—Ei‘*¢(-eFQºí¡(-"Ò>P”–¶Ž¢t9Q”–kAQº¼(J‹TEi)3ŠÒmEi‘ö¢´´u¥Ë‰¢´\ ŠÒåEQZ¤Š(JK™Q”n{(J‹ˆ´¥¥­£(]N¥åZP”./ŠÒ"UDQZÊŒ¢tÛCQZD¤} (-mEér¢(-ׂ¢tyQ”©"ŠÒRf¥ÛŠÒ""íEiië(J—Ei¹¥Ë‹¢´Há†{æÌ™âd‘R ¸Ù6¡ÞêêꊓED¤DÐÏ#J‹´Uð1¥ËccEii)<ÌP”.'ŠÒ"U„îñãÇã“~M+›1Ø=yòäÓµ×PoDçiš¦iå1úù-[¶Ä~¿8OÓÚ‚ác⳧kmÛ§·-5­¹vìØ1Eé’¢(-RE¥7mÚ6oÞ¬iš¦iš¦iš¦iš¦µÀÐT¥Ë‰¢´HI¢´¦•Ñ’ƒ µ-³Þ4MÓÚ‡ÙßkmÝlÃ峼Ȩi-1Eér¢(-RE¥kkkã«gšV6«©©‰¯§k7¶QoÅiš¦iZùÌþ^kˆi.ŸY¯ÚµšŠ¢t9Q”©"ˆÒä?"9ý4­L†ƒpâĉ+¦k7®ñ‘ê<Åyš¦iZy,åã¥ß/ÎÓ´Ý;ácÒ†‹ó´¶k)709¥‹ó4­1£Ý=zTQº¤(J‹TEi­Ì¦(ÝöLQZÓ4­}˜¢´Ö–MQºœ¦(­]‹)J—Ei‘*¢(­•Ù¥Ûž)Jkš¦µS”ÖÚ²)J—Ó¥µk1Eér£(-RE¥µ2›¢tÛ3EiMÓ´öaŠÒZ[6Eérš¢´v-¦(]n¥Eªˆ¢´VfS”n{¦(­išÖ>LQZk˦(]NS”֮ťˢ´HQ”ÖÊlŠÒmÏ¥5MÓÚ‡)JkmÙ¥ËiŠÒÚµ˜¢t¹Q”©"ŠÒZ™MQºí™¢´¦iZû0Ei­-›¢t9MQZ»S”.7ŠÒ"UDQZ+³)J·=S”Ö4Mk¦(­µeS”.§)Jk×bŠÒåFQZ¤Š(Jke6Eé¶gŠÒš¦iíÃ¥µ¶lŠÒå4EiíZLQºÜ(J‹T‘æˆÒçÎ uuuÑŠËñ›éçÏŸ+­›æ5¶.ÆrÅu+Ys÷åZÊkl´Lqz%klÙæl'Ycç%?=oW+¯=[KDiÎ/m§Ø®´Öµö JŸËÚX8.œ­0¯¥vöì™ëRN2ÊcßÎ×Uÿ:hɶ8F–åÜçµÔ(‹m§7×Bݹ¸/y»Pw}ö­)Kç«ÚÛ©dçÏÇXœ®iÄš#Jk5oõ•©?¦ç]ãº)öy«´OE«;W¿ü-ÎÓ~zÖRQº±1BÑÒX†¿yŸ4?–È[q™Jë­ÒX)ùÁ•ö1•[œ~áÂ…ŠË·eã<(Jk-5Ú¢tyQ”©"W¥q4öíÛV¯^6lØ—MÎÊ©S§Â–-[ÂÊ•+ÃŽ;b9©,æïÞ½;ÎÛ¼ystÜ’DÇ›6m «V­ {ö칪SÃüýû÷Ç}Y¿~}8räÈeû’Êc{»víºÂÙ*ë:t(¬]»6ÚÁƒ¯ØÊà¸Øÿæ”ǹâ<§>|8nƒce™âvòËrŽ7nÜ—Ý»woƒcÉñR_ÜìòÆ9¦ŠeiÍ¥S[;wnt(»´ê[ÙEésçΆŽ…‹·†#'N~`Á¥)1§¥Æ¾:v"¼·ðý°pÓÞª õÛ:†fÛZt•m±ì±§ÂÈ%ÛÂÖ}Gâ9,.Ó;}ºÞŠÓ›c'OWìçm ï^´Áó·„™ëv…£Ù>VK0æÎÎç`÷ÁcUÛN%£n×ì 3Öîºbž¦}kŽ(ÝÔ¼–X¥kžëhÏ¡ãaPv ¯Þz É~èZmÚš]aÀœÚ‹}ÅûaÈ‚÷cÿñÎÜ-Y_»5^×MÝا-{Ž„AÙ:5»Weµk³–ˆÒø™ŒuŠc„by¾ýŠ+BmmmôïÓx‹±M¥qãÖc¹;wÆu·0.*Ž]XŽënñâÅqü–æ§qÔüùóâE‹.óÁ˜Ç12==RÖ„ ⸫±qM[4EiíZLQºÜ(J‹T‘¦DiŒyóæ…{ï½7<øàƒ¡k×®¡W¯^q$ž¾öÚk¡sçÎáᇎGŽÙàMœ81ÜsÏ=qŸ}öÙAǪwïÞ±LÊîÒ¥K˜9sæÎS~_–.]ºuëxà¸O={öŒû“{Äêgžy&–wÿý÷‡»ï¾;Lž<¹Éòp Ùö}÷ÝË}衇¢ØD`àÙÏ×_=þ.–ƒ¥eqûõëž{锅ÓǾ² ö}DZ+:pIâ‰'.;çsæ̉‘8ŠœcÊa>ûÝ©S§X6icÇÚž­¹¢t:¿·ß~{§óç2 Šë4ei–®§5_”&º A¡8˜ÏG]1¯)—Éý_éwSFY•ö)YŒÈ»p.Š{¿qÏØ°mÿÑ¡±uÓöë÷òéX÷«Bß)³6\x6}lÅyX8_†/ÚþŸŸ Ÿ{|j8zü’ÈZé8ûÂØ)½iXß—ýW”ÖZj´Eéò¢(-RE¥q6Ž‹!nÎ81Ñ8(3fÌ;vŒOÛqÊÆ:tèjjj¢{çw†ñãÇÇ'ï8,8Jï½÷^Ü.âm÷îÝãö)ûwÞ‰ó‰"(Š«ìË°|ß¾}£3E‡hK98BC‡ÎŽ!ó²Jåa8UˆÍÀ8lD$ócÓ8îG}4Ü|óÍá­·ÞŠû],'íÇEY8‡/½ôRÃ<¶ŠpŒØyàÀxþpì¸q1Ë/Ë<ÄiÊ}ûí·c º³<ÇG=püD^ã”âTV:F­e¢ô‚ bûM¢4 øËõAÎ;Ü©ÎÒúùßi9ê 'ß” -³æˆÒ ÊyÕÑôÄÉÓõ¯@_|3=¾ê|á\4r™–^“fZýﺸ.ÓX6¿~^ˆMå¥mççç-•äx¶ß§ëç÷—߯lkÞÆ=Q\L¢t|-›¾$[÷Ô©úãÉo?Í?Îúçë=Íû×çç„Û,Éz§ú×j9¶ÁyÄ]Ò±±­úsS·“L‚Æ¿¿4/ŠÄ½w\˜¸rGþ°^þõv–%íë±oq?ãvëë03ûœŽ±AtçÈ0ieýq°lª‡Ø õÀ¹c›'NžªŸŸšâþ^`qŒÒE é3a]ø»¬œxjFxjüÚx^F¾2FBw¹*üÇËóÂìø’èÂñ¬Þv FpO[³3 Æ?xmAÜ„N—…W²sÉr¬søøÉlû‹Ãøì\?;i}xmææðüä ñs—Öì ß~u~¶ÏÓ¿dÇ6øâùˆ‚gVƪ­õÛš±fWøZß9ñ<3=ŠÜçêSœt¼<ÖÑ?<53Öáoß7>nÈ]¢…¿–KÎéß<9=ô·6;ÇzÏàeÙòëÃ÷²ýG`ÿzÖVˆNGŒé›ÕÝ'š>Üe\øzvŒ›³sÉuüÚŒÍáûÌŒçíáìüZ„c§ýë¡~ÿë’œ Ogõó\V&õñÙlÿi7<ˆbRÝÙ°òýýñܳï_xzVè?»&ž^Ýÿø}£YûA`ûËG§Æ:ض¯^¤¢ÞïîËöóþ¯ÌKk÷…GǬ‰m矲ã#â4‰×SWïŒåpÿöɱi µ»ÄÓ²ù´[êµëÐYÛÅ*Î Ë~?›Ç>q=ÝÜQø묜ÏguÑcÔª(t+LkW³ÆDi®“ýGN„ogmø×ïþ¨ûäÐ-kƒ´©ÃÇNÄk‡K\ý²k”u¸~x¸õ/ÏÍŽ÷®+ÞÆ@x¾3»/pÍòÁ‰áGo.Š}%×îô¬¯#z™þŠkñÁá«bÿ²=»¦¾Ûo~˜³aOìû¹fÖn?~øÆÂø÷¡«Â°ìZê‘]‹ô#r„\Ü‹²r¸æ7î<smÓ7sÌÍÉýßÞ¬9¢4c„aÆE#Ðÿ°’ÏŽ?{öì(òƈ¡þmM"‘'OžýÒ0Œ‡KÒPäßlƒí±ÝÆDé±cÇÆm!HͺEAc›ˆÕˆëiÌŒoÌ>S>㩧Ÿ~:úÌK•Ž³-˜¢´v-¦(]n¥EªHc¢4¿{×­[.ÒDóòû±Ç ï¾ûn,‡uptáú÷ïó3#êá¼°,¿‰²4hPC4Â/믜!€#'Ñ5ï|áÜðJYŠbÅiÂAÂ#r‘>~ÿøðÑ®ãÂçzM oÎÚîº"üò]£Ã<8)t´,,AT|tj ˆn%j–Ô ˆ lƒ¨¼;Þ©<&¢ù¿Þ2,ŠÌÌC$ù­nãb¹LCTž¼jGøPÇQq[#–l݇ï…UC„¿jfxw^m„Y~-³=‡G!˜ßˆÏˆ/ÿ3ÛßÕ[†õ;Åcc>ç[³kâ~#Fî9x<Ûÿš(ò 89Q‰*|`ØÊð{Ùö.^œ¶1üv·ñQÀAìý»'§G†2žÌÊC¬ÉG>6fM,áçí95áÃÙ²[öÉŽ±.Ü5pYLUQí]EÞ_í<& ¨ßmAø•N£ã¹CÐå¼ýŒ2hÁ–øêù/gÇÁþ§È]Ä}–G˜AtùÍ®cÃûÙ¶>ŽŸ<#À‰ Ddá!Àï=01žãÉ«wÆóJ= XÑØoÎ"+màž~¡SVßlÿ[/Ï ÿ+«Ú™ëv‡ÉÊþÃlÊݟկﳿÏg¿ßÈöåSÝ'EñŸ}EŒ§ýüáC“B×!ËÃæ]‡ÂŸgçagö†=1_옥ÛâoD°ƒGOÆ Ddþa¶ÿodm…íP>#8vÎׇ:Œ w Xdmâ‹ÏÎ Ÿ¸o|l_<ˆ ½ÐæH3rO¶ÍŸ¹mxl;´ý_»gL, QøÅ©³ãÀPî÷ú-¿ÐadÜgÊý³¬Ðv~M´ðýðÍ—æÅu÷fÇSdõûŸ¯/ŒB5õÆu°0»> ~µóèðñ¬í|å¹ÙaJv^»díÁèè#WǺ'7.m%> Ê®7Ú>ÇÊ>ÄÈÖ‹©šÖ˜5&JÓ®x€r²´±‘K·Æå~”µÙ?ÏÚöÛYÿÝ{üºxm¿:}cØuðXlß·e}é˜åÛcÚ ú…›öÆÜõüÈ”(b¿•õmˆÂôÝ7÷_Å\Úê#ÙõIŸÃ5Ìí¿È¶qûÛKêû¬lþãY?óé¬oØ—]?Ÿõ¹ð¯- ƒ³mÝôæÂø;½9Á1Ei"§f×)}ã]—†wç×ÆmÉúÂe[ö5<øüBÖ¯ÐW°ÿùÆÂøà2Š–Ùú<ìûPÇ‘á÷³>ƒYVlÏÊŸ>Òe\ÜÇþ³jÂ;së·A¿·ñÂœx\³ý@˜½~wœÇÃ6îy<Œúß¾?8ž#ºq¯û­¬ ÏȧÒêíj¢t²4F  Á¶8æJÆõ€ØK™ŒSO¥·'“™¶Éx‹·8ñC™Æx‡± ÓŠŸ!<§qNÚF¥‹hÄX‰·½GrˆüR^ilúôéq;ˆ†8G8O_*/¥Å €uS$5†ÀË6p“Ŷ^|ñÅ«ŠÒëààEi¶Å+t)6Â'ÎóX°8¢ùóÍt"#ØÇ1cÆ\挲<HíÁòÅýÐ.YS¢4ç2µû¼(^ƒ>|xÃÇ4yXÒ­[·øª#uƒ#žRÑ RS·¬Ãöˆ¨¦Ý±× mŠ|ÖUóìj¢4b"dŠtÛ¸ópO_šº! ¬¤xøç§gEAƒA69—®ÃÈ0wÃîÀ ›h_–¯ÍÊa>F+ÓˆŒFäFð@A`ûƒ‡'ECð˜¹vWøûÿÿöÎ3È®ã<Ó¿·j÷ÇúÇf—dÙkŠ)ÉZÙr£œÖZ«,9É^¹¼²ÁLf‚bΙ &‰b0È9AäœÓ`ò À¤ÞóôEg&^à9wž§ê«™{BwŸ>©ûí¯¿39ìÚÛtÚá1ñ¤4-º9Œ[º3ŠzkkDᡤnß‘ÓùÝ3aMøÄðšø­•;ö…G&¯b!llŒbbÑxß"š¦}7ÖˆBã«óK¢B#â+ൊ ñÆÂmÑ3šò½4ws ű½ñPø«‡çĸ¦Ù±$Qã> ,Åoß9=Ü5OÁ¶è)H¤¼mj ¿–ùÿ'æG¡1ñ›sú›Þ^~ç®é1M„òBìøÃûfF±§éT^Q¤Í¶§>–žž¾1þž·¾!|jÄøXßé¸'®¨¢)žx%?])lÇÚ]û£È€Â5ðµ;¦Fï¾’skÒÿ(Ë; ;eçø¯›×ã-þ¥‘S¢ˆZšf2Œ_¶+;²º>n}gUøÄ°qaéÖ½QBÈÿ“,-òCØFÄB4ûå¡ãÂÿ{~a´6ÕŒu÷îâ1n5uüì¬QÜi8pô´°Ãvdç´9~Äà'¦oˆåÙ“]7×Çúrvl;öŽùr­—ê¥=ž‡8h’¥…gøwž|/zŒÞˆSJl»dËÞ(þ¯ÏòÁ˜¸îtèåÛ÷Åk–kwsvO!æ1ˆÃõ;ˆãóļpË;+ã5‰O}0òçÏß~r~Ø‘GB°ñÿ\?H©U¿u%Jc1¼NfßxpVøé,®ù=|?sýÄ0;{æ¦û•gÖogϸnï¿&>{HÑvKÃÁx¿üó3 ÃÈwW•î«ìÚ$ ÏfŸÆWíØ>•=cø-¿Ÿ˜¶1ÞÛû›JmfQ L·feb xÔÜÌ<ùÖ©ÁKž9¤Ï !oRHœ¼(Í}óðOœ†üüƒøŒd=3 HwØë<Ó[ã}ɳ/lÊÄL‘_ÍžÌâ8ðÿï׌ /ÍÙSÖ/gÏò>v*¼ Â9¥¥<Úbh“ß»{z,¤ ˆØ̆ ìßÎîñ¯eϵáY9³›?t1`š)¢±ÞŠÒÐnïN”Nýú*x#ô"2ç?ÞN{•ðôÇŠ³Aq~¡ÿ“¾óƒ Mß&í‹%Qšôqö¡mK;•ØÕÉ ºX.òd'é%g#úEô³g¾ÉB߆>Ƭ®B(~ÜMQZ+Ç¥«Ei‘ Ò(ÑØAF襃Çg•G”¦A’¥zè¡.q§dù8G~Äœ4‰ ?Aà Oi:¤5}úô¸Œ)k4ºhä¬X±"6ÎòžÒ,KÂr*i¤ ¬ChdÔ>¥Ç_„c¼\ó£ø*Ö#$žQ:5ü0 Ñ4u؉ø™òÉ×}jŒ2½/}¬‘åiÖsÜìŸìOœCC dþÞÝ3¢pdž[~ðÒûQÄ =‘+^]Óé‡Ò4ñÏ\?!ŠÀL5yÝ›òš~ùÚš(ô ‚"À^tÓ¤8H@]c Äð1¦t{É;á™kàïŸ>1¬&ŠKä÷»YýãaŒâ1áJ(Ïç³ë¯tÎÂéé<|=«/á•Û÷euµ8\…@ÞVŠóŠøÍyFœÆËó{£Ååy1 ‹1¥ožC—DïÍlÙƒÙõ™<'ÿõùEY¹J^Ó¬{!»' ƒÎö«²¼‰IÍuÉà u¯»§Kׂ^Ͳ±¨-†‰¢\Ì¢àçyĬ‹²cc}ÝþÃñ˜I“ÏãŠì9Ã{<mÂ{dÃîñùÏ,™?ÊžuYšüfDz–k­¢tjKÒ¾$]ÂÒa–¾‰ƒ£P¾OE;“v*ý.î+úw´]Ù6¥¾ 3Ž2äAÌjâ]#ÂH?ˆ~Xg÷õ@0Ei­S”®n¥E*Hw¢4éÓ§G1/f¶£a’Mˆn£F:½Œ†§ä=Í_„Ub?§é^)]>„ˆÐM㌗>Û“fjìÆ<…³X½zu ñ'ÀSÁ›Sâˆ;ÒÁh<.„rçÓCt$Ê–>èñ?yP'çC”¦œ”—|ð¦–á‰KÃéuäË2òaÿ­[·ÆP¤A9Òñäó`ªqܨïÎΛvƺ¥©7¦6R—üôx˜†k ®®=P¸žRÜrΞ%¬Cx¦AŸÒ`?îÒàڦᮧtï­;QšiÐ3VïÈwWFZ„ƒß¿{FxffÉÓ Ï¶×n/‰i™áÙIg2~#J "¬ò‘;Dz‹?6u} ç õLú>ˆAÓW×…¯þdjLQ1žxQžæ“ÍÑ{›xäŸDUc¦ŒãÉGÈ<¶ŸÈå…WK¶î"¢(¢M[ õÜC{Px "s¾?NŒøxžÒ\ü „(ÃoÊ·b9ÇŽ(ý›:Æ3ý‰¬ ¤Eac>Ø…U{9Î$J#*Ý|Jbyü¨YvŒï¼u 1Í7ms³ó@ÜZÎøì<ìk:c¶ÆSžu ÓR|ØÅn6¸¶RÌÙ%[öÄs‹(Ív r¢4±´ó¢ôµ£ñp,­ãZe€QŒíÄ¥©/bÌR‡÷Ž_½óYF¼D7fåa;àò#Æú7›Ÿk<‡ˆC×2ÏffS™¿^÷:1ê×G/âØÖÍÊÄ;â±i”©5Æç¹È€ÔQº-ŠÖ„:ýAÝ‚(—5r|§dÛ§mx¿A9ö'ü3V(#å^eÆ žÒ„ûá¸y–0ûæ«·Oƒ€¤ÕxðH Ãó"ƒt¹<!ÄûŒgïÔÿu{é™Éû„çúÜ;#~ÛtBEº£s;¢4m|Ú‡N8ÌÈK+¾m“fàѯ@|Æa(9ñžñN} ½ô}Ÿ”gÞa&µiYG¨ú@ô“RºÉò¢4åäú€É³š~yУ\ô™hOÓOˆm_Ei­S”®n¥E*HW¢tŽ¤ ¯Áïä…œf>g1Ûž¦4Bð0EK1xÙ—}Ø—40bõòAŒ”7ë’•o %CȽå–[ÂèÑ£ão¦†1¥tØañ˜|h<‚qäÓ¢<äKãŠxÑé˜ùŸexÏæoåŠÒìK,mDfÄMÒ¤lˆÑ4(ñ0àå…ç3Só؇:aŽß©¾ù?ÕNꡧòh]‹Òœ+ðxt B¿¹ž9ÿ ¤0˜Àù¼Ú¹җʹ†éXp½à‰Âvœ[¼¦YÆ9ƒôñLDëb#_ëܺ¥éð#Âá1Ãw„ö(„á9÷ܬM±ãŒp½Âr¢4q:£ÛÞÆ/ß=«1âé><¾ÿ¢á&ˆ3¨ˆÂo/Þ;öxß!‚“Öœu qº8 b'Â,å¡ãÏTh9”˜×xø¶Ærø“„ï8t4 'xÝEï½l_bŸâ±Š—.žðˆLÅŽå_¶3®'¼LJЂ×[kö?"ñ”ï&G;Þµmñ#ˆ6ˆÒx~sŒõûÄ4XN8 Äl>êÇ>IÔDôø¬>cŒí,=ÂRP¿ˆìˆ7ñãÙñsìˆ-1ÏlÙO³ã'ÌDü˜V–Þ{ 0ŘÌä•D’”uÏ´sB‘j„úpâÚ(´`wÖ¬ ÿãÚqQäåÃ}ÄÍŽÞïYÏZWÏé+ó·Äsƒèƒ·ónB¦deBè'œÆ}Ò^úð߃³¢˜ÄùÅS2Š¹Q€m!,¢(ýFç¢ôˆ7—‡ãÙ5ÉTx–q-\tˤ¸=-¢åçØ8Îgfl ýèÜX§x'Eé$ÿ•˜âÄåf_B ~…Xã1ÜÀ9ŠÒˆÚ|Üqø¥á-'9æ¶x}›–Ð%ÔEQúú·VDOSê“óÊ5ˆHFý‘Œ³¯ÿ½ñ¹_8§ŠÒZwÖQšÌhkka—¾Ý/ñ±?žÙõHœ~„Û9këbh¼™¡ùÄÉø›g,÷ïfÊðÜâfÀÌæŸCÄtG”æÚæz&ÏoãÈóAE¶¡LáYÆLJÚžôÍñ˜¾ý‚Ô‡ O†cK>¤}ú;´S•~³QùÆNgžÒˆÒ´i“c íXÚ¯ô§ø?ï“Dif¦þWœ1f²/@;Ç#õää3LQZ+Ç¥«Ei‘ Ò•(M#ƒ—^zilÐ0ý‹ÐÄ2F`¥±‚˜‡ Ç(<Â-‚ÓÉØ/Ô+®¸"6\Ø/íKøRÆìGX´Ë7²òF£†ò ðÑ ÃS˜† 9 AÄoò ”ÿÓË ÌÅô )džñ?Ëò£ú”…Æ ÇÅtŠirÜxð›|Y†ÀIÚiñ?¢2i#ž2$~H°ùˆØùz¡lxP.¶¡Ñ9}zéKÜÅ2h­;QšA<íi€s­pm¥qÒ¨§ž¹öYÇ6\¯4²¹F)’?øÁNÇöÆR®LhÀãQÍ}’÷TѺ¶îDi<¥ñbÅÃ!÷ïž~/üÙƒ³¢ø[YGAe>`;ð™æâo'öqÉó”ô‰-L(‰ï<5? ÄCfŠ6écx¥þÇËÞÞdtÚ‰íü_¯¾óäüx%èÀó›íñ°Cà =>¤E˜ ÄBŒ0µ›æ!t°=uxü"Z#J3ýúW†ÿ'Û¸¼Ǭ'd^©ˆƒÃßx`vÜŸP×fåç8CO‚@øýçE>UÂDüÅCsbYØþÑ)ëcý!ÐàŽØ‚àñK—¿=êžœ¾!Æ&ŽtŠ•Qx–# —>BØþõ¹EáßÿøíèíHŽççc¤®ðÇÃ]ààñ‡·áéù(žáxãåó¢o|kEô–F &F7Ûá•ûÍ̈àã…|tk€íÇ™ÒÎ4~Är¦šoklŠ!+¨C<ÊÉŸó‹ \Š¥ÝÞX´-|b踸?Û3À@X „h¦üSg fp¼W½º4 ÿEQáèï›ë” âN¨“«<«Î+×yàIš V&éÅ8²§Déuµ£H„PƵ…È_ª»ÑS‘k”ú§¬l7]ýiQúîñkâñr=.nKÛ336ƺO¢4a`ØŸiúx£sÝP¿|ø‘kY}|úúRˆŽ$J#¦þ$^ÙñÝS³:ž÷66†û³2¤ûéï³úFDpÒÚçWÓ:³nEi,»çøx ƒsC³ç᎘‘ÂÀÏLî/î3Âz>öaøVöÜ$”÷÷#ƒ# Z!榮òTú¾ò“)qà$ïÌýÈài°Ž&Äæÿ™]ßÖ¥±ÏbÓs/¤øÔ1¿l¿ÌîC0y0à—îî#ž“_99ŠÒ< п~ÏôŸžg÷=÷1³xÞ®£Cdee€’ð¬gëóÙý–D醃G³çݤ8¸IÙ(ïª,bH3`™òÀ#š<(Ï^Þ‰ÿáÒ·ã`(ïâýÿ»¾=±óßÐ:Z_EifçѦOmAÀé%…Žãû=ô›pÂÁQá˜vg ÷x4ÓÍ·mIþ¡ÿhsÒ¢ý‰xÍLͼ0œDiò ]êóЧÀ©ˆØÐôy’CLZ‡×5³aq¤Ás›²&¯lþ§_Fÿ…}i;“>ílEim°˜¢tu£(-RAº¥ùÍ1D6B Öb„! d4Dˆ6iÒ¤°dÉ’8‚N…Pì›öÃØŽ†â¤ÅõXwÂËñÀ¦a4sæÌØLþÒxÀƒ€²,[¶ìtYŠé$cq¤ýÇø¿¸=y"PRî⺢±íÖ­[c;Ë!ñ9Õ_jÈñÅk^`x&ðgõBc35$©g<Ç»ª'íŒu%JcÔ'ç báqMÒA¦^1®a΃ Ü'¤Å9L$œ öÁ3$]œWÎ/× ݤ“A:\¯ÝgÚÙÖ(! "„!n"\"D#žáYF‡|ÂòÚÓÞÌزlÝœõ §c@³Œ©Ó›MÝ=îH#}ȉíÞ]²3¦KÜa„1Kw†¶îxÂàqMlè(dÛãaJèq<ƒñÆeÛ#p’çô5uQà F3q¦çmh/Ïß½o/9%†~ ]¼ØÏ á‘ŽÏ8B1 22MœãÀ3á1:~pîÔuGØ BŠp½"¢¼µxG:ñÊEì=«®³cET¡Ì1Fp¶o|3Û²'/XÊ»xËž¸=¿©#„Pêè‘)ë¢pB~”mþ†,¯µçµ­ñPÌ >žœ¶1¼8gK´ItX‡ÇôÓ×ÇãÄó—°ÏÎÞc¨GAvJVWœ<Ý÷7å¢Ì”u£n‹õ‡(ƒw0q¨I“:æ¡LœsÂz°×çž-þbÑöðZ¶2Îb|ò*äïªû³Yù¸ÎéÂy ÒÃ[1ýæãˆc—îŠeimmŽé¬É®‡³óÇ5±;;&,„'¶£~âù'ì\s^É“tãy8µlŽ“2³ OzöbTö›2sMáMèM>ºkïá˜?ÛQ>Ò¢~ •Pº‡NFÁŽsŇبÒ" ×8÷÷±^ÒZOÖ(qííÌ®G®µÉ+wÇeÜ[x–ƒÐ|À“{CDæã†Äÿ'¤ÏB–§g%Ïuž{xH/Èž/éÙ˜Œm×ì:·á~æ7Ï7îyÖs­ólà7ÏÂôÌ£LäÌ=õVöÌÙ—Ýoy/c¶åþ䣭éC´<y¶Ž^¸=>+¹ï¸ÿ©šìy·rG©üÅ2¾Ÿ¥Ã3ŸwÏEÞ1ɳœôyn§²Éc[̃wƒ¥)]¶#¼«fÏ¢ôܧ®xëH;c}¥S_g”â2Úõ´éËж¤ÝOÛtÞ¼yÑ)%õuØžþ9¥O8é°/mÐâ‡1ÊÌ}‡¸L¹óíÒÔÏ¢/‚¸–ÖÅ6C–¢5ý+Ê•ïñ?¡éçá}M_ŒvsOý¦«qÜŠÒZ_ëFQºzQ”© ]‰Ò ™o9oEOâ´Œ¿)ô»h¤É6ÅåÉ:+GÞ’GjJ«³òËÒ屘^1ÏâòÎŒ4:Û6_¶üúTÔcoê%¿}1ílëN”ÆÒùJ–¿fŠ×Z:‡=‹ü~¬KyôæzÔz¥±ä=Á²tôæŒ1—K;xàâQÚVlJi´ŸI£Ãúæ’7hªcžüNéÆíO—éT9Nýî°>æWŠEÍþI(ÏOô\;å‹€ü¹›&•>*—¦o·•¼“S9JÇÜÓMËRúù2¤|X÷k)ms:¿B%c=Û%aâtÙs‚#ûæëëX'¥|Ò¶½Í«³:KBK1íR\ç–è-ILi>b'q*¼°ræ|äŽëTÞ)ÝTÇÕMÚ·ƒµuÌãô¾§ÎM> þÏ×WºÞÒþ¥ßg×]q»XnŽûTÚÅtã1åòM×J:wéXSéXÓµ‘÷ -ÖÙ{¤p>NYþÕ´®¬'QKÏ…üõWº·J–¿æÓ; t]ŸyÞYwæ~ͧ—·â3®³mù]|oó.>°Tîü²üñ±ît¾…cëÿ©gKzn¥åñÝÖÞñxVÜŸ…£”NZ^|kíºÞŠÒXg}„®Ú–ùvd~ûb?"oÅ}»ê×`]­OûÛ¬©¿ÇòÎö+ö;Ûf Ǩ(­õÕ¸n¥«Ei‘ Ò(­iÝz¥µŸõF” †H@˜b®ÞY £PÜFëh*ˆÒóø¼è¥ÞÞ6p;Åš6¬7¢´¦}\­¯¢´60LQZ+Ç¥«Ei‘ ¢(­U³)JcLÉn…)ðzëjÚÇÛ¥µlŠÒÕiŠÒZ9¦(]Ý(J‹TEi­šMQzà™¢tGsútß:+NO×4íãgŠÒÚ@6Eéê4Ei­S”®n¥E*ˆ¢´Vͦ(=ðLQZÓ4mp˜¢´6MQº:MQZ+Ç¥«Ei‘ ¢(­U³)JvíÚ¥õÒêjkÃìe›ÂEÃ'„‹G|töÙkk›sÖ„†ºÚ³Ê8P­®®.ŽHu£(-"""""""2ˆAÞ¹sg8zôhô––Þ±¾®)\<|Bøâu?2ûì55aÖÚÆbÑ,­­­áÀQœnii)®–*BQZDDDDDDDd‚][[Ž;V\%= (]9öìÙöíÛW\,U„¢´ˆˆˆˆˆˆˆÈ …¾x¥¶··WI(JWIã!Õ‹¢´ˆˆˆˆˆˆˆÈ …ؽˆÒÒw¥+ÇñãÇ£¿T/ŠÒ""""""""ƒEéòQ”®ŠÒÕ¢´ˆˆˆˆˆˆˆÈ EQº|¥+‡¢tõ£(-"""""""2HQ”.EéÊ¡(]ý(J‹ˆˆˆˆˆˆˆ R¥ËGQºr(JW?ŠÒ""""""""ƒEéòQ”®ŠÒÕ¢´H?°gÏž0sæÌðæ›o†‰'†­[·†ööö¸Ž¿kÖ¬ …½Î?aíÚµ¡­­­¸ê¼ÂñÖÕÕW Žë\ó£N6lئOŸÞã¹Ø¾}{hjj*.îW¸fêëë‹‹EDDDDDú„¢tùœQú ÙþÖòÞš¢´ T¥E*ÌŠ+ 7Ü~ò“Ÿ„Gy$ÜyçáÚk¯â4BhKKK\6iÒ¤â®çiÓ¦…‘#G†?ü°¸ê¼qâĉpÿý÷‡W^y¥¸ªsçÎ cÆŒ9-Η /ªûî»/¼úê«ÅU}bùòåaèСñ\pκb÷îÝáÆo K–,)®:o,[¶,ü⿈ÃΠÎn¹å–ðî»ïW‰ˆˆˆˆˆô‰rEiúbô]è}ðÁÑV®\ûLçÚÏëŒC‡Åô»ê'}œ«(}Ñ°ñá[Ì ?õ~¸dÄÙë{cåŠÒ8“qî>\\6nÜé>J¥«Ei‘ ÂËòî»ïO<ñD8zôh\†h‹Çô•W^½{yY³Í”)S {Ÿ<¶mÛVOi^hˆÌ=ôP¸üòË£¨ÚÔÁ£>Z\ÜkðŽ&¿|0æ÷Æo7écÇŽâvWuƒ5y0¸põÕWÇÆW¥˜1cF¸ãŽ;BkkkqÕih ìÛ·¯¸XDDDDD¤O”+J3Ãôšk® ·Þzk¸ë®»¢ƒÏÍ7ß—Ñßë®?S«V­ ÇïqfkR®(}Ɉ á3×Ô„/gÿ/Ù²?lm<=¦‹ÛõÆÊ¥éw2$¼üòËg "<ùä“á§?ýi‡eå€þðÒK/Å…¾¢(]ý(J‹Týû÷G¯ÚwÞy§ÃrÄħžz*lÞ¼9>üï½÷ÞðöÛo‡ &„§Ÿ~:üìg?‹!"ÍÍÍaÞ¼yáù矣FŠ!&x@aAfÍšuú%²iÓ¦0zôè(@"ø¸qãâHç–-[ÂÔ©Sã¶ü?yòä°nݺðóŸÿ<¾pð¤&¯/Äågžy&®c䛼:ƒô{ì±Xþë®».–¡+È—F ^ÛxTïÝ»7.§!„ðKYØŸ4»¢˜ß믿^ܤ4ˆÞÿýðâ‹/†gŸ}6z¦9r$®£> Á‹’}xið?¢6ûór&m^ô¼ˆ¹ˆ½¼hQä$J3ÅŠtó£Á„êàåƒ÷0/kí<4·§4P&¼Èq–ÚâÇ ·Aâ³pyz¥©g께û:hˆ00ÐSܳވÒÔ#·$mJü%/Ê€§/ó|ÌmÎ-ÇÏñB '’B´š„óšÄy„éôqLäÅË>Ÿ_Šã¶cÇŽxž aùìÙ³¤EDDDDä4ç"JÓÇIÞ³ÌðDFà¥?C_ïdfä":Óï#̳w…ôYˆ²é¦òàðC_†´=ÉÿÕyG­}\ºpØ"MìüqáüCçÊœïsá4Fÿ:ýî-åŠÒyAyöÚN”æ\pÌxµSŸ (@^”¦ïŠSUê+£Aàp…‡4 jsþÐ|àor¦¢^9gžxoP”®~¥E*°Äù-NBôåaFb•‰?œ¨¯¯?-J#$#X¦øÇ ^ x`^¾ˆÏóçÏB/#–ˆÊŒX“OŠåU¥I7_6^¼œÐYÇK;¢f%DiDY„Þ:$ÁÈ*/ÆäåÛ=‰Òˆ½Œ’ÅXòÇkΧ(ÍË™‘cŽ‘)cx:3ŒsLCŒƒ4P©¾»¥i0²œ^ìyQÏg[4R^üe"5(h´!Fs½!`“WÞc@DDDDD/ç"JÓÏ¢oJ”þýú§é;68f!\§oìÐWe¶Å¡C°¦ß‡PhMkÊ”)q{úk8GA8é“‘}"ú¶‰äì„G4¢4}Ÿ¢(Í>é{C©ô—é³Ó?.'Äa5ˆÒh Ô Bt õHß‘º§ïI?}!ïE=ÓgMýgfe_vÙeqç0AŸ”sÖSŸ¾3¥«Ei‘ ‚¨|ùå—ÇQÇ"¼¸yÈ'¯ç¢(MÈöçÅËÈ2£Â ÌLI¢1#Á<è­DéåÚQÑ9…ɇê`;¶ïÍx{#J3ŠžÂwð¢C¨-¦¸ËT ž=‰ÒŒÎÓÈIžé@C†:¦1çS”æåÎ6ùr3’Ox;÷$JÓˆb "5°¥óƒI”Æ#žôØ7ïÕÎÈ?õˇ)K>ô b4×NÞ3ADDDDD/ç"JwS:Rô»R bò¡o†¸Iÿ§þâT“>äN¿”>Î5ôsS؇¼(MŠ>SrÜÄÐ伃ƒ¢t^ÔLá*’ó>ˆßôK¨é÷õƒ|Õ J'M€sD_‘I”F ŸBF&'(œ®˜ÕúÏôM¹8?ù:W”–îP”© ˆƒŒ"2ýˆ8/Q^ؼ€Æ^'Nœxz¿˜8 fäxB#\òòåeš<&yI02™¼cÉï‡?üa ýà¥ÃHpúÐaš•@HNÞËäÃÔ'¼±‰eŒ`K™{ó^bL%/î®à%F#21¢`L~Ñ/;+Œ¸ö‡4u¨38^ÊG1éññ?9x£§|òcYoDiö#® ñÄH2uÉväÉËx^ê äÏ ¢4ç‚ ¾ñ& 1EZ4ÊÒy^4øpí>\S4ØŸ8Ëø¡›Á  \éÚ‘Á͹ŠÒùï¡OKÿ‰>.Ð÷A¤f&g‚þ,Ž88éÑg¢ï›wòÈ&´D‚Ð"ôÅè㜔…ðMº;½3ªI”Ùèé«£[šz@¶KÞì &p.p”BÀÆ’ðO½¦„¾¢(]ý(J‹TW¼Ry™"TòÐǃ9½ì)™ŽÄK5A¼%¦=¥,pDW„gŒ—ÞÔyx! ^'¡š—?¹‘ „_öExå…ÎègÞ›3ïÁ‹70e%FK)g wÑx-ãC¡ÑÃ1"(§ ¢,Â0 „íüÈkw ¦ç=Í;ƒmM§ñÊÎ7`Ž)wO¶ Ôû¦ég]G2õËàÇCÞ©®ùÊtúÀ"0]ŠmÓÇ/Eæw:4ð|N Js©`ðƒóÏö pÐ`Hiñ"çÜqÌ”…mð°ïM½ŠˆˆˆˆHõS®(MŸtÈ!Ý:¼0ãç ä) 8× |Ò‡a_úWüNe <ô I;9â"2b)ý%úÙôíèkápÅìdfÒ¤„øͶ8÷ЯM3d{Ò§oËzö£ïŒ#eèk˜Ãó!JÏ_¿'ìØ{¬ßEiDdføæEiú–xŽ_zé¥1¬ öS7ôcÑ18'ˆÿô¡Ñ¨÷´-³½q.Ûš~+ýf¶Å ¬¯×˜¢tõ£(-ÒOà­Ë»_º//‰þÁ“—{‚¼yçÅÑJ‘>èW ci„ôKåÇK¿«P'œ?Ö)F‹ˆˆˆˆHžrEi<–q2êÎaG œ‘øà{‚>³=q˜Áq§š|¸J ?ŠÀLlèÄÖ­[cXCòÄi¼Ÿ“ó3SóÔÃ)‹p–äSÞ½lOßg#œ›è㲞P˜Ý͆íŠs¥/¸v|xyÞ¶0me}Ùé”+Jã†ÓY±‰0Ï @~@gœà¨+ê'+ºpx‘-~YÇþxÒ#,3‹š„¼#^oP”®~¥E¤SxaöcWpªIaηò<̃SNwypìùóõ•s¥±‹GL8§4ʥ˺ꭣS>Þw9(JW?ŠÒ"Ò)Œz2͉eF£ ?ÑשL""""""òñ棥«ó!JŸ«õ§(ÝŸ(JW?ŠÒ""""""""ƒEéòQ”®ŠÒÕ¢´ˆˆˆˆˆˆˆÈ EQº|¥+‡¢tõ£(-"""""""2HQ”.EéÊ¡(]ý(J‹ˆˆˆˆˆˆˆ R¥ËGQºr(JW?ŠÒ""""""""ƒEéòQ”®ŠÒÕ¢´ˆˆˆˆˆˆˆÈ EQº|¥+‡¢tõ£(-"""""""2HinnŽ¢tkkkq•ô€¢tå8|øp¨¯¯/.–*BQZDDDDDDDdÒÞÞÅ¿WIl¨?ü±¥g¯ÛS,Ú€¦­­-zI755WI¡(-"""""""2ˆI!<ñcÇbè­{;qâxX±mO¸àÚšpáÐñ™ýÚ•cÃäe»ÂÉg—q ÒÒ{öì‰&R½(J‹ˆˆˆˆˆˆˆ rZZZ¾}ûÂîÝ»£@­uoõ»kÃû«·†|ü½ðÝ'>:ûÛGæ…É‹7††úÚ³Ê8 ¯}<¤¤«Eié7¥EDDDDDDDDD¤ßP”‘~CQZDDDDDDDDDDú Eié7¥EDDDDDDDDD¤ßP”éN4·†»†•Ûö†U۫ܲc\³c_Øøx±DDDDDDDDD¥E*ÍŠ­{Â?=0)|éê_„ †¼.¸¬Ê-;Æ/\ñZø³[Þ £go(V‡ˆˆˆˆˆˆˆˆ r¥E*Ȧºƒáox+|öŠ7Ã%#&„/^7qÐ؅׌ Ÿ¿ìÕðú\…i9ƒ¢´H¹ýµ…áÓ—½¾xݤ³DÛê·Iáó׌ yÛØÐtìd±jDDDDDDDDd¢(-RAþáÞ‰Q˜=[°!¼·aOؾçhø­›§Ä4ÓºK:Ùö¥Y[”uá7oš.ÉmÛÙö}5EiÉ£(-RAÊ¥÷ 7í ¿{ëÔð§wÍûȼð`ͺÐÒÖÖÕ6…¯ÝRš±Ï];>\4lBª/¸¶â¡yÎÚÆ°©þðiQú¢¸¾&\œ­c»$læšš0kMCر÷hô˜¾hxiyJ¯´}MÌ£XÖÞ˜¢´ˆˆˆˆˆˆˆˆäQ”© åŠÒ;ö ³Öî ¿~Õ¸(0#óÿã“6Äto½<|öšš(Z¿8{kX¾ý`Xºuxl↮ƒíó¢ôCkÂ÷ž\¦¯ª«v 3×4†z?Æ«~}ÁŽpðèÉðáÉÖ°x˾ðÝÇæÇ´‡<¿8Ì[¿'n?yy]\ÞÛÕyS”‘<ŠÒ"ä\DéÙëöDq8-Ghþîãï…öö^›·=\4|BX¸io8ÙÒEãùë÷Æ<_™»-z6'QúË×O ÿ”íwìDkضçHxwñ®èýýŸ.ŒÛ:ÖŽ7—Déo=0'\ö¡¯ìÝMaÌâÚ°çÐñpàèɸŽ|‹eîÎ¥EDDDDDDD$¢´H9Ÿ¢4ÓÿûÞYQ\ûAmøá¨E1+_\>9dLøÕ+ƆQÓ7…æÖ¶ð÷Ì 3V7DQú7n˜=ž¤‰ýÉËÆÄ¿»÷ sו¼±§¯¬[Ž„¯Ü8)†ùØX×½¯ùŸí¿qçŒÐt¬9¼¹pgøܵgÊÔS”‘<ŠÒ"ä|ŠÒxJÿÝ#óBKk{ôn~jÊƘâñâ-ûâÍûÂÆúÃqÙeÏ/ÓVÕÇß„øÀ3zÿ‘“1Nõ[KÛ=Ñ?¤øõÛ¦†i+ëÃÖÆ#á«7Mq÷ÌÐôashl:~fûì/ù®Üq0|)+ q«‹åîÊ¥EDDDDDDD$¢´H9Ÿ¢ô§¯nüÅŠ˜îÈ7V†çglŽ¡&ƇFœþöÃs£P½›GN Ã^Yóøþ3‹bˆ_ɶòBéã„߸sf˜y*|ÇWn˜–n;6Ô5E/ç_½bL Áñê¼máÅY[Â…Ã&Äm¥ õAþˆØÄ«þÌÕãbXö»¤6<2q}Vž¾}ìPQZDDDDDDDDò(J‹Ts¥;=›'.Ûæ®ßŽh‰ñ¤ñ–F¤þúmÓbhBm<3mSx~æ–ø±BbGí–)QœÞ¾çhøòõ“•/-‰åYºõ@x`ÜÚ0}UCü=jÆæSzÌû»B[{{÷AmøË{f…;Þ^×Ï]×î˶óþøûö·W‡ †*J‹ˆˆˆˆˆˆˆHù(J‹TrEi¼’‰ß¼¶¶)Úòí› v„xt~øü)Qø¢a·˜æÚý¥ð3×4dËæDïê—fo 5KkÃWoœ.>!ŠÙkv ‡ŽG±ú¹›£g4±ªÿå©áƒ-û³¼…ï=¹ ¦}ßØ5aCÝá虽¹áH§ ÛїИ¢´ˆˆˆˆˆˆˆˆäQ”© åˆÒØEÃÆG±ø´ Åh–wØnø„¸üwn™~ÿ¶iq[–¥uùíÙî Ù²?ȶ#v4¿Ó YNì(Îò´þï˜Åëø»“²ödŠÒ""""""""’GQZ¤‚”+J÷Õ’±âòÎ,‰ÎÅå"4‚uo·ïÌ¥EDDDDDDD$¢´Hé/QúãlŠÒ""""""""’GQZ¤‚(J+J‹ˆˆˆˆˆˆˆHG¥E*ˆ¢´¢´ˆˆˆˆˆˆˆˆtDQZ¤‚(J+J‹ˆˆˆˆˆˆˆHG¥E*ˆ¢´¢´ˆˆˆˆˆˆˆˆtDQZ¤‚|û.Eé‹¥EDDDDDDD$‡¢´Hù·Ç¦… ®z'|ñºIg‰µƒÁ.>>üæÐ7BíÞ#Ū‘AŠ¢´H={}øô_ —Œ˜JÂôà²O_öz¸ô©¡¥µ­X5""""""""2HQ”© 'š[à ?›>sé«ásW¾.¸úÝAbï„O¾yûØ°¥þP±ZDDDDDDDDd£(-Raš[ÚÂÛó7…+ž™¾÷Ðäð/W¹eÇøoM [+V‡ˆˆˆˆˆˆˆˆ r¥EDDDDDDDDD¤ßP”‘~CQZDDDDDDDDDDú Eié7¥EDDDDDDDDD¤ßP”‘~CQZDDDDDDDDDDú Eié7¥EDDDDDDDDD¤ßP”‘~CQZDDDDDDDDDDú Eié7þ?‹ÃKÎ×èIEND®B`‚apache-arrow-java-19.0.0/docs/source/developers/img/conbench_ui.png000066400000000000000000002113001515451615200252310ustar00rootroot00000000000000‰PNG  IHDRˆ$u"l€IDATx^ì½w]W~çGy·ÊþÃÚuùO—·Ê¹Ê*WÙëµk­­™‘VÚ•5I#i‚4I£ÍŒ4 ÁÀœrˆœ3 ‚$ˆH‚$‘ˆHd"ˆÄ4sÜŸþZ·_w¿×ýx÷ùU}ª_ß{ò9÷Þs¾÷œs¯úôÓO“ˆˆˆˆˆˆˆˆ4/W•ˆˆˆˆˆˆˆˆHs¡@$"""""""Òä(‰ˆˆˆˆˆˆˆ49 D""""""""MŽ‘ˆˆˆˆˆˆˆH“£@$"""""""Òä(‰ˆˆˆˆˆˆˆ49 D""""""""MŽ‘ˆˆˆˆˆˆˆtÈgŸ}–¾üòËôÅ_´9'W݈h4"—¢±DÃüüóÏÛœk=}"""""""ÁØ6þž}Ò«¯¾šÿ/»¹œž•+Wæô½ùæ› —>‘Žˆ‰&LHO>ùd‡üñtóÍ7§;î¸#½ôÒKiݺuiÔ¨Qy\^öß[@àbfÔùóçÛœkbuÒ¥±jˆ˜9ôÜsϥ믿>ÝÿýYÙ·o_Ú¿ZµjU8p`fÏžÝÆ/±j…£Jn)œ 6d1jñâÅ­…UvW‰p×YáVë7åxIÏ;３Ó÷Ö[oÕ”>‘Ë ãØ÷ß?ÝrË-馛njý ˆDÀ1taÆ¥3gÎt:~n4Ð6Ðú÷ïŸöìÙsI¶Ë©4„ǧ÷Þ{/>}ú’”oMöÁäÙB<ð@Ú½{wúío›¿:”|ðÁtà 7d·QÈdŽF†BÇ1Ž••0~GX€’WtC byFcÄ]{êZÄM¡=z4O‡+Ïì ±'Ü¡–ò·Rx#}ÇÏD¼q."fa'Nœh“‘F…ñ13ˆB ºûî»ÓèÑ£Ó!CòÊ!@(b¢ËÍÊþÆæï¾ûnºñÆÓöíÛÛ«‡NÁùJú@èºBQl*ž/ûí£ÿú e{çwf}¢½ôÕ“š"4cÆŒÖW7*j*ñâFE¥aD„OúÈnpËÄ1f^µWv"""""""—ƒ÷?ýôÓy2cXÆ¿±z7!z0¹#ö$4hP@ÚçgД¹\ÂRgéB«`ú(gš™8T5I&§OŸžÅ fü ,RlbMAõíÛ7g£rXšÆì6ܳ•ðO!•¢˜Î†;„(fÿPÁ[·nÍlj±ñ) U èã?Î_cCÈ¡ñǾAØÆsXTFø¸»ï¾û.J‚écfe3ˆ(+ħpGúcöUOW¬ˆˆˆˆˆˆH-0~g«•Áƒç±kD|¥»¸ŒŒñ,ãäØ“èšk®É+rÊã\ÂcÆ«…ØèzóæÍÁ˜»<ƒæRÑ™@„ÄäÊw¬~B‹@¬aeîY1Äÿä äõ×_ÏeÂD4VFáŸò+û§<‡Øü;"ʽ…™CüF;`ᣡôt9U-Ål‡¢ÕQâb©fS%„›XN†Î1“‘å‘ð)4f¡VR˜¸#* áˆÊ ˆÆÊl b%âç7ÞÈç^xá…,þ !Ü n‘–›qœóÄEEÓ QæXFž*päÈ‘9f6•…g(¬ˆBˆK`!=ûì³­îðhE>X××hŠˆˆˆˆˆÈ•cd¶_a<@Ä ƾŒc1–A1ÎG a¼Ìì¡ØŸ·NüÏd Âa<›]3ff;V U3ù¤Þt$…@úÐH#6þü<©…2`âIŒûC¡<øºÛ'Ÿ|’ý?ñÄùG‹ ì(/V•"¶º Í„:@ÏhovV=©Z "3d¡‡ÌÀJ•H¦˜nEFhTT6•OÆ‹¢î(°Ø“ˆÂgDT·¨þ²NÆ„ú†Å&Õe…2Ä$Ü2sˆÙ;±¬‹Bei%b4H>χð…@4eÊ”Va*ËÓŽ¨ *ž™A¨œåÆCÞ€ã±I5bU1¨¦¤ LHDDDDDDÆ¿Œ™”Áxžqõ´iÓò~ÁŒ×™å_6CÂÊâPÀqÂc™Zln?Â`¯žâØúRÒ‘@„ÖÁä4„åË—_4®Gë@O@¿@)Žû£ÌˆÐE˜`S\q¢Q C°¥NG+¡Ð/HK%ý¥ÞT- &£¨~lÊ\Iäà›3÷éÓ' 8¨jT>ÓÏ* DdžY:E¨,¬´'±P±ù—J4>D­ØÐ:xî¹çò1¦jÑ(ˆði, DTÂ~Šñâ–is䇆_þÌ}¤MHDDDDDDÆ¿|̉ñpÌøa,[é73dÊþËE"–¤]nqŠQìwû)s.f…A0¶ÈAbB åÃø>ü.+˜bZG¬„ ÿLRAôaeiØ“a‰89Æì$Ü`!5Ü " Áˆ>ˆd‚åTÅDòÅ‹¨‘vØ´‰Â£à(àp@›Cóõ±â³j"fÿ ¶…Ò†[ 5–!þ0㉸xˆ/ CC dÙ{u&‘>”¿˜u„…;þN™2%— Kïb³m"ém0þeü›+#d0›&f±ˆI ˆ+lÈŒÐÞ ¢€óOŒ“™Ìq9Å!ÁŠ¿hÌêVC‘'–—1Û‡UDL A› ˜%Dñµô؃q'–ÒQ&è|ÉóáÛ¶mÙ?úšç€(âCkˆ/¢c¤ ÿì×ÄÖå<Ô›š"*Ñ„9$š•€(ØP»È " ›-Qø™9sfÞŒa„cnqÃqfE櫈øŸÆúæ›oæ836©¦‚iÌ#ˆU(•„Ç1öB¢Âˆ ¥0ö êH "fCÑ8Pôð‡[ÒG(˜%å "éÍ0&~ÿý÷óx±ƒ=‰ç2 ƒ?1Ffl^ëò'ÆƵ¸ï)H“;Ë#°,@`ò šAø’³žÐ˜àBžCA×à8å‚6øŽ¯˜ñµ2üçbïcüs­#¾âŽö@Øè ±ÄŒI5ñµõ†úŠY‚ KÂPãÓlüæçhD1=‹ÆÃ1…RtËofÝ„°ÂÒ´òD¨l¸;wn.¤õë×çÿãóòlrM¡£ÄÅgîc é  ‰¨(âc#)ÒÂT8Ü Ö!*¡–"¾DÆ1fK‘=Dâ¦ÒÙ€ ¿4"¾l†½ýöÛ9e¡ Õt¯‘ˆˆˆˆˆˆ42Œ—søËÌ>gϤŽ1fn±§+~V¡7• ˜Ü˺˜¹Ãyô Æô±µ aÄl#4üÚÄ”)SZËÿ„YôeÆ_ÂàÛøàŸ¥wÄ3²?ÒÓÙ,­îR³@d” —b˜Z†RÆÚ:Ô®øzXyéÊ›<£–!ñÙ7¦MEáÄ6Öê¡B†ÿyظ‰%a+1øgæ‚ Â3‡¨Ä%üEÁáŸedˆ2I¸# Ì(b6yam%S昂þISÆسˆc€ÊˆÐ„èÅÂÌ#”½ĨtòÁ4³b>XÊF>.…ò'""""""RO#3Ž.ï­„ÈU‰qÊnBkà÷¬Y³ò„V¡ƒ,[¶,O Y¹råEÚB%ÿZA1ÎåÊç/…×%bYëäP²XjÆqŽUJx¸ç73†âÿbáı²Üp<Ä–âÿü¥f\¹À‹q³¼,*(Ü•Ã/Uéj ¿‹~«Í‡ˆˆˆˆˆˆHo!&M”7#Œ÷Y-Å$˜4Ãò0&Í QôÆrê²@i ª-€jÝUKµ´ZwÕRïðDDDDDDD¤w€Âj#–…1‹ˆ•Vµh#F·"‘f1(¶›)¯"êm(‰ˆˆˆˆˆˆˆ49 D""""""""MŽ‘ˆˆˆˆˆˆˆH“£@$"""""""Òä(‰ˆˆˆˆˆˆˆ49 D""""""""MŽ‘ˆˆˆˆˆˆˆH“£@$"""""""Òä(‰ˆˆˆˆˆˆˆ49 D""""""""MŽ‘ˆˆˆˆˆˆˆH“£@$"""""""Òä(‰ˆˆˆˆˆˆˆ49 D""""""""MŽ‘ˆˆˆˆˆˆˆH“sÕ‘#G’ˆˆˆˆˆˆˆˆ4/W}öÙgIDDDDDDDDš—«’¦iš¦iš¦iš¦išÖÔ¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖä¦@¤iš¦iš¦iš¦išÖäÖeè‹/¾HgÏžMŸ|òI:uꔈˆˆˆˆˆˆˆ\FÐhÐjÐljµ. DçÎKÇOüqþ+""""""""—ŸÐjÐnj±š""Pi\Ðnj‰jˆ¾üòË6ŠˆˆˆˆˆˆˆHc‚–SÕ$9{HDDDDDDD¤wPË,¢š¢Ó§O+‰ˆˆˆˆˆˆˆôÐpÐrª±š"vÃV i|ÐpÐrª1"‘+"‘&GHDDDDDDD¤ÉQ irˆDDDDDDDDš"‘&GHDDDDDDD¤ÉQ irˆDDDDDDDDš"‘&GHDDDDDDD¤ÉQ irˆDDDDDDDDš"‘&GHDDDDDDD¤ÉQ irˆDDDDDDDDš"‘&GHDDDDDDD¤ÉQ irˆDDDDDDDDš"‘&GHDDDDDDD¤ÉQ irˆDDDDDDDDš"‘&GHDDDDDDD¤ÉQ irˆDDDDDDDDš"‘&GHDDDDDDD¤ÉQ irˆDDDDDDDDš"‘&GHDDDDDDD¤ÉQ irˆDDDDDDDDš"‘&GHDDDDDDD¤ÉQ irˆDDDDDDDDš"‘&GHDDDDDDD¤ÉQ ê!Èÿ©S§òï'NäßãwÙí¥€øOŸ>Nž<Ùæ\£@ÚHc”›ˆˆˆˆˆˆˆ\ˆê BÇÚµkÓO<‘ÆŽ›Æ—ž|òÉôöÛo§cÇŽ]r‘†4½ùæ›é¥—^J6l¸äñWiÚ¼ysNãÊ•+Ûœ‘žC¨Ž Ä|ôÑGY ºå–[Òu×]—n¼ñÆtà 7¤k¯½6Ý|óÍiÔ¨QiÏž=—t– âËàÁƒSŸ>}ÒÓO?ë¦ìærCš,X~ýë_§¡C‡f!írͶ©'ôDZ."RD¨ÎLž<9 BˆA&LHË–-ËLš4)ÝvÛmùÜÈ‘#³T@b)Zùx%Â-›ÎÜãføðáéúë¯OÏ=÷\®ŽUW¸«fÖQ5a¶—nÒôòË/gQòA Â]µq‹ˆˆˆˆ4"ôy·oßžžyæ™´|ùòKú¢øŠæÄ©nàøB¤ Du‚ü²„Œ™CˆCO=õTëž:±¯³wˆjÞyçÖ·!ª|øá‡y™Õ¾}ûZã8zôhëµC‡e·;vìÈÇËn‹¢^x!Ç‹¿­[·f¡ª,Ö„ŽíÝ»7½ÿþûùoÙi!nŽQ×<ø¶lÙ’<Ø&=ø…ÇçxqKÜ‘÷¢@Ä,+ŽQÄÍŒ+þ/§QDDDD¤Ñ¡»jÕªÜg;gu—–1É‘ýéüúqé³wïoáÁôéªZào%8÷@«;üœß2/øøh…°Eš¢:A^GŒ‘ Aƒ¥#GŽ\4ó…ßû÷ïÏ{ 80½úê«­3dAFî»ï¾Ô¯_¿tï½÷æûŬÜ<ôÐCY<áÍÃ#<’g$Ýu×]iØ°aY,*‹2Ÿ ånS§NM'NL·ß~{ºãŽ;r˜W¤7˜ˆ¸î¹çžiš;wn« „»ñãǧ|0Ïâ7i0-ZÔî¿fÎœ™xà÷wÞ™{ì±´xñâì¦(QŽÌtºûî»[˃t†"‘ˆˆˆˆô6èëÒ_ýõ×Ûí³7 ôçOÌæ)þ_†óá¿'N¥“·¥/Ÿÿý”f´ QgÖH‹Ÿ/—ü48öQKz[ˆQ ª |ðA2ažþù}zö‡ŸM›6µ¦©ÿþ9þgÃí˜5„F~|süñœö›nº)Ýzë­iݺu9ÌŠL¤—ðH/îð¿páÂtöìÙ,qŒðH#BâWìç4oÞ<߸ˆˆˆˆH¯Cè_8yè`:µoo û.°ÿ«¿íÁùìfo:yà@bÑɃÛÓ—ó¿‘Ò¬–!êìiñóå²T ©€Qà&Ѓ ƒÀÁÃÚ2BÀ‰eYB'ˆ5;wîl]ZÇadøú¢ áO™2¥uy3‘BTaéñ²ÄkÅŠù!Dx!ᡇ¯„qœ%bÏ>ûl«ð_8CÀAA4bfar Aç7ÞÈq#òð?³ˆˆ›c,c–ÇYΆ¸E¡‘º"îÙ³g§¾}û¦G}4ç1"!D£(#öo¢,ؼº8ÓIDDDD¤7°mÛ¶Ÿ>5ý¦ß­é7wÜž~sgþ/S<û­éó±£Óñc§“‡ŠÑïÔ€‘HG(ÕPY[Ø‹¨=ˆç˜á³ŽEmÎœ9“Å΂b â ÿ³ü‹0€sˆ>ÄËqÂ@xá«i5»wïÎîbb ÿ‡p„ÐÄ,!Îñ1Ü3{ˆxÈÿ“NöbîØ|›0ˆøŸ¯¶EÚÉ›Lsœ?Ì bù‚é*ºc_#D!–™!½òÊ+­³–Øc 7”Åk¯½ÖzüÀÙ¹\EDDDDú®lA™¾o{ã„žâر[úÚìZ;ø=þqÛ0»/É[Ƈ_šÒßý$¥Ÿý´6þþoÓÞŸN9¦@$ÒC(Õ˜9à fÎ ´TZbì×ìիWg!‡M£‹KÂØo‡}†·,ûÂ-{…ÛˆA‘)"öúA B|* DóçÏ¿è¡Äo„œY³fµ¦ ±‹å`E¨b ü"13a'Â$í1û ˆÿÙs wK—.½(nÎ! ‘fÊ«ü3¦¸ë­·ˆDDDD¤×Û8ðâóRÎ :uêdšÿìÊ4ðágÓÐþ/ÔÀóið£Ï¥ñ£¦²EÛ°k&¢ÃSúéߥô?«øûôÅÃ*‰ô DuÁ‚0,CA@!ïE!ƒ³ps®¹æšü¾ø…Ãò0¦"ܲt ˆ½|ªˆZðËVö/*.1CdáóšEñŠ0Ø;afÎœ99„ f/1ã‡Í¢ñ,c%løíL zâ‰'rø,“ÃËÆ"nÒÅr1òM~ø¿=ˆ%{ D""""Ò[¡O{9ö :}úTòØóéÏÿä±ôÃï ¬‰ï³ú矎M{v·ô¿OÖ¡ÿý•@ô…‘Hâ@T'È/3dŠ{èp,@aæPœGÌa–Pˆ@ˆ,̦‰Y5ìóò1Ä#ö&b‰Y5Çb™ZQpB B¤aÓgÄÎá†ýŽHC8,ûbæq!L}þùçy¡óçÏgÁ ‘ˆ¿„Ù™@KÌb†#ÝÌâ‰@Æ #6²&"¹’ ïÊŠ¶]ˆ1BÙMO@4tÀ éûßþú{ƒkâGߘúü|\Ú³§¾‘3ˆD¢:³ˆ7M>Ø›‡M™×¬Y“?Ï1ÎñÉvü ”0kñ‘‡ågˆAìýÃ×ÀKø$=îÖ¯__µ@TN[ñ+f@œ„Á &D›Ø¼š°ð?yòäœ&¾”FúIê!V1û‰ßÔo5âKĈ—™IL»víÊᲄŽ}òµXRÆþF‘/‹O×ó3#f±wÓ¯ýëüµ²¢@ÄÙïá‡V ‘^}Ú˵Ĭû‘3ˆDš¢:ƒ@‚¸Á‚qãÆe¡…™6|=Œ™7Ì2ÂMÑ=_cö ³…X†ÀÄLÄ ¸A´AdaVËÂB`â7{ý̘1#oH] ;@PAÄÁ Ó[ù=jÔ¨œ.>iÏžBø á…8YþÆ~CäÁ†4±„Ž¼E/¾øbš>}zž!ǃ©³Ì b™aq x N™2%ç‘)¶ì“„ y$üQ±‡Ç™ÅqÒCY)‰ˆˆˆHo‚>,_-î•3ˆv7Î ¢/z@H¤Q ê0˜9Ã_öÖáÀƒ Žµç7! •Ýò›Y6ñ™øÎŽ—!<ÜGÔ éâ\¥uÐÅ<°™4qWŒ£fÑ/îÚ;ä‘øñ_´p_4+úÃM¥ã"""""ýçU«Vå´¼H-÷‘{¢a݈¾šAÄ—Ì t"‘ÆF¨áa”ϵG-n»C-i«ÆM-T¯ˆˆˆˆHo¤8;?þß°aÃeYbÖmhÏt²%/ÇIG÷úö~w±§ãW ipˆDDDDDDêÈÖ­[ó– ñb”±Û&ðñ•Þ&íÝs0 D›Ç/NoüjRzãŸ&_àŸ¿úÛ +~91­ºc^:¼}_:qò¸‘H£@$"""""RƒØôñÇÏ3†:t( FìGÊž±×fÙoOPOh]ÿÒ¢o J‹¾3¸&^ýæÀ´âÓ¡Í{²@ôÅÈa D" Š‘ˆˆˆˆˆH@øAâC/ÌbÍ;wæ/óAŽ³GioˆNœ8žÖ œŸwpZü½¡5HÄÌ£C[.DÎ i\ˆDDDDDDêÄÁƒóË6mÚÔº\á3÷|å—qRÙOOÑHÑa"‘†GHDDDDD¤NTšÄØèµ×^«z¢“'O¤Ó-þºÆ©œŽz D'OOëë$¹ÄL¤qQ é!jÖ¬Y“ÆŒ“?s_Í ¢S'O¦•ë7§s_L£žXP£¿bä¼iÜÓ¯¤{÷§³gOw[ ÚSgÈD"‹‘ˆˆˆˆˆH@´oß¾,±?Q53ˆÎ9†Ìz!ýÎ7~œþõýmMü«?üIú¿ýËôúš éÓógº-ÕeÑ/'åMªO*‰44 D""""""=›Uoܸ1-[¶¬èùtÕ×þ*]õS_ÿqúï¿ùóôZ]¢y¢EßÒFê "‘Þƒ‘ˆˆˆˆˆHrôèѼqõÒ¥Kkˆ¾þ×éª?üIm|ãoÒï~ëK úê+fÌDR i\ˆDDDDDDz–˜±9õ]wÝ•^|ñŪö R ª€‘H£@$"""""Òƒ;v,­Zµ*-_¾¼×Í º°Iu7"—˜‰ô ˆDDDDDDzÆFï¿ÿ~¯ˆÜ¤ZäÊGHDDDDD¤Ù¿zùå—Ó¢E‹z©@Ôݯ˜MT é(õ¬5&ÿL'å7_.¨f½q=ááq’~¿Ën/År)ŸŽs¾½4’§âCõäÉ“óÄ1ʼ½0Ç;:élæ6,""""݇~å|æÌ™“^}õÕªÆ DP éqˆzD>c9xðà4iÒ¤´nݺ´`Á‚<¥”/”Ý÷'Lce­3é9tèPZ¼xqÞoÏž=m•Kéà¡Ø¿ÿôúë¯_ôp Afúôéiøðáiûöí‰7œ‡÷Þ{/=ñÄiìرiôèÑiöìÙéÍ7ßÌÂyÆ-áPÖóçÏOk×®­(­X±"ŸçwZýÓùÊ+¯ät>óÌ3ý‹ˆˆˆˆTK,1£\MßR¨ D"=ŽQ@Þg̘‘®¹æš4mÚ´üÕ‚>}ú¤Ç{,>|ø’ˆ3¤qc̘1éìÙ³i×®]éî»ïN}ûöÍ"K5¦z‚³sçÎœ†;ï¼3ÿÞ»woÐ{Hß™3g²0C§L™ÒZNñ÷©§žJ·ÞzkºîºëÒ7Þ˜á÷M7Ý”óyàÀbÑ#<’Ë|êÔ©y¦P1-„7hР|þñÇo#T!¦=úè£9ìeË–Uõ–GDDDD¤=è{oÚ´IHH¤¡Q ª3<ð@ºù曳à —믿>‹e(fÆÃàÿXJU>WtSŒ³ì–òöÙg³ø‚@´{÷îœ&Äf4q?øm/Žˆ§RøÅóE!§½tsìÉ'ŸL×^{mžõCš¶nÝšî¸ãŽ,ôlÙ²% 9L?üp‚xËá!²áî–[nI&LÈ3€`òäÉé¶ÛnK7ÜpCš5kVv‹@4pàÀ\æu•¢aÆåó'N¼Hâ÷Ê•+s9=ôÐCéÃ?l3ÃHDDDD¤Zè{²ÄŒYð̦¿ÔÑкDÝûŠÙ "‘^Q!ßÌ:A|@¤  xSPˆb¦Ë‘#Gò±Ø«ˆ" ‚  çÊ"ËÔðÇoÂ߶m[~#Q3øËò*DÒTˆXþF|0,å"¬öDâÚ±cGkøewøâ#­„IÚù?Îáá'D3ÒD˜wÝuW}øM^ð3wîÜ<{‡¿¸Ã?BN”!ñ1Û(öâ³µî¿ÿþqvU ".f5ã/úéŒr¿šþ=c¨I úš‘‘È¥C¨5jTžþùtþüù<û¥(!X ¢ bÜsÏ=y9ÕG}”Ë‹½Š˜¹Ò¯_¿tûí·gQ‡¥b!,Á¸qãÒ}÷Ý—÷â73hp[¾Ž$>!°„@„ÃÌ"âf–3xn–,YrуŒß¤›eq¤ƒ8ð¿páÂÖ°y2`À€¼W{þà6ÂäXì³r˜åÃqò‚øÂR3ÂÊs<«W¯Î³…‡)þÇŸý“VD­Øô;„&6ý{ú駳Ê1"f}úé§Ù}€¸Ä>Ge(òDZÈdž ªz€‹ˆˆˆˆôEéÃÇBô)y‘ÊËÛjú— DP éqˆêHÜøTI˜ÙÃRª¢@ÄÌ„–[92íÛ·/—{ìàÁƒ}u~8{òÀ!C²€Â9Ä –\±ç³ƒ€Ê3—"fñÜ{ï½Y\B„!,üÆW€OpFx#FŒÈB“.Ü>÷ÜsYd!1ˆs„EZHî|ðÁüIODÄD3Îó?âîÉ~šx˜²é#\Ê“¸Hǘ)Ä]D'–ó‘?D ÜQ>±ÄŒ4vÄXE(/Â+ DüE`#”Ìê*׳ˆˆˆˆH{ÐGå¥j¼ldLD¿’>úÒ¥K«Ú¾hð¬çB Ú«@$Ò4(ÕòB›#.p,¢¡C‡æå]!NàAaƒe\làŒXÂlÂãáñÖ[oeÁ a„™5„G8!À ÒpŒMŸùŸp_BôˆpŠKÌØ€™ø8Nüˆ$1;‡ž‹{(!Å, „fÖ 1Ó&fÛ&{±¹3îH3îð¿~ýú,!„Q”ÿS6ˆ?ä9–˜…Ʋ6#ÒDÙážVløû QÆ!t6ÓuñK÷ˆ8q¸/Bœ¤»(áñtòP//Mé ^vòáú§ôóé‡Ó·¤ïÏ6l—ÐÙKÈz DÝݤ:"7©¹òQ ªÜäH˜ƒ€{ÿ„@„ ƒè"îcO~#¦ ” TPvÀ9fïÅ"þg â'–¶qœåVí D¤ƒ¥iáð6oÞœ—{!˜ ÚðéwÂAÈ"=ˆIÀžH‘?r<üB B¼Š=˜È'Ç ‹´ <¸Õ_” âVD„i&½ä÷/½ôR>oZ(£yóæå‡k¤wäkæÌ™­å†@Ä9fQ~”MâÅODÀÞH[䉼UóvGDDDD$ ïÎ8€¶¼\¥oûŽÒ_¦^Mó‚@ÔKÌê1ƒHH¤w @T'ÈïÛo¿EfâĆÑ!!V >IJ-Ä ¦™ržǻﾛE f°°B ð;Ä 6O¨(ž?a0‹§3(6©Ž%hEa‹8Þxã¼Ì w±RÒÎ Īˆ˜¥ÃžJ!v±¬Œ2 òE; „ŸÇ-¢PDá?Ò{±ì÷ˆf€Ê€ßÌšâ!‹ØKç(SÒ{MŸ>½õÁžÙŠ_W#o,ç«fm¸ˆˆˆˆHú¡ôk飲DÌ¢oI¼æÏÜ7€@T=ˆˆDz DuQ±¡üõ«²@4uêÔ<ûq†Ù<3øE”aÉû¢R#®Ä¹º+­Y³¦õ<-–•±ìŒ8X3ÍðÙ,›·Á /¼7Ñ~ñÅóLf…@Ä,¤ö"âgÿ%ÂdÉZ5Q|eŒø˜iDxÌ"}ø'Ü;w.¯ïfÏ'ÒO™p¼ø³Ø›((†Oqž4Ǭ'Ò\,?‘j‰—Äåed¼ÈŒYEe?e®/""""Ò¡_ËG]ØÚ@HH¤QQ ’K3t˜ÄÞE,é*Ÿ§½°I5x—ωˆˆˆˆôF˜5¿råÊüÅà… VõT¨ D"=Ž‘\R‰b3ëò¹8ßÞ9‘Þ}[^‚.^¼Ø%f D" ‘ˆˆˆˆˆHÂØÈÏÜ+‰4: D""""""=_Í}ûí·›wÑ/ˆDz D""""""=KÌV¯^F÷ jJÈD"½"‘:QiÏMþ_·n]í3ˆ¾Ö(щî DÎ é(‰ˆˆˆˆˆÔ‰]»v¥wÞy'íÝ»·õ,§OŸN7n¬ý3÷ $­8?-úî6Pg´ D›÷¤“'ˆD"‘:€{ -Y²$Ï:xð`Ú²eKZ´hQzþùç󩳯ö^‘‘3ˆD"‘:€ð³aÆ´`Á‚, 1[èƒ>HÏ<óL3fL5jT:|øð%ˆ†ÕE êæD_ D'ˆD"‘:ÀŒ¡yó楱cÇf¡(–˜={6 Fo½õV'•ý•©§@Ô03ˆ\b&Òð(‰ˆˆˆˆˆÔ >iÿá‡æ¿qŒ±ÑæÍ›Ó믿^ÛD ó3"‘f@HDDDDD¤NĬ¡XFÆß÷Þ{/?>ïCÔ[¢ K̺(ù™{‘^‘ˆˆˆˆˆHÂŒ¢¥K—Öþ³†ˆº¹‘‘H¯@¨Žp³ç–¬=æ/eP^c\<nã\|³øÐ÷ÅOdr>ÂŽså´ˆˆˆˆˆHc@¿ýý÷ßïÅK̈Dš¢:H³iÓ¦´bÅŠ´cÇŽ´lÙ²üK>sŸ²äa°wïÞüÙË^x!-^¼8íÞ½;‹@œß·o_~³Àúä8¶uëÖ|ŒsÈ~ÙànõêÕ9ŒW_}5íß¿ÿ"¡IDDDDD‡£G¦wÞy'÷å›Y :áD" QàFϧ+o¾ùætÿý÷§»îº+ÝrË-馛nÊ_+8sæLZ»vmzðÁ³›;ï¼3Ýxã顇JÛ¶mËç‚®¾úê¼69fM:5]wÝuYBlâ÷<î¾ûîtûí·çÿ§L™¢@$""""Ò€ÐOgæÐ=÷Ü“,XÐf…A%®d苑ÈD¢:ÀLŸC‡¥Ç{, BÏ<óLþjopž~úé<ûaáˆÙ?GŽIóçÏO}ûöM³gÏΟ¾œ1cF6nÜØZ~ƒ ÊþðÿÔSO¥ë¯¿> 6,íܹ3Ã9â%~—š‰ˆˆˆˆ4Ì b ÐìK̈D¢:ÀMžea·Ýv[1bDj˜ÄÌ!"f±œìÚk¯M .Ìç(ŸíÛ·§~ýú¥±cÇæ0œî»ï¾¼ ó@Ì5jT>OØÄÁúe%––! :4;vLHDDDD¤¡oß»?s¯@$=ĉ“éøÉÓ5ó1O0Ï:¬' Du€¼†ôâ‹/¶n$=oÞ¼<ãgÆ yvKËØS(6™fÙØ­·Þš—ˆ±ÇÓN"LÞ2Üpà éÉ'ŸL‡Îb"o 8¿~ýú?8¦M›–Ã\¹re~ã€ÀD8ñ4ö, ÿü_N—ˆˆˆˆˆ\^|ðÁ¹ŸÏfªé·+U@èŠãã“gÓ¹u#ÒoŸüŸÒoŸúŸ ü»Òÿåsÿ.¥'[~?û¿¥Ó;¥ã':]¥:ˆêbÎÈ‘#ÓwÜ‘oþäAlV=dÈüP@Äa‰Ùã?žg=÷ÜsYðéß¿Þ¯hݺuY b"ûqž%e|¯£á?*Ä{ñõ´jÞDˆˆˆˆˆÈ¥‡þ;ÛD¼öÚkUõÛ¯dȯ˜I€@t~퀔f^¨—šÀÏÜ“Nïxé«¥fm×ÚQ ª±?{±ù4yG´á+c3gÎÌ3~X †àŒ=:‹?¸?xð`™BböÑ£>š ——±ÌŒåhÌ4Š‡ q>üðÃyÏ"¿b&""""Ò˜Ðg|мÑÄ,T º²` ÚUNœøJ øU=ÖX—ø™÷?¤3(Õ¢:ÁWÆyb£hö bÂÿÇq6žf¶ÐîÝ»óÆÒñ€à<ø[>ÇÿÅsEÿ•ÂªtLDDDDD ^ê²å{ˆ*)uÆÖ»wH;¶ïKÔ~>üðp›0»LK>ýÞ{é·7ߘ~{ã é·7•àXP>7\—ξùFúøÌg D †‘ˆˆˆˆˆHˆ¸Å±P——˜}­ûÑ0¢!ý¼W DÌ:xà£t×-3ÓO8<ýìÇ#kâ§?‘fO[ÞÒöê³-ÉÇ-cýÓ«V]¨—Ÿý´vZêÿܲ¥ D ˆ‘ˆˆˆˆˆH@b»ˆ¹sç¶n?Á˜ˆm"ø ³ˆªˆ+õ @ô³^'øðpºéêÉé-uùWßT?øÖÀ4um¯ÎÑ/þ±m=UÃÏ~šÎ-_Vèߦ3nR]WˆDDDDDDê‚Ю]»Òĉó‡dƒØ{ô¥—^JcÇŽMÆ Ëµél›ˆÆˆN|% i#uFC~Åliïˆn¾fJúá·¶©«ÎøaK]NëD êl‹“âïƈ¤O|Ò&½Ò5ˆDDDDDDêT>RóÔSOe(f! ­^½:-^¼8ï7ZöW¦‘–˜íqÑe§ÑÔ ‹;ˆj! D«W¥ôËƈœAT?ˆDDDDDDêËÉ6oÞ|ÑL!fmܸ1­X±¢ê%fõÚ¤º>QÌ ºR¢Þ·Q÷¢AiúÄŽgÕ¿D¿Héçêª3ê*¹Ä¬ž(‰ˆˆˆˆˆÔ‘J³„Ö®][û&Õ $­8?-ún÷–˜u[ z°^QïúŠY«@Ô·kѾ3(͘T?èxˆV§ô+¢h[WQ7ÈMªë‘ˆˆˆˆˆHÁ,¢mÛ¶¥iÓ¦¥… æqRÙM™î D?N¿ûÍ_¤×ˆ*Ðû¢tU š¼´±¢e Dˆ‘ˆˆˆˆˆHrôèÑ´råÊK:ƒ¨á¢-{ò^F DµÓhÑE{uu‰™QC¢@$"""""Òƒ06b_¢7Þxã’ D+V×G Š¯˜uK ú•Qw¨‹@4é+ˆ¯}}r²…S]£%-uÙ¤Z¨!Q éAøÔýË/¿œ–/_~ ¢ KÌê+u÷+fuˆê¶Iuoˆ>J7÷Ú=èô©td÷Á´m抴eÊò´ejmlž´,í_¹5oiŸ DW& D""""""={mÙ²%Íž=;-k7¥@ô«Iéð–=-e¡@Ôê"M^š>9óI:°vGZúýáéÕoL¯~kPM¼ò_ú§Í“–¦çΤӫV¥ô ¢+ "‘„±"Ñ¥YbÖ˜‘3ˆºN]¢I¢ƒë>HËþzTZügµ/|õ›ƒÒæ©Ëˆ®`ˆDDDDDDzK+õÄD DÝ™SÌ$ë ø¿""?sß(‰ˆˆˆˆˆô'OžLï¾ûn6lXZ´hQ'•Ý”©‡@T߯˜)u•¡e\ÜUŽÿ¸>›TµÄìàºõˆº³Q«@4@¨ Í)‰ˆˆˆˆˆôÌ9pà@þÄýŠ+.Ù ¢ˆ†*¸´_Ûµ÷Ãtï¸9韛~Ý\MüêѱiÚü%yR}¢Ñ݈ÞeÑÏÛÖS5\$õïB=þŽQ¡@$"""""҃ij7ß|ó’ D¯ÕI ªË³&ýŠÙ'§N¦u[v¤ÿõ¯®KWýþÒU_ûëÚhñó«GƦã-åv0/1k (Ï R ºÒP éAŽ9’—™Õ¼"AYê ¢¸ôÑú-;ÒïýäÆtÕ×ܶ®:£ÅÏ??6î‚@t ›Q]÷ :›N·´gg]y(‰ˆˆˆˆˆô,1[»vm7n\Z²dI¯ˆê²Äì—Ó¡Í{Ò‰“݈lB¨ÿW3ˆ®™’~Øè‡ß”¦×U ªç³nnR½ã%¢:¢@$"""""R'„ؘ:¾@lØ°!/1ã\ÙO™z Dõؤš=pÖµ„Óu¨N3ˆˆº,M›´ä ˆþm:³cQQ ©ˆB{÷îÍËÉDX¹ð©òÓ§O§Í›7§·Þz+ÿF$êˆOÏMCçÌÿJTøÛÊü"B…ãßøIúÝoÿ"­xoSúâósiøÀÓ¾90ýøχ´á¯¿ƒÛ‡¿úΠtõ?ŽOû÷N§NLZÂYúgCÒÒ?Ö†%ßš)Ïç¾3$½õO“Ó‘mû[ƈ'ÓoGŽÈbOJü6„ çÒC¦SÇŽ§O>ú ¥߸ *Ì©@ˆAs+kaùÏÓ©ãG[ÊúT›²¯'gN’6nÛ™~ïooj©Ë¿i[W™Ÿ´_—-~úôŸN¶”ÿáƒGÓ-}§¦}{P›ººP—Ä òqøÑw§é“—¥3çΤÃv§×~<&-­T_ß–… %ß»ð»|~ñ7§­Ó^O§>=—ή^sá+fåz*Öe…㙟ÿCúôõ×Ò‰s_¤ÏÞÐ~=Ρÿ›öëñ‰“Îí\ÐR>gÚ”}{”¯Y¹"‘:€´`Á‚´|ùòÖý†ŠÞÿý´xñâôì³ÏæßÌ&êˆm[¶¤;‡NHWýûo¥«þïï¶åÿüÓtÕïýaºêÿù³¶çþÃwÒûõ¤'^\˜vlßšîºmBú“¯Ý–¾õGw\Ä·ÿèÎô_¿~Cúƒÿ÷ŸóïòùoþA¿ô“ï?œÞzkUÚ´qCz¹ßÄ4óïH³þ讋˜óÇw§aÿ©oêÿÿ)Íýã{ÛœŸù‡w¦§~ôXZ½ä­´qÓÆ´ëÞ{ÒÞ?ÿ³´÷yû~ðýôÚŸþé½ï~'íÿá÷ÛœßûßK»¯¿.m|o}Ú´zyÚ=é?¤½Ã®J{‡_̾W¥Åwü«´úÁßIûG¶=Ÿ]³ÿ"m\¿¦¥¬7¶)ûz²ùýMiáÒ×ÒÿòŸ§«þ¯o·­+êïÿø¯éªÿý?W®Ëÿíôã[N7nL«ÞY“~ñwÒŸ~£_›ºúöß™þà÷ÿ)ýÉ7n¬X—ÚR—ƒ•6oÝœV½üzšû­ûÓ¬ÿ|g›ºšÝR—üÇ_¤ ߸%Íþ£»ÛœŸñõÛÓ’óÒÆíÛÒ–žO{+ÕS {¾ÿ—iÁùã´ëûÑæ\æ/ÿiÏŒ…;v¤9sæ¤éÓ§§iÓ¦¥Ù³g§Y³funÆMœœ:"õ/1pø¨tǽ¦ŸÿêŸ[~ns?†JS§MÏáŒ796.~1£FNH?40õëwwþ]>£GMH3fÌÌišþøä4¥åX™i£&¦»oº=ÝrõõiƘÊn¦¶„3súŒ4»%œ9ãǧ9#G¦9£.fîèQéîë®K#î¿/=1ft›ó™ÇO³f¶¤gæô4gòð4gÒÀ6Ìmá‘{nL#Ü•æMÔæ|fÊÈ–´¯Îë¢X'ÕÔ]ÙÏ´–zפ#*×åèÇ'\oæ¬4ftåz¢þúö½1ÝwïÃíÖå„ S.¤úÌ4eäø6u”1>Ýus¿4òÑ!iêˆ mÏ·0}ÂÔ eÑÒ–ÛÔõ8jTš5bxºþ‘f šÿ/»ÙS¦¤Y³ç¤ÙÓǵ­êqòÀ4mÌ#éú_ÿ}š3q@›ó”fϘÔNçõ2³¥Ý¼öÚkm®Y¹"‘:€ ôÎ;ïä)3…Z—9“9{ölëïÎÀíùóçÒùsóù知íÛ¶¥gž~:}ùÅçmÎg áœ;×–O?=ŸvïÞÕ’Þ•ùwù|К¦–pΞkËg-iXþúké¹çŸO_þö7mÎÅpδ¤¯Ì¹O?M//\˜6¶”ÛùÏ>ks>S,¿³ü¾ gÏ}šÞywUÚ¾cg:wþ³6ç/p®MywKã–WC®“ uIý­Y³:½´`AúÍ—_´9Ÿs…<—ë&ø¬¥¼^ziAÚºuK»uYl{åº ÎµÄG]îÿp:ÿ)eÙÖÍEuP®ŸΞ?ŸNµ”Õô3ÒÇ-mŸÿËn.®KÂ,×Ï…z<úñ‰4múŒtúLåºÎà¿B™W‚ú+_³r1 D"" By3K‘F¢Ú{TµîºKìíR>.Ò:t¨ÇÚ'‚ÓîÝ»«Þðº=b¿$„¬î¤•etëÖ­K+W®ìÖœ4¬Zµ*/ên¾¶lÙ’öìÙÓ­|•9vìX®×òñî@Ù‘VDÅj¾n×”û^íܹ³[yÆïúõëÓ‡~ØåpðGY±ÌòèÑ£mÎW á|ôÑG9u…K‡Q¡£,]½¨¤ñà¦õÚ›w-p Ù–.¦xu§Ò(D~h_Ü/y˜F[»Tí¬»\iur9(Þ_:‚6Ñ[綠˜òù+b]•ÏU¢7];¤•{}îSåó@½F>ÂM<ÇÚk·ÝiûÄé)ŸëŒr­3½~šÞÒ?êé´qýÔã~J:»+Æð—´Ô#=„Ñôéwp¿a?©zqë‘gòZ<×3œzÜ;ëŽTQ ñ¢(³®‘ è–,YÒ-åU‡xK³hѢ̚5kº}¯Þ íéÀMß–È?å°lÙ²\.”OwÞJt‡îvãa·ÿþüV„iècÇŽM£GοÙÜ’ôºÛi«ÅüÓCúx;ÇýŽ:Y½zu¿Ò1”çÖ­[sùQŽ•àðú믧÷Þ{/?WêÑq»”G¾Ü÷P~7B»®'Ü‹âù_ÍSêAFñ~Öžðr¹!­GŽIS¦LILO?ýôEç£.yãüÜsÏ¥I“&¥‘#G¦‰'¦gžy&—÷îr¼5ï¬íS®´}î—„ƒ—·ß~; <8 ><Ï~¨ö^ŒÂ*ÇÕ¤úi¦>m£BÝqm±ñ3톶qøðá^u?¬ìiĽ”Ù?åsÕBÙñü!œ®ö3cüõ „SvS-¤‡k™pº3ûŠç$yÚ¶m[¾wu5œJÔS "]ÌŽ!­Ô¿»’VÒÂ8…rcæUwž%Ì>Ú´iSn_ÝÉ#÷Xê’ôt5òʇüu5©¢:AÇ„‹»_¿~éÆoL7ÜpCÞœ®Ú‹4. Šù$éµ×^›?~|>VvWOx@<öØcéúë¯Ï퉛~³·%òÏC‚kŒryôÑG»Õè*ÄÇC/Ä»Ø_ ì®=ðK—.M=ôPºé¦›ró7~sÿ Ÿ Èx8^κ'­t"¿ÅÁ=÷yN¤—k‘+ü”ÑÊP†Ï?ÿ|.¿b(sóÍ7§Ûn»-·™W^y%?_{K9“GD…¸‡"°@Ù]o…z`ÀpçwækáŽ;îhý´uÙmâuÊýìá‡îòà §¡þ¸þãÞĽ+Äþ2 ˜0aBkÿ'Ú1àž¿<ÏÃðy¤ Ъiû„ýàƒ¦—^z)·}î‰ î½÷ÞtÝuץǼê>'~³‹éìòÀu‡0шõÓLpÝ0ð¼çž{òus×]wåvPË3¸·}–smß¾=-\¸0ÿ­µ pO8/¾øbw §+˱‡´p_ ú Ýÿ*{®1^Ž±LŠk½+âîYæ†ˆÈ ]DÄzŠDõˆHùãåb Ëêº*DãçÕW_Íy&¿Û] ‡ú#-Œ{¸Ov¥ÜðƒpÉ8˜vA8µ¶OÂØ·o_~é@û¢ßË‚ZÑ®¡@T'B âauë­·¦[n¹%ª.çÀNêb7í¾}û¦k®¹&7î’D¼©¥cL{¢m5{["ÿ ¨¸Æ(—\ˆtÄ †ÉSO=Uu{ˆ´â‡~©ë¤7CΟ??„ âùQ„öÀùøŸA+¥ÞÒQ"´›ðÓá»"®aꇿ bVõÍu…шéám<uÇs)Ò¢}œ#/´U„2îeÜ«‹í8ÚmÜËhÑ.ªmûÜ™aIùq¢-.çØ¿£!ˆE~éLF~¹6èÄQ¸¡3Vö_¦³øêEwË¥ºWQ â¹A{ ƒF'9 Í<ñÄyÐÏ…öÚ9MåsÑ•¼tæ§#¨#íÑ™ŸÎηç¶Z?eÿ\µ DœCxáSØ\;t¤;ÛK§Ö4Öê¾\ë b6 XjÊ=‹çÇã~E;åyIÛä 2m‰r 7¼æ¾Rˆ8ÿÈ#TlûO>ùd¢í#>a0ð»ûî»sÌ"*§¿Eˆ0´æÍ›—Û%é)CýPÕO{å\¬ƒöÜ´Ggþ::W‰ÎÂëiº?~¨»fˆ(Ú ¸yqÃ̾®.ébŒ3u¸ÖººhíÚµ¹ÃÌ®ê£ì¦3È÷LúÜCø|:âUÙ]5ĬœXÒLyu%_íQOˆ—±4—û"aw%\],‡åX-á®+Ê-îÁ2„Q« CXä‰zòF8µ¤¸¦iç¼àùˆxIص†#]C¨NtW ¢ÑG™ñvŽ‹FùáW~Ð↿ÜñÇ ?eAÙ=qâ6ŽíÅY¯£ó‘7~óF‰¯½2Š|á.üP>QíÅåüÅ4Õöâê âãÓP|CYLC1ürýUJc±\ŠõeÒ™@Äÿ”G¥:¬W9¤­štÑ>»~#>üóÀiïMo„‰›H¿)û¸ÆÚˆÊqË´œ·r#Þh£üÅ]1~Y b S¿ á3`¢coƇ ’;±ÑfÔ 0`Á ƒÜÓy(·½rð7nêåø«ÉoÄ_Ž£,ó‹û¸6ŠBY9¾·³6WôSÎCù|%7f5÷Ñrxåsåóe7”M9®ho•âªa„@D}3 ¿”%ç€ß”/ƒåh3Ó¬ÒW^"ÿÕÜûÊy‹öT.·r¾ËqEþ‹×v1]ü_ˆÎŸ?_U\å4â¦øLˆtÄßâõÐQ]pŒt‘ÞbYEú;JG1.ü‘ÿŽ"þŸYÃ/¿¹–â³ÃÅÅßjÛX¸/ß«ƒbØAy 6,_ûÜ{¢ìâ¾À=Š%` N¢î#Ïä‹û6÷½Å¼ÉQ ¢í³—P¥¶îܹüv½ØöÒ#ŸÌê%ÜÛo¿½ª~çC "L„Wê‹xÊŸCâÞV®›Èk”uô{9V©}ÑÎË×G9Üð[|ŽEß'ÒÏßJõ[Îk¤±|ÝTjçå4”Ãéè|{Ç‹ñóSmÿÿË×ËN:ˆÊmüJ‹0hfà€‚8Ô^½weÃÌš@é_”ÝT Ï%úÌfâú«Ô:÷ÔâbB×we÷í[Ú53ru¹_ÕF5ÔC Â/×#åøÁL©²›jˆ<õˆ ÖÕe‚”Õœ9sr=Ò'á:«µmᡉ0H÷~– ¶÷‚³=pKà¥È¹:‚ð¸QáŸ2à¦Ç1òMÜ'ï¸!|ê´õGyб)ÖC<y²§uN:É+7ynÌø­$Ey¸óÆ”vBùP¶3gÎÌåYÌ#îyëuJóÆhРA­åŠPÁ =„ÅbþñϱJñÑf褔˔·âÄÇù؈yèС9´ÑQ£FµÎ6)ÞðÉoÕ©w69¸80@ ^y[\I "´5:Ä…? Ÿ¤‡z± ICê”òàÖì•A:©Î'a혺¦íré˜1cF¾ÞŠé)CQÆ1 ŠAN¹ì"/tü(+òʲ ÂA$ç£CÅ5å„{êˆtF½E~iS¼•#O´‡È/{9‘_ê“:ãÞ‚?ü>ù%‘_òN;"=”7³º¸—à—p£¹F‰ ÿ\„K['­´9ê…¹ƒK‘°“°Ë× ×LÜ‹xûY¬SÊ’Î3é+Þk¸–h•6Ê-ÆÇuÖQ|1óãÀ5?Ê0â"_\ÃÕîEzŠƒçJ3â>D»ážÀ`˜k½XǸ‰{5÷<ÒÔÞ½÷\Sœ£ý’^Ê“kÔÂ@¼­”ö ~ã™E\#FŒÈí<6ÓŽ<"ê‚늶Úѽ’6Î1ÒHÛ¥ 6m0ÒÈs…tðœbZ=û`çqÇó³ÜÆHWtb¹ÏDYqý~ÚqÑîyCN¹F;‰ôã{ÏJ~—¢€ûe”7îI÷HŽÑñ€vEž‰¥$´iò×9m¬R¾¢¬ãùG~âÚæYK?‚¿QžݯŠuL[‹¶IÄý‰~ùEœ!Oíµwî[!ª÷pòŽ0Vˆho•Ú>ñÑ>‹m?úWÔ;m4î“ÜSÊ×xü•¢J}¥2ø‹¥=Ô XÊ”6AÝPÞ\£¤‰öOs=QNѾx–Åó&âã/×"aRŽ´/î½”þhÏ”=÷Â&ÜFödzŒ´”ó@šI å\ìK‘úr”7~pí¸/áʌ•hÃÅMÁùKYD°è—º ~ê…k<®MÚ$yg`Ô9¿y~E:( â¢_…?®iŽÅ³±(É~IÏ‚5ËõÙÛàþÀóè?QÍTìê‰ûå/ˆÊn:‚2¥þè#Èðü¨5 À××ý ®gòZvׄÃ=´Ðæ>ØÙ½ +tW "1Öâ™ÅµÌ8®+e¤…>Ï­îìÄÌ#ž‰<Ÿ»²§U„ÃõJ»¢x~t5=” ÏGÚéiïù"=ƒQ ávE â"@¡3oèéä°æ=6wä/ðب˜²eÀ;øáAO—=BbScâÇSÎÅŽ.7:Ñ)Ã]ÄI<1]›sÄI'Žøb:ªÅ™¤‡›apžNMt"¹‘Òéé„[ÌÿÓ±¢Óe…?::tÔ#Eté\Ä1:>ÑÁ:et$c0eçíB-7>òZi“jÒËš<“F&ÔÔù‹òe ê7P>Üø¢\ŠõtäâkQ Â/á0˜à8y ‘Gü•Δ%nhtÒŠñEs ±¡ØQŽøÈg{ñÑQ‹s”fÜðV™r!lʉcÄñ1ø+Ö;q#VG¸Á~ùM{®4ƒ(üRΤ)òW¼†h—Ñ×´Â'½ÄiÄ} |pÃj9t:úôé“ÇÞüfÿ–)töÆ…Áhä‹M]£=”ÝEÙó€dà@›åAigpÁ=#Â*¶sþ'í Zb \Ë\;¸c@Ñ^~É×#mœ0*å—½¸¸OpïâÚˆk?6©Æ/ŒØ|–Aî+] |B€àºˆ¶Š,ÅvE>¸ŸD8´âŠ2¤sCú+µŽÑfè4F˜ÄÇý3â+^7×iÄG‡:î} x)³â5QŒ‹Aa5_V* D\ƒä‡ãø{>÷a:òĉ;ÂÙA‘:êÝûhKä)üÐ1¥Þ¸¶€Q¿Ñ&ŠuÄ@.â¨nËùg©³HÄE=PO•ÚéˆA+÷îW¸á^Âõm­Øf¹çÑA¦}Ãä<¿¹gS×q¯@œ‰ûL±ÿ§Æ žôÓá%l¤<â9‰{ž§ RÀ¢¸ö¸"¸ý£\QæMª¹ÏÄgçÚ!¼b½ÄósÑ¢½ò6ŸzÁ]ä%ê†røpCþÚ»ÿiá‹:¥Æ³k3žO”{ñº*CÜéGÐÜ’î²@ĵØ^ÛçH\¸£¿ûtà&î夓0h?å4ÁOY "<òK|e"_äþåI} >P¦Ñkž°xÞtÔ¾8ÏZÂGЈöÅ9þ–Û9× ÏbD‘hÄík¥(.éá¹Y¼'D:øŸxyŽ’7©ãa†_:ää‰ÎxÀÓ¡( D d£C'>7ëÈaFéä0€ˆ¼Ñ!*N$¯Ñ1¢#Λ üÄ›ã„ÇïÈiBüŠó¤Ÿ² ³MxQVœ§“›Î \näÑA¡~¢“LØQ„Ký‘VnjtÞ‹é$mÑá/ÂŨ–!Vp.ÊŒ¿!á—A” uÉ žtP6ÑH¢~(SáÄühg´=ʆ²"ÂåxÔññ°­¨b|¼AŒöË㢠ҡ'Ô!â~ˆºNq†P×yÂ2â.¶ùˆðK{‹6C¼´eâã¢<9†_ü ˜’/\ ¼bÃy:Ø” "i‹ãÔ%ù⺣ý“¿H#uG™ðPlï-ÇbFP´AØÅë¨QçYÂ" Q'Ô×åuÉqÎSn1H¥ Ež"_ ’©K:äqïâx= 舋r‹ürÝ’ßøÔ5×ÇñG:8F¾h7ÑÆ#\ÿ´GŽE{Œë‘z,¶ÕJ‚ õJ¸¸¡B\ø¥£×>P—Ô?eˈ a%„<â£n#¾²Ð‰î–¡ÐÖ£nXÅ}w"âc©yƒè ‘FÂdÍm ÜŠynòñ§ó΃0Ê„pé:ÞÅø›ŽfÄÇ@;®ùhG1x@Œ‰z'/Ea†N(ç˜!CÜÑ!¤l sĉ»ðG¸!qŽŽ|Ô-ÏX‰_Ú8ísÔÒQˆ8‡¸Â.ÒIû‰vH;%êœpisq=NÜwô`æ\”%ágYü¨Ü#bDÙ1Èæ¡颮yÓåÈß(ߢ@D~¹/Ä€´8;èÔÇ!ÂAù%ï‘_Îѱk¿=âzÄ/m'„Ü ,à‡ú®U "ý\ËÜ9tv¸þ“vNÞ(E:ã^D|µ DŸ~úi.«xCÍu†öÔ!-\;ÔóÿÏÞþJwWbÿ?‚ A¤¤ t' ùÐ0:v'H:iÇm±%Ëš,Éšeͳ¬y°æy°¨y5“¢æ‰)ŠÔ@‰¦(Q²(’")’Oô+õz½yXuoÕ}ë¾oÝ[k uosöøìáYgï]ÚÙ¸ýa^—¡ôE JY¥? 㔥Þ9ýâLÿ5ïûb«ú¾«ˆ„iì"*+3uDeÍ9OÛa Óê3évM–þÅ5}˜@” véÌ2áí0[ÏV…µŠò:ÚÅa¶/íâÍØœtêO¤Á}ž9Ì9tmþ+f©§9¥-B§|QêZ½æ|ãwV1ºQ-öe.áYad®$=ˆR7æ~î74^û>e‘2p]~Ôµï=¯<Ùƒ0S¶â2oÉø&=ÚéøÂBº„—>Aÿ¦ßŽéóŒgõ»â!d¢~8gd‚R¾ÏvtñK‹¸Òh«Ú4ŒQÚ†:õŒ°çgù_½¤Î2.¦]œt²qý–™þ'Ƥ<6isl&?®?dW›žWã^vήõ)ú`vv˜@¼ŒÒC1†Q˜G ÇœC¾ÌI¤g~߶x6‘gRêO?¤>7 Kû2†—ÌKŒ³› tÊÙüDøÆvµiù',éÑOh÷æ‰/Ëè^õfN-OÆ=}Æ&a”Ûc¢-Qc=Š@¤1pœL$sÖƒ2“ ×dÃ@hŒH<VvœeŸ¥Â™„ø^‡6Nãì'=&qöâ\ž@$ýžÓгôÑdÇý®™ Äaö©ƒ.‘#q8ƸäѤ$ù?çÌ€•ç¼=—nÏfBâYŽMžËö:ÂË2êÀ={˜@„&i©·”‹zÍsD;éô¶$<'Ò€8Nd„ÍñÈäØ}©‡„'ýÞ’/8õ'½y˜ú“¦8Ý©¿ÔAìPÙHî1P;N¾ï DóøLŠ‰ ™”åL”Dh’éÞ±½ÝwÍsl„]°“¬ò–p~’0lÕí>Ž.'?N„²Óéoö™I¹°9žéDÊŠ=¦„É.ò–UÞ²=MyΩmt=7öÒrЖŒ‘êg¬+õ›3¡÷x]˜V¸.îa;£@$¿£ÝI߸‚]æmTì<ùÕÞ’_Ï&)sö—ôùT—&rcŸ’²ØT JÛ]°ÿñü›Ägòâ:rnôG®m*‰Oàe©N¥Cù™ y^üì;KA΢ÑASÆ®‹3+¸”µ¼Å&Ôý¼ïs_Di6>¸gç}­|(ëØQÚú¯Ueíou¯cÏQd´êAÙŽö•‘î±âEúFÈsSné'å‰íFô_–þÃ9F]ê{Ò—¥ï6 ÏsúÄÄÈ5c²ö“¶'/sH›‰ƒë~åö†ž9L ’o/ FñÅuãiê-+>9~é_´_í}Œ/ÏÅÖˆäi£s¿²RG©Oi˜?¿ŒéËƼÌ¢¹Ý/³}ïX&Ê2ý¢üE|œÇJÿ\ JyÎéå–ÕtsH"qˆ]ÅŒI±/s;ßö廌E©[×"Àˆ“ƒ5‹²ZNÝŽsKaG(Eõ⚶`lº†tf\$tg;–çÌO„en©¼}Ÿí®æ µùÔKÆZùuMúÍ”“v–Õ„ÒmœçÒÈž2ð¼6"]éãÓÖŒ•žíy.aûd™+)Ÿ¬@œ×ûI¥<”­6éÿqn´ Õ™y”qü¨çº¤ïÒŸ˜GKÛª~ä jS^JèïØñò("‡¹¢¾<[ËÆc5Žƒg#I—>$ç5¥_˜ßwÝ/ÆLe§½¬Z½~Ý/æOÄíO¾ŽŽ~-h‹ÙB|ØÜgÎØ;W—ÒsœõXÌ D[¢F¾©@øР§qœ8ÿ9Ó&á-ˆÐòàq0ÌÄ9“*“t 8“Q±·‚ó.ýÍÉÉ[%€³ˆ’7qq¾äÍD›ƒ0nw‰@$¬qYVå¥LMH34Ð"ÖxŽ3ä{ß…:Ro±’N×¥“``R¡¬GúNÇ)ìu"u6/—”š€™ü˜(eò¿l J¼¼½W>YA¤®Åå;e(Oc Ž9»È³™(Çéö]Ve›šç‘ÿI|ÊK؈|Ÿ­£¥þ<ãMÉYU⟗eÊ3Ω°#"•kzéÓ(ìL°}š,[Ñcâ,,6?ª ¸"·!i[œ¤#e5¿w$û4Èb–ëƆ…1ÿe3”–&±1ßï´¹LòGgÚ§=å¸MM`Fç`ìÃ\'j¤ì"qqæõ‘ Lœ õ™7S› Dé«Ç­ ®é38’œ;å¯|¥ã°É–2\µÅÌäõ¯&]Ù>œ8õ/£øâ;‚ʼïs=[:Üçùˆ@ä9ô±Üæu¨ßÖ‰Óý¾;è`âqy¹<&(Ms{°*+×î(iÛ£½»nŒW.Œ“äyœÚlœÌ8ÊÚ´ºR¯êÍX)¼8ðË"ßgeQÒïïDê&Îj!m>Þ&aV»å9iÉŠu§/u“&íÂR×óí'âÓç%¾u"ý$»v¿¶“3ßÔ—I|}ãí|[F÷¤?còš:‡6Åöc÷!Û—Ï”«ûÇ"> œ)3Û¡æ}ÀHå5n1S&Ò‘s3FrÂŒ£Ò+ÌDIﲕ±/}¾þ‘#;ôÅ£}-ˆ”é˜~õ¦žôÝâ4†¨—Äé>ó*᱇ˆ«æ O=Ë~Ç>éw#bÊ‹´# ŠK?æeJlKŸà{éÔgÈ«p¤=cR^HDÜ3®ÄvÆr’vm=õ¦¬GHž²•-Ϥ<2ïHÃqµï¼ÞO"åÝX¬ßUÏæ@GqšÓGh'êИ²Ìv#ÉÙYú¶z”U?î׶µ'õn¬?J¾<£š_kkÙ¢vœ<Š@$¿ÊÛ¸®ñ(ç¡<+w}†<åÐÄ«Mjßò¡?4šÎï=ˆÂ†ô™ô+G9w(é1Þé+•Ó²>£8¤¡gIùÙ Dò.½&,y³êàAy– G÷DÒir› ØøS¶cYsÚâØdË‹x|gÂ#|éΙC¥A'$D€lé;弎@”tÌËe™@!Æs&V™äŽTó_1s߸õL^æùK¥Ý!–œ.i·˜Í›¤ÕIf„«uâWâ3Hkܪ8oê8+\Ôq¢ñ ° õ¼\<—U?îÉÊ9“§/âccó4ªï¤SÚÜ([SþÛˆæáynn/sŠËDŽóExð·ï†¶³l Vnqå-Îâ(éKÆ7E>S ÊöʤqìÃ\?Hû$(&®¹øü±gùNÙe¢½©@”¶oÅ™ö!¬Äü¸WYÊã¼~æT†‡R)s"QÚ)ÑÖ½é×éûîw¿û-ú>}f"e’mLI×¼9˜wÜqÇ¢ÌÓ–²lž'i’oŸÉË2±fŒK_•ò[&±ç±ïò9 Dì|ÓÄ9 -‰Ós¶³êãr°±þA¿!ß$á\Àç\ R¦±+q©G[ð÷:‘rëÅßsÈwDŒŒïÒ=ŸI£ñ8BÃa‘ô‚=œ4rêħü 7Ï K[Ófظþt.n.;¤:¶/ +ƒ3FŒ[š]O¿À.;ßM˜óCª³½Éss¦üýµ¿ea«WydKî[e_Ë"̹R)»Q RÖê1e=¦IÙD ç|‹Õ‘4Ìû„„nìVÇ3º´9c´°ÙŒqDÛÑŸ°=åF¿Oÿ+/Ï¥|Õs{Dv4öÝY¥˜ºÌwó>yˆ’¿±­é×ëwOÕ…¼kÛúx}¯68¿ï0 ‡Ah5Ör^ûÕ£þÉ܈$:á™Àdb±j‘IÜ89NYg%û"™hfÂ*ÏóÕ.y“î~Ÿâg£&œl‰MÞäKØÛˆlïHÝ)Ÿe teIw"a› 绌;Ï£ï’GÏIç(e{ÉŸ°Çtyk']qF×o¾ÅlÈÀøMƲL½sʳ Nù5ÆmuœE°´,K'²?ÏŒbÓq Dó|‹3í²ú ñŒgxåoíÓÀ<ïc¤%ç‰'o÷O‹@4 6YA4:ÔV6ÌE ùSÖ£˜à ï0([^ÝÈ÷âðiûƒ2ÑV„Ÿþ?udÒ3¯£‘ÊpˆäÕj qaõ©ÂWî¿ç}wXßço}Ÿx"ek ª‹y_‚Ò”¶ipÏ.DÄD¢jlöœ"ò‘·ØÊ7Žæ*H2¶ùŒåˆPß!Üq…ƒñ<åÅv¶%) “ó”E¶DÍÇw}À&+ˆ\Ï‹+ó‚¼ÑŸ¾G|±wÎçª~Ð÷ÊU~¤Y?d.3ˆrÖϲ4ÅFÜ'¾Q°R&Ùò†›¬ –ùÇøë«(̹3>ãoㄺOžô×úl™Pw©§ãˆ”‡±?õ£îæ+ˆÒ'°5ŸVI‹ðÆ-cÚ…U'Æ mLZ Ìþ×vØtDòü(BÚf^n*ƒeÛMØzúó´«qêsGûOyŒ+—Ù‚4öxÒ·™© õl|1ΩWvq¿J}Û„£Ýk—µ±ƒ˜9‚±ÃXÁžºÂC?ma[ìrœ{mJmÝ<Û8{ÔôlÊM"¶o•Ÿ.?¸£ëoæLh3úãáÜ[‡éÛµmîlÎbŸú ýˆ1`~Ï:Tú«gG9O©Ü>+m‰:à¹@”3M4¢e4\ ‚î7©à´zÆ5“'“Ù8Äg#ù^‡0ÿ…Šñ9ƒµpÅ%£@¤ágrÇšÇ7ß6Ç)ea‚bÂà^ñ“šž_¢ˆ@$\“ŒLÞMƸ”³|pdãsz8‰‰¬1!ã̈+§ÎTëœó¶ŒêŒ–1¦m DÒi€Î›Eu>ßzà9ŽCl Õƒ`Úü¦‘¸LrSÄ(ÏŒƒ–4Ɖô\¶˜¹'«­Ä'-ÂÓ©LÙ—kqrHµçÎ…@”v’t&ïËž•f‚]Þ»7Ûkr«p îŸ×¥²ôÜIˆr/»õM:Å[6áåPÄùô'3vŸç„ÁåKZµ­lÓð¦?ñÍ ×­H|ˆ”%1Ah¿ò–2TìFxÒ´¬½ß:‘ÿÅcBç^ùPòjž²Ög.k§Ëú¾M¢Üé{—9ôêQôéú}uH˜ÙHúå)/-ŒUYñ 9_*ö|THš6žË6]ñq¨³•·% K{MX>s pž.‡,¶¾Ž@4nËTŸÙꆕ¶¡žØÿÜéñ?ɹs±qv+MëDþgÃÊ%öª¯Š@„Ù¶,œMÎ JÞ¶!mF¸ò“9j‰WÇ%ež”ñÍå%ÏVÙuµŽºä{JþjÓƒÊCùÐæ•¿yÎQ¶Ì¹ß¸aL†Ïôéó{ËsË D[â2(*¨Ž~=—¥»È12ðt¢ðâÔù~SH]èHÆ%Ú&'ZÏzƤ [0L`FH§æ™L4è<—p3ù3a~Ô8ñKÇäš0ÇCcMÐtRù>q™€d‘ø”'õÜõL²8¬ùÅ7sÏqL8qÊJ|ž•fB‚|™¸˜°d`J‡4§kžÝ–@$=ê0Ûøs˜²ô)^“EáÊã(åp\4i‹s{!b=gRh`rm\Adb5–OuŸ³§Äkb¢L#h zŸÇG2ŸÏ › DâÒ^’_“>6Ì—°sþNìЩÜL4RÊN£­Ä2áͪ­mDHüàxŽ“Øe&V»“Vùd‡Úrl4vÀ6³RHú8bÒã:a0ˆôI×X'òÅþ3iW—Ê`\AxTýüâ6"ô¶:}›çFÈ/Á%o"×Ùó(¨FxK¹êÿâ@bñEF”1¾ñ sñ™³Ué’nÏ°½<£<88£}ÌžÓ3ëD¾Ó&òæ<åíùQüR÷Ò=ö}ÚIú>å”­»› Dò©Íçü8íUÛío\™'>"ˆïwA "¥î’'¶œºÃqûrÓöýºQÆ-íÉség|ïÙl5ç¶"ß³ñ¬¾E/rÆqZŸ’Dë DÒ“_”NéMù'®¥}°ÇQ‘u¨]‰7‚LÎ_íÂó«"ßq*óbH¾ »‚Å‘íÐlnÙ6ê‘®‡@$?ùå.dkãœÎóãX•-u®m[ RfãjìˆÛÊ=m•¥ üí{áŽÏ¦ ;rô#¯ÓWê&¿ª‡ä¼pKÝ˧gQþôÊk]Hú|ª;÷ä—qc#™ Ìëq×)Íúhé7¯ôDÌï;Œiú4u§ßmg] C#¤‰8qQ}‰_~¼(%R°¥MÃq¿øú <—õ¬.rvÏüÚœìÒKó&mŒÔ7­¢ü™çÃ*.¾À:ñLŸN$4çÐW˜+Í9Œé³ôÉY½E¸šß·.=«/0ŸÓwí½<¬@´%j,£@„ﺑïG'\¼&¥:.o)L¾ È™D v° æ:iž¿3ùfV™fɯ°L}gà3èËOžs0¤£BÎŒ´dâ å!/qhFHÅÉáÉdSÞt^ùIPN¡ç¼ ³­dÙDt…½-(“#¬t¤¬³´Ò¤Õ}Âó}êJ=xÎ@­|“Goßr0ŸA(ot-ÿŽ“%㶟ìÏÀ`–ƃ®r€y&Þã>ñ)“¬Ü_&¦ò¾©@$oÒÉ!K½¬¬Z2˜±¼iÌ„vˆ º?”—IŒ:W÷&Ï&¯®‰S~#¤E R‡ì7Û|8P7årØ@«”ǘVùQç9äƒgÀuŽ¿2Hɳ¼Çîü-Oî‘ŽK®GXòìQ"ùÍ/L¥Ü“_ÏmS ’Ne›~Hz ê‘S™óÉbÇ&i‹ø´)é×µ)}æ¸UR_cbãñ #ñ±‡ƒâ“6"ƒ´çq™lyF?©lÆDÚúØ/Ì) ˆ<ÃV¥—èúŸý§?EeQG=èûÆ­gcß7þ,zÎnI™v¸Ž@!h,ku¤\ô%ú ô ÂM\Â݈x—¶+íúíNÛ`Äö‘>FÛÎVÄQ$]G Jû’Žñ¼äq´·)ù>‚¾k>Õ…vÀAdW¢²Ž@”¾ÚsÂW渻K¸Ò£`7úum++‡\wŸûóü(%M«l_;Kù¹?çÝ ‡Ó™ùD”ƒF¶s‘ëüÔµ2w?ûÒ?ëÛGû"Øå,—m DéKµÇqaÃìÜ<Œ¥íø_˜ “Ý/ögÞ.Sïêj´%ñÏç›lUìûÄïšö(mÂf›Dʘȱ*vÇVÖù1ˆ]$ÀœYÝuõòä|ë¯ômÂ;Êv./føl[ýEœßwóËÀÒ 46ù~ßA”vý¤zÕwÈ—2ʵùýÇÁu¢ôMÚ³1žßvPŸ»Š±y}žzЇšoŽñQŸ Pö9fcÓpÌA#ì÷;æ>óû£|µÙäK_yÛ*·Ï D[¢Ž˜ag²g`òi ZF¢‰§3bH&|MÌ2ø t¡I±‰ZPó8YUöϺžeÌ™ÐqX„/®¤%s(Ç_„’'yÓ`M`2±ö\&:¾3ùð¿0³¥Í$¤<3æMºÆ s¨5AÇüJPò!S4!0aL:LdäY @Ê5Ïa&–q>LÊÆ7¡‡QÙêŒåSÞq”’žy=dÛæì&u ƒ%Xy6Îu&Rù.“'§P‡lr”ús=«’GÎR&¼â‹@”)ŸLŽëã¯ÃˆÏÄV™¯ŠOÞØð8iˆ€{ËÚ=&ŠIƒÁS|îQ/ìϳ _šbþÎÿœ·ñL ƒLl0å)ÉŸçM&´÷ˆLzÝQ,e&!9#K¸ˆ¤Õ¤+iDa+£l™ÛÐHÏœÕáXÿáÜ&äkœàJ»ç95©“Ôeþ÷©#„HSl×´­eQ¶| +‘¸Ç­°®«séWF¹]Ï¡­ìÜ„(¶??Èõˆ§®g› Š7öøR¾‹îoŽ·°<§þ²ýn´´}Ïû[ºÒ.|šèľSŸi¾‹cæo¯|È·ö+í¹æ™ÑÞ\“žÃÎEP¾&ðé'çö02å,_D·Ô]ú>ÂlÚUÊmìûع‰›ða2÷Çáëh¬CÎ…¼‹Ë¤W??Úªü§¾üŸ•Š±!`Â2±ŸÛƒI~ìÁX#Ù‚’º\&¥ÍËû\ Ï§Õ ÚJ}§œ2–ùN™âÛÓ? Ní)¬Ï´/”.uê1‚´ò‰Íh“ˆÆt‹w¬—Œ+Ò¦¼ŒGIƒOãPê…‘ʽ\ÈK§Ã¢Ô =á°Ó1-é'm±¹ØhlÎÿIkV-§”ה뺶¯^-yÁ“1*víkó¼Œô (se±‰@”¼åÌ\O;ÌW£}IŸ|ú?óØ—±EX©óe‘zó¬¾j.#V¥)+®Ó.¥#cEÒÈñŸQ>£QÜѾ²Š… 7õ&Þ1͘~‰Í$ ñù;u•ú6OÔ^ÜÏ&6êsG›Ky¤ÏV~ˆÄ)„éØŒqð(2Ÿ/¦¿o㩺™ß³…æÍQ‰¨¿9¬ÍÏ©-ér†û>JYŠS¿)_GYÅ4’@®\ÄFã@¶MÏï=®+±[" áÃÇüö(e‡ú|¾›¹QYÝÿHñzNZÔ£BG‡ÜO(–ýîQË=c‡— ê‘ïwû,ˆ¶Ä8œJ“CŽƒHHÊ/EiþŽ €&Ââ,3YE¡NÂ=hB0Ì™0I‡ëã/m¸æÍ¥‰6ÇDFzLN &*Yýa" ×ÛDiÏÄÁ=w ‡ ñÅ 7pû?ŽŠô˧xM‚Q™Hå KiŸ7IÒf |+trãÊqÒ*Nƒ…|Üî'½ÊC§i^‡«(=Ê;âVÊ[(yS&ó7]þö–#Žã¸ÕDyšpò¤+õ+½ò*\i®ëcš}¤•EÊ5õåYñdk™ûu¶£@d2É1Ìù)Î#gÌCâóiU|D—1>ä< ŸmÍÅ8y—Fa©£ñ ŸûØ ÛmFúØ'–øågƒ¨‘ÁD¸&ýì7ÎjìMÞLLâT¸Wž|/„ѹ`¢~Ä%loœ„²o6(ÿž—žœ¶Î@~å“šèÊß8ž>A~–Ùê¼Nb?>¥‡Ó(ü<+MéŸ\笎ÂEkžüf{â8é&$¿Ê4?+î¢g¤ÃäÃýìÜ.¶?w¨åߤ aoâ¥Y¾¥%}¢¼‰C™cÚgmަĥLÓ¯x^8ìB;ž×Oâ“Ÿ1>íNÿ¤_L|ùU@a¨kŽ*ÛíM¼ÊYÛ^çŒe¡ïÆ|œ)Oò íR7c>¤Iÿ”¾MÍû¾´5Ïj»òæúü¼'ùÌËÔÏFÔ·¶Gøû¯ØŸö›mlÉ£‰`ì!ÂÞ— pìA¹z–Ý_|/ÿã¡ì>Åá{×õ9sH<ó8SßÊr´a°9}Œ1+åÇæõ cú³]'é§{ô'ÊÞçؾÄÉ‘‘Vu)áek¯g´IéWÆIwî›+úiáH‹úmÀó®s’Ù¦ûضúã(Ë<ûþ0G-e,e¡OšÏëò¼ ^¤-d|fò¨ÝÎÛƒ2L¹Îí}™íge¨8¯0ä-É\PXFù2RÊž`·ÎªáBÒß±ÿù3îáŽe‘>]=è³ôϱ/ùÓWçe÷~§¯Ñ§ˆ×83¾ÌŸ²_–¦Üc~æùyŸh|[¶='ã16áêÿs_úomÙû*áÀwĉy.-êUÍîó,ûÕçÆþõ¹c}¦<”ƒ6k^0ž}â^ýcl_Ú…·,m»Fyàp³v¬.ÍÁj£«ÈÆÕ»q˜=©OßmVú5a¨á™lFÂa³Â@}½ôlŽ:ŽU…æúñ|?¿÷8¹Ž@$oúÛ¬¤Tvã|u]Ê_Hæ/9CjÓpÐ C}¥ñÅ1¿¾Åkö"B_(=GIKÚ©þßrœoÎï-Ï+m‘òk¤C]‡YZ©Qx–#­#×ñ™Èç ×4HÉ‘ÁT•pt´ó´¸ž´ŒK8MöLŒ„—gói`æ›4™¸˜Ôo^ÒxML¬bÒ9ä­Ò<_yƧüÈ—üIƒ´e°ò¿¼Y•!{iÔ™™80…› ·tšÀgj‚3Ÿ@¸O^LL¶°|c& óòVùÞ=ë>‡Ò"­ÊX:u¶>å_Ç°çm)åJÉW±5ã€)žD&f9Ó‰s¢N<ë»ÔÉ<c|D†ÃâC6°*í8æm¾·>íÍ°1q™0 S\Ëì:t]™x²7=¶.Ïãà3¶×eaa®ûó;2IPÊŤuþüA”Ï þHŽÇÓäIÞ#²ÎŸŸ÷)~ƒ½ç|ª×ØÖxÿÙæWxËò»ª/m_}ÏãZõÆþÕcúDýNê7ÏÍÛ\êEÚØ´ú÷¼p„·ª<}/ ›Ä§<\WúAÏ¡~K9Ííí ŽíaÕIò°*ó¾OŸrPß—þž·eéš×alB`⪲cåæú<®ƒÂÂUö’¾dn“8ÚÐ2›^gê[½i7ÊJ»qzó©ÏQ~úyXËÊjl_>çו×8?ÞUmrUºÃemGXì7oÔ…›²FÎ!ã0ˆ"‡9i® 7+2<»ì‹û”›r±òD™êËôiÚ¡4¬j›Ø~Úܯ8³ÊŒ˜¡ç}ß2Že¸Ì¶Vq•ŽLÿ•-Üúeýæh_æ8q°Ä}P[ 4#Å€IDATœÛÖ<­‡¥Iz<'NíÔ¸j¾™ôÌïKÓX~ógçñ§<ô‘Ê„]è§]›ÛŦmm^Âç ‡¥o(ò­ýX‘Gb3ó²Y—æÕÄ/‚W‘m°"½붭9Õ±1BžGÉWÚ:»!Šòçõ~®x˜@¤ÑüÃK3}áQÄ!÷‹Ç|„@«Ý®Šó fìñrØvsŸó{Ö¥°ØÑU{Þ4O ÃœD:äËœé(á”ÇK¶Vh‹dä›p|NÇ«Ü0ý\×Áú~>9š‡³*-þfÎáAµ0M³uÀó¬d *, Lç8¸Ïã[õÌ‘hžF‡TglÙë:çì;Õ¡7ËÊNz“P¨,Ëòô0Ž‰s¥AÆÎqÜÌønõ-t“qÓ|Îj/mð€,œÆó=†H?gÅ\FšòËwóûÊò¤Pf×VÍÅò¬U&¶ð™Csâ7iç ÃJ+`l%r^aÓpäE²ÉŠ&é9j¾v˜¿:–á|ø&¡~–˜wP^¤9[?’V}™²# XUv”UWîg D5b n~&è:Œ@'œœ=t”|¹ßsÂ` V¥Í˜(w‡ˆöŒ(R=r¬ÎóËBóûʲ,Ë“ËŒÙÞ´[Y“ñÝÁÎíó9?U>ö0šÈs1&öç{~—IÆ5ÎO~er~_Yž4eî'žó­=XMx”³‡ÜoõsÙ„¥/YvÔaÔ‰„¡Ì?×ý‘ÒãX ‚•¾‡àq”íZÛ¤ò8L £¦ÑsVÉ3ÁŸ¿3¿g Çö­ì9ª°¦ÎôùÊÞ]øÜ4÷+7iá‡X…duÛ¦á”玈ö”&W¨ÁZ:j@ˆ1eòxÔ)ÅŽ¤O: ù¾ËúL}Ÿoá¯,˲<^óÛ¸™ƒ‹óýüþuiü8Û0¶ÉÌ:®•åÏ+Nr¹ 8j›ðœ6•¾7 ÇýžuvŽÕ&hÅȦ?aî^ý¡Ä¬nSNì’}ÊãüÞr7Y¨,˲,˲,Ëò~ÅW,¶E$âØo*Èl›D!«`œ]s”³¬¢±jŶ&+’¯ù}ç“ÛˆˆBV]Ù–çì¥]¨G$èØ&h…Ú¦é‘öéy+Á^ýµÉ“Å DeY–eY–eY–;NηÕVé¼ÿýï_ülø®ø[$tlº-Ìs¶”9‡`"ó{vÛˆ²êê /\Ô¥¼[9µ+"ŠtdõÏüÚAôœòyë[ߺ8àš}æ—ìæ÷–»Ë DeY–eY–eY–;Ì8Ù~>þÕ¯~õb¥Í.:ßGI±Äê¡ .¸`qþÕDG 縹 (+£ýÚ×¾v!¦Åv1¿›R¬j“/SˆvmXy8+•eY–eY–eYî0ì×½lÝyó›ß¼8#æ4øZò`Ò;ßùÎ…P²Ë‡Ÿ­@D@!~]}õÕÓ?øÁé oxÃâPo+¦Nº@”UQVC±QBæ.U®Ï DeY–eY–eY–;Ȉ ŸøÄ'A;ÈÙ¯Bñ³Nƒ¨ð­o}k‘7+ˆl1Ûå'ÛˆˆCêÑ õ˜C¹ç÷ž$Fr¨õg?ûÙ…ø¥^7=àºÜ V *˲,˲,˲ÜAr¾m%³­ÌÁÔÄ”]ÜZ¶)‰"V 9ƒçU¯zÕbå‰_ºÚå|­@D0±òëE/zÑôÑ~tq¨÷I‡Pù;[ËÞõ®wM—]vّ˨<ÿ¬@T–eY–eY–e¹£´µÌ,Ñ¥—^z¯ë'•¶!9KÉ/ŸYQC0:­Q„>‡‹;À™ÐgkÝ.çw]Úþ¨þ¬²ýÑʨӯ}e¢²,˲,˲,Ë#'Û/I9ÈØY5VÛ(NºómÕÌ•W^¹Jl™³ªF>m¥›ß»K<Èê!?o•Íë^÷ºE¾OÃ,¶øñ|!]tÑE‹•óûN ×^{íâ@j¢—üÉëI½ºlµÌ&>®|ñ‹­ò mÎZrF28 y>ˆÊák_ûÚbU”º¼êª«yßWž,V *˲,˲,˲Ü!^~ùå Ú²CPp®Ëi8¯†ôÍo~szÙË^6½ño\Ú|R„’M"ùr®’z$¢X vÉ%—œšU`¶½¬l#`ž†³±Ê DeY–eY–eY–;CÎ÷…^¸ØŽtÁ,Ϋ9 ?…Ž¶–Ù’ô¶·½mñy’~ ý(‘û‰|ÎR—Γ:Iy^Ee`õáK]úÕ²Ó `–ˆÊ²,˲,˲,wŠÎã±ÒÆ Û“æ×Osž’-V§–/+Næ÷í27ˆP¾åÓ=Víú9Këd»#Óã~Áì4ˆ—åïX¨,˲,˲,Ër‡Èáæˆó§NƒO%??go%Ï“&*H¿-c›ÖGêò4¬Bù±¢Øg› !³«‡N+•eY–eY–eY–ÇNB‚s‡¾ñoœ8_Ñ¡áGˆN##xYIä—Ëæ×Ë“Ë DeY–eY–eY–å9áI]mRèžW¹Í¯•'—ˆÊ²,˲,˲,Ëòå ¢²=ö±žúÔ§NoûÛ~Ṇ´¼îu¯;+?«Qq\8Vè½ï}ïôЇ>tÑ¿ò•¯LŸüä'§¿ú«¿Z4ÌÄ£w¾ó³'ˆÖÃqDV“©¶° â|æ3Ÿ9ýüç?Ÿ_*Š¢(Š¢(Š¢X-¼þõ¯Ÿõ¨GMýèG~é‡>ô¡éáøô7ó7ÓwÜ1äXÁ}ô£=}á _˜_ZˆŠã± D·ÝvÛô”§L|9÷9悃ý×½xá^V¿ímo›ùÈG.ü>Îk_ûÚ…Ïëy~̇?üá…/Ãâ‡ò…Ãõœç Hîç'‰Ë½V¹þñ|?ÇŽ ×WᤠDʈÿüf~ß|MÇÒŒÂÝÙÀn¥¯ýëó¯vQu•-Ä*7 dÎ5ŽM ‚O}êS‹-fœz ˆðw÷wg®G ¢èº×µ÷½ï} %—â¦!jlAŽNi½üò˧‡<ä!‹{ÁµûßÿþÓ»ßýîÅÿÄ ÷™Fœ­@Ä…ûû¿ÿûÓŸÿùŸŸéTÂMzÕ«^µ….½ôÒéßþÛ{^ rt‚ĹUJ$uuWV DV|éuJî'Ôgˆ=VT=èAZ5ßùÎwê/±ç%/yÉ¢±36‚Ñ[Þò–E½1zNÖwW^y墼H"Ê~RÔwÏ»Ÿ¥¼…¥¡Ô‹ëñüô¦7½ižä¢(Š¢(Š¢(öü2Î:“0ÃWøÌg>sf% nø+|>?åe/{Ùâ%µkîçw~éK_Z<÷ùÏ~ß…Ä7âÓ¾ùÍo^ø™Äž¿ø‹¿˜>ñ‰O,|ÑOúÓ ¿Ñ*&⎗æÄ#•¸ø±þ·°‚,¦U8ª@D|â ÏWGùŸO; h––¬ÊµìÚÉ.¾eîsOŽ’ñ½0Ç•@ÊÎwüB>%d¾ÛÇ5bžzð¬ûç .’ù³â÷ýøãX|U~¹ï–åäi K;—ü*)¬ïÉ*2aÉ_Ä0¶„麴,[%& vsªÆ2),¦ ²êlYYmŠcˆÀäU"ÔS‚Ñ+_ùÊ3«C4BÿD bƒ†¥@4¼Q¤PˆDaªdbqIE‹ãE/zÑBtb“ðçx¶!Kãf¸ög¶ˆ#ð°0怴+‚CÔyTæйײ: – DÊ\ÇJ„“oÆKäQ'¶iæuö‘|d!òiìÔx‘8t†:Å€ˆEE¾º•f«¸|'<¶D0u_|ÒF8œ ‚EQEQEQì78½^:óø›ü´·¾õ­ …ÿ@à!ÅÇq?ßÕ¢ ø|N8a‡äØN;ßä5¯yÍ™¸ìrÿb°8ò^jÕ2^˜~äAØT âó«ˆü.ùÎGñmýo—†kî#2ð×­zúà?¸ð{]“~xî•naûN~¼ðõ«_½ðõˆdüN‹'"”X0¢Ìlïsþ“URVP±«ˆM+à×óeå—ئl¥A~,J~¬ïÄ)mʤM>ˆwÒâG²æ?†dуç¢Qö¬i @åDâOÇœ¥ EQEQEQì7²Z# 2ðU¬ðxà¸Døu|¾çès’ˆD ?áǪ";$¬(²2…@à>«Š€dC^ˆÏ!-Âxœ};(ÖÝVµ©@DЊÏÌô²ŸN\ \nˆbâ'È[øb/~ñ‹¾a„èá>¯<*?äßñíùm„ þ#?ÜsîS~â rDD‘~iðÿx D "‚ðùv„$i³c…BÐs_? GäØ"=4ù‹xÅ÷õ¬z'åx”€dÅŽçÅOoPVÊè%-ⳓ%çTÉ·Å ´‚–#T„­Ž‰ZêXyñQù­êË3YàBÄ’Ïq••°¤SzØ—´)CiPÊ‹¨F<ã_ÓWä¡ü¥CÞº’èØ"BFæsj˜ý™­‘eKHC$ÈV#•(T+…¨|KôT€e ÉyF Í÷ËÙ¶" ç0ˆò‚¼àÌÿ»Fž-]Œr™8Ë¢Ô…˜XDqö)»_½‘0ˆ2D Ê¡åˆñk„ê‘š®ãDuêžqÉ¡ï/¹ä’Åê&oÆe…EQEQEQì'8ׄŸqKð!ø«Pø0†ø8ÈGá“Xa# ¾'ÿ…#îåtVˆâÏðC, žŒàÛº‡ÿ2ˆÄuœ{ô÷ˆ|l‚ae\QÃwË*bÿì "Jd‹AŒáß@€oì¹øƒ|c¾;A‡™…„•ø„#Ü'<+v¢±D\ÊŠhä|«u"†(oB‰ó‡¥5>ºëãY@D©e¾ï…^¸ø•;õKB¾-?WÉ׸ʋݰå'Ê.Wù¶X$ñÄA0" y~Ù0¥Ls<ñ¯x†åùlk øŵ£®":6HC¢ºƒÆ@E5DË"G¦V¦($§p-çËVÂ& h”Y¦Ç0üO|Z¶Çð¸"ùïþÝ¿»Ç'+[4º]ƒr“¶ÿïÿýB ^†e‘NÈÁÛÔãˆ0 W©›t® D:L†¼J RïÊ5çHåºøÙƒUCÄGõ¸·Óò:u?_uVEQEQÅþÃœ,9ó…_I€àgpØ}Ï'!Dt!~xiôé8ëV}æ‘°­þàóF”â/‰‹¿š¸V DâçÉ›Dî!0ÈÏø‘ˆïÆ/ÏßáÓe5!#ß˃ß2ˆˆ27ñ QË" :²«dÄxQ@°R&Ä&õH?N>¥k î!I·{ø«>G_s•@$„@¾%› 7ø—YdUÑèçÒ6äUˆ²­ Ä#„0ùçL!º²Èv¿q…D$‹®`•ÑxŒ•YòHÓ }äLfà_Kç¸Èf›@2‘_£bI(Q‚R›‚u5 ¢Õ#1`ʚƔó…üÍ(FáGc%ÔØ*ìŸüÉŸÜKµ ŽK Ò¬ ŠP%Î$Zf€»ÖØXæX&ÅÖ6>±Hã4@Æô€<àÌžVÐðu­“¸¤“Ö`tØ:ÔÑà¥É9FâO‡¦Ã´T‘°uœ%z®"ÙaIã<ê’º¢(Š¢(Š¢(N¼°æ3ða¬±ó„ßÀ¿ŒãNðð_5~'Ç<Û”l)óœÝ!v,ðsóÃ<îý¾—â„ ~°O~/¿‘sÏßÍ–4 ðˆ"|ªü:ó*l"1†’•*|)>Kú²úG^ù^„BÁq D„[Ûæˆ@4®¦’vâ‹m{VöÒ!ˆ“² nêP¹ÒÖˆ”©ü±º1Jþåß,ŸÄ˜q%ôä×é,b…•R-kê€?,ÏYÝã“?s”‚”AŽNQV£x)+ˆÔŸï•C@x$P.;{«@ Ó •ì…SHã»ú(Œìy ‡F-ãîŸ;ÿD _”O†AQ\µŠd[‹ 4î›Tù~¹L§qßûÞwÑYµrÎ7V DÀp.FµUÖÉ£–QŽ*´ÎÁ’ÀÔ›ð …Wc¦¼uåoõ™SÞudê‹’«A…çyÅ3?”¼(Š¢(Š¢(ŠýF~hˆÂyçsò=GX‰ÁïäWâWø'¾·…‰ÿGDßÅ “¹ïébˆg|F¼à¿ˆ›¯ð‰„ϧ'ÿæ '}Sˆ?F´áóß¼Xç7üEøZʆ¸ /Ä‹q«‰@’2á f§ÏxX4?QøY0 n :”5¡Ís@¤qß(¸€û¬¬ª! ñåYšè„ñJ;=AÞˆ9'”¬ÌQNò4Šw%õ¶ D U:”(~,-Bü„.GÜø;+½r|K  UÄE6%é–NvH<›ÛŽøè'D"þ¸ô(ß©'yŽÐF˜#åWá¤é E ‡áØ¢]Ķ"…§bÆŸžÓH‰` DÍÝ:I8H *Š¢(Š¢(Š¢(Î=6ˆ€ïJ(ñ¢~Ü’äe¼ð|X¢QvÇxYOŠ¯Kà"F$N¢aÆÿ¾Ï"÷{. ²•èA(‰0–Ý'ó3‹Ýç»qQ‰xˆ(ââŸòÁ=+­ç¤W™œlY–giXµýŠ&­Á<›_ _Ä<‹ENÊpÜ*¦ŒGÐê0áe‹™?Ýçõ_ž>0\+Š¢(Š¢(Š¢(Š¢(¶‹# Dÿæ%÷ˆnýÕ]Óýßø•Åv³[n¿kúÐe×Oÿõc>4]÷óÛÎÜóô\1ýóg}jñ÷ßýÍé?{øû§ñÜÏL¾àÒé_¿ðâé?yØû¦‹®ºaúÕ¿™~ÿ·qü³§}|zƾ5ý¯ýÒô=ä½Ó;¾ôÃ3aEQEQEQEQÛÃÆÑÿñÒK¦ÿêÑšþðU_œþßW~aúƒW|aúçÏúäo¿ûàô¡oühqχ.½~úoÿ‘éG7þƒ@ô¬}k!ÁSÞwùô_>òÓ¥?¼qñÿÍ·Ý5ýOùØôœ_9ýä¦Û§ô¨N/¸ðª3Ï>ûCWN¯»ø{gþ/Š¢(Š¢(Š¢(Š¢(¶‡¢ó’K¦ÿîñÿîoN{×e þß/ÿÜBÔyûXås˜@ô„ß>û{Ï»èÌ5ø__xÑôØw^6ýæ7wO÷ÓW¦ÿôaïŸþõ‹.^¬<òËiÝaVEQEQEQEq<ØX ZœA4ÛböëßÜ=ý_/½dú_žó©…c‹Ù÷„L×ÿâ¢g~ð"âÒ¿záÅg®Á¿|Áï¢à£ß¼~zÔ—Nÿó3>±Ø~ö´÷]ñÛ°«EQEQEQEQlGˆþõ’Cª_xáUÓ?{êǧ_ÝùëéãWüdú/ñ隟þòÌõ?xÅ爞ðžoN?üÙ­ ¡è¦ÛwÀ5QèAoùÚâÒ~þË;îñLQEQEQEQEqöØX ú½\4ýË™¸oúÜ5Ó?zä§ën¼múî ¿\l»ïë¾4}òŠŸLOzÏå‹C©ïy¿ˆ@¿÷ü{n1û¿½öØw]6ÝpÓíÓ?~쇧?|å¿röÉoýd!,ýþ‹>;ÝùëßÜ㙢(Š¢(Š¢(Š¢(Šâì±±@ôà·~mzà[¾6ÿzúüw6ýOOÿÄôé+oXüÿš‹¾?ýOýøôOŸò±é~oøÊâÐéû½ñ«‹k/ûÄÕÓýßô»¿ƒû½é+Ó‹?~õâo¿fö¯^pñôO~û¬çÿèU_˜¾}ýÍ÷¸¿(Š¢(Š¢(Š¢(Š¢Ø6ˆÃxLÐíwþzúåíwýÃâ¦Ûî<«ç‹¢(Š¢(Š¢(Š¢(ŠÃ±u¨(Š¢(Š¢(Š¢(Š¢8Y¨@TEQEQEQE±ç¨@TEQEQEQE±ç¨@TEQEQEQE±ç¨@TEQEQEQE±ç¨@TEQEQEQE±ç¨@TEQEQEQE±ç¨@TEQEQEQE±ç¨@TEQEQEQE±ç¨@TEQEQEQE±ç¨@TEQEQEQE±çØH ºóÎ;§_ýêWeyNyÇwÜ뻲,OÛ†Ë]$»¬m–e¹M¶O)÷‰µ÷“EZÎ:ØH ºîºë¦k®¹fºöÚkËòœñ{ßûÞôÃþð^ß—e¹û4f|ÿûß¿×÷ey¾YÛ,ËrÛÔ§ÔW*÷…|´üà÷ú¾Ü=ê—h9ë`#èïþîï¦Ûn»mþuQ+5Õ³(Š“‡o¼qúÉO~2ÿº(Î;þþïÿ~úéO:ÿº(ŠâÈøñ<ÝtÓMó¯‹âTâG?úÑôË_þrþu±ƒ¸õÖ[ZÎ:¨@Tì<*ÅÉE¢bWQ¨(Šm£Q±O¨@trP¨8U¨@T'ˆŠ]E¢¢(¶ DÅ>¡ÑÉA¢âT¡QQœ\T *vˆŠ¢Ø6*û„ D'ˆŠS… DEqrQ¨ØUT *ŠbÛ¨@Tì*T *N*ÅÉE¢bWQ¨(Šm£Q±O¨@trP¨8U¨@T'ˆŠ]E¢¢(¶ DÅ>¡ÑÉA¢âT¡QQœ\T *vˆŠ¢Ø6*û„ D'ˆŠS… DEqrQ¨ØUT *ŠbÛ¨@Tì*T *Î+~ó›ßœáªïGÞ}÷Ý÷¸oŽ DEqn¡M†g‹ DÅ6±-»„ DEQls¼ƒ DÅ.cÛö^èx±Íúª@Tœ70à«®ºjúÜç>7]zé¥Ó]wݵøþ׿þõtÙe—M_øÂîÁ/~ñ‹Ó/~ñ‹Y(÷D¢¢8w¸óÎ;§n¸a!êLæ¸ù曧Ÿÿüç÷€W¡Q±-Ü~ûí ÛÄ;î¸ã׌=lmǬQQú cÔÏ~ö³Å\u„9¬ï7ñ}*» >{7þÍçqæìÝX».*ÔúP_ËúsôeýÖ*T *ÎLÒ¿þõ¯Oüà§Ï~ö³‹Î ²þÿЇ>4}øÃ>Ã|ä#‡NÐ+ŹÇ[_ÿÃþpºå–[îqÍÿ®i“D¢uÞhT *¶“¶‡ó ,;Ë5"æ:¨@T…±ÌxÇÑ-ΙþÁµë®»î^}Î*T *v6}ýõ×ßC ⧙«¥-¬ëwU :>蔯:ÑOÐé—\ÓOÍže8¯‘~ë[ßš¾ò•¯L_ýêWÏh ëª\ÅɇñßøÆBü¹ä’KÎDêÝŠ"+‹>ÿùÏŸ¡ÿM즑ÎïË_þòôµ¯}míý hx_úÒ—¦ï~÷»óKEqê@ 2èhw£@d` ]×v‰ÖˆÄ©ÝŽã^}õÕ;5ùÈjHýËaãa±]drƒéÛÙŸ7¢ìq´Ï¹¸¹ › Dî¿è¢‹¦w¾ó‹1îŠ+®8³Jv— ÿVçjSEQ Ž—þƒÃ%âPúŸë:ÍëDú ㈹¥~eÄw¾óÅ÷âÜ&ôÂýÁ~0¿TìøHl‹FT`÷æi£½kó•ºË°Ž@¤Ýðçó;߉§Xõ¢|òR60ÿQîc}©×ÃD¢ó*1¦ç<ç9Óƒüàé/þâ/¦>ð úÿa{Øô†7¼a+N{±›X%1Z*Ì`{˜ƒ¹©@ôÆ7¾qzЃ´°9iØï{ßû¦—¼ä%Óå—_¾ø_ög6½â¯˜ÝY§Ë¢¼YÒf〻ÇÀuX¾®@ô|`zÀ0=ä!Y´_ã†Oÿ?ãÏØ™I­¾èÑ~ô"}ßÿþ÷ý×ßþíßNó7³ø¿8>,ˆÔ;mÓß&¿VaÈK¯'>ñ‰ ›d§±Uãº+–ÎXüÇ<½ýío_üOÌ|ñ‹_<]xá…³;‹¢X&iÓúãÜد¬ã„­#¹þWõWÓŸÿùŸ/^ŒŒxùË_>Ýï~÷›Þö¶·Ýãû³Å{ÞóžE¸o~ó›ç—Š=Â2ˆ=.³wcäa>Ú:‘9à£õ¨Å˜É7ËüÎßøÃó¿ÃÚÕ>b™@ñz>÷ÿa/Æλ@ôÜç>w!½ð…/\l)zÿûß?½èE/šþò/ÿraGqÚ‹“UÓj![Ê>úÑ.>?õ©OÚ©À&‘ô„'O}êS+JŠãÃ2(¶™íè>ßf›ë DV±=âXLfÿú¯ÿza›DAB!¡èMozÓ¡qK\sÍ5Ó§?ýé3«^?ö±M÷½ï}§×½îu³;‹¢X&é?ü­póel\g ZW 2Ö7ˆÏ#Œuú•¼Û_L¸o}ë[ç—Š=Â2(öžUsÚÿm{_W zìc»ðÍÌëøˆï~÷» J|gÉ÷¾÷½ùc{e¨ýQ®é§Ü{X}í„@dBÿñüÌ÷ íYÏzÖâûw½ë]‹Œ˜d騲¼Ò§I—NL8–r¿å-oYLº9 /{ÙË1…Qì&ˆüŸs‡|~ò“Ÿ<´SM"N›I¼ŽÈäý‘|ä=–Ù[vïíÉ7¿ùÍÅÿÒ+=ìФÚÙIüãÖóŸÿüéŸøÄtñÅ/Uƒ¶ç}‹,ŒSàÓRIQŠÝr*×Ä«l„ç ‘-—E±kX&Úáªk«°‰@¤Yé0â½ï}ïbÜxæ3Ÿy¦0PúÞê mÎ;mÑÛXíÏwÞ˜j&ÚóAQûÔî]†mÐãd‰ôš×¼fzéK_ºøÔö}¯/ÓO˜ÜÓj "– Òþ™Ï|fFÆ0ÏjïÆ=20þycöŽw¼c±êȧûŒþÎÁýúOùô&yžþ}Ä2hk«°Ž@¤ÎÙôÒk´{umbû¸Ç=îLYy`|3N ¿³Âˆí¨w6c^ÃþØ°—&lÚ'[± 6«ÑäGý‹ËK azŒaÚ¡qÅ3æM9TÞx#l÷°ãg?ûÙ‹1í)OyÊ¢-¶­»(ö Ë¢u®­Â¦‘±g„ùå(郌kÚ½ñÈu¢oú9ó^ý€ˆ1Ÿ5–hãé;ô_ñ·ÄwÁ,Vò§ŒQ ï0F‰3}¢ùª{ÌYÁ˜il•a_‰  Ÿq¯~8­O2W0†ùÿÕ¯~õ"ÝüÂuûçâx°L ZçÚ*l"ñÏØéø}^f%xÍßØ»ñ‘ÿ?Š#æüGöd·Çk_ûÚÅe»·0±l”n`®;¶n5­ç…o>—ñ[>Üoœ6~¾þõ¯_„ãÿÃôãÀ*h~mÝ1}g"U&Ô È6žŽD8O~ò“+3¢šyÎ “0ë~ƃÞàùŸâ˜‰P±[X%±Q$dØëœß°®@$šÐ`È~òw #±]Lcéåx²+ Ž€¨ÄÙÓ™™¨@­FŠžØ©e‘±]ƒµ¥“ø¤'=iq?{&pòØ:GÂ5+”²…­(v ç[ â„‹ÇàDZ6©Í >íUXOúÓÏŒtíŽs &°ÂцµUmN?`2n“×µEß Ãßú-Ðn…ïÙ¼à‹>ÁV«¥Ë÷ž1É~Þóž·hëúßyôkD'a¸&Iƒ²ÓWùNº³Í i•öˆL&äÒåÚº€ÓŒƒD ƒ®­Â:‘pÕ‘zȪÒ@=›·Ø,{5ÿQ×Æô·UGBã!›2žø›-7|Ç^Å}î×éħŸàŒ{ÆêWý˜sôE¾ãlƒùð}îsŸ…3þíoû̘ÉÓ? ۜ۸L 2ŽyÁê{épÆÑ(ýa¶f¹×™ÓǃƒD ƒ®­Â¦ßGØæšEv‘Dl‰ ä>ãVlÍåW׈6|(ãŸùYÝ˦i ^ˆ°S㮶Ä.-P‘F×Ùl´…̹d 5Ž‰‡My"t˜§e~ÇÎ3žò×´3«‰ˆ¨ÆJ«ÙâÊ£üZ…äy/òÍWÙ´0Œ«ìÖJ9vLœ]·L¶…S+é|LœMˆr8• çbHEæ`E“>†¥3Í3Åna•@¤áçgîѪœÃ Öˆ Jy[k…±Æà¦ó¹òÊ+÷$eEÎBg@‚œAÄÉ âì~„-LNm&aœ[ 8Žƒ¥¹œ»Œó)å­‘1Ãøíž„/´׌Ú›1ÂRv“\Ïê{L´Wƒ}àÍ’6­msÀÅa|'5ý€åþÞ æ0EŽ ŠÕCÆ?åcb/μ¨à˜Kg¶®Ê‹7\ã “g/NLzôâ3‰Œ±»”YÆÒ®6üº¶ ëDî!²¹ùS&‹iâÏÖlÉc]Þš³öËV–{Ù¦ 3°q“׌Iîs¿ç295AfãÆT°dŸ=²KÏpþ¤‹ÓXÅ=qtеUØD Ò¿xø99Ÿs¾‚Ș㥃ö¬1.znˆ´ïñl!ý‰ï¬r Œ‰¾Ëá×cý§YÆEiÒÿxNÆEùÖ÷…Œmò¼‘ã™´ ÛxˆG>Ü›-·ÆS}Ó¶ÏW*ÖÇA"ÐA×VaSÈ|ŽM˜·"»Œ‡ã¼ Ý˶ÌùÀÜÓB6Ä/ËJ\öæZ„Pó-öë¾hl^XÏ€­g~h>À¦#œÁUš´‹u|ÑmâT D [Gâ“5oy³ÏÐ$^8ˆ2É&éæQöÈféX¢Ý†,QTŸl€Ø‚ä:ÛË`ˆmÄI‹c‡þgcy?ˆt€ãd"tHuC‘ ¼ëž‹ø“UÒ"ýY"Ü•Å®ƒp»J:nH›²¥ËdÃØa†Iïm÷ÔgX"™nH1Š8¡Ùø#Nâ>ÂýÄíXŸ0 DÚ{¶&Ã3nšÆÄQHÒègô-Æ%«,8ÙÆŽµ8³íšxå^~œþGÚô9£@dܵÚ"È/ ‘Zúã×ç£4a}Ü‘±M²Qã!›5ß±úM#V²3v•3‚Ø7›±¢ÇÜ)>»âSYM$ÜøWV´ Ç‹/Ù ðÍÁ²Úh|‰ãÅ^^ÒH§g´œ_ëåÞùˆÔÁ©ˆ¼µa™¡ÎÐ÷&:Qù¼­e¤ÞÐúßÄyˆÒéU Úm0^ Šƒå-ˆjÒl®¾5<†ìÿuZXG Òáèd ^?N¢•lÈ÷V6ePÓ9è$¤(©Ó;MF÷$™ˆk Ùž»ärl­/ç·Q±Ë0i7¦É‹ |ÊI2¦Ë‹}DV7ú“Ÿl3ÐÇì;Ø;Tö±MN‰ïb+Ç%º3&ø86ù¥sß›8j?¶=rÔØ™t¡¿#Öȇù͹ˆ8wlHuØ®(öÆý†Ú³±M­Äq>Ñ|»!Eû6‡ÎËxdÕ)?h.N/qƘç é6vçŒÎDÄ€¬êÐyÞ˜¨¿ñy­·ú6}’4å|¢ìø˜ DVòY¹˜óÁÊi*.Î Ø{/–ŸEÇËs!s¶£°?/å¼tqñNZÁýùá³)/åòâ,/þÍ¥½e+¤v¦°Ms«ìæ0†æWÚlóö8ÌWwA Ò^ôCêGùÊ—~J™$ß'R "âè¬RAg}\vŸ·mœh—É—I‘·³ Å„Éý„#È*ËÝg)?cŸå²£ ®†¨qrÇ_ÿZ‡ DÂò惭Œ"hP„› l¶ùÛ€k5[ËŠ£l±JÀ=ÇüzŒƒ­çKg}g®g –ìZœl;«,íõ3Ãë¢ØEl2qˆ£Is¤ãˆ¼Ò¦¼ùWj·ÆŽµ~Ä3ŒÚ¤AÛuíÙõœ…`’ëš±E×æÓÆã0k§Æ%í7“i“\핈ÃéÏr}²ÉT¶¾™$C^~X¢Ï)7Y2)ñ@Æ2“ii fë³¼±•Ng5ŒÈ„Éý‹S°Ï`™$¶iÈÜÃÛHㄉl~4#už ¯4š\fµ»ð·9ŽkÒš1È3^¤°ßöês(öÆ^‰Dy ˆÐ£­d‹±‰=ŠÛuãÃe3>¹OYu°N{,ŠÓý´~dÞ¯èÎ…@¤½š·æ×›}öm,“Æl)375×ÌaÏ9sÌ÷Ú{^¦CDc}~(G)è8ÄÀ9—¾FßÀ6^óÅÜgì2¦F”æ{ ËK/IÜ#-Dsbc渂HXÒè>ý¡UöÒ Ï\×I-¶s¶ÑÖÑK|6{?.Hxê «¦ÕÆ»´c\l’HÄî2ÖzNÚ´6ª-GÍ«\'ÏÝÏÎüÏVsè45/øؤUrì9+ã”E¶” #6J´’îñ—tê#u1öMáìÄ D*Wç£Â£4&B*\…£2qVYî7igœ|ñ˜hûž *F…›ôGÙ+v šL$"ùô6"oR½íäŽ?I½ÿAR§G”1Î.¨ÊìŽÓ'öÇÙäüÌØ«·#9}&à6×,ÿõü¸½À›ß±i`¯Ä ßq¤…°™s–¼%Ö¹d¹nQì"Ø«<ƒ’~6ƒ§@¤Ð~´×q"nP7¹ä@gä8û_[Ó–µÓ\ó¼ÉÑ/Úw~QŒSÚ®ï…!ÞñgÂMZôUÚ,±Æ¶ o=Å!mD%i28gŸˆ¤ÏÈøf¢áÇ^¤7+–Äoœ#†EdaÂocB½î$í´C°½L–”ý¸õî8"`ûÆ4ufìȶ “ÇêSÝfŒqOêXzÙ¤0²µÞ˜eNÃnŒ‹ì‹Pi á°—Ld³* ˆ–l6ö”6Â.ÁxÃórÍä-šO±ym¼(ŠßmG0ýÁØç—@¤Mj¯Úäü%y©ö]æÌúˆÌ-Í5µwýLÆ*í}>ÿå;ékV?„ãÓºýOVxŒÈ^ûu@`7ÙÉ’ëqË@ñgŠæu¢Œó†mn«°‰@4Âö°:•†uÓq® ¬²õ¢ØWÄi6G·}„ñŽ<¿¶ ëD›"cÑ||9 î?Ìy_/ǾëB½­tG‡:4'[&†‚ïØû²ó¸Vaè(`g«ü©uìɽêyÆqhÛ„>'+ûç}ŠkÊ>«ÀÖÁ‰ˆŠÓFk¿­ÉñqDEQ¬‘hY_ŸÛDu݃“S Z+%þäOþdñVê¤8ô§ú§ q¨çÅ,‡qeÙx@ˆc—¸®sT¨(ŠÓ$Çi™¤¿Ñ§÷Öu˜C *Šmá {7ç`ïÄ—UÂÊÇ%¿ƒùŒúš×GÄ_õµ®¿]¨8U¨@T'çZ ò6ÅêÅ“¸ Y9Ù&`fÇÚãG¢¢(¶ DÅ>¡ÑÉA¢âT¡QQœ\œk¨(ÖE¢¢(¶ DÅ>¡ÑÉA¢âT¡QQœ\T *vˆŠ¢Ø6*û„ D'ˆŠS… DEqrQ¨ØUT *ŠbÛ¨@Tì*T *N*ÅÉE¢bWQ¨(Šm£Q±O¨@trP¨8U¨@T'ˆŠ]E¢¢(¶ DÅ>¡ÑÉA¢âT¡QQœ\T *vˆŠ¢Ø6*û„ D'ˆŠS… DEqrQ¨ØUT *ŠbÛ¨@Tì*T *N*ÅÉE¢bWQ¨(Šm£Q±O¨@trp¬Ñí·ß>ÿº(Ž¢;î¸cþuQ'ˆŠ]E¢¢(¶ ÑÍ7ß<ÿº(N%D„‡b÷a‘ϱ DTñÿ¿½ó€³ªÈò¿Î¬³³³ûŸÙ™¸;†s "95t% ˆ æH¢€‚@DT’$É©É9çÜš Mìîó¯o½®æöåu|¯û”ŸŸô»·ò­[çԯΩ‹5‡Bq£°k×.+lý× Eèƒ8 „ÿºBñsãСCv1ç¿®P(W ää³ÿºB‘7ÀF ÿº"ôpüøñëCíÝ»WvìØ¡PÜPÄÄÄ\vM¡P„xõV„"tl*Šœ†Î+Š¼ïḜì„+&ˆ¶oß~Ya Åõ„N< EøB•E¨BǦB¡Èiè¼¢ÈKб>€Ã¹nQ||¼=F¡¸Qˆµnfþë …"ô̈‹‹»ìºBñsƒ±©:B¡ÈI ït^Qä°FÛ¹sçeסæ¥ëFqÑÅ‹åÂ… Å äЩS§tÜ)aÞÙÇËþýûõýU„œAÄê:6 EN€¹ÄA¤óŠ"·ƒ1Ι6œm£ã=´Áóù®ñU)…âFÀD:îŠðï¬#ˆôýU„Ž Ò±©P(rÌ%Ž ÒyE‘ÛÁw‘Ž÷ÐÏG "E®‚D ExB "E¨B "…B‘ÓP‚H‘— Qø@ "E®ƒD ExB "E¨B "…B‘ÓP‚H‘— Qø@ "E®Ãõ$ˆ.¹¦P(rJ)BJ)Šœ†DŠ¼%ˆÂJ)r®• ò’@äáÀïÄ3g ÎʹsÜ»<­B¡¸z(A¤U(A¤P(rJ)ò” (A¤Èu¸Z‚ÂXÈü>sö¬œJL”c'OËé³çåȉS¿÷°ìÚwTâö“CG%))I’-’¯¸<…B‘J)BJ)Šœ†DŠ¼%ˆÂJ)r®† ºxá¼?yF¶ï<&[bËÉÄDI8z\ö:"Ûv%HLÜ>Ù¼f½¬Ÿ-[Í‘uKÉƵ[dãö2qö|Y³i³Í'ébÒey+ŠìA "E¨B "…B‘ÓP‚H‘— Qø@ "E®Ã•DX AíØs\f¬Ø/b˦˜Y¿=Nböî“ûwËúYÓdãØA2§O™Ó£¬ìÿž¬ÒU–}û±ŒèÑIºth.Z6–Oû})$%%å²r EÖP‚HªP‚H¡Pä4” Rä%(A>øÙ ¢ä‹ÄüÏ"ébúx,ÞÝ=sÓþöÞOIº”ö¢¯ ~»{)&­¿Ü`Ȭ.þü²sX1y¸4äí¿ïâd·~ÅMWNRðrÒ·ãR\ïu/‚•êÈ.AÄx:{öœlŒ;&?Ή‘/~\"Ã&Ï–QSæJ‹nŸÊ£‡KÜ´dEŸ¶²¨G ×ñ-™ø~CYúE+YÙ·½ÌúôÿaÑá ù¼ÍëR÷µçåíŽÍd{ìv%‰Š«@¨DNùeЕ"3Y md‡ ²r3%Éþ›‘Lu÷üiƒ¸~]äF¹oÛ‘|Ì;#˜t­uÎnzÊV·PºØÖŸ¼:svÇXFÈhüfÖ&Úìâeõ ³£çd‡ JNN¾ oàw5 ŸÌòÊê~f ºðÅ‹—ÚêÚ’Q»³ {ÔC> VîÕÀÕõJêé¯KNµÕáZžKvq=ÊP‚(|ð³Dš˜ÇeüÊ]2yÍÙwôTšðA81‹÷…ÛÊèe;e鎃rîÜ9{ŲiÏQ»b§ÌX¿WŽœHL#G¸wìd¢ÌÚ¸OÆ,ß)ëw‘¤Ôtþ:xëwð„LXµK&­Þ#{ŸL« ‚ìÄé32;5¿Õq Vg¦œ@¾ì9rÒæEž±_&Xù½mï1™·y¿™È‚÷‘C²‰»+á¤Ì4må|ïõƒÇNÉ”5»e¬©ÛfÓ'´?X[¹~ÔôK´iý¶q×{²;-ó·°ýí°xûAYXÏœ š_¨";‘kÏòM 2,:NÚ5R^ïøÔkÞZ^lÑZš¾ÛJæé.‹{´”‘­_–M^ïÞyAƵm “Ú¿(S:¼*³?l*ó?n):6’žoÖ‘ÖÍ^”*õ*Ií7È–˜íF¹V’H¡¸„"A„Ì9xô´ü°0VŽŸÎtÞÏ È‘¾ÑÛŒ<‹·ó¶ÿ¾"´‘AääIŸ[¥^ÿÅÒîÇ5²ëÐ ùnA¬ôŸ½]úÍ `àÜ2ÝÈð“‰g3•ùä‡Üýqi¼l12=«ÅvN½dð¼y®ï"i:t•2cÞK6P—­{Ê÷¦mÇNI{\‡/Š3:Òá+®3éÏoìò]ôYlv¡ž6ýè¿  í³7î—©k÷\vïz!p~â9±$NºMÜ`ôÙãWü ¼yí7úøÂm—ë„«ŒþËsò§Açݱÿ˜üdôyÚ}ÀêóÁËgŒñ~LZ½[&ì4ûõ㼂Ì"~'&&Ê’%KdÆŒ2}út‰ŽŽ–U«VÙE6‹wš+ÊêÕ«eíÚµAóq÷×­[ô~f n¤›6mš] zÉ×® ȼyóääÉ“Wœ¿ùR×?€¿é3æêk%‰œN2wî\9vìX–õäþéÓ§eÑ¢Eéê4kÖ,Ù´i“œ9Ãœ™yY616nÜxÍíËä»råÊ/ƒ¶+Aàùü,QJòEi”Ÿ‡ß*EºFK¾÷§ËÓÌ”¥F¡<ŸL<#o Y.÷¶,%»Ï–{ÚL–ÃWY‹IßYÛäþvS¤øG³äáŽÓ¤ì§s,Ù’œ Ñs\*ôœ'½7MŠ™ûäñù´-î¨)*ÓÖí‘GL>M tž!O,1uIIN’ýGNIµ/ØûÔõî6“ä£IÅ’DAò£~+bÉ“]fH~Ó®‚Î4uœ*“ ¤Ýv—'%IŽåªfŸò\¿Åö·?Ÿ´ú¥\´ínôý )×cnÚ´ wÖï:"%?š-ù;M—B¦n˜>ùvþŽ´§¼¡(ìØ\Ê}6'­_îk;E¾œ¹ÕôY’mÿƒ¦Úúò,hÿ?[M²ýqèØé«V4~d‡ ›bŽÉع{¤ëwóå¹æ­¤V£×¤Æ› ¥n£†2¬WG™óqùüåšòNå(iýL¤ô©_IF7~V–ÐXâúw‘£S‡Ë‰ÅÓäà܉2¾{GéøæóòÂÏJáÊ¥¥A³×äÀÁ9ξ+¹NËAÄ|Î"ÛOظ9ÙÞÏ ŒÒ¦»Ÿš–EÎêØÃro»É–TG^Ø´AÒ‘_°2Ù ØkäÈSf^]¸õ€ÝÌ }°:z󥬥„'oƒ¦Òþºú+²Ï"+‚¨FŸ…òìW‹ìÆЊ‡äF†6r9Êè#‘Fv—1zÊ?ZM”—-µòßÿ,ÓÈó/÷Ÿê<]†.ŽM§xãùáƃÿùf4ÖÒî{î1¶m;(Ǭ“ûŒ>ÁâÝ+ÿÑG†.Š•»Œï¹Gz6 j?EzNßb4À¤Kã6õ~fõ&>Q²û¬´ôÕ‹ »:}ÙÍ;o¾½ƒ^«‡ÿš¿LÇwß×ݧî­G®‘×/\ؽedô|¼yùëâ-ÃÝç2riœÜÑr¢¼:h™lÜ}$ˆÎ(mFó V=Vî’Çß›ž¦¢Ïþ³õ$yüýiV?÷ÆOI¾ ãVì”ͳ/ô¡Ñç>Š®¼vçáËÈpž›¢.Oð¸ÑI¹–¡”›Á»—AÄÂœ…ô‡~(;v”Ï?ÿ\zôè!;w–6mÚÈرc-ÑâMgß勬3Ò_ V.Ö=ƒ–ï¾ûî2‡ßX¼ôïßßÆq÷ƒå „!C†Èo¼!#FŒ°y¹{èÅ[¶l‘f͚ɻï¾+ÈvþV&zî‘ï7ß|#-Z´ž={Z|öÙgÒ®];ùàƒì:Ô«‡»üýùf”?iwìØ!­[·–øøø´¼2ªyó,yFï½÷ž}fÔ©[·nÒ¼ysùá‡ÌzîlºrüeúûA›¿úê+›—¿m™¥Ëì>õöö eôéÓG†š£ëÊW‚(<Àó¹áŠÀ>£,#hÚ^cw‚=mIš_.4“V’ŒZ/ÿ4ŠÔŒu{¬Â‚åÎí-'X« íûŽÉFPõ2JÄùóâã©.3¥ÙÐU¦Ôd»ë±{ี¶ºËÄß´ûèe ¡yÒ(¥?ž-/~³Ä*ìWé5_ªeøŸLÙ$uœ&ÛL¹‰¦®*äg­u‚4'i#?›+6ùŸµ»ÄG³ìß\k9b•1‚ôÇÉ+æ^0‚È>SŸ÷Ç®³$Ùß'Õú,H#ˆX¨4ºRŠv‹¶$Öé3çŒR·V tša­Š¼JBùÍ!Ë¥ˆ‰»óÐI›G÷I›,yFò v™ë»NÚg³q×Q)dúð“ɛŠYD´'fÏ ½[¾´YvýZª¿õªTyùe©XÿUiÖ²‘|óN=iùL9y«|¤´y&R¾oö’¬ùº»œ0DNÏ!çf“ ËgɹõKäÜÖårxÅÓóCy£a)ól„}¦¸téÑ% èίG8áz˜š*~ðÎfE1g³€aÞ†¼f.óî>»û§ÏÊ ³Èöº;Øy3õi™÷øíÒ²@GO$Ú8ÎbMÜaù×»Síù9“æpªUEZ¾çùR—c'ezÝ&˜ßÍÛ!ϹB['Úz<Ñ–E™þg gó}Zþçl½œ;4$‚«gúô©ý™ú·w1Çoòt¿Ùã1+‚ý¥ç´ÍVÿY›` –[„c} [›OéF·qcÓ>K3n°Ì<ËATäÙ2ÜÄgü0^¸‡N‘âuÿ"g!̘"®”k§mž–àL­¿gXâ`ÕäÆ‘ÈEÙ°ëˆÑѦ%ˆ†/Ž“Ú§'HGùèJ½Žä"ÚM¬¼Ý8vyyßeêvÑ\‹ülNZz›öB -H¹Èâ.fŒÑê0EŽœdœZ·ä€Ûõ¼ËÁ‚ÄIÿNÚÀ{á?컕â·.=W&u§ß¹‡ÎÈóà þqå¹vÓ÷̽uâ9ñÌȇzx-âé˜[hS0ë3ú«é°•ò¬ƒgþ=këËß~}ŽkŒ;Úê%Œˆ¿Û£®?,Oõs3¾½yÐæ#'N[¢§ñ÷+l>᥌Ný€%¦®éû+'týg¿ZhŸmªõÕ"»É‹~í'Ïr;˜K²"ˆ :&Nœhçâ`É2sæLyçw¬•Š³$bQîK^Xª8R¹áÍ›¿¹Æß3nNsy@:QaРAi÷ 7È+/éãáûï¿·$d i1iĽN:¥D.ÿ£GZ«"¯K–»Ïuú„6»kD½zõ²»ëÛ·o·$Ú”)SÒêÉ¿´ÛYyëïò¢ÿ\Ù\#¿˜˜iß¾½ìܹ3.LûÛí"Ú…åøÄ¥.DÛ¶mKËõ™þtyÓ>žýä. ÄÝ°aÃÒ={òp}ãÕ×½ùÒ6—ÆÕÐ'ô;×]_ôíÛW†ž.¯kù+AàùÜp‚´<æ0ʤ ¿n_ÌØjw°–yû‡•R¥÷|+èœ0*ýÉlùpÂFëR†µ$ÂØd ͇­²»tÄ­øùJÖF#Q쨃És¾Q±È¡,”‰þ³¶Y¢©|¹Òà›¥Á "vj¾»CúÌØ"µû.’Ê© G½þí2©ŽBZ·¡‹âì®Ê YÕ/ØH”rŸÍµ–W~I’ͦïÙÍdç%È)™ä ÉÖàë%¦–ËN(#3‚ˆ¶œN£¥FÃÆRùå$êùRî¹—äÍ×êK£ÈbòJd„4®%_¿ñ¼$Ì.·.• ›ŒÂ³z\X6Wί˜-çVÏ30ÿn˜'['”ÖM_•¨ºR©^„DÔ*%sÍå±\V?Ü„ì¿ÎÄì’éÛrå~Ñ~8¥ÀývBÑÁüé³ ×>¸Òz“žtôGFBSþà¹fF±èaAÑaÌZ)Öm–5ò‚ÅÖ))É2kÃ>KÆ3g—ê>ÛÆi6|•uwH2cè´YX:y“½GZ,*Ç­Üeç>”+cI~‹ì=8=¦¾Nȼ͜ÞÉä[öÓ¹vC W䃳&ú~a¬Ý(nÊ„ô‡Lr–B,ÀªyöÍÜív®Æº¤–Y Q€ègc‚¹½Ík¬…F«‘«åùþ‹-±ô®iï`ó/õÅÍ2ÂKP6¸ïH-e ))…ÿ+FF`ñš”º`„ÀèüÓút‹^EÖ`Œc6ºÚŒZcëß'‚(YV=±n¿ÅöZÑ®³¤±é'ÜÚx× tž.ÿ0zïãœMû¬…Ëúqëš÷xEL‚}¯]ÝíÑ»ØçDûÐßœE ïóÄ)Ó§Žß %LŸ0¼;f?}ƾ›X}gÊngúéó>ï6ÏŽ9Í?Þct4Ž `““>è>y£í;?,ˆ‘(£ï6Ïô£—¡R7È›wŒþ:`ö6ÛO…M?Ué½ P÷Tý•ù„¹  ›áþåˆkýcþå~Àè¬v˜bëqÐÄa3òá«í˜¡=ÍL9>ôú)u<‡”1ºõ´µ{¬Åü¥ñžJ<¦mÌw¯š1æ'׸‡µ=VþÌŸN‡8g»3!âÝ\CÿÇì?n­Ï&­ÙÐE (›ñ´/ÏXË `.ÉA4a„4ÑéS´÷ Ðy±(ÂR¥K—.ÖŠfÍš5váµ îBN¿£¼/¿üÒZÄ@bL:Õ^‡”€ˆêÞ½»Í ,`ø—´¸5aÁÄ=â@N9¢ÉÄR×®]åý÷ß·.Wî<"ÚJX¨ð/ºÝ²eËÒò'm¦N¤A/€¡}XTaÄ\Lp‘Ó- ô$— Â¥üé³Þ½{Û>q u‚´úè£lÙäBž^‚ˆ€Ký‚›__vy9‚Èé«ú½I“&ÖõŽ¾Ä:‡:Q&y²v0`€mßÇ,ãǺ–ñDäé3nÜ8ûLHK8k'°{÷n;V¸GÛFe×æîÞ×_mû„~ço×n%ˆò6x>7œ B¸ ä1×GÁ€ªj„b­¾ "dS£4;B‡ë%çåK­°Æz'°Û–d'„js#É Ÿ}ÎŒp Ù#¨î1‰ëäã”À}Îä±÷÷,‚ˆ³dÛAK!ì­2cêŒ9õP£D|:Gšü°ÒNŒ(9éò3u_ádò[°í€Í‹‚ìÀ•á‡E±2×,ìXì1/B^/ÿÝh¬YðÌ‘±+vYk \˜…› ßš¹ Î…Y°õ€]ôàN iÏB×v³9C£çôÍò·fãí‚Ùq·Y±1‹D¬Xß·ÎÎÍ#Ì‚™:‘/ò„ÝqƒÈ9y, °ÈdŠü¢ÌçÌv憽öœWýÌR±Ç<»€b3Ÿzü_ó RÀÔÝJ³°-gÅ,ê&¬Þe˸CamÁ"³nËÉYˆµ 2‹#êñ¹É33™£¸WCÝkäè(ó|÷9)ñOX7øw†¯²ÄËîÃ'åÔ™sRý‹6ÝÌ û¬®‚ëü3NϘEcèµAˬ•ç55ãÅ7$òB ÂŒùñ«ç¿0¾Úþ¸Öê'5¾\`7Û8cb‡Í­ÏŽ¡ ÑÁØÀB¢"Ý}€ú_Adòˆê9WþÒô'yß¼'œ½ÔÜ´ýÎVeö¦}ö½ÀâBƒ1€@"ÐçJå]ð|n8A(8 ,ž˜âBüpѲ0ÂIv§ Òà þ}ãÛeVY K‹ÂUÉÄEqß¾?p(Ÿ3¿']£tc%ƒ«‚ ‚dL쯉O°;,Ó’ŽRæv|É›«{ŒG˜9%‚Ã÷Ü(ïìÖ±¨äcó3ù²;Ãß(èG:»kràxà"vM‚è¢ÝÌ÷þ4«°ˆò÷Q|&¥E¿s'pä¿!“îhpÓsäñ1ÿeaƒ"‡ç¯G8 #‚ˆß\oØ¡½”©û²D6x]"ž¯'eëÖ’ruž“25êI¥/ÈëÍHÃjå¥bÁBR¿X!YúQS¹°|¢[>UN*G g‘33ÇHâœåøÜ1rdÞX‰™-«ÖÑ—XÍd°EP²ÛÃdÌ5'Иü0Þö(:tê¿ío·î:;<ì!´Ù…B @øíÒCÈ T©—¿Oå»îî±ØGQAøQ;N÷œ‚ÁÕ‹Ý>iP|üýáMã×2«Ÿ"´ÀsÊŒ b!„Cȉ£fþm;z­]äà.6fy¼]ˆA¢ˆµªL¶ f\‚!¼YdaeÀ"÷„Y˜a™‰e'ñBÐp>2ÁîlC(™Å'äϺø#–Lg¡ÄâŸ|qiÆê€9¶¬Yœ·2‹HŸ¸‚¬ŒIHµa—œ€5ÖJˆrÉÿëÙÛ%Á·þ>Ñyº´1‹uæg¬J9«r‰€Õ'‹ì9¶Íb­XÄb áÚø“Y A`Aö³Øg‘~âôYk}‚uò5l0xÄÈÙ-{òÞNýµ‚ñx%n‰ŒžD–#|û'KñÜp3#zã†ñ…º V)X3CôÝgâM]³Ç’0ŒžS7Ë£f¼bÅ2d^ŒÕ…³ûŽœ´g&¾2p™•õèG•zͷ粇õ6Öϸ=Ôaª|· ÆæGÙ#ÇÚ…:ÖÇè1×JqաÌ;E¬š!3!Ìh;XàJ²:úZ_ó~àn9°þ»` ¬)›bu)ή±îv&,±y±Ä£/éêÿ%ÖHýŠ¿Ñ‘Šu‹¶”X$úâ7ÞºêßoÚ5Ì´Ïõ –?#äW½O€€±n¤¦m^µÇO5ºì ;ßpfÓa«ìYSÌKX9ABÙyÅ€9kC欸8§ªãØÙL¾ˆˆ~ù|ÿ%öÆâ‰çÇÞè–þñIº#W[‚˜|°tgž°–ñ©ó„óÝâ­d€éã¿7Ÿ:—&µ(䮇%Ìxl;ê’>ž®\s òŽ±È˜r›¦XÍ3¾ìÇPRuL®/Ú¸Î|Äuæ!?Æ‚µfRFnsÉÕDü YqÁ‚kø\gó ’¿]<ä*yB @p ‡ÐËX«Aj`‰â‚+’½ ² ‰¿Ñãbcc-ØÌs¤ V(>”nNý v (¨3„ó*m@÷DEçC¿… ‚Ô Ä Ä‰ ܧ®Ä¥HôW€•„ ÇÅÅÙþ‚x¡­ÎˆkX¡ßRoW– `\s–=èÚ”G¾< G„QwÚìÚNŸáªåÎúä“O,è[ôyÈÒFGGÛß®LH9ÊpVQ4®¹çæ}öDKô?m‡tµwÛ¶mí‚ü©3¤” ‹/¶„"cƒ2Ü4m¡ŸÈ› QÞÏçg#ˆ (¶OvžaÝŸ€V‘0`÷´É+Ò„ÿb6ÍN0Åš/‚@*Õþj‘UÜŽ J¿ÙUcƒx”Gìæå3‚¬`—öß^Ó6[AÆnŽ×‚áz—ü좥°ÀL­3Š ;U,8ÎÝnâ0JÊÔ´ü8[…ŸûFJ^Ö‚È((T^77þ½Z‚ˆz À@X½ðõâ4ÿþA¦­}v¦IãW—ŠfÍ> ­ò1Ä(>ÖüùÂ¥ûhX|±"\uF¤Ç 3Ù¯Z]ž,_]ò•«iþ­!%«×”²µž“ÈšÏK×ëÉÛï¼)Íë<'U‹”Ï^¨& czÊŽ¡=dj›×d@ÊòeÕÒ2­aMÙß÷=ÙðQcX¯¬ z©²Lz÷eéðZ-)UXŠV,,%*—ÇJå—^>·ïƒ¿žLöì4°£âü¾çÏŸo…Bœ8´Åù‘ßíHÐV$“;“8ï›ó#v“z€$LW6‚ àꇹ©7=qùò‹v']ZW&uv;)ÁÚùÅnWð!´Ø}¤˜8þú¹:º ŠêµbÅŠt¾ìüKzâ8“YwëÙõ—W„²"ˆ˜÷ö$œ´‹){NœYlb!„û 1¬qXˆ372·Ë Èy¾&Ébwáª_ÒrV,B!qÞ3{ ºä´3ˆ eÜ UÜ_Xò¡ˆúˆOfKi³˜*c‘l:|3g»]äaÙ‰Lqg ±h{qà©Ö{¾Tì9×.Š;ŒYgó…l²§.&Ylü`Fà<ó»ûäMvÁÈÂØÉ$ô,ºpÕ -ìæÏ5e`!‚ÕÑëß.·u8o‡µLÀúÖµQ‘=\)Aä,ˆXˆóÑ çXÎ0n#Ìø8käxßè­ò÷w&XÒÁt OH ¬A°êåÀk\ÁˆK<Æ.oè:,Â} +lιµíÁË­N™À¦ÏæX" ë%ÜÙì"ò#_Æc‡…@(eDAš¢Ï;¤Œ3„ëŸÎ¶ã—qkÓšÛYkÉOÆ,g¡#1dŠ;ƒ(Ylá=YjÞ“æ=™gÉ.ã7Ø{?­Üi ]ÈÞ\ËÐkÜ;H{¨c—ŸÖÛº¸úóþ®ÛyÄÞÃkÞæ¶y!NØ܃˜s}‚ŽE? ÏáÖå§ ¶pñÂEk ,—°ðÃ=Õ.D¯ Z˜Ì}ÜF½ý€¥RóA|a‘dßy3@8QæñÓ‰–¨c¾¢ ¬½ •Ð1ýÖ>¤ÅÖY-Q¤û ³ž‚TÃú’ùK*7xóJËÓÔ%Úè¬ô{À]ïr8èñô½+èàÑ©D[S¿¬K\Ò3¶ á¼DŒ>àtf¹Ì%WCXäCB ûRÒ  §S¢;ÂBÒ«u‚#ˆ¸ ƒnæ6)…ûüeu|À ÝÝ + ô<¬Œ¾ýö[›7Ä dú!ñ± B¶gëÖ­ieÓvêG|H Èt`H!êD<ïWɨ#©#iÐûÝB貟~ú©= šy›ôX®»¶°ŒB†@£Üf«·ß!­È‹¾ ÃJ ¢ÈMüMÐvˆ¬ðyô=›šäM»–.]jûÝ–ÀÆ,ù—€Eõuz4ÿBêP.›÷Ù{ "ˆ,"È)·‰‹ÞK»q9œ4i’íKï×Ó\ x0è?ʧ ô'A ¢¼ ·þºá‚sjve±ba·¡æˆœ·†,·æ²V¡F¸e ð£×Z4Lu˜ß-Œ±Ê™%C ’>æR%;ªvçêb@AXsîîV› ø½Ö(ÿi;-)3ƒF- |‘a‹‰Ëùn×™Åù±3¹“h^Dʳù™¸î`B”$»3–š$æ¼Y„'}@¯– ¢ÿ°JB³ðA!Á¢e)§–s®lÖª*)p r7_zsýâ-—7,‘ìgxƒ<·p@F‘82høòhÉH)P¶šÚÍŽ  …ˆ>¥ä‰åW=/9¦=0Î3"ˆ˜ç!6pï`Á‹+מ#§¬¤‹K¬-˜{!IìùtfNc1Åb0zÃ^붌›Åäµ{$ÎÌÓöÜê罱ë¥By¶,\šy›…ÑÃïNKGARñJ\·ØäÀMcãž#Öz‰ÅÏD3/Çšyk ,8Îœ=o­2±’Å’„….ñ¨k> .Y™Åsà ãAÄb—2k)aæy¬?ù²Ä“]Œ™yë(d)›ä;6V œ{DzÏØ* ̼ÿªY°Z# y£¸ŒÇ+%ˆ jì¢ Žä€«8&X{a51lq¬Ñ3¦Ú8Œ¬‘!E¦®ÝmöX„qÿ>3N±8‚0 .nðSÖì¶ç¶@p@T,5c‹tœ5óâ7K­Å : ç­ßuØêœQÈxcÌP¿a‹b­5Ÿ/Ç’‰q‹^ĘJ™÷ Wù»ZO–;ÒÆ: q!hï dÏ„ŽÕíÌ¿MLÝ°pbCðÞ¶S,qJ”Áy<¸…õ3º!›Ywªõ¶nœ9 ë,mADÛ(÷/\ÖpWâ,úˆþÁZÊ«ËPÜ× êø<ç:ÑwÔþç=b#})Ð'íyQœ7 e¿°–Úfo>v21í“íÔÏ}øk,ˆ¸ÿD§TK¦´~H’×/·ñ!³9‰ ;G_ ˆh3e0Ç ¯âªV¤ëLKZa!è%vÒDÒ¢‘gÈ âkxÎÁŒ4χ1âÝHô‚4X>ñ™_½ñÜ<° bLC:ºqn7p;N³ q¯`ûþcökhëÊ€~Ì;idçW™—À\’‚ÈRí6Å ЛЕøKHˆÈtôCt%Œ†T@ï‚ ×A„ÜuV'®€kš³bgþ£>èmè§lbrÏS ©AùôJÒÐqÑ!-ÐÇh:,º!m‡¼Aï={¶ÍRÝ«((H%®Ñvô@Ö—ÔÕ¹Éy©H5¬¨ˆ‹~Hô3v£#:·1Êuîrô¤e8‚Ë#ˆ+çZ‡emvmwDt@êå@»I‡.Ïóâ9ÜùGn#˜€{}‰ãÒsÏKQ&–NHäË3cNŒŽÎóFGwG:›¾‡¨¡ô„º3}©D Àó¹á‘µ¨1“i{ ã…Àa‰zü±"ì~`É2/Õู!©Ë3½X+&êËÉʧîôP7nÌaq?@¸yósŠT¥Ïç›EÁ<ûu„¸=@Ò,LÈÃ)+äuµ‘5ý=0ý}cð2Ûoó‹ãåö¬B‰1Î(—]ê…{ÖVÔ aïí'øÙ1³¦ÐYÔ'”‘A”œœ$í>ì.•,'—®,•®*OFV“ÒÕkJéjµå™u¤Ö›5eÚ[oÈÞvoHÿ«É¶ŸÉÙ•ãdïè^²ê“2£õ;2¿mk‰ëõ¾$/˜! óçÉœme^禲cT/éÞöM©];Rê5¨"uëU’ ÕŠK¥:åeï¾ôŸÚL_¯KD°ýXÖ01;a@þ…8AHá/Œàag‚ ¿îXÚ T<>v PP¸‡ À§sT̃Q<àüMŸñ;::Ú–G9( OÈ € Cp¡´¿ÛñAyû6cD=ÔsƒÛbÁEý¾Î7áŠÒà„¨s)CxrŸëVÊ…HÂêÊ™@#ÜQ.PŠ†ô/ýJ;½;1ŠÐÏ8#‚rž…+Vˆ7ó× ß,±Ÿg®ç,?7'OÚhç,¡ÄY,¸½Äý¾qw€øÀ²ÇY~ý`žãYÜØŒ%ôˆ"\Ÿ žè4C>¶®ò=ã½q륆ÑsÈŸ·%M¹ä‰Ë£aQƒ›b`<¥XK9M¬_8Ç ‹"·É†>†•R‚y2"ˆø›w"–2 WÑ=‡¼OX4B%ßû3¬k#Ĥ{´ƒMAˆ{®$&_H…ß7'ƒæÇØ÷ '5æ=Y¾ã%_8ψßD¼“;ž°}Ò}âFûÎÑ~Ú éÅoëê9|ðbMLJR8c‚‚kœÆÙ¼óv3Ðô n¨œï¹’F™üyÛZcϳ_€3ÝùµÕkä%Ó6êÖмŸX†áBG?Ø–˜giD?ðœýn¡œ­™Ö’-%°èå}Ƴ Jk“‡ ¢Lú…¾f^a®ƒÚŠ´<¿>fþˆÊ„ rz9ç¹9ëÆ´{æYAðáî‹>ɼƒ%gS3víWçÌØ¥ÏÐ9ɇ£F-5‹~3–øÚ..wµM¿¡‡,1Ù˜DÍÈš)·‚¹$+‚r"ÆY ¡ã²ÁÇWÀÐiл Z m&l¢!O è…ü†4rÇ 8‚ˆ¼(’ ½yŽ²Ñû °´aÃ+Ê# ‹¡÷ 0 ”q -ÁDè\,¤G‡e³ŽÀFõ†PAoC·C¯#>º :SÚŒ~KP.qБé~Cj@z‘qùâÝMQ·©‰J=¨ã–-[,)„nH|¬ª¨}GžKè¸Ä'/çÊÇs!Ow$„î ",¦‚é›~‚ˆ8èÑô5D™{îôÏý—çHŸ@¾‘}½~ ?ä?{öl›¤ä}±åò¥_HO0† —´•z+A¤<ŸN!hØ á0Bv,ƒ :0Æ,š8( :0d÷õþv“÷L^̈?ÊBŦý0Úîñu””®ÒF@q°CƒïvPYs A‡àäKi(˜ËÎÙÈIîIÖꋤ‚]fÚº<`”¾¾YøHŽâ€²†B^´Û,+`ÙÁ*É+dù7ºZ©ç*ùóIWG÷ãI›l^î "¬«&%ê‘ŽSmû Ýø¢Í…‹˜Q–¿6o•±½‡¾õ(uÞ~¡­Ì¢‡v²CÆ‚ˆjo=à DbûR“wäÑòXé ’/¢’*ÿŒ”«YSÊ×®%Õ_©.•_¬"ÃÚ·’¤.-ec²\O9¿n¢œß0C—L”c“†Ë±Ÿ¾—óã¾™3G.nØ"G'Á=õ{9½|‚ôíÚVj=[^ž­]^jÔŒ”¨ÊŤxTÙØð×pÉ® Àn “·KC[w„7AЄ b§raŠ²À"†û;ò p ìÎBÈ@"9_tBA±fú‹ã¾ì@ÀªÈ™œŸ³÷<$oûP.Ø!¡mÔ™º’—s]suä=1‚€Eð{ "êA<„ ÐÚO_xwF èv¡ü‡}+B ™D,¶ T*÷žoçCÉ·µ˜`Ïa1ˆ»ðÓÌ°÷È\bøBO× lZθã>ç¡° cÑÌB òˆE(_FÃýW,08ƒÈ-|AÄ"’3IXŒaÍÉHÃ<Œ,á0bΈc1ÎÜj-9 °ŒàÚSä _&ƒ˜gîåì"võY:‚¨‡iW‰f§s)ë3s‹‘C“íB ƒ¼âLä›s9»µùĀ˜ù}øD¢ÿ9@×Zä±EXNàJ "Ðÿh9)àr”*O-¹bÆ!›AiÏ£b!oŸe×YÖUkfÎÇy²ótëD~X!¿Kš±Â¸æo®¡+aQ‡kV×,´qäêº_/2}yÁZlã"‡{$Ä!_ƒrˆ¶í"Ü 9xžzq^d/ãƒù‚ö41p8¹·M¤mdôÉÚ}¦µ+vthÆK ï¦Xò‚­‡§¥?âÌ—à‘«·µ˜(ì×/é¨üÝbÄj{6šŽøÆ2’¸0K±(JÝàý¯7F›1gÇ2®°èén,Ó&Hº`úynGv "È ô7t@ˆô Ò,ÆÕK¬±!Ü™7èPÄA?r:  ÇÆÆÚMBHÈÈÊ…8A?$ú yQ`–ã€@|êᬑ ÐÙÐÑ îl$tXtHòeã½zàà e£C‡tîüÈ0È$úÆÁA„%:mC7Eß…¢Nä!ÉÃ}ú’<Ñm)ƒôXßp“~pçø ÏÒôy°gƳ¤o¼_Pó‚<Ñùy®ô«ÓmÑÍÑé{Ú¡Å+~~õÕW-IDÀ Pº>äÏŽ~#.768ÛåKò¦ ô+í¥Ÿ7®íέþX ö|¯ÔG ¢ðÏç†Dkj¬¦®Û+SŒ`“ n;` knœxF¦™ûß΋±î(È\Ç„86Φ5ŠÐ²˜CVØÎÙ¼_¦­¿”¯Ëû oÇÅ „=.b(c¸`ëüÖZ(øH|å %Ûëžåiw˜6_oó¤Îþ][ê9ƒsFõòÆÅ]›¾œ€«ÙËâ­"Á9‰ÀE³6í³÷ù›óüý‚«»ÞÄGY¡ïPœ²ªO(##‚ˆ‰»æ+oÊ#%ËJþˆ R \Ey:ª²”«U]"ëT•gTˆZå¥w›V"CÉÁïÉÁQÝäÜÜrqî÷raþ(ƒqrnÊ09?| œ;JNMœ §§Œó³FJâôAòåûÍ¥jÕ²R¡B)‰ˆ,&¥¢ I©ò…dÛöKd^‚ËÞ&i&p&gÚÂÎŽ ƒeBŽ BnРA–XBP ˆÝ×HÏdŒ å:‰¾àƒÝ!‚#ˆ †çÛMZŒH^gòLý)‹ºøÓCèQ?ê ñļ€àE¡x8ESb„0–Dü&_?ADžÔ‘\ñ ´Í•I:yrßï/l>R„x6™DÌK,šp«e.f'žÝ{\&p%a,n'0X^ ]+›÷æ^Ò²heQÊb 5vÈ‘#Ìéì|#{°¾v€Ü`ÁæÚ€ Û¾÷˜=ˆ–\œ™ŸY„s 4V X0×3'æoæõ´s¹(Ï\_¿ó°=L—6CyW´ƒ¾A¾89…ûÑ:³¨öÖE‘}0³Cõ°Ö_)v‰/[aaãä)ñÐ8÷†spKgyÄ~±Ž1“ªy¿%oIDAT{`õ±ÈèC»ò˜kX^3®:EàlGvØëfLqv@|mO¤369iÈ‚Xù΀1ëBÆ –u|á îQ–u[ÜÜ‚ÈÁêfœRnïé[lÞ´‰òÈŸw3ˆ nÉ—wòŹFºq_²åú—}7SŸ'ý }ðÍÜV/Ee3Šç°:þÒüÀx@Ÿå¬!Îzd/ÿ2çÐÏö7yÓ¾yæýuqh'_N䋴̹à'yIËûN}Ýx´cÆÌ5#–ú—9‹ö¥ŸW.ï@þè„|]yÉ«’ž¼f™9”öò>ÐŒ6yžô í&„ Ï=:0–/Úzñ\™Ã ™ü_GË+ ï2"ˆzVÐî0dï":/::$¤: ÜGG ŠÍ7oÙÎ=ŠßèM5Xd£?9·)tjÒ£ÇQ>t,ô3ÚpsåQW÷ÝÏ•ïd=z¦Óï b° qgNBb`qÎ}êF~”ëÊvº1×½ýaB¾ÎÂÇ•Ïæ)z3e ³’¯«?ñ¨ŸûJívz2z#}ê,ä¹™F<o»¹žK¿g¤Çp  W=ïy›üKÐ÷ôòƵ‘¶Ð6úŒòéž ÷ˆC~¤Á%Ž>纃7_ˆ0—/ù8 ~÷¼éw~s>Àª(X®ä¥QxÀýJ„ž;#' ©¿]«¨xî;ˆÐ –¡Cœt×=÷³"=¦.¾ŸÌ±õ R—ÌàõÑö’˜.NR@Xú¯ƒË/£ëiuK½îú‚ºgÞ/—Ç÷—NÈŒ ªóZcy¤Dy¼LÉ_ª‚Ž¬,ejT‘Jõ*JùÚe¥dõrRç…geA›–ßµ©ÙYŽè"§¾ï,Æ)¦ –ÓC{ËöNÍd|í 2å¹ ²»wGIŠþAöýBÞyµ¶”,SHJ•),%" I‰²OKD…¬ "çbæ¬sÏ9Ó_êAÄŽÁK!„þÆŠ†'ÀÙÉ$qçó8Áùã>ié"„?ï.yc^J]P^ —œîpŽÈAp#ˆœŸ¼3}vÂŒ@^aÄá™±ã®y’ÿBAN¹C°èôuDCˆQ7îÂk^yòe状fô??™D ˜\°›E9gŸìøl½à’‘øHÀ¥üÓσƒ«CÚÂ>5žwÞ÷ÏÝ—Õ+õZ°ùÝÊ°ÔxÌ·N¾Ø<’ÒËwÏŸùº<üò Úuézà÷åù(²ÆcVQ> ¥é°•²Í,¾ù‘Þá6=K7>Üoï¸ñ¦÷ê'öÜÆÔ±è%b.ÝO?Ò§ Ô—ƒÝ9›«¶Œ"[Ÿî‘ö]D%?ž%=§oµï¦‹ãïoîoWžÿž÷=Ik—§ýïwVd¨·ÎÕÉ{ßÕÃׯ^ª[óƒ§nüöö±‹ë~{ŸóC0‹Ÿ`eï˜ñÎ…Í+^¸þ ¶êòõ¶#£qlÛÃsò\˨^y Ì%™DÅ GøóñÆq¿AøÛ¯ÿ8ÝÉýF÷ò–ã½ï/Û.üyûËwe¹¶¸û€{îw°ûþ|.êEvúÈ_ß`ù»tÞøþ¶øáÚà¿îàÚî¿î­Ÿ7=e;ݘßþ¾õ?·ìæëOç}&þ2rä«Qx€çó³D ÅõBF]Öíå_ÅJËã¥ËK>ƒ‚å*JÙš•¤Fƒ(©P;BJW+'E+GHÕJQòi“ç%npÙß¿…êÛBÎ ë.gG÷’S?ö‘õ][ÉZåä»ç¢d{ŸNrvjYûýGRó™H)Rü )]ö))[¾°DTªVZâwÆ›ò/€ÉØû37Aã'ŽÙ,»„Œ",Xâ`ÊÎ$|‹Ýù;Îm‹û삸þ¼Á*ˆz8‚S^v`8³‡] Ä ¦¹WJñ7&ª9ìÀ°[Â5íÅ"Š ˆ"êïêæÔcW¢Œzð›Ý1,­ ˆØÙàf¼”EL“qŃÂlÖùËsŸ2œ¿¼ÿ(BŒÉ¬¢Œ€[ nVuû-²çÊù¢ ŵ€ñ˜AÄyW¸{¿6h™µlËÎFR¨E>g`ï6˞ÈÛ~0‚ 3Ðn,¾9_Ë\Øýq Eö"…"·@ ¢ðDŠ\‡Œ¢””déöyy¨H ÉWªœ<^2R EU¨:¤j½²Y­””ªRJŠV*#%*—•º ªÊÂM%®G#ÙÿUs96ø=‰ïÛFuxIæ4­-sÕ“i¯×‘‰«ÊÚž eÂ'ïHTTq)^òI)_þi)W¾)_ê¾XÓÿ€Y¨¿®‚÷' GQwÈ‹Ãò¢‚ȹxAÌà³ A‚ò |š9Ìš¿ dÓᶅ5 Äq°(‚,½Œü)‡²Ý!Õ,Ò!‹œo2éœ_6gÿà«ì¬}¨?D .cþ3~ø“eêD= p°n‚lÂ2ˆøä÷Ö[oY¿h¬œ\!ªÜ¹I´‡´¤ˆÂ²ˆùEÛpCsõÅ´—¾ãŒ$êKyÔ½cÇŽ–T£¾þg  \ A8ðë!dõßS(®™D¸4áJAé¿N€ÜâìÞ%ÿ½+ýðÈ™ô. …â” Rä%(A>P‚H‘ëAD?uš©‹óM†x¢<ê —\]üýN}iu€|‚‰²éòážó9\wõÀ’ˆ³‰(ÿjw6eá Ni;×IÃuç/Õi3ó—W„xn×B9·.ÿu…âZÁxÌŠ Âr†ñ—Ñýpkǵ¾KäŽVT Ås…DŠ¼ƸDážDŠ\…Œ¢ä¤$‰‰‹“¢ž‘ü%#$²jy®~%ƒ²R½~i©ò\)]©˜-_TJT*!O.&m+—¡¯WÅž—ÅmŸ“µ]ê˾^Ídc›Wdü³d~“:²»O™Ô¡®Ô,_JŠ•Ì/O}X"£—¨Š¥@‘‡eÄèáö}ð×Ó çVìºó;öûÓ>ç+ÌoçíÀuÇ—÷¾í£äKþÉ®N.¾ûíý7£ºø¬.®k£ÿ~°z+ÛÛî`urPײÐÏìZ"…âzñ˜A¤P(Wæ%ˆyŒq%ˆÂ<%ˆ¹ DàÂ…‹Òó‹÷äí–%¥ÓG¥íûRûÅ"U½”«^XŠE•Bå I±¨¢’ï©R§ð“òqµ"Ò±bAX¿”LmVE¶}òºÄôh( Ú½ ›>{[vôm*M*‘Ûn»M Ë/%Ê<,å+>&‘IÙÊÅeמ]JN(Ù€DŠP…D …"§¡‘"/A ¢ðDŠ\‡Ì¢¤¤‰[.Ÿö)#=¿,-÷(%Õ_((Å*<)¥+=%%Ê’QŤtd¤<^ „|ìii\ê)yµÐƒRé‘Û¥IÄc2øÍŠ² Ók²ºk#YÖå-é÷Z¹ëo”[~õky,ßcR¤ä£R.êy¢È=Ò­gWk㯇B¡¸J)BJ)Šœ†DŠ¼%ˆÂJ)r2#ˆVDÓft‘~_—/û—‘:/”R•ŠJ©¨ÒR¢lY)S¶Š<ýt9yà¡‚òøãÅ%ÿCÊ‹…’úOÜ)Åîú“¹ó¤ì}‘êß.­#‘E‘[~y‹üòßn‘ÿõˆ<]<¿”,õ”©X\âvÅ[÷& ÅåP‚HªP‚H¡Pä4” Rä%(A>P‚H‘ëAtñb’?²GÆl"ïv(/ϽP^J–®(ý«°üëá"’¿@i¹÷¾rϽù —ßþáåïý›TËw·Ô/p‡T¼ÿÏRüÎßIÍü·É§µ‹J©î”›nºY~yË-rÛ?ï’ü JÙ²…d̸í{à/_¡P‡DŠP…D …"§¡‘"/A ¢ðDŠ\‡¬"€Ûצu¤fzRà‰RrÿCOËm·?,÷Þÿ„µºãŸÿ’û( ÿû÷»åæ[~#·üÇïåÿý‰¸ÿÿ¤QÑ»¤eăÒù™'¥uäò·ßÿ·Ütó/ä׿ùÜzû}rëË'ŸõJw³B¡ÈJ)BJ)Šœ†DŠ¼%ˆÂJ)r²Cóß„‰ÓåáG‹Èm·?$7¸÷þr÷=ùäÖÛî—»ï{Lþü×Ûåæ_þ›Üô‹_Z+¡?üö?åù‚wË{•óÉ»•Ÿ”Èo“_ýêßíýÿøÏÿ'üËÒ¶}g9sæ¬~R]¡¸B(A¤U(A¤P(rJ)ò” (A¤ÈuÈ.AsæÌ—' –’ßþþV¹ãŸËíÿxHþò·;åÏû§ü×oÿ(¿€ ºé&¹éæ›-ôô=·É»•ž–·J>&·þþ7{ÿï’n÷°ù*9¤P\9” R„*” R(9 %ˆy J…” Rä:\ A11±òú›MäïwÜ/¿ûŸ[å¿~ûWùÝþ*ÿùÛ?ʯþý?­õÅÍ¿âÞ)Ýj–‘â÷ý#zª`A™:ušu]SrH¡¸:(A¤U(A¤P(rJ)ò” (A¤Èu¸R‚¤¤$KRÒEY°p±¼úúÛrÇÈ/~õrÓ¿ýZnþå¯å¦›p1 Aÿý¯Rûɇåÿô')R¤ˆ|óÍ7v²#øóU(Ù‡DŠP…D …"§¡‘"/A ¢ðDŠ\‡«!ˆñ XaQ4jôiÛ®½<[»®D”« E‹•’ò*KãÆoËÀdùŠröìY›F­†Šk‡DŠP…D …"§¡‘"/A ¢ðDŠ\‡«%ˆ¼à dÙ ×R†B¡H%ˆ¡ %ˆ ENC "E^‚Dá%ˆ¹9A)Š%ˆ¡ %ˆ ENC "E^‚Dá%ˆ¹J)á %ˆ¡ %ˆ ENC "E^‚Dá%ˆ¹J)á %ˆ¡ %ˆ ENC "E^‚DáƒëN>þ²ë ÅÏ › …"§gçÿu…"7‚ñ®Ü@xÀ=§ì„+"ˆØ>{ö¬ÿ² ×5@á?©Aƒ†ð ø¦ãÆ£AC¨,SÙQÓ Aƒ†œ ÎzHƒ†¼°ÂMLLô_Ö‚áÌ™3–ËÉN¸b‚ˆÌ5h¸‘‚èܹsþË4hƒpìØ1%ˆ4„dP‚Hƒ 9 ˆpéР!/"ΉÕú"O " ¹&(A¤ACø%ˆ4„jP‚Hƒ 9” Ò—‚Dá” Ò«‚D4„oP‚HC¨%ˆ4hÐÓA " y)(A>A " ¹*(A¤ACø%ˆ4„jP‚Hƒ 9” Ò—‚Dá” Ò«Âõ"ˆ’““åü… ráB’¤¤¤øokР!‚DB5(A¤Aƒ†œJiÈKA ¢ð JiÈUáZ ¢óç“äbR’œ=w^N™ñ{ôÄ)óïY9|ü„Äí>$»ö‘¸=G%áh¢E4äpP‚HC¨%ˆ4hÐÓA " y)(A>A " ¹*\-A”””"q{OÊæØÃrüÔ)9pTöHí»dÇÎÙ´a›¬]0_¶¯\&›Ö¯—Í[öHÌî½2où9lµ4h¸ö ‘†P JiР!§ƒDòRP‚(|•DÿC8f]Ÿ]í[IEND®B`‚apache-arrow-java-19.0.0/docs/source/developers/index.rst000066400000000000000000000016421515451615200233420ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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 apache-arrow-java-19.0.0/docs/source/flight.rst000066400000000000000000000216101515451615200213350ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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 apache-arrow-java-19.0.0/docs/source/flight_sql.rst000066400000000000000000000024101515451615200222110ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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 apache-arrow-java-19.0.0/docs/source/flight_sql_jdbc_driver.rst000066400000000000000000000212201515451615200245460ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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 11+. 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 apache-arrow-java-19.0.0/docs/source/index.rst000066400000000000000000000025351515451615200211740ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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 apache-arrow-java-19.0.0/docs/source/install.rst000066400000000000000000000217641515451615200215400ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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 11 and above. Currently, JDK versions 11, 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. apache-arrow-java-19.0.0/docs/source/ipc.rst000066400000000000000000000174301515451615200206400ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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(); } apache-arrow-java-19.0.0/docs/source/jdbc.rst000066400000000000000000000332541515451615200207710ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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. apache-arrow-java-19.0.0/docs/source/memory.rst000066400000000000000000000616141515451615200214000ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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/11/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. apache-arrow-java-19.0.0/docs/source/overview.rst000066400000000000000000000064001515451615200217260ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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 apache-arrow-java-19.0.0/docs/source/quickstartguide.rst000066400000000000000000000262771515451615200233060ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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 apache-arrow-java-19.0.0/docs/source/reference/000077500000000000000000000000001515451615200212645ustar00rootroot00000000000000apache-arrow-java-19.0.0/docs/source/reference/index.rst000066400000000000000000000016511515451615200231300ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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. apache-arrow-java-19.0.0/docs/source/substrait.rst000066400000000000000000000217551515451615200221120ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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 apache-arrow-java-19.0.0/docs/source/table.rst000066400000000000000000000441051515451615200211530ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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 apache-arrow-java-19.0.0/docs/source/vector.rst000066400000000000000000000435711515451615200213740ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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]. apache-arrow-java-19.0.0/docs/source/vector_schema_root.rst000066400000000000000000000164001515451615200237460ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT 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 apache-arrow-java-19.0.0/flight/000077500000000000000000000000001515451615200163535ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/000077500000000000000000000000001515451615200205565ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/README.md000066400000000000000000000022021515451615200220310ustar00rootroot00000000000000 # Arrow Flight Java Package Exposing Apache Arrow data on the wire. [Protocol Description Slides](https://www.slideshare.net/JacquesNadeau5/apache-arrow-flight-overview) [GRPC Protocol Definition](https://github.com/apache/arrow/blob/main/format/Flight.proto) ## Example Usage See the [Arrow Cookbook](https://arrow.apache.org/cookbook/java/flight.html). apache-arrow-java-19.0.0/flight/flight-core/pom.xml000066400000000000000000000154141515451615200221000ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-flight 19.0.0 flight-core jar Arrow Flight Core An RPC mechanism for transferring ValueVectors. 1 --add-opens=org.apache.arrow.flight.core/org.apache.arrow.flight.perf.impl=protobuf.java --add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED org.apache.arrow arrow-format org.apache.arrow arrow-vector ${arrow.vector.classifier} org.apache.arrow arrow-memory-core org.apache.arrow arrow-memory-netty runtime io.grpc grpc-netty io.grpc grpc-core io.grpc grpc-protobuf io.netty netty-tcnative-boringssl-static io.netty netty-buffer io.netty netty-handler io.netty netty-transport com.google.guava guava io.grpc grpc-stub com.google.protobuf protobuf-java com.google.protobuf protobuf-java-util io.grpc grpc-api io.grpc grpc-services test io.grpc grpc-inprocess test com.fasterxml.jackson.core jackson-databind org.slf4j slf4j-api javax.annotation javax.annotation-api org.immutables value-annotations com.google.api.grpc proto-google-common-protos 2.66.0 test org.apache.arrow arrow-vector ${project.version} tests test-jar test maven-surefire-plugin --add-opens=org.apache.arrow.flight.core/org.apache.arrow.flight.perf.impl=protobuf.java --add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED false ${project.basedir}/../../testing/data org.xolstice.maven.plugins protobuf-maven-plugin src compile compile-custom ${basedir}/../../arrow-format/ test test-compile test-compile-custom org.apache.maven.plugins maven-dependency-plugin analyze analyze-only verify io.netty:netty-tcnative-boringssl-static:* apache-arrow-java-19.0.0/flight/flight-core/src/000077500000000000000000000000001515451615200213455ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/000077500000000000000000000000001515451615200222715ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/000077500000000000000000000000001515451615200232125ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/module-info.java000066400000000000000000000032351515451615200262760ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ module org.apache.arrow.flight.core { exports org.apache.arrow.flight; 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; requires com.fasterxml.jackson.databind; requires com.google.common; requires com.google.errorprone.annotations; requires com.google.protobuf; requires com.google.protobuf.util; requires io.grpc; requires io.grpc.internal; requires io.grpc.netty; requires io.grpc.protobuf; requires io.grpc.stub; requires io.netty.buffer; requires io.netty.common; requires io.netty.handler; requires io.netty.transport; requires org.apache.arrow.format; requires org.apache.arrow.memory.core; requires org.apache.arrow.vector; requires org.slf4j; } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/000077500000000000000000000000001515451615200240015ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/000077500000000000000000000000001515451615200252225ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200263545ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/000077500000000000000000000000001515451615200276315ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/Action.java000066400000000000000000000032321515451615200317110ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.google.protobuf.ByteString; import org.apache.arrow.flight.impl.Flight; /** * An opaque action for the service to perform. * *

    This is a POJO wrapper around the message of the same name in Flight.proto. */ public class Action { private final String type; private final byte[] body; public Action(String type) { this(type, null); } public Action(String type, byte[] body) { this.type = type; this.body = body == null ? new byte[0] : body; } Action(Flight.Action action) { this(action.getType(), action.getBody().toByteArray()); } public String getType() { return type; } public byte[] getBody() { return body; } Flight.Action toProtocol() { return Flight.Action.newBuilder() .setType(getType()) .setBody(ByteString.copyFrom(getBody())) .build(); } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/ActionType.java000066400000000000000000000036771515451615200325700ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import org.apache.arrow.flight.impl.Flight; /** POJO wrapper around protocol specifics for Flight actions. */ public class ActionType { private final String type; private final String description; /** * Construct a new instance. * * @param type The type of action to perform * @param description The description of the type. */ public ActionType(String type, String description) { super(); this.type = type; this.description = description; } /** Constructs a new instance from the corresponding protocol buffer object. */ ActionType(Flight.ActionType type) { this.type = type.getType(); this.description = type.getDescription(); } public String getType() { return type; } public String getDescription() { return description; } /** Converts the POJO to the corresponding protocol buffer type. */ Flight.ActionType toProtocol() { return Flight.ActionType.newBuilder().setType(type).setDescription(description).build(); } @Override public String toString() { return "ActionType{" + "type='" + type + '\'' + ", description='" + description + '\'' + '}'; } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java000066400000000000000000000556101515451615200331020ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.google.common.collect.ImmutableList; import com.google.common.collect.Iterables; import com.google.common.io.ByteStreams; import com.google.protobuf.ByteString; import com.google.protobuf.CodedInputStream; import com.google.protobuf.CodedOutputStream; import com.google.protobuf.WireFormat; import io.grpc.Drainable; import io.grpc.MethodDescriptor.Marshaller; import io.grpc.protobuf.ProtoUtils; import io.netty.buffer.ByteBuf; import io.netty.buffer.ByteBufInputStream; import io.netty.buffer.CompositeByteBuf; import io.netty.buffer.Unpooled; import io.netty.buffer.UnpooledByteBufAllocator; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; import org.apache.arrow.flight.grpc.AddWritableBuffer; import org.apache.arrow.flight.grpc.GetReadableBuffer; import org.apache.arrow.flight.impl.Flight.FlightData; import org.apache.arrow.flight.impl.Flight.FlightDescriptor; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.ipc.message.IpcOption; import org.apache.arrow.vector.ipc.message.MessageMetadataResult; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.MetadataVersion; import org.apache.arrow.vector.types.pojo.Schema; /** The in-memory representation of FlightData used to manage a stream of Arrow messages. */ class ArrowMessage implements AutoCloseable { // If true, deserialize Arrow data by giving Arrow a reference to the underlying gRPC buffer // instead of copying the data. Defaults to true. public static final boolean ENABLE_ZERO_COPY_READ; // If true, serialize Arrow data by giving gRPC a reference to the underlying Arrow buffer // instead of copying the data. Defaults to false. public static final boolean ENABLE_ZERO_COPY_WRITE; static { String zeroCopyReadFlag = System.getProperty("arrow.flight.enable_zero_copy_read"); if (zeroCopyReadFlag == null) { zeroCopyReadFlag = System.getenv("ARROW_FLIGHT_ENABLE_ZERO_COPY_READ"); } String zeroCopyWriteFlag = System.getProperty("arrow.flight.enable_zero_copy_write"); if (zeroCopyWriteFlag == null) { zeroCopyWriteFlag = System.getenv("ARROW_FLIGHT_ENABLE_ZERO_COPY_WRITE"); } ENABLE_ZERO_COPY_READ = !"false".equalsIgnoreCase(zeroCopyReadFlag); ENABLE_ZERO_COPY_WRITE = "true".equalsIgnoreCase(zeroCopyWriteFlag); } private static final int DESCRIPTOR_TAG = (FlightData.FLIGHT_DESCRIPTOR_FIELD_NUMBER << 3) | WireFormat.WIRETYPE_LENGTH_DELIMITED; private static final int BODY_TAG = (FlightData.DATA_BODY_FIELD_NUMBER << 3) | WireFormat.WIRETYPE_LENGTH_DELIMITED; private static final int HEADER_TAG = (FlightData.DATA_HEADER_FIELD_NUMBER << 3) | WireFormat.WIRETYPE_LENGTH_DELIMITED; private static final int APP_METADATA_TAG = (FlightData.APP_METADATA_FIELD_NUMBER << 3) | WireFormat.WIRETYPE_LENGTH_DELIMITED; private static final Marshaller NO_BODY_MARSHALLER = ProtoUtils.marshaller(FlightData.getDefaultInstance()); /** * Get the application-specific metadata in this message. The ArrowMessage retains ownership of * the buffer. */ public ArrowBuf getApplicationMetadata() { return appMetadata; } /** Types of messages that can be sent. */ public enum HeaderType { NONE, SCHEMA, DICTIONARY_BATCH, RECORD_BATCH, TENSOR; public static HeaderType getHeader(byte b) { switch (b) { case 0: return NONE; case 1: return SCHEMA; case 2: return DICTIONARY_BATCH; case 3: return RECORD_BATCH; case 4: return TENSOR; default: throw new UnsupportedOperationException("unknown type: " + b); } } } // Pre-allocated buffers for padding serialized ArrowMessages. private static final List PADDING_BUFFERS = Arrays.asList( null, Unpooled.copiedBuffer(new byte[] {0}), Unpooled.copiedBuffer(new byte[] {0, 0}), Unpooled.copiedBuffer(new byte[] {0, 0, 0}), Unpooled.copiedBuffer(new byte[] {0, 0, 0, 0}), Unpooled.copiedBuffer(new byte[] {0, 0, 0, 0, 0}), Unpooled.copiedBuffer(new byte[] {0, 0, 0, 0, 0, 0}), Unpooled.copiedBuffer(new byte[] {0, 0, 0, 0, 0, 0, 0})); private final IpcOption writeOption; private final FlightDescriptor descriptor; private final MessageMetadataResult message; private final ArrowBuf appMetadata; private final List bufs; private final boolean tryZeroCopyWrite; public ArrowMessage(FlightDescriptor descriptor, Schema schema, IpcOption option) { this.writeOption = option; ByteBuffer serializedMessage = MessageSerializer.serializeMetadata(schema, writeOption); this.message = MessageMetadataResult.create(serializedMessage.slice(), serializedMessage.remaining()); bufs = ImmutableList.of(); this.descriptor = descriptor; this.appMetadata = null; this.tryZeroCopyWrite = false; } /** * Create an ArrowMessage from a record batch and app metadata. * * @param batch The record batch. * @param appMetadata The app metadata. May be null. Takes ownership of the buffer otherwise. * @param tryZeroCopy Whether to enable the zero-copy optimization. */ public ArrowMessage( ArrowRecordBatch batch, ArrowBuf appMetadata, boolean tryZeroCopy, IpcOption option) { this.writeOption = option; ByteBuffer serializedMessage = MessageSerializer.serializeMetadata(batch, writeOption); this.message = MessageMetadataResult.create(serializedMessage.slice(), serializedMessage.remaining()); this.bufs = ImmutableList.copyOf(batch.getBuffers()); this.descriptor = null; this.appMetadata = appMetadata; this.tryZeroCopyWrite = tryZeroCopy; } public ArrowMessage(ArrowDictionaryBatch batch, IpcOption option) { this.writeOption = option; ByteBuffer serializedMessage = MessageSerializer.serializeMetadata(batch, writeOption); serializedMessage = serializedMessage.slice(); this.message = MessageMetadataResult.create(serializedMessage, serializedMessage.remaining()); // asInputStream will free the buffers implicitly, so increment the reference count batch.getDictionary().getBuffers().forEach(buf -> buf.getReferenceManager().retain()); this.bufs = ImmutableList.copyOf(batch.getDictionary().getBuffers()); this.descriptor = null; this.appMetadata = null; this.tryZeroCopyWrite = false; } /** * Create an ArrowMessage containing only application metadata. * * @param appMetadata The application-provided metadata buffer. */ public ArrowMessage(ArrowBuf appMetadata) { // No need to take IpcOption as it's not used to serialize this kind of message. this.writeOption = IpcOption.DEFAULT; this.message = null; this.bufs = ImmutableList.of(); this.descriptor = null; this.appMetadata = appMetadata; this.tryZeroCopyWrite = false; } public ArrowMessage(FlightDescriptor descriptor) { // No need to take IpcOption as it's not used to serialize this kind of message. this.writeOption = IpcOption.DEFAULT; this.message = null; this.bufs = ImmutableList.of(); this.descriptor = descriptor; this.appMetadata = null; this.tryZeroCopyWrite = false; } private ArrowMessage( FlightDescriptor descriptor, MessageMetadataResult message, ArrowBuf appMetadata, ArrowBuf buf) { // No need to take IpcOption as this is used for deserialized ArrowMessage coming from the wire. this.writeOption = message != null ? // avoid writing legacy ipc format by default new IpcOption(false, MetadataVersion.fromFlatbufID(message.getMessage().version())) : IpcOption.DEFAULT; this.message = message; this.descriptor = descriptor; this.appMetadata = appMetadata; this.bufs = buf == null ? ImmutableList.of() : ImmutableList.of(buf); this.tryZeroCopyWrite = false; } public MessageMetadataResult asSchemaMessage() { return message; } public FlightDescriptor getDescriptor() { return descriptor; } public HeaderType getMessageType() { if (message == null) { // Null message occurs for metadata-only messages (in DoExchange) return HeaderType.NONE; } return HeaderType.getHeader(message.headerType()); } public Schema asSchema() { Preconditions.checkArgument(bufs.size() == 0); Preconditions.checkArgument(getMessageType() == HeaderType.SCHEMA); return MessageSerializer.deserializeSchema(message); } public ArrowRecordBatch asRecordBatch() throws IOException { Preconditions.checkArgument( bufs.size() == 1, "A batch can only be consumed if it contains a single ArrowBuf."); Preconditions.checkArgument(getMessageType() == HeaderType.RECORD_BATCH); ArrowBuf underlying = bufs.get(0); underlying.getReferenceManager().retain(); return MessageSerializer.deserializeRecordBatch(message, underlying); } public ArrowDictionaryBatch asDictionaryBatch() throws IOException { Preconditions.checkArgument( bufs.size() == 1, "A batch can only be consumed if it contains a single ArrowBuf."); Preconditions.checkArgument(getMessageType() == HeaderType.DICTIONARY_BATCH); ArrowBuf underlying = bufs.get(0); // Retain a reference to keep the batch alive when the message is closed underlying.getReferenceManager().retain(); // Do not set drained - we still want to release our reference return MessageSerializer.deserializeDictionaryBatch(message, underlying); } public Iterable getBufs() { return Iterables.unmodifiableIterable(bufs); } private static ArrowMessage frame(BufferAllocator allocator, final InputStream stream) { try { FlightDescriptor descriptor = null; MessageMetadataResult header = null; ArrowBuf body = null; ArrowBuf appMetadata = null; while (stream.available() > 0) { final int tagFirstByte = stream.read(); if (tagFirstByte == -1) { break; } int tag = readRawVarint32(tagFirstByte, stream); switch (tag) { case DESCRIPTOR_TAG: { int size = readRawVarint32(stream); byte[] bytes = new byte[size]; ByteStreams.readFully(stream, bytes); descriptor = FlightDescriptor.parseFrom(bytes); break; } case HEADER_TAG: { int size = readRawVarint32(stream); byte[] bytes = new byte[size]; ByteStreams.readFully(stream, bytes); header = MessageMetadataResult.create(ByteBuffer.wrap(bytes), size); break; } case APP_METADATA_TAG: { int size = readRawVarint32(stream); appMetadata = allocator.buffer(size); GetReadableBuffer.readIntoBuffer(stream, appMetadata, size, ENABLE_ZERO_COPY_READ); break; } case BODY_TAG: if (body != null) { // only read last body. body.getReferenceManager().release(); body = null; } int size = readRawVarint32(stream); body = allocator.buffer(size); GetReadableBuffer.readIntoBuffer(stream, body, size, ENABLE_ZERO_COPY_READ); break; default: // ignore unknown fields. } } // Protobuf implementations can omit empty fields, such as body; for some message types, like // RecordBatch, // this will fail later as we still expect an empty buffer. In those cases only, fill in an // empty buffer here - // in other cases, like Schema, having an unexpected empty buffer will also cause failures. // We don't fill in defaults for fields like header, for which there is no reasonable default, // or for appMetadata // or descriptor, which are intended to be empty in some cases. if (header != null) { switch (HeaderType.getHeader(header.headerType())) { case SCHEMA: // Ignore 0-length buffers in case a Protobuf implementation wrote it out if (body != null && body.capacity() == 0) { body.close(); body = null; } break; case DICTIONARY_BATCH: case RECORD_BATCH: // A Protobuf implementation can skip 0-length bodies, so ensure we fill it in here if (body == null) { body = allocator.getEmpty(); } break; case NONE: case TENSOR: default: // Do nothing break; } } return new ArrowMessage(descriptor, header, appMetadata, body); } catch (Exception ioe) { throw new RuntimeException(ioe); } } 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); } /** * Convert the ArrowMessage to an InputStream. * *

    Implicitly, this transfers ownership of the contained buffers to the InputStream. * * @return InputStream */ private InputStream asInputStream() { if (message == null) { // If we have no IPC message, it's a pure-metadata message final FlightData.Builder builder = FlightData.newBuilder(); if (descriptor != null) { builder.setFlightDescriptor(descriptor); } if (appMetadata != null) { builder.setAppMetadata(ByteString.copyFrom(appMetadata.nioBuffer())); } return NO_BODY_MARSHALLER.stream(builder.build()); } try { final ByteString bytes = ByteString.copyFrom(message.getMessageBuffer(), message.bytesAfterMessage()); if (getMessageType() == HeaderType.SCHEMA) { final FlightData.Builder builder = FlightData.newBuilder().setDataHeader(bytes); if (descriptor != null) { builder.setFlightDescriptor(descriptor); } Preconditions.checkArgument(bufs.isEmpty()); return NO_BODY_MARSHALLER.stream(builder.build()); } Preconditions.checkArgument( getMessageType() == HeaderType.RECORD_BATCH || getMessageType() == HeaderType.DICTIONARY_BATCH); // There may be no buffers in the case that we write only a null array Preconditions.checkArgument( descriptor == null, "Descriptor should only be included in the schema message."); ByteArrayOutputStream baos = new ByteArrayOutputStream(); CodedOutputStream cos = CodedOutputStream.newInstance(baos); cos.writeBytes(FlightData.DATA_HEADER_FIELD_NUMBER, bytes); if (appMetadata != null && appMetadata.capacity() > 0) { // Must call slice() as CodedOutputStream#writeByteBuffer writes -capacity- bytes, not // -limit- bytes cos.writeByteBuffer(FlightData.APP_METADATA_FIELD_NUMBER, appMetadata.nioBuffer().slice()); } cos.writeTag(FlightData.DATA_BODY_FIELD_NUMBER, WireFormat.WIRETYPE_LENGTH_DELIMITED); int size = 0; List allBufs = new ArrayList<>(); for (ArrowBuf b : bufs) { // [ARROW-11066] This creates a Netty buffer whose refcnt is INDEPENDENT of the backing // Arrow buffer. This is susceptible to use-after-free, so we subclass CompositeByteBuf // below to tie the Arrow buffer refcnt to the Netty buffer refcnt allBufs.add(Unpooled.wrappedBuffer(b.nioBuffer()).retain()); size += (int) b.readableBytes(); // [ARROW-4213] These buffers must be aligned to an 8-byte boundary in order to be readable // from C++. if (b.readableBytes() % 8 != 0) { int paddingBytes = (int) (8 - (b.readableBytes() % 8)); assert paddingBytes > 0 && paddingBytes < 8; size += paddingBytes; allBufs.add(PADDING_BUFFERS.get(paddingBytes).retain()); } } // rawvarint is used for length definition. cos.writeUInt32NoTag(size); cos.flush(); ByteBuf initialBuf = Unpooled.buffer(baos.size()); initialBuf.writeBytes(baos.toByteArray()); final CompositeByteBuf bb; final ImmutableList byteBufs = ImmutableList.builder().add(initialBuf).addAll(allBufs).build(); // See: https://github.com/apache/arrow/issues/40039 // CompositeByteBuf requires us to pass maxNumComponents to constructor. // This number will be used to decide when to stop adding new components as separate buffers // and instead merge existing components into a new buffer by performing a memory copy. // We want to avoind memory copies as much as possible so we want to set the limit that won't // be reached. // At a first glance it seems reasonable to set limit to byteBufs.size() + 1, // because it will be enough to avoid merges of byteBufs that we pass to constructor. // But later this buffer will be written to socket by Netty // and DefaultHttp2ConnectionEncoder uses CoalescingBufferQueue to combine small buffers into // one. // Method CoalescingBufferQueue.compose will check if current buffer is already a // CompositeByteBuf // and if it's the case it will just add a new component to this buffer. // But in out case if we set maxNumComponents=byteBufs.size() + 1 it will happen on the first // attempt // to write data to socket because header message is small and Netty will always try to // compine it with the // large CompositeByteBuf we're creating here. // We never want additional memory copies so setting the limit to Integer.MAX_VALUE final int maxNumComponents = Integer.MAX_VALUE; if (tryZeroCopyWrite) { bb = new ArrowBufRetainingCompositeByteBuf(maxNumComponents, byteBufs, bufs); } else { // Don't retain the buffers in the non-zero-copy path since we're copying them bb = new CompositeByteBuf( UnpooledByteBufAllocator.DEFAULT, /* direct */ true, maxNumComponents, byteBufs); } return new DrainableByteBufInputStream(bb, tryZeroCopyWrite); } catch (Exception ex) { throw new RuntimeException("Unexpected IO Exception", ex); } } /** * ARROW-11066: enable the zero-copy optimization and protect against use-after-free. * *

    When you send a message through gRPC, the following happens: 1. gRPC immediately serializes * the message, eventually calling asInputStream above. 2. gRPC buffers the serialized message for * sending. 3. Later, gRPC will actually write out the message. * *

    The problem with this is that when the zero-copy optimization is enabled, Flight * "serializes" the message by handing gRPC references to Arrow data. That means we need a way to * keep the Arrow buffers valid until gRPC actually writes them, else, we'll read invalid data or * segfault. gRPC doesn't know anything about Arrow buffers, either. * *

    This class solves that issue by bridging Arrow and Netty/gRPC. We increment the refcnt on a * set of Arrow backing buffers and decrement them once the Netty buffers are freed by gRPC. */ private static final class ArrowBufRetainingCompositeByteBuf extends CompositeByteBuf { // Arrow buffers that back the Netty ByteBufs here; ByteBufs held by this class are // either slices of one of the ArrowBufs or independently allocated. final List backingBuffers; boolean freed; ArrowBufRetainingCompositeByteBuf( int maxNumComponents, Iterable buffers, List backingBuffers) { super(UnpooledByteBufAllocator.DEFAULT, /* direct */ true, maxNumComponents, buffers); this.backingBuffers = backingBuffers; this.freed = false; // N.B. the Netty superclass avoids enhanced-for to reduce GC pressure, so follow that here for (int i = 0; i < backingBuffers.size(); i++) { backingBuffers.get(i).getReferenceManager().retain(); } } @Override protected void deallocate() { super.deallocate(); if (freed) { return; } freed = true; for (int i = 0; i < backingBuffers.size(); i++) { backingBuffers.get(i).getReferenceManager().release(); } } } private static class DrainableByteBufInputStream extends ByteBufInputStream implements Drainable { private final CompositeByteBuf buf; private final boolean isZeroCopy; public DrainableByteBufInputStream(CompositeByteBuf buffer, boolean isZeroCopy) { super(buffer, buffer.readableBytes(), true); this.buf = buffer; this.isZeroCopy = isZeroCopy; } @Override public int drainTo(OutputStream target) throws IOException { int size = buf.readableBytes(); AddWritableBuffer.add(buf, target, isZeroCopy); return size; } @Override public void close() { buf.release(); } } public static Marshaller createMarshaller(BufferAllocator allocator) { return new ArrowMessageHolderMarshaller(allocator); } private static class ArrowMessageHolderMarshaller implements Marshaller { private final BufferAllocator allocator; public ArrowMessageHolderMarshaller(BufferAllocator allocator) { this.allocator = allocator; } @Override public InputStream stream(ArrowMessage value) { return value.asInputStream(); } @Override public ArrowMessage parse(InputStream stream) { return ArrowMessage.frame(allocator, stream); } } @Override public void close() throws Exception { AutoCloseables.close(Iterables.concat(bufs, Collections.singletonList(appMetadata))); } } AsyncPutListener.java000066400000000000000000000042141515451615200336720ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; import org.apache.arrow.flight.grpc.StatusUtils; /** * A handler for server-sent application metadata messages during a Flight DoPut operation. * *

    To handle messages, create an instance of this class overriding {@link #onNext(PutResult)}. * The other methods should not be overridden. */ public class AsyncPutListener implements FlightClient.PutListener { private CompletableFuture completed; public AsyncPutListener() { completed = new CompletableFuture<>(); } /** * Wait for the stream to finish on the server side. You must call this to be notified of any * errors that may have happened during the upload. */ @Override public final void getResult() { try { completed.get(); } catch (ExecutionException e) { throw StatusUtils.fromThrowable(e.getCause()); } catch (InterruptedException e) { throw StatusUtils.fromThrowable(e); } } @Override public void onNext(PutResult val) {} @Override public final void onError(Throwable t) { completed.completeExceptionally(StatusUtils.fromThrowable(t)); } @Override public final void onCompleted() { completed.complete(null); } @Override public boolean isCancelled() { return completed.isDone(); } } BackpressureStrategy.java000066400000000000000000000124041515451615200345720ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.google.common.base.Preconditions; import org.apache.arrow.vector.VectorSchemaRoot; /** * Helper interface to dynamically handle backpressure when implementing FlightProducers. This must * only be used in FlightProducer implementations that are non-blocking. */ public interface BackpressureStrategy { /** The state of the client after a call to waitForListener. */ enum WaitResult { /** Listener is ready. */ READY, /** Listener was cancelled by the client. */ CANCELLED, /** Timed out waiting for the listener to change state. */ TIMEOUT, /** Indicates that the wait was interrupted for a reason unrelated to the listener itself. */ OTHER } /** * Set up operations to work against the given listener. * *

    This must be called exactly once and before any calls to {@link #waitForListener(long)} and * {@link OutboundStreamListener#start(VectorSchemaRoot)} * * @param listener The listener this strategy applies to. */ void register(FlightProducer.ServerStreamListener listener); /** * Waits for the listener to be ready or cancelled up to the given timeout. * * @param timeout The timeout in milliseconds. Infinite if timeout is <= 0. * @return The result of the wait. */ WaitResult waitForListener(long timeout); /** * A back pressure strategy that uses callbacks to notify when the client is ready or cancelled. */ class CallbackBackpressureStrategy implements BackpressureStrategy { private final Object lock = new Object(); private FlightProducer.ServerStreamListener listener; @Override public void register(FlightProducer.ServerStreamListener listener) { this.listener = listener; listener.setOnReadyHandler(this::onReady); listener.setOnCancelHandler(this::onCancel); } @Override public WaitResult waitForListener(long timeout) { Preconditions.checkNotNull(listener); long remainingTimeout = timeout; final long startTime = System.currentTimeMillis(); synchronized (lock) { while (!listener.isReady() && !listener.isCancelled()) { try { lock.wait(remainingTimeout); if (timeout != 0) { // If timeout was zero explicitly, we should never report timeout. remainingTimeout = startTime + timeout - System.currentTimeMillis(); if (remainingTimeout <= 0) { return WaitResult.TIMEOUT; } } if (!shouldContinueWaiting(listener, remainingTimeout)) { return WaitResult.OTHER; } } catch (InterruptedException ex) { Thread.currentThread().interrupt(); return WaitResult.OTHER; } } if (listener.isReady()) { return WaitResult.READY; } else if (listener.isCancelled()) { return WaitResult.CANCELLED; } else if (System.currentTimeMillis() > startTime + timeout) { return WaitResult.TIMEOUT; } throw new RuntimeException("Invalid state when waiting for listener."); } } /** * Interrupt waiting on the listener to change state. * *

    This method can be used in conjunction with {@link * #shouldContinueWaiting(FlightProducer.ServerStreamListener, long)} to allow FlightProducers * to terminate streams internally and notify clients. */ public void interruptWait() { synchronized (lock) { lock.notifyAll(); } } /** * Callback function to run to check if the listener should continue to be waited on if it * leaves the waiting state without being cancelled, ready, or timed out. * *

    This method should be used to determine if the wait on the listener was interrupted * explicitly using a call to {@link #interruptWait()} or if it was woken up due to a spurious * wake. */ protected boolean shouldContinueWaiting( FlightProducer.ServerStreamListener listener, long remainingTimeout) { return true; } /** Callback to execute when the listener becomes ready. */ protected void readyCallback() {} /** Callback to execute when the listener is cancelled. */ protected void cancelCallback() {} private void onReady() { synchronized (lock) { readyCallback(); lock.notifyAll(); } } private void onCancel() { synchronized (lock) { cancelCallback(); lock.notifyAll(); } } } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/CallHeaders.java000066400000000000000000000042031515451615200326420ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.util.Set; /** A set of metadata key value pairs for a call (request or response). */ public interface CallHeaders { /** Get the value of a metadata key. If multiple values are present, then get the last one. */ String get(String key); /** 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. * * @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. * * @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); /** * Insert a metadata pair with the given value. * *

    Duplicate metadata are permitted. */ void insert(String key, String value); /** * Insert a metadata pair with the given value. * *

    Duplicate metadata are permitted. */ void insert(String key, byte[] value); /** Get a set of all the metadata keys. */ Set keys(); /** Check whether the given metadata key is present. */ boolean containsKey(String key); } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/CallInfo.java000066400000000000000000000021121515451615200321570ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** A description of a Flight call for middleware to inspect. */ public final class CallInfo { private final FlightMethod method; public CallInfo(FlightMethod method) { this.method = method; } public FlightMethod method() { return method; } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/CallOption.java000066400000000000000000000017051515451615200325430ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** Per-call RPC options. These are hints to the underlying RPC layer and may not be respected. */ public interface CallOption {} apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/CallOptions.java000066400000000000000000000035111515451615200327230ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import io.grpc.stub.AbstractStub; import java.util.concurrent.TimeUnit; /** Common call options. */ public class CallOptions { public static CallOption timeout(long duration, TimeUnit unit) { return new Timeout(duration, unit); } static > T wrapStub(T stub, CallOption[] options) { for (CallOption option : options) { if (option instanceof GrpcCallOption) { stub = ((GrpcCallOption) option).wrapStub(stub); } } return stub; } private static class Timeout implements GrpcCallOption { long timeout; TimeUnit timeoutUnit; Timeout(long timeout, TimeUnit timeoutUnit) { this.timeout = timeout; this.timeoutUnit = timeoutUnit; } @Override public > T wrapStub(T stub) { return stub.withDeadlineAfter(timeout, timeoutUnit); } } /** CallOptions specific to GRPC stubs. */ public interface GrpcCallOption extends CallOption { > T wrapStub(T stub); } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/CallStatus.java000066400000000000000000000116701515451615200325600ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.util.Objects; import org.apache.arrow.flight.FlightProducer.ServerStreamListener; import org.apache.arrow.flight.FlightProducer.StreamListener; /** * The result of a Flight RPC, consisting of a status code with an optional description and/or * exception that led to the status. * *

    If raised or sent through {@link StreamListener#onError(Throwable)} or {@link * ServerStreamListener#error(Throwable)}, the client call will raise the same error (a {@link * FlightRuntimeException} with the same {@link FlightStatusCode} and description). The exception * within, if present, will not be sent to the client. */ public class CallStatus { private final FlightStatusCode code; private final Throwable cause; private final String description; private final ErrorFlightMetadata metadata; public static final CallStatus UNKNOWN = FlightStatusCode.UNKNOWN.toStatus(); public static final CallStatus INTERNAL = FlightStatusCode.INTERNAL.toStatus(); public static final CallStatus INVALID_ARGUMENT = FlightStatusCode.INVALID_ARGUMENT.toStatus(); public static final CallStatus TIMED_OUT = FlightStatusCode.TIMED_OUT.toStatus(); public static final CallStatus NOT_FOUND = FlightStatusCode.NOT_FOUND.toStatus(); public static final CallStatus ALREADY_EXISTS = FlightStatusCode.ALREADY_EXISTS.toStatus(); public static final CallStatus CANCELLED = FlightStatusCode.CANCELLED.toStatus(); public static final CallStatus UNAUTHENTICATED = FlightStatusCode.UNAUTHENTICATED.toStatus(); public static final CallStatus UNAUTHORIZED = FlightStatusCode.UNAUTHORIZED.toStatus(); public static final CallStatus UNIMPLEMENTED = FlightStatusCode.UNIMPLEMENTED.toStatus(); public static final CallStatus UNAVAILABLE = FlightStatusCode.UNAVAILABLE.toStatus(); public static final CallStatus RESOURCE_EXHAUSTED = FlightStatusCode.RESOURCE_EXHAUSTED.toStatus(); /** * Create a new status. * * @param code The status code. * @param cause An exception that resulted in this status (or null). * @param description A description of the status (or null). */ public CallStatus( FlightStatusCode code, Throwable cause, String description, ErrorFlightMetadata metadata) { this.code = Objects.requireNonNull(code); this.cause = cause; this.description = description == null ? "" : description; this.metadata = metadata == null ? new ErrorFlightMetadata() : metadata; } /** * Create a new status with no cause or description. * * @param code The status code. */ public CallStatus(FlightStatusCode code) { this(code, /* no cause */ null, /* no description */ null, /* no metadata */ null); } /** The status code describing the result of the RPC. */ public FlightStatusCode code() { return code; } /** The exception that led to this result. May be null. */ public Throwable cause() { return cause; } /** A description of the result. */ public String description() { return description; } /** * Metadata associated with the exception. * *

    May be null. */ public ErrorFlightMetadata metadata() { return metadata; } /** Return a copy of this status with an error message. */ public CallStatus withDescription(String message) { return new CallStatus(code, cause, message, metadata); } /** * Return a copy of this status with the given exception as the cause. This will not be sent over * the wire. */ public CallStatus withCause(Throwable t) { return new CallStatus(code, t, description, metadata); } /** Return a copy of this status with associated exception metadata. */ public CallStatus withMetadata(ErrorFlightMetadata metadata) { return new CallStatus(code, cause, description, metadata); } /** Convert the status to an equivalent exception. */ public FlightRuntimeException toRuntimeException() { return new FlightRuntimeException(this); } @Override public String toString() { return "CallStatus{" + "code=" + code + ", cause=" + cause + ", description='" + description + "', metadata='" + metadata + '\'' + '}'; } } CancelFlightInfoRequest.java000066400000000000000000000047101515451615200351270ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.io.IOException; import java.net.URISyntaxException; import java.nio.ByteBuffer; import java.util.Objects; import org.apache.arrow.flight.impl.Flight; /** A request to cancel a FlightInfo. */ public class CancelFlightInfoRequest { private final FlightInfo info; public CancelFlightInfoRequest(FlightInfo info) { this.info = Objects.requireNonNull(info); } CancelFlightInfoRequest(Flight.CancelFlightInfoRequest proto) throws URISyntaxException { this(new FlightInfo(proto.getInfo())); } public FlightInfo getInfo() { return info; } Flight.CancelFlightInfoRequest toProtocol() { Flight.CancelFlightInfoRequest.Builder b = Flight.CancelFlightInfoRequest.newBuilder(); b.setInfo(info.toProtocol()); return b.build(); } /** * Get the serialized form of this protocol message. * *

    Intended to help interoperability by allowing non-Flight services to still return Flight * types. */ public ByteBuffer serialize() { return ByteBuffer.wrap(toProtocol().toByteArray()); } /** * Parse the serialized form of this protocol message. * *

    Intended to help interoperability by allowing Flight clients to obtain stream info from * non-Flight services. * * @param serialized The serialized form of the message, as returned by {@link #serialize()}. * @return The deserialized message. * @throws IOException if the serialized form is invalid. */ public static CancelFlightInfoRequest deserialize(ByteBuffer serialized) throws IOException, URISyntaxException { return new CancelFlightInfoRequest(Flight.CancelFlightInfoRequest.parseFrom(serialized)); } } CancelFlightInfoResult.java000066400000000000000000000073371515451615200347650ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.io.IOException; import java.nio.ByteBuffer; import java.util.Objects; import org.apache.arrow.flight.impl.Flight; /** The result of cancelling a FlightInfo. */ public class CancelFlightInfoResult { private final CancelStatus status; public CancelFlightInfoResult(CancelStatus status) { this.status = status; } CancelFlightInfoResult(Flight.CancelFlightInfoResult proto) { switch (proto.getStatus()) { case CANCEL_STATUS_UNSPECIFIED: status = CancelStatus.UNSPECIFIED; break; case CANCEL_STATUS_CANCELLED: status = CancelStatus.CANCELLED; break; case CANCEL_STATUS_CANCELLING: status = CancelStatus.CANCELLING; break; case CANCEL_STATUS_NOT_CANCELLABLE: status = CancelStatus.NOT_CANCELLABLE; break; default: throw new IllegalArgumentException(""); } } public CancelStatus getStatus() { return status; } Flight.CancelFlightInfoResult toProtocol() { Flight.CancelFlightInfoResult.Builder b = Flight.CancelFlightInfoResult.newBuilder(); switch (status) { case UNSPECIFIED: b.setStatus(Flight.CancelStatus.CANCEL_STATUS_UNSPECIFIED); break; case CANCELLED: b.setStatus(Flight.CancelStatus.CANCEL_STATUS_CANCELLED); break; case CANCELLING: b.setStatus(Flight.CancelStatus.CANCEL_STATUS_CANCELLING); break; case NOT_CANCELLABLE: b.setStatus(Flight.CancelStatus.CANCEL_STATUS_NOT_CANCELLABLE); break; default: // Not possible throw new AssertionError(); } return b.build(); } /** * Get the serialized form of this protocol message. * *

    Intended to help interoperability by allowing non-Flight services to still return Flight * types. */ public ByteBuffer serialize() { return ByteBuffer.wrap(toProtocol().toByteArray()); } /** * Parse the serialized form of this protocol message. * *

    Intended to help interoperability by allowing Flight clients to obtain stream info from * non-Flight services. * * @param serialized The serialized form of the message, as returned by {@link #serialize()}. * @return The deserialized message. * @throws IOException if the serialized form is invalid. */ public static CancelFlightInfoResult deserialize(ByteBuffer serialized) throws IOException { return new CancelFlightInfoResult(Flight.CancelFlightInfoResult.parseFrom(serialized)); } @Override public boolean equals(Object o) { if (this == o) { return true; } if (!(o instanceof CancelFlightInfoResult)) { return false; } CancelFlightInfoResult that = (CancelFlightInfoResult) o; return status == that.status; } @Override public int hashCode() { return Objects.hash(status); } @Override public String toString() { return "CancelFlightInfoResult{" + "status=" + status + '}'; } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/CancelStatus.java000066400000000000000000000027331515451615200330720ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** The result of cancelling a FlightInfo. */ public enum CancelStatus { /** * The cancellation status is unknown. Servers should avoid using this value (send a NOT_FOUND * error if the requested query is not known). Clients can retry the request. */ UNSPECIFIED, /** * The cancellation request is complete. Subsequent requests with the same payload may return * CANCELLED or a NOT_FOUND error. */ CANCELLED, /** The cancellation request is in progress. The client may retry the cancellation request. */ CANCELLING, /** The query is not cancellable. The client should not retry the cancellation request. */ NOT_CANCELLABLE, ; } CloseSessionRequest.java000066400000000000000000000040561515451615200344040ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.io.IOException; import java.nio.ByteBuffer; import org.apache.arrow.flight.impl.Flight; /** A request to close/invalidate a server session context. */ public class CloseSessionRequest { public CloseSessionRequest() {} CloseSessionRequest(Flight.CloseSessionRequest proto) {} Flight.CloseSessionRequest toProtocol() { return Flight.CloseSessionRequest.getDefaultInstance(); } /** * Get the serialized form of this protocol message. * *

    Intended to help interoperability by allowing non-Flight services to still return Flight * types. */ public ByteBuffer serialize() { return ByteBuffer.wrap(toProtocol().toByteArray()); } /** * Parse the serialized form of this protocol message. * *

    Intended to help interoperability by allowing Flight clients to obtain stream info from * non-Flight services. * * @param serialized The serialized form of the message, as returned by {@link #serialize()}. * @return The deserialized message. * @throws IOException if the serialized form is invalid. */ public static CloseSessionRequest deserialize(ByteBuffer serialized) throws IOException { return new CloseSessionRequest(Flight.CloseSessionRequest.parseFrom(serialized)); } } CloseSessionResult.java000066400000000000000000000062571515451615200342370ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.io.IOException; import java.nio.ByteBuffer; import org.apache.arrow.flight.impl.Flight; /** The result of attempting to close/invalidate a server session context. */ public class CloseSessionResult { /** Close operation result status values. */ public enum Status { /** * The session close status is unknown. Servers should avoid using this value (send a NOT_FOUND * error if the requested session is not known). Clients can retry the request. */ UNSPECIFIED, /** The session close request is complete. */ CLOSED, /** The session close request is in progress. The client may retry the request. */ CLOSING, /** The session is not closeable. */ NOT_CLOSABLE, ; public static Status fromProtocol(Flight.CloseSessionResult.Status proto) { return values()[proto.getNumber()]; } public Flight.CloseSessionResult.Status toProtocol() { return Flight.CloseSessionResult.Status.values()[ordinal()]; } } private final Status status; public CloseSessionResult(Status status) { this.status = status; } CloseSessionResult(Flight.CloseSessionResult proto) { status = Status.fromProtocol(proto.getStatus()); if (status == null) { // Unreachable throw new IllegalArgumentException(""); } } public Status getStatus() { return status; } Flight.CloseSessionResult toProtocol() { Flight.CloseSessionResult.Builder b = Flight.CloseSessionResult.newBuilder(); b.setStatus(status.toProtocol()); return b.build(); } /** * Get the serialized form of this protocol message. * *

    Intended to help interoperability by allowing non-Flight services to still return Flight * types. */ public ByteBuffer serialize() { return ByteBuffer.wrap(toProtocol().toByteArray()); } /** * Parse the serialized form of this protocol message. * *

    Intended to help interoperability by allowing Flight clients to obtain stream info from * non-Flight services. * * @param serialized The serialized form of the message, as returned by {@link #serialize()}. * @return The deserialized message. * @throws IOException if the serialized form is invalid. */ public static CloseSessionResult deserialize(ByteBuffer serialized) throws IOException { return new CloseSessionResult(Flight.CloseSessionResult.parseFrom(serialized)); } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/Criteria.java000066400000000000000000000032131515451615200322350ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.google.protobuf.ByteString; import org.apache.arrow.flight.impl.Flight; /** * An opaque object that can be used to filter a list of streams available from a server. * *

    This is a POJO wrapper around the protobuf Criteria message. */ public class Criteria { public static Criteria ALL = new Criteria((byte[]) null); private final byte[] bytes; public Criteria(byte[] bytes) { this.bytes = bytes; } Criteria(Flight.Criteria criteria) { this.bytes = criteria.getExpression().toByteArray(); } /** Get the contained filter criteria. */ public byte[] getExpression() { return bytes; } Flight.Criteria asCriteria() { Flight.Criteria.Builder b = Flight.Criteria.newBuilder(); if (bytes != null) { b.setExpression(ByteString.copyFrom(bytes)); } return b.build(); } } DictionaryUtils.java000066400000000000000000000127621515451615200335530ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.util.ArrayList; import java.util.Collections; import java.util.HashSet; import java.util.List; import java.util.Set; import java.util.function.Consumer; import java.util.stream.Collectors; import org.apache.arrow.flight.impl.Flight; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch; import org.apache.arrow.vector.ipc.message.IpcOption; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.DictionaryUtility; import org.apache.arrow.vector.validate.MetadataV4UnionChecker; /** Utilities to work with dictionaries in Flight. */ final class DictionaryUtils { private DictionaryUtils() { throw new UnsupportedOperationException("Do not instantiate this class."); } /** * Generate all the necessary Flight messages to send a schema and associated dictionaries. * * @throws Exception if there was an error closing {@link ArrowMessage} objects. This is not * generally expected. */ static Schema generateSchemaMessages( final Schema originalSchema, final FlightDescriptor descriptor, final DictionaryProvider provider, final IpcOption option, final Consumer messageCallback) throws Exception { final Set dictionaryIds = new HashSet<>(); final Schema schema = generateSchema(originalSchema, provider, dictionaryIds); MetadataV4UnionChecker.checkForUnion(schema.getFields().iterator(), option.metadataVersion); // Send the schema message final Flight.FlightDescriptor protoDescriptor = descriptor == null ? null : descriptor.toProtocol(); try (final ArrowMessage message = new ArrowMessage(protoDescriptor, schema, option)) { messageCallback.accept(message); } // Create and write dictionary batches for (Long id : dictionaryIds) { final Dictionary dictionary = provider.lookup(id); final FieldVector vector = dictionary.getVector(); final int count = vector.getValueCount(); // Do NOT close this root, as it does not actually own the vector. final VectorSchemaRoot dictRoot = new VectorSchemaRoot( Collections.singletonList(vector.getField()), Collections.singletonList(vector), count); final VectorUnloader unloader = new VectorUnloader(dictRoot); try (final ArrowDictionaryBatch dictionaryBatch = new ArrowDictionaryBatch(id, unloader.getRecordBatch()); final ArrowMessage message = new ArrowMessage(dictionaryBatch, option)) { messageCallback.accept(message); } } return schema; } static void closeDictionaries(final Schema schema, final DictionaryProvider provider) throws Exception { // Close dictionaries final Set dictionaryIds = new HashSet<>(); schema .getFields() .forEach(field -> DictionaryUtility.toMessageFormat(field, provider, dictionaryIds)); final List dictionaryVectors = dictionaryIds.stream() .map(id -> (AutoCloseable) provider.lookup(id).getVector()) .collect(Collectors.toList()); AutoCloseables.close(dictionaryVectors); } /** * Generates the schema to send with flight messages. If the schema contains no field with a * dictionary, it will return the schema as is. Otherwise, it will return a newly created a new * schema after converting the fields. * * @param originalSchema the original schema. * @param provider the dictionary provider. * @param dictionaryIds dictionary IDs that are used. * @return the schema to send with the flight messages. */ static Schema generateSchema( final Schema originalSchema, final DictionaryProvider provider, Set dictionaryIds) { // first determine if a new schema needs to be created. boolean createSchema = false; for (Field field : originalSchema.getFields()) { if (DictionaryUtility.needConvertToMessageFormat(field)) { createSchema = true; break; } } if (!createSchema) { return originalSchema; } else { final List fields = new ArrayList<>(originalSchema.getFields().size()); for (final Field field : originalSchema.getFields()) { fields.add(DictionaryUtility.toMessageFormat(field, provider, dictionaryIds)); } return new Schema(fields, originalSchema.getCustomMetadata()); } } } ErrorFlightMetadata.java000066400000000000000000000044251515451615200343120ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.google.common.collect.Iterables; import com.google.common.collect.LinkedListMultimap; import com.google.common.collect.Multimap; import java.nio.charset.StandardCharsets; import java.util.Set; import java.util.stream.Collectors; import java.util.stream.StreamSupport; /** metadata container specific to the binary metadata held in the grpc trailer. */ public class ErrorFlightMetadata implements CallHeaders { private final Multimap metadata = LinkedListMultimap.create(); public ErrorFlightMetadata() {} @Override public String get(String key) { return new String(getByte(key), StandardCharsets.US_ASCII); } @Override public byte[] getByte(String key) { return Iterables.getLast(metadata.get(key)); } @Override public Iterable getAll(String key) { return StreamSupport.stream(getAllByte(key).spliterator(), false) .map(b -> new String(b, StandardCharsets.US_ASCII)) .collect(Collectors.toList()); } @Override public Iterable getAllByte(String key) { return metadata.get(key); } @Override public void insert(String key, String value) { metadata.put(key, value.getBytes(StandardCharsets.UTF_8)); } @Override public void insert(String key, byte[] value) { metadata.put(key, value); } @Override public Set keys() { return metadata.keySet(); } @Override public boolean containsKey(String key) { return metadata.containsKey(key); } } FlightBindingService.java000066400000000000000000000170441515451615200344540ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.google.common.collect.ImmutableSet; import io.grpc.BindableService; import io.grpc.MethodDescriptor; import io.grpc.MethodDescriptor.MethodType; import io.grpc.ServerMethodDefinition; import io.grpc.ServerServiceDefinition; import io.grpc.ServiceDescriptor; import io.grpc.protobuf.ProtoUtils; import io.grpc.stub.ServerCalls; import io.grpc.stub.StreamObserver; import java.util.Set; import java.util.concurrent.ExecutorService; import org.apache.arrow.flight.auth.ServerAuthHandler; import org.apache.arrow.flight.impl.Flight; import org.apache.arrow.flight.impl.Flight.PutResult; import org.apache.arrow.flight.impl.FlightServiceGrpc; import org.apache.arrow.memory.BufferAllocator; /** Extends the basic flight service to override some methods for more efficient implementations. */ class FlightBindingService implements BindableService { private static final String DO_GET = MethodDescriptor.generateFullMethodName(FlightConstants.SERVICE, "DoGet"); private static final String DO_PUT = MethodDescriptor.generateFullMethodName(FlightConstants.SERVICE, "DoPut"); private static final String DO_EXCHANGE = MethodDescriptor.generateFullMethodName(FlightConstants.SERVICE, "DoExchange"); private static final Set OVERRIDE_METHODS = ImmutableSet.of(DO_GET, DO_PUT, DO_EXCHANGE); private final FlightService delegate; private final BufferAllocator allocator; public FlightBindingService( BufferAllocator allocator, FlightProducer producer, ServerAuthHandler authHandler, ExecutorService executor) { this.allocator = allocator; this.delegate = new FlightService(allocator, producer, authHandler, executor); } public static MethodDescriptor getDoGetDescriptor( BufferAllocator allocator) { return MethodDescriptor.newBuilder() .setType(io.grpc.MethodDescriptor.MethodType.SERVER_STREAMING) .setFullMethodName(DO_GET) .setSampledToLocalTracing(false) .setRequestMarshaller(ProtoUtils.marshaller(Flight.Ticket.getDefaultInstance())) .setResponseMarshaller(ArrowMessage.createMarshaller(allocator)) .setSchemaDescriptor(FlightServiceGrpc.getDoGetMethod().getSchemaDescriptor()) .build(); } public static MethodDescriptor getDoPutDescriptor( BufferAllocator allocator) { return MethodDescriptor.newBuilder() .setType(MethodType.BIDI_STREAMING) .setFullMethodName(DO_PUT) .setSampledToLocalTracing(false) .setRequestMarshaller(ArrowMessage.createMarshaller(allocator)) .setResponseMarshaller(ProtoUtils.marshaller(Flight.PutResult.getDefaultInstance())) .setSchemaDescriptor(FlightServiceGrpc.getDoPutMethod().getSchemaDescriptor()) .build(); } public static MethodDescriptor getDoExchangeDescriptor( BufferAllocator allocator) { return MethodDescriptor.newBuilder() .setType(MethodType.BIDI_STREAMING) .setFullMethodName(DO_EXCHANGE) .setSampledToLocalTracing(false) .setRequestMarshaller(ArrowMessage.createMarshaller(allocator)) .setResponseMarshaller(ArrowMessage.createMarshaller(allocator)) .setSchemaDescriptor(FlightServiceGrpc.getDoExchangeMethod().getSchemaDescriptor()) .build(); } @Override public ServerServiceDefinition bindService() { final ServerServiceDefinition baseDefinition = delegate.bindService(); final MethodDescriptor doGetDescriptor = getDoGetDescriptor(allocator); final MethodDescriptor doPutDescriptor = getDoPutDescriptor(allocator); final MethodDescriptor doExchangeDescriptor = getDoExchangeDescriptor(allocator); // Make sure we preserve SchemaDescriptor fields on methods so that gRPC reflection still works. final ServiceDescriptor.Builder serviceDescriptorBuilder = ServiceDescriptor.newBuilder(FlightConstants.SERVICE) .setSchemaDescriptor(baseDefinition.getServiceDescriptor().getSchemaDescriptor()); serviceDescriptorBuilder.addMethod(doGetDescriptor); serviceDescriptorBuilder.addMethod(doPutDescriptor); serviceDescriptorBuilder.addMethod(doExchangeDescriptor); for (MethodDescriptor definition : baseDefinition.getServiceDescriptor().getMethods()) { if (OVERRIDE_METHODS.contains(definition.getFullMethodName())) { continue; } serviceDescriptorBuilder.addMethod(definition); } final ServiceDescriptor serviceDescriptor = serviceDescriptorBuilder.build(); ServerServiceDefinition.Builder serviceBuilder = ServerServiceDefinition.builder(serviceDescriptor); serviceBuilder.addMethod( doGetDescriptor, ServerCalls.asyncServerStreamingCall(new DoGetMethod(delegate))); serviceBuilder.addMethod( doPutDescriptor, ServerCalls.asyncBidiStreamingCall(new DoPutMethod(delegate))); serviceBuilder.addMethod( doExchangeDescriptor, ServerCalls.asyncBidiStreamingCall(new DoExchangeMethod(delegate))); // copy over not-overridden methods. for (ServerMethodDefinition definition : baseDefinition.getMethods()) { if (OVERRIDE_METHODS.contains(definition.getMethodDescriptor().getFullMethodName())) { continue; } serviceBuilder.addMethod(definition); } return serviceBuilder.build(); } private static class DoGetMethod implements ServerCalls.ServerStreamingMethod { private final FlightService delegate; public DoGetMethod(FlightService delegate) { this.delegate = delegate; } @Override public void invoke(Flight.Ticket request, StreamObserver responseObserver) { delegate.doGetCustom(request, responseObserver); } } private static class DoPutMethod implements ServerCalls.BidiStreamingMethod { private final FlightService delegate; public DoPutMethod(FlightService delegate) { this.delegate = delegate; } @Override public StreamObserver invoke(StreamObserver responseObserver) { return delegate.doPutCustom(responseObserver); } } private static class DoExchangeMethod implements ServerCalls.BidiStreamingMethod { private final FlightService delegate; public DoExchangeMethod(FlightService delegate) { this.delegate = delegate; } @Override public StreamObserver invoke(StreamObserver responseObserver) { return delegate.doExchangeCustom(responseObserver); } } } FlightCallHeaders.java000066400000000000000000000070251515451615200337260ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.google.common.base.Preconditions; import com.google.common.collect.ArrayListMultimap; import com.google.common.collect.Iterables; import com.google.common.collect.Multimap; import io.grpc.Metadata; import java.nio.charset.StandardCharsets; import java.util.Collection; import java.util.Set; import java.util.stream.Collectors; /** An implementation of the Flight headers interface for headers. */ public class FlightCallHeaders implements CallHeaders { private final Multimap keysAndValues; public FlightCallHeaders() { this.keysAndValues = ArrayListMultimap.create(); } @Override public String get(String key) { final Collection values = this.keysAndValues.get(key); if (values.isEmpty()) { return null; } if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) { return new String((byte[]) Iterables.get(values, 0), StandardCharsets.UTF_8); } return (String) Iterables.get(values, 0); } @Override public byte[] getByte(String key) { final Collection values = this.keysAndValues.get(key); if (values.isEmpty()) { return null; } if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) { return (byte[]) Iterables.get(values, 0); } return ((String) Iterables.get(values, 0)).getBytes(StandardCharsets.UTF_8); } @Override public Iterable getAll(String key) { if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) { return this.keysAndValues.get(key).stream() .map(o -> new String((byte[]) o, StandardCharsets.UTF_8)) .collect(Collectors.toList()); } return (Collection) (Collection) this.keysAndValues.get(key); } @Override public Iterable getAllByte(String key) { if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) { return (Collection) (Collection) this.keysAndValues.get(key); } return this.keysAndValues.get(key).stream() .map(o -> ((String) o).getBytes(StandardCharsets.UTF_8)) .collect(Collectors.toList()); } @Override public void insert(String key, String value) { this.keysAndValues.put(key, value); } @Override public void insert(String key, byte[] value) { Preconditions.checkArgument( key.endsWith("-bin"), "Binary header is named %s. It must end with %s", key, "-bin"); Preconditions.checkArgument(key.length() > "-bin".length(), "empty key name"); this.keysAndValues.put(key, value); } @Override public Set keys() { return this.keysAndValues.keySet(); } @Override public boolean containsKey(String key) { return this.keysAndValues.containsKey(key); } @Override public String toString() { return this.keysAndValues.toString(); } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightClient.java000066400000000000000000000745111515451615200330600ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import io.grpc.Channel; import io.grpc.ClientCall; import io.grpc.ClientInterceptor; import io.grpc.ClientInterceptors; import io.grpc.ManagedChannel; import io.grpc.MethodDescriptor; import io.grpc.StatusRuntimeException; 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 java.io.IOException; import java.io.InputStream; import java.net.URISyntaxException; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Iterator; import java.util.List; import java.util.Optional; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import java.util.function.BooleanSupplier; import org.apache.arrow.flight.FlightProducer.StreamListener; import org.apache.arrow.flight.auth.BasicClientAuthHandler; import org.apache.arrow.flight.auth.ClientAuthHandler; import org.apache.arrow.flight.auth.ClientAuthInterceptor; import org.apache.arrow.flight.auth.ClientAuthWrapper; import org.apache.arrow.flight.auth2.BasicAuthCredentialWriter; import org.apache.arrow.flight.auth2.ClientBearerHeaderHandler; import org.apache.arrow.flight.auth2.ClientHandshakeWrapper; 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; import org.apache.arrow.flight.impl.FlightServiceGrpc; import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceBlockingStub; import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceStub; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider; /** Client for Flight services. */ public class FlightClient implements AutoCloseable { private static final int PENDING_REQUESTS = 5; private final BufferAllocator allocator; private final ManagedChannel channel; private final FlightServiceBlockingStub blockingStub; private final FlightServiceStub asyncStub; private final ClientAuthInterceptor authInterceptor = new ClientAuthInterceptor(); private final MethodDescriptor doGetDescriptor; private final MethodDescriptor doPutDescriptor; private final MethodDescriptor doExchangeDescriptor; private final List middleware; /** Create a Flight client from an allocator and a gRPC channel. */ FlightClient( BufferAllocator incomingAllocator, ManagedChannel channel, List middleware) { this.allocator = incomingAllocator.newChildAllocator("flight-client", 0, Long.MAX_VALUE); this.channel = channel; // 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(this.middleware)}; // Create a channel with interceptors pre-applied for DoGet and DoPut Channel interceptedChannel = ClientInterceptors.intercept(channel, interceptors); blockingStub = FlightServiceGrpc.newBlockingStub(interceptedChannel); asyncStub = FlightServiceGrpc.newStub(interceptedChannel); doGetDescriptor = FlightBindingService.getDoGetDescriptor(allocator); doPutDescriptor = FlightBindingService.getDoPutDescriptor(allocator); doExchangeDescriptor = FlightBindingService.getDoExchangeDescriptor(allocator); } /** * Get a list of available flights. * * @param criteria Criteria for selecting flights * @param options RPC-layer hints for the call. * @return FlightInfo Iterable */ public Iterable listFlights(Criteria criteria, CallOption... options) { final Iterator flights; try { flights = CallOptions.wrapStub(blockingStub, options).listFlights(criteria.asCriteria()); } catch (StatusRuntimeException sre) { throw StatusUtils.fromGrpcRuntimeException(sre); } return () -> StatusUtils.wrapIterator( flights, t -> { try { return new FlightInfo(t); } catch (URISyntaxException e) { // We don't expect this will happen for conforming Flight implementations. For // instance, a Java server // itself wouldn't be able to construct an invalid Location. throw new RuntimeException(e); } }); } /** * Lists actions available on the Flight service. * * @param options RPC-layer hints for the call. */ public Iterable listActions(CallOption... options) { final Iterator actions; try { actions = CallOptions.wrapStub(blockingStub, options).listActions(Empty.getDefaultInstance()); } catch (StatusRuntimeException sre) { throw StatusUtils.fromGrpcRuntimeException(sre); } return () -> StatusUtils.wrapIterator(actions, ActionType::new); } /** * Performs an action on the Flight service. * * @param action The action to perform. * @param options RPC-layer hints for this call. * @return An iterator of results. */ public Iterator doAction(Action action, CallOption... options) { return StatusUtils.wrapIterator( CallOptions.wrapStub(blockingStub, options).doAction(action.toProtocol()), Result::new); } /** Authenticates with a username and password. */ public void authenticateBasic(String username, String password) { BasicClientAuthHandler basicClient = new BasicClientAuthHandler(username, password); authenticate(basicClient); } /** * Authenticates against the Flight service. * * @param options RPC-layer hints for this call. * @param handler The auth mechanism to use. */ public void authenticate(ClientAuthHandler handler, CallOption... options) { Preconditions.checkArgument(!authInterceptor.hasAuthHandler(), "Auth already completed."); ClientAuthWrapper.doClientAuth(handler, CallOptions.wrapStub(asyncStub, options)); authInterceptor.setAuthHandler(handler); } /** * Authenticates with a username and password. * * @param username the username. * @param password the password. * @return a CredentialCallOption containing a bearer token if the server emitted one, or empty if * no bearer token was returned. This can be used in subsequent API calls. */ public Optional authenticateBasicToken(String username, String password) { final ClientIncomingAuthHeaderMiddleware.Factory clientAuthMiddleware = new ClientIncomingAuthHeaderMiddleware.Factory(new ClientBearerHeaderHandler()); middleware.add(clientAuthMiddleware); handshake(new CredentialCallOption(new BasicAuthCredentialWriter(username, password))); return Optional.ofNullable(clientAuthMiddleware.getCredentialCallOption()); } /** * Executes the handshake against the Flight service. * * @param options RPC-layer hints for this call. */ public void handshake(CallOption... options) { ClientHandshakeWrapper.doClientHandshake(CallOptions.wrapStub(asyncStub, options)); } /** * Create or append a descriptor with another stream. * * @param descriptor FlightDescriptor the descriptor for the data * @param root VectorSchemaRoot the root containing data * @param metadataListener A handler for metadata messages from the server. This will be passed * buffers that will be freed after {@link StreamListener#onNext(Object)} is called! * @param options RPC-layer hints for this call. * @return ClientStreamListener an interface to control uploading data */ public ClientStreamListener startPut( FlightDescriptor descriptor, VectorSchemaRoot root, PutListener metadataListener, CallOption... options) { return startPut(descriptor, root, new MapDictionaryProvider(), metadataListener, options); } /** * Create or append a descriptor with another stream. * * @param descriptor FlightDescriptor the descriptor for the data * @param root VectorSchemaRoot the root containing data * @param metadataListener A handler for metadata messages from the server. * @param options RPC-layer hints for this call. * @return ClientStreamListener an interface to control uploading data. {@link * ClientStreamListener#start(VectorSchemaRoot, DictionaryProvider)} will already have been * called. */ public ClientStreamListener startPut( FlightDescriptor descriptor, VectorSchemaRoot root, DictionaryProvider provider, PutListener metadataListener, CallOption... options) { Preconditions.checkNotNull(root, "root must not be null"); Preconditions.checkNotNull(provider, "provider must not be null"); final ClientStreamListener writer = startPut(descriptor, metadataListener, options); writer.start(root, provider); return writer; } /** * Create or append a descriptor with another stream. * * @param descriptor FlightDescriptor the descriptor for the data * @param metadataListener A handler for metadata messages from the server. * @param options RPC-layer hints for this call. * @return ClientStreamListener an interface to control uploading data. {@link * ClientStreamListener#start(VectorSchemaRoot, DictionaryProvider)} will NOT already have * been called. */ public ClientStreamListener startPut( FlightDescriptor descriptor, PutListener metadataListener, CallOption... options) { Preconditions.checkNotNull(descriptor, "descriptor must not be null"); Preconditions.checkNotNull(metadataListener, "metadataListener must not be null"); try { final ClientCall call = asyncStubNewCall(doPutDescriptor, options); final SetStreamObserver resultObserver = new SetStreamObserver(allocator, metadataListener); ClientCallStreamObserver observer = (ClientCallStreamObserver) ClientCalls.asyncBidiStreamingCall(call, resultObserver); return new PutObserver( descriptor, observer, metadataListener::isCancelled, metadataListener::getResult); } catch (StatusRuntimeException sre) { throw StatusUtils.fromGrpcRuntimeException(sre); } } /** * Get info on a stream. * * @param descriptor The descriptor for the stream. * @param options RPC-layer hints for this call. */ public FlightInfo getInfo(FlightDescriptor descriptor, CallOption... options) { try { return new FlightInfo( CallOptions.wrapStub(blockingStub, options).getFlightInfo(descriptor.toProtocol())); } catch (URISyntaxException e) { // We don't expect this will happen for conforming Flight implementations. For instance, a // Java server // itself wouldn't be able to construct an invalid Location. throw new RuntimeException(e); } catch (StatusRuntimeException sre) { throw StatusUtils.fromGrpcRuntimeException(sre); } } /** * Start or get info on execution of a long-running query. * * @param descriptor The descriptor for the stream. * @param options RPC-layer hints for this call. * @return Metadata about execution. */ public PollInfo pollInfo(FlightDescriptor descriptor, CallOption... options) { try { return new PollInfo( CallOptions.wrapStub(blockingStub, options).pollFlightInfo(descriptor.toProtocol())); } catch (URISyntaxException e) { throw new RuntimeException(e); } catch (StatusRuntimeException sre) { throw StatusUtils.fromGrpcRuntimeException(sre); } } /** * Get schema for a stream. * * @param descriptor The descriptor for the stream. * @param options RPC-layer hints for this call. */ public SchemaResult getSchema(FlightDescriptor descriptor, CallOption... options) { try { return SchemaResult.fromProtocol( CallOptions.wrapStub(blockingStub, options).getSchema(descriptor.toProtocol())); } catch (StatusRuntimeException sre) { throw StatusUtils.fromGrpcRuntimeException(sre); } } /** * Retrieve a stream from the server. * * @param ticket The ticket granting access to the data stream. * @param options RPC-layer hints for this call. */ public FlightStream getStream(Ticket ticket, CallOption... options) { final ClientCall call = asyncStubNewCall(doGetDescriptor, options); FlightStream stream = new FlightStream( allocator, PENDING_REQUESTS, (String message, Throwable cause) -> call.cancel(message, cause), (count) -> call.request(count)); final StreamObserver delegate = stream.asObserver(); ClientResponseObserver clientResponseObserver = new ClientResponseObserver() { @Override public void beforeStart( ClientCallStreamObserver requestStream) { requestStream.disableAutoInboundFlowControl(); } @Override public void onNext(ArrowMessage value) { delegate.onNext(value); } @Override public void onError(Throwable t) { delegate.onError(StatusUtils.toGrpcException(t)); } @Override public void onCompleted() { delegate.onCompleted(); } }; ClientCalls.asyncServerStreamingCall(call, ticket.toProtocol(), clientResponseObserver); return stream; } /** * Initiate a bidirectional data exchange with the server. * * @param descriptor A descriptor for the data stream. * @param options RPC call options. * @return A pair of a readable stream and a writable stream. */ public ExchangeReaderWriter doExchange(FlightDescriptor descriptor, CallOption... options) { Preconditions.checkNotNull(descriptor, "descriptor must not be null"); try { final ClientCall call = asyncStubNewCall(doExchangeDescriptor, options); final FlightStream stream = new FlightStream(allocator, PENDING_REQUESTS, call::cancel, call::request); final ClientCallStreamObserver observer = (ClientCallStreamObserver) ClientCalls.asyncBidiStreamingCall(call, stream.asObserver()); final ClientStreamListener writer = new PutObserver( descriptor, observer, stream.cancelled::isDone, () -> { try { stream.completed.get(); } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw CallStatus.INTERNAL .withDescription("Client error: interrupted while completing call") .withCause(e) .toRuntimeException(); } catch (ExecutionException e) { throw CallStatus.INTERNAL .withDescription("Client error: internal while completing call") .withCause(e) .toRuntimeException(); } }); // Send the descriptor to start. try (final ArrowMessage message = new ArrowMessage(descriptor.toProtocol())) { observer.onNext(message); } catch (Exception e) { throw CallStatus.INTERNAL .withCause(e) .withDescription("Could not write descriptor " + descriptor) .toRuntimeException(); } return new ExchangeReaderWriter(stream, writer); } catch (StatusRuntimeException sre) { throw StatusUtils.fromGrpcRuntimeException(sre); } } /** A pair of a reader and a writer for a DoExchange call. */ public static class ExchangeReaderWriter implements AutoCloseable { private final FlightStream reader; private final ClientStreamListener writer; ExchangeReaderWriter(FlightStream reader, ClientStreamListener writer) { this.reader = reader; this.writer = writer; } /** Get the reader for the call. */ public FlightStream getReader() { return reader; } /** Get the writer for the call. */ public ClientStreamListener getWriter() { return writer; } /** * Make sure stream is drained. You must call this to be notified of any errors that may have * happened after the exchange is complete. This should be called after * `getWriter().completed()` and instead of `getWriter().getResult()`. */ public void getResult() { // After exchange is complete, make sure stream is drained to propagate errors through reader while (reader.next()) {} } /** Shut down the streams in this call. */ @Override public void close() throws Exception { reader.close(); } } /** A stream observer for Flight.PutResult. */ private static class SetStreamObserver implements StreamObserver { private final BufferAllocator allocator; private final StreamListener listener; SetStreamObserver(BufferAllocator allocator, StreamListener listener) { super(); this.allocator = allocator; this.listener = listener == null ? NoOpStreamListener.getInstance() : listener; } @Override public void onNext(Flight.PutResult value) { try (final PutResult message = PutResult.fromProtocol(allocator, value)) { listener.onNext(message); } } @Override public void onError(Throwable t) { listener.onError(StatusUtils.fromThrowable(t)); } @Override public void onCompleted() { listener.onCompleted(); } } /** The implementation of a {@link ClientStreamListener} for writing data to a Flight server. */ static class PutObserver extends OutboundStreamListenerImpl implements ClientStreamListener { private final BooleanSupplier isCancelled; private final Runnable getResult; /** * Create a new client stream listener. * * @param descriptor The descriptor for the stream. * @param observer The write-side gRPC StreamObserver. * @param isCancelled A flag to check if the call has been cancelled. * @param getResult A flag that blocks until the overall call completes. */ PutObserver( FlightDescriptor descriptor, ClientCallStreamObserver observer, BooleanSupplier isCancelled, Runnable getResult) { super(descriptor, observer); Preconditions.checkNotNull(descriptor, "descriptor must be provided"); Preconditions.checkNotNull(isCancelled, "isCancelled must be provided"); Preconditions.checkNotNull(getResult, "getResult must be provided"); this.isCancelled = isCancelled; this.getResult = getResult; this.unloader = null; } @Override protected void waitUntilStreamReady() { // Check isCancelled as well to avoid inadvertently blocking forever // (so long as PutListener properly implements it) while (!responseObserver.isReady() && !isCancelled.getAsBoolean()) { /* busy wait */ } } @Override public void getResult() { getResult.run(); } } /** * Cancel execution of a distributed query. * * @param request The query to cancel. * @param options Call options. * @return The server response. */ public CancelFlightInfoResult cancelFlightInfo( CancelFlightInfoRequest request, CallOption... options) { Action action = new Action(FlightConstants.CANCEL_FLIGHT_INFO.getType(), request.serialize().array()); Iterator results = doAction(action, options); if (!results.hasNext()) { throw CallStatus.INTERNAL .withDescription("Server did not return a response") .toRuntimeException(); } CancelFlightInfoResult result; try { result = CancelFlightInfoResult.deserialize(ByteBuffer.wrap(results.next().getBody())); } catch (IOException e) { throw CallStatus.INTERNAL .withDescription("Failed to parse server response: " + e) .withCause(e) .toRuntimeException(); } results.forEachRemaining((ignored) -> {}); return result; } /** * Request the server to extend the lifetime of a query result set. * * @param request The result set partition. * @param options Call options. * @return The new endpoint with an updated expiration time. */ public FlightEndpoint renewFlightEndpoint( RenewFlightEndpointRequest request, CallOption... options) { Action action = new Action(FlightConstants.RENEW_FLIGHT_ENDPOINT.getType(), request.serialize().array()); Iterator results = doAction(action, options); if (!results.hasNext()) { throw CallStatus.INTERNAL .withDescription("Server did not return a response") .toRuntimeException(); } FlightEndpoint result; try { result = FlightEndpoint.deserialize(ByteBuffer.wrap(results.next().getBody())); } catch (IOException | URISyntaxException e) { throw CallStatus.INTERNAL .withDescription("Failed to parse server response: " + e) .withCause(e) .toRuntimeException(); } results.forEachRemaining((ignored) -> {}); return result; } /** * Set server session option(s) by name/value. * *

    Sessions are generally persisted via HTTP cookies. * * @param request The session options to set on the server. * @param options Call options. * @return The result containing per-value error statuses, if any. */ public SetSessionOptionsResult setSessionOptions( SetSessionOptionsRequest request, CallOption... options) { Action action = new Action(FlightConstants.SET_SESSION_OPTIONS.getType(), request.serialize().array()); Iterator results = doAction(action, options); if (!results.hasNext()) { throw CallStatus.INTERNAL .withDescription("Server did not return a response") .toRuntimeException(); } SetSessionOptionsResult result; try { result = SetSessionOptionsResult.deserialize(ByteBuffer.wrap(results.next().getBody())); } catch (IOException e) { throw CallStatus.INTERNAL .withDescription("Failed to parse server response: " + e) .withCause(e) .toRuntimeException(); } results.forEachRemaining((ignored) -> {}); return result; } /** * Get the current server session options. * *

    The session is generally accessed via an HTTP cookie. * * @param request The (empty) GetSessionOptionsRequest. * @param options Call options. * @return The result containing the set of session options configured on the server. */ public GetSessionOptionsResult getSessionOptions( GetSessionOptionsRequest request, CallOption... options) { Action action = new Action(FlightConstants.GET_SESSION_OPTIONS.getType(), request.serialize().array()); Iterator results = doAction(action, options); if (!results.hasNext()) { throw CallStatus.INTERNAL .withDescription("Server did not return a response") .toRuntimeException(); } GetSessionOptionsResult result; try { result = GetSessionOptionsResult.deserialize(ByteBuffer.wrap(results.next().getBody())); } catch (IOException e) { throw CallStatus.INTERNAL .withDescription("Failed to parse server response: " + e) .withCause(e) .toRuntimeException(); } results.forEachRemaining((ignored) -> {}); return result; } /** * Close/invalidate the current server session. * *

    The session is generally accessed via an HTTP cookie. * * @param request The (empty) CloseSessionRequest. * @param options Call options. * @return The result containing the status of the close operation. */ public CloseSessionResult closeSession(CloseSessionRequest request, CallOption... options) { Action action = new Action(FlightConstants.CLOSE_SESSION.getType(), request.serialize().array()); Iterator results = doAction(action, options); if (!results.hasNext()) { throw CallStatus.INTERNAL .withDescription("Server did not return a response") .toRuntimeException(); } CloseSessionResult result; try { result = CloseSessionResult.deserialize(ByteBuffer.wrap(results.next().getBody())); } catch (IOException e) { throw CallStatus.INTERNAL .withDescription("Failed to parse server response: " + e) .withCause(e) .toRuntimeException(); } results.forEachRemaining((ignored) -> {}); return result; } /** Interface for writers to an Arrow data stream. */ public interface ClientStreamListener extends OutboundStreamListener { /** * Wait for the stream to finish on the server side. You must call this to be notified of any * errors that may have happened during the upload. */ void getResult(); } /** * A handler for server-sent application metadata messages during a Flight DoPut operation. * *

    Generally, instead of implementing this yourself, you should use {@link AsyncPutListener} or * {@link SyncPutListener}. */ public interface PutListener extends StreamListener { /** * Wait for the stream to finish on the server side. You must call this to be notified of any * errors that may have happened during the upload. */ void getResult(); /** * Called when a message from the server is received. * * @param val The application metadata. This buffer will be reclaimed once onNext returns; you * must retain a reference to use it outside this method. */ @Override void onNext(PutResult val); /** * Check if the call has been cancelled. * *

    By default, this always returns false. Implementations should provide an appropriate * implementation, as otherwise, a DoPut operation may inadvertently block forever. */ default boolean isCancelled() { return false; } } /** Shut down this client. */ @Override public void close() throws InterruptedException { channel.shutdown().awaitTermination(5, TimeUnit.SECONDS); allocator.close(); } /** Create a builder for a Flight client. */ public static Builder builder() { return new Builder(); } /** * Create a builder for a Flight client. * * @param allocator The allocator to use for the client. * @param location The location to connect to. */ public static Builder builder(BufferAllocator allocator, Location location) { return new Builder(allocator, location); } /** A builder for Flight clients. */ public static final class Builder { private final NettyClientBuilder builder; private Builder() { this.builder = new NettyClientBuilder(); } private Builder(BufferAllocator allocator, Location location) { this.builder = new NettyClientBuilder(allocator, location); } /** Force the client to connect over TLS. */ public Builder useTls() { builder.useTls(); return this; } /** Override the hostname checked for TLS. Use with caution in production. */ public Builder overrideHostname(final String hostname) { builder.overrideHostname(hostname); return this; } /** Set the maximum inbound message size. */ public Builder maxInboundMessageSize(int maxSize) { builder.maxInboundMessageSize(maxSize); return this; } /** Set the trusted TLS certificates. */ public Builder trustedCertificates(final InputStream stream) { builder.trustedCertificates(stream); return this; } /** Set the trusted TLS certificates. */ public Builder clientCertificate( final InputStream clientCertificate, final InputStream clientKey) { builder.clientCertificate(clientCertificate, clientKey); return this; } public Builder allocator(BufferAllocator allocator) { builder.allocator(allocator); return this; } public Builder location(Location location) { builder.location(location); return this; } public Builder intercept(FlightClientMiddleware.Factory factory) { builder.intercept(factory); return this; } public Builder verifyServer(boolean verifyServer) { builder.verifyServer(verifyServer); return this; } /** Create the client from this builder. */ public FlightClient build() { final NettyChannelBuilder channelBuilder = builder.build(); return new FlightClient(builder.allocator(), channelBuilder.build(), builder.middleware()); } } /** * Helper method to create a call from the asyncStub, method descriptor, and list of calling * options. */ private ClientCall asyncStubNewCall( MethodDescriptor descriptor, CallOption... options) { FlightServiceStub wrappedStub = CallOptions.wrapStub(asyncStub, options); return wrappedStub.getChannel().newCall(descriptor, wrappedStub.getCallOptions()); } } FlightClientMiddleware.java000066400000000000000000000033321515451615200347700ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * Client-side middleware for Flight. * *

    Middleware are instantiated per-call and should store state in the middleware instance. */ public interface FlightClientMiddleware { /** A callback used before request headers are sent. The headers may be manipulated. */ void onBeforeSendingHeaders(CallHeaders outgoingHeaders); /** A callback called after response headers are received. The headers may be manipulated. */ void onHeadersReceived(CallHeaders incomingHeaders); /** A callback called after the call completes. */ void onCallCompleted(CallStatus status); /** A factory for client middleware instances. */ interface Factory { /** * Create a new middleware instance for the given call. * * @throws FlightRuntimeException if the middleware wants to reject the call with the given * status */ FlightClientMiddleware onCallStarted(CallInfo info); } } FlightConstants.java000066400000000000000000000047021515451615200335320ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** String constants relevant to flight implementations. */ public interface FlightConstants { String SERVICE = "arrow.flight.protocol.FlightService"; FlightServerMiddleware.Key HEADER_KEY = FlightServerMiddleware.Key.of("org.apache.arrow.flight.ServerHeaderMiddleware"); ActionType CANCEL_FLIGHT_INFO = new ActionType( "CancelFlightInfo", "Explicitly cancel a running FlightInfo.\n" + "Request Message: CancelFlightInfoRequest\n" + "Response Message: CancelFlightInfoResult"); ActionType RENEW_FLIGHT_ENDPOINT = new ActionType( "RenewFlightEndpoint", "Extend expiration time of the given FlightEndpoint.\n" + "Request Message: RenewFlightEndpointRequest\n" + "Response Message: Renewed FlightEndpoint"); ActionType SET_SESSION_OPTIONS = new ActionType( "SetSessionOptions", "Set client session options by name/value pairs.\n" + "Request Message: SetSessionOptionsRequest\n" + "Response Message: SetSessionOptionsResult"); ActionType GET_SESSION_OPTIONS = new ActionType( "GetSessionOptions", "Get current client session options\n" + "Request Message: GetSessionOptionsRequest\n" + "Response Message: GetSessionOptionsResult"); ActionType CLOSE_SESSION = new ActionType( "CloseSession", "Explicitly close/invalidate the cookie-specified client session.\n" + "Request Message: CloseSessionRequest\n" + "Response Message: CloseSessionResult"); } FlightDescriptor.java000066400000000000000000000121321515451615200336700ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.google.common.base.Joiner; import com.google.common.collect.ImmutableList; import com.google.protobuf.ByteString; import java.io.IOException; import java.nio.ByteBuffer; import java.util.Arrays; import java.util.List; import org.apache.arrow.flight.impl.Flight; import org.apache.arrow.flight.impl.Flight.FlightDescriptor.DescriptorType; import org.apache.arrow.util.Preconditions; /** * An identifier for a particular set of data. This can either be an opaque command that generates * the data or a static "path" to the data. This is a POJO wrapper around the protobuf message with * the same name. */ public class FlightDescriptor { private boolean isCmd; private List path; private byte[] cmd; private FlightDescriptor(boolean isCmd, List path, byte[] cmd) { super(); this.isCmd = isCmd; this.path = path; this.cmd = cmd; } public static FlightDescriptor command(byte[] cmd) { return new FlightDescriptor(true, null, cmd); } public static FlightDescriptor path(Iterable path) { return new FlightDescriptor(false, ImmutableList.copyOf(path), null); } public static FlightDescriptor path(String... path) { return new FlightDescriptor(false, ImmutableList.copyOf(path), null); } FlightDescriptor(Flight.FlightDescriptor descriptor) { if (descriptor.getType() == DescriptorType.CMD) { isCmd = true; cmd = descriptor.getCmd().toByteArray(); } else if (descriptor.getType() == DescriptorType.PATH) { isCmd = false; path = descriptor.getPathList(); } else { throw new UnsupportedOperationException(); } } public boolean isCommand() { return isCmd; } public List getPath() { Preconditions.checkArgument(!isCmd); return path; } public byte[] getCommand() { Preconditions.checkArgument(isCmd); return cmd; } Flight.FlightDescriptor toProtocol() { Flight.FlightDescriptor.Builder b = Flight.FlightDescriptor.newBuilder(); if (isCmd) { return b.setType(DescriptorType.CMD).setCmd(ByteString.copyFrom(cmd)).build(); } return b.setType(DescriptorType.PATH).addAllPath(path).build(); } /** * Get the serialized form of this protocol message. * *

    Intended to help interoperability by allowing non-Flight services to still return Flight * types. */ public ByteBuffer serialize() { return ByteBuffer.wrap(toProtocol().toByteArray()); } /** * Parse the serialized form of this protocol message. * *

    Intended to help interoperability by allowing Flight clients to obtain stream info from * non-Flight services. * * @param serialized The serialized form of the FlightDescriptor, as returned by {@link * #serialize()}. * @return The deserialized FlightDescriptor. * @throws IOException if the serialized form is invalid. */ public static FlightDescriptor deserialize(ByteBuffer serialized) throws IOException { return new FlightDescriptor(Flight.FlightDescriptor.parseFrom(serialized)); } @Override public String toString() { if (isCmd) { return toHex(cmd); } else { return Joiner.on('.').join(path); } } private String toHex(byte[] bytes) { StringBuilder sb = new StringBuilder(); for (byte b : bytes) { sb.append(String.format("%02X ", b)); } return sb.toString(); } @Override public int hashCode() { final int prime = 31; int result = 1; result = prime * result + ((cmd == null) ? 0 : Arrays.hashCode(cmd)); result = prime * result + (isCmd ? 1231 : 1237); result = prime * result + ((path == null) ? 0 : path.hashCode()); return result; } @Override public boolean equals(Object obj) { if (this == obj) { return true; } if (obj == null) { return false; } if (!(obj instanceof FlightDescriptor)) { return false; } FlightDescriptor other = (FlightDescriptor) obj; if (cmd == null) { if (other.cmd != null) { return false; } } else if (!Arrays.equals(cmd, other.cmd)) { return false; } if (isCmd != other.isCmd) { return false; } if (path == null) { if (other.path != null) { return false; } } else if (!path.equals(other.path)) { return false; } return true; } } FlightEndpoint.java000066400000000000000000000172641515451615200333450ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.google.protobuf.ByteString; import com.google.protobuf.Timestamp; import com.google.protobuf.util.Timestamps; import java.io.IOException; import java.net.URISyntaxException; import java.nio.ByteBuffer; import java.time.Instant; import java.util.ArrayList; import java.util.Arrays; import java.util.Base64; import java.util.Collections; import java.util.List; import java.util.Objects; import java.util.Optional; import org.apache.arrow.flight.impl.Flight; /** POJO to convert to/from the underlying protobuf FlightEndpoint. */ public class FlightEndpoint { private final List locations; private final Ticket ticket; private final Instant expirationTime; private final byte[] appMetadata; /** * Constructs a new endpoint with no expiration time. * * @param ticket A ticket that describe the key of a data stream. * @param locations The possible locations the stream can be retrieved from. */ public FlightEndpoint(Ticket ticket, Location... locations) { this(ticket, /*expirationTime*/ null, locations); } /** * Constructs a new endpoint with an expiration time. * * @param ticket A ticket that describe the key of a data stream. * @param expirationTime (optional) When this endpoint expires. * @param locations The possible locations the stream can be retrieved from. */ public FlightEndpoint(Ticket ticket, Instant expirationTime, Location... locations) { this( ticket, expirationTime, null, Collections.unmodifiableList(new ArrayList<>(Arrays.asList(locations)))); } /** Private constructor with all parameters. Should only be called by Builder. */ private FlightEndpoint( Ticket ticket, Instant expirationTime, byte[] appMetadata, List locations) { Objects.requireNonNull(ticket); this.locations = locations; this.expirationTime = expirationTime; this.ticket = ticket; this.appMetadata = appMetadata; } /** Constructs from the protocol buffer representation. */ FlightEndpoint(Flight.FlightEndpoint flt) throws URISyntaxException { this.locations = new ArrayList<>(); for (final Flight.Location location : flt.getLocationList()) { this.locations.add(new Location(location.getUri())); } if (flt.hasExpirationTime()) { this.expirationTime = Instant.ofEpochSecond( flt.getExpirationTime().getSeconds(), Timestamps.toNanos(flt.getExpirationTime())); } else { this.expirationTime = null; } this.appMetadata = (flt.getAppMetadata().isEmpty() ? null : flt.getAppMetadata().toByteArray()); this.ticket = new Ticket(flt.getTicket()); } public List getLocations() { return locations; } public Ticket getTicket() { return ticket; } public Optional getExpirationTime() { return Optional.ofNullable(expirationTime); } public byte[] getAppMetadata() { return appMetadata; } /** Converts to the protocol buffer representation. */ Flight.FlightEndpoint toProtocol() { Flight.FlightEndpoint.Builder b = Flight.FlightEndpoint.newBuilder().setTicket(ticket.toProtocol()); for (Location l : locations) { b.addLocation(l.toProtocol()); } if (expirationTime != null) { b.setExpirationTime( Timestamp.newBuilder() .setSeconds(expirationTime.getEpochSecond()) .setNanos(expirationTime.getNano()) .build()); } if (appMetadata != null) { b.setAppMetadata(ByteString.copyFrom(appMetadata)); } return b.build(); } /** * Get the serialized form of this protocol message. * *

    Intended to help interoperability by allowing non-Flight services to still return Flight * types. */ public ByteBuffer serialize() { return ByteBuffer.wrap(toProtocol().toByteArray()); } /** * Parse the serialized form of this protocol message. * *

    Intended to help interoperability by allowing Flight clients to obtain stream info from * non-Flight services. * * @param serialized The serialized form of the message, as returned by {@link #serialize()}. * @return The deserialized message. * @throws IOException if the serialized form is invalid. * @throws URISyntaxException if the serialized form contains an unsupported URI format. */ public static FlightEndpoint deserialize(ByteBuffer serialized) throws IOException, URISyntaxException { return new FlightEndpoint(Flight.FlightEndpoint.parseFrom(serialized)); } @Override public boolean equals(Object o) { if (this == o) { return true; } if (!(o instanceof FlightEndpoint)) { return false; } FlightEndpoint that = (FlightEndpoint) o; return locations.equals(that.locations) && ticket.equals(that.ticket) && Objects.equals(expirationTime, that.expirationTime) && Arrays.equals(appMetadata, that.appMetadata); } @Override public int hashCode() { return Objects.hash(locations, ticket, expirationTime, Arrays.hashCode(appMetadata)); } @Override public String toString() { return "FlightEndpoint{" + "locations=" + locations + ", ticket=" + ticket + ", expirationTime=" + (expirationTime == null ? "(none)" : expirationTime.toString()) + ", appMetadata=" + (appMetadata == null ? "(none)" : Base64.getEncoder().encodeToString(appMetadata)) + '}'; } /** * Create a builder for FlightEndpoint. * * @param ticket A ticket that describe the key of a data stream. * @param locations The possible locations the stream can be retrieved from. */ public static Builder builder(Ticket ticket, Location... locations) { return new Builder(ticket, locations); } /** Builder for FlightEndpoint. */ public static final class Builder { private final Ticket ticket; private final List locations; private Instant expirationTime = null; private byte[] appMetadata = null; private Builder(Ticket ticket, Location... locations) { this.ticket = ticket; this.locations = Collections.unmodifiableList(new ArrayList<>(Arrays.asList(locations))); } /** * Set expiration time for the endpoint. Default is null, which means don't expire. * * @param expirationTime (optional) When this endpoint expires. */ public Builder setExpirationTime(Instant expirationTime) { this.expirationTime = expirationTime; return this; } /** * Set the app metadata to send along with the flight. Default is null; * * @param appMetadata Metadata to send along with the flight */ public Builder setAppMetadata(byte[] appMetadata) { this.appMetadata = appMetadata; return this; } /** Build FlightEndpoint object. */ public FlightEndpoint build() { return new FlightEndpoint(ticket, expirationTime, appMetadata, locations); } } } FlightGrpcUtils.java000066400000000000000000000132241515451615200334710ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import io.grpc.BindableService; import io.grpc.CallOptions; import io.grpc.ClientCall; import io.grpc.ConnectivityState; 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; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.VisibleForTesting; /** Exposes Flight GRPC service & client. */ public class FlightGrpcUtils { /** Proxy class for ManagedChannel that makes closure a no-op. */ @VisibleForTesting static class NonClosingProxyManagedChannel extends ManagedChannel { private final ManagedChannel channel; private boolean isShutdown; NonClosingProxyManagedChannel(ManagedChannel channel) { this.channel = channel; this.isShutdown = channel.isShutdown(); } @Override public ManagedChannel shutdown() { isShutdown = true; return this; } @Override public boolean isShutdown() { if (this.channel.isShutdown()) { // If the underlying channel is shut down, ensure we're updated to match. shutdown(); } return isShutdown; } @Override public boolean isTerminated() { return this.isShutdown(); } @Override public ManagedChannel shutdownNow() { return shutdown(); } @Override public boolean awaitTermination(long l, TimeUnit timeUnit) { // Don't actually await termination, since it'll be a no-op, so simply return whether or not // the channel has been shut down already. return this.isShutdown(); } @Override public ClientCall newCall( MethodDescriptor methodDescriptor, CallOptions callOptions) { if (this.isShutdown()) { throw new IllegalStateException("Channel has been shut down."); } return this.channel.newCall(methodDescriptor, callOptions); } @Override public String authority() { return this.channel.authority(); } @Override public ConnectivityState getState(boolean requestConnection) { if (this.isShutdown()) { return ConnectivityState.SHUTDOWN; } return this.channel.getState(requestConnection); } @Override public void notifyWhenStateChanged(ConnectivityState source, Runnable callback) { // The proxy has no insight into the underlying channel state changes, so we'll have to leak // the abstraction // a bit here and simply pass to the underlying channel, even though it will never transition // to shutdown via // the proxy. This should be fine, since it's mainly targeted at the FlightClient and there's // no getter for // the channel. this.channel.notifyWhenStateChanged(source, callback); } @Override public void resetConnectBackoff() { this.channel.resetConnectBackoff(); } @Override public void enterIdle() { this.channel.enterIdle(); } } private FlightGrpcUtils() {} /** * Creates a Flight service. * * @param allocator Memory allocator * @param producer Specifies the service api * @param authHandler Authentication handler * @param executor Executor service * @return FlightBindingService */ public static BindableService createFlightService( BufferAllocator allocator, FlightProducer producer, ServerAuthHandler authHandler, ExecutorService executor) { return new FlightBindingService(allocator, producer, authHandler, executor); } /** * 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) { 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. * * @param incomingAllocator Memory allocator * @param channel provides a connection to a gRPC server. Will not be closed on closure of the * returned FlightClient. */ public static FlightClient createFlightClientWithSharedChannel( BufferAllocator incomingAllocator, ManagedChannel channel) { return new FlightClient( incomingAllocator, new NonClosingProxyManagedChannel(channel), Collections.emptyList()); } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightInfo.java000066400000000000000000000304211515451615200325250ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.fasterxml.jackson.databind.util.ByteBufferBackedInputStream; import com.google.protobuf.ByteString; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.net.URISyntaxException; import java.nio.ByteBuffer; import java.nio.channels.Channels; import java.util.ArrayList; import java.util.Arrays; import java.util.Base64; import java.util.Collections; import java.util.List; import java.util.Objects; import java.util.Optional; import java.util.stream.Collectors; import org.apache.arrow.flight.impl.Flight; import org.apache.arrow.vector.ipc.ReadChannel; import org.apache.arrow.vector.ipc.WriteChannel; import org.apache.arrow.vector.ipc.message.IpcOption; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.validate.MetadataV4UnionChecker; /** A POJO representation of a FlightInfo, metadata associated with a set of data records. */ public class FlightInfo { private final Schema schema; private final FlightDescriptor descriptor; private final List endpoints; private final long bytes; private final long records; private final boolean ordered; private final IpcOption option; private final byte[] appMetadata; /** * Constructs a new instance. * * @param schema The schema of the Flight * @param descriptor An identifier for the Flight. * @param endpoints A list of endpoints that have the flight available. * @param bytes The number of bytes in the flight * @param records The number of records in the flight. */ public FlightInfo( Schema schema, FlightDescriptor descriptor, List endpoints, long bytes, long records) { this(schema, descriptor, endpoints, bytes, records, /*ordered*/ false, IpcOption.DEFAULT); } /** * Constructs a new instance. * * @param schema The schema of the Flight * @param descriptor An identifier for the Flight. * @param endpoints A list of endpoints that have the flight available. * @param bytes The number of bytes in the flight * @param records The number of records in the flight. * @param option IPC write options. */ public FlightInfo( Schema schema, FlightDescriptor descriptor, List endpoints, long bytes, long records, IpcOption option) { this(schema, descriptor, endpoints, bytes, records, /*ordered*/ false, option); } /** * Constructs a new instance. * * @param schema The schema of the Flight * @param descriptor An identifier for the Flight. * @param endpoints A list of endpoints that have the flight available. * @param bytes The number of bytes in the flight * @param records The number of records in the flight. * @param ordered Whether the endpoints in this flight are ordered. * @param option IPC write options. */ public FlightInfo( Schema schema, FlightDescriptor descriptor, List endpoints, long bytes, long records, boolean ordered, IpcOption option) { this(schema, descriptor, endpoints, bytes, records, ordered, option, null); } /** * Constructs a new instance. * * @param schema The schema of the Flight * @param descriptor An identifier for the Flight. * @param endpoints A list of endpoints that have the flight available. * @param bytes The number of bytes in the flight * @param records The number of records in the flight. * @param ordered Whether the endpoints in this flight are ordered. * @param option IPC write options. * @param appMetadata Metadata to send along with the flight */ public FlightInfo( Schema schema, FlightDescriptor descriptor, List endpoints, long bytes, long records, boolean ordered, IpcOption option, byte[] appMetadata) { Objects.requireNonNull(descriptor); Objects.requireNonNull(endpoints); if (schema != null) { MetadataV4UnionChecker.checkForUnion(schema.getFields().iterator(), option.metadataVersion); } this.schema = schema; this.descriptor = descriptor; this.endpoints = endpoints; this.bytes = bytes; this.records = records; this.ordered = ordered; this.option = option; this.appMetadata = appMetadata; } /** Constructs from the protocol buffer representation. */ FlightInfo(Flight.FlightInfo pbFlightInfo) throws URISyntaxException { try { final ByteBuffer schemaBuf = pbFlightInfo.getSchema().asReadOnlyByteBuffer(); schema = pbFlightInfo.getSchema().size() > 0 ? MessageSerializer.deserializeSchema( new ReadChannel(Channels.newChannel(new ByteBufferBackedInputStream(schemaBuf)))) : null; } catch (IOException e) { throw new RuntimeException(e); } descriptor = new FlightDescriptor(pbFlightInfo.getFlightDescriptor()); endpoints = new ArrayList<>(); for (final Flight.FlightEndpoint endpoint : pbFlightInfo.getEndpointList()) { endpoints.add(new FlightEndpoint(endpoint)); } bytes = pbFlightInfo.getTotalBytes(); records = pbFlightInfo.getTotalRecords(); ordered = pbFlightInfo.getOrdered(); appMetadata = (pbFlightInfo.getAppMetadata().size() == 0 ? null : pbFlightInfo.getAppMetadata().toByteArray()); option = IpcOption.DEFAULT; } public Optional getSchemaOptional() { return Optional.ofNullable(schema); } /** * Returns the schema, or an empty schema if no schema is present. * * @deprecated Deprecated. Use {@link #getSchemaOptional()} instead. */ @Deprecated public Schema getSchema() { return schema != null ? schema : new Schema(Collections.emptyList()); } public long getBytes() { return bytes; } public long getRecords() { return records; } public FlightDescriptor getDescriptor() { return descriptor; } public List getEndpoints() { return endpoints; } public boolean getOrdered() { return ordered; } public byte[] getAppMetadata() { return appMetadata; } /** Converts to the protocol buffer representation. */ Flight.FlightInfo toProtocol() { Flight.FlightInfo.Builder builder = Flight.FlightInfo.newBuilder() .addAllEndpoint( endpoints.stream().map(t -> t.toProtocol()).collect(Collectors.toList())) .setFlightDescriptor(descriptor.toProtocol()) .setTotalBytes(FlightInfo.this.bytes) .setTotalRecords(records) .setOrdered(ordered); if (schema != null) { ByteArrayOutputStream baos = new ByteArrayOutputStream(); try { MessageSerializer.serialize(new WriteChannel(Channels.newChannel(baos)), schema, option); builder.setSchema(ByteString.copyFrom(baos.toByteArray())); } catch (IOException e) { throw new RuntimeException(e); } } if (appMetadata != null) { builder.setAppMetadata(ByteString.copyFrom(appMetadata)); } return builder.build(); } /** * Get the serialized form of this protocol message. * *

    Intended to help interoperability by allowing non-Flight services to still return Flight * types. */ public ByteBuffer serialize() { return ByteBuffer.wrap(toProtocol().toByteArray()); } /** * Parse the serialized form of this protocol message. * *

    Intended to help interoperability by allowing Flight clients to obtain stream info from * non-Flight services. * * @param serialized The serialized form of the FlightInfo, as returned by {@link #serialize()}. * @return The deserialized FlightInfo. * @throws IOException if the serialized form is invalid. * @throws URISyntaxException if the serialized form contains an unsupported URI format. */ public static FlightInfo deserialize(ByteBuffer serialized) throws IOException, URISyntaxException { return new FlightInfo(Flight.FlightInfo.parseFrom(serialized)); } @Override public boolean equals(Object o) { if (this == o) { return true; } if (!(o instanceof FlightInfo)) { return false; } FlightInfo that = (FlightInfo) o; return bytes == that.bytes && records == that.records && schema.equals(that.schema) && descriptor.equals(that.descriptor) && endpoints.equals(that.endpoints) && ordered == that.ordered && Arrays.equals(appMetadata, that.appMetadata); } @Override public int hashCode() { return Objects.hash( schema, descriptor, endpoints, bytes, records, ordered, Arrays.hashCode(appMetadata)); } @Override public String toString() { return "FlightInfo{" + "schema=" + schema + ", descriptor=" + descriptor + ", endpoints=" + endpoints + ", bytes=" + bytes + ", records=" + records + ", ordered=" + ordered + ", appMetadata=" + (appMetadata == null ? "(none)" : Base64.getEncoder().encodeToString(appMetadata)) + '}'; } /** * Create a builder for FlightInfo. * * @param schema The schema of the Flight * @param descriptor An identifier for the Flight. * @param endpoints A list of endpoints that have the flight available. */ public static Builder builder( Schema schema, FlightDescriptor descriptor, List endpoints) { return new Builder(schema, descriptor, endpoints); } /** Builder for FlightInfo. */ public static final class Builder { private final Schema schema; private final FlightDescriptor descriptor; private final List endpoints; private long bytes = -1; private long records = -1; private boolean ordered = false; private IpcOption option = IpcOption.DEFAULT; private byte[] appMetadata = null; private Builder(Schema schema, FlightDescriptor descriptor, List endpoints) { this.schema = schema; this.descriptor = descriptor; this.endpoints = endpoints; } /** * Set the number of bytes for the flight. Default to -1 for unknown. * * @param bytes The number of bytes in the flight */ public Builder setBytes(long bytes) { this.bytes = bytes; return this; } /** * Set the number of records for the flight. Default to -1 for unknown. * * @param records The number of records in the flight. */ public Builder setRecords(long records) { this.records = records; return this; } /** * Set whether the flight endpoints are ordered. Default is false. * * @param ordered Whether the endpoints in this flight are ordered. */ public Builder setOrdered(boolean ordered) { this.ordered = ordered; return this; } /** * Set IPC write options. Default is IpcOption.DEFAULT * * @param option IPC write options. */ public Builder setOption(IpcOption option) { this.option = option; return this; } /** * Set the app metadata to send along with the flight. Default is null. * * @param appMetadata Metadata to send along with the flight */ public Builder setAppMetadata(byte[] appMetadata) { this.appMetadata = appMetadata; return this; } /** Build FlightInfo object. */ public FlightInfo build() { return new FlightInfo( schema, descriptor, endpoints, bytes, records, ordered, option, appMetadata); } } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightMethod.java000066400000000000000000000051361515451615200330570ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import org.apache.arrow.flight.impl.FlightServiceGrpc; /** All the RPC methods available in Flight. */ public enum FlightMethod { HANDSHAKE, LIST_FLIGHTS, GET_FLIGHT_INFO, GET_SCHEMA, DO_GET, DO_PUT, DO_ACTION, LIST_ACTIONS, DO_EXCHANGE, POLL_FLIGHT_INFO, ; /** * Convert a method name string into a {@link FlightMethod}. * * @throws IllegalArgumentException if the method name is not valid. */ public static FlightMethod fromProtocol(final String methodName) { if (FlightServiceGrpc.getHandshakeMethod().getFullMethodName().equals(methodName)) { return HANDSHAKE; } else if (FlightServiceGrpc.getListFlightsMethod().getFullMethodName().equals(methodName)) { return LIST_FLIGHTS; } else if (FlightServiceGrpc.getGetFlightInfoMethod().getFullMethodName().equals(methodName)) { return GET_FLIGHT_INFO; } else if (FlightServiceGrpc.getGetSchemaMethod().getFullMethodName().equals(methodName)) { return GET_SCHEMA; } else if (FlightServiceGrpc.getDoGetMethod().getFullMethodName().equals(methodName)) { return DO_GET; } else if (FlightServiceGrpc.getDoPutMethod().getFullMethodName().equals(methodName)) { return DO_PUT; } else if (FlightServiceGrpc.getDoActionMethod().getFullMethodName().equals(methodName)) { return DO_ACTION; } else if (FlightServiceGrpc.getListActionsMethod().getFullMethodName().equals(methodName)) { return LIST_ACTIONS; } else if (FlightServiceGrpc.getDoExchangeMethod().getFullMethodName().equals(methodName)) { return DO_EXCHANGE; } else if (FlightServiceGrpc.getPollFlightInfoMethod().getFullMethodName().equals(methodName)) { return POLL_FLIGHT_INFO; } throw new IllegalArgumentException("Not a Flight method name in gRPC: " + methodName); } } FlightProducer.java000066400000000000000000000143601515451615200333420ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.util.Map; /** API to Implement an Arrow Flight producer. */ public interface FlightProducer { /** * Return data for a stream. * * @param context Per-call context. * @param ticket The application-defined ticket identifying this stream. * @param listener An interface for sending data back to the client. */ void getStream(CallContext context, Ticket ticket, ServerStreamListener listener); /** * List available data streams on this service. * * @param context Per-call context. * @param criteria Application-defined criteria for filtering streams. * @param listener An interface for sending data back to the client. */ void listFlights(CallContext context, Criteria criteria, StreamListener listener); /** * Get information about a particular data stream. * * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return Metadata about the stream. */ FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor); /** * Begin or get an update on execution of a long-running query. * *

    If the descriptor would begin a query, the server should return a response immediately to * not block the client. Otherwise, the server should not return an update until progress is made * to not spam the client with inactionable updates. * * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return Metadata about execution. */ default PollInfo pollFlightInfo(CallContext context, FlightDescriptor descriptor) { FlightInfo info = getFlightInfo(context, descriptor); return new PollInfo(info, null, null, null); } /** * Get schema for a particular data stream. * * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return Schema for the stream. */ default SchemaResult getSchema(CallContext context, FlightDescriptor descriptor) { FlightInfo info = getFlightInfo(context, descriptor); return new SchemaResult( info.getSchemaOptional() .orElseThrow( () -> CallStatus.INVALID_ARGUMENT .withDescription("No schema is present in FlightInfo") .toRuntimeException())); } /** * Accept uploaded data for a particular stream. * * @param context Per-call context. * @param flightStream The data stream being uploaded. */ Runnable acceptPut( CallContext context, FlightStream flightStream, StreamListener ackStream); /** * This method is used to perform a bidirectional data exchange between a client and a server. * * @param context Per-call context. * @param reader The FlightStream from which data is read. * @param writer The ServerStreamListener to which data is written. * @throws RuntimeException if the method is not implemented. */ default void doExchange(CallContext context, FlightStream reader, ServerStreamListener writer) { throw CallStatus.UNIMPLEMENTED .withDescription("DoExchange is unimplemented") .toRuntimeException(); } /** * Generic handler for application-defined RPCs. * * @param context Per-call context. * @param action Client-supplied parameters. * @param listener A stream of responses. */ void doAction(CallContext context, Action action, StreamListener listener); /** * List available application-defined RPCs. * * @param context Per-call context. * @param listener An interface for sending data back to the client. */ void listActions(CallContext context, StreamListener listener); /** An interface for sending Arrow data back to a client. */ interface ServerStreamListener extends OutboundStreamListener { /** Check whether the call has been cancelled. If so, stop sending data. */ boolean isCancelled(); /** * Set a callback for when the client cancels a call, i.e. {@link #isCancelled()} has become * true. * *

    Note that this callback may only be called some time after {@link #isCancelled()} becomes * true, and may never be called if all executor threads on the server are busy, or the RPC * method body is implemented in a blocking fashion. */ void setOnCancelHandler(Runnable handler); } /** * Callbacks for pushing objects to a receiver. * * @param Type of the values in the stream. */ interface StreamListener { /** Send the next value to the client. */ void onNext(T val); /** * Indicate an error to the client. * *

    Terminates the stream; do not call {@link #onCompleted()}. */ void onError(Throwable t); /** Indicate that the transmission is finished. */ void onCompleted(); } /** Call-specific context. */ interface CallContext { /** The identity of the authenticated peer. May be the empty string if unknown. */ String peerIdentity(); /** Whether the call has been cancelled by the client. */ boolean isCancelled(); /** * Get the middleware instance of the given type for this call. * *

    Returns null if not found. */ T getMiddleware(FlightServerMiddleware.Key key); /** Get an immutable map of middleware for this call. */ Map, FlightServerMiddleware> getMiddleware(); } } FlightRuntimeException.java000066400000000000000000000027701515451615200350630ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * An exception raised from a Flight RPC. * *

    In service implementations, raising an instance of this exception will provide clients with a * more detailed message and error code. */ public class FlightRuntimeException extends RuntimeException { private final CallStatus status; /** Create a new exception from the given status. */ FlightRuntimeException(CallStatus status) { super(status.description(), status.cause()); this.status = status; } public CallStatus status() { return status; } @Override public String toString() { String s = getClass().getName(); return String.format("%s: %s: %s", s, status.code(), status.description()); } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightServer.java000066400000000000000000000463411515451615200331100ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.google.common.util.concurrent.ThreadFactoryBuilder; import io.grpc.Server; import io.grpc.ServerInterceptors; import io.grpc.netty.GrpcSslContexts; import io.grpc.netty.NettyServerBuilder; import io.netty.channel.EventLoopGroup; import io.netty.channel.ServerChannel; import io.netty.handler.ssl.ClientAuth; import io.netty.handler.ssl.SslContext; import io.netty.handler.ssl.SslContextBuilder; import java.io.File; import java.io.FileInputStream; import java.io.IOException; import java.io.InputStream; import java.lang.reflect.InvocationTargetException; import java.net.URI; import java.net.URISyntaxException; import java.util.ArrayList; import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Set; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.function.Consumer; import javax.net.ssl.SSLException; import org.apache.arrow.flight.auth.ServerAuthHandler; import org.apache.arrow.flight.auth.ServerAuthInterceptor; import org.apache.arrow.flight.auth2.Auth2Constants; import org.apache.arrow.flight.auth2.CallHeaderAuthenticator; import org.apache.arrow.flight.auth2.ServerCallHeaderAuthMiddleware; import org.apache.arrow.flight.grpc.ServerBackpressureThresholdInterceptor; import org.apache.arrow.flight.grpc.ServerInterceptorAdapter; import org.apache.arrow.flight.grpc.ServerInterceptorAdapter.KeyFactory; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.util.VisibleForTesting; /** * Generic server of flight data that is customized via construction with delegate classes for the * actual logic. The server currently uses GRPC as its transport mechanism. */ public class FlightServer implements AutoCloseable { private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(FlightServer.class); private final Location location; private final Server server; // The executor used by the gRPC server. We don't use it here, but we do need to clean it up with // the server. // May be null, if a user-supplied executor was provided (as we do not want to clean that up) @VisibleForTesting final ExecutorService grpcExecutor; /** The maximum size of an individual gRPC message. This effectively disables the limit. */ static final int MAX_GRPC_MESSAGE_SIZE = Integer.MAX_VALUE; /** The default number of bytes that can be queued on an output stream before blocking. */ public static final int DEFAULT_BACKPRESSURE_THRESHOLD = 10 * 1024 * 1024; // 10MB /** Create a new instance from a gRPC server. For internal use only. */ private FlightServer(Location location, Server server, ExecutorService grpcExecutor) { this.location = location; this.server = server; this.grpcExecutor = grpcExecutor; } /** Start the server. */ public FlightServer start() throws IOException { server.start(); return this; } /** Get the port the server is running on (if applicable). */ public int getPort() { return server.getPort(); } /** Get the location for this server. */ public Location getLocation() { if (location.getUri().getPort() == 0) { // If the server was bound to port 0, replace the port in the location with the real port. final URI uri = location.getUri(); try { return new Location( new URI( uri.getScheme(), uri.getUserInfo(), uri.getHost(), getPort(), uri.getPath(), uri.getQuery(), uri.getFragment())); } catch (URISyntaxException e) { // We don't expect this to happen throw new RuntimeException(e); } } return location; } /** Block until the server shuts down. */ public void awaitTermination() throws InterruptedException { server.awaitTermination(); } /** Request that the server shut down. */ public void shutdown() { server.shutdown(); if (grpcExecutor != null) { grpcExecutor.shutdown(); } } /** * Wait for the server to shut down with a timeout. * * @return true if the server shut down successfully. */ public boolean awaitTermination(final long timeout, final TimeUnit unit) throws InterruptedException { return server.awaitTermination(timeout, unit); } /** Shutdown the server, waits for up to 6 seconds for successful shutdown before returning. */ @Override public void close() throws InterruptedException { shutdown(); final boolean terminated = awaitTermination(3000, TimeUnit.MILLISECONDS); if (terminated) { logger.debug("Server was terminated within 3s"); return; } // get more aggressive in termination. server.shutdownNow(); int count = 0; while (!server.isTerminated() && count < 30) { count++; logger.debug("Waiting for termination"); Thread.sleep(100); } if (!server.isTerminated()) { logger.warn("Couldn't shutdown server, resources likely will be leaked."); } } /** Create a builder for a Flight server. */ public static Builder builder() { return new Builder(); } /** Create a builder for a Flight server. */ public static Builder builder( BufferAllocator allocator, Location location, FlightProducer producer) { return new Builder(allocator, location, producer); } /** A builder for Flight servers. */ public static final class Builder { private BufferAllocator allocator; private Location location; private FlightProducer producer; private final Map builderOptions; private ServerAuthHandler authHandler = ServerAuthHandler.NO_OP; private CallHeaderAuthenticator headerAuthenticator = CallHeaderAuthenticator.NO_OP; private ExecutorService executor = null; private int maxInboundMessageSize = MAX_GRPC_MESSAGE_SIZE; private int maxHeaderListSize = MAX_GRPC_MESSAGE_SIZE; private int backpressureThreshold = DEFAULT_BACKPRESSURE_THRESHOLD; private InputStream certChain; private InputStream key; private InputStream mTlsCACert; private SslContext sslContext; private final List> interceptors; // Keep track of inserted interceptors private final Set interceptorKeys; Builder() { builderOptions = new HashMap<>(); interceptors = new ArrayList<>(); interceptorKeys = new HashSet<>(); } Builder(BufferAllocator allocator, Location location, FlightProducer producer) { this(); this.allocator = Preconditions.checkNotNull(allocator); this.location = Preconditions.checkNotNull(location); this.producer = Preconditions.checkNotNull(producer); } /** Create the server for this builder. */ public FlightServer build() { // Add the auth middleware if applicable. if (headerAuthenticator != CallHeaderAuthenticator.NO_OP) { this.middleware( FlightServerMiddleware.Key.of(Auth2Constants.AUTHORIZATION_HEADER), new ServerCallHeaderAuthMiddleware.Factory(headerAuthenticator)); } this.middleware(FlightConstants.HEADER_KEY, new ServerHeaderMiddleware.Factory()); final NettyServerBuilder builder; switch (location.getUri().getScheme()) { case LocationSchemes.GRPC_DOMAIN_SOCKET: { // The implementation is platform-specific, so we have to find the classes at runtime builder = NettyServerBuilder.forAddress(location.toSocketAddress()); try { try { // Linux builder.channelType( Class.forName("io.netty.channel.epoll.EpollServerDomainSocketChannel") .asSubclass(ServerChannel.class)); final EventLoopGroup elg = Class.forName("io.netty.channel.epoll.EpollEventLoopGroup") .asSubclass(EventLoopGroup.class) .getConstructor() .newInstance(); builder.bossEventLoopGroup(elg).workerEventLoopGroup(elg); } catch (ClassNotFoundException e) { // BSD builder.channelType( Class.forName("io.netty.channel.kqueue.KQueueServerDomainSocketChannel") .asSubclass(ServerChannel.class)); final EventLoopGroup elg = Class.forName("io.netty.channel.kqueue.KQueueEventLoopGroup") .asSubclass(EventLoopGroup.class) .getConstructor() .newInstance(); builder.bossEventLoopGroup(elg).workerEventLoopGroup(elg); } } catch (ClassNotFoundException | InstantiationException | IllegalAccessException | NoSuchMethodException | InvocationTargetException e) { throw new UnsupportedOperationException( "Could not find suitable Netty native transport implementation for domain socket address."); } break; } case LocationSchemes.GRPC: case LocationSchemes.GRPC_INSECURE: { builder = NettyServerBuilder.forAddress(location.toSocketAddress()); break; } case LocationSchemes.GRPC_TLS: { if (certChain == null) { throw new IllegalArgumentException( "Must provide a certificate and key to serve gRPC over TLS"); } builder = NettyServerBuilder.forAddress(location.toSocketAddress()); break; } default: throw new IllegalArgumentException( "Scheme is not supported: " + location.getUri().getScheme()); } if (certChain != null) { SslContextBuilder sslContextBuilder = GrpcSslContexts.forServer(certChain, key); if (mTlsCACert != null) { sslContextBuilder.clientAuth(ClientAuth.REQUIRE).trustManager(mTlsCACert); } try { sslContext = sslContextBuilder.build(); } catch (SSLException e) { throw new RuntimeException(e); } finally { closeMTlsCACert(); closeCertChain(); closeKey(); } builder.sslContext(sslContext); } // Share one executor between the gRPC service, DoPut, and Handshake final ExecutorService exec; // We only want to have FlightServer close the gRPC executor if we created it here. We should // not close // user-supplied executors. final ExecutorService grpcExecutor; if (executor != null) { exec = executor; grpcExecutor = null; } else { exec = Executors.newCachedThreadPool( // Name threads for better debuggability new ThreadFactoryBuilder() .setNameFormat("flight-server-default-executor-%d") .build()); grpcExecutor = exec; } final FlightBindingService flightService = new FlightBindingService(allocator, producer, authHandler, exec); builder .executor(exec) .maxInboundMessageSize(maxInboundMessageSize) .maxInboundMetadataSize(maxHeaderListSize) .addService( ServerInterceptors.intercept( flightService, new ServerBackpressureThresholdInterceptor(backpressureThreshold), new ServerAuthInterceptor(authHandler))); // Allow hooking into the gRPC builder. This is not guaranteed to be available on all Arrow // versions or // Flight implementations. builderOptions.computeIfPresent( "grpc.builderConsumer", (key, builderConsumer) -> { final Consumer consumer = (Consumer) builderConsumer; consumer.accept(builder); return null; }); // Allow explicitly setting some Netty-specific options builderOptions.computeIfPresent( "netty.channelType", (key, channelType) -> { builder.channelType((Class) channelType); return null; }); builderOptions.computeIfPresent( "netty.bossEventLoopGroup", (key, elg) -> { builder.bossEventLoopGroup((EventLoopGroup) elg); return null; }); builderOptions.computeIfPresent( "netty.workerEventLoopGroup", (key, elg) -> { builder.workerEventLoopGroup((EventLoopGroup) elg); return null; }); builder.intercept(new ServerInterceptorAdapter(interceptors)); return new FlightServer(location, builder.build(), grpcExecutor); } public Builder setMaxHeaderListSize(int maxHeaderListSize) { this.maxHeaderListSize = maxHeaderListSize; return this; } /** * Set the maximum size of a message. Defaults to "unlimited", depending on the underlying * transport. */ public Builder maxInboundMessageSize(int maxMessageSize) { this.maxInboundMessageSize = maxMessageSize; return this; } /** * Set the number of bytes that may be queued on a server output stream before writes are * blocked. */ public Builder backpressureThreshold(int backpressureThreshold) { Preconditions.checkArgument(backpressureThreshold > 0); this.backpressureThreshold = backpressureThreshold; return this; } /** * A small utility function to ensure that InputStream attributes. are closed if they are not * null * * @param stream The InputStream to close (if it is not null). */ private void closeInputStreamIfNotNull(InputStream stream) { if (stream != null) { try { stream.close(); } catch (IOException expected) { // stream closes gracefully, doesn't expect an exception. } } } /** * A small utility function to ensure that the certChain attribute is closed if it is not null. * It then sets the attribute to null. */ private void closeCertChain() { closeInputStreamIfNotNull(certChain); certChain = null; } /** * A small utility function to ensure that the key attribute is closed if it is not null. It * then sets the attribute to null. */ private void closeKey() { closeInputStreamIfNotNull(key); key = null; } /** * A small utility function to ensure that the mTlsCACert attribute is closed if it is not null. * It then sets the attribute to null. */ private void closeMTlsCACert() { closeInputStreamIfNotNull(mTlsCACert); mTlsCACert = null; } /** * Enable TLS on the server. * * @param certChain The certificate chain to use. * @param key The private key to use. */ public Builder useTls(final File certChain, final File key) throws IOException { closeCertChain(); this.certChain = new FileInputStream(certChain); closeKey(); this.key = new FileInputStream(key); return this; } /** * Enable Client Verification via mTLS on the server. * * @param mTlsCACert The CA certificate to use for verifying clients. */ public Builder useMTlsClientVerification(final File mTlsCACert) throws IOException { closeMTlsCACert(); this.mTlsCACert = new FileInputStream(mTlsCACert); return this; } /** * Enable TLS on the server. * * @param certChain The certificate chain to use. * @param key The private key to use. */ public Builder useTls(final InputStream certChain, final InputStream key) throws IOException { closeCertChain(); this.certChain = certChain; closeKey(); this.key = key; return this; } /** * Enable mTLS on the server. * * @param mTlsCACert The CA certificate to use for verifying clients. */ public Builder useMTlsClientVerification(final InputStream mTlsCACert) throws IOException { closeMTlsCACert(); this.mTlsCACert = mTlsCACert; return this; } /** * Set the executor used by the server. * *

    Flight will NOT take ownership of the executor. The application must clean it up if one is * provided. (If not provided, Flight will use a default executor which it will clean up.) */ public Builder executor(ExecutorService executor) { this.executor = executor; return this; } /** Set the authentication handler. */ public Builder authHandler(ServerAuthHandler authHandler) { this.authHandler = authHandler; return this; } /** Set the header-based authentication mechanism. */ public Builder headerAuthenticator(CallHeaderAuthenticator headerAuthenticator) { this.headerAuthenticator = headerAuthenticator; return this; } /** Provide a transport-specific option. Not guaranteed to have any effect. */ public Builder transportHint(final String key, Object option) { builderOptions.put(key, option); return this; } /** * Add a Flight middleware component to inspect and modify requests to this service. * * @param key An identifier for this middleware component. Service implementations can retrieve * the middleware instance for the current call using {@link * org.apache.arrow.flight.FlightProducer.CallContext}. * @param factory A factory for the middleware. * @param The middleware type. * @throws IllegalArgumentException if the key already exists */ public Builder middleware( final FlightServerMiddleware.Key key, final FlightServerMiddleware.Factory factory) { if (interceptorKeys.contains(key.key)) { throw new IllegalArgumentException("Key already exists: " + key.key); } interceptors.add(new KeyFactory<>(key, factory)); interceptorKeys.add(key.key); return this; } public Builder allocator(BufferAllocator allocator) { this.allocator = Preconditions.checkNotNull(allocator); return this; } public Builder location(Location location) { this.location = Preconditions.checkNotNull(location); return this; } public Builder producer(FlightProducer producer) { this.producer = Preconditions.checkNotNull(producer); return this; } } } FlightServerMiddleware.java000066400000000000000000000076621515451615200350320ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.util.Objects; /** * Server-side middleware for Flight calls. * *

    Middleware are instantiated per-call. * *

    Methods are not guaranteed to be called on any particular thread, relative to the thread that * Flight requests are executed on. Do not depend on thread-local storage; instead, use state on the * middleware instance. Service implementations may communicate with middleware implementations * through {@link org.apache.arrow.flight.FlightProducer.CallContext#getMiddleware(Key)}. Methods on * the middleware instance are non-reentrant, that is, a particular RPC will not make multiple * concurrent calls to methods on a single middleware instance. However, methods on the factory * instance are expected to be thread-safe, and if the factory instance returns the same middleware * object more than once, then that middleware object must be thread-safe. */ public interface FlightServerMiddleware { /** * A factory for Flight server middleware. * * @param The middleware type. */ interface Factory { /** * A callback for when the call starts. * * @param info Details about the call. * @param incomingHeaders A mutable set of request headers. * @param context Context about the current request. * @throws FlightRuntimeException if the middleware wants to reject the call with the given * status */ T onCallStarted(CallInfo info, CallHeaders incomingHeaders, RequestContext context); } /** * A key for Flight server middleware. On a server, middleware instances are identified by this * key. * *

    Keys use reference equality, so instances should be shared. * * @param The middleware class stored in this key. This provides a compile-time check when * retrieving instances. */ class Key { final String key; Key(String key) { this.key = Objects.requireNonNull(key, "Key must not be null."); } /** Create a new key for the given type. */ public static Key of(String key) { return new Key<>(key); } } /** * Callback for when the underlying transport is about to send response headers. * * @param outgoingHeaders A mutable set of response headers. These can be manipulated to send * different headers to the client. */ void onBeforeSendingHeaders(CallHeaders outgoingHeaders); /** * Callback for when the underlying transport has completed a call. * * @param status Whether the call completed successfully or not. */ void onCallCompleted(CallStatus status); /** * Callback for when an RPC method implementation throws an uncaught exception. * *

    May be called multiple times, and may be called before or after {@link * #onCallCompleted(CallStatus)}. Generally, an uncaught exception will end the call with a error * {@link CallStatus}, and will be reported to {@link #onCallCompleted(CallStatus)}, but not * necessarily this method. * * @param err The exception that was thrown. */ void onCallErrored(Throwable err); } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightService.java000066400000000000000000000407171515451615200332430ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.google.common.base.Strings; import io.grpc.stub.ServerCallStreamObserver; import io.grpc.stub.StreamObserver; import java.util.Collections; import java.util.Map; import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; import java.util.function.BooleanSupplier; import java.util.function.Consumer; import org.apache.arrow.flight.FlightProducer.ServerStreamListener; import org.apache.arrow.flight.auth.AuthConstants; import org.apache.arrow.flight.auth.ServerAuthHandler; import org.apache.arrow.flight.auth.ServerAuthWrapper; import org.apache.arrow.flight.auth2.Auth2Constants; import org.apache.arrow.flight.grpc.ContextPropagatingExecutorService; import org.apache.arrow.flight.grpc.RequestContextAdapter; import org.apache.arrow.flight.grpc.ServerInterceptorAdapter; import org.apache.arrow.flight.grpc.StatusUtils; import org.apache.arrow.flight.impl.Flight; import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceImplBase; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.AutoCloseables; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** GRPC service implementation for a flight server. */ class FlightService extends FlightServiceImplBase { private static final Logger logger = LoggerFactory.getLogger(FlightService.class); private static final int PENDING_REQUESTS = 5; private final BufferAllocator allocator; private final FlightProducer producer; private final ServerAuthHandler authHandler; private final ExecutorService executors; FlightService( BufferAllocator allocator, FlightProducer producer, ServerAuthHandler authHandler, ExecutorService executors) { this.allocator = allocator; this.producer = producer; this.authHandler = authHandler; this.executors = new ContextPropagatingExecutorService(executors); } private CallContext makeContext(ServerCallStreamObserver responseObserver) { // Try to get the peer identity from middleware first (using the auth2 interfaces). final RequestContext context = RequestContextAdapter.REQUEST_CONTEXT_KEY.get(); String peerIdentity = null; if (context != null) { peerIdentity = context.get(Auth2Constants.PEER_IDENTITY_KEY); } if (Strings.isNullOrEmpty(peerIdentity)) { // Try the legacy auth interface, which defaults to empty string. peerIdentity = AuthConstants.PEER_IDENTITY_KEY.get(); } return new CallContext(peerIdentity, responseObserver::isCancelled); } @Override public StreamObserver handshake( StreamObserver responseObserver) { // This method is not meaningful with the auth2 interfaces. Authentication would already // have happened by header/middleware with the auth2 classes. return ServerAuthWrapper.wrapHandshake(authHandler, responseObserver, executors); } @Override public void listFlights( Flight.Criteria criteria, StreamObserver responseObserver) { final StreamPipe listener = StreamPipe.wrap( responseObserver, FlightInfo::toProtocol, this::handleExceptionWithMiddleware); try { final CallContext context = makeContext((ServerCallStreamObserver) responseObserver); producer.listFlights(context, new Criteria(criteria), listener); } catch (Exception ex) { listener.onError(ex); } // Do NOT call StreamPipe#onCompleted, as the FlightProducer implementation may be asynchronous } public void doGetCustom( Flight.Ticket ticket, StreamObserver responseObserverSimple) { final ServerCallStreamObserver responseObserver = (ServerCallStreamObserver) responseObserverSimple; final GetListener listener = new GetListener(responseObserver, this::handleExceptionWithMiddleware); try { producer.getStream(makeContext(responseObserver), new Ticket(ticket), listener); } catch (Exception ex) { listener.error(ex); } // Do NOT call GetListener#completed, as the implementation of getStream may be asynchronous } @Override public void doAction(Flight.Action request, StreamObserver responseObserver) { final StreamPipe listener = StreamPipe.wrap(responseObserver, Result::toProtocol, this::handleExceptionWithMiddleware); try { final CallContext context = makeContext((ServerCallStreamObserver) responseObserver); producer.doAction(context, new Action(request), listener); } catch (Exception ex) { listener.onError(ex); } // Do NOT call StreamPipe#onCompleted, as the FlightProducer implementation may be asynchronous } @Override public void listActions( Flight.Empty request, StreamObserver responseObserver) { final StreamPipe listener = StreamPipe.wrap( responseObserver, ActionType::toProtocol, this::handleExceptionWithMiddleware); try { final CallContext context = makeContext((ServerCallStreamObserver) responseObserver); producer.listActions(context, listener); } catch (Exception ex) { listener.onError(ex); } // Do NOT call StreamPipe#onCompleted, as the FlightProducer implementation may be asynchronous } private static class GetListener extends OutboundStreamListenerImpl implements ServerStreamListener { private final ServerCallStreamObserver serverCallResponseObserver; private final Consumer errorHandler; private Runnable onCancelHandler = null; private Runnable onReadyHandler = null; private boolean completed; public GetListener( ServerCallStreamObserver responseObserver, Consumer errorHandler) { super(null, responseObserver); this.errorHandler = errorHandler; this.completed = false; this.serverCallResponseObserver = responseObserver; this.serverCallResponseObserver.setOnCancelHandler(this::onCancel); this.serverCallResponseObserver.setOnReadyHandler(this::onReady); this.serverCallResponseObserver.disableAutoInboundFlowControl(); } private void onCancel() { logger.debug("Stream cancelled by client."); if (onCancelHandler != null) { onCancelHandler.run(); } } private void onReady() { if (onReadyHandler != null) { onReadyHandler.run(); } } @Override public void setOnCancelHandler(Runnable handler) { this.onCancelHandler = handler; } @Override public void setOnReadyHandler(Runnable handler) { this.onReadyHandler = handler; } @Override public boolean isCancelled() { return serverCallResponseObserver.isCancelled(); } @Override protected void waitUntilStreamReady() { // Don't do anything - service implementations are expected to manage backpressure themselves } @Override public void error(Throwable ex) { if (!completed) { completed = true; super.error(ex); } else { errorHandler.accept(ex); } } @Override public void completed() { if (!completed) { completed = true; super.completed(); } else { errorHandler.accept(new IllegalStateException("Tried to complete already-completed call")); } } } public StreamObserver doPutCustom( final StreamObserver responseObserverSimple) { ServerCallStreamObserver responseObserver = (ServerCallStreamObserver) responseObserverSimple; responseObserver.disableAutoInboundFlowControl(); responseObserver.request(1); final StreamPipe ackStream = StreamPipe.wrap( responseObserver, PutResult::toProtocol, this::handleExceptionWithMiddleware); final FlightStream fs = new FlightStream( allocator, PENDING_REQUESTS, /* server-upload streams are not cancellable */ null, responseObserver::request); // When the ackStream is completed, the FlightStream will be closed with it ackStream.setAutoCloseable(fs); final StreamObserver observer = fs.asObserver(); Future unused = executors.submit( () -> { try { producer.acceptPut(makeContext(responseObserver), fs, ackStream).run(); } catch (Throwable ex) { ackStream.onError(ex); } finally { // ARROW-6136: Close the stream if and only if acceptPut hasn't closed it itself // We don't do this for other streams since the implementation may be asynchronous ackStream.ensureCompleted(); } }); return observer; } @Override public void getFlightInfo( Flight.FlightDescriptor request, StreamObserver responseObserver) { final FlightInfo info; try { info = producer.getFlightInfo( makeContext((ServerCallStreamObserver) responseObserver), new FlightDescriptor(request)); } catch (Exception ex) { // Don't capture exceptions from onNext or onCompleted with this block - because then we can't // call onError responseObserver.onError(StatusUtils.toGrpcException(ex)); return; } responseObserver.onNext(info.toProtocol()); responseObserver.onCompleted(); } @Override public void pollFlightInfo( Flight.FlightDescriptor request, StreamObserver responseObserver) { final PollInfo info; try { info = producer.pollFlightInfo( makeContext((ServerCallStreamObserver) responseObserver), new FlightDescriptor(request)); } catch (Exception ex) { // Don't capture exceptions from onNext or onCompleted with this block - because then we can't // call onError responseObserver.onError(StatusUtils.toGrpcException(ex)); return; } responseObserver.onNext(info.toProtocol()); responseObserver.onCompleted(); } /** Broadcast the given exception to all registered middleware. */ private void handleExceptionWithMiddleware(Throwable t) { final Map, FlightServerMiddleware> middleware = ServerInterceptorAdapter.SERVER_MIDDLEWARE_KEY.get(); if (middleware == null || middleware.isEmpty()) { logger.error("Uncaught exception in Flight method body", t); return; } middleware.forEach((k, v) -> v.onCallErrored(t)); } @Override public void getSchema( Flight.FlightDescriptor request, StreamObserver responseObserver) { try { SchemaResult result = producer.getSchema( makeContext((ServerCallStreamObserver) responseObserver), new FlightDescriptor(request)); responseObserver.onNext(result.toProtocol()); responseObserver.onCompleted(); } catch (Exception ex) { responseObserver.onError(StatusUtils.toGrpcException(ex)); } } /** Ensures that other resources are cleaned up when the service finishes its call. */ private static class ExchangeListener extends GetListener { private AutoCloseable resource; private boolean closed = false; private Runnable onCancelHandler = null; public ExchangeListener( ServerCallStreamObserver responseObserver, Consumer errorHandler) { super(responseObserver, errorHandler); this.resource = null; super.setOnCancelHandler( () -> { try { if (onCancelHandler != null) { onCancelHandler.run(); } } finally { cleanup(); } }); } private void cleanup() { if (closed) { // Prevent double-free. gRPC will call the OnCancelHandler even on a normal call end, which // means that // we'll double-free without this guard. return; } closed = true; try { AutoCloseables.close(resource); } catch (Exception e) { throw CallStatus.INTERNAL .withCause(e) .withDescription("Server internal error cleaning up resources") .toRuntimeException(); } } @Override public void error(Throwable ex) { try { this.cleanup(); } finally { super.error(ex); } } @Override public void completed() { try { this.cleanup(); } finally { super.completed(); } } @Override public void setOnCancelHandler(Runnable handler) { onCancelHandler = handler; } } public StreamObserver doExchangeCustom( StreamObserver responseObserverSimple) { final ServerCallStreamObserver responseObserver = (ServerCallStreamObserver) responseObserverSimple; final ExchangeListener listener = new ExchangeListener(responseObserver, this::handleExceptionWithMiddleware); final FlightStream fs = new FlightStream( allocator, PENDING_REQUESTS, /* server-upload streams are not cancellable */ null, responseObserver::request); // When service completes the call, this cleans up the FlightStream listener.resource = fs; responseObserver.disableAutoInboundFlowControl(); responseObserver.request(1); final StreamObserver observer = fs.asObserver(); try { Future unused = executors.submit( () -> { try { producer.doExchange(makeContext(responseObserver), fs, listener); } catch (Exception ex) { listener.error(ex); } // We do not clean up or close anything here, to allow long-running asynchronous // implementations. // It is the service's responsibility to call completed() or error(), which will // then clean up the FlightStream. }); } catch (Exception ex) { listener.error(ex); } return observer; } /** Call context for the service. */ static class CallContext implements FlightProducer.CallContext { private final String peerIdentity; private final BooleanSupplier isCancelled; CallContext(final String peerIdentity, BooleanSupplier isCancelled) { this.peerIdentity = peerIdentity; this.isCancelled = isCancelled; } @Override public String peerIdentity() { return peerIdentity; } @Override public boolean isCancelled() { return this.isCancelled.getAsBoolean(); } @Override public T getMiddleware(FlightServerMiddleware.Key key) { final Map, FlightServerMiddleware> middleware = ServerInterceptorAdapter.SERVER_MIDDLEWARE_KEY.get(); if (middleware == null) { return null; } final FlightServerMiddleware m = middleware.get(key); if (m == null) { return null; } @SuppressWarnings("unchecked") final T result = (T) m; return result; } @Override public Map, FlightServerMiddleware> getMiddleware() { final Map, FlightServerMiddleware> middleware = ServerInterceptorAdapter.SERVER_MIDDLEWARE_KEY.get(); if (middleware == null) { return Collections.emptyMap(); } // This is an unmodifiable map return middleware; } } } FlightStatusCode.java000066400000000000000000000045241515451615200336360ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** A status code describing the result of a Flight call. */ public enum FlightStatusCode { /** The call completed successfully. Generally clients will not see this, but middleware may. */ OK, /** * An unknown error occurred. This may also be the result of an implementation error on the * server-side; by default, unhandled server exceptions result in this code. */ UNKNOWN, /** An internal/implementation error occurred. */ INTERNAL, /** One or more of the given arguments was invalid. */ INVALID_ARGUMENT, /** The operation timed out. */ TIMED_OUT, /** The operation describes a resource that does not exist. */ NOT_FOUND, /** The operation creates a resource that already exists. */ ALREADY_EXISTS, /** The operation was cancelled. */ CANCELLED, /** The client was not authenticated. */ UNAUTHENTICATED, /** The client did not have permission to make the call. */ UNAUTHORIZED, /** The requested operation is not implemented. */ UNIMPLEMENTED, /** * The server cannot currently handle the request. This should be used for retriable requests, * i.e. the server should send this code only if it has not done any work. */ UNAVAILABLE, /** * Some resource has been exhausted, perhaps a per-user quota, or perhaps the entire file system * is out of space. (see: https://grpc.github.io/grpc/core/md_doc_statuscodes.html) */ RESOURCE_EXHAUSTED; /** Create a blank {@link CallStatus} with this code. */ public CallStatus toStatus() { return new CallStatus(this); } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java000066400000000000000000000463721515451615200331010ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.google.common.util.concurrent.SettableFuture; import io.grpc.stub.StreamObserver; import java.io.IOException; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.atomic.AtomicInteger; import org.apache.arrow.flight.ArrowMessage.HeaderType; import org.apache.arrow.flight.grpc.StatusUtils; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.VisibleForTesting; import org.apache.arrow.vector.FieldVector; 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.message.ArrowDictionaryBatch; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.types.MetadataVersion; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.DictionaryUtility; import org.apache.arrow.vector.validate.MetadataV4UnionChecker; /** An adaptor between protobuf streams and flight data streams. */ public class FlightStream implements AutoCloseable { // Use AutoCloseable sentinel objects to simplify logic in #close private final AutoCloseable DONE = new AutoCloseable() { @Override public void close() throws Exception {} }; private final AutoCloseable DONE_EX = new AutoCloseable() { @Override public void close() throws Exception {} }; private final BufferAllocator allocator; private final Cancellable cancellable; private final LinkedBlockingQueue queue = new LinkedBlockingQueue<>(); private final SettableFuture root = SettableFuture.create(); private final SettableFuture descriptor = SettableFuture.create(); private final int pendingTarget; private final Requestor requestor; // The completion flags. // This flag is only updated as the user iterates through the data, i.e. it tracks whether the // user has read all the // data and closed the stream final CompletableFuture completed; // This flag is immediately updated when gRPC signals that the server has ended the call. This is // used to make sure // we don't block forever trying to write to a server that has rejected a call. final CompletableFuture cancelled; private final AtomicInteger pending = new AtomicInteger(); private volatile VectorSchemaRoot fulfilledRoot; private DictionaryProvider.MapDictionaryProvider dictionaries; private volatile VectorLoader loader; private volatile Throwable ex; private volatile ArrowBuf applicationMetadata = null; @VisibleForTesting volatile MetadataVersion metadataVersion = null; /** * Constructs a new instance. * * @param allocator The allocator to use for creating/reallocating buffers for Vectors. * @param pendingTarget Target number of messages to receive. * @param cancellable Used to cancel mid-stream requests. * @param requestor A callback to determine how many pending items there are. */ public FlightStream( BufferAllocator allocator, int pendingTarget, Cancellable cancellable, Requestor requestor) { Objects.requireNonNull(allocator); Objects.requireNonNull(requestor); this.allocator = allocator; this.pendingTarget = pendingTarget; this.cancellable = cancellable; this.requestor = requestor; this.dictionaries = new DictionaryProvider.MapDictionaryProvider(); this.completed = new CompletableFuture<>(); this.cancelled = new CompletableFuture<>(); } /** Get the schema for this stream. Blocks until the schema is available. */ public Schema getSchema() { return getRoot().getSchema(); } /** * Get the provider for dictionaries in this stream. * *

    Does NOT retain a reference to the underlying dictionaries. Dictionaries may be updated as * the stream is read. This method is intended for stream processing, where the application code * will not retain references to values after the stream is closed. * * @throws IllegalStateException if {@link #takeDictionaryOwnership()} was called * @see #takeDictionaryOwnership() */ public DictionaryProvider getDictionaryProvider() { if (dictionaries == null) { throw new IllegalStateException("Dictionary ownership was claimed by the application."); } return dictionaries; } /** * Get an owned reference to the dictionaries in this stream. Should be called after finishing * reading the stream, but before closing. * *

    If called, the client is responsible for closing the dictionaries in this provider. Can only * be called once. * * @return The dictionary provider for the stream. * @throws IllegalStateException if called more than once. */ public DictionaryProvider takeDictionaryOwnership() { if (dictionaries == null) { throw new IllegalStateException("Dictionary ownership was claimed by the application."); } // Swap out the provider so it is not closed final DictionaryProvider provider = dictionaries; dictionaries = null; return provider; } /** * Get the descriptor for this stream. Only applicable on the server side of a DoPut operation. * Will block until the client sends the descriptor. */ public FlightDescriptor getDescriptor() { // This blocks until the first message from the client is received. try { return descriptor.get(); } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw CallStatus.INTERNAL.withCause(e).withDescription("Interrupted").toRuntimeException(); } catch (ExecutionException e) { throw CallStatus.INTERNAL .withCause(e) .withDescription("Error getting descriptor") .toRuntimeException(); } } /** * Closes the stream (freeing any existing resources). * *

    If the stream isn't complete and is cancellable, this method will cancel and drain the * stream first. */ @Override public void close() throws Exception { final List closeables = new ArrayList<>(); Throwable suppressor = null; if (cancellable != null) { // Client-side stream. Cancel the call, to help ensure gRPC doesn't deliver a message after // close() ends. // On the server side, we can't rely on draining the stream , because this gRPC bug means the // completion callback // may never run https://github.com/grpc/grpc-java/issues/5882 try { synchronized (cancellable) { if (!cancelled.isDone()) { // Only cancel if the call is not done on the gRPC side cancellable.cancel("Stream closed before end", /* no exception to report */ null); } } // Drain the stream without the lock (as next() implicitly needs the lock) while (next()) {} } catch (FlightRuntimeException e) { suppressor = e; } } // Perform these operations under a lock. This way the observer can't enqueue new messages while // we're in the // middle of cleanup. This should only be a concern for server-side streams since client-side // streams are drained // by the lambda above. synchronized (completed) { try { if (fulfilledRoot != null) { closeables.add(fulfilledRoot); } closeables.add(applicationMetadata); closeables.addAll(queue); if (dictionaries != null) { dictionaries .getDictionaryIds() .forEach(id -> closeables.add(dictionaries.lookup(id).getVector())); } if (suppressor != null) { AutoCloseables.close(suppressor, closeables); } else { AutoCloseables.close(closeables); } // Remove any metadata after closing to prevent negative refcnt applicationMetadata = null; } finally { // The value of this CompletableFuture is meaningless, only whether it's completed (or has // an exception) // No-op if already complete completed.complete(null); } } } /** * Blocking request to load next item into list. * * @return Whether or not more data was found. */ public boolean next() { try { if (completed.isDone() && queue.isEmpty()) { return false; } pending.decrementAndGet(); requestOutstanding(); Object data = queue.take(); if (DONE == data) { queue.put(DONE); // Other code ignores the value of this CompletableFuture, only whether it's completed (or // has an exception) completed.complete(null); return false; } else if (DONE_EX == data) { queue.put(DONE_EX); if (ex instanceof Exception) { throw (Exception) ex; } else { throw new Exception(ex); } } else { try (ArrowMessage msg = ((ArrowMessage) data)) { if (msg.getMessageType() == HeaderType.NONE) { updateMetadata(msg); // We received a message without data, so erase any leftover data if (fulfilledRoot != null) { fulfilledRoot.clear(); } } else if (msg.getMessageType() == HeaderType.RECORD_BATCH) { checkMetadataVersion(msg); // Ensure we have the root root.get().clear(); try (ArrowRecordBatch arb = msg.asRecordBatch()) { loader.load(arb); } updateMetadata(msg); } else if (msg.getMessageType() == HeaderType.DICTIONARY_BATCH) { checkMetadataVersion(msg); // Ensure we have the root root.get().clear(); try (ArrowDictionaryBatch arb = msg.asDictionaryBatch()) { final long id = arb.getDictionaryId(); if (dictionaries == null) { throw new IllegalStateException( "Dictionary ownership was claimed by the application."); } final Dictionary dictionary = dictionaries.lookup(id); if (dictionary == null) { throw new IllegalArgumentException("Dictionary not defined in schema: ID " + id); } final FieldVector vector = dictionary.getVector(); final VectorSchemaRoot dictionaryRoot = new VectorSchemaRoot( Collections.singletonList(vector.getField()), Collections.singletonList(vector), 0); final VectorLoader dictionaryLoader = new VectorLoader(dictionaryRoot); dictionaryLoader.load(arb.getDictionary()); } return next(); } else { throw new UnsupportedOperationException( "Message type is unsupported: " + msg.getMessageType()); } return true; } } } catch (RuntimeException e) { throw e; } catch (ExecutionException e) { throw StatusUtils.fromThrowable(e.getCause()); } catch (Exception e) { throw new RuntimeException(e); } } /** Update our metadata reference with a new one from this message. */ private void updateMetadata(ArrowMessage msg) { if (this.applicationMetadata != null) { this.applicationMetadata.close(); } this.applicationMetadata = msg.getApplicationMetadata(); if (this.applicationMetadata != null) { this.applicationMetadata.getReferenceManager().retain(); } } /** Ensure the Arrow metadata version doesn't change mid-stream. */ private void checkMetadataVersion(ArrowMessage msg) { if (msg.asSchemaMessage() == null) { return; } MetadataVersion receivedVersion = MetadataVersion.fromFlatbufID(msg.asSchemaMessage().getMessage().version()); if (this.metadataVersion != receivedVersion) { throw new IllegalStateException( "Metadata version mismatch: stream started as " + this.metadataVersion + " but got message with version " + receivedVersion); } } /** * Get the current vector data from the stream. * *

    The data in the root may change at any time. Clients should NOT modify the root, but instead * unload the data into their own root. * * @throws FlightRuntimeException if there was an error reading the schema from the stream. */ public VectorSchemaRoot getRoot() { try { return root.get(); } catch (InterruptedException e) { throw CallStatus.INTERNAL.withCause(e).toRuntimeException(); } catch (ExecutionException e) { throw StatusUtils.fromThrowable(e.getCause()); } } /** * Check if there is a root (i.e. whether the other end has started sending data). * *

    Updated by calls to {@link #next()}. * * @return true if and only if the other end has started sending data. */ public boolean hasRoot() { return root.isDone(); } /** * Get the most recent metadata sent from the server. This may be cleared by calls to {@link * #next()} if the server sends a message without metadata. This does NOT take ownership of the * buffer - call retain() to create a reference if you need the buffer after a call to {@link * #next()}. * * @return the application metadata. May be null. */ public ArrowBuf getLatestMetadata() { return applicationMetadata; } private synchronized void requestOutstanding() { if (pending.get() < pendingTarget) { requestor.request(pendingTarget - pending.get()); pending.set(pendingTarget); } } private class Observer implements StreamObserver { Observer() { super(); } /** Helper to add an item to the queue under the appropriate lock. */ private void enqueue(AutoCloseable message) { synchronized (completed) { if (completed.isDone()) { // The stream is already closed (RPC ended), discard the message AutoCloseables.closeNoChecked(message); } else { queue.add(message); } } } @Override public void onNext(ArrowMessage msg) { // Operations here have to be under a lock so that we don't add a message to the queue while // in the middle of // close(). requestOutstanding(); switch (msg.getMessageType()) { case NONE: { // No IPC message - pure metadata or descriptor if (msg.getDescriptor() != null) { descriptor.set(new FlightDescriptor(msg.getDescriptor())); } if (msg.getApplicationMetadata() != null) { enqueue(msg); } break; } case SCHEMA: { Schema schema = msg.asSchema(); // if there is app metadata in the schema message, make sure // that we don't leak it. ArrowBuf meta = msg.getApplicationMetadata(); if (meta != null) { meta.close(); } final List fields = new ArrayList<>(); final Map dictionaryMap = new HashMap<>(); for (final Field originalField : schema.getFields()) { final Field updatedField = DictionaryUtility.toMemoryFormat(originalField, allocator, dictionaryMap); fields.add(updatedField); } for (final Map.Entry entry : dictionaryMap.entrySet()) { dictionaries.put(entry.getValue()); } schema = new Schema(fields, schema.getCustomMetadata()); metadataVersion = MetadataVersion.fromFlatbufID(msg.asSchemaMessage().getMessage().version()); try { MetadataV4UnionChecker.checkRead(schema, metadataVersion); } catch (IOException e) { ex = e; enqueue(DONE_EX); break; } synchronized (completed) { if (!completed.isDone()) { fulfilledRoot = VectorSchemaRoot.create(schema, allocator); loader = new VectorLoader(fulfilledRoot); if (msg.getDescriptor() != null) { descriptor.set(new FlightDescriptor(msg.getDescriptor())); } root.set(fulfilledRoot); } } break; } case RECORD_BATCH: case DICTIONARY_BATCH: enqueue(msg); break; case TENSOR: default: ex = new UnsupportedOperationException( "Unable to handle message of type: " + msg.getMessageType()); enqueue(DONE_EX); } } @Override public void onError(Throwable t) { ex = StatusUtils.fromThrowable(t); queue.add(DONE_EX); cancelled.complete(null); root.setException(ex); } @Override public void onCompleted() { // Depends on gRPC calling onNext and onCompleted non-concurrently cancelled.complete(null); queue.add(DONE); } } /** * Cancels sending the stream to a client. * *

    Callers should drain the stream (with {@link #next()}) to ensure all messages sent before * cancellation are received and to wait for the underlying transport to acknowledge cancellation. */ public void cancel(String message, Throwable exception) { if (cancellable == null) { throw new UnsupportedOperationException( "Streams cannot be cancelled that are produced by client. " + "Instead, server should reject incoming messages."); } cancellable.cancel(message, exception); // Do not mark the stream as completed, as gRPC may still be delivering messages. } StreamObserver asObserver() { return new Observer(); } /** Provides a callback to cancel a process that is in progress. */ @FunctionalInterface public interface Cancellable { void cancel(String message, Throwable exception); } /** Provides a interface to request more items from a stream producer. */ @FunctionalInterface public interface Requestor { /** Requests count more messages from the instance of this object. */ void request(int count); } } GetSessionOptionsRequest.java000066400000000000000000000041151515451615200354260ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.io.IOException; import java.nio.ByteBuffer; import org.apache.arrow.flight.impl.Flight; /** A request to get current session options. */ public class GetSessionOptionsRequest { public GetSessionOptionsRequest() {} GetSessionOptionsRequest(Flight.GetSessionOptionsRequest proto) {} Flight.GetSessionOptionsRequest toProtocol() { return Flight.GetSessionOptionsRequest.getDefaultInstance(); } /** * Get the serialized form of this protocol message. * *

    Intended to help interoperability by allowing non-Flight services to still return Flight * types. */ public ByteBuffer serialize() { return ByteBuffer.wrap(toProtocol().toByteArray()); } /** * Parse the serialized form of this protocol message. * *

    Intended to help interoperability by allowing Flight clients to obtain stream info from * non-Flight services. * * @param serialized The serialized form of the message, as returned by {@link #serialize()}. * @return The deserialized message. * @throws IOException if the serialized form is invalid. */ public static GetSessionOptionsRequest deserialize(ByteBuffer serialized) throws IOException { return new GetSessionOptionsRequest(Flight.GetSessionOptionsRequest.parseFrom(serialized)); } } GetSessionOptionsResult.java000066400000000000000000000062421515451615200352570ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.io.IOException; import java.nio.ByteBuffer; import java.util.Collections; import java.util.HashMap; import java.util.Map; import java.util.stream.Collectors; import org.apache.arrow.flight.impl.Flight; /** A request to view the currently-set options for the current server session. */ public class GetSessionOptionsResult { private final Map sessionOptions; public GetSessionOptionsResult(Map sessionOptions) { this.sessionOptions = Collections.unmodifiableMap(new HashMap(sessionOptions)); } GetSessionOptionsResult(Flight.GetSessionOptionsResult proto) { sessionOptions = Collections.unmodifiableMap( proto.getSessionOptionsMap().entrySet().stream() .collect( Collectors.toMap( Map.Entry::getKey, (e) -> SessionOptionValueFactory.makeSessionOptionValue(e.getValue())))); } /** * Get the session options map contained in the request. * * @return An immutable view of the session options map. */ public Map getSessionOptions() { return sessionOptions; } Flight.GetSessionOptionsResult toProtocol() { Flight.GetSessionOptionsResult.Builder b = Flight.GetSessionOptionsResult.newBuilder(); b.putAllSessionOptions( sessionOptions.entrySet().stream() .collect(Collectors.toMap(Map.Entry::getKey, (e) -> e.getValue().toProtocol()))); return b.build(); } /** * Get the serialized form of this protocol message. * *

    Intended to help interoperability by allowing non-Flight services to still return Flight * types. */ public ByteBuffer serialize() { return ByteBuffer.wrap(toProtocol().toByteArray()); } /** * Parse the serialized form of this protocol message. * *

    Intended to help interoperability by allowing Flight clients to obtain stream info from * non-Flight services. * * @param serialized The serialized form of the message, as returned by {@link #serialize()}. * @return The deserialized message. * @throws IOException if the serialized form is invalid. */ public static GetSessionOptionsResult deserialize(ByteBuffer serialized) throws IOException { return new GetSessionOptionsResult(Flight.GetSessionOptionsResult.parseFrom(serialized)); } } HeaderCallOption.java000066400000000000000000000040151515451615200335720ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import io.grpc.Metadata; import io.grpc.stub.AbstractStub; import io.grpc.stub.MetadataUtils; /** Method option for supplying headers to method calls. */ public class HeaderCallOption implements CallOptions.GrpcCallOption { private final Metadata propertiesMetadata = new Metadata(); /** * Header property constructor. * * @param headers the headers that should be sent across. If a header is a string, it should only * be valid ASCII characters. Binary headers should end in "-bin". */ public HeaderCallOption(CallHeaders headers) { for (String key : headers.keys()) { if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) { final Metadata.Key metaKey = Metadata.Key.of(key, Metadata.BINARY_BYTE_MARSHALLER); headers.getAllByte(key).forEach(v -> propertiesMetadata.put(metaKey, v)); } else { final Metadata.Key metaKey = Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER); headers.getAll(key).forEach(v -> propertiesMetadata.put(metaKey, v)); } } } @Override public > T wrapStub(T stub) { return stub.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(propertiesMetadata)); } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/Location.java000066400000000000000000000117151515451615200322510ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.lang.reflect.InvocationTargetException; import java.net.InetSocketAddress; import java.net.SocketAddress; import java.net.URI; import java.net.URISyntaxException; import java.util.Objects; import org.apache.arrow.flight.impl.Flight; /** A URI where a Flight stream is available. */ public class Location { private final URI uri; /** * Constructs a new instance. * * @param uri the URI of the Flight service * @throws IllegalArgumentException if the URI scheme is unsupported */ public Location(String uri) throws URISyntaxException { this(new URI(uri)); } /** * Construct a new instance from an existing URI. * * @param uri the URI of the Flight service */ public Location(URI uri) { super(); Objects.requireNonNull(uri); this.uri = uri; } public URI getUri() { return uri; } /** * Helper method to turn this Location into a SocketAddress. * * @return null if could not be converted */ public SocketAddress toSocketAddress() { switch (uri.getScheme()) { case LocationSchemes.GRPC: case LocationSchemes.GRPC_TLS: case LocationSchemes.GRPC_INSECURE: { return new InetSocketAddress(uri.getHost(), uri.getPort()); } case LocationSchemes.GRPC_DOMAIN_SOCKET: { try { // This dependency is not available on non-Unix platforms. return Class.forName("io.netty.channel.unix.DomainSocketAddress") .asSubclass(SocketAddress.class) .getConstructor(String.class) .newInstance(uri.getPath()); } catch (InstantiationException | ClassNotFoundException | InvocationTargetException | NoSuchMethodException | IllegalAccessException e) { return null; } } default: { return null; } } } /** Convert this Location into its protocol-level representation. */ Flight.Location toProtocol() { return Flight.Location.newBuilder().setUri(uri.toString()).build(); } /** * Construct a special URI to indicate to clients that they may fetch data by reusing an existing * connection to a Flight RPC server. */ public static Location reuseConnection() { try { return new Location(new URI(LocationSchemes.REUSE_CONNECTION, "", "", "", null)); } catch (URISyntaxException e) { // This should never happen. throw new IllegalArgumentException(e); } } /** * Construct a URI for a Flight+gRPC server without transport security. * * @throws IllegalArgumentException if the constructed URI is invalid. */ public static Location forGrpcInsecure(String host, int port) { try { return new Location( new URI(LocationSchemes.GRPC_INSECURE, null, host, port, null, null, null)); } catch (URISyntaxException e) { throw new IllegalArgumentException(e); } } /** * Construct a URI for a Flight+gRPC server with transport security. * * @throws IllegalArgumentException if the constructed URI is invalid. */ public static Location forGrpcTls(String host, int port) { try { return new Location(new URI(LocationSchemes.GRPC_TLS, null, host, port, null, null, null)); } catch (URISyntaxException e) { throw new IllegalArgumentException(e); } } /** * Construct a URI for a Flight+gRPC server over a Unix domain socket. * * @throws IllegalArgumentException if the constructed URI is invalid. */ public static Location forGrpcDomainSocket(String path) { try { return new Location(new URI(LocationSchemes.GRPC_DOMAIN_SOCKET, null, path, null)); } catch (URISyntaxException e) { throw new IllegalArgumentException(e); } } @Override public String toString() { return "Location{" + "uri=" + uri + '}'; } @Override public boolean equals(Object o) { if (this == o) { return true; } if (!(o instanceof Location)) { return false; } Location location = (Location) o; return uri.equals(location.uri); } @Override public int hashCode() { return Objects.hash(uri); } } LocationSchemes.java000066400000000000000000000025001515451615200334720ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** Constants representing well-known URI schemes for Flight services. */ public final class LocationSchemes { public static final String GRPC = "grpc"; public static final String GRPC_INSECURE = "grpc+tcp"; public static final String GRPC_DOMAIN_SOCKET = "grpc+unix"; public static final String GRPC_TLS = "grpc+tls"; public static final String REUSE_CONNECTION = "arrow-flight-reuse-connection"; private LocationSchemes() { throw new AssertionError("Do not instantiate this class."); } } NoOpFlightProducer.java000066400000000000000000000043671515451615200341440ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** A {@link FlightProducer} that throws on all operations. */ public class NoOpFlightProducer implements FlightProducer { @Override public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { listener.error( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } @Override public void listFlights( CallContext context, Criteria criteria, StreamListener listener) { listener.onError( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } @Override public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public Runnable acceptPut( CallContext context, FlightStream flightStream, StreamListener ackStream) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public void doAction(CallContext context, Action action, StreamListener listener) { listener.onError( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } @Override public void listActions(CallContext context, StreamListener listener) { listener.onError( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } } NoOpSessionOptionValueVisitor.java000066400000000000000000000040211515451615200363770ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * A helper to facilitate easier anonymous subclass declaration. * *

    Implementations need only override callbacks for types they wish to do something with. * * @param Return type of the visit operation. */ public class NoOpSessionOptionValueVisitor implements SessionOptionValueVisitor { /** A callback to handle SessionOptionValue containing a String. */ public T visit(String value) { return null; } /** A callback to handle SessionOptionValue containing a boolean. */ public T visit(boolean value) { return null; } /** A callback to handle SessionOptionValue containing a long. */ public T visit(long value) { return null; } /** A callback to handle SessionOptionValue containing a double. */ public T visit(double value) { return null; } /** A callback to handle SessionOptionValue containing an array of String. */ public T visit(String[] value) { return null; } /** * A callback to handle SessionOptionValue containing no value. * *

    By convention, an attempt to set a valueless SessionOptionValue should attempt to unset or * clear the named option value on the server. */ public T visit(Void value) { return null; } } NoOpStreamListener.java000066400000000000000000000030651515451615200341560ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import org.apache.arrow.flight.FlightProducer.StreamListener; /** * A {@link StreamListener} that does nothing for all callbacks. * * @param The type of the callback object. */ public class NoOpStreamListener implements StreamListener { private static NoOpStreamListener INSTANCE = new NoOpStreamListener(); /** Ignores the value received. */ @Override public void onNext(T val) {} /** Ignores the error received. */ @Override public void onError(Throwable t) {} /** Ignores the stream completion event. */ @Override public void onCompleted() {} @SuppressWarnings("unchecked") public static StreamListener getInstance() { // Safe because we never use T return (StreamListener) INSTANCE; } } OutboundStreamListener.java000066400000000000000000000116221515451615200351000ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.ipc.message.IpcOption; /** An interface for writing data to a peer, client or server. */ public interface OutboundStreamListener { /** * A hint indicating whether the client is ready to receive data without excessive buffering. * *

    Writers should poll this flag before sending data to respect backpressure from the client * and avoid sending data faster than the client can handle. Ignoring this flag may mean that the * server will start consuming excessive amounts of memory, as it may buffer messages in memory. */ boolean isReady(); /** * Set a callback for when the listener is ready for new calls to putNext(), i.e. {@link * #isReady()} has become true. * *

    Note that this callback may only be called some time after {@link #isReady()} becomes true, * and may never be called if all executor threads on the server are busy, or the RPC method body * is implemented in a blocking fashion. Note that isReady() must still be checked after the * callback is run as it may have been run spuriously. */ default void setOnReadyHandler(Runnable handler) { throw new UnsupportedOperationException("Not yet implemented."); } /** * Start sending data, using the schema of the given {@link VectorSchemaRoot}. * *

    This method must be called before all others, except {@link #putMetadata(ArrowBuf)}. */ default void start(VectorSchemaRoot root) { start(root, null, IpcOption.DEFAULT); } /** * Start sending data, using the schema of the given {@link VectorSchemaRoot}. * *

    This method must be called before all others, except {@link #putMetadata(ArrowBuf)}. */ default void start(VectorSchemaRoot root, DictionaryProvider dictionaries) { start(root, dictionaries, IpcOption.DEFAULT); } /** * Start sending data, using the schema of the given {@link VectorSchemaRoot}. * *

    This method must be called before all others, except {@link #putMetadata(ArrowBuf)}. */ void start(VectorSchemaRoot root, DictionaryProvider dictionaries, IpcOption option); /** * Send the current contents of the associated {@link VectorSchemaRoot}. * *

    This will not necessarily block until the message is actually sent; it may buffer messages * in memory. Use {@link #isReady()} to check if there is backpressure and avoid excessive * buffering. */ void putNext(); /** * Send the current contents of the associated {@link VectorSchemaRoot} alongside * application-defined metadata. * * @param metadata The metadata to send. Ownership of the buffer is transferred to the Flight * implementation. */ void putNext(ArrowBuf metadata); /** * Send a pure metadata message without any associated data. * *

    This may be called without starting the stream. */ void putMetadata(ArrowBuf metadata); /** * Indicate an error to the client. Terminates the stream; do not call {@link #completed()} * afterwards. */ void error(Throwable ex); /** Indicate that transmission is finished. */ void completed(); /** * Toggle whether to use the zero-copy write optimization. * *

    By default or when disabled, Arrow may copy data into a buffer for the underlying * implementation to send. When enabled, Arrow will instead try to directly enqueue the Arrow * buffer for sending. Not all implementations support this optimization, so even if enabled, you * may not see a difference. * *

    In this mode, buffers must not be reused after they are written with {@link #putNext()}. For * example, you would have to call {@link VectorSchemaRoot#allocateNew()} after every call to * {@link #putNext()}. Hence, this is not enabled by default. * *

    The default value can be toggled globally by setting the JVM property * arrow.flight.enable_zero_copy_write or the environment variable * ARROW_FLIGHT_ENABLE_ZERO_COPY_WRITE. */ default void setUseZeroCopy(boolean enabled) {} } OutboundStreamListenerImpl.java000066400000000000000000000120731515451615200357230ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import io.grpc.stub.CallStreamObserver; import org.apache.arrow.flight.grpc.StatusUtils; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.ipc.message.IpcOption; /** A base class for writing Arrow data to a Flight stream. */ abstract class OutboundStreamListenerImpl implements OutboundStreamListener { private final FlightDescriptor descriptor; // nullable protected final CallStreamObserver responseObserver; protected volatile VectorUnloader unloader; // null until stream started protected IpcOption option; // null until stream started protected boolean tryZeroCopy = ArrowMessage.ENABLE_ZERO_COPY_WRITE; OutboundStreamListenerImpl( FlightDescriptor descriptor, CallStreamObserver responseObserver) { Preconditions.checkNotNull(responseObserver, "responseObserver must be provided"); this.descriptor = descriptor; this.responseObserver = responseObserver; this.unloader = null; } @Override public boolean isReady() { return responseObserver.isReady(); } @Override public void setOnReadyHandler(Runnable handler) { responseObserver.setOnReadyHandler(handler); } @Override public void start(VectorSchemaRoot root, DictionaryProvider dictionaries, IpcOption option) { this.option = option; try { DictionaryUtils.generateSchemaMessages( root.getSchema(), descriptor, dictionaries, option, responseObserver::onNext); } catch (RuntimeException e) { // Propagate runtime exceptions, like those raised when trying to write unions with V4 // metadata throw e; } catch (Exception e) { // Only happens if closing buffers somehow fails - indicates application is an unknown state // so propagate // the exception throw new RuntimeException("Could not generate and send all schema messages", e); } // We include the null count and align buffers to be compatible with Flight/C++ unloader = new VectorUnloader(root, /* includeNullCount */ true, /* alignBuffers */ true); } @Override public void putNext() { putNext(null); } /** * Busy-wait until the stream is ready. * *

    This is overridable as client/server have different behavior. */ protected abstract void waitUntilStreamReady(); @Override public void putNext(ArrowBuf metadata) { if (unloader == null) { throw CallStatus.INTERNAL .withDescription("Stream was not started, call start()") .toRuntimeException(); } waitUntilStreamReady(); // close is a no-op if the message has been written to gRPC, otherwise frees the associated // buffers // in some code paths (e.g. if the call is cancelled), gRPC does not write the message, so we // need to clean up // ourselves. Normally, writing the ArrowMessage will transfer ownership of the data to // gRPC/Netty. try (final ArrowMessage message = new ArrowMessage(unloader.getRecordBatch(), metadata, tryZeroCopy, option)) { responseObserver.onNext(message); } catch (Exception e) { // This exception comes from ArrowMessage#close, not responseObserver#onNext. // Generally this should not happen - ArrowMessage's implementation only closes non-throwing // things. // The user can't reasonably do anything about this, but if something does throw, we shouldn't // let // execution continue since other state (e.g. allocators) may be in an odd state. throw new RuntimeException("Could not free ArrowMessage", e); } } @Override public void putMetadata(ArrowBuf metadata) { waitUntilStreamReady(); try (final ArrowMessage message = new ArrowMessage(metadata)) { responseObserver.onNext(message); } catch (Exception e) { throw StatusUtils.fromThrowable(e); } } @Override public void error(Throwable ex) { responseObserver.onError(StatusUtils.toGrpcException(ex)); } @Override public void completed() { responseObserver.onCompleted(); } @Override public void setUseZeroCopy(boolean enabled) { tryZeroCopy = enabled; } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/PollInfo.java000066400000000000000000000131741515451615200322240ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.google.protobuf.Timestamp; import com.google.protobuf.util.Timestamps; import java.io.IOException; import java.net.URISyntaxException; import java.nio.ByteBuffer; import java.time.Instant; import java.util.Objects; import java.util.Optional; import org.apache.arrow.flight.impl.Flight; /** A POJO representation of the execution of a long-running query. */ public class PollInfo { private final FlightInfo flightInfo; private final FlightDescriptor flightDescriptor; private final Double progress; private final Instant expirationTime; /** * Create a new PollInfo. * * @param flightInfo The FlightInfo (must not be null). * @param flightDescriptor The descriptor used to poll for more information; null if and only if * query is finished. * @param progress Optional progress info in [0.0, 1.0]. * @param expirationTime An expiration time, after which the server may no longer recognize the * descriptor. */ public PollInfo( FlightInfo flightInfo, FlightDescriptor flightDescriptor, Double progress, Instant expirationTime) { this.flightInfo = Objects.requireNonNull(flightInfo); this.flightDescriptor = flightDescriptor; this.progress = progress; this.expirationTime = expirationTime; } PollInfo(Flight.PollInfo flt) throws URISyntaxException { this.flightInfo = new FlightInfo(flt.getInfo()); this.flightDescriptor = flt.hasFlightDescriptor() ? new FlightDescriptor(flt.getFlightDescriptor()) : null; this.progress = flt.hasProgress() ? flt.getProgress() : null; this.expirationTime = flt.hasExpirationTime() ? Instant.ofEpochSecond( flt.getExpirationTime().getSeconds(), Timestamps.toNanos(flt.getExpirationTime())) : null; } /** * The FlightInfo describing the result set of the execution of a query. * *

    This is always present and always contains all endpoints for the query execution so far, not * just new endpoints that completed execution since the last call to {@link * FlightClient#pollInfo(FlightDescriptor, CallOption...)}. */ public FlightInfo getFlightInfo() { return flightInfo; } /** * The FlightDescriptor that should be used to get further updates on this query. * *

    It is present if and only if the query is still running. If present, it should be passed to * {@link FlightClient#pollInfo(FlightDescriptor, CallOption...)} to get an update. */ public Optional getFlightDescriptor() { return Optional.ofNullable(flightDescriptor); } /** * The progress of the query. * *

    If present, should be a value in [0.0, 1.0]. It is not necessarily monotonic or * non-decreasing. */ public Optional getProgress() { return Optional.ofNullable(progress); } /** * The expiration time of the query execution. * *

    After this passes, the server may not recognize the descriptor anymore and the client will * not be able to track the query anymore. */ public Optional getExpirationTime() { return Optional.ofNullable(expirationTime); } Flight.PollInfo toProtocol() { Flight.PollInfo.Builder b = Flight.PollInfo.newBuilder(); b.setInfo(flightInfo.toProtocol()); if (flightDescriptor != null) { b.setFlightDescriptor(flightDescriptor.toProtocol()); } if (progress != null) { b.setProgress(progress); } if (expirationTime != null) { b.setExpirationTime( Timestamp.newBuilder() .setSeconds(expirationTime.getEpochSecond()) .setNanos(expirationTime.getNano()) .build()); } return b.build(); } public ByteBuffer serialize() { return ByteBuffer.wrap(toProtocol().toByteArray()); } public static PollInfo deserialize(ByteBuffer serialized) throws IOException, URISyntaxException { return new PollInfo(Flight.PollInfo.parseFrom(serialized)); } @Override public boolean equals(Object o) { if (this == o) { return true; } if (!(o instanceof PollInfo)) { return false; } PollInfo pollInfo = (PollInfo) o; return Objects.equals(getFlightInfo(), pollInfo.getFlightInfo()) && Objects.equals(getFlightDescriptor(), pollInfo.getFlightDescriptor()) && Objects.equals(getProgress(), pollInfo.getProgress()) && Objects.equals(getExpirationTime(), pollInfo.getExpirationTime()); } @Override public int hashCode() { return Objects.hash(getFlightInfo(), getFlightDescriptor(), getProgress(), getExpirationTime()); } @Override public String toString() { return "PollInfo{" + "flightInfo=" + flightInfo + ", flightDescriptor=" + flightDescriptor + ", progress=" + progress + ", expirationTime=" + expirationTime + '}'; } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/PutResult.java000066400000000000000000000061551515451615200324520ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.google.protobuf.ByteString; import org.apache.arrow.flight.impl.Flight; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReferenceManager; /** * A message from the server during a DoPut operation. * *

    This object owns an {@link ArrowBuf} and should be closed when you are done with it. */ public class PutResult implements AutoCloseable { private ArrowBuf applicationMetadata; private PutResult(ArrowBuf metadata) { applicationMetadata = metadata; } /** * Create a PutResult with application-specific metadata. * *

    This method assumes ownership of the {@link ArrowBuf}. */ public static PutResult metadata(ArrowBuf metadata) { if (metadata == null) { return empty(); } return new PutResult(metadata); } /** Create an empty PutResult. */ public static PutResult empty() { return new PutResult(null); } /** * Get the metadata in this message. May be null. * *

    Ownership of the {@link ArrowBuf} is retained by this object. Call {@link * ReferenceManager#retain()} to preserve a reference. */ public ArrowBuf getApplicationMetadata() { return applicationMetadata; } Flight.PutResult toProtocol() { if (applicationMetadata == null) { return Flight.PutResult.getDefaultInstance(); } return Flight.PutResult.newBuilder() .setAppMetadata(ByteString.copyFrom(applicationMetadata.nioBuffer())) .build(); } /** * Construct a PutResult from a Protobuf message. * * @param allocator The allocator to use for allocating application metadata memory. The result * object owns the allocated buffer, if any. * @param message The gRPC/Protobuf message. */ static PutResult fromProtocol(BufferAllocator allocator, Flight.PutResult message) { final ArrowBuf buf = allocator.buffer(message.getAppMetadata().size()); message .getAppMetadata() .asReadOnlyByteBufferList() .forEach( bb -> { buf.setBytes(buf.writerIndex(), bb); buf.writerIndex(buf.writerIndex() + bb.limit()); }); return new PutResult(buf); } @Override public void close() { if (applicationMetadata != null) { applicationMetadata.close(); } } } RenewFlightEndpointRequest.java000066400000000000000000000050731515451615200357120ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.io.IOException; import java.net.URISyntaxException; import java.nio.ByteBuffer; import java.util.Objects; import org.apache.arrow.flight.impl.Flight; /** A request to extend the expiration time of a FlightEndpoint. */ public class RenewFlightEndpointRequest { private final FlightEndpoint endpoint; public RenewFlightEndpointRequest(FlightEndpoint endpoint) { this.endpoint = Objects.requireNonNull(endpoint); } RenewFlightEndpointRequest(Flight.RenewFlightEndpointRequest proto) throws URISyntaxException { this(new FlightEndpoint(proto.getEndpoint())); } public FlightEndpoint getFlightEndpoint() { return endpoint; } Flight.RenewFlightEndpointRequest toProtocol() { Flight.RenewFlightEndpointRequest.Builder b = Flight.RenewFlightEndpointRequest.newBuilder(); b.setEndpoint(endpoint.toProtocol()); return b.build(); } /** * Get the serialized form of this protocol message. * *

    Intended to help interoperability by allowing non-Flight services to still return Flight * types. */ public ByteBuffer serialize() { return ByteBuffer.wrap(toProtocol().toByteArray()); } /** * Parse the serialized form of this protocol message. * *

    Intended to help interoperability by allowing Flight clients to obtain stream info from * non-Flight services. * * @param serialized The serialized form of the message, as returned by {@link #serialize()}. * @return The deserialized message. * @throws IOException if the serialized form is invalid. */ public static RenewFlightEndpointRequest deserialize(ByteBuffer serialized) throws IOException, URISyntaxException { return new RenewFlightEndpointRequest(Flight.RenewFlightEndpointRequest.parseFrom(serialized)); } } RequestContext.java000066400000000000000000000030731515451615200334150ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.util.Set; /** Tracks variables about the current request. */ public interface RequestContext { /** * Register a variable and a value. * * @param key the variable name. * @param value the value. */ void put(String key, String value); /** * Retrieve a registered variable. * * @param key the variable name. * @return the value, or null if not found. */ String get(String key); /** * Retrieves the keys that have been registered to this context. * * @return the keys used in this context. */ Set keySet(); /** * Deletes a registered variable. * * @return the value associated with the deleted variable, or null if the key doesn't exist. */ String remove(String key); } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/Result.java000066400000000000000000000026311515451615200317540ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.google.protobuf.ByteString; import org.apache.arrow.flight.impl.Flight; /** * Opaque result returned after executing an action. * *

    POJO wrapper around the Flight protocol buffer message sharing the same name. */ public class Result { private final byte[] body; public Result(byte[] body) { this.body = body; } Result(Flight.Result result) { this.body = result.getBody().toByteArray(); } public byte[] getBody() { return body; } Flight.Result toProtocol() { return Flight.Result.newBuilder().setBody(ByteString.copyFrom(body)).build(); } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/SchemaResult.java000066400000000000000000000064321515451615200331000ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.fasterxml.jackson.databind.util.ByteBufferBackedInputStream; import com.google.common.collect.ImmutableList; import com.google.protobuf.ByteString; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.channels.Channels; import java.util.Objects; import org.apache.arrow.flight.impl.Flight; import org.apache.arrow.vector.ipc.ReadChannel; import org.apache.arrow.vector.ipc.WriteChannel; import org.apache.arrow.vector.ipc.message.IpcOption; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.validate.MetadataV4UnionChecker; /** * Opaque result returned after executing a getSchema request. * *

    POJO wrapper around the Flight protocol buffer message sharing the same name. */ public class SchemaResult { private final Schema schema; private final IpcOption option; public SchemaResult(Schema schema) { this(schema, IpcOption.DEFAULT); } /** Create a schema result with specific IPC options for serialization. */ public SchemaResult(Schema schema, IpcOption option) { Objects.requireNonNull(schema); MetadataV4UnionChecker.checkForUnion(schema.getFields().iterator(), option.metadataVersion); this.schema = schema; this.option = option; } public Schema getSchema() { return schema; } /** Converts to the protocol buffer representation. */ Flight.SchemaResult toProtocol() { // Encode schema in a Message payload ByteArrayOutputStream baos = new ByteArrayOutputStream(); try { MessageSerializer.serialize(new WriteChannel(Channels.newChannel(baos)), schema, option); } catch (IOException e) { throw new RuntimeException(e); } return Flight.SchemaResult.newBuilder() .setSchema(ByteString.copyFrom(baos.toByteArray())) .build(); } /** Converts from the protocol buffer representation. */ static SchemaResult fromProtocol(Flight.SchemaResult pbSchemaResult) { try { final ByteBuffer schemaBuf = pbSchemaResult.getSchema().asReadOnlyByteBuffer(); Schema schema = pbSchemaResult.getSchema().size() > 0 ? MessageSerializer.deserializeSchema( new ReadChannel(Channels.newChannel(new ByteBufferBackedInputStream(schemaBuf)))) : new Schema(ImmutableList.of()); return new SchemaResult(schema); } catch (IOException e) { throw new RuntimeException(e); } } } ServerHeaderMiddleware.java000066400000000000000000000035541515451615200350010ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** Middleware that's used to extract and pass headers to the server during requests. */ public class ServerHeaderMiddleware implements FlightServerMiddleware { /** Factory for accessing ServerHeaderMiddleware. */ public static class Factory implements FlightServerMiddleware.Factory { /** Construct a factory for receiving call headers. */ public Factory() {} @Override public ServerHeaderMiddleware onCallStarted( CallInfo callInfo, CallHeaders incomingHeaders, RequestContext context) { return new ServerHeaderMiddleware(incomingHeaders); } } private final CallHeaders headers; private ServerHeaderMiddleware(CallHeaders incomingHeaders) { this.headers = incomingHeaders; } /** Retrieve the headers for this call. */ public CallHeaders headers() { return headers; } @Override public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) {} @Override public void onCallCompleted(CallStatus status) {} @Override public void onCallErrored(Throwable err) {} } ServerSessionMiddleware.java000066400000000000000000000163511515451615200352330ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.util.Collections; import java.util.HashMap; import java.util.Map; import java.util.concurrent.Callable; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; /** Middleware for handling Flight SQL Sessions including session cookie handling. */ public class ServerSessionMiddleware implements FlightServerMiddleware { Factory factory; boolean existingSession; private Session session; private String closedSessionId = null; public static final String sessionCookieName = "arrow_flight_session_id"; /** Factory for managing and accessing ServerSessionMiddleware. */ public static class Factory implements FlightServerMiddleware.Factory { private final ConcurrentMap sessionStore = new ConcurrentHashMap<>(); private final Callable idGenerator; /** * Construct a factory for ServerSessionMiddleware. * *

    Factory manages and accesses persistent sessions based on HTTP cookies. * * @param idGenerator A Callable returning unique session id Strings. */ public Factory(Callable idGenerator) { this.idGenerator = idGenerator; } private synchronized Session createNewSession() { String id; try { id = idGenerator.call(); } catch (Exception ignored) { // Most impls aren't going to throw so don't make caller handle a nonexistent checked // exception throw CallStatus.INTERNAL.withDescription("Session creation error").toRuntimeException(); } Session newSession = new Session(id); if (sessionStore.putIfAbsent(id, newSession) != null) { // Collision, should never happen throw CallStatus.INTERNAL.withDescription("Session creation error").toRuntimeException(); } return newSession; } private void closeSession(String id) { if (sessionStore.remove(id) == null) { throw CallStatus.NOT_FOUND .withDescription("Session id '" + id + "' not found.") .toRuntimeException(); } } @Override public ServerSessionMiddleware onCallStarted( CallInfo callInfo, CallHeaders incomingHeaders, RequestContext context) { String sessionId = null; final Iterable it = incomingHeaders.getAll("cookie"); 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; } } } if (sessionId == null) { // No session cookie, create middleware instance without session. return new ServerSessionMiddleware(this, incomingHeaders, null); } Session session = sessionStore.get(sessionId); // Cookie provided by caller, but invalid if (session == null) { // Can't soft-fail/proceed here, clients will get unexpected behaviour without options they // thought were set. throw CallStatus.NOT_FOUND .withDescription("Invalid " + sessionCookieName + " cookie.") .toRuntimeException(); } return new ServerSessionMiddleware(this, incomingHeaders, session); } } /** A thread-safe container for named SessionOptionValues. */ public static class Session { public final String id; private ConcurrentMap sessionData = new ConcurrentHashMap(); /** * Construct a new Session with the given id. * * @param id The Session's id string, which is used as the session cookie value. */ private Session(String id) { this.id = id; } /** Get session option by name, or null if it does not exist. */ public SessionOptionValue getSessionOption(String name) { return sessionData.get(name); } /** Get an immutable copy of the session options map. */ public Map getSessionOptions() { return Collections.unmodifiableMap(new HashMap(sessionData)); } /** Set session option by name to given value. */ public void setSessionOption(String name, SessionOptionValue value) { sessionData.put(name, value); } /** Idempotently remove name from this session. */ public void eraseSessionOption(String name) { sessionData.remove(name); } } private final CallHeaders headers; private ServerSessionMiddleware( ServerSessionMiddleware.Factory factory, CallHeaders incomingHeaders, Session session) { this.factory = factory; headers = incomingHeaders; this.session = session; existingSession = (session != null); } /** * Check if there is an open session associated with this call. * * @return True iff there is an open session associated with this call. */ public boolean hasSession() { return session != null; } /** * Get the existing or new session value map for this call. * * @return The session option value map, or null in case of an id generation collision. */ public synchronized Session getSession() { if (session == null) { session = factory.createNewSession(); } return session; } /** * Close the current session. * *

    It is an error to call this without a valid session specified via cookie or equivalent. */ public synchronized void closeSession() { if (session == null) { throw CallStatus.NOT_FOUND .withDescription("No session found for the current call.") .toRuntimeException(); } factory.closeSession(session.id); closedSessionId = session.id; session = null; } public CallHeaders getCallHeaders() { return headers; } @Override public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) { if (!existingSession && session != null) { outgoingHeaders.insert("set-cookie", sessionCookieName + "=" + session.id); } if (closedSessionId != null) { outgoingHeaders.insert( "set-cookie", sessionCookieName + "=" + closedSessionId + "; Max-Age=0"); } } @Override public void onCallCompleted(CallStatus status) {} @Override public void onCallErrored(Throwable err) {} } SessionOptionValue.java000066400000000000000000000052151515451615200342310ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.util.Arrays; import org.apache.arrow.flight.impl.Flight; /** A union-like container interface for supported session option value types. */ public abstract class SessionOptionValue { SessionOptionValue() {} /** Value access via a caller-provided visitor/functor. */ public abstract T acceptVisitor(SessionOptionValueVisitor v); Flight.SessionOptionValue toProtocol() { Flight.SessionOptionValue.Builder b = Flight.SessionOptionValue.newBuilder(); SessionOptionValueToProtocolVisitor visitor = new SessionOptionValueToProtocolVisitor(b); this.acceptVisitor(visitor); return b.build(); } /** Check whether the SessionOptionValue is empty/valueless. */ public boolean isEmpty() { return false; } private class SessionOptionValueToProtocolVisitor implements SessionOptionValueVisitor { final Flight.SessionOptionValue.Builder b; SessionOptionValueToProtocolVisitor(Flight.SessionOptionValue.Builder b) { this.b = b; } @Override public Void visit(String value) { b.setStringValue(value); return null; } @Override public Void visit(boolean value) { b.setBoolValue(value); return null; } @Override public Void visit(long value) { b.setInt64Value(value); return null; } @Override public Void visit(double value) { b.setDoubleValue(value); return null; } @Override public Void visit(String[] value) { Flight.SessionOptionValue.StringListValue.Builder pbSLVBuilder = Flight.SessionOptionValue.StringListValue.newBuilder(); pbSLVBuilder.addAllValues(Arrays.asList(value)); b.setStringListValue(pbSLVBuilder.build()); return null; } @Override public Void visit(Void ignored) { b.clearOptionValue(); return null; } } } SessionOptionValueFactory.java000066400000000000000000000167471515451615200355750ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.nio.charset.StandardCharsets; import java.util.Arrays; import org.apache.arrow.flight.impl.Flight; /** Abstract factory for concrete SessionOptionValue instances. */ public class SessionOptionValueFactory { public static SessionOptionValue makeSessionOptionValue(String value) { return new SessionOptionValueString(value); } public static SessionOptionValue makeSessionOptionValue(boolean value) { return new SessionOptionValueBoolean(value); } public static SessionOptionValue makeSessionOptionValue(long value) { return new SessionOptionValueLong(value); } public static SessionOptionValue makeSessionOptionValue(double value) { return new SessionOptionValueDouble(value); } public static SessionOptionValue makeSessionOptionValue(String[] value) { return new SessionOptionValueStringList(value); } public static SessionOptionValue makeEmptySessionOptionValue() { return new SessionOptionValueEmpty(); } /** Construct a SessionOptionValue from its Protobuf object representation. */ public static SessionOptionValue makeSessionOptionValue(Flight.SessionOptionValue proto) { switch (proto.getOptionValueCase()) { case STRING_VALUE: return new SessionOptionValueString(proto.getStringValue()); case BOOL_VALUE: return new SessionOptionValueBoolean(proto.getBoolValue()); case INT64_VALUE: return new SessionOptionValueLong(proto.getInt64Value()); case DOUBLE_VALUE: return new SessionOptionValueDouble(proto.getDoubleValue()); case STRING_LIST_VALUE: // Using ByteString::toByteArray() here otherwise we still somehow get `ByteArray`s with // broken .equals(String) return new SessionOptionValueStringList( proto.getStringListValue().getValuesList().asByteStringList().stream() .map((e) -> new String(e.toByteArray(), StandardCharsets.UTF_8)) .toArray(String[]::new)); case OPTIONVALUE_NOT_SET: return new SessionOptionValueEmpty(); default: // Unreachable throw new IllegalArgumentException(""); } } private static class SessionOptionValueString extends SessionOptionValue { private final String value; SessionOptionValueString(String value) { this.value = value; } @Override public T acceptVisitor(SessionOptionValueVisitor v) { return v.visit(value); } @Override public boolean equals(Object o) { if (this == o) { return true; } if (o == null || getClass() != o.getClass()) { return false; } SessionOptionValueString that = (SessionOptionValueString) o; return value.equals(that.value); } @Override public int hashCode() { return value.hashCode(); } @Override public String toString() { return '"' + value + '"'; } } private static class SessionOptionValueBoolean extends SessionOptionValue { private final boolean value; SessionOptionValueBoolean(boolean value) { this.value = value; } @Override public T acceptVisitor(SessionOptionValueVisitor v) { return v.visit(value); } @Override public boolean equals(Object o) { if (this == o) { return true; } if (o == null || getClass() != o.getClass()) { return false; } SessionOptionValueBoolean that = (SessionOptionValueBoolean) o; return value == that.value; } @Override public int hashCode() { return Boolean.hashCode(value); } @Override public String toString() { return String.valueOf(value); } } private static class SessionOptionValueLong extends SessionOptionValue { private final long value; SessionOptionValueLong(long value) { this.value = value; } @Override public T acceptVisitor(SessionOptionValueVisitor v) { return v.visit(value); } @Override public boolean equals(Object o) { if (this == o) { return true; } if (o == null || getClass() != o.getClass()) { return false; } SessionOptionValueLong that = (SessionOptionValueLong) o; return value == that.value; } @Override public int hashCode() { return Long.hashCode(value); } @Override public String toString() { return String.valueOf(value); } } private static class SessionOptionValueDouble extends SessionOptionValue { private final double value; SessionOptionValueDouble(double value) { this.value = value; } @Override public T acceptVisitor(SessionOptionValueVisitor v) { return v.visit(value); } @Override public boolean equals(Object o) { if (this == o) { return true; } if (o == null || getClass() != o.getClass()) { return false; } SessionOptionValueDouble that = (SessionOptionValueDouble) o; return value == that.value; } @Override public int hashCode() { return Double.hashCode(value); } @Override public String toString() { return String.valueOf(value); } } private static class SessionOptionValueStringList extends SessionOptionValue { private final String[] value; SessionOptionValueStringList(String[] value) { this.value = value.clone(); } @Override public T acceptVisitor(SessionOptionValueVisitor v) { return v.visit(value); } @Override public boolean equals(Object o) { if (this == o) { return true; } if (o == null || getClass() != o.getClass()) { return false; } SessionOptionValueStringList that = (SessionOptionValueStringList) o; return Arrays.deepEquals(value, that.value); } @Override public int hashCode() { return Arrays.deepHashCode(value); } @Override public String toString() { if (value.length == 0) { return "[]"; } return "[\"" + String.join("\", \"", value) + "\"]"; } } private static class SessionOptionValueEmpty extends SessionOptionValue { @Override public T acceptVisitor(SessionOptionValueVisitor v) { return v.visit((Void) null); } @Override public boolean isEmpty() { return true; } @Override public boolean equals(Object o) { if (this == o) { return true; } if (o == null || getClass() != o.getClass()) { return false; } return true; } @Override public int hashCode() { return SessionOptionValueEmpty.class.hashCode(); } @Override public String toString() { return ""; } } } SessionOptionValueVisitor.java000066400000000000000000000033411515451615200356070ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * A visitor interface to access SessionOptionValue's contained value. * * @param Return type of the visit operation. */ public interface SessionOptionValueVisitor { /** A callback to handle SessionOptionValue containing a String. */ T visit(String value); /** A callback to handle SessionOptionValue containing a boolean. */ T visit(boolean value); /** A callback to handle SessionOptionValue containing a long. */ T visit(long value); /** A callback to handle SessionOptionValue containing a double. */ T visit(double value); /** A callback to handle SessionOptionValue containing an array of String. */ T visit(String[] value); /** * A callback to handle SessionOptionValue containing no value. * *

    By convention, an attempt to set a valueless SessionOptionValue should attempt to unset or * clear the named option value on the server. */ T visit(Void value); } SetSessionOptionsRequest.java000066400000000000000000000063501515451615200354450ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.io.IOException; import java.nio.ByteBuffer; import java.util.Collections; import java.util.HashMap; import java.util.Map; import java.util.stream.Collectors; import org.apache.arrow.flight.impl.Flight; /** A request to set option(s) in an existing or implicitly-created server session. */ public class SetSessionOptionsRequest { private final Map sessionOptions; public SetSessionOptionsRequest(Map sessionOptions) { this.sessionOptions = Collections.unmodifiableMap(new HashMap(sessionOptions)); } SetSessionOptionsRequest(Flight.SetSessionOptionsRequest proto) { sessionOptions = Collections.unmodifiableMap( proto.getSessionOptionsMap().entrySet().stream() .collect( Collectors.toMap( Map.Entry::getKey, (e) -> SessionOptionValueFactory.makeSessionOptionValue(e.getValue())))); } /** * Get the session option map from the request. * * @return An immutable view of the session options map. */ public Map getSessionOptions() { return Collections.unmodifiableMap(sessionOptions); } Flight.SetSessionOptionsRequest toProtocol() { Flight.SetSessionOptionsRequest.Builder b = Flight.SetSessionOptionsRequest.newBuilder(); b.putAllSessionOptions( sessionOptions.entrySet().stream() .collect(Collectors.toMap(Map.Entry::getKey, (e) -> e.getValue().toProtocol()))); return b.build(); } /** * Get the serialized form of this protocol message. * *

    Intended to help interoperability by allowing non-Flight services to still return Flight * types. */ public ByteBuffer serialize() { return ByteBuffer.wrap(toProtocol().toByteArray()); } /** * Parse the serialized form of this protocol message. * *

    Intended to help interoperability by allowing Flight clients to obtain stream info from * non-Flight services. * * @param serialized The serialized form of the message, as returned by {@link #serialize()}. * @return The deserialized message. * @throws IOException if the serialized form is invalid. */ public static SetSessionOptionsRequest deserialize(ByteBuffer serialized) throws IOException { return new SetSessionOptionsRequest(Flight.SetSessionOptionsRequest.parseFrom(serialized)); } } SetSessionOptionsResult.java000066400000000000000000000112611515451615200352700ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.io.IOException; import java.nio.ByteBuffer; import java.util.Collections; import java.util.HashMap; import java.util.Map; import java.util.stream.Collectors; import org.apache.arrow.flight.impl.Flight; /** The result of attempting to set a set of session options. */ public class SetSessionOptionsResult { /** Error status value for per-option errors. */ public enum ErrorValue { /** * The status of setting the option is unknown. Servers should avoid using this value (send a * NOT_FOUND error if the requested session is not known). Clients can retry the request. */ UNSPECIFIED, /** The given session option name is invalid. */ INVALID_NAME, /** The session option value or type is invalid. */ INVALID_VALUE, /** The session option cannot be set. */ ERROR, ; static ErrorValue fromProtocol(Flight.SetSessionOptionsResult.ErrorValue s) { return values()[s.getNumber()]; } Flight.SetSessionOptionsResult.ErrorValue toProtocol() { return Flight.SetSessionOptionsResult.ErrorValue.values()[ordinal()]; } } /** Per-option extensible error response container. */ public static class Error { public ErrorValue value; public Error(ErrorValue value) { this.value = value; } Error(Flight.SetSessionOptionsResult.Error e) { value = ErrorValue.fromProtocol(e.getValue()); } Flight.SetSessionOptionsResult.Error toProtocol() { Flight.SetSessionOptionsResult.Error.Builder b = Flight.SetSessionOptionsResult.Error.newBuilder(); b.setValue(value.toProtocol()); return b.build(); } @Override public boolean equals(Object o) { if (this == o) { return true; } if (o == null || getClass() != o.getClass()) { return false; } Error that = (Error) o; return value == that.value; } @Override public int hashCode() { return value.hashCode(); } } private final Map errors; public SetSessionOptionsResult(Map errors) { this.errors = Collections.unmodifiableMap(new HashMap(errors)); } SetSessionOptionsResult(Flight.SetSessionOptionsResult proto) { errors = Collections.unmodifiableMap( proto.getErrors().entrySet().stream() .collect(Collectors.toMap(Map.Entry::getKey, (e) -> new Error(e.getValue())))); } /** Report whether the error map has nonzero length. */ public boolean hasErrors() { return errors.size() > 0; } /** * Get the error status map from the result object. * * @return An immutable view of the error status map. */ public Map getErrors() { return errors; } Flight.SetSessionOptionsResult toProtocol() { Flight.SetSessionOptionsResult.Builder b = Flight.SetSessionOptionsResult.newBuilder(); b.putAllErrors( errors.entrySet().stream() .collect(Collectors.toMap(Map.Entry::getKey, (e) -> e.getValue().toProtocol()))); return b.build(); } /** * Get the serialized form of this protocol message. * *

    Intended to help interoperability by allowing non-Flight services to still return Flight * types. */ public ByteBuffer serialize() { return ByteBuffer.wrap(toProtocol().toByteArray()); } /** * Parse the serialized form of this protocol message. * *

    Intended to help interoperability by allowing Flight clients to obtain stream info from * non-Flight services. * * @param serialized The serialized form of the message, as returned by {@link #serialize()}. * @return The deserialized message. * @throws IOException if the serialized form is invalid. */ public static SetSessionOptionsResult deserialize(ByteBuffer serialized) throws IOException { return new SetSessionOptionsResult(Flight.SetSessionOptionsResult.parseFrom(serialized)); } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/StreamPipe.java000066400000000000000000000070621515451615200325520ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import io.grpc.stub.StreamObserver; import java.util.function.Consumer; import java.util.function.Function; import org.apache.arrow.flight.FlightProducer.StreamListener; import org.apache.arrow.flight.grpc.StatusUtils; import org.apache.arrow.util.AutoCloseables; /** * Shim listener to avoid exposing GRPC internals. * * @param From Type * @param To Type */ class StreamPipe implements StreamListener { private final StreamObserver delegate; private final Function mapFunction; private final Consumer errorHandler; private AutoCloseable resource; private boolean closed = false; /** * Wrap the given gRPC StreamObserver with a transformation function. * * @param delegate The {@link StreamObserver} to wrap. * @param func The transformation function. * @param errorHandler A handler for uncaught exceptions (e.g. if something tries to double-close * this stream). * @param The source type. * @param The output type. * @return A wrapped listener. */ public static StreamPipe wrap( StreamObserver delegate, Function func, Consumer errorHandler) { return new StreamPipe<>(delegate, func, errorHandler); } public StreamPipe( StreamObserver delegate, Function func, Consumer errorHandler) { super(); this.delegate = delegate; this.mapFunction = func; this.errorHandler = errorHandler; this.resource = null; } /** Set an AutoCloseable resource to be cleaned up when the gRPC observer is to be completed. */ void setAutoCloseable(AutoCloseable ac) { resource = ac; } @Override public void onNext(FROM val) { delegate.onNext(mapFunction.apply(val)); } @Override public void onError(Throwable t) { if (closed) { errorHandler.accept(t); return; } try { AutoCloseables.close(resource); } catch (Exception e) { errorHandler.accept(e); } finally { // Set closed to true in case onError throws, so that we don't try to close again closed = true; delegate.onError(StatusUtils.toGrpcException(t)); } } @Override public void onCompleted() { if (closed) { errorHandler.accept(new IllegalStateException("Tried to complete already-completed call")); return; } try { AutoCloseables.close(resource); } catch (Exception e) { errorHandler.accept(e); } finally { // Set closed to true in case onCompleted throws, so that we don't try to close again closed = true; delegate.onCompleted(); } } /** Ensure this stream has been completed. */ void ensureCompleted() { if (!closed) { onCompleted(); } } } SyncPutListener.java000066400000000000000000000101141515451615200335250ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; import org.apache.arrow.flight.grpc.StatusUtils; import org.apache.arrow.memory.ArrowBuf; /** * A listener for server-sent application metadata messages during a Flight DoPut. This class wraps * the messages in a synchronous interface. */ public final class SyncPutListener implements FlightClient.PutListener, AutoCloseable { private final LinkedBlockingQueue queue; private final CompletableFuture completed; private static final Object DONE = new Object(); private static final Object DONE_WITH_EXCEPTION = new Object(); public SyncPutListener() { queue = new LinkedBlockingQueue<>(); completed = new CompletableFuture<>(); } private PutResult unwrap(Object queueItem) throws InterruptedException, ExecutionException { if (queueItem == DONE) { queue.put(queueItem); return null; } else if (queueItem == DONE_WITH_EXCEPTION) { queue.put(queueItem); completed.get(); } return (PutResult) queueItem; } /** * Get the next message from the server, blocking until it is available. * * @return The next message, or null if the server is done sending messages. The caller assumes * ownership of the metadata and must remember to close it. * @throws InterruptedException if interrupted while waiting. * @throws ExecutionException if the server sent an error, or if there was an internal error. */ public PutResult read() throws InterruptedException, ExecutionException { return unwrap(queue.take()); } /** * Get the next message from the server, blocking for the specified amount of time until it is * available. * * @return The next message, or null if the server is done sending messages or no message arrived * before the timeout. The caller assumes ownership of the metadata and must remember to close * it. * @throws InterruptedException if interrupted while waiting. * @throws ExecutionException if the server sent an error, or if there was an internal error. */ public PutResult poll(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException { return unwrap(queue.poll(timeout, unit)); } @Override public void getResult() { try { completed.get(); } catch (ExecutionException e) { throw StatusUtils.fromThrowable(e.getCause()); } catch (InterruptedException e) { throw StatusUtils.fromThrowable(e); } } @Override public void onNext(PutResult val) { final ArrowBuf metadata = val.getApplicationMetadata(); metadata.getReferenceManager().retain(); queue.add(PutResult.metadata(metadata)); } @Override public void onError(Throwable t) { completed.completeExceptionally(StatusUtils.fromThrowable(t)); queue.add(DONE_WITH_EXCEPTION); } @Override public void onCompleted() { completed.complete(null); queue.add(DONE); } @Override public void close() { queue.forEach( o -> { if (o instanceof PutResult) { ((PutResult) o).close(); } }); } @Override public boolean isCancelled() { return completed.isDone(); } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/Ticket.java000066400000000000000000000053241515451615200317230ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import com.google.protobuf.ByteString; import java.io.IOException; import java.nio.ByteBuffer; import java.util.Arrays; import org.apache.arrow.flight.impl.Flight; /** Endpoint for a particular stream. */ public class Ticket { private final byte[] bytes; public Ticket(byte[] bytes) { super(); this.bytes = bytes; } public byte[] getBytes() { return bytes; } Ticket(org.apache.arrow.flight.impl.Flight.Ticket ticket) { this.bytes = ticket.getTicket().toByteArray(); } Flight.Ticket toProtocol() { return Flight.Ticket.newBuilder().setTicket(ByteString.copyFrom(bytes)).build(); } /** * Get the serialized form of this protocol message. * *

    Intended to help interoperability by allowing non-Flight services to still return Flight * types. */ public ByteBuffer serialize() { return ByteBuffer.wrap(toProtocol().toByteArray()); } /** * Parse the serialized form of this protocol message. * *

    Intended to help interoperability by allowing Flight clients to obtain stream info from * non-Flight services. * * @param serialized The serialized form of the Ticket, as returned by {@link #serialize()}. * @return The deserialized Ticket. * @throws IOException if the serialized form is invalid. */ public static Ticket deserialize(ByteBuffer serialized) throws IOException { return new Ticket(Flight.Ticket.parseFrom(serialized)); } @Override public int hashCode() { final int prime = 31; int result = 1; result = prime * result + Arrays.hashCode(bytes); return result; } @Override public boolean equals(Object obj) { if (this == obj) { return true; } if (obj == null) { return false; } if (!(obj instanceof Ticket)) { return false; } Ticket other = (Ticket) obj; if (!Arrays.equals(bytes, other.bytes)) { return false; } return true; } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth/000077500000000000000000000000001515451615200305725ustar00rootroot00000000000000AuthConstants.java000066400000000000000000000035201515451615200341540ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth; import io.grpc.Context; import io.grpc.Metadata; import io.grpc.Metadata.BinaryMarshaller; import io.grpc.MethodDescriptor; import org.apache.arrow.flight.FlightConstants; /** Constants used in authorization of flight connections. */ public final class AuthConstants { public static final String HANDSHAKE_DESCRIPTOR_NAME = MethodDescriptor.generateFullMethodName(FlightConstants.SERVICE, "Handshake"); public static final String TOKEN_NAME = "Auth-Token-bin"; public static final Metadata.Key TOKEN_KEY = Metadata.Key.of( TOKEN_NAME, new BinaryMarshaller() { @Override public byte[] toBytes(byte[] value) { return value; } @Override public byte[] parseBytes(byte[] serialized) { return serialized; } }); public static final Context.Key PEER_IDENTITY_KEY = Context.keyWithDefault("arrow-flight-peer-identity", ""); private AuthConstants() {} } BasicClientAuthHandler.java000066400000000000000000000032641515451615200356630ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth; import java.util.Iterator; import org.apache.arrow.flight.impl.Flight.BasicAuth; /** A client auth handler that supports username and password. */ public class BasicClientAuthHandler implements ClientAuthHandler { private final String name; private final String password; private byte[] token = null; public BasicClientAuthHandler(String name, String password) { this.name = name; this.password = password; } @Override public void authenticate(ClientAuthSender outgoing, Iterator incoming) { BasicAuth.Builder builder = BasicAuth.newBuilder(); if (name != null) { builder.setUsername(name); } if (password != null) { builder.setPassword(password); } outgoing.send(builder.build().toByteArray()); this.token = incoming.next(); } @Override public byte[] getCallToken() { return token; } } BasicServerAuthHandler.java000066400000000000000000000045161515451615200357140ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth; import com.google.protobuf.InvalidProtocolBufferException; import java.util.Iterator; import java.util.Optional; import org.apache.arrow.flight.impl.Flight.BasicAuth; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** A ServerAuthHandler for username/password authentication. */ public class BasicServerAuthHandler implements ServerAuthHandler { private static final Logger logger = LoggerFactory.getLogger(BasicServerAuthHandler.class); private final BasicAuthValidator authValidator; public BasicServerAuthHandler(BasicAuthValidator authValidator) { super(); this.authValidator = authValidator; } /** Interface that this handler delegates for determining if credentials are valid. */ public interface BasicAuthValidator { byte[] getToken(String username, String password) throws Exception; Optional isValid(byte[] token); } @Override public boolean authenticate(ServerAuthSender outgoing, Iterator incoming) { byte[] bytes = incoming.next(); try { BasicAuth auth = BasicAuth.parseFrom(bytes); byte[] token = authValidator.getToken(auth.getUsername(), auth.getPassword()); outgoing.send(token); return true; } catch (InvalidProtocolBufferException e) { logger.debug("Failure parsing auth message.", e); } catch (Exception e) { logger.debug("Unknown error during authorization.", e); } return false; } @Override public Optional isValid(byte[] token) { return authValidator.isValid(token); } } ClientAuthHandler.java000066400000000000000000000036231515451615200347200ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth; import java.util.Iterator; import org.apache.arrow.flight.FlightClient; /** * Implement authentication for Flight on the client side. * * @deprecated As of 14.0.0. This implements a stateful "login" flow that does not play well with * distributed or stateless systems. It will not be removed, but should not be used. Instead see * {@link FlightClient#authenticateBasicToken(String, String)}. */ @Deprecated public interface ClientAuthHandler { /** * Handle the initial handshake with the server. * * @param outgoing A channel to send data to the server. * @param incoming An iterator of incoming data from the server. */ void authenticate(ClientAuthSender outgoing, Iterator incoming); /** Get the per-call authentication token. */ byte[] getCallToken(); /** A communication channel to the server during initial connection. */ interface ClientAuthSender { /** Send the server a message. */ void send(byte[] payload); /** Signal an error to the server and abort the authentication attempt. */ void onError(Throwable cause); } } ClientAuthInterceptor.java000066400000000000000000000046011515451615200356360ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth; import io.grpc.CallOptions; import io.grpc.Channel; import io.grpc.ClientCall; import io.grpc.ClientInterceptor; import io.grpc.ForwardingClientCall.SimpleForwardingClientCall; import io.grpc.Metadata; import io.grpc.MethodDescriptor; /** GRPC client intercepter that handles authentication with the server. */ public class ClientAuthInterceptor implements ClientInterceptor { private volatile ClientAuthHandler authHandler = null; public void setAuthHandler(ClientAuthHandler authHandler) { this.authHandler = authHandler; } public ClientAuthInterceptor() {} public boolean hasAuthHandler() { return authHandler != null; } @Override public ClientCall interceptCall( MethodDescriptor methodDescriptor, CallOptions callOptions, Channel next) { ClientCall call = next.newCall(methodDescriptor, callOptions); // once we have an auth header, add that to the calls. if (authHandler != null) { call = new HeaderAttachingClientCall<>(call); } return call; } private final class HeaderAttachingClientCall extends SimpleForwardingClientCall { private HeaderAttachingClientCall(ClientCall call) { super(call); } @Override public void start(Listener responseListener, Metadata headers) { final Metadata authHeaders = new Metadata(); authHeaders.put(AuthConstants.TOKEN_KEY, authHandler.getCallToken()); headers.merge(authHeaders); super.start(responseListener, headers); } } } ClientAuthWrapper.java000066400000000000000000000121771515451615200347670ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth; import com.google.protobuf.ByteString; import io.grpc.StatusRuntimeException; import io.grpc.stub.StreamObserver; import java.util.Iterator; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; import java.util.concurrent.LinkedBlockingQueue; import org.apache.arrow.flight.auth.ClientAuthHandler.ClientAuthSender; import org.apache.arrow.flight.grpc.StatusUtils; import org.apache.arrow.flight.impl.Flight.HandshakeRequest; import org.apache.arrow.flight.impl.Flight.HandshakeResponse; import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceStub; /** Utility class for performing authorization over using a GRPC stub. */ public class ClientAuthWrapper { /** * Do client auth for a client. The stub will be authenticated after this method returns. * * @param authHandler The handler to use. * @param stub The service stub. */ public static void doClientAuth(ClientAuthHandler authHandler, FlightServiceStub stub) { AuthObserver observer = new AuthObserver(); try { observer.responseObserver = stub.handshake(observer); authHandler.authenticate(observer.sender, observer.iter); if (!observer.sender.errored) { observer.responseObserver.onCompleted(); } } catch (StatusRuntimeException sre) { throw StatusUtils.fromGrpcRuntimeException(sre); } try { if (!observer.completed.get()) { // TODO: ARROW-5681 throw new RuntimeException("Unauthenticated"); } } catch (InterruptedException e) { throw new RuntimeException(e); } catch (ExecutionException e) { throw StatusUtils.fromThrowable(e.getCause()); } } private static class AuthObserver implements StreamObserver { private volatile StreamObserver responseObserver; private final LinkedBlockingQueue messages = new LinkedBlockingQueue<>(); private final AuthSender sender = new AuthSender(); private CompletableFuture completed; public AuthObserver() { super(); completed = new CompletableFuture<>(); } @Override public void onNext(HandshakeResponse value) { ByteString payload = value.getPayload(); messages.add(payload.toByteArray()); } private Iterator iter = new Iterator() { @Override public byte[] next() { while (!completed.isDone() || !messages.isEmpty()) { byte[] bytes = messages.poll(); if (bytes == null) { // busy wait. continue; } else { return bytes; } } if (completed.isCompletedExceptionally()) { // Preserve prior exception behavior // TODO: with ARROW-5681, throw an appropriate Flight exception if gRPC raised an // exception try { completed.get(); } catch (InterruptedException e) { throw new RuntimeException(e); } catch (ExecutionException e) { if (e.getCause() instanceof StatusRuntimeException) { throw (StatusRuntimeException) e.getCause(); } throw new RuntimeException(e); } } throw new IllegalStateException( "You attempted to retrieve messages after there were none."); } @Override public boolean hasNext() { return !messages.isEmpty(); } }; @Override public void onError(Throwable t) { completed.completeExceptionally(t); } private class AuthSender implements ClientAuthSender { private boolean errored = false; @Override public void send(byte[] payload) { try { responseObserver.onNext( HandshakeRequest.newBuilder().setPayload(ByteString.copyFrom(payload)).build()); } catch (StatusRuntimeException sre) { throw StatusUtils.fromGrpcRuntimeException(sre); } } @Override public void onError(Throwable cause) { this.errored = true; responseObserver.onError(StatusUtils.toGrpcException(cause)); } } @Override public void onCompleted() { completed.complete(true); } } } ServerAuthHandler.java000066400000000000000000000050701515451615200347460ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth; import java.util.Iterator; import java.util.Optional; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.auth2.CallHeaderAuthenticator; /** * Interface for Server side authentication handlers. * * @deprecated As of 14.0.0. This implements a stateful "login" flow that does not play well with * distributed or stateless systems. It will not be removed, but should not be used. Instead, * see {@link FlightServer.Builder#headerAuthenticator(CallHeaderAuthenticator)} and {@link * CallHeaderAuthenticator}. */ @Deprecated public interface ServerAuthHandler { /** * Validate the client token provided on each call. * * @return An empty optional if the client is not authenticated; the peer identity otherwise (may * be the empty string). */ Optional isValid(byte[] token); /** * Handle the initial handshake with the client. * * @param outgoing A writer to send messages to the client. * @param incoming An iterator of messages from the client. * @return true if client is authenticated, false otherwise. */ boolean authenticate(ServerAuthSender outgoing, Iterator incoming); /** * Interface for a server implementations to send back authentication messages back to the client. */ interface ServerAuthSender { void send(byte[] payload); void onError(Throwable cause); } /** An auth handler that does nothing. */ ServerAuthHandler NO_OP = new ServerAuthHandler() { @Override public Optional isValid(byte[] token) { return Optional.of(""); } @Override public boolean authenticate(ServerAuthSender outgoing, Iterator incoming) { return true; } }; } ServerAuthInterceptor.java000066400000000000000000000062121515451615200356660ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth; import io.grpc.Context; import io.grpc.Contexts; import io.grpc.Metadata; import io.grpc.ServerCall; import io.grpc.ServerCall.Listener; import io.grpc.ServerCallHandler; import io.grpc.ServerInterceptor; import io.grpc.Status; import io.grpc.StatusRuntimeException; import java.util.Optional; import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.flight.grpc.StatusUtils; /** GRPC Interceptor for performing authentication. */ public class ServerAuthInterceptor implements ServerInterceptor { private final ServerAuthHandler authHandler; public ServerAuthInterceptor(ServerAuthHandler authHandler) { this.authHandler = authHandler; } @Override public Listener interceptCall( ServerCall call, Metadata headers, ServerCallHandler next) { if (!call.getMethodDescriptor() .getFullMethodName() .equals(AuthConstants.HANDSHAKE_DESCRIPTOR_NAME)) { final Optional peerIdentity; // Allow customizing the response code by throwing FlightRuntimeException try { peerIdentity = isValid(headers); } catch (FlightRuntimeException e) { final Status grpcStatus = StatusUtils.toGrpcStatus(e.status()); call.close(grpcStatus, new Metadata()); return new NoopServerCallListener<>(); } catch (StatusRuntimeException e) { Metadata trailers = e.getTrailers(); call.close(e.getStatus(), trailers == null ? new Metadata() : trailers); return new NoopServerCallListener<>(); } if (!peerIdentity.isPresent()) { // Send back a description along with the status code call.close( Status.UNAUTHENTICATED.withDescription( "Unauthenticated (invalid or missing auth token)"), new Metadata()); return new NoopServerCallListener<>(); } return Contexts.interceptCall( Context.current().withValue(AuthConstants.PEER_IDENTITY_KEY, peerIdentity.get()), call, headers, next); } return next.startCall(call, headers); } private Optional isValid(Metadata headers) { byte[] token = headers.get(AuthConstants.TOKEN_KEY); return authHandler.isValid(token); } private static class NoopServerCallListener extends ServerCall.Listener {} } ServerAuthWrapper.java000066400000000000000000000111271515451615200350110ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth; import com.google.protobuf.ByteString; import io.grpc.stub.StreamObserver; import java.util.Iterator; import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingQueue; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.auth.ServerAuthHandler.ServerAuthSender; import org.apache.arrow.flight.grpc.StatusUtils; import org.apache.arrow.flight.impl.Flight.HandshakeRequest; import org.apache.arrow.flight.impl.Flight.HandshakeResponse; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** Contains utility methods for integrating authorization into a GRPC stream. */ public class ServerAuthWrapper { private static final Logger LOGGER = LoggerFactory.getLogger(ServerAuthWrapper.class); /** * Wrap the auth handler for handshake purposes. * * @param authHandler Authentication handler * @param responseObserver Observer for handshake response * @param executors ExecutorService * @return AuthObserver */ public static StreamObserver wrapHandshake( ServerAuthHandler authHandler, StreamObserver responseObserver, ExecutorService executors) { // stream started. AuthObserver observer = new AuthObserver(responseObserver); final Runnable r = () -> { try { if (authHandler.authenticate(observer.sender, observer.iter)) { responseObserver.onCompleted(); return; } responseObserver.onError( StatusUtils.toGrpcException(CallStatus.UNAUTHENTICATED.toRuntimeException())); } catch (Exception ex) { LOGGER.error("Error during authentication", ex); responseObserver.onError(StatusUtils.toGrpcException(ex)); } }; observer.future = executors.submit(r); return observer; } private static class AuthObserver implements StreamObserver { private final StreamObserver responseObserver; private volatile Future future; private volatile boolean completed = false; private final LinkedBlockingQueue messages = new LinkedBlockingQueue<>(); private final AuthSender sender = new AuthSender(); public AuthObserver(StreamObserver responseObserver) { super(); this.responseObserver = responseObserver; } @Override public void onNext(HandshakeRequest value) { ByteString payload = value.getPayload(); messages.add(payload.toByteArray()); } private Iterator iter = new Iterator() { @Override public byte[] next() { while (!completed || !messages.isEmpty()) { byte[] bytes = messages.poll(); if (bytes == null) { // busy wait. continue; } return bytes; } throw new IllegalStateException("Requesting more messages than client sent."); } @Override public boolean hasNext() { return !messages.isEmpty(); } }; @Override public void onError(Throwable t) { completed = true; while (future == null) { /* busy wait */ } future.cancel(true); } @Override public void onCompleted() { completed = true; } private class AuthSender implements ServerAuthSender { @Override public void send(byte[] payload) { responseObserver.onNext( HandshakeResponse.newBuilder().setPayload(ByteString.copyFrom(payload)).build()); } @Override public void onError(Throwable cause) { responseObserver.onError(StatusUtils.toGrpcException(cause)); } } } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/000077500000000000000000000000001515451615200306545ustar00rootroot00000000000000Auth2Constants.java000066400000000000000000000023161515451615200343220ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth2; /** Constants used in authorization of flight connections. */ public final class Auth2Constants { public static final String PEER_IDENTITY_KEY = "arrow-flight-peer-identity"; public static final String BEARER_PREFIX = "Bearer "; public static final String BASIC_PREFIX = "Basic "; public static final String AUTHORIZATION_HEADER = "Authorization"; private Auth2Constants() {} } AuthUtilities.java000066400000000000000000000031751515451615200342430ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth2; import org.apache.arrow.flight.CallHeaders; /** Utility class for completing the auth process. */ public final class AuthUtilities { private AuthUtilities() {} /** * Helper method for retrieving a value from the Authorization header. * * @param headers The headers to inspect. * @param valuePrefix The prefix within the value portion of the header to extract away. * @return The header value. */ public static String getValueFromAuthHeader(CallHeaders headers, String valuePrefix) { final String authHeaderValue = headers.get(Auth2Constants.AUTHORIZATION_HEADER); if (authHeaderValue != null) { if (authHeaderValue.regionMatches(true, 0, valuePrefix, 0, valuePrefix.length())) { return authHeaderValue.substring(valuePrefix.length()); } } return null; } } BasicAuthCredentialWriter.java000066400000000000000000000031661515451615200365010ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth2; import java.nio.charset.StandardCharsets; import java.util.Base64; import java.util.function.Consumer; import org.apache.arrow.flight.CallHeaders; /** Client credentials that use a username and password. */ public final class BasicAuthCredentialWriter implements Consumer { private final String name; private final String password; public BasicAuthCredentialWriter(String name, String password) { this.name = name; this.password = password; } @Override public void accept(CallHeaders outputHeaders) { outputHeaders.insert( Auth2Constants.AUTHORIZATION_HEADER, Auth2Constants.BASIC_PREFIX + Base64.getEncoder() .encodeToString( String.format("%s:%s", name, password).getBytes(StandardCharsets.UTF_8))); } } BasicCallHeaderAuthenticator.java000066400000000000000000000070261515451615200371260ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth2; import java.io.UnsupportedEncodingException; import java.nio.charset.StandardCharsets; import java.util.Base64; import org.apache.arrow.flight.CallHeaders; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightRuntimeException; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** A ServerAuthHandler for username/password authentication. */ public class BasicCallHeaderAuthenticator implements CallHeaderAuthenticator { private static final Logger logger = LoggerFactory.getLogger(BasicCallHeaderAuthenticator.class); private final CredentialValidator authValidator; public BasicCallHeaderAuthenticator(CredentialValidator authValidator) { this.authValidator = authValidator; } @Override public AuthResult authenticate(CallHeaders incomingHeaders) { try { final String authEncoded = AuthUtilities.getValueFromAuthHeader(incomingHeaders, Auth2Constants.BASIC_PREFIX); if (authEncoded == null) { throw CallStatus.UNAUTHENTICATED.toRuntimeException(); } // The value has the format Base64(:) final String authDecoded = new String(Base64.getDecoder().decode(authEncoded), StandardCharsets.UTF_8); final int colonPos = authDecoded.indexOf(':'); if (colonPos == -1) { throw CallStatus.UNAUTHENTICATED.toRuntimeException(); } final String user = authDecoded.substring(0, colonPos); final String password = authDecoded.substring(colonPos + 1); return authValidator.validate(user, password); } catch (UnsupportedEncodingException ex) { // Note: Intentionally discarding the exception cause when reporting back to the client for // security purposes. logger.error("Authentication failed due to missing encoding.", ex); throw CallStatus.INTERNAL.toRuntimeException(); } catch (FlightRuntimeException ex) { throw ex; } catch (Exception ex) { // Note: Intentionally discarding the exception cause when reporting back to the client for // security purposes. logger.error("Authentication failed.", ex); throw CallStatus.UNAUTHENTICATED.toRuntimeException(); } } /** Interface that this handler delegates to for validating the incoming headers. */ public interface CredentialValidator { /** * Validate the supplied credentials (username/password) and return the peer identity. * * @param username The username to validate. * @param password The password to validate. * @return The peer identity if the supplied credentials are valid. * @throws Exception If the supplied credentials are not valid. */ AuthResult validate(String username, String password) throws Exception; } } BearerCredentialWriter.java000066400000000000000000000024761515451615200360410ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth2; import java.util.function.Consumer; import org.apache.arrow.flight.CallHeaders; /** Client credentials that use a bearer token. */ public final class BearerCredentialWriter implements Consumer { private final String bearer; public BearerCredentialWriter(String bearer) { this.bearer = bearer; } @Override public void accept(CallHeaders outputHeaders) { outputHeaders.insert( Auth2Constants.AUTHORIZATION_HEADER, Auth2Constants.BEARER_PREFIX + bearer); } } BearerTokenAuthenticator.java000066400000000000000000000046351515451615200364040ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth2; import org.apache.arrow.flight.CallHeaders; /** * Partial implementation of {@link CallHeaderAuthenticator} for bearer-token based authentication. */ public abstract class BearerTokenAuthenticator implements CallHeaderAuthenticator { final CallHeaderAuthenticator initialAuthenticator; public BearerTokenAuthenticator(CallHeaderAuthenticator initialAuthenticator) { this.initialAuthenticator = initialAuthenticator; } @Override public AuthResult authenticate(CallHeaders incomingHeaders) { // Check if headers contain a bearer token and if so, validate the token. final String bearerToken = AuthUtilities.getValueFromAuthHeader(incomingHeaders, Auth2Constants.BEARER_PREFIX); if (bearerToken != null) { return validateBearer(bearerToken); } // Delegate to the basic auth handler to do the validation. final CallHeaderAuthenticator.AuthResult result = initialAuthenticator.authenticate(incomingHeaders); return getAuthResultWithBearerToken(result); } /** * Callback to run when the initial authenticator succeeds. * * @param authResult A successful initial authentication result. * @return an alternate AuthResult based on the original AuthResult that will write a bearer token * to output headers. */ protected abstract AuthResult getAuthResultWithBearerToken(AuthResult authResult); /** * Validate the bearer token. * * @param bearerToken The bearer token to validate. * @return A successful AuthResult if validation succeeded. */ protected abstract AuthResult validateBearer(String bearerToken); } CallHeaderAuthenticator.java000066400000000000000000000070011515451615200361550ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth2; import org.apache.arrow.flight.CallHeaders; import org.apache.arrow.flight.FlightRuntimeException; /** * Interface for Server side authentication handlers. * *

    A CallHeaderAuthenticator is used by {@link ServerCallHeaderAuthMiddleware} to validate * headers sent by a Flight client for authentication purposes. The headers validated do not * necessarily have to be Authorization headers. * *

    The workflow is that the FlightServer will intercept headers on a request, validate the * headers, and either send back an UNAUTHENTICATED error, or succeed and potentially send back * additional headers to the client. * *

    Implementations of CallHeaderAuthenticator should take care not to provide leak confidential * details (such as indicating if usernames are valid or not) for security reasons when reporting * errors back to clients. * *

    Example CallHeaderAuthenticators provided include: The {@link BasicCallHeaderAuthenticator} * will authenticate basic HTTP credentials. * *

    The {@link BearerTokenAuthenticator} will authenticate basic HTTP credentials initially, then * also send back a bearer token that the client can use for subsequent requests. The {@link * GeneratedBearerTokenAuthenticator} will provide internally generated bearer tokens and maintain a * cache of them. */ public interface CallHeaderAuthenticator { /** * Encapsulates the result of the {@link CallHeaderAuthenticator} analysis of headers. * *

    This includes the identity of the incoming user and any outbound headers to send as a * response to the client. */ interface AuthResult { /** * The peer identity that was determined by the handshake process based on the authentication * credentials supplied by the client. * * @return The peer identity. */ String getPeerIdentity(); /** * Appends a header to the outgoing call headers. * * @param outgoingHeaders The outgoing headers. */ default void appendToOutgoingHeaders(CallHeaders outgoingHeaders) {} } /** * Validate the auth headers sent by the client. * * @param incomingHeaders The incoming headers to authenticate. * @return an auth result containing a peer identity and optionally a bearer token. * @throws FlightRuntimeException with CallStatus.UNAUTHENTICATED if credentials were not supplied * or if credentials were supplied but were not valid. */ AuthResult authenticate(CallHeaders incomingHeaders); /** An auth handler that does nothing. */ CallHeaderAuthenticator NO_OP = new CallHeaderAuthenticator() { @Override public AuthResult authenticate(CallHeaders incomingHeaders) { return () -> ""; } }; } ClientBearerHeaderHandler.java000066400000000000000000000027201515451615200364070ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth2; import org.apache.arrow.flight.CallHeaders; import org.apache.arrow.flight.grpc.CredentialCallOption; /** A client header handler that parses the incoming headers for a bearer token. */ public class ClientBearerHeaderHandler implements ClientHeaderHandler { @Override public CredentialCallOption getCredentialCallOptionFromIncomingHeaders( CallHeaders incomingHeaders) { final String bearerValue = AuthUtilities.getValueFromAuthHeader(incomingHeaders, Auth2Constants.BEARER_PREFIX); if (bearerValue != null) { return new CredentialCallOption(new BearerCredentialWriter(bearerValue)); } return null; } } ClientHandshakeWrapper.java000066400000000000000000000066611515451615200360370ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth2; import io.grpc.StatusRuntimeException; import io.grpc.stub.StreamObserver; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.flight.grpc.StatusUtils; import org.apache.arrow.flight.impl.Flight.HandshakeRequest; import org.apache.arrow.flight.impl.Flight.HandshakeResponse; import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceStub; /** Utility class for executing a handshake with a FlightServer. */ public class ClientHandshakeWrapper { private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(ClientHandshakeWrapper.class); /** * Do handshake for a client. The stub will be authenticated after this method returns. * * @param stub The service stub. */ public static void doClientHandshake(FlightServiceStub stub) { final HandshakeObserver observer = new HandshakeObserver(); try { observer.requestObserver = stub.handshake(observer); observer.requestObserver.onNext(HandshakeRequest.newBuilder().build()); observer.requestObserver.onCompleted(); try { if (!observer.completed.get()) { // TODO: ARROW-5681 throw CallStatus.UNAUTHENTICATED.toRuntimeException(); } } catch (InterruptedException ex) { Thread.currentThread().interrupt(); throw ex; } catch (ExecutionException ex) { final FlightRuntimeException wrappedException = StatusUtils.fromThrowable(ex.getCause()); logger.error("Failed on completing future", wrappedException); throw wrappedException; } } catch (StatusRuntimeException sre) { logger.error("Failed with SRE", sre); throw StatusUtils.fromGrpcRuntimeException(sre); } catch (Throwable ex) { logger.error("Failed with unknown", ex); if (ex instanceof FlightRuntimeException) { throw (FlightRuntimeException) ex; } throw StatusUtils.fromThrowable(ex); } } private static class HandshakeObserver implements StreamObserver { private volatile StreamObserver requestObserver; private final CompletableFuture completed; public HandshakeObserver() { super(); completed = new CompletableFuture<>(); } @Override public void onNext(HandshakeResponse value) {} @Override public void onError(Throwable t) { completed.completeExceptionally(t); } @Override public void onCompleted() { completed.complete(true); } } } ClientHeaderHandler.java000066400000000000000000000032101515451615200352610ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth2; import org.apache.arrow.flight.CallHeaders; import org.apache.arrow.flight.grpc.CredentialCallOption; /** Interface for client side header parsing and conversion to CredentialCallOption. */ public interface ClientHeaderHandler { /** * Parses the incoming headers and converts them into a CredentialCallOption. * * @param incomingHeaders Incoming headers to parse. * @return An instance of CredentialCallOption. */ CredentialCallOption getCredentialCallOptionFromIncomingHeaders(CallHeaders incomingHeaders); /** An client header handler that does nothing. */ ClientHeaderHandler NO_OP = new ClientHeaderHandler() { @Override public CredentialCallOption getCredentialCallOptionFromIncomingHeaders( CallHeaders incomingHeaders) { return null; } }; } ClientIncomingAuthHeaderMiddleware.java000066400000000000000000000051401515451615200402730ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth2; import org.apache.arrow.flight.CallHeaders; import org.apache.arrow.flight.CallInfo; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightClientMiddleware; import org.apache.arrow.flight.grpc.CredentialCallOption; /** Middleware for capturing bearer tokens sent back from the Flight server. */ public class ClientIncomingAuthHeaderMiddleware implements FlightClientMiddleware { private final Factory factory; /** Factory used within FlightClient. */ public static class Factory implements FlightClientMiddleware.Factory { private final ClientHeaderHandler headerHandler; private CredentialCallOption credentialCallOption = null; /** * Construct a factory with the given header handler. * * @param headerHandler The header handler that will be used for handling incoming headers from * the flight server. */ public Factory(ClientHeaderHandler headerHandler) { this.headerHandler = headerHandler; } @Override public FlightClientMiddleware onCallStarted(CallInfo info) { return new ClientIncomingAuthHeaderMiddleware(this); } void setCredentialCallOption(CredentialCallOption callOption) { this.credentialCallOption = callOption; } public CredentialCallOption getCredentialCallOption() { return credentialCallOption; } } private ClientIncomingAuthHeaderMiddleware(Factory factory) { this.factory = factory; } @Override public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) {} @Override public void onHeadersReceived(CallHeaders incomingHeaders) { factory.setCredentialCallOption( factory.headerHandler.getCredentialCallOptionFromIncomingHeaders(incomingHeaders)); } @Override public void onCallCompleted(CallStatus status) {} } GeneratedBearerTokenAuthenticator.java000066400000000000000000000123641515451615200402210ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth2; import com.google.common.base.Strings; import com.google.common.cache.Cache; import com.google.common.cache.CacheBuilder; import io.grpc.Metadata; import java.nio.ByteBuffer; import java.util.Base64; import java.util.UUID; import java.util.concurrent.TimeUnit; import org.apache.arrow.flight.CallHeaders; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.grpc.MetadataAdapter; /** Generates and caches bearer tokens from user credentials. */ public class GeneratedBearerTokenAuthenticator extends BearerTokenAuthenticator { private final Cache bearerToIdentityCache; /** * Generate bearer tokens for the given basic call authenticator. * * @param authenticator The authenticator to initial validate inputs with. */ public GeneratedBearerTokenAuthenticator(CallHeaderAuthenticator authenticator) { this(authenticator, CacheBuilder.newBuilder().expireAfterAccess(2, TimeUnit.HOURS)); } /** * Generate bearer tokens for the given basic call authenticator. * * @param authenticator The authenticator to initial validate inputs with. * @param timeoutMinutes The time before tokens expire after being accessed. */ public GeneratedBearerTokenAuthenticator( CallHeaderAuthenticator authenticator, int timeoutMinutes) { this( authenticator, CacheBuilder.newBuilder().expireAfterAccess(timeoutMinutes, TimeUnit.MINUTES)); } /** * Generate bearer tokens for the given basic call authenticator. * * @param authenticator The authenticator to initial validate inputs with. * @param cacheBuilder The configuration of the cache of bearer tokens. */ public GeneratedBearerTokenAuthenticator( CallHeaderAuthenticator authenticator, CacheBuilder cacheBuilder) { super(authenticator); bearerToIdentityCache = cacheBuilder.build(); } @Override protected AuthResult validateBearer(String bearerToken) { final String peerIdentity = bearerToIdentityCache.getIfPresent(bearerToken); if (peerIdentity == null) { throw CallStatus.UNAUTHENTICATED.toRuntimeException(); } return new AuthResult() { @Override public String getPeerIdentity() { return peerIdentity; } @Override public void appendToOutgoingHeaders(CallHeaders outgoingHeaders) { if (null == AuthUtilities.getValueFromAuthHeader( outgoingHeaders, Auth2Constants.BEARER_PREFIX)) { outgoingHeaders.insert( Auth2Constants.AUTHORIZATION_HEADER, Auth2Constants.BEARER_PREFIX + bearerToken); } } }; } @Override protected AuthResult getAuthResultWithBearerToken(AuthResult authResult) { // We generate a dummy header and call appendToOutgoingHeaders with it. // We then inspect the dummy header and parse the bearer token if present in the header // and generate a new bearer token if a bearer token is not present in the header. final CallHeaders dummyHeaders = new MetadataAdapter(new Metadata()); authResult.appendToOutgoingHeaders(dummyHeaders); String bearerToken = AuthUtilities.getValueFromAuthHeader(dummyHeaders, Auth2Constants.BEARER_PREFIX); final AuthResult authResultWithBearerToken; if (Strings.isNullOrEmpty(bearerToken)) { // Generate a new bearer token and return an AuthResult that can write it. final UUID uuid = UUID.randomUUID(); final ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[16]); byteBuffer.putLong(uuid.getMostSignificantBits()); byteBuffer.putLong(uuid.getLeastSignificantBits()); final String newToken = Base64.getEncoder().encodeToString(byteBuffer.array()); bearerToken = newToken; authResultWithBearerToken = new AuthResult() { @Override public String getPeerIdentity() { return authResult.getPeerIdentity(); } @Override public void appendToOutgoingHeaders(CallHeaders outgoingHeaders) { authResult.appendToOutgoingHeaders(outgoingHeaders); outgoingHeaders.insert( Auth2Constants.AUTHORIZATION_HEADER, Auth2Constants.BEARER_PREFIX + newToken); } }; } else { // Use the bearer token supplied by the original auth result. authResultWithBearerToken = authResult; } bearerToIdentityCache.put(bearerToken, authResult.getPeerIdentity()); return authResultWithBearerToken; } } ServerCallHeaderAuthMiddleware.java000066400000000000000000000051261515451615200374370ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/auth2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth2; import static org.apache.arrow.flight.auth2.CallHeaderAuthenticator.AuthResult; import org.apache.arrow.flight.CallHeaders; import org.apache.arrow.flight.CallInfo; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightServerMiddleware; import org.apache.arrow.flight.RequestContext; /** * Middleware that's used to validate credentials during the handshake and verify the bearer token * in subsequent requests. */ public class ServerCallHeaderAuthMiddleware implements FlightServerMiddleware { /** Factory for accessing ServerAuthMiddleware. */ public static class Factory implements FlightServerMiddleware.Factory { private final CallHeaderAuthenticator authHandler; /** * Construct a factory with the given auth handler. * * @param authHandler The auth handler what will be used for authenticating requests. */ public Factory(CallHeaderAuthenticator authHandler) { this.authHandler = authHandler; } @Override public ServerCallHeaderAuthMiddleware onCallStarted( CallInfo callInfo, CallHeaders incomingHeaders, RequestContext context) { final AuthResult result = authHandler.authenticate(incomingHeaders); context.put(Auth2Constants.PEER_IDENTITY_KEY, result.getPeerIdentity()); return new ServerCallHeaderAuthMiddleware(result); } } private final AuthResult authResult; public ServerCallHeaderAuthMiddleware(AuthResult authResult) { this.authResult = authResult; } @Override public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) { authResult.appendToOutgoingHeaders(outgoingHeaders); } @Override public void onCallCompleted(CallStatus status) {} @Override public void onCallErrored(Throwable err) {} } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/client/000077500000000000000000000000001515451615200311075ustar00rootroot00000000000000ClientCookieMiddleware.java000066400000000000000000000116521515451615200362460ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/client/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.client; import java.net.HttpCookie; import java.util.List; import java.util.Locale; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import java.util.stream.Collectors; import org.apache.arrow.flight.CallHeaders; import org.apache.arrow.flight.CallInfo; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightClientMiddleware; import org.apache.arrow.util.VisibleForTesting; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * A client middleware for receiving and sending cookie information. Note that this class will not * persist permanent cookies beyond the lifetime of this session. * *

    This middleware will automatically remove cookies that have expired. Note: Negative * max-age values currently do not get marked as expired due to a JDK issue. Use max-age=0 to * explicitly remove an existing cookie. */ public class ClientCookieMiddleware implements FlightClientMiddleware { private static final Logger LOGGER = LoggerFactory.getLogger(ClientCookieMiddleware.class); private static final String SET_COOKIE_HEADER = "Set-Cookie"; private static final String COOKIE_HEADER = "Cookie"; private final Factory factory; @VisibleForTesting ClientCookieMiddleware(Factory factory) { this.factory = factory; } /** Factory used within FlightClient. */ public static class Factory implements FlightClientMiddleware.Factory { // Use a map to track the most recent version of a cookie from the server. // Note that cookie names are case-sensitive (but header names aren't). private ConcurrentMap cookies = new ConcurrentHashMap<>(); @Override public ClientCookieMiddleware onCallStarted(CallInfo info) { return new ClientCookieMiddleware(this); } private void updateCookies(Iterable newCookieHeaderValues) { // Note: Intentionally overwrite existing cookie values. // A cookie defined once will continue to be used in all subsequent // requests on the client instance. The server can send the same cookie again // with a different value and the client will use the new value in future requests. // The server can also update a cookie to have an Expiry in the past or negative age // to signal that the client should stop using the cookie immediately. newCookieHeaderValues.forEach( headerValue -> { try { final List parsedCookies = HttpCookie.parse(headerValue); parsedCookies.forEach( parsedCookie -> { final String cookieNameLc = parsedCookie.getName().toLowerCase(Locale.ENGLISH); if (parsedCookie.hasExpired()) { cookies.remove(cookieNameLc); } else { cookies.put(parsedCookie.getName().toLowerCase(Locale.ENGLISH), parsedCookie); } }); } catch (IllegalArgumentException ex) { LOGGER.warn( "Skipping incorrectly formatted Set-Cookie header with value '{}'.", headerValue); } }); } } @Override public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) { final String cookieValue = getValidCookiesAsString(); if (!cookieValue.isEmpty()) { outgoingHeaders.insert(COOKIE_HEADER, cookieValue); } } @Override public void onHeadersReceived(CallHeaders incomingHeaders) { factory.updateCookies(incomingHeaders.getAll(SET_COOKIE_HEADER)); } @Override public void onCallCompleted(CallStatus status) {} /** Discards expired cookies and returns the valid cookies as a String delimited by ';'. */ @VisibleForTesting String getValidCookiesAsString() { // Discard expired cookies. factory.cookies.entrySet().removeIf(cookieEntry -> cookieEntry.getValue().hasExpired()); // Cookie header value format: // [=; = cookie.getValue().toString()) .collect(Collectors.joining("; ")); } } apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/000077500000000000000000000000001515451615200305645ustar00rootroot00000000000000AddWritableBuffer.java000066400000000000000000000106251515451615200346700ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.netty.buffer.ByteBuf; import java.io.IOException; import java.io.OutputStream; import java.lang.reflect.Constructor; import java.lang.reflect.Field; import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Method; import java.util.List; /** * Allow a user to add a ByteBuf based InputStream directly into GRPC WritableBuffer to avoid an * extra copy. This could be solved in GRPC by adding a ByteBufListable interface on InputStream and * letting BufferChainOutputStream take advantage of it. */ public class AddWritableBuffer { private static final Constructor bufConstruct; private static final Field bufferList; private static final Field current; private static final Method listAdd; private static final Class bufChainOut; static { Constructor tmpConstruct = null; Field tmpBufferList = null; Field tmpCurrent = null; Class tmpBufChainOut = null; Method tmpListAdd = null; try { Class nwb = Class.forName("io.grpc.netty.NettyWritableBuffer"); Constructor tmpConstruct2 = nwb.getDeclaredConstructor(ByteBuf.class); tmpConstruct2.setAccessible(true); Class tmpBufChainOut2 = Class.forName("io.grpc.internal.MessageFramer$BufferChainOutputStream"); Field tmpBufferList2 = tmpBufChainOut2.getDeclaredField("bufferList"); tmpBufferList2.setAccessible(true); Field tmpCurrent2 = tmpBufChainOut2.getDeclaredField("current"); tmpCurrent2.setAccessible(true); Method tmpListAdd2 = List.class.getDeclaredMethod("add", Object.class); // output fields last. tmpConstruct = tmpConstruct2; tmpBufferList = tmpBufferList2; tmpCurrent = tmpCurrent2; tmpListAdd = tmpListAdd2; tmpBufChainOut = tmpBufChainOut2; } catch (Exception ex) { new RuntimeException("Failed to initialize AddWritableBuffer, falling back to slow path", ex) .printStackTrace(); } bufConstruct = tmpConstruct; bufferList = tmpBufferList; current = tmpCurrent; listAdd = tmpListAdd; bufChainOut = tmpBufChainOut; } /** * Add the provided ByteBuf to the gRPC BufferChainOutputStream if possible, else copy the buffer * to the stream. * * @param buf The buffer to add. * @param stream The Candidate OutputStream to add to. * @param tryZeroCopy If true, try to zero-copy append the buffer to the stream. This may not * succeed. * @return True if buffer was zero-copy added to the stream. False if the buffer was copied. * @throws IOException if the fast path is not enabled and there was an error copying the buffer * to the stream. */ public static boolean add(ByteBuf buf, OutputStream stream, boolean tryZeroCopy) throws IOException { if (!tryZeroCopy || !tryAddBuffer(buf, stream)) { buf.getBytes(0, stream, buf.readableBytes()); return false; } return true; } private static boolean tryAddBuffer(ByteBuf buf, OutputStream stream) throws IOException { if (bufChainOut == null) { return false; } if (!stream.getClass().equals(bufChainOut)) { return false; } try { if (current.get(stream) != null) { return false; } buf.retain(); Object obj = bufConstruct.newInstance(buf); Object list = bufferList.get(stream); listAdd.invoke(list, obj); current.set(stream, obj); return true; } catch (IllegalAccessException | IllegalArgumentException | InvocationTargetException | InstantiationException e) { e.printStackTrace(); return false; } } } CallCredentialAdapter.java000066400000000000000000000037751515451615200355330ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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; /** Adapter class to utilize a CredentialWriter to implement Grpc CallCredentials. */ public class CallCredentialAdapter extends CallCredentials { private final Consumer credentialWriter; public CallCredentialAdapter(Consumer credentialWriter) { this.credentialWriter = credentialWriter; } @Override public void applyRequestMetadata( RequestInfo requestInfo, Executor executor, MetadataApplier metadataApplier) { executor.execute( () -> { 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())); } }); } @Override public void thisUsesUnstableApi() { // Mandatory to override this to acknowledge that CallCredentials is Experimental. } } ClientInterceptorAdapter.java000066400000000000000000000132471515451615200363150ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.CallOptions; import io.grpc.Channel; import io.grpc.ClientCall; import io.grpc.ClientInterceptor; import io.grpc.ForwardingClientCall.SimpleForwardingClientCall; import io.grpc.ForwardingClientCallListener.SimpleForwardingClientCallListener; import io.grpc.Metadata; import io.grpc.MethodDescriptor; import io.grpc.Status; import io.grpc.StatusRuntimeException; import java.util.ArrayList; import java.util.List; import org.apache.arrow.flight.CallInfo; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightClientMiddleware; import org.apache.arrow.flight.FlightMethod; import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.flight.FlightStatusCode; /** * An adapter between Flight client middleware and gRPC interceptors. * *

    This is implemented as a single gRPC interceptor that runs all Flight client middleware * sequentially. */ public class ClientInterceptorAdapter implements ClientInterceptor { private final List factories; public ClientInterceptorAdapter(List factories) { this.factories = factories; } @Override public ClientCall interceptCall( MethodDescriptor method, CallOptions callOptions, Channel next) { final List middleware = new ArrayList<>(); final CallInfo info = new CallInfo(FlightMethod.fromProtocol(method.getFullMethodName())); try { for (final FlightClientMiddleware.Factory factory : factories) { middleware.add(factory.onCallStarted(info)); } } catch (FlightRuntimeException e) { // Explicitly propagate throw e; } catch (StatusRuntimeException e) { throw StatusUtils.fromGrpcRuntimeException(e); } catch (RuntimeException e) { throw StatusUtils.fromThrowable(e); } return new FlightClientCall<>(next.newCall(method, callOptions), middleware); } /** * The ClientCallListener which hooks into the gRPC request cycle and actually runs middleware at * certain points. */ private static class FlightClientCallListener extends SimpleForwardingClientCallListener { private final List middleware; boolean receivedHeaders; public FlightClientCallListener( ClientCall.Listener responseListener, List middleware) { super(responseListener); this.middleware = middleware; receivedHeaders = false; } @Override public void onHeaders(Metadata headers) { receivedHeaders = true; final MetadataAdapter adapter = new MetadataAdapter(headers); try { middleware.forEach(m -> m.onHeadersReceived(adapter)); } finally { // Make sure to always call the gRPC callback to avoid interrupting the gRPC request cycle super.onHeaders(headers); } } @Override public void onClose(Status status, Metadata trailers) { try { if (!receivedHeaders) { // gRPC doesn't always send response headers if the call errors or completes immediately, // but instead // consolidates them with the trailers. If we never got headers, assume this happened and // run the header // callback with the trailers. final MetadataAdapter adapter = new MetadataAdapter(trailers); middleware.forEach(m -> m.onHeadersReceived(adapter)); } final CallStatus flightStatus = StatusUtils.fromGrpcStatusAndTrailers(status, trailers); middleware.forEach(m -> m.onCallCompleted(flightStatus)); } finally { // Make sure to always call the gRPC callback to avoid interrupting the gRPC request cycle super.onClose(status, trailers); } } } /** * The gRPC ClientCall which hooks into the gRPC request cycle and injects our ClientCallListener. */ private static class FlightClientCall extends SimpleForwardingClientCall { private final List middleware; public FlightClientCall( ClientCall clientCall, List middleware) { super(clientCall); this.middleware = middleware; } @Override public void start(Listener responseListener, Metadata headers) { final MetadataAdapter metadataAdapter = new MetadataAdapter(headers); middleware.forEach(m -> m.onBeforeSendingHeaders(metadataAdapter)); super.start(new FlightClientCallListener<>(responseListener, middleware), headers); } @Override public void cancel(String message, Throwable cause) { final CallStatus flightStatus = new CallStatus(FlightStatusCode.CANCELLED, cause, message, null); middleware.forEach(m -> m.onCallCompleted(flightStatus)); super.cancel(message, cause); } } } ContextPropagatingExecutorService.java000066400000000000000000000077031515451615200402370ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.Context; import java.util.Collection; import java.util.List; import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.stream.Collectors; /** * An {@link ExecutorService} that propagates the {@link Context}. * *

    Context is used to propagate per-call state, like the authenticated user, between threads (as * gRPC makes no guarantees about what thread things execute on). This wrapper makes it easy to * preserve this when using an Executor. The Context itself is immutable, so it is thread-safe. */ public class ContextPropagatingExecutorService implements ExecutorService { private final ExecutorService delegate; public ContextPropagatingExecutorService(ExecutorService delegate) { this.delegate = delegate; } // These are just delegate methods. @Override public void shutdown() { delegate.shutdown(); } @Override public List shutdownNow() { return delegate.shutdownNow(); } @Override public boolean isShutdown() { return delegate.isShutdown(); } @Override public boolean isTerminated() { return delegate.isTerminated(); } @Override public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { return delegate.awaitTermination(timeout, unit); } // These are delegate methods that wrap the submitted task in the current gRPC Context. @Override public Future submit(Callable task) { return delegate.submit(Context.current().wrap(task)); } @Override public Future submit(Runnable task, T result) { return delegate.submit(Context.current().wrap(task), result); } @Override public Future submit(Runnable task) { return delegate.submit(Context.current().wrap(task)); } @Override public List> invokeAll(Collection> tasks) throws InterruptedException { return delegate.invokeAll( tasks.stream().map(Context.current()::wrap).collect(Collectors.toList())); } @Override public List> invokeAll( Collection> tasks, long timeout, TimeUnit unit) throws InterruptedException { return delegate.invokeAll( tasks.stream().map(Context.current()::wrap).collect(Collectors.toList()), timeout, unit); } @Override public T invokeAny(Collection> tasks) throws InterruptedException, ExecutionException { return delegate.invokeAny( tasks.stream().map(Context.current()::wrap).collect(Collectors.toList())); } @Override public T invokeAny(Collection> tasks, long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException { return delegate.invokeAny( tasks.stream().map(Context.current()::wrap).collect(Collectors.toList()), timeout, unit); } @Override public void execute(Runnable command) { delegate.execute(Context.current().wrap(command)); } } CredentialCallOption.java000066400000000000000000000027031515451615200354110ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.stub.AbstractStub; import java.util.function.Consumer; import org.apache.arrow.flight.CallHeaders; import org.apache.arrow.flight.CallOptions; /** Method option for supplying credentials to method calls. */ public class CredentialCallOption implements CallOptions.GrpcCallOption { private final Consumer credentialWriter; public CredentialCallOption(Consumer credentialWriter) { this.credentialWriter = credentialWriter; } @Override public > T wrapStub(T stub) { return stub.withCallCredentials(new CallCredentialAdapter(credentialWriter)); } } GetReadableBuffer.java000066400000000000000000000066701515451615200346520ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 com.google.common.base.Throwables; import com.google.common.io.ByteStreams; import io.grpc.internal.ReadableBuffer; import java.io.IOException; import java.io.InputStream; import java.lang.reflect.Field; import org.apache.arrow.memory.ArrowBuf; /** * Enable access to ReadableBuffer directly to copy data from a BufferInputStream into a target * ByteBuffer/ByteBuf. * *

    This could be solved by BufferInputStream exposing Drainable. */ public class GetReadableBuffer { private static final Field READABLE_BUFFER; private static final Class BUFFER_INPUT_STREAM; static { Field tmpField = null; Class tmpClazz = null; try { Class clazz = Class.forName("io.grpc.internal.ReadableBuffers$BufferInputStream"); Field f = clazz.getDeclaredField("buffer"); f.setAccessible(true); // don't set until we've gotten past all exception cases. tmpField = f; tmpClazz = clazz; } catch (Exception e) { new RuntimeException("Failed to initialize GetReadableBuffer, falling back to slow path", e) .printStackTrace(); } READABLE_BUFFER = tmpField; BUFFER_INPUT_STREAM = tmpClazz; } /** * Extracts the ReadableBuffer for the given input stream. * * @param is Must be an instance of io.grpc.internal.ReadableBuffers$BufferInputStream or null * will be returned. */ public static ReadableBuffer getReadableBuffer(InputStream is) { if (BUFFER_INPUT_STREAM == null || !is.getClass().equals(BUFFER_INPUT_STREAM)) { return null; } try { return (ReadableBuffer) READABLE_BUFFER.get(is); } catch (Exception ex) { throw Throwables.propagate(ex); } } /** * Helper method to read a gRPC-provided InputStream into an ArrowBuf. * * @param stream The stream to read from. Should be an instance of {@link #BUFFER_INPUT_STREAM}. * @param buf The buffer to read into. * @param size The number of bytes to read. * @param fastPath Whether to enable the fast path (i.e. detect whether the stream is a {@link * #BUFFER_INPUT_STREAM}). * @throws IOException if there is an error reading form the stream */ 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(heapBytes, 0, size); } else { ByteStreams.readFully(stream, heapBytes); } buf.writeBytes(heapBytes); buf.writerIndex(size); } } MetadataAdapter.java000066400000000000000000000065371515451615200344040ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.Metadata; import java.nio.charset.StandardCharsets; import java.util.Collections; import java.util.HashSet; import java.util.Set; import java.util.stream.Collectors; import java.util.stream.StreamSupport; import org.apache.arrow.flight.CallHeaders; /** * A mutable adapter between the gRPC Metadata object and the Flight headers interface. * *

    This allows us to present the headers (metadata) from gRPC without copying to/from our own * object. */ public class MetadataAdapter implements CallHeaders { private final Metadata metadata; public MetadataAdapter(Metadata metadata) { this.metadata = metadata; } @Override public String get(String key) { return this.metadata.get(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER)); } @Override public byte[] getByte(String key) { if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) { return this.metadata.get(Metadata.Key.of(key, Metadata.BINARY_BYTE_MARSHALLER)); } return get(key).getBytes(StandardCharsets.UTF_8); } @Override public Iterable getAll(String key) { 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)) { 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) .collect(Collectors.toList()); } @Override public void insert(String key, String value) { this.metadata.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value); } @Override public void insert(String key, byte[] value) { this.metadata.put(Metadata.Key.of(key, Metadata.BINARY_BYTE_MARSHALLER), value); } @Override public Set keys() { return new HashSet<>(this.metadata.keys()); } @Override public boolean containsKey(String key) { if (key.endsWith("-bin")) { final Metadata.Key grpcKey = Metadata.Key.of(key, Metadata.BINARY_BYTE_MARSHALLER); return this.metadata.containsKey(grpcKey); } final Metadata.Key grpcKey = Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER); return this.metadata.containsKey(grpcKey); } @Override public String toString() { return this.metadata.toString(); } } NettyClientBuilder.java000066400000000000000000000210221515451615200351160ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; } } RequestContextAdapter.java000066400000000000000000000032771515451615200356570ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.Context; import java.util.HashMap; import java.util.Set; import org.apache.arrow.flight.RequestContext; /** Adapter for holding key value pairs. */ public class RequestContextAdapter implements RequestContext { public static final Context.Key REQUEST_CONTEXT_KEY = Context.key("arrow-flight-request-context"); private final HashMap map = new HashMap<>(); @Override public void put(String key, String value) { if (map.putIfAbsent(key, value) != null) { throw new IllegalArgumentException( "Duplicate write to a RequestContext at key " + key + " not allowed."); } } @Override public String get(String key) { return map.get(key); } @Override public Set keySet() { return map.keySet(); } @Override public String remove(String key) { return map.remove(key); } } ServerBackpressureThresholdInterceptor.java000066400000000000000000000030551515451615200412670ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.Metadata; import io.grpc.ServerCall; import io.grpc.ServerCallHandler; import io.grpc.ServerInterceptor; /** * An interceptor for specifying the number of bytes that can be queued before a call with an output * stream gets blocked by backpressure. */ public class ServerBackpressureThresholdInterceptor implements ServerInterceptor { private final int numBytes; public ServerBackpressureThresholdInterceptor(int numBytes) { this.numBytes = numBytes; } @Override public ServerCall.Listener interceptCall( ServerCall call, Metadata headers, ServerCallHandler next) { call.setOnReadyThreshold(numBytes); return next.startCall(call, headers); } } ServerInterceptorAdapter.java000066400000000000000000000143171515451615200363440ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.Context; import io.grpc.Contexts; import io.grpc.ForwardingServerCall.SimpleForwardingServerCall; import io.grpc.Metadata; import io.grpc.ServerCall; import io.grpc.ServerCall.Listener; import io.grpc.ServerCallHandler; import io.grpc.ServerInterceptor; import io.grpc.Status; import java.util.ArrayList; import java.util.Collections; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import org.apache.arrow.flight.CallInfo; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightConstants; import org.apache.arrow.flight.FlightMethod; import org.apache.arrow.flight.FlightProducer.CallContext; import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.flight.FlightServerMiddleware; import org.apache.arrow.flight.FlightServerMiddleware.Key; /** * An adapter between Flight middleware and a gRPC interceptor. * *

    This is implemented as a single gRPC interceptor that runs all Flight server middleware * sequentially. Flight middleware instances are stored in the gRPC Context so their state is * accessible later. */ public class ServerInterceptorAdapter implements ServerInterceptor { /** * A combination of a middleware Key and factory. * * @param The middleware type. */ public static class KeyFactory { private final FlightServerMiddleware.Key key; private final FlightServerMiddleware.Factory factory; public KeyFactory( FlightServerMiddleware.Key key, FlightServerMiddleware.Factory factory) { this.key = key; this.factory = factory; } } /** * The {@link Context.Key} that stores the Flight middleware active for a particular call. * *

    Applications should not use this directly. Instead, see {@link * CallContext#getMiddleware(Key)}. */ public static final Context.Key, FlightServerMiddleware>> SERVER_MIDDLEWARE_KEY = Context.key("arrow.flight.server_middleware"); private final List> factories; public ServerInterceptorAdapter(List> factories) { this.factories = factories; } @Override public Listener interceptCall( ServerCall call, Metadata headers, ServerCallHandler next) { if (!FlightConstants.SERVICE.equals(call.getMethodDescriptor().getServiceName())) { return Contexts.interceptCall(Context.current(), call, headers, next); } final CallInfo info = new CallInfo(FlightMethod.fromProtocol(call.getMethodDescriptor().getFullMethodName())); final List middleware = new ArrayList<>(); // Use LinkedHashMap to preserve insertion order final Map, FlightServerMiddleware> middlewareMap = new LinkedHashMap<>(); final MetadataAdapter headerAdapter = new MetadataAdapter(headers); final RequestContextAdapter requestContextAdapter = new RequestContextAdapter(); for (final KeyFactory factory : factories) { final FlightServerMiddleware m; try { m = factory.factory.onCallStarted(info, headerAdapter, requestContextAdapter); } catch (FlightRuntimeException e) { // Cancel call call.close(StatusUtils.toGrpcStatus(e.status()), new Metadata()); return new Listener() {}; } middleware.add(m); middlewareMap.put(factory.key, m); } // Inject the middleware into the context so RPC method implementations can communicate with // middleware instances final Context contextWithMiddlewareAndRequestsOptions = Context.current() .withValue(SERVER_MIDDLEWARE_KEY, Collections.unmodifiableMap(middlewareMap)) .withValue(RequestContextAdapter.REQUEST_CONTEXT_KEY, requestContextAdapter); final SimpleForwardingServerCall forwardingServerCall = new SimpleForwardingServerCall(call) { boolean sentHeaders = false; @Override public void sendHeaders(Metadata headers) { sentHeaders = true; try { final MetadataAdapter headerAdapter = new MetadataAdapter(headers); middleware.forEach(m -> m.onBeforeSendingHeaders(headerAdapter)); } finally { // Make sure to always call the gRPC callback to avoid interrupting the gRPC request // cycle super.sendHeaders(headers); } } @Override public void close(Status status, Metadata trailers) { try { if (!sentHeaders) { // gRPC doesn't always send response headers if the call errors or completes // immediately final MetadataAdapter headerAdapter = new MetadataAdapter(trailers); middleware.forEach(m -> m.onBeforeSendingHeaders(headerAdapter)); } } finally { // Make sure to always call the gRPC callback to avoid interrupting the gRPC request // cycle super.close(status, trailers); } final CallStatus flightStatus = StatusUtils.fromGrpcStatus(status); middleware.forEach(m -> m.onCallCompleted(flightStatus)); } }; return Contexts.interceptCall( contextWithMiddlewareAndRequestsOptions, forwardingServerCall, headers, next); } } StatusUtils.java000066400000000000000000000215241515451615200336600ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.InternalMetadata; import io.grpc.Metadata; import io.grpc.Status; import io.grpc.Status.Code; import io.grpc.StatusException; import io.grpc.StatusRuntimeException; import java.nio.charset.StandardCharsets; import java.util.Iterator; import java.util.Objects; import java.util.function.Function; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.ErrorFlightMetadata; import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.flight.FlightStatusCode; /** * Utilities to adapt gRPC and Flight status objects. * *

    NOT A PUBLIC CLASS, interface is not guaranteed to remain stable. */ public class StatusUtils { private StatusUtils() { throw new AssertionError("Do not instantiate this class."); } /** Convert from a Flight status code to a gRPC status code. */ public static Status.Code toGrpcStatusCode(FlightStatusCode code) { switch (code) { case OK: return Code.OK; case UNKNOWN: return Code.UNKNOWN; case INTERNAL: return Code.INTERNAL; case INVALID_ARGUMENT: return Code.INVALID_ARGUMENT; case TIMED_OUT: return Code.DEADLINE_EXCEEDED; case NOT_FOUND: return Code.NOT_FOUND; case ALREADY_EXISTS: return Code.ALREADY_EXISTS; case CANCELLED: return Code.CANCELLED; case UNAUTHENTICATED: return Code.UNAUTHENTICATED; case UNAUTHORIZED: return Code.PERMISSION_DENIED; case UNIMPLEMENTED: return Code.UNIMPLEMENTED; case UNAVAILABLE: return Code.UNAVAILABLE; case RESOURCE_EXHAUSTED: return Code.RESOURCE_EXHAUSTED; default: return Code.UNKNOWN; } } /** Convert from a gRPC status code to a Flight status code. */ public static FlightStatusCode fromGrpcStatusCode(Status.Code code) { switch (code) { case OK: return FlightStatusCode.OK; case CANCELLED: return FlightStatusCode.CANCELLED; case UNKNOWN: return FlightStatusCode.UNKNOWN; case INVALID_ARGUMENT: return FlightStatusCode.INVALID_ARGUMENT; case DEADLINE_EXCEEDED: return FlightStatusCode.TIMED_OUT; case NOT_FOUND: return FlightStatusCode.NOT_FOUND; case ALREADY_EXISTS: return FlightStatusCode.ALREADY_EXISTS; case PERMISSION_DENIED: return FlightStatusCode.UNAUTHORIZED; case RESOURCE_EXHAUSTED: return FlightStatusCode.RESOURCE_EXHAUSTED; case FAILED_PRECONDITION: return FlightStatusCode.INVALID_ARGUMENT; case ABORTED: return FlightStatusCode.INTERNAL; case OUT_OF_RANGE: return FlightStatusCode.INVALID_ARGUMENT; case UNIMPLEMENTED: return FlightStatusCode.UNIMPLEMENTED; case INTERNAL: return FlightStatusCode.INTERNAL; case UNAVAILABLE: return FlightStatusCode.UNAVAILABLE; case DATA_LOSS: return FlightStatusCode.INTERNAL; case UNAUTHENTICATED: return FlightStatusCode.UNAUTHENTICATED; default: return FlightStatusCode.UNKNOWN; } } /** Create Metadata Key for binary metadata. */ static Metadata.Key keyOfBinary(String name) { return Metadata.Key.of(name, Metadata.BINARY_BYTE_MARSHALLER); } /** Create Metadata Key for ascii metadata. */ static Metadata.Key keyOfAscii(String name) { // Use InternalMetadata for keys that start with ":", e.g. ":status". See ARROW-14014. return InternalMetadata.keyOf(name, Metadata.ASCII_STRING_MARSHALLER); } /** Convert from a gRPC Status & trailers to a Flight status. */ public static CallStatus fromGrpcStatusAndTrailers(Status status, Metadata trailers) { // gRPC may not always have trailers - this happens when the server internally generates an // error, which is rare, // but can happen. final ErrorFlightMetadata errorMetadata = trailers == null ? null : parseTrailers(trailers); return new CallStatus( fromGrpcStatusCode(status.getCode()), status.getCause(), status.getDescription(), errorMetadata); } /** Convert from a gRPC status to a Flight status. */ public static CallStatus fromGrpcStatus(Status status) { return new CallStatus( fromGrpcStatusCode(status.getCode()), status.getCause(), status.getDescription(), null); } /** Convert from a Flight status to a gRPC status. */ public static Status toGrpcStatus(CallStatus status) { return toGrpcStatusCode(status.code()) .toStatus() .withDescription(status.description()) .withCause(status.cause()); } /** Convert from a gRPC exception to a Flight exception. */ public static FlightRuntimeException fromGrpcRuntimeException(StatusRuntimeException sre) { return fromGrpcStatusAndTrailers(sre.getStatus(), sre.getTrailers()).toRuntimeException(); } /** Convert gRPC trailers into Flight error metadata. */ private static ErrorFlightMetadata parseTrailers(Metadata trailers) { ErrorFlightMetadata metadata = new ErrorFlightMetadata(); for (String key : trailers.keys()) { if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) { metadata.insert(key, trailers.get(keyOfBinary(key))); } else { metadata.insert( key, Objects.requireNonNull(trailers.get(keyOfAscii(key))).getBytes(StandardCharsets.UTF_8)); } } return metadata; } /** Convert arbitrary exceptions to a {@link FlightRuntimeException}. */ public static FlightRuntimeException fromThrowable(Throwable t) { if (t instanceof StatusRuntimeException) { return fromGrpcRuntimeException((StatusRuntimeException) t); } else if (t instanceof FlightRuntimeException) { return (FlightRuntimeException) t; } return CallStatus.UNKNOWN.withCause(t).withDescription(t.getMessage()).toRuntimeException(); } /** * Convert arbitrary exceptions to a {@link StatusRuntimeException} or {@link StatusException}. * *

    Such exceptions can be passed to {@link io.grpc.stub.StreamObserver#onError(Throwable)} and * will give the client a reasonable error message. */ public static Throwable toGrpcException(Throwable ex) { if (ex instanceof StatusRuntimeException) { return ex; } else if (ex instanceof StatusException) { return ex; } else if (ex instanceof FlightRuntimeException) { final FlightRuntimeException fre = (FlightRuntimeException) ex; if (fre.status().metadata() != null) { Metadata trailers = toGrpcMetadata(fre.status().metadata()); return new StatusRuntimeException(toGrpcStatus(fre.status()), trailers); } return toGrpcStatus(fre.status()).asRuntimeException(); } return Status.INTERNAL .withCause(ex) .withDescription("There was an error servicing your request.") .asRuntimeException(); } private static Metadata toGrpcMetadata(ErrorFlightMetadata metadata) { final Metadata trailers = new Metadata(); for (final String key : metadata.keys()) { if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) { trailers.put(keyOfBinary(key), metadata.getByte(key)); } else { trailers.put(keyOfAscii(key), metadata.get(key)); } } return trailers; } /** * Maps a transformation function to the elements of an iterator, while wrapping exceptions in * {@link FlightRuntimeException}. */ public static Iterator wrapIterator( Iterator fromIterator, Function transformer) { Objects.requireNonNull(fromIterator); Objects.requireNonNull(transformer); return new Iterator() { @Override public boolean hasNext() { try { return fromIterator.hasNext(); } catch (StatusRuntimeException e) { throw fromGrpcRuntimeException(e); } } @Override public TO next() { try { return transformer.apply(fromIterator.next()); } catch (StatusRuntimeException e) { throw fromGrpcRuntimeException(e); } } }; } } apache-arrow-java-19.0.0/flight/flight-core/src/test/000077500000000000000000000000001515451615200223245ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/000077500000000000000000000000001515451615200232455ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/000077500000000000000000000000001515451615200240345ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/000077500000000000000000000000001515451615200252555ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200264075ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/000077500000000000000000000000001515451615200276645ustar00rootroot00000000000000FlightTestUtil.java000066400000000000000000000072121515451615200333650ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import java.io.File; import java.lang.reflect.InvocationTargetException; import java.nio.file.Files; import java.nio.file.Path; import java.util.Arrays; import java.util.List; import org.apache.arrow.vector.test.util.ArrowTestDataUtil; import org.junit.jupiter.api.function.Executable; /** Utility methods and constants for testing flight servers. */ public class FlightTestUtil { public static final String LOCALHOST = "localhost"; static Path getFlightTestDataRoot() { return ArrowTestDataUtil.getTestDataRoot().resolve("flight"); } static Path exampleTlsRootCert() { return getFlightTestDataRoot().resolve("root-ca.pem"); } static List exampleTlsCerts() { final Path root = getFlightTestDataRoot(); final Path cert0Pem = root.resolve("cert0.pem"); if (!Files.exists(cert0Pem)) { throw new RuntimeException( cert0Pem + " doesn't exist. Make sure submodules are initialized (see https://arrow.apache.org/docs/dev/developers/java/building.html#building)"); } return Arrays.asList( new CertKeyPair(cert0Pem.toFile(), root.resolve("cert0.pkcs1").toFile()), new CertKeyPair(root.resolve("cert1.pem").toFile(), root.resolve("cert1.pkcs1").toFile())); } static boolean isEpollAvailable() { try { Class epoll = Class.forName("io.netty.channel.epoll.Epoll"); return (Boolean) epoll.getMethod("isAvailable").invoke(null); } catch (ClassNotFoundException | NoSuchMethodException | IllegalAccessException | InvocationTargetException e) { return false; } } static boolean isKqueueAvailable() { try { Class kqueue = Class.forName("io.netty.channel.kqueue.KQueue"); return (Boolean) kqueue.getMethod("isAvailable").invoke(null); } catch (ClassNotFoundException | NoSuchMethodException | IllegalAccessException | InvocationTargetException e) { return false; } } static boolean isNativeTransportAvailable() { return isEpollAvailable() || isKqueueAvailable(); } /** * Assert that the given runnable fails with a Flight exception of the given code. * * @param code The expected Flight status code. * @param r The code to run. * @return The thrown status. */ public static CallStatus assertCode(FlightStatusCode code, Executable r) { final FlightRuntimeException ex = assertThrows(FlightRuntimeException.class, r); assertEquals(code, ex.status().code()); return ex.status(); } public static class CertKeyPair { public final File cert; public final File key; public CertKeyPair(File cert, File key) { this.cert = cert; this.key = key; } } private FlightTestUtil() {} } TestActionType.java000066400000000000000000000023551515451615200333740ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.junit.jupiter.api.Assertions.assertEquals; import org.junit.jupiter.api.Test; public class TestActionType { @Test void testActionType() { final String type = "MyType"; final String description = "My Description"; final ActionType actionType = new ActionType(type, description); assertEquals(type, actionType.getType()); assertEquals(description, actionType.getDescription()); } } TestApplicationMetadata.java000066400000000000000000000335521515451615200352240ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; 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.assertNotNull; import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.Collections; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import java.util.function.BiConsumer; import org.apache.arrow.flight.FlightClient.PutListener; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; 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.Schema; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; /** Tests for application-specific metadata support in Flight. */ public class TestApplicationMetadata { // The command used to trigger the test for ARROW-6136. private static final byte[] COMMAND_ARROW_6136 = "ARROW-6136".getBytes(StandardCharsets.UTF_8); // The expected error message. private static final String MESSAGE_ARROW_6136 = "The stream should not be double-closed."; /** Ensure that a client can read the metadata sent from the server. */ @Test // This test is consistently flaky on CI, unfortunately. @Disabled public void retrieveMetadata() { test( (allocator, client) -> { try (final FlightStream stream = client.getStream(new Ticket(new byte[0]))) { byte i = 0; while (stream.next()) { final IntVector vector = (IntVector) stream.getRoot().getVector("a"); assertEquals(1, vector.getValueCount()); assertEquals(10, vector.get(0)); assertEquals(i, stream.getLatestMetadata().getByte(0)); i++; } } catch (Exception e) { throw new RuntimeException(e); } }); } /** * ARROW-6136: make sure that the Flight implementation doesn't double-close the server-to-client * stream. */ @Test public void arrow6136() { final Schema schema = new Schema(Collections.emptyList()); test( (allocator, client) -> { try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final FlightDescriptor descriptor = FlightDescriptor.command(COMMAND_ARROW_6136); final PutListener listener = new SyncPutListener(); final FlightClient.ClientStreamListener writer = client.startPut(descriptor, root, listener); // Must attempt to retrieve the result to get any server-side errors. final CallStatus status = FlightTestUtil.assertCode(FlightStatusCode.INTERNAL, writer::getResult); assertEquals(MESSAGE_ARROW_6136, status.description()); } catch (Exception e) { throw new RuntimeException(e); } }); } /** Ensure that a client can send metadata to the server. */ @Test @Disabled public void uploadMetadataAsync() { final Schema schema = new Schema(Collections.singletonList(Field.nullable("a", new ArrowType.Int(32, true)))); test( (allocator, client) -> { try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final FlightDescriptor descriptor = FlightDescriptor.path("test"); final PutListener listener = new AsyncPutListener() { int counter = 0; @Override public void onNext(PutResult val) { assertNotNull(val); assertEquals(counter, val.getApplicationMetadata().getByte(0)); counter++; } }; final FlightClient.ClientStreamListener writer = client.startPut(descriptor, root, listener); root.allocateNew(); for (byte i = 0; i < 10; i++) { final IntVector vector = (IntVector) root.getVector("a"); final ArrowBuf metadata = allocator.buffer(1); metadata.writeByte(i); vector.set(0, 10); vector.setValueCount(1); root.setRowCount(1); writer.putNext(metadata); } writer.completed(); // Must attempt to retrieve the result to get any server-side errors. writer.getResult(); } catch (Exception e) { throw new RuntimeException(e); } }); } /** Ensure that a client can send metadata to the server. Uses the synchronous API. */ @Test @Disabled public void uploadMetadataSync() { final Schema schema = new Schema(Collections.singletonList(Field.nullable("a", new ArrowType.Int(32, true)))); test( (allocator, client) -> { try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); final SyncPutListener listener = new SyncPutListener()) { final FlightDescriptor descriptor = FlightDescriptor.path("test"); final FlightClient.ClientStreamListener writer = client.startPut(descriptor, root, listener); root.allocateNew(); for (byte i = 0; i < 10; i++) { final IntVector vector = (IntVector) root.getVector("a"); final ArrowBuf metadata = allocator.buffer(1); metadata.writeByte(i); vector.set(0, 10); vector.setValueCount(1); root.setRowCount(1); writer.putNext(metadata); try (final PutResult message = listener.poll(5000, TimeUnit.SECONDS)) { assertNotNull(message); assertEquals(i, message.getApplicationMetadata().getByte(0)); } catch (InterruptedException | ExecutionException e) { throw new RuntimeException(e); } } writer.completed(); // Must attempt to retrieve the result to get any server-side errors. writer.getResult(); } }); } /** Make sure that a {@link SyncPutListener} properly reclaims memory if ignored. */ @Test @Disabled public void syncMemoryReclaimed() { final Schema schema = new Schema(Collections.singletonList(Field.nullable("a", new ArrowType.Int(32, true)))); test( (allocator, client) -> { try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); final SyncPutListener listener = new SyncPutListener()) { final FlightDescriptor descriptor = FlightDescriptor.path("test"); final FlightClient.ClientStreamListener writer = client.startPut(descriptor, root, listener); root.allocateNew(); for (byte i = 0; i < 10; i++) { final IntVector vector = (IntVector) root.getVector("a"); final ArrowBuf metadata = allocator.buffer(1); metadata.writeByte(i); vector.set(0, 10); vector.setValueCount(1); root.setRowCount(1); writer.putNext(metadata); } writer.completed(); // Must attempt to retrieve the result to get any server-side errors. writer.getResult(); } }); } /** * ARROW-9221: Flight copies metadata from the byte buffer of a Protobuf ByteString, which is in * big-endian by default, thus mangling metadata. */ @Test public void testMetadataEndianness() throws Exception { try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); final BufferAllocator serverAllocator = allocator.newChildAllocator("flight-server", 0, Long.MAX_VALUE); final FlightServer server = FlightServer.builder( serverAllocator, forGrpcInsecure(LOCALHOST, 0), new EndianFlightProducer(serverAllocator)) .build() .start(); final FlightClient client = FlightClient.builder(allocator, server.getLocation()).build()) { final Schema schema = new Schema(Collections.emptyList()); final FlightDescriptor descriptor = FlightDescriptor.command(new byte[0]); try (final SyncPutListener reader = new SyncPutListener(); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final FlightClient.ClientStreamListener writer = client.startPut(descriptor, root, reader); writer.completed(); try (final PutResult metadata = reader.read()) { assertEquals(16, metadata.getApplicationMetadata().readableBytes()); byte[] bytes = new byte[16]; metadata.getApplicationMetadata().readBytes(bytes); assertArrayEquals(EndianFlightProducer.EXPECTED_BYTES, bytes); } writer.getResult(); } } } private void test(BiConsumer fun) { try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); final FlightServer s = FlightServer.builder( allocator, forGrpcInsecure(LOCALHOST, 0), new MetadataFlightProducer(allocator)) .build() .start(); final FlightClient client = FlightClient.builder(allocator, s.getLocation()).build()) { fun.accept(allocator, client); } catch (Exception e) { throw new RuntimeException(e); } } /** A FlightProducer that always produces a fixed data stream with metadata on the side. */ private static class MetadataFlightProducer extends NoOpFlightProducer { private final BufferAllocator allocator; public MetadataFlightProducer(BufferAllocator allocator) { this.allocator = allocator; } @Override public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { final Schema schema = new Schema(Collections.singletonList(Field.nullable("a", new ArrowType.Int(32, true)))); try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { root.allocateNew(); listener.start(root); for (byte i = 0; i < 10; i++) { final IntVector vector = (IntVector) root.getVector("a"); vector.set(0, 10); vector.setValueCount(1); root.setRowCount(1); final ArrowBuf metadata = allocator.buffer(1); metadata.writeByte(i); listener.putNext(metadata); } listener.completed(); } } @Override public Runnable acceptPut( CallContext context, FlightStream stream, StreamListener ackStream) { return () -> { // Wait for the descriptor to be sent stream.getRoot(); if (stream.getDescriptor().isCommand() && Arrays.equals(stream.getDescriptor().getCommand(), COMMAND_ARROW_6136)) { // ARROW-6136: Try closing the stream ackStream.onError( CallStatus.INTERNAL.withDescription(MESSAGE_ARROW_6136).toRuntimeException()); return; } try { byte current = 0; while (stream.next()) { final ArrowBuf metadata = stream.getLatestMetadata(); if (current != metadata.getByte(0)) { ackStream.onError( CallStatus.INVALID_ARGUMENT .withDescription( String.format( "Metadata does not match expected value; got %d but expected %d.", metadata.getByte(0), current)) .toRuntimeException()); return; } ackStream.onNext(PutResult.metadata(metadata)); current++; } if (current != 10) { throw CallStatus.INVALID_ARGUMENT .withDescription("Wrong number of messages sent.") .toRuntimeException(); } } catch (Exception e) { throw CallStatus.INTERNAL.withCause(e).withDescription(e.toString()).toRuntimeException(); } }; } } private static class EndianFlightProducer extends NoOpFlightProducer { static final byte[] EXPECTED_BYTES = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}; private final BufferAllocator allocator; private EndianFlightProducer(BufferAllocator allocator) { this.allocator = allocator; } @Override public Runnable acceptPut( CallContext context, FlightStream flightStream, StreamListener ackStream) { return () -> { while (flightStream.next()) { // Ignore any data } try (final ArrowBuf buf = allocator.buffer(16)) { buf.writeBytes(EXPECTED_BYTES); ackStream.onNext(PutResult.metadata(buf)); } ackStream.onCompleted(); }; } } } apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/TestAuth.java000066400000000000000000000100501515451615200322640ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; import static org.apache.arrow.flight.Location.forGrpcInsecure; import static org.junit.jupiter.api.Assertions.assertThrows; import java.util.Iterator; import java.util.Optional; import org.apache.arrow.flight.auth.ClientAuthHandler; import org.apache.arrow.flight.auth.ServerAuthHandler; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.Test; public class TestAuth { /** An auth handler that does not send messages should not block the server forever. */ @Test public void noMessages() throws Exception { assertThrows( RuntimeException.class, () -> { try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); final FlightServer s = FlightServer.builder( allocator, forGrpcInsecure(LOCALHOST, 0), new NoOpFlightProducer()) .authHandler(new OneshotAuthHandler()) .build() .start(); final FlightClient client = FlightClient.builder(allocator, s.getLocation()).build()) { client.authenticate( new ClientAuthHandler() { @Override public void authenticate(ClientAuthSender outgoing, Iterator incoming) {} @Override public byte[] getCallToken() { return new byte[0]; } }); } }); } /** An auth handler that sends an error should not block the server forever. */ @Test public void clientError() throws Exception { assertThrows( RuntimeException.class, () -> { try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); final FlightServer s = FlightServer.builder( allocator, forGrpcInsecure(LOCALHOST, 0), new NoOpFlightProducer()) .authHandler(new OneshotAuthHandler()) .build() .start(); final FlightClient client = FlightClient.builder(allocator, s.getLocation()).build()) { client.authenticate( new ClientAuthHandler() { @Override public void authenticate(ClientAuthSender outgoing, Iterator incoming) { outgoing.send(new byte[0]); // Ensure the server-side runs incoming.next(); outgoing.onError(new RuntimeException("test")); } @Override public byte[] getCallToken() { return new byte[0]; } }); } }); } private static class OneshotAuthHandler implements ServerAuthHandler { @Override public Optional isValid(byte[] token) { return Optional.of("test"); } @Override public boolean authenticate(ServerAuthSender outgoing, Iterator incoming) { incoming.next(); outgoing.send(new byte[0]); return false; } } } TestBackPressure.java000066400000000000000000000227701515451615200337110ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; import static org.apache.arrow.flight.Location.forGrpcInsecure; import static org.junit.jupiter.api.Assertions.assertTrue; import com.google.common.collect.ImmutableList; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.atomic.AtomicLong; import java.util.function.Function; import org.apache.arrow.flight.perf.PerformanceTestServer; import org.apache.arrow.flight.perf.TestPerf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.Types.MinorType; 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 TestBackPressure { private static final int BATCH_SIZE = 4095; /** Make sure that failing to consume one stream doesn't block other streams. */ @Disabled @Test public void ensureIndependentSteams() throws Exception { ensureIndependentSteams((b) -> (location -> new PerformanceTestServer(b, location))); } /** Make sure that failing to consume one stream doesn't block other streams. */ @Disabled @Test public void ensureIndependentSteamsWithCallbacks() throws Exception { ensureIndependentSteams( (b) -> (location -> new PerformanceTestServer( b, location, new BackpressureStrategy.CallbackBackpressureStrategy(), true))); } /** Test to make sure stream doesn't go faster than the consumer is consuming. */ @Disabled @Test public void ensureWaitUntilProceed() throws Exception { ensureWaitUntilProceed(new PollingBackpressureStrategy(), false); } /** * Test to make sure stream doesn't go faster than the consumer is consuming using a * callback-based backpressure strategy. */ @Disabled @Test public void ensureWaitUntilProceedWithCallbacks() throws Exception { ensureWaitUntilProceed(new RecordingCallbackBackpressureStrategy(), true); } /** Make sure that failing to consume one stream doesn't block other streams. */ private static void ensureIndependentSteams( Function> serverConstructor) throws Exception { try (final BufferAllocator a = new RootAllocator(Long.MAX_VALUE); final PerformanceTestServer server = serverConstructor.apply(a).apply(forGrpcInsecure(LOCALHOST, 0)).start(); final FlightClient client = FlightClient.builder(a, server.getLocation()).build()) { try (FlightStream fs1 = client.getStream( client .getInfo(TestPerf.getPerfFlightDescriptor(110L * BATCH_SIZE, BATCH_SIZE, 1)) .getEndpoints() .get(0) .getTicket())) { consume(fs1, 10); // stop consuming fs1 but make sure we can consume a large amount of fs2. try (FlightStream fs2 = client.getStream( client .getInfo(TestPerf.getPerfFlightDescriptor(200L * BATCH_SIZE, BATCH_SIZE, 1)) .getEndpoints() .get(0) .getTicket())) { consume(fs2, 100); consume(fs1, 100); consume(fs2, 100); consume(fs1); consume(fs2); } } } } /** Make sure that a stream doesn't go faster than the consumer is consuming. */ private static void ensureWaitUntilProceed( SleepTimeRecordingBackpressureStrategy bpStrategy, boolean isNonBlocking) throws Exception { // request some values. final long wait = 3000; final long epsilon = 1000; try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { final FlightProducer producer = new NoOpFlightProducer() { @Override public void getStream( CallContext context, Ticket ticket, ServerStreamListener listener) { bpStrategy.register(listener); final Runnable loadData = () -> { int batches = 0; final Schema pojoSchema = new Schema( ImmutableList.of(Field.nullable("a", MinorType.BIGINT.getType()))); try (VectorSchemaRoot root = VectorSchemaRoot.create(pojoSchema, allocator)) { listener.start(root); while (true) { bpStrategy.waitForListener(0); if (batches > 100) { root.clear(); listener.completed(); return; } root.allocateNew(); root.setRowCount(4095); listener.putNext(); batches++; } } }; if (!isNonBlocking) { loadData.run(); } else { final ExecutorService service = Executors.newSingleThreadExecutor(); Future unused = service.submit(loadData); service.shutdown(); } } }; try (BufferAllocator serverAllocator = allocator.newChildAllocator("server", 0, Long.MAX_VALUE); FlightServer server = FlightServer.builder(serverAllocator, forGrpcInsecure(LOCALHOST, 0), producer) .build() .start(); BufferAllocator clientAllocator = allocator.newChildAllocator("client", 0, Long.MAX_VALUE); FlightClient client = FlightClient.builder(clientAllocator, server.getLocation()).build(); FlightStream stream = client.getStream(new Ticket(new byte[1]))) { VectorSchemaRoot root = stream.getRoot(); root.clear(); Thread.sleep(wait); while (stream.next()) { root.clear(); } long expected = wait - epsilon; assertTrue( bpStrategy.getSleepTime() > expected, String.format( "Expected a sleep of at least %dms but only slept for %d", expected, bpStrategy.getSleepTime())); } } } private static void consume(FlightStream stream) { VectorSchemaRoot root = stream.getRoot(); while (stream.next()) { root.clear(); } } private static void consume(FlightStream stream, int batches) { VectorSchemaRoot root = stream.getRoot(); while (batches > 0 && stream.next()) { root.clear(); batches--; } } private interface SleepTimeRecordingBackpressureStrategy extends BackpressureStrategy { /** * Returns the total time spent waiting on the listener to be ready. * * @return the total time spent waiting on the listener to be ready. */ long getSleepTime(); } /** * Implementation of a backpressure strategy that polls on isReady and records amount of time * spent in Thread.sleep(). */ private static class PollingBackpressureStrategy implements SleepTimeRecordingBackpressureStrategy { private final AtomicLong sleepTime = new AtomicLong(0); private FlightProducer.ServerStreamListener listener; @Override public long getSleepTime() { return sleepTime.get(); } @Override public void register(FlightProducer.ServerStreamListener listener) { this.listener = listener; } @Override public WaitResult waitForListener(long timeout) { while (!listener.isReady()) { try { Thread.sleep(1); sleepTime.addAndGet(1L); } catch (InterruptedException expected) { // it is expected and no action needed } } return WaitResult.READY; } } /** * Implementation of a backpressure strategy that uses callbacks to detect changes in client * readiness state and records spent time waiting. */ private static class RecordingCallbackBackpressureStrategy extends BackpressureStrategy.CallbackBackpressureStrategy implements SleepTimeRecordingBackpressureStrategy { private final AtomicLong sleepTime = new AtomicLong(0); @Override public long getSleepTime() { return sleepTime.get(); } @Override public WaitResult waitForListener(long timeout) { final long startTime = System.currentTimeMillis(); final WaitResult result = super.waitForListener(timeout); sleepTime.addAndGet(System.currentTimeMillis() - startTime); return result; } } } TestBasicOperation.java000066400000000000000000000572041515451615200342220ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; 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.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; import com.google.common.base.Charsets; import com.google.protobuf.ByteString; import io.grpc.MethodDescriptor; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.InputStream; import java.net.InetSocketAddress; import java.net.URI; import java.net.URISyntaxException; import java.nio.channels.Channels; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.function.BiConsumer; import java.util.function.Consumer; import org.apache.arrow.flight.FlightClient.ClientStreamListener; import org.apache.arrow.flight.impl.Flight; import org.apache.arrow.flight.impl.Flight.FlightDescriptor.DescriptorType; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; import org.apache.arrow.vector.ipc.WriteChannel; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.ipc.message.IpcOption; import org.apache.arrow.vector.ipc.message.MessageSerializer; 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; import org.junit.jupiter.api.condition.DisabledOnOs; import org.junit.jupiter.api.condition.OS; /** Test the operations of a basic flight service. */ public class TestBasicOperation { @Test public void fastPathDefaults() { assertTrue(ArrowMessage.ENABLE_ZERO_COPY_READ); assertFalse(ArrowMessage.ENABLE_ZERO_COPY_WRITE); } @Test public void fallbackLocation() { assertEquals( "arrow-flight-reuse-connection://?", Location.reuseConnection().getUri().toString()); } /** ARROW-6017: we should be able to construct locations for unknown schemes. */ @Test public void unknownScheme() throws URISyntaxException { final Location location = new Location("s3://unknown"); assertEquals("s3", location.getUri().getScheme()); } @Test public void unknownSchemeRemote() throws Exception { test( c -> { try { final FlightInfo info = c.getInfo(FlightDescriptor.path("test")); assertEquals( new URI("https://example.com"), info.getEndpoints().get(0).getLocations().get(0).getUri()); } catch (URISyntaxException e) { throw new RuntimeException(e); } }); } @Test public void roundTripTicket() throws Exception { final Ticket ticket = new Ticket(new byte[] {0, 1, 2, 3, 4, 5}); assertEquals(ticket, Ticket.deserialize(ticket.serialize())); } @Test public void roundTripInfo() throws Exception { final Map metadata = new HashMap<>(); metadata.put("foo", "bar"); final Schema schema = new Schema( Arrays.asList( Field.nullable("a", new ArrowType.Int(32, true)), Field.nullable("b", new ArrowType.FixedSizeBinary(32))), metadata); final FlightInfo info1 = FlightInfo.builder(schema, FlightDescriptor.path(), Collections.emptyList()) .setAppMetadata("foo".getBytes(StandardCharsets.UTF_8)) .build(); final FlightInfo info2 = new FlightInfo( schema, FlightDescriptor.command(new byte[2]), Collections.singletonList( FlightEndpoint.builder( new Ticket(new byte[10]), Location.forGrpcDomainSocket("/tmp/test.sock")) .setAppMetadata("bar".getBytes(StandardCharsets.UTF_8)) .build()), 200, 500); final FlightInfo info3 = new FlightInfo( schema, FlightDescriptor.path("a", "b"), Arrays.asList( new FlightEndpoint( new Ticket(new byte[10]), Location.forGrpcDomainSocket("/tmp/test.sock")), new FlightEndpoint( new Ticket(new byte[10]), Location.forGrpcDomainSocket("/tmp/test.sock"), forGrpcInsecure("localhost", 50051))), 200, 500); final FlightInfo info4 = new FlightInfo( schema, FlightDescriptor.path("a", "b"), Arrays.asList( new FlightEndpoint( new Ticket(new byte[10]), Location.forGrpcDomainSocket("/tmp/test.sock")), new FlightEndpoint( new Ticket(new byte[10]), Location.forGrpcDomainSocket("/tmp/test.sock"), forGrpcInsecure("localhost", 50051))), 200, 500, /*ordered*/ true, IpcOption.DEFAULT); assertEquals(info1, FlightInfo.deserialize(info1.serialize())); assertEquals(info2, FlightInfo.deserialize(info2.serialize())); assertEquals(info3, FlightInfo.deserialize(info3.serialize())); assertEquals(info4, FlightInfo.deserialize(info4.serialize())); assertNotEquals(info3, info4); assertFalse(info1.getOrdered()); assertFalse(info2.getOrdered()); assertFalse(info3.getOrdered()); assertTrue(info4.getOrdered()); } @Test public void roundTripDescriptor() throws Exception { final FlightDescriptor cmd = FlightDescriptor.command("test command".getBytes(StandardCharsets.UTF_8)); assertEquals(cmd, FlightDescriptor.deserialize(cmd.serialize())); final FlightDescriptor path = FlightDescriptor.path("foo", "bar", "test.arrow"); assertEquals(path, FlightDescriptor.deserialize(path.serialize())); } @Test public void getDescriptors() throws Exception { test( c -> { int count = 0; for (FlightInfo unused : c.listFlights(Criteria.ALL)) { count += 1; } assertEquals(1, count); }); } @Test public void getDescriptorsWithCriteria() throws Exception { test( c -> { int count = 0; for (FlightInfo unused : c.listFlights(new Criteria(new byte[] {1}))) { count += 1; } assertEquals(0, count); }); } @Test public void getDescriptor() throws Exception { test( c -> { System.out.println(c.getInfo(FlightDescriptor.path("hello")).getDescriptor()); }); } @Test public void getSchema() throws Exception { test( c -> { System.out.println(c.getSchema(FlightDescriptor.path("hello")).getSchema()); }); } @Test public void listActions() throws Exception { test( c -> { for (ActionType at : c.listActions()) { System.out.println(at.getType()); } }); } @Test public void doAction() throws Exception { test( c -> { Iterator stream = c.doAction(new Action("hello")); assertTrue(stream.hasNext()); Result r = stream.next(); assertArrayEquals("world".getBytes(Charsets.UTF_8), r.getBody()); }); test( c -> { Iterator stream = c.doAction(new Action("hellooo")); assertTrue(stream.hasNext()); Result r = stream.next(); assertArrayEquals("world".getBytes(Charsets.UTF_8), r.getBody()); assertTrue(stream.hasNext()); r = stream.next(); assertArrayEquals("!".getBytes(Charsets.UTF_8), r.getBody()); assertFalse(stream.hasNext()); }); } @Test public void putStream() throws Exception { test( (c, a) -> { final int size = 10; IntVector iv = new IntVector("c1", a); try (VectorSchemaRoot root = VectorSchemaRoot.of(iv)) { ClientStreamListener listener = c.startPut(FlightDescriptor.path("hello"), root, new AsyncPutListener()); // batch 1 root.allocateNew(); for (int i = 0; i < size; i++) { iv.set(i, i); } iv.setValueCount(size); root.setRowCount(size); listener.putNext(); // batch 2 root.allocateNew(); for (int i = 0; i < size; i++) { iv.set(i, i + size); } iv.setValueCount(size); root.setRowCount(size); listener.putNext(); root.clear(); listener.completed(); // wait for ack to avoid memory leaks. listener.getResult(); } }); } @Test public void propagateErrors() throws Exception { test( client -> { FlightTestUtil.assertCode( FlightStatusCode.UNIMPLEMENTED, () -> { client.doAction(new Action("invalid-action")).forEachRemaining(action -> fail()); }); }); } @Test public void getStream() throws Exception { test( c -> { try (final FlightStream stream = c.getStream(new Ticket(new byte[0]))) { VectorSchemaRoot root = stream.getRoot(); IntVector iv = (IntVector) root.getVector("c1"); int value = 0; while (stream.next()) { for (int i = 0; i < root.getRowCount(); i++) { assertEquals(value, iv.get(i)); value++; } } } catch (Exception e) { throw new RuntimeException(e); } }); } /** Ensure the client is configured to accept large messages. */ @Test @DisabledOnOs( value = {OS.WINDOWS}, disabledReason = "https://github.com/apache/arrow/issues/33237: flaky test") public void getStreamLargeBatch() throws Exception { test( c -> { try (final FlightStream stream = c.getStream(new Ticket(Producer.TICKET_LARGE_BATCH))) { assertEquals(128, stream.getRoot().getFieldVectors().size()); assertTrue(stream.next()); assertEquals(65536, stream.getRoot().getRowCount()); assertTrue(stream.next()); assertEquals(65536, stream.getRoot().getRowCount()); assertFalse(stream.next()); } catch (Exception e) { throw new RuntimeException(e); } }); } /** Ensure the server is configured to accept large messages. */ @Test public void startPutLargeBatch() throws Exception { try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) { final List vectors = new ArrayList<>(); for (int col = 0; col < 128; col++) { final BigIntVector vector = new BigIntVector("f" + col, allocator); for (int row = 0; row < 65536; row++) { vector.setSafe(row, row); } vectors.add(vector); } test( c -> { try (final VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { root.setRowCount(65536); final ClientStreamListener stream = c.startPut(FlightDescriptor.path(""), root, new SyncPutListener()); stream.putNext(); stream.putNext(); stream.completed(); stream.getResult(); } catch (Exception e) { throw new RuntimeException(e); } }); } } private void test(Consumer consumer) throws Exception { test( (c, a) -> { consumer.accept(c); }); } private void test(BiConsumer consumer) throws Exception { try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE); Producer producer = new Producer(a); FlightServer s = FlightServer.builder(a, forGrpcInsecure(LOCALHOST, 0), producer).build().start()) { try (FlightClient c = FlightClient.builder(a, s.getLocation()).build()) { try (BufferAllocator testAllocator = a.newChildAllocator("testcase", 0, Long.MAX_VALUE)) { consumer.accept(c, testAllocator); } } } } /** Helper method to convert an ArrowMessage into a Protobuf message. */ private Flight.FlightData arrowMessageToProtobuf( MethodDescriptor.Marshaller marshaller, ArrowMessage message) throws IOException { final ByteArrayOutputStream baos = new ByteArrayOutputStream(); try (final InputStream serialized = marshaller.stream(message)) { final byte[] buf = new byte[1024]; while (true) { int read = serialized.read(buf); if (read < 0) { break; } baos.write(buf, 0, read); } } final byte[] serializedMessage = baos.toByteArray(); return Flight.FlightData.parseFrom(serializedMessage); } /** * ARROW-10962: accept FlightData messages generated by Protobuf (which can omit empty fields). */ @Test public void testProtobufRecordBatchCompatibility() throws Exception { final Schema schema = new Schema(Collections.singletonList(Field.nullable("foo", new ArrowType.Int(32, true)))); try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final VectorUnloader unloader = new VectorUnloader(root); root.setRowCount(0); final MethodDescriptor.Marshaller marshaller = ArrowMessage.createMarshaller(allocator); try (final ArrowMessage message = new ArrowMessage( unloader.getRecordBatch(), /* appMetadata */ null, /* tryZeroCopy */ false, IpcOption.DEFAULT)) { assertEquals(ArrowMessage.HeaderType.RECORD_BATCH, message.getMessageType()); // Should have at least one empty body buffer (there may be multiple for e.g. data and // validity) Iterator iterator = message.getBufs().iterator(); assertTrue(iterator.hasNext()); while (iterator.hasNext()) { assertEquals(0, iterator.next().capacity()); } final Flight.FlightData protobufData = arrowMessageToProtobuf(marshaller, message).toBuilder().clearDataBody().build(); assertEquals(0, protobufData.getDataBody().size()); ArrowMessage parsedMessage = marshaller.parse(new ByteArrayInputStream(protobufData.toByteArray())); // Should have an empty body buffer Iterator parsedIterator = parsedMessage.getBufs().iterator(); assertTrue(parsedIterator.hasNext()); assertEquals(0, parsedIterator.next().capacity()); // Should have only one (the parser synthesizes exactly one); in the case of empty buffers, // this is equivalent assertFalse(parsedIterator.hasNext()); // Should not throw final ArrowRecordBatch rb = parsedMessage.asRecordBatch(); assertEquals(rb.computeBodyLength(), 0); } } } /** * ARROW-10962: accept FlightData messages generated by Protobuf (which can omit empty fields). */ @Test public void testProtobufSchemaCompatibility() throws Exception { final Schema schema = new Schema(Collections.singletonList(Field.nullable("foo", new ArrowType.Int(32, true)))); try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) { final MethodDescriptor.Marshaller marshaller = ArrowMessage.createMarshaller(allocator); Flight.FlightDescriptor descriptor = FlightDescriptor.command(new byte[0]).toProtocol(); try (final ArrowMessage message = new ArrowMessage(descriptor, schema, IpcOption.DEFAULT)) { assertEquals(ArrowMessage.HeaderType.SCHEMA, message.getMessageType()); // Should have no body buffers assertFalse(message.getBufs().iterator().hasNext()); final Flight.FlightData protobufData = arrowMessageToProtobuf(marshaller, message).toBuilder() .setDataBody(ByteString.EMPTY) .build(); assertEquals(0, protobufData.getDataBody().size()); final ArrowMessage parsedMessage = marshaller.parse(new ByteArrayInputStream(protobufData.toByteArray())); // Should have no body buffers assertFalse(parsedMessage.getBufs().iterator().hasNext()); // Should not throw parsedMessage.asSchema(); } } } @Test public void testGrpcInsecureLocation() throws Exception { Location location = Location.forGrpcInsecure(LOCALHOST, 9000); assertEquals( new URI(LocationSchemes.GRPC_INSECURE, null, LOCALHOST, 9000, null, null, null), location.getUri()); assertEquals(new InetSocketAddress(LOCALHOST, 9000), location.toSocketAddress()); } @Test public void testGrpcTlsLocation() throws Exception { Location location = Location.forGrpcTls(LOCALHOST, 9000); assertEquals( new URI(LocationSchemes.GRPC_TLS, null, LOCALHOST, 9000, null, null, null), location.getUri()); assertEquals(new InetSocketAddress(LOCALHOST, 9000), location.toSocketAddress()); } /** An example FlightProducer for test purposes. */ public static class Producer implements FlightProducer, AutoCloseable { static final byte[] TICKET_LARGE_BATCH = "large-batch".getBytes(StandardCharsets.UTF_8); private final BufferAllocator allocator; public Producer(BufferAllocator allocator) { super(); this.allocator = allocator; } @Override public void listFlights( CallContext context, Criteria criteria, StreamListener listener) { if (criteria.getExpression().length > 0) { // Don't send anything if criteria are set listener.onCompleted(); } Flight.FlightInfo getInfo = Flight.FlightInfo.newBuilder() .setFlightDescriptor( Flight.FlightDescriptor.newBuilder() .setType(DescriptorType.CMD) .setCmd(ByteString.copyFrom("cool thing", Charsets.UTF_8))) .build(); try { listener.onNext(new FlightInfo(getInfo)); } catch (URISyntaxException e) { listener.onError(e); return; } listener.onCompleted(); } @Override public Runnable acceptPut( CallContext context, FlightStream flightStream, StreamListener ackStream) { return () -> { while (flightStream.next()) { // Drain the stream } }; } @Override public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { if (Arrays.equals(TICKET_LARGE_BATCH, ticket.getBytes())) { getLargeBatch(listener); return; } final int size = 10; IntVector iv = new IntVector("c1", allocator); VectorSchemaRoot root = VectorSchemaRoot.of(iv); listener.start(root); // batch 1 root.allocateNew(); for (int i = 0; i < size; i++) { iv.set(i, i); } iv.setValueCount(size); root.setRowCount(size); listener.putNext(); // batch 2 root.allocateNew(); for (int i = 0; i < size; i++) { iv.set(i, i + size); } iv.setValueCount(size); root.setRowCount(size); listener.putNext(); root.clear(); listener.completed(); } private void getLargeBatch(ServerStreamListener listener) { final List vectors = new ArrayList<>(); for (int col = 0; col < 128; col++) { final BigIntVector vector = new BigIntVector("f" + col, allocator); for (int row = 0; row < 65536; row++) { vector.setSafe(row, row); } vectors.add(vector); } try (final VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { root.setRowCount(65536); listener.start(root); listener.putNext(); listener.putNext(); listener.completed(); } } @Override public void close() throws Exception { allocator.close(); } @Override public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) { try { Flight.FlightInfo getInfo = Flight.FlightInfo.newBuilder() .setSchema(schemaToByteString(new Schema(Collections.emptyList()))) .setFlightDescriptor( Flight.FlightDescriptor.newBuilder() .setType(DescriptorType.CMD) .setCmd(ByteString.copyFrom("cool thing", Charsets.UTF_8))) .addEndpoint( Flight.FlightEndpoint.newBuilder() .addLocation(new Location("https://example.com").toProtocol())) .build(); return new FlightInfo(getInfo); } catch (URISyntaxException e) { throw new RuntimeException(e); } } private static ByteString schemaToByteString(Schema schema) { try (ByteArrayOutputStream baos = new ByteArrayOutputStream()) { MessageSerializer.serialize( new WriteChannel(Channels.newChannel(baos)), schema, IpcOption.DEFAULT); return ByteString.copyFrom(baos.toByteArray()); } catch (IOException e) { throw new RuntimeException(e); } } @Override public void doAction(CallContext context, Action action, StreamListener listener) { switch (action.getType()) { case "hello": { listener.onNext(new Result("world".getBytes(Charsets.UTF_8))); listener.onCompleted(); break; } case "hellooo": { listener.onNext(new Result("world".getBytes(Charsets.UTF_8))); listener.onNext(new Result("!".getBytes(Charsets.UTF_8))); listener.onCompleted(); break; } default: listener.onError( CallStatus.UNIMPLEMENTED .withDescription("Action not implemented: " + action.getType()) .toRuntimeException()); } } @Override public void listActions(CallContext context, StreamListener listener) { listener.onNext(new ActionType("get", "")); listener.onNext(new ActionType("put", "")); listener.onNext(new ActionType("hello", "")); listener.onCompleted(); } } } TestCallOptions.java000066400000000000000000000162141515451615200335430ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; 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; import io.grpc.Metadata; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.time.Duration; import java.time.Instant; import java.util.Iterator; import java.util.concurrent.TimeUnit; import java.util.function.Consumer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; public class TestCallOptions { @Test @Disabled public void timeoutFires() { // Ignored due to CI flakiness test( (client) -> { Instant start = Instant.now(); Iterator results = client.doAction(new Action("hang"), CallOptions.timeout(1, TimeUnit.SECONDS)); try { results.next(); fail("Call should have failed"); } catch (RuntimeException e) { assertTrue(e.getMessage().contains("deadline exceeded"), e.getMessage()); } Instant end = Instant.now(); assertTrue( Duration.between(start, end).toMillis() < 1500, "Call took over 1500 ms despite timeout"); }); } @Test @Disabled public void underTimeout() { // Ignored due to CI flakiness test( (client) -> { Instant start = Instant.now(); // This shouldn't fail and it should complete within the timeout Iterator results = client.doAction(new Action("fast"), CallOptions.timeout(2, TimeUnit.SECONDS)); assertArrayEquals(new byte[] {42, 42}, results.next().getBody()); Instant end = Instant.now(); assertTrue( Duration.between(start, end).toMillis() < 2500, "Call took over 2500 ms despite timeout"); }); } @Test public void singleProperty() { final FlightCallHeaders headers = new FlightCallHeaders(); headers.insert("key", "value"); testHeaders(headers); } @Test public void multipleProperties() { final FlightCallHeaders headers = new FlightCallHeaders(); headers.insert("key", "value"); headers.insert("key2", "value2"); testHeaders(headers); } @Test public void binaryProperties() { final FlightCallHeaders headers = new FlightCallHeaders(); headers.insert("key-bin", "value".getBytes(StandardCharsets.UTF_8)); headers.insert("key3-bin", "ëfßæ".getBytes(StandardCharsets.UTF_8)); testHeaders(headers); } @Test public void mixedProperties() { final FlightCallHeaders headers = new FlightCallHeaders(); headers.insert("key", "value"); headers.insert("key3-bin", "ëfßæ".getBytes(StandardCharsets.UTF_8)); 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(); FlightServer s = FlightServer.builder(a, forGrpcInsecure(LOCALHOST, 0), producer).build().start(); FlightClient client = FlightClient.builder(a, s.getLocation()).build()) { assertFalse(client.doAction(new Action(""), new HeaderCallOption(headers)).hasNext()); final CallHeaders incomingHeaders = producer.headers(); for (String key : headers.keys()) { if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) { assertArrayEquals(headers.getByte(key), incomingHeaders.getByte(key)); } else { assertEquals(headers.get(key), incomingHeaders.get(key)); } } } catch (InterruptedException | IOException e) { throw new RuntimeException(e); } } void test(Consumer testFn) { try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE); Producer producer = new Producer(); FlightServer s = FlightServer.builder(a, forGrpcInsecure(LOCALHOST, 0), producer).build().start(); FlightClient client = FlightClient.builder(a, s.getLocation()).build()) { testFn.accept(client); } catch (InterruptedException | IOException e) { throw new RuntimeException(e); } } static class HeaderProducer extends NoOpFlightProducer implements AutoCloseable { CallHeaders headers; @Override public void close() {} public CallHeaders headers() { return headers; } @Override public void doAction(CallContext context, Action action, StreamListener listener) { this.headers = context.getMiddleware(FlightConstants.HEADER_KEY).headers(); listener.onCompleted(); } } static class Producer extends NoOpFlightProducer implements AutoCloseable { Producer() {} @Override public void close() {} @Override public void doAction(CallContext context, Action action, StreamListener listener) { switch (action.getType()) { case "hang": { try { Thread.sleep(25000); } catch (InterruptedException e) { throw new RuntimeException(e); } listener.onNext(new Result(new byte[] {})); listener.onCompleted(); return; } case "fast": { try { Thread.sleep(500); } catch (InterruptedException e) { throw new RuntimeException(e); } listener.onNext(new Result(new byte[] {42, 42})); listener.onCompleted(); return; } default: { throw new UnsupportedOperationException(action.getType()); } } } } } TestClientMiddleware.java000066400000000000000000000326661515451615200345410ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; 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.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.function.BiConsumer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.Test; /** A basic test of client middleware using a simplified OpenTracing-like example. */ public class TestClientMiddleware { /** * Test that a client middleware can fail a call before it starts by throwing a {@link * FlightRuntimeException}. */ @Test public void clientMiddleware_failCallBeforeSending() { test( new NoOpFlightProducer(), null, Collections.singletonList(new CallRejector.Factory()), (allocator, client) -> { FlightTestUtil.assertCode(FlightStatusCode.UNAVAILABLE, client::listActions); }); } /** * Test an OpenTracing-like scenario where client and server middleware work together to propagate * a request ID without explicit intervention from the service implementation. */ @Test public void middleware_propagateHeader() { final Context context = new Context("span id"); test( new NoOpFlightProducer(), new TestServerMiddleware.ServerMiddlewarePair<>( FlightServerMiddleware.Key.of("test"), new ServerSpanInjector.Factory()), Collections.singletonList(new ClientSpanInjector.Factory(context)), (allocator, client) -> { FlightTestUtil.assertCode( FlightStatusCode.UNIMPLEMENTED, () -> client.listActions().forEach(actionType -> {})); }); assertEquals(context.outgoingSpanId, context.incomingSpanId); assertNotNull(context.finalStatus); assertEquals(FlightStatusCode.UNIMPLEMENTED, context.finalStatus.code()); } /** * Ensure both server and client can send and receive multi-valued headers (both binary and text * values). */ @Test public void testMultiValuedHeaders() { final MultiHeaderClientMiddlewareFactory clientFactory = new MultiHeaderClientMiddlewareFactory(); test( new NoOpFlightProducer(), new TestServerMiddleware.ServerMiddlewarePair<>( FlightServerMiddleware.Key.of("test"), new MultiHeaderServerMiddlewareFactory()), Collections.singletonList(clientFactory), (allocator, client) -> { FlightTestUtil.assertCode( FlightStatusCode.UNIMPLEMENTED, () -> client.listActions().forEach(actionType -> {})); }); // The server echoes the headers we send back to us, so ensure all the ones we sent are present // with the correct // values in the correct order. for (final Map.Entry> entry : EXPECTED_BINARY_HEADERS.entrySet()) { // Compare header values entry-by-entry because byte arrays don't compare via equals final List receivedValues = clientFactory.lastBinaryHeaders.get(entry.getKey()); assertNotNull(receivedValues, "Missing for header: " + entry.getKey()); assertEquals( entry.getValue().size(), receivedValues.size(), "Missing or wrong value for header: " + entry.getKey()); for (int i = 0; i < entry.getValue().size(); i++) { assertArrayEquals(entry.getValue().get(i), receivedValues.get(i)); } } for (final Map.Entry> entry : EXPECTED_TEXT_HEADERS.entrySet()) { assertEquals( entry.getValue(), clientFactory.lastTextHeaders.get(entry.getKey()), "Missing or wrong value for header: " + entry.getKey()); } } private static void test( FlightProducer producer, TestServerMiddleware.ServerMiddlewarePair serverMiddleware, List clientMiddleware, BiConsumer body) { try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) { final FlightServer.Builder serverBuilder = FlightServer.builder(allocator, forGrpcInsecure(LOCALHOST, 0), producer); if (serverMiddleware != null) { serverBuilder.middleware(serverMiddleware.key, serverMiddleware.factory); } final FlightServer server = serverBuilder.build().start(); FlightClient.Builder clientBuilder = FlightClient.builder(allocator, server.getLocation()); clientMiddleware.forEach(clientBuilder::intercept); try (final FlightServer ignored = server; final FlightClient client = clientBuilder.build()) { body.accept(allocator, client); } } catch (InterruptedException | IOException e) { throw new RuntimeException(e); } } /** * A server middleware component that reads a request ID from incoming headers and sends the * request ID back on outgoing headers. */ static class ServerSpanInjector implements FlightServerMiddleware { private final String spanId; public ServerSpanInjector(String spanId) { this.spanId = spanId; } @Override public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) { outgoingHeaders.insert("x-span", spanId); } @Override public void onCallCompleted(CallStatus status) {} @Override public void onCallErrored(Throwable err) {} static class Factory implements FlightServerMiddleware.Factory { @Override public ServerSpanInjector onCallStarted( CallInfo info, CallHeaders incomingHeaders, RequestContext context) { return new ServerSpanInjector(incomingHeaders.get("x-span")); } } } /** * A client middleware component that, given a mock OpenTracing-like "request context", sends the * request ID in the context on outgoing headers and reads it from incoming headers. */ static class ClientSpanInjector implements FlightClientMiddleware { private final Context context; public ClientSpanInjector(Context context) { this.context = context; } @Override public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) { outgoingHeaders.insert("x-span", context.outgoingSpanId); } @Override public void onHeadersReceived(CallHeaders incomingHeaders) { context.incomingSpanId = incomingHeaders.get("x-span"); } @Override public void onCallCompleted(CallStatus status) { context.finalStatus = status; } static class Factory implements FlightClientMiddleware.Factory { private final Context context; Factory(Context context) { this.context = context; } @Override public FlightClientMiddleware onCallStarted(CallInfo info) { return new ClientSpanInjector(context); } } } /** A mock OpenTracing-like "request context". */ static class Context { final String outgoingSpanId; String incomingSpanId; CallStatus finalStatus; Context(String spanId) { this.outgoingSpanId = spanId; } } /** A client middleware that fails outgoing calls. */ static class CallRejector implements FlightClientMiddleware { @Override public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) {} @Override public void onHeadersReceived(CallHeaders incomingHeaders) {} @Override public void onCallCompleted(CallStatus status) {} static class Factory implements FlightClientMiddleware.Factory { @Override public FlightClientMiddleware onCallStarted(CallInfo info) { throw CallStatus.UNAVAILABLE.withDescription("Rejecting call.").toRuntimeException(); } } } // Used to test that middleware can send and receive multi-valued text and binary headers. static Map> EXPECTED_BINARY_HEADERS = new HashMap>(); static Map> EXPECTED_TEXT_HEADERS = new HashMap>(); static { EXPECTED_BINARY_HEADERS.put("x-binary-bin", Arrays.asList(new byte[] {0}, new byte[] {1})); EXPECTED_TEXT_HEADERS.put("x-text", Arrays.asList("foo", "bar")); } static class MultiHeaderServerMiddlewareFactory implements FlightServerMiddleware.Factory { @Override public MultiHeaderServerMiddleware onCallStarted( CallInfo info, CallHeaders incomingHeaders, RequestContext context) { // Echo the headers back to the client. Copy values out of CallHeaders since the underlying // gRPC metadata // object isn't safe to use after this function returns. Map> binaryHeaders = new HashMap<>(); Map> textHeaders = new HashMap<>(); for (final String key : incomingHeaders.keys()) { if (key.endsWith("-bin")) { binaryHeaders.compute( key, (ignored, values) -> { if (values == null) { values = new ArrayList<>(); } incomingHeaders.getAllByte(key).forEach(values::add); return values; }); } else { textHeaders.compute( key, (ignored, values) -> { if (values == null) { values = new ArrayList<>(); } incomingHeaders.getAll(key).forEach(values::add); return values; }); } } return new MultiHeaderServerMiddleware(binaryHeaders, textHeaders); } } static class MultiHeaderServerMiddleware implements FlightServerMiddleware { private final Map> binaryHeaders; private final Map> textHeaders; MultiHeaderServerMiddleware( Map> binaryHeaders, Map> textHeaders) { this.binaryHeaders = binaryHeaders; this.textHeaders = textHeaders; } @Override public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) { binaryHeaders.forEach( (key, values) -> values.forEach(value -> outgoingHeaders.insert(key, value))); textHeaders.forEach( (key, values) -> values.forEach(value -> outgoingHeaders.insert(key, value))); } @Override public void onCallCompleted(CallStatus status) {} @Override public void onCallErrored(Throwable err) {} } static class MultiHeaderClientMiddlewareFactory implements FlightClientMiddleware.Factory { Map> lastBinaryHeaders = null; Map> lastTextHeaders = null; @Override public FlightClientMiddleware onCallStarted(CallInfo info) { return new MultiHeaderClientMiddleware(this); } } static class MultiHeaderClientMiddleware implements FlightClientMiddleware { private final MultiHeaderClientMiddlewareFactory factory; public MultiHeaderClientMiddleware(MultiHeaderClientMiddlewareFactory factory) { this.factory = factory; } @Override public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) { for (final Map.Entry> entry : EXPECTED_BINARY_HEADERS.entrySet()) { entry.getValue().forEach((value) -> outgoingHeaders.insert(entry.getKey(), value)); assertTrue(outgoingHeaders.containsKey(entry.getKey())); } for (final Map.Entry> entry : EXPECTED_TEXT_HEADERS.entrySet()) { entry.getValue().forEach((value) -> outgoingHeaders.insert(entry.getKey(), value)); assertTrue(outgoingHeaders.containsKey(entry.getKey())); } } @Override public void onHeadersReceived(CallHeaders incomingHeaders) { factory.lastBinaryHeaders = new HashMap<>(); factory.lastTextHeaders = new HashMap<>(); incomingHeaders .keys() .forEach( header -> { if (header.endsWith("-bin")) { final List values = new ArrayList<>(); incomingHeaders.getAllByte(header).forEach(values::add); factory.lastBinaryHeaders.put(header, values); } else { final List values = new ArrayList<>(); incomingHeaders.getAll(header).forEach(values::add); factory.lastTextHeaders.put(header, values); } }); } @Override public void onCallCompleted(CallStatus status) {} } } TestDictionaryUtils.java000066400000000000000000000071551515451615200344460ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertSame; import com.google.common.collect.ImmutableList; import java.util.TreeSet; 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.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.types.pojo.Schema; import org.junit.jupiter.api.Test; /** Test cases for {@link DictionaryUtils}. */ public class TestDictionaryUtils { @Test public void testReuseSchema() { FieldType varcharType = new FieldType(true, new ArrowType.Utf8(), null); FieldType intType = new FieldType(true, new ArrowType.Int(32, true), null); ImmutableList build = ImmutableList.of( new Field("stringCol", varcharType, null), new Field("intCol", intType, null)); Schema schema = new Schema(build); Schema newSchema = DictionaryUtils.generateSchema(schema, null, new TreeSet<>()); // assert that no new schema is created. assertSame(schema, newSchema); } @Test public void testCreateSchema() { try (BufferAllocator allocator = new RootAllocator(1024)) { DictionaryEncoding dictionaryEncoding = new DictionaryEncoding(0, true, new ArrowType.Int(8, true)); VarCharVector dictVec = new VarCharVector("dict vector", allocator); Dictionary dictionary = new Dictionary(dictVec, dictionaryEncoding); DictionaryProvider dictProvider = new DictionaryProvider.MapDictionaryProvider(dictionary); TreeSet dictionaryUsed = new TreeSet<>(); FieldType encodedVarcharType = new FieldType(true, new ArrowType.Int(8, true), dictionaryEncoding); FieldType intType = new FieldType(true, new ArrowType.Int(32, true), null); ImmutableList build = ImmutableList.of( new Field("stringCol", encodedVarcharType, null), new Field("intCol", intType, null)); Schema schema = new Schema(build); Schema newSchema = DictionaryUtils.generateSchema(schema, dictProvider, dictionaryUsed); // assert that a new schema is created. assertNotSame(schema, newSchema); // assert the column is converted as expected ArrowType newColType = newSchema.getFields().get(0).getType(); assertEquals(new ArrowType.Utf8(), newColType); assertEquals(1, dictionaryUsed.size()); assertEquals(0, dictionaryUsed.first()); } } } TestDoExchange.java000066400000000000000000000660561515451615200333320ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; 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 java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.Collections; import java.util.stream.IntStream; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; 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.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; public class TestDoExchange { static byte[] EXCHANGE_DO_GET = "do-get".getBytes(StandardCharsets.UTF_8); static byte[] EXCHANGE_DO_PUT = "do-put".getBytes(StandardCharsets.UTF_8); static byte[] EXCHANGE_ECHO = "echo".getBytes(StandardCharsets.UTF_8); static byte[] EXCHANGE_METADATA_ONLY = "only-metadata".getBytes(StandardCharsets.UTF_8); static byte[] EXCHANGE_TRANSFORM = "transform".getBytes(StandardCharsets.UTF_8); static byte[] EXCHANGE_CANCEL = "cancel".getBytes(StandardCharsets.UTF_8); static byte[] EXCHANGE_ERROR = "error".getBytes(StandardCharsets.UTF_8); private BufferAllocator allocator; private FlightServer server; private FlightClient client; @BeforeEach public void setUp() throws Exception { allocator = new RootAllocator(Integer.MAX_VALUE); final Location serverLocation = Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, 0); server = FlightServer.builder(allocator, serverLocation, new Producer(allocator)).build(); server.start(); final Location clientLocation = Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, server.getPort()); client = FlightClient.builder(allocator, clientLocation).build(); } @AfterEach public void tearDown() throws Exception { AutoCloseables.close(client, server, allocator); } /** Test a pure-metadata flow. */ @Test public void testDoExchangeOnlyMetadata() throws Exception { // Send a particular descriptor to the server and check for a particular response pattern. try (final FlightClient.ExchangeReaderWriter stream = client.doExchange(FlightDescriptor.command(EXCHANGE_METADATA_ONLY))) { final FlightStream reader = stream.getReader(); // Server starts by sending a message without data (hence no VectorSchemaRoot should be // present) assertTrue(reader.next()); assertFalse(reader.hasRoot()); assertEquals(42, reader.getLatestMetadata().getInt(0)); // Write a metadata message to the server (without sending any data) ArrowBuf buf = allocator.buffer(4); buf.writeInt(84); stream.getWriter().putMetadata(buf); // Check that the server echoed the metadata back to us assertTrue(reader.next()); assertFalse(reader.hasRoot()); assertEquals(84, reader.getLatestMetadata().getInt(0)); // Close our write channel and ensure the server also closes theirs stream.getWriter().completed(); assertFalse(reader.next()); } } /** Emulate a DoGet with a DoExchange. */ @Test public void testDoExchangeDoGet() throws Exception { try (final FlightClient.ExchangeReaderWriter stream = client.doExchange(FlightDescriptor.command(EXCHANGE_DO_GET))) { final FlightStream reader = stream.getReader(); VectorSchemaRoot root = reader.getRoot(); IntVector iv = (IntVector) root.getVector("a"); int value = 0; while (reader.next()) { for (int i = 0; i < root.getRowCount(); i++) { assertFalse(iv.isNull(i), String.format("Row %d should not be null", value)); assertEquals(value, iv.get(i)); value++; } } assertEquals(100, value); } } /** Emulate a DoPut with a DoExchange. */ @Test public void testDoExchangeDoPut() throws Exception { final Schema schema = new Schema(Collections.singletonList(Field.nullable("a", new ArrowType.Int(32, true)))); try (final FlightClient.ExchangeReaderWriter stream = client.doExchange(FlightDescriptor.command(EXCHANGE_DO_PUT)); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { IntVector iv = (IntVector) root.getVector("a"); iv.allocateNew(); stream.getWriter().start(root); int counter = 0; for (int i = 0; i < 10; i++) { ValueVectorDataPopulator.setVector( iv, IntStream.range(0, i).boxed().toArray(Integer[]::new)); root.setRowCount(i); counter += i; stream.getWriter().putNext(); assertTrue(stream.getReader().next()); assertFalse(stream.getReader().hasRoot()); // For each write, the server sends back a metadata message containing the index of the last // written batch final ArrowBuf metadata = stream.getReader().getLatestMetadata(); assertEquals(counter, metadata.getInt(0)); } stream.getWriter().completed(); while (stream.getReader().next()) { // Drain the stream. Otherwise closing the stream sends a CANCEL which seriously screws with // the server. // CANCEL -> runs onCancel handler -> closes the FlightStream early } } } /** Test a DoExchange that echoes the client message. */ @Test public void testDoExchangeEcho() throws Exception { final Schema schema = new Schema(Collections.singletonList(Field.nullable("a", new ArrowType.Int(32, true)))); try (final FlightClient.ExchangeReaderWriter stream = client.doExchange(FlightDescriptor.command(EXCHANGE_ECHO)); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final FlightStream reader = stream.getReader(); // First try writing metadata without starting the Arrow data stream ArrowBuf buf = allocator.buffer(4); buf.writeInt(42); stream.getWriter().putMetadata(buf); buf = allocator.buffer(4); buf.writeInt(84); stream.getWriter().putMetadata(buf); // Ensure that the server echoes the metadata back, also without starting its data stream assertTrue(reader.next()); assertFalse(reader.hasRoot()); assertEquals(42, reader.getLatestMetadata().getInt(0)); assertTrue(reader.next()); assertFalse(reader.hasRoot()); assertEquals(84, reader.getLatestMetadata().getInt(0)); // Write data and check that it gets echoed back. IntVector iv = (IntVector) root.getVector("a"); iv.allocateNew(); stream.getWriter().start(root); for (int i = 0; i < 10; i++) { iv.setSafe(0, i); root.setRowCount(1); stream.getWriter().putNext(); assertTrue(reader.next()); assertNull(reader.getLatestMetadata()); assertEquals(root.getSchema(), reader.getSchema()); assertEquals(i, ((IntVector) reader.getRoot().getVector("a")).get(0)); } // Complete the stream so that the server knows not to expect any more messages from us. stream.getWriter().completed(); // The server will end its side of the call, so this shouldn't block or indicate that // there is more data. assertFalse(reader.next(), "We should not be waiting for any messages"); } } /** Write some data, have it transformed, then read it back. */ @Test public void testTransform() throws Exception { final Schema schema = new Schema( Arrays.asList( Field.nullable("a", new ArrowType.Int(32, true)), Field.nullable("b", new ArrowType.Int(32, true)))); try (final FlightClient.ExchangeReaderWriter stream = client.doExchange(FlightDescriptor.command(EXCHANGE_TRANSFORM))) { // Write ten batches of data to the stream, where batch N contains N rows of data (N in [0, // 10)) final FlightStream reader = stream.getReader(); final FlightClient.ClientStreamListener writer = stream.getWriter(); try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { writer.start(root); for (int batchIndex = 0; batchIndex < 10; batchIndex++) { for (final FieldVector rawVec : root.getFieldVectors()) { final IntVector vec = (IntVector) rawVec; ValueVectorDataPopulator.setVector( vec, IntStream.range(0, batchIndex).boxed().toArray(Integer[]::new)); } root.setRowCount(batchIndex); writer.putNext(); } } // Indicate that we're done writing so that the server does not expect more data. writer.completed(); // Read back data. We expect the server to double each value in each row of each batch. assertEquals(schema, reader.getSchema()); final VectorSchemaRoot root = reader.getRoot(); for (int batchIndex = 0; batchIndex < 10; batchIndex++) { assertTrue(reader.next(), "Didn't receive batch #" + batchIndex); assertEquals(batchIndex, root.getRowCount()); for (final FieldVector rawVec : root.getFieldVectors()) { final IntVector vec = (IntVector) rawVec; for (int row = 0; row < batchIndex; row++) { assertEquals(2 * row, vec.get(row)); } } } // The server also sends back a metadata-only message containing the message count assertTrue(reader.next(), "There should be one extra message"); assertEquals(10, reader.getLatestMetadata().getInt(0)); assertFalse(reader.next(), "There should be no more data"); } } /** Write some data, have it transformed, then read it back. Use the zero-copy optimization. */ @Test public void testTransformZeroCopy() throws Exception { final int rowsPerBatch = 4096; final Schema schema = new Schema( Arrays.asList( Field.nullable("a", new ArrowType.Int(32, true)), Field.nullable("b", new ArrowType.Int(32, true)))); try (final FlightClient.ExchangeReaderWriter stream = client.doExchange(FlightDescriptor.command(EXCHANGE_TRANSFORM))) { // Write ten batches of data to the stream, where batch N contains 1024 rows of data (N in [0, // 10)) final FlightStream reader = stream.getReader(); final FlightClient.ClientStreamListener writer = stream.getWriter(); try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { writer.start(root); // Enable the zero-copy optimization writer.setUseZeroCopy(true); for (int batchIndex = 0; batchIndex < 100; batchIndex++) { for (final FieldVector rawVec : root.getFieldVectors()) { final IntVector vec = (IntVector) rawVec; for (int row = 0; row < rowsPerBatch; row++) { // Use a value that'll be different per batch, so we can detect if we accidentally // reuse a buffer (and overwrite a buffer that hasn't yet been sent over the network) vec.setSafe(row, batchIndex + row); } } root.setRowCount(rowsPerBatch); writer.putNext(); // Allocate new buffers every time since we don't know if gRPC has written the buffer // to the network yet root.allocateNew(); } } // Indicate that we're done writing so that the server does not expect more data. writer.completed(); // Read back data. We expect the server to double each value in each row of each batch. assertEquals(schema, reader.getSchema()); final VectorSchemaRoot root = reader.getRoot(); for (int batchIndex = 0; batchIndex < 100; batchIndex++) { assertTrue(reader.next(), "Didn't receive batch #" + batchIndex); assertEquals(rowsPerBatch, root.getRowCount()); for (final FieldVector rawVec : root.getFieldVectors()) { final IntVector vec = (IntVector) rawVec; for (int row = 0; row < rowsPerBatch; row++) { assertEquals(2 * (batchIndex + row), vec.get(row)); } } } // The server also sends back a metadata-only message containing the message count assertTrue(reader.next(), "There should be one extra message"); assertEquals(100, reader.getLatestMetadata().getInt(0)); assertFalse(reader.next(), "There should be no more data"); } } /** Have the server immediately cancel; ensure the client doesn't hang. */ @Test public void testServerCancel() throws Exception { try (final FlightClient.ExchangeReaderWriter stream = client.doExchange(FlightDescriptor.command(EXCHANGE_CANCEL))) { final FlightStream reader = stream.getReader(); final FlightClient.ClientStreamListener writer = stream.getWriter(); final FlightRuntimeException fre = assertThrows(FlightRuntimeException.class, reader::next); assertEquals(FlightStatusCode.CANCELLED, fre.status().code()); assertEquals("expected", fre.status().description()); // Before, this would hang forever, because the writer checks if the stream is ready and not // cancelled. // However, the cancellation flag (was) only updated by reading, and the stream is never ready // once the call ends. // The test looks weird since normally, an application shouldn't try to write after the read // fails. However, // an application that isn't reading data wouldn't notice, and would instead get stuck on the // write. // Here, we read first to avoid a race condition in the test itself. writer.putMetadata(allocator.getEmpty()); } } /** Have the server immediately cancel; ensure the server cleans up the FlightStream. */ @Test public void testServerCancelLeak() throws Exception { try (final FlightClient.ExchangeReaderWriter stream = client.doExchange(FlightDescriptor.command(EXCHANGE_CANCEL))) { final FlightStream reader = stream.getReader(); final FlightClient.ClientStreamListener writer = stream.getWriter(); try (final VectorSchemaRoot root = VectorSchemaRoot.create(Producer.SCHEMA, allocator)) { writer.start(root); final IntVector ints = (IntVector) root.getVector("a"); for (int i = 0; i < 128; i++) { for (int row = 0; row < 1024; row++) { ints.setSafe(row, row); } root.setRowCount(1024); writer.putNext(); } } final FlightRuntimeException fre = assertThrows(FlightRuntimeException.class, reader::next); assertEquals(FlightStatusCode.CANCELLED, fre.status().code()); assertEquals("expected", fre.status().description()); } } /** Have the client cancel without reading; ensure memory is not leaked. */ @Test @Disabled public void testClientCancel() throws Exception { try (final FlightClient.ExchangeReaderWriter stream = client.doExchange(FlightDescriptor.command(EXCHANGE_DO_GET))) { final FlightStream reader = stream.getReader(); reader.cancel("", null); // Cancel should be idempotent reader.cancel("", null); } } /** Test a DoExchange error handling. */ @Test public void testDoExchangeError() throws Exception { final Schema schema = new Schema(Collections.singletonList(Field.nullable("a", new ArrowType.Int(32, true)))); try (final FlightClient.ExchangeReaderWriter stream = client.doExchange(FlightDescriptor.command(EXCHANGE_ERROR)); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final FlightStream reader = stream.getReader(); // Write data and check that it gets echoed back. IntVector iv = (IntVector) root.getVector("a"); iv.allocateNew(); stream.getWriter().start(root); for (int i = 0; i < 10; i++) { iv.setSafe(0, i); root.setRowCount(1); stream.getWriter().putNext(); assertTrue(reader.next()); assertEquals(root.getSchema(), reader.getSchema()); assertEquals(i, ((IntVector) reader.getRoot().getVector("a")).get(0)); } // Complete the stream so that the server knows not to expect any more messages from us. stream.getWriter().completed(); // Must call reader.next() to get any errors after exchange, will return false if no error final FlightRuntimeException fre = assertThrows(FlightRuntimeException.class, stream::getResult); assertEquals("error completing exchange", fre.status().description()); } } /** Have the client close the stream without reading; ensure memory is not leaked. */ @Test public void testClientClose() throws Exception { try (final FlightClient.ExchangeReaderWriter stream = client.doExchange(FlightDescriptor.command(EXCHANGE_DO_GET))) { assertEquals(Producer.SCHEMA, stream.getReader().getSchema()); } // Intentionally leak the allocator in this test. gRPC has a bug where it does not wait for all // calls to complete // when shutting down the server, so this test will fail otherwise because it closes the // allocator while the // server-side call still has memory allocated. // TODO(ARROW-9586): fix this once we track outstanding RPCs outside of gRPC. // https://stackoverflow.com/questions/46716024/ allocator = null; client = null; } /** Test closing with Metadata can't lead to error. */ @Test public void testCloseWithMetadata() throws Exception { // Send a particular descriptor to the server and check for a particular response pattern. try (final FlightClient.ExchangeReaderWriter stream = client.doExchange(FlightDescriptor.command(EXCHANGE_METADATA_ONLY))) { final FlightStream reader = stream.getReader(); // Server starts by sending a message without data (hence no VectorSchemaRoot should be // present) assertTrue(reader.next()); assertFalse(reader.hasRoot()); assertEquals(42, reader.getLatestMetadata().getInt(0)); // Write a metadata message to the server (without sending any data) ArrowBuf buf = allocator.buffer(4); buf.writeInt(84); stream.getWriter().putMetadata(buf); // Check that the server echoed the metadata back to us assertTrue(reader.next()); assertFalse(reader.hasRoot()); assertEquals(84, reader.getLatestMetadata().getInt(0)); // Close our write channel and ensure the server also closes theirs stream.getWriter().completed(); stream.getResult(); // Not necessary to close reader here, but check closing twice doesn't lead to negative refcnt // from metadata stream.getReader().close(); } } static class Producer extends NoOpFlightProducer { static final Schema SCHEMA = new Schema(Collections.singletonList(Field.nullable("a", new ArrowType.Int(32, true)))); private final BufferAllocator allocator; Producer(BufferAllocator allocator) { this.allocator = allocator; } @Override public void doExchange(CallContext context, FlightStream reader, ServerStreamListener writer) { if (Arrays.equals(reader.getDescriptor().getCommand(), EXCHANGE_METADATA_ONLY)) { metadataOnly(context, reader, writer); } else if (Arrays.equals(reader.getDescriptor().getCommand(), EXCHANGE_DO_GET)) { doGet(context, reader, writer); } else if (Arrays.equals(reader.getDescriptor().getCommand(), EXCHANGE_DO_PUT)) { doPut(context, reader, writer); } else if (Arrays.equals(reader.getDescriptor().getCommand(), EXCHANGE_ECHO)) { echo(context, reader, writer); } else if (Arrays.equals(reader.getDescriptor().getCommand(), EXCHANGE_TRANSFORM)) { transform(context, reader, writer); } else if (Arrays.equals(reader.getDescriptor().getCommand(), EXCHANGE_CANCEL)) { cancel(context, reader, writer); } else if (Arrays.equals(reader.getDescriptor().getCommand(), EXCHANGE_ERROR)) { error(context, reader, writer); } else { writer.error( CallStatus.UNIMPLEMENTED .withDescription("Command not implemented") .toRuntimeException()); } } /** Emulate DoGet. */ private void doGet( CallContext unusedContext, FlightStream unusedReader, ServerStreamListener writer) { try (VectorSchemaRoot root = VectorSchemaRoot.create(SCHEMA, allocator)) { writer.start(root); root.allocateNew(); IntVector iv = (IntVector) root.getVector("a"); for (int i = 0; i < 100; i += 2) { iv.set(0, i); iv.set(1, i + 1); root.setRowCount(2); writer.putNext(); } } writer.completed(); } /** Emulate DoPut. */ private void doPut( CallContext unusedContext, FlightStream reader, ServerStreamListener writer) { int counter = 0; while (reader.next()) { if (!reader.hasRoot()) { writer.error( CallStatus.INVALID_ARGUMENT .withDescription("Message has no data") .toRuntimeException()); return; } counter += reader.getRoot().getRowCount(); final ArrowBuf pong = allocator.buffer(4); pong.writeInt(counter); writer.putMetadata(pong); } writer.completed(); } /** Exchange metadata without ever exchanging data. */ private void metadataOnly( CallContext unusedContext, FlightStream reader, ServerStreamListener writer) { final ArrowBuf buf = allocator.buffer(4); buf.writeInt(42); writer.putMetadata(buf); assertTrue(reader.next()); assertNotNull(reader.getLatestMetadata()); reader.getLatestMetadata().getReferenceManager().retain(); writer.putMetadata(reader.getLatestMetadata()); writer.completed(); } /** Echo the client's response back to it. */ private void echo(CallContext unusedContext, FlightStream reader, ServerStreamListener writer) { VectorSchemaRoot root = null; VectorLoader loader = null; while (reader.next()) { if (reader.hasRoot()) { if (root == null) { root = VectorSchemaRoot.create(reader.getSchema(), allocator); loader = new VectorLoader(root); writer.start(root); } VectorUnloader unloader = new VectorUnloader(reader.getRoot()); try (final ArrowRecordBatch arb = unloader.getRecordBatch()) { loader.load(arb); } if (reader.getLatestMetadata() != null) { reader.getLatestMetadata().getReferenceManager().retain(); writer.putNext(reader.getLatestMetadata()); } else { writer.putNext(); } } else { // Pure metadata reader.getLatestMetadata().getReferenceManager().retain(); writer.putMetadata(reader.getLatestMetadata()); } } if (root != null) { root.close(); } writer.completed(); } /** Accept a set of messages, then return some result. */ private void transform( CallContext unusedContext, FlightStream reader, ServerStreamListener writer) { final Schema schema = reader.getSchema(); for (final Field field : schema.getFields()) { if (!(field.getType() instanceof ArrowType.Int)) { writer.error( CallStatus.INVALID_ARGUMENT .withDescription("Invalid type: " + field) .toRuntimeException()); return; } final ArrowType.Int intType = (ArrowType.Int) field.getType(); if (!intType.getIsSigned() || intType.getBitWidth() != 32) { writer.error( CallStatus.INVALID_ARGUMENT .withDescription("Must be i32: " + field) .toRuntimeException()); return; } } int batches = 0; try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { writer.start(root); writer.setUseZeroCopy(true); final VectorLoader loader = new VectorLoader(root); final VectorUnloader unloader = new VectorUnloader(reader.getRoot()); while (reader.next()) { try (final ArrowRecordBatch batch = unloader.getRecordBatch()) { loader.load(batch); } batches++; for (final FieldVector rawVec : root.getFieldVectors()) { final IntVector vec = (IntVector) rawVec; for (int i = 0; i < root.getRowCount(); i++) { if (!vec.isNull(i)) { vec.set(i, vec.get(i) * 2); } } } writer.putNext(); } } final ArrowBuf count = allocator.buffer(4); count.writeInt(batches); writer.putMetadata(count); writer.completed(); } /** Immediately cancel the call. */ private void cancel( CallContext unusedContext, FlightStream unusedReader, ServerStreamListener writer) { writer.error(CallStatus.CANCELLED.withDescription("expected").toRuntimeException()); } private void error( CallContext unusedContext, FlightStream reader, ServerStreamListener writer) { VectorSchemaRoot root = null; VectorLoader loader = null; while (reader.next()) { if (root == null) { root = VectorSchemaRoot.create(reader.getSchema(), allocator); loader = new VectorLoader(root); writer.start(root); } VectorUnloader unloader = new VectorUnloader(reader.getRoot()); try (final ArrowRecordBatch arb = unloader.getRecordBatch()) { loader.load(arb); } writer.putNext(); } if (root != null) { root.close(); } // An error occurs before completing the writer writer.error( CallStatus.INTERNAL.withDescription("error completing exchange").toRuntimeException()); } } } TestErrorMetadata.java000066400000000000000000000155721515451615200340540ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; 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; import com.google.protobuf.Any; import com.google.protobuf.InvalidProtocolBufferException; import com.google.rpc.Status; import io.grpc.Metadata; import io.grpc.StatusRuntimeException; import io.grpc.protobuf.ProtoUtils; import io.grpc.protobuf.StatusProto; import java.nio.charset.StandardCharsets; import org.apache.arrow.flight.perf.impl.PerfOuterClass; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.Test; public class TestErrorMetadata { private static final Metadata.BinaryMarshaller marshaller = ProtoUtils.metadataMarshaller(Status.getDefaultInstance()); /** Ensure metadata attached to a gRPC error is propagated. */ @Test public void testGrpcMetadata() throws Exception { PerfOuterClass.Perf perf = PerfOuterClass.Perf.newBuilder() .setStreamCount(12) .setRecordsPerBatch(1000) .setRecordsPerStream(1000000L) .build(); StatusRuntimeExceptionProducer producer = new StatusRuntimeExceptionProducer(perf); try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); final FlightServer s = FlightServer.builder(allocator, forGrpcInsecure(LOCALHOST, 0), producer) .build() .start(); final FlightClient client = FlightClient.builder(allocator, s.getLocation()).build()) { final CallStatus flightStatus = FlightTestUtil.assertCode( FlightStatusCode.CANCELLED, () -> { FlightStream stream = client.getStream(new Ticket("abs".getBytes(StandardCharsets.UTF_8))); stream.next(); }); PerfOuterClass.Perf newPerf = null; ErrorFlightMetadata metadata = flightStatus.metadata(); assertNotNull(metadata); assertEquals(2, metadata.keys().size()); assertTrue(metadata.containsKey("grpc-status-details-bin")); Status status = marshaller.parseBytes(metadata.getByte("grpc-status-details-bin")); for (Any details : status.getDetailsList()) { if (details.is(PerfOuterClass.Perf.class)) { try { newPerf = details.unpack(PerfOuterClass.Perf.class); } catch (InvalidProtocolBufferException e) { fail(); } } } assertNotNull(newPerf); assertEquals(perf, newPerf); } } /** Ensure metadata attached to a Flight error is propagated. */ @Test public void testFlightMetadata() throws Exception { try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); final FlightServer s = FlightServer.builder(allocator, forGrpcInsecure(LOCALHOST, 0), new CallStatusProducer()) .build() .start(); final FlightClient client = FlightClient.builder(allocator, s.getLocation()).build()) { CallStatus flightStatus = FlightTestUtil.assertCode( FlightStatusCode.INVALID_ARGUMENT, () -> { FlightStream stream = client.getStream(new Ticket(new byte[0])); stream.next(); }); ErrorFlightMetadata metadata = flightStatus.metadata(); assertNotNull(metadata); assertEquals("foo", metadata.get("x-foo")); assertArrayEquals(new byte[] {1}, metadata.getByte("x-bar-bin")); flightStatus = FlightTestUtil.assertCode( FlightStatusCode.INVALID_ARGUMENT, () -> { client.getInfo(FlightDescriptor.command(new byte[0])); }); metadata = flightStatus.metadata(); assertNotNull(metadata); assertEquals("foo", metadata.get("x-foo")); assertArrayEquals(new byte[] {1}, metadata.getByte("x-bar-bin")); } } @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; private StatusRuntimeExceptionProducer(PerfOuterClass.Perf perf) { this.perf = perf; } @Override public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { StatusRuntimeException sre = StatusProto.toStatusRuntimeException( Status.newBuilder() .setCode(1) .setMessage("Testing 1 2 3") .addDetails(Any.pack(perf, "arrow/meta/types")) .build()); listener.error(sre); } } private static class CallStatusProducer extends NoOpFlightProducer { ErrorFlightMetadata metadata; CallStatusProducer() { this.metadata = new ErrorFlightMetadata(); metadata.insert("x-foo", "foo"); metadata.insert("x-bar-bin", new byte[] {1}); } @Override public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { listener.error( CallStatus.INVALID_ARGUMENT .withDescription("Failed") .withMetadata(metadata) .toRuntimeException()); } @Override public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) { throw CallStatus.INVALID_ARGUMENT .withDescription("Failed") .withMetadata(metadata) .toRuntimeException(); } } } TestFlightClient.java000066400000000000000000000262471515451615200336770ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; 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.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.charset.StandardCharsets; import java.util.Collections; import java.util.List; import org.apache.arrow.flight.FlightClient.ClientStreamListener; import org.apache.arrow.flight.TestBasicOperation.Producer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; 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.ipc.message.ArrowRecordBatch; 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.types.pojo.Schema; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; public class TestFlightClient { /** ARROW-5063: make sure two clients to the same location can be closed independently. */ @Test public void independentShutdown() throws Exception { try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); final FlightServer server = FlightServer.builder(allocator, forGrpcInsecure(LOCALHOST, 0), new Producer(allocator)) .build() .start()) { final Schema schema = new Schema(Collections.singletonList(Field.nullable("a", new ArrowType.Int(32, true)))); try (final FlightClient client1 = FlightClient.builder(allocator, server.getLocation()).build(); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { // Use startPut as this ensures the RPC won't finish until we want it to final ClientStreamListener listener = client1.startPut(FlightDescriptor.path("test"), root, new AsyncPutListener()); try (final FlightClient client2 = FlightClient.builder(allocator, server.getLocation()).build()) { client2.listActions().forEach(actionType -> assertNotNull(actionType.getType())); } listener.completed(); listener.getResult(); } } } /** * ARROW-5978: make sure that we can properly close a client/stream after requesting dictionaries. */ @Disabled // Unfortunately this test is flaky in CI. @Test public void freeDictionaries() throws Exception { final Schema expectedSchema = new Schema( Collections.singletonList( new Field( "encoded", new FieldType( true, new ArrowType.Int(32, true), new DictionaryEncoding(1L, false, null)), null))); try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); final BufferAllocator serverAllocator = allocator.newChildAllocator("flight-server", 0, Integer.MAX_VALUE); final FlightServer server = FlightServer.builder( serverAllocator, forGrpcInsecure(LOCALHOST, 0), new DictionaryProducer(serverAllocator)) .build() .start()) { try (final FlightClient client = FlightClient.builder(allocator, server.getLocation()).build()) { try (final FlightStream stream = client.getStream(new Ticket(new byte[0]))) { assertTrue(stream.next()); assertNotNull(stream.getDictionaryProvider().lookup(1)); final VectorSchemaRoot root = stream.getRoot(); assertEquals(expectedSchema, root.getSchema()); assertEquals(6, root.getVector("encoded").getValueCount()); try (final ValueVector decoded = DictionaryEncoder.decode( root.getVector("encoded"), stream.getDictionaryProvider().lookup(1))) { assertFalse(decoded.isNull(1)); assertTrue(decoded instanceof VarCharVector); assertArrayEquals( "one".getBytes(StandardCharsets.UTF_8), ((VarCharVector) decoded).get(1)); } assertFalse(stream.next()); } // Closing stream fails if it doesn't free dictionaries; closing dictionaries fails // (refcount goes negative) // if reference isn't retained in ArrowMessage } } } /** ARROW-5978: make sure that dictionary ownership can't be claimed twice. */ @Disabled // Unfortunately this test is flaky in CI. @Test public void ownDictionaries() throws Exception { try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); final BufferAllocator serverAllocator = allocator.newChildAllocator("flight-server", 0, Integer.MAX_VALUE); final FlightServer server = FlightServer.builder( serverAllocator, forGrpcInsecure(LOCALHOST, 0), new DictionaryProducer(serverAllocator)) .build() .start()) { try (final FlightClient client = FlightClient.builder(allocator, server.getLocation()).build()) { try (final FlightStream stream = client.getStream(new Ticket(new byte[0]))) { assertTrue(stream.next()); assertFalse(stream.next()); final DictionaryProvider provider = stream.takeDictionaryOwnership(); assertThrows(IllegalStateException.class, stream::takeDictionaryOwnership); assertThrows(IllegalStateException.class, stream::getDictionaryProvider); DictionaryUtils.closeDictionaries(stream.getSchema(), provider); } } } } /** ARROW-5978: make sure that dictionaries can be used after closing the stream. */ @Disabled // Unfortunately this test is flaky in CI. @Test public void useDictionariesAfterClose() throws Exception { try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); final BufferAllocator serverAllocator = allocator.newChildAllocator("flight-server", 0, Integer.MAX_VALUE); final FlightServer server = FlightServer.builder( serverAllocator, forGrpcInsecure(LOCALHOST, 0), new DictionaryProducer(serverAllocator)) .build() .start()) { try (final FlightClient client = FlightClient.builder(allocator, server.getLocation()).build()) { final VectorSchemaRoot root; final DictionaryProvider provider; try (final FlightStream stream = client.getStream(new Ticket(new byte[0]))) { final VectorUnloader unloader = new VectorUnloader(stream.getRoot()); root = VectorSchemaRoot.create(stream.getSchema(), allocator); final VectorLoader loader = new VectorLoader(root); while (stream.next()) { try (final ArrowRecordBatch arb = unloader.getRecordBatch()) { loader.load(arb); } } provider = stream.takeDictionaryOwnership(); } try (final ValueVector decoded = DictionaryEncoder.decode(root.getVector("encoded"), provider.lookup(1))) { assertFalse(decoded.isNull(1)); assertTrue(decoded instanceof VarCharVector); assertArrayEquals( "one".getBytes(StandardCharsets.UTF_8), ((VarCharVector) decoded).get(1)); } root.close(); DictionaryUtils.closeDictionaries(root.getSchema(), provider); } } } static class DictionaryProducer extends NoOpFlightProducer { private final BufferAllocator allocator; public DictionaryProducer(BufferAllocator allocator) { this.allocator = allocator; } @Override public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { final byte[] zero = "zero".getBytes(StandardCharsets.UTF_8); final byte[] one = "one".getBytes(StandardCharsets.UTF_8); final byte[] two = "two".getBytes(StandardCharsets.UTF_8); try (final VarCharVector dictionaryVector = newVarCharVector("dictionary", allocator)) { final DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); dictionaryVector.allocateNew(512, 3); dictionaryVector.setSafe(0, zero, 0, zero.length); dictionaryVector.setSafe(1, one, 0, one.length); dictionaryVector.setSafe(2, two, 0, two.length); dictionaryVector.setValueCount(3); final Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); provider.put(dictionary); final FieldVector encodedVector; try (final VarCharVector unencoded = newVarCharVector("encoded", allocator)) { unencoded.allocateNewSafe(); unencoded.set(1, one); unencoded.set(2, two); unencoded.set(3, zero); unencoded.set(4, two); unencoded.setValueCount(6); encodedVector = (FieldVector) DictionaryEncoder.encode(unencoded, dictionary); } final List fields = Collections.singletonList(encodedVector.getField()); final List vectors = Collections.singletonList(encodedVector); try (final VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, encodedVector.getValueCount())) { listener.start(root, provider); listener.putNext(); listener.completed(); } } } private static VarCharVector newVarCharVector(String name, BufferAllocator allocator) { return (VarCharVector) FieldType.nullable(new ArrowType.Utf8()).createNewSingleVector(name, allocator, null); } } } TestFlightGrpcUtils.java000066400000000000000000000172011515451615200343630ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; 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 static org.junit.jupiter.api.Assertions.fail; import com.google.protobuf.Empty; import io.grpc.BindableService; import io.grpc.ConnectivityState; import io.grpc.ManagedChannel; import io.grpc.Server; import io.grpc.inprocess.InProcessChannelBuilder; import io.grpc.inprocess.InProcessServerBuilder; import io.grpc.stub.StreamObserver; import java.io.IOException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import org.apache.arrow.flight.auth.ServerAuthHandler; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Unit test which adds 2 services to same server end point. */ public class TestFlightGrpcUtils { private Server server; private BufferAllocator allocator; private String serverName; @BeforeEach public void setup() throws IOException { // Defines flight service allocator = new RootAllocator(Integer.MAX_VALUE); final NoOpFlightProducer producer = new NoOpFlightProducer(); final ServerAuthHandler authHandler = ServerAuthHandler.NO_OP; final ExecutorService exec = Executors.newCachedThreadPool(); final BindableService flightBindingService = FlightGrpcUtils.createFlightService(allocator, producer, authHandler, exec); // initializes server with 2 services - FlightBindingService & TestService serverName = InProcessServerBuilder.generateName(); server = InProcessServerBuilder.forName(serverName) .directExecutor() .addService(flightBindingService) .addService(new TestServiceAdapter()) .build(); server.start(); } @AfterEach public void cleanup() { server.shutdownNow(); } /** * This test checks if multiple gRPC services can be added to the same server endpoint and if they * can be used by different clients via the same channel. * * @throws IOException If server fails to start. */ @Test public void testMultipleGrpcServices() throws IOException { // Initializes channel so that multiple clients can communicate with server final ManagedChannel managedChannel = InProcessChannelBuilder.forName(serverName).directExecutor().build(); // Defines flight client and calls service method. Since we use a NoOpFlightProducer we expect // the service // to throw a RunTimeException final FlightClient flightClient = FlightGrpcUtils.createFlightClient(allocator, managedChannel); final Iterable actionTypes = flightClient.listActions(); assertThrows( FlightRuntimeException.class, () -> actionTypes.forEach(actionType -> System.out.println(actionType.toString()))); // Define Test client as a blocking stub and call test method which correctly returns an empty // protobuf object final TestServiceGrpc.TestServiceBlockingStub blockingStub = TestServiceGrpc.newBlockingStub(managedChannel); assertEquals(Empty.newBuilder().build(), blockingStub.test(Empty.newBuilder().build())); } @Test public void testShutdown() throws IOException, InterruptedException { // Initializes channel so that multiple clients can communicate with server final ManagedChannel managedChannel = InProcessChannelBuilder.forName(serverName).directExecutor().build(); // Defines flight client and calls service method. Since we use a NoOpFlightProducer we expect // the service // to throw a RunTimeException final FlightClient flightClient = FlightGrpcUtils.createFlightClientWithSharedChannel(allocator, managedChannel); // Should be a no-op. flightClient.close(); assertFalse(managedChannel.isShutdown()); assertFalse(managedChannel.isTerminated()); assertEquals(ConnectivityState.IDLE, managedChannel.getState(false)); managedChannel.shutdownNow(); } @Test public void testProxyChannel() throws IOException, InterruptedException { // Initializes channel so that multiple clients can communicate with server final ManagedChannel managedChannel = InProcessChannelBuilder.forName(serverName).directExecutor().build(); final FlightGrpcUtils.NonClosingProxyManagedChannel proxyChannel = new FlightGrpcUtils.NonClosingProxyManagedChannel(managedChannel); assertFalse(proxyChannel.isShutdown()); assertFalse(proxyChannel.isTerminated()); proxyChannel.shutdown(); assertTrue(proxyChannel.isShutdown()); assertTrue(proxyChannel.isTerminated()); assertEquals(ConnectivityState.SHUTDOWN, proxyChannel.getState(false)); try { proxyChannel.newCall(null, null); fail(); } catch (IllegalStateException e) { // This is expected, since the proxy channel is shut down. } assertFalse(managedChannel.isShutdown()); assertFalse(managedChannel.isTerminated()); assertEquals(ConnectivityState.IDLE, managedChannel.getState(false)); managedChannel.shutdownNow(); } @Test public void testProxyChannelWithClosedChannel() throws IOException, InterruptedException { // Initializes channel so that multiple clients can communicate with server final ManagedChannel managedChannel = InProcessChannelBuilder.forName(serverName).directExecutor().build(); final FlightGrpcUtils.NonClosingProxyManagedChannel proxyChannel = new FlightGrpcUtils.NonClosingProxyManagedChannel(managedChannel); assertFalse(proxyChannel.isShutdown()); assertFalse(proxyChannel.isTerminated()); managedChannel.shutdownNow(); assertTrue(proxyChannel.isShutdown()); assertTrue(proxyChannel.isTerminated()); assertEquals(ConnectivityState.SHUTDOWN, proxyChannel.getState(false)); try { proxyChannel.newCall(null, null); fail(); } catch (IllegalStateException e) { // This is expected, since the proxy channel is shut down. } assertTrue(managedChannel.isShutdown()); assertTrue(managedChannel.isTerminated()); assertEquals(ConnectivityState.SHUTDOWN, managedChannel.getState(false)); } /** Private class used for testing purposes that overrides service behavior. */ private static class TestServiceAdapter extends TestServiceGrpc.TestServiceImplBase { /** * gRPC service that receives an empty object & returns and empty protobuf object. * * @param request google.protobuf.Empty * @param responseObserver google.protobuf.Empty */ @Override public void test(Empty request, StreamObserver responseObserver) { responseObserver.onNext(Empty.newBuilder().build()); responseObserver.onCompleted(); } } } TestFlightService.java000066400000000000000000000205721515451615200340540ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; 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.assertThrows; import static org.junit.jupiter.api.Assertions.fail; import io.grpc.stub.ServerCallStreamObserver; import java.nio.charset.StandardCharsets; import java.util.Collections; import java.util.Optional; import java.util.Random; import org.apache.arrow.flight.impl.Flight; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.ipc.message.IpcOption; 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; public class TestFlightService { private BufferAllocator allocator; @BeforeEach public void setup() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void cleanup() throws Exception { AutoCloseables.close(allocator); } @Test public void testFlightServiceWithNoAuthHandlerOrInterceptors() { // This test is for ARROW-10491. There was a bug where FlightService would try to access the // RequestContext, // but the RequestContext was getting set to null because no interceptors were active to // initialize it // when using FlightService directly rather than starting up a FlightServer. // Arrange final FlightProducer producer = new NoOpFlightProducer() { @Override public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { listener.completed(); } }; // This response observer notifies that the test failed if onError() is called. final ServerCallStreamObserver observer = new ServerCallStreamObserver() { @Override public boolean isCancelled() { return false; } @Override public void setOnCancelHandler(Runnable runnable) {} @Override public void setCompression(String s) {} @Override public boolean isReady() { return false; } @Override public void setOnReadyHandler(Runnable runnable) {} @Override public void disableAutoInboundFlowControl() {} @Override public void request(int i) {} @Override public void setMessageCompression(boolean b) {} @Override public void onNext(ArrowMessage arrowMessage) {} @Override public void onError(Throwable throwable) { fail(throwable); } @Override public void onCompleted() {} }; final FlightService flightService = new FlightService(allocator, producer, null, null); // Act flightService.doGetCustom(Flight.Ticket.newBuilder().build(), observer); // fail() would have been called if an error happened during doGetCustom(), so this test passed. } @Test public void supportsNullSchemas() throws Exception { final FlightProducer producer = new NoOpFlightProducer() { @Override public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) { return new FlightInfo( null, descriptor, Collections.emptyList(), 0, 0, false, IpcOption.DEFAULT, "foo".getBytes(StandardCharsets.UTF_8)); } }; try (final FlightServer s = FlightServer.builder(allocator, forGrpcInsecure(LOCALHOST, 0), producer) .build() .start(); final FlightClient client = FlightClient.builder(allocator, s.getLocation()).build()) { FlightInfo flightInfo = client.getInfo(FlightDescriptor.path("test")); assertEquals(Optional.empty(), flightInfo.getSchemaOptional()); assertEquals(new Schema(Collections.emptyList()), flightInfo.getSchema()); assertArrayEquals(flightInfo.getAppMetadata(), "foo".getBytes(StandardCharsets.UTF_8)); Exception e = assertThrows( FlightRuntimeException.class, () -> client.getSchema(FlightDescriptor.path("test"))); assertEquals("No schema is present in FlightInfo", e.getMessage()); } } /** * Test for GH-41584 where flight defaults for header size was not in sync b\w client and server. */ @Test public void testHeaderSizeExchangeInService() throws Exception { final FlightProducer producer = new NoOpFlightProducer() { @Override public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) { String longHeader = context.getMiddleware(FlightConstants.HEADER_KEY).headers().get("long-header"); return new FlightInfo( null, descriptor, Collections.emptyList(), 0, 0, false, IpcOption.DEFAULT, longHeader.getBytes(StandardCharsets.UTF_8)); } }; String headerVal = generateRandom(1024 * 10); FlightCallHeaders callHeaders = new FlightCallHeaders(); callHeaders.insert("long-header", headerVal); // sever with default header limit same as client try (final FlightServer s = FlightServer.builder(allocator, forGrpcInsecure(LOCALHOST, 0), producer) .build() .start(); final FlightClient client = FlightClient.builder(allocator, s.getLocation()).build()) { FlightInfo flightInfo = client.getInfo(FlightDescriptor.path("test"), new HeaderCallOption(callHeaders)); assertEquals(Optional.empty(), flightInfo.getSchemaOptional()); assertEquals(new Schema(Collections.emptyList()), flightInfo.getSchema()); assertArrayEquals(flightInfo.getAppMetadata(), headerVal.getBytes(StandardCharsets.UTF_8)); } // server with 15kb header limit try (final FlightServer s = FlightServer.builder(allocator, forGrpcInsecure(LOCALHOST, 0), producer) .setMaxHeaderListSize(1024 * 15) .build() .start(); final FlightClient client = FlightClient.builder(allocator, s.getLocation()).build()) { FlightInfo flightInfo = client.getInfo(FlightDescriptor.path("test"), new HeaderCallOption(callHeaders)); assertEquals(Optional.empty(), flightInfo.getSchemaOptional()); assertEquals(new Schema(Collections.emptyList()), flightInfo.getSchema()); assertArrayEquals(flightInfo.getAppMetadata(), headerVal.getBytes(StandardCharsets.UTF_8)); callHeaders.insert("another-header", headerVal + headerVal); FlightRuntimeException e = assertThrows( FlightRuntimeException.class, () -> client.getInfo(FlightDescriptor.path("test"), new HeaderCallOption(callHeaders))); assertEquals("http2 exception", e.getMessage()); } } private static String generateRandom(int size) { String aToZ = "ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"; Random random = new Random(); StringBuilder res = new StringBuilder(); for (int i = 0; i < size; i++) { int randIndex = random.nextInt(aToZ.length()); res.append(aToZ.charAt(randIndex)); } return res.toString(); } } TestLargeMessage.java000066400000000000000000000133031515451615200336470ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static com.google.common.collect.ImmutableList.toImmutableList; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; import static org.apache.arrow.flight.Location.forGrpcInsecure; import static org.junit.jupiter.api.Assertions.assertEquals; import java.util.Arrays; import java.util.List; import java.util.stream.Stream; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; 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 org.junit.jupiter.api.Test; public class TestLargeMessage { /** Make sure a Flight client accepts large message payloads by default. */ @Test public void getLargeMessage() throws Exception { try (final BufferAllocator a = new RootAllocator(Long.MAX_VALUE); final Producer producer = new Producer(a); final FlightServer s = FlightServer.builder(a, forGrpcInsecure(LOCALHOST, 0), producer).build().start()) { try (FlightClient client = FlightClient.builder(a, s.getLocation()).build()) { try (FlightStream stream = client.getStream(new Ticket(new byte[] {})); VectorSchemaRoot root = stream.getRoot()) { while (stream.next()) { for (final Field field : root.getSchema().getFields()) { int value = 0; final IntVector iv = (IntVector) root.getVector(field.getName()); for (int i = 0; i < root.getRowCount(); i++) { assertEquals(value, iv.get(i)); value++; } } } } } } } /** Make sure a Flight server accepts large message payloads by default. */ @Test public void putLargeMessage() throws Exception { try (final BufferAllocator a = new RootAllocator(Long.MAX_VALUE); final Producer producer = new Producer(a); final FlightServer s = FlightServer.builder(a, forGrpcInsecure(LOCALHOST, 0), producer).build().start()) { try (FlightClient client = FlightClient.builder(a, s.getLocation()).build(); BufferAllocator testAllocator = a.newChildAllocator("testcase", 0, Long.MAX_VALUE); VectorSchemaRoot root = generateData(testAllocator)) { final FlightClient.ClientStreamListener listener = client.startPut(FlightDescriptor.path("hello"), root, new AsyncPutListener()); listener.putNext(); listener.completed(); listener.getResult(); } } } private static VectorSchemaRoot generateData(BufferAllocator allocator) { final int size = 128 * 1024; final List fieldNames = Arrays.asList("c1", "c2", "c3", "c4", "c5", "c6", "c7", "c8", "c9", "c10"); final Stream fields = fieldNames.stream() .map( fieldName -> new Field(fieldName, FieldType.nullable(new ArrowType.Int(32, true)), null)); final Schema schema = new Schema(fields.collect(toImmutableList()), null); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); root.allocateNew(); for (final String fieldName : fieldNames) { final IntVector iv = (IntVector) root.getVector(fieldName); iv.setValueCount(size); for (int i = 0; i < size; i++) { iv.set(i, i); } } root.setRowCount(size); return root; } private static class Producer implements FlightProducer, AutoCloseable { private final BufferAllocator allocator; Producer(BufferAllocator allocator) { this.allocator = allocator; } @Override public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { try (VectorSchemaRoot root = generateData(allocator)) { listener.start(root); listener.putNext(); listener.completed(); } } @Override public void listFlights( CallContext context, Criteria criteria, StreamListener listener) {} @Override public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) { return null; } @Override public Runnable acceptPut( CallContext context, FlightStream flightStream, StreamListener ackStream) { return () -> { try (VectorSchemaRoot root = flightStream.getRoot()) { while (flightStream.next()) { ; } } }; } @Override public void doAction(CallContext context, Action action, StreamListener listener) { listener.onCompleted(); } @Override public void listActions(CallContext context, StreamListener listener) {} @Override public void close() throws Exception { allocator.close(); } } } apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/TestLeak.java000066400000000000000000000177261515451615200322600ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; import static org.apache.arrow.flight.Location.forGrpcInsecure; import java.util.Arrays; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.FloatingPointPrecision; 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; /** Tests for scenarios where Flight could leak memory. */ public class TestLeak { private static final int ROWS = 2048; private static Schema getSchema() { return new Schema( Arrays.asList( Field.nullable("0", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), Field.nullable("1", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), Field.nullable("2", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), Field.nullable("3", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), Field.nullable("4", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), Field.nullable("5", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), Field.nullable("6", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), Field.nullable("7", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), Field.nullable("8", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), Field.nullable("9", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), Field.nullable("10", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)))); } /** * Ensure that if the client cancels, the server does not leak memory. * *

    In gRPC, canceling the stream from the client sends an event to the server. Once processed, * gRPC will start silently rejecting messages sent by the server. However, Flight depends on gRPC * processing these messages in order to free the associated memory. */ @Test public void testCancelingDoGetDoesNotLeak() throws Exception { final CountDownLatch callFinished = new CountDownLatch(1); try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); final FlightServer s = FlightServer.builder( allocator, forGrpcInsecure(LOCALHOST, 0), new LeakFlightProducer(allocator, callFinished)) .build() .start(); final FlightClient client = FlightClient.builder(allocator, s.getLocation()).build()) { final FlightStream stream = client.getStream(new Ticket(new byte[0])); stream.getRoot(); stream.cancel("Cancel", null); // Wait for the call to finish. (Closing the allocator while a call is ongoing is a guaranteed // leak.) callFinished.await(60, TimeUnit.SECONDS); s.shutdown(); s.awaitTermination(); } } @Test public void testCancelingDoPutDoesNotBlock() throws Exception { final CountDownLatch callFinished = new CountDownLatch(1); try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); final FlightServer s = FlightServer.builder( allocator, forGrpcInsecure(LOCALHOST, 0), new LeakFlightProducer(allocator, callFinished)) .build() .start(); final FlightClient client = FlightClient.builder(allocator, s.getLocation()).build()) { try (final VectorSchemaRoot root = VectorSchemaRoot.create(getSchema(), allocator)) { final FlightDescriptor descriptor = FlightDescriptor.command(new byte[0]); final SyncPutListener listener = new SyncPutListener(); final FlightClient.ClientStreamListener stream = client.startPut(descriptor, root, listener); // Wait for the server to cancel callFinished.await(60, TimeUnit.SECONDS); for (int col = 0; col < 11; col++) { final Float8Vector vector = (Float8Vector) root.getVector(Integer.toString(col)); vector.allocateNew(); for (int row = 0; row < ROWS; row++) { vector.setSafe(row, 10.); } } root.setRowCount(ROWS); // Unlike DoGet, this method fairly reliably will write the message to the stream, so even // without the fix // for ARROW-7343, this won't leak memory. // However, it will block if FlightClient doesn't check for cancellation. stream.putNext(); stream.completed(); } s.shutdown(); s.awaitTermination(); } } /** A FlightProducer that always produces a fixed data stream with metadata on the side. */ private static class LeakFlightProducer extends NoOpFlightProducer { private final BufferAllocator allocator; private final CountDownLatch callFinished; public LeakFlightProducer(BufferAllocator allocator, CountDownLatch callFinished) { this.allocator = allocator; this.callFinished = callFinished; } @Override public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { BufferAllocator childAllocator = allocator.newChildAllocator("foo", 0, Long.MAX_VALUE); VectorSchemaRoot root = VectorSchemaRoot.create(TestLeak.getSchema(), childAllocator); root.allocateNew(); listener.start(root); // We can't poll listener#isCancelled since gRPC has two distinct "is cancelled" flags. // TODO: should we continue leaking gRPC semantics? Can we even avoid this? listener.setOnCancelHandler( () -> { try { for (int col = 0; col < 11; col++) { final Float8Vector vector = (Float8Vector) root.getVector(Integer.toString(col)); vector.allocateNew(); for (int row = 0; row < ROWS; row++) { vector.setSafe(row, 10.); } } root.setRowCount(ROWS); // Once the call is "really cancelled" (setOnCancelListener has run/is running), this // call is actually a // no-op on the gRPC side and will leak the ArrowMessage unless Flight checks for // this. listener.putNext(); listener.completed(); } finally { try { root.close(); childAllocator.close(); } finally { // Don't let the test hang if we throw above callFinished.countDown(); } } }); } @Override public Runnable acceptPut( CallContext context, FlightStream flightStream, StreamListener ackStream) { return () -> { flightStream.getRoot(); ackStream.onError(CallStatus.CANCELLED.withDescription("CANCELLED").toRuntimeException()); callFinished.countDown(); ackStream.onCompleted(); }; } } } TestMetadataVersion.java000066400000000000000000000322061515451615200344010ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; 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.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.Collections; import java.util.Optional; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.message.IpcOption; import org.apache.arrow.vector.types.MetadataVersion; import org.apache.arrow.vector.types.UnionMode; 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.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; /** Test clients/servers with different metadata versions. */ public class TestMetadataVersion { private static BufferAllocator allocator; private static Schema schema; private static IpcOption optionV4; private static IpcOption optionV5; private static Schema unionSchema; @BeforeAll public static void setUpClass() { allocator = new RootAllocator(Integer.MAX_VALUE); schema = new Schema(Collections.singletonList(Field.nullable("foo", new ArrowType.Int(32, true)))); unionSchema = new Schema( Collections.singletonList( Field.nullable("union", new ArrowType.Union(UnionMode.Dense, new int[] {0})))); // avoid writing legacy ipc format by default optionV4 = new IpcOption(false, MetadataVersion.V4); optionV5 = IpcOption.DEFAULT; } @AfterAll public static void tearDownClass() { allocator.close(); } @Test public void testGetFlightInfoV4() throws Exception { try (final FlightServer server = startServer(optionV4); final FlightClient client = connect(server)) { final FlightInfo result = client.getInfo(FlightDescriptor.command(new byte[0])); assertEquals(Optional.of(schema), result.getSchemaOptional()); } } @Test public void testGetSchemaV4() throws Exception { try (final FlightServer server = startServer(optionV4); final FlightClient client = connect(server)) { final SchemaResult result = client.getSchema(FlightDescriptor.command(new byte[0])); assertEquals(schema, result.getSchema()); } } @Test public void testUnionCheck() throws Exception { assertThrows(IllegalArgumentException.class, () -> new SchemaResult(unionSchema, optionV4)); assertThrows( IllegalArgumentException.class, () -> new FlightInfo( unionSchema, FlightDescriptor.command(new byte[0]), Collections.emptyList(), -1, -1, optionV4)); try (final FlightServer server = startServer(optionV4); final FlightClient client = connect(server); final FlightStream stream = client.getStream(new Ticket("union".getBytes(StandardCharsets.UTF_8)))) { final FlightRuntimeException err = assertThrows(FlightRuntimeException.class, stream::next); assertTrue( err.getMessage().contains("Cannot write union with V4 metadata"), err.getMessage()); } try (final FlightServer server = startServer(optionV4); final FlightClient client = connect(server); final VectorSchemaRoot root = VectorSchemaRoot.create(unionSchema, allocator)) { final FlightDescriptor descriptor = FlightDescriptor.command(new byte[0]); final SyncPutListener reader = new SyncPutListener(); final FlightClient.ClientStreamListener listener = client.startPut(descriptor, reader); final IllegalArgumentException err = assertThrows(IllegalArgumentException.class, () -> listener.start(root, null, optionV4)); assertTrue( err.getMessage().contains("Cannot write union with V4 metadata"), err.getMessage()); } } @Test public void testPutV4() throws Exception { try (final FlightServer server = startServer(optionV4); final FlightClient client = connect(server); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { generateData(root); final FlightDescriptor descriptor = FlightDescriptor.command(new byte[0]); final SyncPutListener reader = new SyncPutListener(); final FlightClient.ClientStreamListener listener = client.startPut(descriptor, reader); listener.start(root, null, optionV4); listener.putNext(); listener.completed(); listener.getResult(); } } @Test public void testGetV4() throws Exception { try (final FlightServer server = startServer(optionV4); final FlightClient client = connect(server); final FlightStream stream = client.getStream(new Ticket(new byte[0]))) { assertTrue(stream.next()); assertEquals(optionV4.metadataVersion, stream.metadataVersion); validateRoot(stream.getRoot()); assertFalse(stream.next()); } } @Test public void testExchangeV4ToV5() throws Exception { try (final FlightServer server = startServer(optionV5); final FlightClient client = connect(server); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); final FlightClient.ExchangeReaderWriter stream = client.doExchange(FlightDescriptor.command(new byte[0]))) { stream.getWriter().start(root, null, optionV4); generateData(root); stream.getWriter().putNext(); stream.getWriter().completed(); assertTrue(stream.getReader().next()); assertEquals(optionV5.metadataVersion, stream.getReader().metadataVersion); validateRoot(stream.getReader().getRoot()); assertFalse(stream.getReader().next()); } } @Test public void testExchangeV5ToV4() throws Exception { try (final FlightServer server = startServer(optionV4); final FlightClient client = connect(server); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); final FlightClient.ExchangeReaderWriter stream = client.doExchange(FlightDescriptor.command(new byte[0]))) { stream.getWriter().start(root, null, optionV5); generateData(root); stream.getWriter().putNext(); stream.getWriter().completed(); assertTrue(stream.getReader().next()); assertEquals(optionV4.metadataVersion, stream.getReader().metadataVersion); validateRoot(stream.getReader().getRoot()); assertFalse(stream.getReader().next()); } } @Test public void testExchangeV4ToV4() throws Exception { try (final FlightServer server = startServer(optionV4); final FlightClient client = connect(server); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); final FlightClient.ExchangeReaderWriter stream = client.doExchange(FlightDescriptor.command(new byte[0]))) { stream.getWriter().start(root, null, optionV4); generateData(root); stream.getWriter().putNext(); stream.getWriter().completed(); assertTrue(stream.getReader().next()); assertEquals(optionV4.metadataVersion, stream.getReader().metadataVersion); validateRoot(stream.getReader().getRoot()); assertFalse(stream.getReader().next()); } } private static void generateData(VectorSchemaRoot root) { assertEquals(schema, root.getSchema()); final IntVector vector = (IntVector) root.getVector("foo"); vector.setSafe(0, 0); vector.setSafe(1, 1); vector.setSafe(2, 4); root.setRowCount(3); } private static void validateRoot(VectorSchemaRoot root) { assertEquals(schema, root.getSchema()); assertEquals(3, root.getRowCount()); final IntVector vector = (IntVector) root.getVector("foo"); assertEquals(0, vector.get(0)); assertEquals(1, vector.get(1)); assertEquals(4, vector.get(2)); } FlightServer startServer(IpcOption option) throws Exception { Location location = Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, 0); VersionFlightProducer producer = new VersionFlightProducer(allocator, option); final FlightServer server = FlightServer.builder(allocator, location, producer).build(); server.start(); return server; } FlightClient connect(FlightServer server) { Location location = Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, server.getPort()); return FlightClient.builder(allocator, location).build(); } static final class VersionFlightProducer extends NoOpFlightProducer { private final BufferAllocator allocator; private final IpcOption option; VersionFlightProducer(BufferAllocator allocator, IpcOption option) { this.allocator = allocator; this.option = option; } @Override public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) { return new FlightInfo(schema, descriptor, Collections.emptyList(), -1, -1, option); } @Override public SchemaResult getSchema(CallContext context, FlightDescriptor descriptor) { return new SchemaResult(schema, option); } @Override public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { if (Arrays.equals("union".getBytes(StandardCharsets.UTF_8), ticket.getBytes())) { try (final VectorSchemaRoot root = VectorSchemaRoot.create(unionSchema, allocator)) { listener.start(root, null, option); } catch (IllegalArgumentException e) { listener.error( CallStatus.INTERNAL .withCause(e) .withDescription(e.getMessage()) .toRuntimeException()); return; } listener.error( CallStatus.INTERNAL .withDescription("Expected exception not raised") .toRuntimeException()); return; } try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { listener.start(root, null, option); generateData(root); listener.putNext(); listener.completed(); } } @Override public Runnable acceptPut( CallContext context, FlightStream flightStream, StreamListener ackStream) { return () -> { try { assertTrue(flightStream.next()); assertEquals(option.metadataVersion, flightStream.metadataVersion); validateRoot(flightStream.getRoot()); } catch (AssertionError err) { // gRPC doesn't propagate stack traces across the wire. err.printStackTrace(); ackStream.onError( CallStatus.INVALID_ARGUMENT .withCause(err) .withDescription("Server assertion failed: " + err) .toRuntimeException()); return; } catch (RuntimeException err) { err.printStackTrace(); ackStream.onError( CallStatus.INTERNAL .withCause(err) .withDescription("Server assertion failed: " + err) .toRuntimeException()); return; } ackStream.onCompleted(); }; } @Override public void doExchange(CallContext context, FlightStream reader, ServerStreamListener writer) { try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { try { assertTrue(reader.next()); validateRoot(reader.getRoot()); assertFalse(reader.next()); } catch (AssertionError err) { // gRPC doesn't propagate stack traces across the wire. err.printStackTrace(); writer.error( CallStatus.INVALID_ARGUMENT .withCause(err) .withDescription("Server assertion failed: " + err) .toRuntimeException()); return; } catch (RuntimeException err) { err.printStackTrace(); writer.error( CallStatus.INTERNAL .withCause(err) .withDescription("Server assertion failed: " + err) .toRuntimeException()); return; } writer.start(root, null, option); generateData(root); writer.putNext(); writer.completed(); } } } } TestServerMiddleware.java000066400000000000000000000317171515451615200345650ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; import static org.apache.arrow.flight.Location.forGrpcInsecure; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.fail; import java.io.IOException; import java.util.Collections; import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.function.BiConsumer; import org.apache.arrow.flight.FlightClient.ClientStreamListener; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.api.Test; public class TestServerMiddleware { /** Make sure errors in DoPut are intercepted. */ @Test public void doPutErrors() { test( new ErrorProducer(new RuntimeException("test")), (allocator, client) -> { final FlightDescriptor descriptor = FlightDescriptor.path("test"); try (final VectorSchemaRoot root = VectorSchemaRoot.create(new Schema(Collections.emptyList()), allocator)) { final ClientStreamListener listener = client.startPut(descriptor, root, new SyncPutListener()); listener.completed(); FlightTestUtil.assertCode(FlightStatusCode.INTERNAL, listener::getResult); } }, (recorder) -> { final CallStatus status = recorder.statusFuture.get(); assertNotNull(status); assertNotNull(status.cause()); assertEquals(FlightStatusCode.INTERNAL, status.code()); }); // Check the status after server shutdown (to make sure gRPC finishes pending calls on the // server side) } /** Make sure custom error codes in DoPut are intercepted. */ @Test public void doPutCustomCode() { test( new ErrorProducer( CallStatus.UNAVAILABLE.withDescription("description").toRuntimeException()), (allocator, client) -> { final FlightDescriptor descriptor = FlightDescriptor.path("test"); try (final VectorSchemaRoot root = VectorSchemaRoot.create(new Schema(Collections.emptyList()), allocator)) { final ClientStreamListener listener = client.startPut(descriptor, root, new SyncPutListener()); listener.completed(); FlightTestUtil.assertCode(FlightStatusCode.UNAVAILABLE, listener::getResult); } }, (recorder) -> { final CallStatus status = recorder.statusFuture.get(); assertNotNull(status); assertNull(status.cause()); assertEquals(FlightStatusCode.UNAVAILABLE, status.code()); assertEquals("description", status.description()); }); } /** Make sure uncaught exceptions in DoPut are intercepted. */ @Test public void doPutUncaught() { test( new ServerErrorProducer(new RuntimeException("test")), (allocator, client) -> { final FlightDescriptor descriptor = FlightDescriptor.path("test"); try (final VectorSchemaRoot root = VectorSchemaRoot.create(new Schema(Collections.emptyList()), allocator)) { final ClientStreamListener listener = client.startPut(descriptor, root, new SyncPutListener()); listener.completed(); listener.getResult(); } }, (recorder) -> { final CallStatus status = recorder.statusFuture.get(); final Throwable err = recorder.errFuture.get(); assertNotNull(status); assertEquals(FlightStatusCode.OK, status.code()); assertNull(status.cause()); assertNotNull(err); assertEquals("test", err.getMessage()); }); } @Test public void listFlightsUncaught() { test( new ServerErrorProducer(new RuntimeException("test")), (allocator, client) -> client.listFlights(new Criteria(new byte[0])).forEach((action) -> {}), (recorder) -> { final CallStatus status = recorder.statusFuture.get(); final Throwable err = recorder.errFuture.get(); assertNotNull(status); assertEquals(FlightStatusCode.OK, status.code()); assertNull(status.cause()); assertNotNull(err); assertEquals("test", err.getMessage()); }); } @Test public void doActionUncaught() { test( new ServerErrorProducer(new RuntimeException("test")), (allocator, client) -> client.doAction(new Action("test")).forEachRemaining(result -> {}), (recorder) -> { final CallStatus status = recorder.statusFuture.get(); final Throwable err = recorder.errFuture.get(); assertNotNull(status); assertEquals(FlightStatusCode.OK, status.code()); assertNull(status.cause()); assertNotNull(err); assertEquals("test", err.getMessage()); }); } @Test public void listActionsUncaught() { test( new ServerErrorProducer(new RuntimeException("test")), (allocator, client) -> client.listActions().forEach(result -> {}), (recorder) -> { final CallStatus status = recorder.statusFuture.get(); final Throwable err = recorder.errFuture.get(); assertNotNull(status); assertEquals(FlightStatusCode.OK, status.code()); assertNull(status.cause()); assertNotNull(err); assertEquals("test", err.getMessage()); }); } @Test public void getFlightInfoUncaught() { test( new ServerErrorProducer(new RuntimeException("test")), (allocator, client) -> { FlightTestUtil.assertCode( FlightStatusCode.INTERNAL, () -> client.getInfo(FlightDescriptor.path("test"))); }, (recorder) -> { final CallStatus status = recorder.statusFuture.get(); assertNotNull(status); assertEquals(FlightStatusCode.INTERNAL, status.code()); assertNotNull(status.cause()); assertEquals(new RuntimeException("test").getMessage(), status.cause().getMessage()); }); } @Test public void doGetUncaught() { test( new ServerErrorProducer(new RuntimeException("test")), (allocator, client) -> { try (final FlightStream stream = client.getStream(new Ticket(new byte[0]))) { while (stream.next()) {} } catch (Exception e) { fail(e.toString()); } }, (recorder) -> { final CallStatus status = recorder.statusFuture.get(); final Throwable err = recorder.errFuture.get(); assertNotNull(status); assertEquals(FlightStatusCode.OK, status.code()); assertNull(status.cause()); assertNotNull(err); assertEquals("test", err.getMessage()); }); } /** A middleware that records the last error on any call. */ static class ErrorRecorder implements FlightServerMiddleware { CompletableFuture statusFuture = new CompletableFuture<>(); CompletableFuture errFuture = new CompletableFuture<>(); @Override public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) {} @Override public void onCallCompleted(CallStatus status) { statusFuture.complete(status); } @Override public void onCallErrored(Throwable err) { errFuture.complete(err); } static class Factory implements FlightServerMiddleware.Factory { ErrorRecorder instance = new ErrorRecorder(); @Override public ErrorRecorder onCallStarted( CallInfo info, CallHeaders incomingHeaders, RequestContext context) { return instance; } } } /** A producer that throws the given exception on a call. */ static class ErrorProducer extends NoOpFlightProducer { final RuntimeException error; ErrorProducer(RuntimeException t) { error = t; } @Override public Runnable acceptPut( CallContext context, FlightStream flightStream, StreamListener ackStream) { return () -> { // Drain queue to avoid FlightStream#close cancelling the call while (flightStream.next()) {} throw error; }; } } /** * A producer that throws the given exception on a call, but only after sending a success to the * client. */ static class ServerErrorProducer extends NoOpFlightProducer { final RuntimeException error; ServerErrorProducer(RuntimeException t) { error = t; } @Override public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); final VectorSchemaRoot root = VectorSchemaRoot.create(new Schema(Collections.emptyList()), allocator)) { listener.start(root); listener.completed(); } throw error; } @Override public void listFlights( CallContext context, Criteria criteria, StreamListener listener) { listener.onCompleted(); throw error; } @Override public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) { throw error; } @Override public Runnable acceptPut( CallContext context, FlightStream flightStream, StreamListener ackStream) { return () -> { while (flightStream.next()) {} ackStream.onCompleted(); throw error; }; } @Override public void doAction(CallContext context, Action action, StreamListener listener) { listener.onCompleted(); throw error; } @Override public void listActions(CallContext context, StreamListener listener) { listener.onCompleted(); throw error; } } static class ServerMiddlewarePair { final FlightServerMiddleware.Key key; final FlightServerMiddleware.Factory factory; ServerMiddlewarePair( FlightServerMiddleware.Key key, FlightServerMiddleware.Factory factory) { this.key = key; this.factory = factory; } } /** * Spin up a service with the given middleware and producer. * * @param producer The Flight producer to use. * @param middleware A list of middleware to register. * @param body A function to run as the body of the test. * @param The middleware type. */ static void test( FlightProducer producer, List> middleware, BiConsumer body) { try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) { final FlightServer.Builder builder = FlightServer.builder(allocator, forGrpcInsecure(LOCALHOST, 0), producer); middleware.forEach(pair -> builder.middleware(pair.key, pair.factory)); final FlightServer server = builder.build().start(); try (final FlightServer ignored = server; final FlightClient client = FlightClient.builder(allocator, server.getLocation()).build()) { body.accept(allocator, client); } } catch (InterruptedException | IOException e) { throw new RuntimeException(e); } } static void test( FlightProducer producer, BiConsumer body, ErrorConsumer verify) { final ErrorRecorder.Factory factory = new ErrorRecorder.Factory(); final List> middleware = Collections.singletonList( new ServerMiddlewarePair<>(FlightServerMiddleware.Key.of("m"), factory)); test( producer, middleware, (allocator, client) -> { body.accept(allocator, client); try { verify.accept(factory.instance); } catch (Exception e) { throw new RuntimeException(e); } }); } @FunctionalInterface interface ErrorConsumer { void accept(T obj) throws Exception; } } TestServerOptions.java000066400000000000000000000214201515451615200341310ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; import static org.apache.arrow.flight.Location.forGrpcInsecure; 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.assertTrue; import io.grpc.Channel; import io.grpc.MethodDescriptor; import io.grpc.ServerServiceDefinition; import io.grpc.health.v1.HealthCheckRequest; import io.grpc.health.v1.HealthCheckResponse; import io.grpc.health.v1.HealthGrpc; import io.grpc.netty.NettyChannelBuilder; import io.grpc.netty.NettyServerBuilder; import io.grpc.protobuf.services.HealthStatusManager; import java.io.File; import java.util.HashMap; import java.util.Map; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Consumer; import org.apache.arrow.flight.TestBasicOperation.Producer; import org.apache.arrow.flight.auth.ServerAuthHandler; import org.apache.arrow.flight.impl.FlightServiceGrpc; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.junit.jupiter.api.Assumptions; import org.junit.jupiter.api.Test; public class TestServerOptions { @Test public void builderConsumer() throws Exception { final AtomicBoolean consumerCalled = new AtomicBoolean(); final Consumer consumer = (builder) -> consumerCalled.set(true); try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE); Producer producer = new Producer(a); FlightServer s = FlightServer.builder(a, forGrpcInsecure(LOCALHOST, 0), producer) .transportHint("grpc.builderConsumer", consumer) .build() .start()) { assertTrue(consumerCalled.get()); } } /** * Make sure that if Flight supplies a default executor to gRPC, then it is closed along with the * server. */ @Test public void defaultExecutorClosed() throws Exception { final ExecutorService executor; try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE); FlightServer server = FlightServer.builder(a, forGrpcInsecure(LOCALHOST, 0), new NoOpFlightProducer()) .build() .start()) { assertNotNull(server.grpcExecutor); executor = server.grpcExecutor; } assertTrue(executor.isShutdown()); } /** Make sure that if the user provides an executor to gRPC, then Flight does not close it. */ @Test public void suppliedExecutorNotClosed() throws Exception { final ExecutorService executor = Executors.newSingleThreadExecutor(); try { try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE); FlightServer server = FlightServer.builder(a, forGrpcInsecure(LOCALHOST, 0), new NoOpFlightProducer()) .executor(executor) .build() .start()) { assertNull(server.grpcExecutor); } assertFalse(executor.isShutdown()); } finally { executor.shutdown(); } } @Test public void domainSocket() throws Exception { Assumptions.assumeTrue( FlightTestUtil.isNativeTransportAvailable(), "We have a native transport available"); final File domainSocket = File.createTempFile("flight-unit-test-", ".sock"); assertTrue(domainSocket.delete()); // Domain socket paths have a platform-dependent limit. Set a conservative limit and skip the // test if the temporary // file name is too long. (We do not assume a particular platform-dependent temporary directory // path.) Assumptions.assumeTrue( domainSocket.getAbsolutePath().length() < 100, "The domain socket path is not too long"); final Location location = Location.forGrpcDomainSocket(domainSocket.getAbsolutePath()); try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE); Producer producer = new Producer(a); FlightServer s = FlightServer.builder(a, location, producer).build().start(); ) { try (FlightClient c = FlightClient.builder(a, location).build()) { try (FlightStream stream = c.getStream(new Ticket(new byte[0]))) { VectorSchemaRoot root = stream.getRoot(); IntVector iv = (IntVector) root.getVector("c1"); int value = 0; while (stream.next()) { for (int i = 0; i < root.getRowCount(); i++) { assertEquals(value, iv.get(i)); value++; } } } } } } @Test public void checkReflectionMetadata() { // This metadata is needed for gRPC reflection to work. final ExecutorService executorService = Executors.newSingleThreadExecutor(); try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) { final FlightBindingService service = new FlightBindingService( allocator, new NoOpFlightProducer(), ServerAuthHandler.NO_OP, executorService); final ServerServiceDefinition definition = service.bindService(); assertEquals( FlightServiceGrpc.getServiceDescriptor().getSchemaDescriptor(), definition.getServiceDescriptor().getSchemaDescriptor()); final Map> definedMethods = new HashMap<>(); final Map> serviceMethods = new HashMap<>(); // Make sure that the reflection metadata object is identical across all the places where it's // accessible definition .getMethods() .forEach( method -> definedMethods.put( method.getMethodDescriptor().getFullMethodName(), method.getMethodDescriptor())); definition .getServiceDescriptor() .getMethods() .forEach(method -> serviceMethods.put(method.getFullMethodName(), method)); for (final MethodDescriptor descriptor : FlightServiceGrpc.getServiceDescriptor().getMethods()) { final String methodName = descriptor.getFullMethodName(); assertTrue( definedMethods.containsKey(methodName), "Method is missing from ServerServiceDefinition: " + methodName); assertTrue( definedMethods.containsKey(methodName), "Method is missing from ServiceDescriptor: " + methodName); assertEquals( descriptor.getSchemaDescriptor(), definedMethods.get(methodName).getSchemaDescriptor()); assertEquals( descriptor.getSchemaDescriptor(), serviceMethods.get(methodName).getSchemaDescriptor()); } } finally { executorService.shutdown(); } } /* * This is an extension of builderConsumer test. * Test that Flight interceptors don't break other registered services */ @Test public void addHealthCheckService() throws Exception { 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()); assertEquals(response.getStatus(), HealthCheckResponse.ServingStatus.SERVING); } } } apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/TestTls.java000066400000000000000000000125511515451615200321350ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; import static org.apache.arrow.flight.Location.forGrpcInsecure; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import java.io.FileInputStream; import java.io.IOException; import java.io.InputStream; import java.nio.charset.StandardCharsets; import java.util.Iterator; import java.util.function.Consumer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.Test; /** Tests for TLS in Flight. */ public class TestTls { /** Test a basic request over TLS. */ @Test public void connectTls() { test( (builder) -> { try (final InputStream roots = new FileInputStream(FlightTestUtil.exampleTlsRootCert().toFile()); final FlightClient client = builder.trustedCertificates(roots).build()) { final Iterator responses = client.doAction(new Action("hello-world")); final byte[] response = responses.next().getBody(); assertEquals("Hello, world!", new String(response, StandardCharsets.UTF_8)); assertFalse(responses.hasNext()); } catch (InterruptedException | IOException e) { throw new RuntimeException(e); } }); } /** Make sure that connections are rejected when the root certificate isn't trusted. */ @Test public void rejectInvalidCert() { test( (builder) -> { try (final FlightClient client = builder.build()) { final Iterator responses = client.doAction(new Action("hello-world")); FlightTestUtil.assertCode( FlightStatusCode.UNAVAILABLE, () -> responses.next().getBody()); } catch (InterruptedException e) { throw new RuntimeException(e); } }); } /** Make sure that connections are rejected when the hostname doesn't match. */ @Test public void rejectHostname() { test( (builder) -> { try (final InputStream roots = new FileInputStream(FlightTestUtil.exampleTlsRootCert().toFile()); final FlightClient client = builder.trustedCertificates(roots).overrideHostname("fakehostname").build()) { final Iterator responses = client.doAction(new Action("hello-world")); FlightTestUtil.assertCode( FlightStatusCode.UNAVAILABLE, () -> responses.next().getBody()); } catch (InterruptedException | IOException e) { throw new RuntimeException(e); } }); } /** Test a basic request over TLS. */ @Test public void connectTlsDisableServerVerification() { test( (builder) -> { try (final FlightClient client = builder.verifyServer(false).build()) { final Iterator responses = client.doAction(new Action("hello-world")); final byte[] response = responses.next().getBody(); assertEquals("Hello, world!", new String(response, StandardCharsets.UTF_8)); assertFalse(responses.hasNext()); } catch (InterruptedException e) { throw new RuntimeException(e); } }); } void test(Consumer testFn) { final FlightTestUtil.CertKeyPair certKey = FlightTestUtil.exampleTlsCerts().get(0); try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE); Producer producer = new Producer(); FlightServer s = FlightServer.builder(a, forGrpcInsecure(LOCALHOST, 0), producer) .useTls(certKey.cert, certKey.key) .build() .start()) { final FlightClient.Builder builder = FlightClient.builder(a, Location.forGrpcTls(FlightTestUtil.LOCALHOST, s.getPort())); testFn.accept(builder); } catch (InterruptedException | IOException e) { throw new RuntimeException(e); } } static class Producer extends NoOpFlightProducer implements AutoCloseable { @Override public void doAction(CallContext context, Action action, StreamListener listener) { if (action.getType().equals("hello-world")) { listener.onNext(new Result("Hello, world!".getBytes(StandardCharsets.UTF_8))); listener.onCompleted(); return; } listener.onError( CallStatus.UNIMPLEMENTED .withDescription("Invalid action " + action.getType()) .toRuntimeException()); } @Override public void close() {} } } apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/000077500000000000000000000000001515451615200306255ustar00rootroot00000000000000TestBasicAuth.java000066400000000000000000000152511515451615200341200ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; import static org.apache.arrow.flight.Location.forGrpcInsecure; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.fail; import com.google.common.collect.ImmutableList; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.Optional; import org.apache.arrow.flight.Criteria; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightStatusCode; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.FlightTestUtil; import org.apache.arrow.flight.NoOpFlightProducer; import org.apache.arrow.flight.Ticket; 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.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; 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; public class TestBasicAuth { private static final String USERNAME = "flight"; private static final String PASSWORD = "woohoo"; private static final byte[] VALID_TOKEN = "my_token".getBytes(StandardCharsets.UTF_8); private FlightClient client; private static FlightServer server; private static BufferAllocator allocator; @Test public void validAuth() { client.authenticateBasic(USERNAME, PASSWORD); assertEquals(0, ImmutableList.copyOf(client.listFlights(Criteria.ALL)).size()); } @Test public void asyncCall() throws Exception { client.authenticateBasic(USERNAME, PASSWORD); client.listFlights(Criteria.ALL); try (final FlightStream s = client.getStream(new Ticket(new byte[1]))) { while (s.next()) { assertEquals(4095, s.getRoot().getRowCount()); } } } @Test public void invalidAuth() { FlightTestUtil.assertCode( FlightStatusCode.UNAUTHENTICATED, () -> { client.authenticateBasic(USERNAME, "WRONG"); }); FlightTestUtil.assertCode( FlightStatusCode.UNAUTHENTICATED, () -> { client.listFlights(Criteria.ALL).forEach(action -> fail()); }); } @Test public void didntAuth() { FlightTestUtil.assertCode( FlightStatusCode.UNAUTHENTICATED, () -> { client.listFlights(Criteria.ALL).forEach(action -> fail()); }); } @BeforeEach public void testSetup() throws IOException { client = FlightClient.builder(allocator, server.getLocation()).build(); } @BeforeAll public static void setup() throws IOException { allocator = new RootAllocator(Long.MAX_VALUE); final BasicServerAuthHandler.BasicAuthValidator validator = new BasicServerAuthHandler.BasicAuthValidator() { @Override public Optional isValid(byte[] token) { if (Arrays.equals(token, VALID_TOKEN)) { return Optional.of(USERNAME); } return Optional.empty(); } @Override public byte[] getToken(String username, String password) { if (USERNAME.equals(username) && PASSWORD.equals(password)) { return VALID_TOKEN; } else { throw new IllegalArgumentException("invalid credentials"); } } }; server = FlightServer.builder( allocator, forGrpcInsecure(LOCALHOST, 0), new NoOpFlightProducer() { @Override public void listFlights( CallContext context, Criteria criteria, StreamListener listener) { if (!context.peerIdentity().equals(USERNAME)) { listener.onError(new IllegalArgumentException("Invalid username")); return; } listener.onCompleted(); } @Override public void getStream( CallContext context, Ticket ticket, ServerStreamListener listener) { if (!context.peerIdentity().equals(USERNAME)) { listener.error(new IllegalArgumentException("Invalid username")); return; } final Schema pojoSchema = new Schema( ImmutableList.of( Field.nullable("a", Types.MinorType.BIGINT.getType()))); try (VectorSchemaRoot root = VectorSchemaRoot.create(pojoSchema, allocator)) { listener.start(root); root.allocateNew(); root.setRowCount(4095); listener.putNext(); listener.completed(); } } }) .authHandler(new BasicServerAuthHandler(validator)) .build() .start(); } @AfterEach public void tearDown() throws Exception { AutoCloseables.close(client); } @AfterAll 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); } } apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/auth2/000077500000000000000000000000001515451615200307075ustar00rootroot00000000000000TestBasicAuth2.java000066400000000000000000000214161515451615200342640ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/auth2/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.auth2; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; import static org.apache.arrow.flight.Location.forGrpcInsecure; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; import com.google.common.base.Strings; import com.google.common.collect.ImmutableList; import java.io.IOException; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.Criteria; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightStatusCode; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.FlightTestUtil; import org.apache.arrow.flight.NoOpFlightProducer; import org.apache.arrow.flight.Ticket; import org.apache.arrow.flight.grpc.CredentialCallOption; 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.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; public class TestBasicAuth2 { private static final String USERNAME_1 = "flight1"; private static final String USERNAME_2 = "flight2"; private static final String NO_USERNAME = ""; private static final String PASSWORD_1 = "woohoo1"; private static final String PASSWORD_2 = "woohoo2"; private static BufferAllocator allocator; private static FlightServer server; private static FlightClient client; private static FlightClient client2; @BeforeAll public static void setup() throws Exception { allocator = new RootAllocator(Long.MAX_VALUE); startServerAndClient(); } private static FlightProducer getFlightProducer() { return new NoOpFlightProducer() { @Override public void listFlights( CallContext context, Criteria criteria, StreamListener listener) { if (!context.peerIdentity().equals(USERNAME_1) && !context.peerIdentity().equals(USERNAME_2)) { listener.onError(new IllegalArgumentException("Invalid username")); return; } listener.onCompleted(); } @Override public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { if (!context.peerIdentity().equals(USERNAME_1) && !context.peerIdentity().equals(USERNAME_2)) { listener.error(new IllegalArgumentException("Invalid username")); return; } final Schema pojoSchema = new Schema(ImmutableList.of(Field.nullable("a", Types.MinorType.BIGINT.getType()))); try (VectorSchemaRoot root = VectorSchemaRoot.create(pojoSchema, allocator)) { listener.start(root); root.allocateNew(); root.setRowCount(4095); listener.putNext(); listener.completed(); } } }; } private static void startServerAndClient() throws IOException { final FlightProducer flightProducer = getFlightProducer(); server = FlightServer.builder(allocator, forGrpcInsecure(LOCALHOST, 0), flightProducer) .headerAuthenticator( new GeneratedBearerTokenAuthenticator( new BasicCallHeaderAuthenticator(TestBasicAuth2::validate))) .build() .start(); client = FlightClient.builder(allocator, server.getLocation()).build(); } @AfterAll public static void shutdown() throws Exception { AutoCloseables.close(client, client2, server); client = null; client2 = null; server = null; allocator.getChildAllocators().forEach(BufferAllocator::close); AutoCloseables.close(allocator); allocator = null; } private void startClient2() throws IOException { client2 = FlightClient.builder(allocator, server.getLocation()).build(); } private static CallHeaderAuthenticator.AuthResult validate(String username, String password) { if (Strings.isNullOrEmpty(username)) { throw CallStatus.UNAUTHENTICATED .withDescription("Credentials not supplied.") .toRuntimeException(); } final String identity; if (USERNAME_1.equals(username) && PASSWORD_1.equals(password)) { identity = USERNAME_1; } else if (USERNAME_2.equals(username) && PASSWORD_2.equals(password)) { identity = USERNAME_2; } else { throw CallStatus.UNAUTHENTICATED .withDescription("Username or password is invalid.") .toRuntimeException(); } return () -> identity; } @Test public void validAuthWithBearerAuthServer() throws IOException { testValidAuth(client); } @Test public void validAuthWithMultipleClientsWithSameCredentialsWithBearerAuthServer() throws IOException { startClient2(); testValidAuthWithMultipleClientsWithSameCredentials(client, client2); } @Test public void validAuthWithMultipleClientsWithDifferentCredentialsWithBearerAuthServer() throws IOException { startClient2(); testValidAuthWithMultipleClientsWithDifferentCredentials(client, client2); } @Test public void asyncCall() throws Exception { final CredentialCallOption bearerToken = client.authenticateBasicToken(USERNAME_1, PASSWORD_1).get(); client.listFlights(Criteria.ALL, bearerToken); try (final FlightStream s = client.getStream(new Ticket(new byte[1]), bearerToken)) { while (s.next()) { assertEquals(4095, s.getRoot().getRowCount()); } } } @Test public void invalidAuthWithBearerAuthServer() throws IOException { testInvalidAuth(client); } @Test public void didntAuthWithBearerAuthServer() throws IOException { didntAuth(client); } private void testValidAuth(FlightClient client) { final CredentialCallOption bearerToken = client.authenticateBasicToken(USERNAME_1, PASSWORD_1).get(); assertTrue(ImmutableList.copyOf(client.listFlights(Criteria.ALL, bearerToken)).isEmpty()); } private void testValidAuthWithMultipleClientsWithSameCredentials( FlightClient client1, FlightClient client2) { final CredentialCallOption bearerToken1 = client1.authenticateBasicToken(USERNAME_1, PASSWORD_1).get(); final CredentialCallOption bearerToken2 = client2.authenticateBasicToken(USERNAME_1, PASSWORD_1).get(); assertTrue(ImmutableList.copyOf(client1.listFlights(Criteria.ALL, bearerToken1)).isEmpty()); assertTrue(ImmutableList.copyOf(client2.listFlights(Criteria.ALL, bearerToken2)).isEmpty()); } private void testValidAuthWithMultipleClientsWithDifferentCredentials( FlightClient client1, FlightClient client2) { final CredentialCallOption bearerToken1 = client1.authenticateBasicToken(USERNAME_1, PASSWORD_1).get(); final CredentialCallOption bearerToken2 = client2.authenticateBasicToken(USERNAME_2, PASSWORD_2).get(); assertTrue(ImmutableList.copyOf(client1.listFlights(Criteria.ALL, bearerToken1)).isEmpty()); assertTrue(ImmutableList.copyOf(client2.listFlights(Criteria.ALL, bearerToken2)).isEmpty()); } private void testInvalidAuth(FlightClient client) { FlightTestUtil.assertCode( FlightStatusCode.UNAUTHENTICATED, () -> client.authenticateBasicToken(USERNAME_1, "WRONG")); FlightTestUtil.assertCode( FlightStatusCode.UNAUTHENTICATED, () -> client.authenticateBasicToken(NO_USERNAME, PASSWORD_1)); FlightTestUtil.assertCode( FlightStatusCode.UNAUTHENTICATED, () -> client.listFlights(Criteria.ALL).forEach(action -> fail())); } private void didntAuth(FlightClient client) { FlightTestUtil.assertCode( FlightStatusCode.UNAUTHENTICATED, () -> client.listFlights(Criteria.ALL).forEach(action -> fail())); } } apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/client/000077500000000000000000000000001515451615200311425ustar00rootroot00000000000000CustomHeaderTest.java000066400000000000000000000161571515451615200351630ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/client/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.client; import static org.junit.jupiter.api.Assertions.assertEquals; import com.google.common.collect.ImmutableMap; import java.util.HashMap; import java.util.Map; import org.apache.arrow.flight.Action; import org.apache.arrow.flight.CallHeaders; import org.apache.arrow.flight.CallInfo; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.Criteria; import org.apache.arrow.flight.FlightCallHeaders; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightClient.ClientStreamListener; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightMethod; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightServerMiddleware; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.FlightTestUtil; import org.apache.arrow.flight.HeaderCallOption; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.NoOpFlightProducer; import org.apache.arrow.flight.RequestContext; import org.apache.arrow.flight.SyncPutListener; import org.apache.arrow.flight.Ticket; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Tests to ensure custom headers are passed along to the server for each command. */ public class CustomHeaderTest { FlightServer server; FlightClient client; BufferAllocator allocator; TestCustomHeaderMiddleware.Factory headersMiddleware; HeaderCallOption headers; Map testHeaders = ImmutableMap.of( "foo", "bar", "bar", "foo", "answer", "42"); @BeforeEach public void setUp() throws Exception { allocator = new RootAllocator(Integer.MAX_VALUE); headersMiddleware = new TestCustomHeaderMiddleware.Factory(); FlightCallHeaders callHeaders = new FlightCallHeaders(); for (Map.Entry entry : testHeaders.entrySet()) { callHeaders.insert(entry.getKey(), entry.getValue()); } headers = new HeaderCallOption(callHeaders); server = FlightServer.builder( allocator, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, /*port*/ 0), new NoOpFlightProducer()) .middleware(FlightServerMiddleware.Key.of("customHeader"), headersMiddleware) .build(); server.start(); client = FlightClient.builder(allocator, server.getLocation()).build(); } @AfterEach public void tearDown() throws Exception { allocator.getChildAllocators().forEach(BufferAllocator::close); AutoCloseables.close(allocator, server, client); } @Test public void testHandshake() { try { client.handshake(headers); } catch (Exception ignored) { } assertHeadersMatch(FlightMethod.HANDSHAKE); } @Test public void testGetSchema() { try { client.getSchema(FlightDescriptor.command(new byte[0]), headers); } catch (Exception ignored) { } assertHeadersMatch(FlightMethod.GET_SCHEMA); } @Test public void testGetFlightInfo() { try { client.getInfo(FlightDescriptor.command(new byte[0]), headers); } catch (Exception ignored) { } assertHeadersMatch(FlightMethod.GET_FLIGHT_INFO); } @Test public void testListActions() { try { client.listActions(headers).iterator().next(); } catch (Exception ignored) { } assertHeadersMatch(FlightMethod.LIST_ACTIONS); } @Test public void testListFlights() { try { client.listFlights(new Criteria(new byte[] {1}), headers).iterator().next(); } catch (Exception ignored) { } assertHeadersMatch(FlightMethod.LIST_FLIGHTS); } @Test public void testDoAction() { try { client.doAction(new Action("test"), headers).next(); } catch (Exception ignored) { } assertHeadersMatch(FlightMethod.DO_ACTION); } @Test public void testStartPut() { try { final ClientStreamListener listener = client.startPut(FlightDescriptor.command(new byte[0]), new SyncPutListener(), headers); listener.getResult(); } catch (Exception ignored) { } assertHeadersMatch(FlightMethod.DO_PUT); } @Test public void testGetStream() { try (final FlightStream stream = client.getStream(new Ticket(new byte[0]), headers)) { stream.next(); } catch (Exception ignored) { } assertHeadersMatch(FlightMethod.DO_GET); } @Test public void testDoExchange() { try (final FlightClient.ExchangeReaderWriter stream = client.doExchange(FlightDescriptor.command(new byte[0]), headers)) { stream.getReader().next(); } catch (Exception ignored) { } assertHeadersMatch(FlightMethod.DO_EXCHANGE); } private void assertHeadersMatch(FlightMethod method) { for (Map.Entry entry : testHeaders.entrySet()) { assertEquals(entry.getValue(), headersMiddleware.getCustomHeader(method, entry.getKey())); } } /** A middleware used to test if customHeaders are being sent to the server properly. */ static class TestCustomHeaderMiddleware implements FlightServerMiddleware { public TestCustomHeaderMiddleware() {} @Override public void onBeforeSendingHeaders(CallHeaders callHeaders) {} @Override public void onCallCompleted(CallStatus callStatus) {} @Override public void onCallErrored(Throwable throwable) {} /** * A factory for the middleware that keeps track of the received headers and provides a way to * check those values for a given Flight Method. */ static class Factory implements FlightServerMiddleware.Factory { private final Map receivedCallHeaders = new HashMap<>(); @Override public TestCustomHeaderMiddleware onCallStarted( CallInfo callInfo, CallHeaders callHeaders, RequestContext requestContext) { receivedCallHeaders.put(callInfo.method(), callHeaders); return new TestCustomHeaderMiddleware(); } public String getCustomHeader(FlightMethod method, String key) { CallHeaders headers = receivedCallHeaders.get(method); if (headers == null) { return null; } return headers.get(key); } } } } TestCookieHandling.java000066400000000000000000000253231515451615200354510ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/client/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.client; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; import static org.apache.arrow.flight.Location.forGrpcInsecure; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.IOException; import org.apache.arrow.flight.CallHeaders; import org.apache.arrow.flight.CallInfo; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.Criteria; import org.apache.arrow.flight.ErrorFlightMetadata; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightMethod; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightServerMiddleware; import org.apache.arrow.flight.NoOpFlightProducer; import org.apache.arrow.flight.RequestContext; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; /** * Tests for correct handling of cookies from the FlightClient using {@link ClientCookieMiddleware}. */ public class TestCookieHandling { private static final String SET_COOKIE_HEADER = "Set-Cookie"; private static final String COOKIE_HEADER = "Cookie"; private BufferAllocator allocator; private FlightServer server; private FlightClient client; private ClientCookieMiddlewareTestFactory testFactory = new ClientCookieMiddlewareTestFactory(); private ClientCookieMiddleware cookieMiddleware = new ClientCookieMiddleware(testFactory); @BeforeEach public void setup() throws Exception { allocator = new RootAllocator(Long.MAX_VALUE); startServerAndClient(); } @AfterEach public void cleanup() throws Exception { testFactory = new ClientCookieMiddlewareTestFactory(); cookieMiddleware = testFactory.onCallStarted(new CallInfo(FlightMethod.DO_ACTION)); AutoCloseables.close(client, server, allocator); client = null; server = null; allocator = null; } @Test public void basicCookie() { CallHeaders headersToSend = new ErrorFlightMetadata(); headersToSend.insert(SET_COOKIE_HEADER, "k=v"); cookieMiddleware = testFactory.onCallStarted(new CallInfo(FlightMethod.DO_ACTION)); cookieMiddleware.onHeadersReceived(headersToSend); assertEquals("k=v", cookieMiddleware.getValidCookiesAsString()); } @Test public void cookieStaysAfterMultipleRequests() { CallHeaders headersToSend = new ErrorFlightMetadata(); headersToSend.insert(SET_COOKIE_HEADER, "k=v"); cookieMiddleware = testFactory.onCallStarted(new CallInfo(FlightMethod.DO_ACTION)); cookieMiddleware.onHeadersReceived(headersToSend); assertEquals("k=v", cookieMiddleware.getValidCookiesAsString()); headersToSend = new ErrorFlightMetadata(); cookieMiddleware = testFactory.onCallStarted(new CallInfo(FlightMethod.DO_ACTION)); cookieMiddleware.onHeadersReceived(headersToSend); assertEquals("k=v", cookieMiddleware.getValidCookiesAsString()); headersToSend = new ErrorFlightMetadata(); cookieMiddleware = testFactory.onCallStarted(new CallInfo(FlightMethod.DO_ACTION)); cookieMiddleware.onHeadersReceived(headersToSend); assertEquals("k=v", cookieMiddleware.getValidCookiesAsString()); } @Disabled @Test public void cookieAutoExpires() { CallHeaders headersToSend = new ErrorFlightMetadata(); headersToSend.insert(SET_COOKIE_HEADER, "k=v; Max-Age=2"); cookieMiddleware = testFactory.onCallStarted(new CallInfo(FlightMethod.DO_ACTION)); cookieMiddleware.onHeadersReceived(headersToSend); // Note: using max-age changes cookie version from 0->1, which quotes values. assertEquals("k=\"v\"", cookieMiddleware.getValidCookiesAsString()); headersToSend = new ErrorFlightMetadata(); cookieMiddleware = testFactory.onCallStarted(new CallInfo(FlightMethod.DO_ACTION)); cookieMiddleware.onHeadersReceived(headersToSend); assertEquals("k=\"v\"", cookieMiddleware.getValidCookiesAsString()); try { Thread.sleep(5000); } catch (InterruptedException ignored) { } // Verify that the k cookie was discarded because it expired. assertTrue(cookieMiddleware.getValidCookiesAsString().isEmpty()); } @Test public void cookieExplicitlyExpires() { CallHeaders headersToSend = new ErrorFlightMetadata(); headersToSend.insert(SET_COOKIE_HEADER, "k=v; Max-Age=2"); cookieMiddleware = testFactory.onCallStarted(new CallInfo(FlightMethod.DO_ACTION)); cookieMiddleware.onHeadersReceived(headersToSend); // Note: using max-age changes cookie version from 0->1, which quotes values. assertEquals("k=\"v\"", cookieMiddleware.getValidCookiesAsString()); // Note: The JDK treats Max-Age < 0 as not expired and treats 0 as expired. // This violates the RFC, which states that less than zero and zero should both be expired. headersToSend = new ErrorFlightMetadata(); headersToSend.insert(SET_COOKIE_HEADER, "k=v; Max-Age=0"); cookieMiddleware = testFactory.onCallStarted(new CallInfo(FlightMethod.DO_ACTION)); cookieMiddleware.onHeadersReceived(headersToSend); // Verify that the k cookie was discarded because the server told the client it is expired. assertTrue(cookieMiddleware.getValidCookiesAsString().isEmpty()); } @Disabled @Test public void cookieExplicitlyExpiresWithMaxAgeMinusOne() { CallHeaders headersToSend = new ErrorFlightMetadata(); headersToSend.insert(SET_COOKIE_HEADER, "k=v; Max-Age=2"); cookieMiddleware = testFactory.onCallStarted(new CallInfo(FlightMethod.DO_ACTION)); cookieMiddleware.onHeadersReceived(headersToSend); // Note: using max-age changes cookie version from 0->1, which quotes values. assertEquals("k=\"v\"", cookieMiddleware.getValidCookiesAsString()); headersToSend = new ErrorFlightMetadata(); // The Java HttpCookie class has a bug where it uses a -1 maxAge to indicate // a persistent cookie, when the RFC spec says this should mean the cookie expires immediately. headersToSend.insert(SET_COOKIE_HEADER, "k=v; Max-Age=-1"); cookieMiddleware = testFactory.onCallStarted(new CallInfo(FlightMethod.DO_ACTION)); cookieMiddleware.onHeadersReceived(headersToSend); // Verify that the k cookie was discarded because the server told the client it is expired. assertTrue(cookieMiddleware.getValidCookiesAsString().isEmpty()); } @Test public void changeCookieValue() { CallHeaders headersToSend = new ErrorFlightMetadata(); headersToSend.insert(SET_COOKIE_HEADER, "k=v"); cookieMiddleware.onHeadersReceived(headersToSend); assertEquals("k=v", cookieMiddleware.getValidCookiesAsString()); headersToSend = new ErrorFlightMetadata(); headersToSend.insert(SET_COOKIE_HEADER, "k=v2"); cookieMiddleware.onHeadersReceived(headersToSend); assertEquals("k=v2", cookieMiddleware.getValidCookiesAsString()); } @Test public void multipleCookiesWithSetCookie() { CallHeaders headersToSend = new ErrorFlightMetadata(); headersToSend.insert(SET_COOKIE_HEADER, "firstKey=firstVal"); headersToSend.insert(SET_COOKIE_HEADER, "secondKey=secondVal"); cookieMiddleware.onHeadersReceived(headersToSend); assertEquals( "firstKey=firstVal; secondKey=secondVal", cookieMiddleware.getValidCookiesAsString()); } @Test public void cookieStaysAfterMultipleRequestsEndToEnd() { client.handshake(); assertEquals("k=v", testFactory.clientCookieMiddleware.getValidCookiesAsString()); client.handshake(); assertEquals("k=v", testFactory.clientCookieMiddleware.getValidCookiesAsString()); client.listFlights(Criteria.ALL); assertEquals("k=v", testFactory.clientCookieMiddleware.getValidCookiesAsString()); } /** A server middleware component that injects SET_COOKIE_HEADER into the outgoing headers. */ static class SetCookieHeaderInjector implements FlightServerMiddleware { private final Factory factory; public SetCookieHeaderInjector(Factory factory) { this.factory = factory; } @Override public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) { if (!factory.receivedCookieHeader) { outgoingHeaders.insert(SET_COOKIE_HEADER, "k=v"); } } @Override public void onCallCompleted(CallStatus status) {} @Override public void onCallErrored(Throwable err) {} static class Factory implements FlightServerMiddleware.Factory { private boolean receivedCookieHeader = false; @Override public SetCookieHeaderInjector onCallStarted( CallInfo info, CallHeaders incomingHeaders, RequestContext context) { receivedCookieHeader = null != incomingHeaders.get(COOKIE_HEADER); return new SetCookieHeaderInjector(this); } } } public static class ClientCookieMiddlewareTestFactory extends ClientCookieMiddleware.Factory { private ClientCookieMiddleware clientCookieMiddleware; @Override public ClientCookieMiddleware onCallStarted(CallInfo info) { this.clientCookieMiddleware = new ClientCookieMiddleware(this); return this.clientCookieMiddleware; } } private void startServerAndClient() throws IOException { final FlightProducer flightProducer = new NoOpFlightProducer() { @Override public void listFlights( CallContext context, Criteria criteria, StreamListener listener) { listener.onCompleted(); } }; this.server = FlightServer.builder(allocator, forGrpcInsecure(LOCALHOST, 0), flightProducer) .middleware( FlightServerMiddleware.Key.of("test"), new SetCookieHeaderInjector.Factory()) .build() .start(); this.client = FlightClient.builder(allocator, server.getLocation()).intercept(testFactory).build(); } } apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/000077500000000000000000000000001515451615200306175ustar00rootroot00000000000000TestMetadataAdapter.java000066400000000000000000000026361515451615200352730ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } TestStatusUtils.java000066400000000000000000000060461515451615200345550ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import io.grpc.Metadata; import io.grpc.Status; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightStatusCode; import org.junit.jupiter.api.Test; public class TestStatusUtils { @Test public void testParseTrailers() { Status status = Status.CANCELLED; Metadata trailers = new Metadata(); // gRPC can have trailers with certain metadata keys beginning with ":", such as ":status". // See https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md trailers.put(StatusUtils.keyOfAscii(":status"), "502"); trailers.put(StatusUtils.keyOfAscii("date"), "Fri, 13 Sep 2015 11:23:58 GMT"); trailers.put(StatusUtils.keyOfAscii("content-type"), "text/html"); CallStatus callStatus = StatusUtils.fromGrpcStatusAndTrailers(status, trailers); assertEquals(FlightStatusCode.CANCELLED, callStatus.code()); assertTrue(callStatus.metadata().containsKey(":status")); assertEquals("502", callStatus.metadata().get(":status")); assertTrue(callStatus.metadata().containsKey("date")); assertEquals("Fri, 13 Sep 2015 11:23:58 GMT", callStatus.metadata().get("date")); assertTrue(callStatus.metadata().containsKey("content-type")); assertEquals("text/html", callStatus.metadata().get("content-type")); } @Test public void testGrpcResourceExhaustedTranslatedToFlightStatus() { Status status = Status.RESOURCE_EXHAUSTED; CallStatus callStatus = StatusUtils.fromGrpcStatus(status); assertEquals(FlightStatusCode.RESOURCE_EXHAUSTED, callStatus.code()); FlightStatusCode flightStatusCode = StatusUtils.fromGrpcStatusCode(status.getCode()); assertEquals(FlightStatusCode.RESOURCE_EXHAUSTED, flightStatusCode); } @Test public void testFlightResourceExhaustedTranslatedToGrpcStatua() { CallStatus callStatus = CallStatus.RESOURCE_EXHAUSTED; Status.Code grpcStatusCode = StatusUtils.toGrpcStatusCode(callStatus.code()); assertEquals(Status.RESOURCE_EXHAUSTED.getCode(), grpcStatusCode); Status grpcStatus = StatusUtils.toGrpcStatus(callStatus); assertEquals(Status.RESOURCE_EXHAUSTED.getCode(), grpcStatus.getCode()); } } apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/perf/000077500000000000000000000000001515451615200306205ustar00rootroot00000000000000PerformanceTestServer.java000066400000000000000000000163541515451615200357050ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/perf/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.perf; import com.google.common.collect.ImmutableList; import com.google.protobuf.InvalidProtocolBufferException; import java.io.IOException; import java.util.ArrayList; import java.util.List; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; import org.apache.arrow.flight.BackpressureStrategy; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.NoOpFlightProducer; import org.apache.arrow.flight.Ticket; import org.apache.arrow.flight.perf.impl.PerfOuterClass.Perf; import org.apache.arrow.flight.perf.impl.PerfOuterClass.Token; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; public class PerformanceTestServer implements AutoCloseable { private final FlightServer flightServer; private final BufferAllocator allocator; private final PerfProducer producer; private final boolean isNonBlocking; public PerformanceTestServer(BufferAllocator incomingAllocator, Location location) { this( incomingAllocator, location, new BackpressureStrategy() { private FlightProducer.ServerStreamListener listener; @Override public void register(FlightProducer.ServerStreamListener listener) { this.listener = listener; } @Override public WaitResult waitForListener(long timeout) { while (!listener.isReady() && !listener.isCancelled()) { // busy wait } return WaitResult.READY; } }, false); } public PerformanceTestServer( BufferAllocator incomingAllocator, Location location, BackpressureStrategy bpStrategy, boolean isNonBlocking) { this.allocator = incomingAllocator.newChildAllocator("perf-server", 0, Long.MAX_VALUE); this.producer = new PerfProducer(bpStrategy); this.flightServer = FlightServer.builder(this.allocator, location, producer).build(); this.isNonBlocking = isNonBlocking; } public Location getLocation() { return flightServer.getLocation(); } public PerformanceTestServer start() throws IOException { flightServer.start(); return this; } @Override public void close() throws Exception { AutoCloseables.close(flightServer, allocator); } private final class PerfProducer extends NoOpFlightProducer { private final BackpressureStrategy bpStrategy; private PerfProducer(BackpressureStrategy bpStrategy) { this.bpStrategy = bpStrategy; } @Override public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { bpStrategy.register(listener); final Runnable loadData = () -> { Token token = null; try { token = Token.parseFrom(ticket.getBytes()); } catch (InvalidProtocolBufferException e) { throw new RuntimeException(e); } Perf perf = token.getDefinition(); Schema schema = Schema.deserializeMessage(perf.getSchema().asReadOnlyByteBuffer()); try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); BigIntVector a = (BigIntVector) root.getVector("a")) { listener.setUseZeroCopy(true); listener.start(root); root.allocateNew(); int current = 0; long i = token.getStart(); while (i < token.getEnd()) { if (listener.isCancelled()) { root.clear(); return; } if (TestPerf.VALIDATE) { a.setSafe(current, i); } i++; current++; if (i % perf.getRecordsPerBatch() == 0) { root.setRowCount(current); bpStrategy.waitForListener(0); if (listener.isCancelled()) { root.clear(); return; } listener.putNext(); current = 0; root.allocateNew(); } } // send last partial batch. if (current != 0) { root.setRowCount(current); listener.putNext(); } listener.completed(); } }; if (!isNonBlocking) { loadData.run(); } else { final ExecutorService service = Executors.newSingleThreadExecutor(); Future unused = service.submit(loadData); service.shutdown(); } } @Override public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) { try { Preconditions.checkArgument(descriptor.isCommand()); Perf exec = Perf.parseFrom(descriptor.getCommand()); final Schema pojoSchema = new Schema( ImmutableList.of( Field.nullable("a", MinorType.BIGINT.getType()), Field.nullable("b", MinorType.BIGINT.getType()), Field.nullable("c", MinorType.BIGINT.getType()), Field.nullable("d", MinorType.BIGINT.getType()))); Token token = Token.newBuilder() .setDefinition(exec) .setStart(0) .setEnd(exec.getRecordsPerStream()) .build(); final Ticket ticket = new Ticket(token.toByteArray()); List endpoints = new ArrayList<>(); for (int i = 0; i < exec.getStreamCount(); i++) { endpoints.add(new FlightEndpoint(ticket, getLocation())); } return new FlightInfo( pojoSchema, descriptor, endpoints, -1, exec.getRecordsPerStream() * exec.getStreamCount()); } catch (InvalidProtocolBufferException e) { throw new RuntimeException(e); } } } } apache-arrow-java-19.0.0/flight/flight-core/src/test/java/org/apache/arrow/flight/perf/TestPerf.java000066400000000000000000000161071515451615200332240ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.perf; import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; import static org.apache.arrow.flight.Location.forGrpcInsecure; import com.google.common.base.MoreObjects; import com.google.common.base.Stopwatch; import com.google.common.collect.ImmutableList; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.ListeningExecutorService; import com.google.common.util.concurrent.MoreExecutors; import com.google.protobuf.ByteString; import java.util.Arrays; import java.util.List; import java.util.concurrent.Callable; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.Ticket; import org.apache.arrow.flight.perf.impl.PerfOuterClass.Perf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.Types.MinorType; 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; @Disabled public class TestPerf { public static final boolean VALIDATE = false; public static FlightDescriptor getPerfFlightDescriptor( long recordCount, int recordsPerBatch, int streamCount) { final Schema pojoSchema = new Schema( ImmutableList.of( Field.nullable("a", MinorType.BIGINT.getType()), Field.nullable("b", MinorType.BIGINT.getType()), Field.nullable("c", MinorType.BIGINT.getType()), Field.nullable("d", MinorType.BIGINT.getType()))); byte[] bytes = pojoSchema.serializeAsMessage(); ByteString serializedSchema = ByteString.copyFrom(bytes); return FlightDescriptor.command( Perf.newBuilder() .setRecordsPerStream(recordCount) .setRecordsPerBatch(recordsPerBatch) .setSchema(serializedSchema) .setStreamCount(streamCount) .build() .toByteArray()); } public static void main(String[] args) throws Exception { new TestPerf().throughput(); } @Test public void throughput() throws Exception { final int numRuns = 10; ListeningExecutorService pool = MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(4)); double[] throughPuts = new double[numRuns]; for (int i = 0; i < numRuns; i++) { try (final BufferAllocator a = new RootAllocator(Long.MAX_VALUE); final PerformanceTestServer server = new PerformanceTestServer(a, forGrpcInsecure(LOCALHOST, 0)).start(); final FlightClient client = FlightClient.builder(a, server.getLocation()).build(); ) { final FlightInfo info = client.getInfo(getPerfFlightDescriptor(50_000_000L, 4095, 2)); List> results = info.getEndpoints().stream() .map(t -> new Consumer(client, t.getTicket())) .map(t -> pool.submit(t)) .collect(Collectors.toList()); final Result r = Futures.whenAllSucceed(results) .call( () -> { Result res = new Result(); for (ListenableFuture f : results) { res.add(f.get()); } return res; }, pool) .get(); double seconds = r.nanos * 1.0d / 1000 / 1000 / 1000; throughPuts[i] = (r.bytes * 1.0d / 1024 / 1024) / seconds; System.out.printf( "Transferred %d records totaling %s bytes at %f MiB/s. %f record/s. %f batch/s.%n", r.rows, r.bytes, throughPuts[i], (r.rows * 1.0d) / seconds, (r.batches * 1.0d) / seconds); } } pool.shutdown(); System.out.println("Summary: "); double average = Arrays.stream(throughPuts).sum() / numRuns; double sqrSum = Arrays.stream(throughPuts).map(val -> val - average).map(val -> val * val).sum(); double stddev = Math.sqrt(sqrSum / numRuns); System.out.printf( "Average throughput: %f MiB/s, standard deviation: %f MiB/s%n", average, stddev); } private static final class Consumer implements Callable { private final FlightClient client; private final Ticket ticket; public Consumer(FlightClient client, Ticket ticket) { super(); this.client = client; this.ticket = ticket; } @Override public Result call() throws Exception { final Result r = new Result(); Stopwatch watch = Stopwatch.createStarted(); try (final FlightStream stream = client.getStream(ticket)) { final VectorSchemaRoot root = stream.getRoot(); try { BigIntVector a = (BigIntVector) root.getVector("a"); while (stream.next()) { int rows = root.getRowCount(); long aSum = r.aSum; for (int i = 0; i < rows; i++) { if (VALIDATE) { aSum += a.get(i); } } r.bytes += rows * 32L; r.rows += rows; r.aSum = aSum; r.batches++; } r.nanos = watch.elapsed(TimeUnit.NANOSECONDS); return r; } finally { root.clear(); } } } } private static final class Result { private long rows; private long aSum; private long bytes; private long nanos; private long batches; public void add(Result r) { rows += r.rows; aSum += r.aSum; bytes += r.bytes; batches += r.batches; nanos = Math.max(nanos, r.nanos); } @Override public String toString() { return MoreObjects.toStringHelper(this) .add("rows", rows) .add("aSum", aSum) .add("batches", batches) .add("bytes", bytes) .add("nanos", nanos) .toString(); } } } apache-arrow-java-19.0.0/flight/flight-core/src/test/proto/000077500000000000000000000000001515451615200234675ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/proto/perf.proto000066400000000000000000000022611515451615200255110ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at *

    * http://www.apache.org/licenses/LICENSE-2.0 *

    * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ syntax = "proto3"; option java_package = "org.apache.arrow.flight.perf.impl"; message Perf { bytes schema = 1; int32 stream_count = 2; int64 records_per_stream = 3; int32 records_per_batch = 4; } /* * Payload of ticket */ message Token { // definition of entire flight. Perf definition = 1; // inclusive start int64 start = 2; // exclusive end int64 end = 3; } apache-arrow-java-19.0.0/flight/flight-core/src/test/proto/test.proto000066400000000000000000000017541515451615200255420ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ syntax = "proto3"; option java_package = "org.apache.arrow.flight"; import "google/protobuf/empty.proto"; service TestService { rpc Test(google.protobuf.Empty) returns (google.protobuf.Empty) {} } apache-arrow-java-19.0.0/flight/flight-core/src/test/resources/000077500000000000000000000000001515451615200243365ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-core/src/test/resources/logback.xml000066400000000000000000000025211515451615200264620ustar00rootroot00000000000000 true 10000 true ${LILITH_HOSTNAME:-localhost} apache-arrow-java-19.0.0/flight/flight-integration-tests/000077500000000000000000000000001515451615200233115ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/pom.xml000066400000000000000000000106531515451615200246330ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-flight 19.0.0 flight-integration-tests jar Arrow Flight Integration Tests Integration tests for Flight RPC. org.apache.arrow arrow-vector org.apache.arrow arrow-memory-core org.apache.arrow flight-core org.apache.arrow flight-sql com.google.protobuf protobuf-java com.google.guava guava commons-cli commons-cli 1.11.0 org.slf4j slf4j-api org.apache.maven.plugins maven-shade-plugin make-assembly shade package false true jar-with-dependencies **/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 apache-arrow-java-19.0.0/flight/flight-integration-tests/src/000077500000000000000000000000001515451615200241005ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/000077500000000000000000000000001515451615200250245ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/000077500000000000000000000000001515451615200257455ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/000077500000000000000000000000001515451615200265345ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/000077500000000000000000000000001515451615200277555ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200311075ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/000077500000000000000000000000001515451615200323645ustar00rootroot00000000000000integration/000077500000000000000000000000001515451615200346305ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flighttests/000077500000000000000000000000001515451615200357725ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integrationAppMetadataFlightInfoEndpointScenario.java000066400000000000000000000061771515451615200461700ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.nio.charset.StandardCharsets; import java.util.Collections; import java.util.List; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.NoOpFlightProducer; import org.apache.arrow.flight.Ticket; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; /** Test app_metadata in FlightInfo and FlightEndpoint. */ final class AppMetadataFlightInfoEndpointScenario implements Scenario { @Override public FlightProducer producer(BufferAllocator allocator, Location location) throws Exception { return new AppMetadataFlightInfoEndpointProducer(); } @Override public void buildServer(FlightServer.Builder builder) throws Exception {} @Override public void client(BufferAllocator allocator, Location location, FlightClient client) throws Exception { byte[] cmd = "foobar".getBytes(StandardCharsets.UTF_8); FlightInfo info = client.getInfo(FlightDescriptor.command(cmd)); IntegrationAssertions.assertEquals(info.getAppMetadata(), cmd); IntegrationAssertions.assertEquals(info.getEndpoints().size(), 1); IntegrationAssertions.assertEquals(info.getEndpoints().get(0).getAppMetadata(), cmd); } /** producer for app_metadata test. */ static class AppMetadataFlightInfoEndpointProducer extends NoOpFlightProducer { @Override public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) { byte[] cmd = descriptor.getCommand(); Schema schema = new Schema( Collections.singletonList( Field.notNullable("number", Types.MinorType.UINT4.getType()))); List endpoints = Collections.singletonList( FlightEndpoint.builder(new Ticket("".getBytes(StandardCharsets.UTF_8))) .setAppMetadata(cmd) .build()); return FlightInfo.builder(schema, descriptor, endpoints).setAppMetadata(cmd).build(); } } } AuthBasicProtoScenario.java000066400000000000000000000100461515451615200432110ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.Optional; import org.apache.arrow.flight.Action; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightStatusCode; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.NoOpFlightProducer; import org.apache.arrow.flight.Result; import org.apache.arrow.flight.auth.BasicClientAuthHandler; import org.apache.arrow.flight.auth.BasicServerAuthHandler; import org.apache.arrow.memory.BufferAllocator; /** A scenario testing the built-in basic authentication Protobuf. */ final class AuthBasicProtoScenario implements Scenario { static final String USERNAME = "arrow"; static final String PASSWORD = "flight"; @Override public FlightProducer producer(BufferAllocator allocator, Location location) { return new NoOpFlightProducer() { @Override public void doAction(CallContext context, Action action, StreamListener listener) { listener.onNext(new Result(context.peerIdentity().getBytes(StandardCharsets.UTF_8))); listener.onCompleted(); } }; } @Override public void buildServer(FlightServer.Builder builder) { builder.authHandler( new BasicServerAuthHandler( new BasicServerAuthHandler.BasicAuthValidator() { @Override public byte[] getToken(String username, String password) throws Exception { if (!USERNAME.equals(username) || !PASSWORD.equals(password)) { throw CallStatus.UNAUTHENTICATED .withDescription("Username or password is invalid.") .toRuntimeException(); } return ("valid:" + username).getBytes(StandardCharsets.UTF_8); } @Override public Optional isValid(byte[] token) { if (token != null) { final String credential = new String(token, StandardCharsets.UTF_8); if (credential.startsWith("valid:")) { return Optional.of(credential.substring(6)); } } return Optional.empty(); } })); } @Override public void client(BufferAllocator allocator, Location location, FlightClient client) { final FlightRuntimeException e = IntegrationAssertions.assertThrows( FlightRuntimeException.class, () -> { client.listActions().forEach(act -> {}); }); if (!FlightStatusCode.UNAUTHENTICATED.equals(e.status().code())) { throw new AssertionError("Expected UNAUTHENTICATED but found " + e.status().code(), e); } client.authenticate(new BasicClientAuthHandler(USERNAME, PASSWORD)); final Result result = client.doAction(new Action("")).next(); if (!USERNAME.equals(new String(result.getBody(), StandardCharsets.UTF_8))) { throw new AssertionError( "Expected " + USERNAME + " but got " + Arrays.toString(result.getBody())); } } } DoExchangeEchoScenario.java000066400000000000000000000100421515451615200431220ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.nio.charset.StandardCharsets; import java.util.Collections; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.Location; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.IntVector; 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.Schema; import org.apache.arrow.vector.util.Validator; /** Test DoExchange by echoing data back to the client. */ final class DoExchangeEchoScenario implements Scenario { public static final byte[] COMMAND = "echo".getBytes(StandardCharsets.UTF_8); @Override public FlightProducer producer(BufferAllocator allocator, Location location) throws Exception { return new DoExchangeProducer(allocator); } @Override public void buildServer(FlightServer.Builder builder) {} @Override public void client(BufferAllocator allocator, Location location, FlightClient client) throws Exception { final Schema schema = new Schema(Collections.singletonList(Field.notNullable("x", new ArrowType.Int(32, true)))); try (final FlightClient.ExchangeReaderWriter stream = client.doExchange(FlightDescriptor.command(COMMAND)); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final FlightStream reader = stream.getReader(); // Write data and check that it gets echoed back. IntVector iv = (IntVector) root.getVector("x"); iv.allocateNew(); stream.getWriter().start(root); int rowCount = 10; for (int batchIdx = 0; batchIdx < 4; batchIdx++) { for (int rowIdx = 0; rowIdx < rowCount; rowIdx++) { iv.setSafe(rowIdx, batchIdx + rowIdx); } root.setRowCount(rowCount); boolean writeMetadata = batchIdx % 2 == 0; final byte[] rawMetadata = Integer.toString(batchIdx).getBytes(StandardCharsets.UTF_8); if (writeMetadata) { final ArrowBuf metadata = allocator.buffer(rawMetadata.length); metadata.writeBytes(rawMetadata); stream.getWriter().putNext(metadata); } else { stream.getWriter().putNext(); } IntegrationAssertions.assertTrue("Unexpected end of reader", reader.next()); if (writeMetadata) { IntegrationAssertions.assertNotNull(reader.getLatestMetadata()); final byte[] readMetadata = new byte[rawMetadata.length]; reader.getLatestMetadata().readBytes(readMetadata); IntegrationAssertions.assertEquals(rawMetadata, readMetadata); } else { IntegrationAssertions.assertNull(reader.getLatestMetadata()); } IntegrationAssertions.assertEquals(root.getSchema(), reader.getSchema()); Validator.compareVectorSchemaRoot(reader.getRoot(), root); } stream.getWriter().completed(); IntegrationAssertions.assertFalse("Expected to reach end of reader", reader.next()); } } } DoExchangeProducer.java000066400000000000000000000060361515451615200423530ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.util.Arrays; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.NoOpFlightProducer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; /** The server used for testing the Flight do_exchange method. */ final class DoExchangeProducer extends NoOpFlightProducer { private final BufferAllocator allocator; DoExchangeProducer(BufferAllocator allocator) { this.allocator = allocator; } @Override public void doExchange(CallContext context, FlightStream reader, ServerStreamListener writer) { FlightDescriptor descriptor = reader.getDescriptor(); if (descriptor.isCommand()) { if (Arrays.equals(DoExchangeEchoScenario.COMMAND, descriptor.getCommand())) { doEcho(reader, writer); } } throw CallStatus.UNIMPLEMENTED .withDescription("Unsupported descriptor: " + descriptor.toString()) .toRuntimeException(); } private void doEcho(FlightStream reader, ServerStreamListener writer) { VectorSchemaRoot root = null; VectorLoader loader = null; while (reader.next()) { if (reader.hasRoot()) { if (root == null) { root = VectorSchemaRoot.create(reader.getSchema(), allocator); loader = new VectorLoader(root); writer.start(root); } VectorUnloader unloader = new VectorUnloader(reader.getRoot()); try (final ArrowRecordBatch arb = unloader.getRecordBatch()) { loader.load(arb); } if (reader.getLatestMetadata() != null) { reader.getLatestMetadata().getReferenceManager().retain(); writer.putNext(reader.getLatestMetadata()); } else { writer.putNext(); } } else { // Pure metadata reader.getLatestMetadata().getReferenceManager().retain(); writer.putMetadata(reader.getLatestMetadata()); } } if (root != null) { root.close(); } writer.completed(); } } ExpirationTimeCancelFlightInfoScenario.java000066400000000000000000000051671515451615200463530ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.nio.charset.StandardCharsets; import org.apache.arrow.flight.CancelFlightInfoRequest; import org.apache.arrow.flight.CancelFlightInfoResult; import org.apache.arrow.flight.CancelStatus; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.Location; import org.apache.arrow.memory.BufferAllocator; /** Test CancelFlightInfo. */ final class ExpirationTimeCancelFlightInfoScenario implements Scenario { @Override public FlightProducer producer(BufferAllocator allocator, Location location) throws Exception { return new ExpirationTimeProducer(allocator); } @Override public void buildServer(FlightServer.Builder builder) {} @Override public void client(BufferAllocator allocator, Location location, FlightClient client) throws Exception { FlightInfo info = client.getInfo(FlightDescriptor.command("expiration".getBytes(StandardCharsets.UTF_8))); CancelFlightInfoRequest request = new CancelFlightInfoRequest(info); CancelFlightInfoResult result = client.cancelFlightInfo(request); IntegrationAssertions.assertEquals(CancelStatus.CANCELLED, result.getStatus()); // All requests should fail for (FlightEndpoint endpoint : info.getEndpoints()) { IntegrationAssertions.assertThrows( FlightRuntimeException.class, () -> { try (FlightStream stream = client.getStream(endpoint.getTicket())) { while (stream.next()) {} } }); } } } ExpirationTimeDoGetScenario.java000066400000000000000000000075401515451615200442130ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.Location; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.UInt4Vector; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; /** Test DoGet with expiration times. */ final class ExpirationTimeDoGetScenario implements Scenario { @Override public FlightProducer producer(BufferAllocator allocator, Location location) throws Exception { return new ExpirationTimeProducer(allocator); } @Override public void buildServer(FlightServer.Builder builder) {} @Override public void client(BufferAllocator allocator, Location location, FlightClient client) throws Exception { FlightInfo info = client.getInfo( FlightDescriptor.command("expiration_time".getBytes(StandardCharsets.UTF_8))); List batches = new ArrayList<>(); try { for (FlightEndpoint endpoint : info.getEndpoints()) { if (batches.size() == 0) { IntegrationAssertions.assertFalse( "endpoints[0] must not have expiration time", endpoint.getExpirationTime().isPresent()); } else { IntegrationAssertions.assertTrue( "endpoints[" + batches.size() + "] must have expiration time", endpoint.getExpirationTime().isPresent()); } try (FlightStream stream = client.getStream(endpoint.getTicket())) { while (stream.next()) { batches.add(new VectorUnloader(stream.getRoot()).getRecordBatch()); } } } // Check data IntegrationAssertions.assertEquals(3, batches.size()); try (final VectorSchemaRoot root = VectorSchemaRoot.create(ExpirationTimeProducer.SCHEMA, allocator)) { final VectorLoader loader = new VectorLoader(root); loader.load(batches.get(0)); IntegrationAssertions.assertEquals(1, root.getRowCount()); IntegrationAssertions.assertEquals(0, ((UInt4Vector) root.getVector(0)).getObject(0)); loader.load(batches.get(1)); IntegrationAssertions.assertEquals(1, root.getRowCount()); IntegrationAssertions.assertEquals(1, ((UInt4Vector) root.getVector(0)).getObject(0)); loader.load(batches.get(2)); IntegrationAssertions.assertEquals(1, root.getRowCount()); IntegrationAssertions.assertEquals(2, ((UInt4Vector) root.getVector(0)).getObject(0)); } } finally { AutoCloseables.close(batches); } } } ExpirationTimeListActionsScenario.java000066400000000000000000000043571515451615200454500ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.util.Iterator; import org.apache.arrow.flight.ActionType; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightConstants; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.Location; import org.apache.arrow.memory.BufferAllocator; /** Test ListActions with expiration times. */ final class ExpirationTimeListActionsScenario implements Scenario { @Override public FlightProducer producer(BufferAllocator allocator, Location location) throws Exception { return new ExpirationTimeProducer(allocator); } @Override public void buildServer(FlightServer.Builder builder) {} @Override public void client(BufferAllocator allocator, Location location, FlightClient client) throws Exception { Iterator actions = client.listActions().iterator(); IntegrationAssertions.assertTrue("Expected 2 actions", actions.hasNext()); ActionType action = actions.next(); IntegrationAssertions.assertEquals( FlightConstants.CANCEL_FLIGHT_INFO.getType(), action.getType()); IntegrationAssertions.assertTrue("Expected 2 actions", actions.hasNext()); action = actions.next(); IntegrationAssertions.assertEquals( FlightConstants.RENEW_FLIGHT_ENDPOINT.getType(), action.getType()); IntegrationAssertions.assertFalse("Expected 2 actions", actions.hasNext()); } } ExpirationTimeProducer.java000066400000000000000000000231511515451615200433040ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.io.IOException; import java.net.URISyntaxException; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.time.Instant; import java.time.temporal.ChronoUnit; import java.util.ArrayList; import java.util.Collections; import java.util.List; import org.apache.arrow.flight.Action; import org.apache.arrow.flight.ActionType; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.CancelFlightInfoRequest; import org.apache.arrow.flight.CancelFlightInfoResult; import org.apache.arrow.flight.CancelStatus; import org.apache.arrow.flight.FlightConstants; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.NoOpFlightProducer; import org.apache.arrow.flight.RenewFlightEndpointRequest; import org.apache.arrow.flight.Result; import org.apache.arrow.flight.Ticket; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.UInt4Vector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; /** * The server used for testing FlightEndpoint.expiration_time. * *

    GetFlightInfo() returns a FlightInfo that has the following three FlightEndpoints: * *

      *
    1. No expiration time *
    2. 5 seconds expiration time *
    3. 6 seconds expiration time *
    * * The client can't read data from the first endpoint multiple times but can read data from the * second and third endpoints. The client can't re-read data from the second endpoint 5 seconds * later. The client can't re-read data from the third endpoint 6 seconds later. * *

    The client can cancel a returned FlightInfo by pre-defined CancelFlightInfo action. The client * can't read data from endpoints even within 6 seconds after the action. * *

    The client can extend the expiration time of a FlightEndpoint in a returned FlightInfo by * pre-defined RenewFlightEndpoint action. The client can read data from endpoints multiple times * within more 10 seconds after the action. */ final class ExpirationTimeProducer extends NoOpFlightProducer { public static final Schema SCHEMA = new Schema( Collections.singletonList(Field.notNullable("number", Types.MinorType.UINT4.getType()))); private final BufferAllocator allocator; private final List statuses; ExpirationTimeProducer(BufferAllocator allocator) { this.allocator = allocator; this.statuses = new ArrayList<>(); } @Override public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) { statuses.clear(); List endpoints = new ArrayList<>(); Instant now = Instant.now(); endpoints.add(addEndpoint("No expiration time", null)); endpoints.add(addEndpoint("5 seconds", now.plus(5, ChronoUnit.SECONDS))); endpoints.add(addEndpoint("6 seconds", now.plus(6, ChronoUnit.SECONDS))); return new FlightInfo(SCHEMA, descriptor, endpoints, -1, -1); } @Override public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { // Obviously, not safe (since we don't lock), but we assume calls are not concurrent int index = parseIndexFromTicket(ticket); EndpointStatus status = statuses.get(index); if (status.cancelled) { listener.error( CallStatus.NOT_FOUND .withDescription( "Invalid flight: cancelled: " + new String(ticket.getBytes(), StandardCharsets.UTF_8)) .toRuntimeException()); return; } else if (status.expirationTime != null && Instant.now().isAfter(status.expirationTime)) { listener.error( CallStatus.NOT_FOUND .withDescription( "Invalid flight: expired: " + new String(ticket.getBytes(), StandardCharsets.UTF_8)) .toRuntimeException()); return; } else if (status.expirationTime == null && status.numGets > 0) { listener.error( CallStatus.NOT_FOUND .withDescription( "Invalid flight: can't read multiple times: " + new String(ticket.getBytes(), StandardCharsets.UTF_8)) .toRuntimeException()); return; } status.numGets++; try (final VectorSchemaRoot root = VectorSchemaRoot.create(SCHEMA, allocator)) { listener.start(root); UInt4Vector vector = (UInt4Vector) root.getVector(0); vector.setSafe(0, index); root.setRowCount(1); listener.putNext(); } listener.completed(); } @Override public void doAction(CallContext context, Action action, StreamListener listener) { try { if (action.getType().equals(FlightConstants.CANCEL_FLIGHT_INFO.getType())) { CancelFlightInfoRequest request = CancelFlightInfoRequest.deserialize(ByteBuffer.wrap(action.getBody())); CancelStatus cancelStatus = CancelStatus.UNSPECIFIED; for (FlightEndpoint endpoint : request.getInfo().getEndpoints()) { int index = parseIndexFromTicket(endpoint.getTicket()); EndpointStatus status = statuses.get(index); if (status.cancelled) { cancelStatus = CancelStatus.NOT_CANCELLABLE; } else { status.cancelled = true; if (cancelStatus == CancelStatus.UNSPECIFIED) { cancelStatus = CancelStatus.CANCELLED; } } } listener.onNext(new Result(new CancelFlightInfoResult(cancelStatus).serialize().array())); } else if (action.getType().equals(FlightConstants.RENEW_FLIGHT_ENDPOINT.getType())) { RenewFlightEndpointRequest request = RenewFlightEndpointRequest.deserialize(ByteBuffer.wrap(action.getBody())); FlightEndpoint endpoint = request.getFlightEndpoint(); int index = parseIndexFromTicket(endpoint.getTicket()); EndpointStatus status = statuses.get(index); if (status.cancelled) { listener.onError( CallStatus.INVALID_ARGUMENT .withDescription("Invalid flight: cancelled: " + index) .toRuntimeException()); return; } String ticketBody = new String(endpoint.getTicket().getBytes(), StandardCharsets.UTF_8); ticketBody += ": renewed (+ 10 seconds)"; Ticket ticket = new Ticket(ticketBody.getBytes(StandardCharsets.UTF_8)); Instant expiration = Instant.now().plus(10, ChronoUnit.SECONDS); status.expirationTime = expiration; FlightEndpoint newEndpoint = new FlightEndpoint( ticket, expiration, endpoint.getLocations().toArray(new Location[0])); listener.onNext(new Result(newEndpoint.serialize().array())); } else { listener.onError( CallStatus.INVALID_ARGUMENT .withDescription("Unknown action: " + action.getType()) .toRuntimeException()); return; } } catch (IOException | URISyntaxException e) { listener.onError( CallStatus.INTERNAL.withCause(e).withDescription(e.toString()).toRuntimeException()); return; } listener.onCompleted(); } @Override public void listActions(CallContext context, StreamListener listener) { listener.onNext(FlightConstants.CANCEL_FLIGHT_INFO); listener.onNext(FlightConstants.RENEW_FLIGHT_ENDPOINT); listener.onCompleted(); } private FlightEndpoint addEndpoint(String ticket, Instant expirationTime) { Ticket flightTicket = new Ticket( String.format("%d: %s", statuses.size(), ticket).getBytes(StandardCharsets.UTF_8)); statuses.add(new EndpointStatus(expirationTime)); return new FlightEndpoint(flightTicket, expirationTime); } private int parseIndexFromTicket(Ticket ticket) { final String contents = new String(ticket.getBytes(), StandardCharsets.UTF_8); int index = contents.indexOf(':'); if (index == -1) { throw CallStatus.INVALID_ARGUMENT .withDescription( "Invalid ticket: " + new String(ticket.getBytes(), StandardCharsets.UTF_8)) .toRuntimeException(); } int endpointIndex = Integer.parseInt(contents.substring(0, index)); if (endpointIndex < 0 || endpointIndex >= statuses.size()) { throw CallStatus.NOT_FOUND.withDescription("Out of bounds").toRuntimeException(); } return endpointIndex; } /** The status of a returned endpoint. */ static final class EndpointStatus { Instant expirationTime; int numGets; boolean cancelled; EndpointStatus(Instant expirationTime) { this.expirationTime = expirationTime; this.numGets = 0; this.cancelled = false; } } } ExpirationTimeRenewFlightEndpointScenario.java000066400000000000000000000051051515451615200471230ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.nio.charset.StandardCharsets; import java.time.Instant; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.RenewFlightEndpointRequest; import org.apache.arrow.memory.BufferAllocator; /** Test RenewFlightEndpoint. */ final class ExpirationTimeRenewFlightEndpointScenario implements Scenario { @Override public FlightProducer producer(BufferAllocator allocator, Location location) throws Exception { return new ExpirationTimeProducer(allocator); } @Override public void buildServer(FlightServer.Builder builder) {} @Override public void client(BufferAllocator allocator, Location location, FlightClient client) throws Exception { FlightInfo info = client.getInfo(FlightDescriptor.command("expiration".getBytes(StandardCharsets.UTF_8))); // Renew all endpoints with expiration time for (FlightEndpoint endpoint : info.getEndpoints()) { if (!endpoint.getExpirationTime().isPresent()) { continue; } Instant expiration = endpoint.getExpirationTime().get(); FlightEndpoint renewed = client.renewFlightEndpoint(new RenewFlightEndpointRequest(endpoint)); IntegrationAssertions.assertTrue( "Renewed FlightEndpoint must have expiration time", renewed.getExpirationTime().isPresent()); IntegrationAssertions.assertTrue( "Renewed FlightEndpoint must have newer expiration time", renewed.getExpirationTime().get().isAfter(expiration)); } } } FlightSqlExtensionScenario.java000066400000000000000000000214751515451615200441240ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.util.Map; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.SchemaResult; import org.apache.arrow.flight.sql.CancelResult; import org.apache.arrow.flight.sql.FlightSqlClient; import org.apache.arrow.flight.sql.FlightSqlProducer; import org.apache.arrow.flight.sql.impl.FlightSql; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.VectorSchemaRoot; /** * Integration test scenario for validating Flight SQL specs across multiple implementations. This * should ensure that RPC objects are being built and parsed correctly for multiple languages and * that the Arrow schemas are returned as expected. */ public class FlightSqlExtensionScenario extends FlightSqlScenario { @Override public void client(BufferAllocator allocator, Location location, FlightClient client) throws Exception { try (final FlightSqlClient sqlClient = new FlightSqlClient(client)) { validateMetadataRetrieval(sqlClient); validateStatementExecution(sqlClient); validatePreparedStatementExecution(allocator, sqlClient); validateTransactions(allocator, sqlClient); } } private void validateMetadataRetrieval(FlightSqlClient sqlClient) throws Exception { validate( FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA, sqlClient.getSqlInfo(), sqlClient, s -> { Map infoValues = readSqlInfoStream(s); IntegrationAssertions.assertEquals( Boolean.FALSE, infoValues.get(FlightSql.SqlInfo.FLIGHT_SQL_SERVER_SQL_VALUE)); IntegrationAssertions.assertEquals( Boolean.TRUE, infoValues.get(FlightSql.SqlInfo.FLIGHT_SQL_SERVER_SUBSTRAIT_VALUE)); IntegrationAssertions.assertEquals( "min_version", infoValues.get(FlightSql.SqlInfo.FLIGHT_SQL_SERVER_SUBSTRAIT_MIN_VERSION_VALUE)); IntegrationAssertions.assertEquals( "max_version", infoValues.get(FlightSql.SqlInfo.FLIGHT_SQL_SERVER_SUBSTRAIT_MAX_VERSION_VALUE)); IntegrationAssertions.assertEquals( FlightSql.SqlSupportedTransaction.SQL_SUPPORTED_TRANSACTION_SAVEPOINT_VALUE, infoValues.get(FlightSql.SqlInfo.FLIGHT_SQL_SERVER_TRANSACTION_VALUE)); IntegrationAssertions.assertEquals( Boolean.TRUE, infoValues.get(FlightSql.SqlInfo.FLIGHT_SQL_SERVER_CANCEL_VALUE)); IntegrationAssertions.assertEquals( 42, infoValues.get(FlightSql.SqlInfo.FLIGHT_SQL_SERVER_STATEMENT_TIMEOUT_VALUE)); IntegrationAssertions.assertEquals( 7, infoValues.get(FlightSql.SqlInfo.FLIGHT_SQL_SERVER_TRANSACTION_TIMEOUT_VALUE)); }); } private void validateStatementExecution(FlightSqlClient sqlClient) throws Exception { FlightInfo info = sqlClient.executeSubstrait(SUBSTRAIT_PLAN); validate(FlightSqlScenarioProducer.getQuerySchema(), info, sqlClient); SchemaResult result = sqlClient.getExecuteSubstraitSchema(SUBSTRAIT_PLAN); validateSchema(FlightSqlScenarioProducer.getQuerySchema(), result); IntegrationAssertions.assertEquals(CancelResult.CANCELLED, sqlClient.cancelQuery(info)); IntegrationAssertions.assertEquals( sqlClient.executeSubstraitUpdate(SUBSTRAIT_PLAN), UPDATE_STATEMENT_EXPECTED_ROWS); } private void validatePreparedStatementExecution( BufferAllocator allocator, FlightSqlClient sqlClient) throws Exception { try (FlightSqlClient.PreparedStatement preparedStatement = sqlClient.prepare(SUBSTRAIT_PLAN); VectorSchemaRoot parameters = VectorSchemaRoot.create(FlightSqlScenarioProducer.getQuerySchema(), allocator)) { parameters.setRowCount(1); preparedStatement.setParameters(parameters); validate(FlightSqlScenarioProducer.getQuerySchema(), preparedStatement.execute(), sqlClient); validateSchema(FlightSqlScenarioProducer.getQuerySchema(), preparedStatement.fetchSchema()); } try (FlightSqlClient.PreparedStatement preparedStatement = sqlClient.prepare(SUBSTRAIT_PLAN)) { IntegrationAssertions.assertEquals( preparedStatement.executeUpdate(), UPDATE_PREPARED_STATEMENT_EXPECTED_ROWS); } } private void validateTransactions(BufferAllocator allocator, FlightSqlClient sqlClient) throws Exception { final FlightSqlClient.Transaction transaction = sqlClient.beginTransaction(); IntegrationAssertions.assertEquals(TRANSACTION_ID, transaction.getTransactionId()); final FlightSqlClient.Savepoint savepoint = sqlClient.beginSavepoint(transaction, SAVEPOINT_NAME); IntegrationAssertions.assertEquals(SAVEPOINT_ID, savepoint.getSavepointId()); FlightInfo info = sqlClient.execute("SELECT STATEMENT", transaction); validate(FlightSqlScenarioProducer.getQueryWithTransactionSchema(), info, sqlClient); info = sqlClient.executeSubstrait(SUBSTRAIT_PLAN, transaction); validate(FlightSqlScenarioProducer.getQueryWithTransactionSchema(), info, sqlClient); SchemaResult schema = sqlClient.getExecuteSchema("SELECT STATEMENT", transaction); validateSchema(FlightSqlScenarioProducer.getQueryWithTransactionSchema(), schema); schema = sqlClient.getExecuteSubstraitSchema(SUBSTRAIT_PLAN, transaction); validateSchema(FlightSqlScenarioProducer.getQueryWithTransactionSchema(), schema); IntegrationAssertions.assertEquals( sqlClient.executeUpdate("UPDATE STATEMENT", transaction), UPDATE_STATEMENT_WITH_TRANSACTION_EXPECTED_ROWS); IntegrationAssertions.assertEquals( sqlClient.executeSubstraitUpdate(SUBSTRAIT_PLAN, transaction), UPDATE_STATEMENT_WITH_TRANSACTION_EXPECTED_ROWS); try (FlightSqlClient.PreparedStatement preparedStatement = sqlClient.prepare("SELECT PREPARED STATEMENT", transaction); VectorSchemaRoot parameters = VectorSchemaRoot.create(FlightSqlScenarioProducer.getQuerySchema(), allocator)) { parameters.setRowCount(1); preparedStatement.setParameters(parameters); validate( FlightSqlScenarioProducer.getQueryWithTransactionSchema(), preparedStatement.execute(), sqlClient); schema = preparedStatement.fetchSchema(); validateSchema(FlightSqlScenarioProducer.getQueryWithTransactionSchema(), schema); } try (FlightSqlClient.PreparedStatement preparedStatement = sqlClient.prepare(SUBSTRAIT_PLAN, transaction); VectorSchemaRoot parameters = VectorSchemaRoot.create(FlightSqlScenarioProducer.getQuerySchema(), allocator)) { parameters.setRowCount(1); preparedStatement.setParameters(parameters); validate( FlightSqlScenarioProducer.getQueryWithTransactionSchema(), preparedStatement.execute(), sqlClient); schema = preparedStatement.fetchSchema(); validateSchema(FlightSqlScenarioProducer.getQueryWithTransactionSchema(), schema); } try (FlightSqlClient.PreparedStatement preparedStatement = sqlClient.prepare("UPDATE PREPARED STATEMENT", transaction)) { IntegrationAssertions.assertEquals( preparedStatement.executeUpdate(), UPDATE_PREPARED_STATEMENT_WITH_TRANSACTION_EXPECTED_ROWS); } try (FlightSqlClient.PreparedStatement preparedStatement = sqlClient.prepare(SUBSTRAIT_PLAN, transaction)) { IntegrationAssertions.assertEquals( preparedStatement.executeUpdate(), UPDATE_PREPARED_STATEMENT_WITH_TRANSACTION_EXPECTED_ROWS); } sqlClient.rollback(savepoint); final FlightSqlClient.Savepoint savepoint2 = sqlClient.beginSavepoint(transaction, SAVEPOINT_NAME); IntegrationAssertions.assertEquals(SAVEPOINT_ID, savepoint2.getSavepointId()); sqlClient.release(savepoint); sqlClient.commit(transaction); final FlightSqlClient.Transaction transaction2 = sqlClient.beginTransaction(); IntegrationAssertions.assertEquals(TRANSACTION_ID, transaction2.getTransactionId()); sqlClient.rollback(transaction); } } FlightSqlIngestionScenario.java000066400000000000000000000111541515451615200441000ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import com.google.common.collect.ImmutableMap; import java.util.HashMap; import java.util.Map; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.sql.FlightSqlClient; import org.apache.arrow.flight.sql.FlightSqlClient.ExecuteIngestOptions; import org.apache.arrow.flight.sql.FlightSqlProducer; import org.apache.arrow.flight.sql.impl.FlightSql; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; /** * Integration test scenario for validating Flight SQL specs across multiple implementations. This * should ensure that RPC objects are being built and parsed correctly for multiple languages and * that the Arrow schemas are returned as expected. */ public class FlightSqlIngestionScenario extends FlightSqlScenario { @Override public FlightProducer producer(BufferAllocator allocator, Location location) throws Exception { FlightSqlScenarioProducer producer = (FlightSqlScenarioProducer) super.producer(allocator, location); producer .getSqlInfoBuilder() .withFlightSqlServerBulkIngestionTransaction(true) .withFlightSqlServerBulkIngestion(true); return producer; } @Override public void client(BufferAllocator allocator, Location location, FlightClient client) throws Exception { try (final FlightSqlClient sqlClient = new FlightSqlClient(client)) { validateMetadataRetrieval(sqlClient); validateIngestion(allocator, sqlClient); } } private void validateMetadataRetrieval(FlightSqlClient sqlClient) throws Exception { validate( FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA, sqlClient.getSqlInfo( FlightSql.SqlInfo.FLIGHT_SQL_SERVER_INGEST_TRANSACTIONS_SUPPORTED, FlightSql.SqlInfo.FLIGHT_SQL_SERVER_BULK_INGESTION), sqlClient, s -> { Map infoValues = readSqlInfoStream(s); IntegrationAssertions.assertEquals( Boolean.TRUE, infoValues.get( FlightSql.SqlInfo.FLIGHT_SQL_SERVER_INGEST_TRANSACTIONS_SUPPORTED_VALUE)); IntegrationAssertions.assertEquals( Boolean.TRUE, infoValues.get(FlightSql.SqlInfo.FLIGHT_SQL_SERVER_BULK_INGESTION_VALUE)); }); } private VectorSchemaRoot getIngestVectorRoot(BufferAllocator allocator) { Schema schema = FlightSqlScenarioProducer.getIngestSchema(); VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); root.setRowCount(3); return root; } private void validateIngestion(BufferAllocator allocator, FlightSqlClient sqlClient) { try (VectorSchemaRoot data = getIngestVectorRoot(allocator)) { TableDefinitionOptions tableDefinitionOptions = TableDefinitionOptions.newBuilder() .setIfExists(TableDefinitionOptions.TableExistsOption.TABLE_EXISTS_OPTION_REPLACE) .setIfNotExist( TableDefinitionOptions.TableNotExistOption.TABLE_NOT_EXIST_OPTION_CREATE) .build(); Map options = new HashMap<>(ImmutableMap.of("key1", "val1", "key2", "val2")); ExecuteIngestOptions executeIngestOptions = new ExecuteIngestOptions( "test_table", tableDefinitionOptions, true, "test_catalog", "test_schema", options); FlightSqlClient.Transaction transaction = new FlightSqlClient.Transaction(BULK_INGEST_TRANSACTION_ID); long updatedRows = sqlClient.executeIngest(data, executeIngestOptions, transaction); IntegrationAssertions.assertEquals(3L, updatedRows); } } } FlightSqlScenario.java000066400000000000000000000267121515451615200422260ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import static java.util.Objects.isNull; import com.google.protobuf.Any; import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.HashMap; import java.util.Map; import java.util.function.Consumer; import org.apache.arrow.flight.CallOption; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.SchemaResult; import org.apache.arrow.flight.Ticket; import org.apache.arrow.flight.sql.FlightSqlClient; import org.apache.arrow.flight.sql.FlightSqlProducer; import org.apache.arrow.flight.sql.FlightSqlUtils; import org.apache.arrow.flight.sql.impl.FlightSql; import org.apache.arrow.flight.sql.util.TableRef; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.UInt4Vector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.types.pojo.Schema; /** * Integration test scenario for validating Flight SQL specs across multiple implementations. This * should ensure that RPC objects are being built and parsed correctly for multiple languages and * that the Arrow schemas are returned as expected. */ public class FlightSqlScenario implements Scenario { public static final long UPDATE_STATEMENT_EXPECTED_ROWS = 10000L; public static final long UPDATE_STATEMENT_WITH_TRANSACTION_EXPECTED_ROWS = 15000L; public static final long UPDATE_PREPARED_STATEMENT_EXPECTED_ROWS = 20000L; public static final long UPDATE_PREPARED_STATEMENT_WITH_TRANSACTION_EXPECTED_ROWS = 25000L; public static final byte[] SAVEPOINT_ID = "savepoint_id".getBytes(StandardCharsets.UTF_8); public static final String SAVEPOINT_NAME = "savepoint_name"; public static final byte[] SUBSTRAIT_PLAN_TEXT = "plan".getBytes(StandardCharsets.UTF_8); public static final String SUBSTRAIT_VERSION = "version"; public static final FlightSqlClient.SubstraitPlan SUBSTRAIT_PLAN = new FlightSqlClient.SubstraitPlan(SUBSTRAIT_PLAN_TEXT, SUBSTRAIT_VERSION); public static final byte[] TRANSACTION_ID = "transaction_id".getBytes(StandardCharsets.UTF_8); public static final byte[] BULK_INGEST_TRANSACTION_ID = "123".getBytes(StandardCharsets.UTF_8); @Override public FlightProducer producer(BufferAllocator allocator, Location location) throws Exception { return new FlightSqlScenarioProducer(allocator); } @Override public void buildServer(FlightServer.Builder builder) throws Exception {} @Override public void client(BufferAllocator allocator, Location location, FlightClient client) throws Exception { try (final FlightSqlClient sqlClient = new FlightSqlClient(client)) { validateMetadataRetrieval(sqlClient); validateStatementExecution(sqlClient); validatePreparedStatementExecution(allocator, sqlClient); } } private void validateMetadataRetrieval(FlightSqlClient sqlClient) throws Exception { final CallOption[] options = new CallOption[0]; validate( FlightSqlProducer.Schemas.GET_CATALOGS_SCHEMA, sqlClient.getCatalogs(options), sqlClient); validateSchema( FlightSqlProducer.Schemas.GET_CATALOGS_SCHEMA, sqlClient.getCatalogsSchema(options)); validate( FlightSqlProducer.Schemas.GET_SCHEMAS_SCHEMA, sqlClient.getSchemas("catalog", "db_schema_filter_pattern", options), sqlClient); validateSchema(FlightSqlProducer.Schemas.GET_SCHEMAS_SCHEMA, sqlClient.getSchemasSchema()); validate( FlightSqlProducer.Schemas.GET_TABLES_SCHEMA, sqlClient.getTables( "catalog", "db_schema_filter_pattern", "table_filter_pattern", Arrays.asList("table", "view"), true, options), sqlClient); validateSchema( FlightSqlProducer.Schemas.GET_TABLES_SCHEMA, sqlClient.getTablesSchema(/*includeSchema*/ true, options)); validateSchema( FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA, sqlClient.getTablesSchema(/*includeSchema*/ false, options)); validate( FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA, sqlClient.getTableTypes(options), sqlClient); validateSchema( FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA, sqlClient.getTableTypesSchema(options)); validate( FlightSqlProducer.Schemas.GET_PRIMARY_KEYS_SCHEMA, sqlClient.getPrimaryKeys(TableRef.of("catalog", "db_schema", "table"), options), sqlClient); validateSchema( FlightSqlProducer.Schemas.GET_PRIMARY_KEYS_SCHEMA, sqlClient.getPrimaryKeysSchema(options)); validate( FlightSqlProducer.Schemas.GET_EXPORTED_KEYS_SCHEMA, sqlClient.getExportedKeys(TableRef.of("catalog", "db_schema", "table"), options), sqlClient); validateSchema( FlightSqlProducer.Schemas.GET_EXPORTED_KEYS_SCHEMA, sqlClient.getExportedKeysSchema(options)); validate( FlightSqlProducer.Schemas.GET_IMPORTED_KEYS_SCHEMA, sqlClient.getImportedKeys(TableRef.of("catalog", "db_schema", "table"), options), sqlClient); validateSchema( FlightSqlProducer.Schemas.GET_IMPORTED_KEYS_SCHEMA, sqlClient.getImportedKeysSchema(options)); validate( FlightSqlProducer.Schemas.GET_CROSS_REFERENCE_SCHEMA, sqlClient.getCrossReference( TableRef.of("pk_catalog", "pk_db_schema", "pk_table"), TableRef.of("fk_catalog", "fk_db_schema", "fk_table"), options), sqlClient); validateSchema( FlightSqlProducer.Schemas.GET_CROSS_REFERENCE_SCHEMA, sqlClient.getCrossReferenceSchema(options)); validate( FlightSqlProducer.Schemas.GET_TYPE_INFO_SCHEMA, sqlClient.getXdbcTypeInfo(options), sqlClient); validateSchema( FlightSqlProducer.Schemas.GET_TYPE_INFO_SCHEMA, sqlClient.getXdbcTypeInfoSchema(options)); FlightInfo sqlInfoFlightInfo = sqlClient.getSqlInfo( new FlightSql.SqlInfo[] { FlightSql.SqlInfo.FLIGHT_SQL_SERVER_NAME, FlightSql.SqlInfo.FLIGHT_SQL_SERVER_READ_ONLY }, options); Ticket ticket = sqlInfoFlightInfo.getEndpoints().get(0).getTicket(); FlightSql.CommandGetSqlInfo requestSqlInfoCommand = FlightSqlUtils.unpackOrThrow( Any.parseFrom(ticket.getBytes()), FlightSql.CommandGetSqlInfo.class); IntegrationAssertions.assertEquals( requestSqlInfoCommand.getInfo(0), FlightSql.SqlInfo.FLIGHT_SQL_SERVER_NAME_VALUE); IntegrationAssertions.assertEquals( requestSqlInfoCommand.getInfo(1), FlightSql.SqlInfo.FLIGHT_SQL_SERVER_READ_ONLY_VALUE); validate(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA, sqlInfoFlightInfo, sqlClient); validateSchema( FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA, sqlClient.getSqlInfoSchema(options)); } private void validateStatementExecution(FlightSqlClient sqlClient) throws Exception { FlightInfo info = sqlClient.execute("SELECT STATEMENT"); validate(FlightSqlScenarioProducer.getQuerySchema(), info, sqlClient); validateSchema( FlightSqlScenarioProducer.getQuerySchema(), sqlClient.getExecuteSchema("SELECT STATEMENT")); IntegrationAssertions.assertEquals( sqlClient.executeUpdate("UPDATE STATEMENT"), UPDATE_STATEMENT_EXPECTED_ROWS); } private void validatePreparedStatementExecution( BufferAllocator allocator, FlightSqlClient sqlClient) throws Exception { try (FlightSqlClient.PreparedStatement preparedStatement = sqlClient.prepare("SELECT PREPARED STATEMENT"); VectorSchemaRoot parameters = VectorSchemaRoot.create(FlightSqlScenarioProducer.getQuerySchema(), allocator)) { parameters.setRowCount(1); preparedStatement.setParameters(parameters); validate(FlightSqlScenarioProducer.getQuerySchema(), preparedStatement.execute(), sqlClient); validateSchema(FlightSqlScenarioProducer.getQuerySchema(), preparedStatement.fetchSchema()); } try (FlightSqlClient.PreparedStatement preparedStatement = sqlClient.prepare("UPDATE PREPARED STATEMENT")) { IntegrationAssertions.assertEquals( preparedStatement.executeUpdate(), UPDATE_PREPARED_STATEMENT_EXPECTED_ROWS); } } protected void validate(Schema expectedSchema, FlightInfo flightInfo, FlightSqlClient sqlClient) throws Exception { validate(expectedSchema, flightInfo, sqlClient, null); } protected void validate( Schema expectedSchema, FlightInfo flightInfo, FlightSqlClient sqlClient, Consumer streamConsumer) throws Exception { Ticket ticket = flightInfo.getEndpoints().get(0).getTicket(); try (FlightStream stream = sqlClient.getStream(ticket)) { Schema actualSchema = stream.getSchema(); IntegrationAssertions.assertEquals(expectedSchema, actualSchema); if (!isNull(streamConsumer)) { streamConsumer.accept(stream); } } } protected void validateSchema(Schema expected, SchemaResult actual) { IntegrationAssertions.assertEquals(expected, actual.getSchema()); } protected Map readSqlInfoStream(FlightStream stream) { Map infoValues = new HashMap<>(); while (stream.next()) { UInt4Vector infoName = (UInt4Vector) stream.getRoot().getVector(0); DenseUnionVector value = (DenseUnionVector) stream.getRoot().getVector(1); for (int i = 0; i < stream.getRoot().getRowCount(); i++) { final int code = infoName.get(i); if (infoValues.containsKey(code)) { throw new AssertionError("Duplicate SqlInfo value: " + code); } Object object; byte typeId = value.getTypeId(i); switch (typeId) { case 0: // string object = Preconditions.checkNotNull( value.getVarCharVector(typeId).getObject(value.getOffset(i))) .toString(); break; case 1: // bool object = value.getBitVector(typeId).getObject(value.getOffset(i)); break; case 2: // int64 object = value.getBigIntVector(typeId).getObject(value.getOffset(i)); break; case 3: // int32 object = value.getIntVector(typeId).getObject(value.getOffset(i)); break; default: throw new AssertionError("Decoding SqlInfo of type code " + typeId); } infoValues.put(code, object); } } return infoValues; } } FlightSqlScenarioProducer.java000066400000000000000000001016551515451615200437320ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import com.google.common.collect.ImmutableMap; import com.google.protobuf.Any; import com.google.protobuf.ByteString; import com.google.protobuf.InvalidProtocolBufferException; import com.google.protobuf.Message; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.Criteria; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.PutResult; import org.apache.arrow.flight.Result; import org.apache.arrow.flight.SchemaResult; import org.apache.arrow.flight.Ticket; import org.apache.arrow.flight.sql.CancelResult; import org.apache.arrow.flight.sql.FlightSqlColumnMetadata; import org.apache.arrow.flight.sql.FlightSqlProducer; import org.apache.arrow.flight.sql.SqlInfoBuilder; import org.apache.arrow.flight.sql.impl.FlightSql; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions.TableExistsOption; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions.TableNotExistOption; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; 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; /** Hardcoded Flight SQL producer used for cross-language integration tests. */ public class FlightSqlScenarioProducer implements FlightSqlProducer { public static final String SERVER_NAME = "Flight SQL Integration Test Server"; private final BufferAllocator allocator; private final SqlInfoBuilder sqlInfoBuilder; /** Constructor. */ public FlightSqlScenarioProducer(BufferAllocator allocator) { this.allocator = allocator; sqlInfoBuilder = new SqlInfoBuilder() .withFlightSqlServerName(SERVER_NAME) .withFlightSqlServerReadOnly(false) .withFlightSqlServerSql(false) .withFlightSqlServerSubstrait(true) .withFlightSqlServerSubstraitMinVersion("min_version") .withFlightSqlServerSubstraitMaxVersion("max_version") .withFlightSqlServerTransaction( FlightSql.SqlSupportedTransaction.SQL_SUPPORTED_TRANSACTION_SAVEPOINT) .withFlightSqlServerCancel(true) .withFlightSqlServerStatementTimeout(42) .withFlightSqlServerTransactionTimeout(7); } /** * Schema to be returned for mocking the statement/prepared statement results. Must be the same * across all languages. */ static Schema getQuerySchema() { return new Schema( Collections.singletonList( new Field( "id", new FieldType( true, new ArrowType.Int(64, true), null, new FlightSqlColumnMetadata.Builder() .tableName("test") .isAutoIncrement(true) .isCaseSensitive(false) .typeName("type_test") .schemaName("schema_test") .isSearchable(true) .catalogName("catalog_test") .precision(100) .remarks("test column") .build() .getMetadataMap()), null))); } /** * The expected schema for queries with transactions. * *

    Must be the same across all languages. */ static Schema getQueryWithTransactionSchema() { return new Schema( Collections.singletonList( new Field( "pkey", new FieldType( true, new ArrowType.Int(32, true), null, new FlightSqlColumnMetadata.Builder() .tableName("test") .isAutoIncrement(true) .isCaseSensitive(false) .typeName("type_test") .schemaName("schema_test") .isSearchable(true) .catalogName("catalog_test") .precision(100) .remarks("test column") .build() .getMetadataMap()), null))); } static Schema getIngestSchema() { return new Schema( Collections.singletonList(Field.nullable("test_field", new ArrowType.Int(64, true)))); } protected SqlInfoBuilder getSqlInfoBuilder() { return sqlInfoBuilder; } @Override public void beginSavepoint( FlightSql.ActionBeginSavepointRequest request, CallContext context, StreamListener listener) { if (!request.getName().equals(FlightSqlScenario.SAVEPOINT_NAME)) { listener.onError( CallStatus.INVALID_ARGUMENT .withDescription( String.format( "Expected name '%s', not '%s'", FlightSqlScenario.SAVEPOINT_NAME, request.getName())) .toRuntimeException()); return; } if (!Arrays.equals( request.getTransactionId().toByteArray(), FlightSqlScenario.TRANSACTION_ID)) { listener.onError( CallStatus.INVALID_ARGUMENT .withDescription( String.format( "Expected transaction ID '%s', not '%s'", Arrays.toString(FlightSqlScenario.TRANSACTION_ID), Arrays.toString(request.getTransactionId().toByteArray()))) .toRuntimeException()); return; } listener.onNext( FlightSql.ActionBeginSavepointResult.newBuilder() .setSavepointId(ByteString.copyFrom(FlightSqlScenario.SAVEPOINT_ID)) .build()); listener.onCompleted(); } @Override public void beginTransaction( FlightSql.ActionBeginTransactionRequest request, CallContext context, StreamListener listener) { listener.onNext( FlightSql.ActionBeginTransactionResult.newBuilder() .setTransactionId(ByteString.copyFrom(FlightSqlScenario.TRANSACTION_ID)) .build()); listener.onCompleted(); } @Override public void cancelQuery( FlightInfo info, CallContext context, StreamListener listener) { final String expectedTicket = "PLAN HANDLE"; if (info.getEndpoints().size() != 1) { listener.onError( CallStatus.INVALID_ARGUMENT .withDescription( String.format("Expected 1 endpoint, got %d", info.getEndpoints().size())) .toRuntimeException()); } final FlightEndpoint endpoint = info.getEndpoints().get(0); try { final Any any = Any.parseFrom(endpoint.getTicket().getBytes()); if (!any.is(FlightSql.TicketStatementQuery.class)) { listener.onError( CallStatus.INVALID_ARGUMENT .withDescription( String.format("Expected TicketStatementQuery, found '%s'", any.getTypeUrl())) .toRuntimeException()); } final FlightSql.TicketStatementQuery ticket = any.unpack(FlightSql.TicketStatementQuery.class); if (!ticket.getStatementHandle().toStringUtf8().equals(expectedTicket)) { listener.onError( CallStatus.INVALID_ARGUMENT .withDescription(String.format("Expected ticket '%s'", expectedTicket)) .toRuntimeException()); } listener.onNext(CancelResult.CANCELLED); listener.onCompleted(); } catch (InvalidProtocolBufferException e) { listener.onError( CallStatus.INVALID_ARGUMENT .withDescription("Invalid Protobuf:" + e) .withCause(e) .toRuntimeException()); } } @Override public void createPreparedStatement( FlightSql.ActionCreatePreparedStatementRequest request, CallContext context, StreamListener listener) { IntegrationAssertions.assertTrue( "Expect to be one of the two queries used on tests", request.getQuery().equals("SELECT PREPARED STATEMENT") || request.getQuery().equals("UPDATE PREPARED STATEMENT")); String text = request.getQuery(); if (!request.getTransactionId().isEmpty()) { text += " WITH TXN"; } text += " HANDLE"; final FlightSql.ActionCreatePreparedStatementResult result = FlightSql.ActionCreatePreparedStatementResult.newBuilder() .setPreparedStatementHandle(ByteString.copyFromUtf8(text)) .build(); listener.onNext(new Result(Any.pack(result).toByteArray())); listener.onCompleted(); } @Override public void createPreparedSubstraitPlan( FlightSql.ActionCreatePreparedSubstraitPlanRequest request, CallContext context, StreamListener listener) { if (!Arrays.equals( request.getPlan().getPlan().toByteArray(), FlightSqlScenario.SUBSTRAIT_PLAN_TEXT)) { listener.onError( CallStatus.INVALID_ARGUMENT .withDescription( String.format( "Expected plan '%s', not '%s'", Arrays.toString(FlightSqlScenario.SUBSTRAIT_PLAN_TEXT), Arrays.toString(request.getPlan().getPlan().toByteArray()))) .toRuntimeException()); return; } if (!FlightSqlScenario.SUBSTRAIT_VERSION.equals(request.getPlan().getVersion())) { listener.onError( CallStatus.INVALID_ARGUMENT .withDescription( String.format( "Expected version '%s', not '%s'", FlightSqlScenario.SUBSTRAIT_VERSION, request.getPlan().getVersion())) .toRuntimeException()); return; } final String handle = request.getTransactionId().isEmpty() ? "PREPARED PLAN HANDLE" : "PREPARED PLAN WITH TXN HANDLE"; final FlightSql.ActionCreatePreparedStatementResult result = FlightSql.ActionCreatePreparedStatementResult.newBuilder() .setPreparedStatementHandle(ByteString.copyFromUtf8(handle)) .build(); listener.onNext(result); listener.onCompleted(); } @Override public void closePreparedStatement( FlightSql.ActionClosePreparedStatementRequest request, CallContext context, StreamListener listener) { final String handle = request.getPreparedStatementHandle().toStringUtf8(); IntegrationAssertions.assertTrue( "Expect to be one of the queries used on tests", handle.equals("SELECT PREPARED STATEMENT HANDLE") || handle.equals("SELECT PREPARED STATEMENT WITH TXN HANDLE") || handle.equals("UPDATE PREPARED STATEMENT HANDLE") || handle.equals("UPDATE PREPARED STATEMENT WITH TXN HANDLE") || handle.equals("PREPARED PLAN HANDLE") || handle.equals("PREPARED PLAN WITH TXN HANDLE")); listener.onCompleted(); } @Override public void endSavepoint( FlightSql.ActionEndSavepointRequest request, CallContext context, StreamListener listener) { switch (request.getAction()) { case END_SAVEPOINT_RELEASE: case END_SAVEPOINT_ROLLBACK: if (!Arrays.equals( request.getSavepointId().toByteArray(), FlightSqlScenario.SAVEPOINT_ID)) { listener.onError( CallStatus.INVALID_ARGUMENT .withDescription( "Unexpected ID: " + Arrays.toString(request.getSavepointId().toByteArray())) .toRuntimeException()); } break; case UNRECOGNIZED: default: { listener.onError( CallStatus.INVALID_ARGUMENT .withDescription("Unknown action: " + request.getAction()) .toRuntimeException()); return; } } listener.onCompleted(); } @Override public void endTransaction( FlightSql.ActionEndTransactionRequest request, CallContext context, StreamListener listener) { switch (request.getAction()) { case END_TRANSACTION_COMMIT: case END_TRANSACTION_ROLLBACK: if (!Arrays.equals( request.getTransactionId().toByteArray(), FlightSqlScenario.TRANSACTION_ID)) { listener.onError( CallStatus.INVALID_ARGUMENT .withDescription( "Unexpected ID: " + Arrays.toString(request.getTransactionId().toByteArray())) .toRuntimeException()); } break; case UNRECOGNIZED: default: { listener.onError( CallStatus.INVALID_ARGUMENT .withDescription("Unknown action: " + request.getAction()) .toRuntimeException()); return; } } listener.onCompleted(); } @Override public FlightInfo getFlightInfoStatement( FlightSql.CommandStatementQuery command, CallContext context, FlightDescriptor descriptor) { IntegrationAssertions.assertEquals(command.getQuery(), "SELECT STATEMENT"); if (command.getTransactionId().isEmpty()) { String handle = "SELECT STATEMENT HANDLE"; FlightSql.TicketStatementQuery ticket = FlightSql.TicketStatementQuery.newBuilder() .setStatementHandle(ByteString.copyFromUtf8(handle)) .build(); return getFlightInfoForSchema(ticket, descriptor, getQuerySchema()); } else { String handle = "SELECT STATEMENT WITH TXN HANDLE"; FlightSql.TicketStatementQuery ticket = FlightSql.TicketStatementQuery.newBuilder() .setStatementHandle(ByteString.copyFromUtf8(handle)) .build(); return getFlightInfoForSchema(ticket, descriptor, getQueryWithTransactionSchema()); } } @Override public FlightInfo getFlightInfoSubstraitPlan( FlightSql.CommandStatementSubstraitPlan command, CallContext context, FlightDescriptor descriptor) { IntegrationAssertions.assertEquals( command.getPlan().getPlan().toByteArray(), FlightSqlScenario.SUBSTRAIT_PLAN_TEXT); IntegrationAssertions.assertEquals( command.getPlan().getVersion(), FlightSqlScenario.SUBSTRAIT_VERSION); String handle = command.getTransactionId().isEmpty() ? "PLAN HANDLE" : "PLAN WITH TXN HANDLE"; FlightSql.TicketStatementQuery ticket = FlightSql.TicketStatementQuery.newBuilder() .setStatementHandle(ByteString.copyFromUtf8(handle)) .build(); return getFlightInfoForSchema(ticket, descriptor, getQuerySchema()); } @Override public FlightInfo getFlightInfoPreparedStatement( FlightSql.CommandPreparedStatementQuery command, CallContext context, FlightDescriptor descriptor) { String handle = command.getPreparedStatementHandle().toStringUtf8(); if (handle.equals("SELECT PREPARED STATEMENT HANDLE") || handle.equals("PREPARED PLAN HANDLE")) { return getFlightInfoForSchema(command, descriptor, getQuerySchema()); } else if (handle.equals("SELECT PREPARED STATEMENT WITH TXN HANDLE") || handle.equals("PREPARED PLAN WITH TXN HANDLE")) { return getFlightInfoForSchema(command, descriptor, getQueryWithTransactionSchema()); } throw CallStatus.INVALID_ARGUMENT .withDescription("Unknown handle: " + handle) .toRuntimeException(); } @Override public SchemaResult getSchemaStatement( FlightSql.CommandStatementQuery command, CallContext context, FlightDescriptor descriptor) { IntegrationAssertions.assertEquals(command.getQuery(), "SELECT STATEMENT"); if (command.getTransactionId().isEmpty()) { return new SchemaResult(getQuerySchema()); } return new SchemaResult(getQueryWithTransactionSchema()); } @Override public SchemaResult getSchemaPreparedStatement( FlightSql.CommandPreparedStatementQuery command, CallContext context, FlightDescriptor descriptor) { String handle = command.getPreparedStatementHandle().toStringUtf8(); if (handle.equals("SELECT PREPARED STATEMENT HANDLE") || handle.equals("PREPARED PLAN HANDLE")) { return new SchemaResult(getQuerySchema()); } else if (handle.equals("SELECT PREPARED STATEMENT WITH TXN HANDLE") || handle.equals("PREPARED PLAN WITH TXN HANDLE")) { return new SchemaResult(getQueryWithTransactionSchema()); } throw CallStatus.INVALID_ARGUMENT .withDescription("Unknown handle: " + handle) .toRuntimeException(); } @Override public SchemaResult getSchemaSubstraitPlan( FlightSql.CommandStatementSubstraitPlan command, CallContext context, FlightDescriptor descriptor) { if (!Arrays.equals( command.getPlan().getPlan().toByteArray(), FlightSqlScenario.SUBSTRAIT_PLAN_TEXT)) { throw CallStatus.INVALID_ARGUMENT .withDescription( String.format( "Expected plan '%s', not '%s'", Arrays.toString(FlightSqlScenario.SUBSTRAIT_PLAN_TEXT), Arrays.toString(command.getPlan().getPlan().toByteArray()))) .toRuntimeException(); } if (!FlightSqlScenario.SUBSTRAIT_VERSION.equals(command.getPlan().getVersion())) { throw CallStatus.INVALID_ARGUMENT .withDescription( String.format( "Expected version '%s', not '%s'", FlightSqlScenario.SUBSTRAIT_VERSION, command.getPlan().getVersion())) .toRuntimeException(); } if (command.getTransactionId().isEmpty()) { return new SchemaResult(getQuerySchema()); } return new SchemaResult(getQueryWithTransactionSchema()); } @Override public void getStreamStatement( FlightSql.TicketStatementQuery ticket, CallContext context, ServerStreamListener listener) { final String handle = ticket.getStatementHandle().toStringUtf8(); if (handle.equals("SELECT STATEMENT HANDLE") || handle.equals("PLAN HANDLE")) { putEmptyBatchToStreamListener(listener, getQuerySchema()); } else if (handle.equals("SELECT STATEMENT WITH TXN HANDLE") || handle.equals("PLAN WITH TXN HANDLE")) { putEmptyBatchToStreamListener(listener, getQueryWithTransactionSchema()); } else { listener.error( CallStatus.INVALID_ARGUMENT .withDescription("Unknown handle: " + handle) .toRuntimeException()); } } @Override public void getStreamPreparedStatement( FlightSql.CommandPreparedStatementQuery command, CallContext context, ServerStreamListener listener) { String handle = command.getPreparedStatementHandle().toStringUtf8(); if (handle.equals("SELECT PREPARED STATEMENT HANDLE") || handle.equals("PREPARED PLAN HANDLE")) { putEmptyBatchToStreamListener(listener, getQuerySchema()); } else if (handle.equals("SELECT PREPARED STATEMENT WITH TXN HANDLE") || handle.equals("PREPARED PLAN WITH TXN HANDLE")) { putEmptyBatchToStreamListener(listener, getQueryWithTransactionSchema()); } else { listener.error( CallStatus.INVALID_ARGUMENT .withDescription("Unknown handle: " + handle) .toRuntimeException()); } } private Runnable acceptPutReturnConstant(StreamListener ackStream, long value) { return () -> { final FlightSql.DoPutUpdateResult build = FlightSql.DoPutUpdateResult.newBuilder().setRecordCount(value).build(); try (final ArrowBuf buffer = allocator.buffer(build.getSerializedSize())) { buffer.writeBytes(build.toByteArray()); ackStream.onNext(PutResult.metadata(buffer)); ackStream.onCompleted(); } }; } @Override public Runnable acceptPutStatement( FlightSql.CommandStatementUpdate command, CallContext context, FlightStream flightStream, StreamListener ackStream) { IntegrationAssertions.assertEquals(command.getQuery(), "UPDATE STATEMENT"); return acceptPutReturnConstant( ackStream, command.getTransactionId().isEmpty() ? FlightSqlScenario.UPDATE_STATEMENT_EXPECTED_ROWS : FlightSqlScenario.UPDATE_STATEMENT_WITH_TRANSACTION_EXPECTED_ROWS); } @Override public Runnable acceptPutStatementBulkIngest( FlightSql.CommandStatementIngest command, CallContext context, FlightStream flightStream, StreamListener ackStream) { IntegrationAssertions.assertEquals( TableExistsOption.TABLE_EXISTS_OPTION_REPLACE, command.getTableDefinitionOptions().getIfExists()); IntegrationAssertions.assertEquals( TableNotExistOption.TABLE_NOT_EXIST_OPTION_CREATE, command.getTableDefinitionOptions().getIfNotExist()); IntegrationAssertions.assertEquals("test_table", command.getTable()); IntegrationAssertions.assertEquals("test_catalog", command.getCatalog()); IntegrationAssertions.assertEquals("test_schema", command.getSchema()); IntegrationAssertions.assertEquals(true, command.getTemporary()); IntegrationAssertions.assertEquals( FlightSqlScenario.BULK_INGEST_TRANSACTION_ID, command.getTransactionId().toByteArray()); Map expectedOptions = new HashMap<>(ImmutableMap.of("key1", "val1", "key2", "val2")); IntegrationAssertions.assertEquals(expectedOptions.size(), command.getOptionsCount()); for (Map.Entry optionEntry : expectedOptions.entrySet()) { String key = optionEntry.getKey(); IntegrationAssertions.assertEquals(optionEntry.getValue(), command.getOptionsOrThrow(key)); } IntegrationAssertions.assertEquals(getIngestSchema(), flightStream.getSchema()); long rowCount = 0; while (flightStream.next()) { rowCount += flightStream.getRoot().getRowCount(); } return acceptPutReturnConstant(ackStream, rowCount); } @Override public Runnable acceptPutSubstraitPlan( FlightSql.CommandStatementSubstraitPlan command, CallContext context, FlightStream flightStream, StreamListener ackStream) { IntegrationAssertions.assertEquals( command.getPlan().getPlan().toByteArray(), FlightSqlScenario.SUBSTRAIT_PLAN_TEXT); IntegrationAssertions.assertEquals( command.getPlan().getVersion(), FlightSqlScenario.SUBSTRAIT_VERSION); return acceptPutReturnConstant( ackStream, command.getTransactionId().isEmpty() ? FlightSqlScenario.UPDATE_STATEMENT_EXPECTED_ROWS : FlightSqlScenario.UPDATE_STATEMENT_WITH_TRANSACTION_EXPECTED_ROWS); } @Override public Runnable acceptPutPreparedStatementUpdate( FlightSql.CommandPreparedStatementUpdate command, CallContext context, FlightStream flightStream, StreamListener ackStream) { final String handle = command.getPreparedStatementHandle().toStringUtf8(); if (handle.equals("UPDATE PREPARED STATEMENT HANDLE") || handle.equals("PREPARED PLAN HANDLE")) { return acceptPutReturnConstant( ackStream, FlightSqlScenario.UPDATE_PREPARED_STATEMENT_EXPECTED_ROWS); } else if (handle.equals("UPDATE PREPARED STATEMENT WITH TXN HANDLE") || handle.equals("PREPARED PLAN WITH TXN HANDLE")) { return acceptPutReturnConstant( ackStream, FlightSqlScenario.UPDATE_PREPARED_STATEMENT_WITH_TRANSACTION_EXPECTED_ROWS); } return () -> { ackStream.onError( CallStatus.INVALID_ARGUMENT .withDescription("Unknown handle: " + handle) .toRuntimeException()); }; } @Override public Runnable acceptPutPreparedStatementQuery( FlightSql.CommandPreparedStatementQuery command, CallContext context, FlightStream flightStream, StreamListener ackStream) { final String handle = command.getPreparedStatementHandle().toStringUtf8(); if (handle.equals("SELECT PREPARED STATEMENT HANDLE") || handle.equals("SELECT PREPARED STATEMENT WITH TXN HANDLE") || handle.equals("PREPARED PLAN HANDLE") || handle.equals("PREPARED PLAN WITH TXN HANDLE")) { IntegrationAssertions.assertEquals(getQuerySchema(), flightStream.getSchema()); return ackStream::onCompleted; } return () -> { ackStream.onError( CallStatus.INVALID_ARGUMENT .withDescription("Unknown handle: " + handle) .toRuntimeException()); }; } @Override public FlightInfo getFlightInfoSqlInfo( FlightSql.CommandGetSqlInfo request, CallContext context, FlightDescriptor descriptor) { return getFlightInfoForSchema(request, descriptor, Schemas.GET_SQL_INFO_SCHEMA); } @Override public void getStreamSqlInfo( FlightSql.CommandGetSqlInfo command, CallContext context, ServerStreamListener listener) { if (command.getInfoCount() == 2 && command.getInfo(0) == FlightSql.SqlInfo.FLIGHT_SQL_SERVER_NAME_VALUE && command.getInfo(1) == FlightSql.SqlInfo.FLIGHT_SQL_SERVER_READ_ONLY_VALUE) { // Integration test for the protocol messages putEmptyBatchToStreamListener(listener, Schemas.GET_SQL_INFO_SCHEMA); return; } sqlInfoBuilder.send(command.getInfoList(), listener); } @Override public void getStreamTypeInfo( FlightSql.CommandGetXdbcTypeInfo request, CallContext context, ServerStreamListener listener) { putEmptyBatchToStreamListener(listener, Schemas.GET_TYPE_INFO_SCHEMA); } @Override public FlightInfo getFlightInfoTypeInfo( FlightSql.CommandGetXdbcTypeInfo request, CallContext context, FlightDescriptor descriptor) { return getFlightInfoForSchema(request, descriptor, Schemas.GET_TYPE_INFO_SCHEMA); } @Override public FlightInfo getFlightInfoCatalogs( FlightSql.CommandGetCatalogs request, CallContext context, FlightDescriptor descriptor) { return getFlightInfoForSchema(request, descriptor, Schemas.GET_CATALOGS_SCHEMA); } private void putEmptyBatchToStreamListener(ServerStreamListener stream, Schema schema) { try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { stream.start(root); stream.putNext(); stream.completed(); } } @Override public void getStreamCatalogs(CallContext context, ServerStreamListener listener) { putEmptyBatchToStreamListener(listener, Schemas.GET_CATALOGS_SCHEMA); } @Override public FlightInfo getFlightInfoSchemas( FlightSql.CommandGetDbSchemas request, CallContext context, FlightDescriptor descriptor) { IntegrationAssertions.assertEquals(request.getCatalog(), "catalog"); IntegrationAssertions.assertEquals( request.getDbSchemaFilterPattern(), "db_schema_filter_pattern"); return getFlightInfoForSchema(request, descriptor, Schemas.GET_SCHEMAS_SCHEMA); } @Override public void getStreamSchemas( FlightSql.CommandGetDbSchemas command, CallContext context, ServerStreamListener listener) { putEmptyBatchToStreamListener(listener, Schemas.GET_SCHEMAS_SCHEMA); } @Override public FlightInfo getFlightInfoTables( FlightSql.CommandGetTables request, CallContext context, FlightDescriptor descriptor) { IntegrationAssertions.assertEquals(request.getCatalog(), "catalog"); IntegrationAssertions.assertEquals( request.getDbSchemaFilterPattern(), "db_schema_filter_pattern"); IntegrationAssertions.assertEquals(request.getTableNameFilterPattern(), "table_filter_pattern"); IntegrationAssertions.assertEquals(request.getTableTypesCount(), 2); IntegrationAssertions.assertEquals(request.getTableTypes(0), "table"); IntegrationAssertions.assertEquals(request.getTableTypes(1), "view"); return getFlightInfoForSchema(request, descriptor, Schemas.GET_TABLES_SCHEMA); } @Override public void getStreamTables( FlightSql.CommandGetTables command, CallContext context, ServerStreamListener listener) { putEmptyBatchToStreamListener(listener, Schemas.GET_TABLES_SCHEMA); } @Override public FlightInfo getFlightInfoTableTypes( FlightSql.CommandGetTableTypes request, CallContext context, FlightDescriptor descriptor) { return getFlightInfoForSchema(request, descriptor, Schemas.GET_TABLE_TYPES_SCHEMA); } @Override public void getStreamTableTypes(CallContext context, ServerStreamListener listener) { putEmptyBatchToStreamListener(listener, Schemas.GET_TABLE_TYPES_SCHEMA); } @Override public FlightInfo getFlightInfoPrimaryKeys( FlightSql.CommandGetPrimaryKeys request, CallContext context, FlightDescriptor descriptor) { IntegrationAssertions.assertEquals(request.getCatalog(), "catalog"); IntegrationAssertions.assertEquals(request.getDbSchema(), "db_schema"); IntegrationAssertions.assertEquals(request.getTable(), "table"); return getFlightInfoForSchema(request, descriptor, Schemas.GET_PRIMARY_KEYS_SCHEMA); } @Override public void getStreamPrimaryKeys( FlightSql.CommandGetPrimaryKeys command, CallContext context, ServerStreamListener listener) { putEmptyBatchToStreamListener(listener, Schemas.GET_PRIMARY_KEYS_SCHEMA); } @Override public FlightInfo getFlightInfoExportedKeys( FlightSql.CommandGetExportedKeys request, CallContext context, FlightDescriptor descriptor) { IntegrationAssertions.assertEquals(request.getCatalog(), "catalog"); IntegrationAssertions.assertEquals(request.getDbSchema(), "db_schema"); IntegrationAssertions.assertEquals(request.getTable(), "table"); return getFlightInfoForSchema(request, descriptor, Schemas.GET_EXPORTED_KEYS_SCHEMA); } @Override public FlightInfo getFlightInfoImportedKeys( FlightSql.CommandGetImportedKeys request, CallContext context, FlightDescriptor descriptor) { IntegrationAssertions.assertEquals(request.getCatalog(), "catalog"); IntegrationAssertions.assertEquals(request.getDbSchema(), "db_schema"); IntegrationAssertions.assertEquals(request.getTable(), "table"); return getFlightInfoForSchema(request, descriptor, Schemas.GET_IMPORTED_KEYS_SCHEMA); } @Override public FlightInfo getFlightInfoCrossReference( FlightSql.CommandGetCrossReference request, CallContext context, FlightDescriptor descriptor) { IntegrationAssertions.assertEquals(request.getPkCatalog(), "pk_catalog"); IntegrationAssertions.assertEquals(request.getPkDbSchema(), "pk_db_schema"); IntegrationAssertions.assertEquals(request.getPkTable(), "pk_table"); IntegrationAssertions.assertEquals(request.getFkCatalog(), "fk_catalog"); IntegrationAssertions.assertEquals(request.getFkDbSchema(), "fk_db_schema"); IntegrationAssertions.assertEquals(request.getFkTable(), "fk_table"); return getFlightInfoForSchema(request, descriptor, Schemas.GET_CROSS_REFERENCE_SCHEMA); } @Override public void getStreamExportedKeys( FlightSql.CommandGetExportedKeys command, CallContext context, ServerStreamListener listener) { putEmptyBatchToStreamListener(listener, Schemas.GET_EXPORTED_KEYS_SCHEMA); } @Override public void getStreamImportedKeys( FlightSql.CommandGetImportedKeys command, CallContext context, ServerStreamListener listener) { putEmptyBatchToStreamListener(listener, Schemas.GET_IMPORTED_KEYS_SCHEMA); } @Override public void getStreamCrossReference( FlightSql.CommandGetCrossReference command, CallContext context, ServerStreamListener listener) { putEmptyBatchToStreamListener(listener, Schemas.GET_CROSS_REFERENCE_SCHEMA); } @Override public void close() throws Exception {} @Override public void listFlights( CallContext context, Criteria criteria, StreamListener listener) {} private FlightInfo getFlightInfoForSchema( final T request, final FlightDescriptor descriptor, final Schema schema) { final Ticket ticket = new Ticket(Any.pack(request).toByteArray()); final List endpoints = Collections.singletonList(new FlightEndpoint(ticket)); return new FlightInfo(schema, descriptor, endpoints, -1, -1); } } IntegrationAssertions.java000066400000000000000000000075021515451615200431770ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.io.PrintWriter; import java.io.StringWriter; import java.util.Arrays; import java.util.Objects; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightRuntimeException; /** Utility methods to implement integration tests without using JUnit assertions. */ final class IntegrationAssertions { /** * Assert that the given code throws the given exception or subclass thereof. * * @param clazz The exception type. * @param body The code to run. * @param The exception type. * @return The thrown exception. */ @SuppressWarnings("unchecked") static T assertThrows(Class clazz, AssertThrows body) { try { body.run(); } catch (Throwable t) { if (clazz.isInstance(t)) { return (T) t; } throw new AssertionError( "Expected exception of class " + clazz + " but got " + t.getClass(), t); } throw new AssertionError("Expected exception of class " + clazz + " but did not throw."); } /** Assert that the two (non-array) objects are equal. */ static void assertEquals(Object expected, Object actual) { if (!Objects.equals(expected, actual)) { throw new AssertionError("Expected:\n" + expected + "\nbut got:\n" + actual); } } /** Assert that the two arrays are equal. */ static void assertEquals(byte[] expected, byte[] actual) { if (!Arrays.equals(expected, actual)) { throw new AssertionError( String.format( "Expected:\n%s\nbut got:\n%s", Arrays.toString(expected), Arrays.toString(actual))); } } /** Assert that the value is false, using the given message as an error otherwise. */ static void assertFalse(String message, boolean value) { if (value) { throw new AssertionError("Expected false: " + message); } } /** Assert that the value is true, using the given message as an error otherwise. */ static void assertTrue(String message, boolean value) { if (!value) { throw new AssertionError("Expected true: " + message); } } static void assertNull(Object actual) { if (actual != null) { throw new AssertionError("Expected: null\n\nbut got: " + actual); } } static void assertNotNull(Object actual) { if (actual == null) { throw new AssertionError("Expected: (not null)\n\nbut got: null\n"); } } /** Convert a throwable into a FlightRuntimeException with error details, for debugging. */ static FlightRuntimeException toFlightRuntimeException(Throwable t) { final StringWriter stringWriter = new StringWriter(); final PrintWriter writer = new PrintWriter(stringWriter); t.printStackTrace(writer); return CallStatus.UNKNOWN .withCause(t) .withDescription("Unknown error: " + t + "\n. Stack trace:\n" + stringWriter.toString()) .toRuntimeException(); } /** An interface used with {@link #assertThrows(Class, AssertThrows)}. */ @FunctionalInterface interface AssertThrows { void run() throws Throwable; } } IntegrationProducer.java000066400000000000000000000172541515451615200426350ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Collections; import java.util.HashSet; import java.util.List; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import java.util.stream.Collectors; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.NoOpFlightProducer; import org.apache.arrow.flight.PutResult; import org.apache.arrow.flight.Ticket; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; import org.apache.arrow.vector.dictionary.DictionaryProvider; 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.apache.arrow.vector.util.DictionaryUtility; /** * A FlightProducer that hosts an in memory store of Arrow buffers. Used for integration testing. */ public class IntegrationProducer extends NoOpFlightProducer implements AutoCloseable { private final ConcurrentMap datasets = new ConcurrentHashMap<>(); private final BufferAllocator allocator; private Location location; /** * Constructs a new instance. * * @param allocator The allocator for creating new Arrow buffers. * @param location The location of the storage. */ public IntegrationProducer(BufferAllocator allocator, Location location) { super(); this.allocator = allocator; this.location = location; } /** * Update the location after server start. * *

    Useful for binding to port 0 to get a free port. */ public void setLocation(Location location) { this.location = location; } @Override public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { try { FlightDescriptor descriptor = FlightDescriptor.deserialize(ByteBuffer.wrap(ticket.getBytes())); Dataset dataset = datasets.get(descriptor); if (dataset == null) { listener.error( CallStatus.NOT_FOUND .withDescription("Unknown ticket: " + descriptor) .toRuntimeException()); return; } dataset.streamTo(allocator, listener); } catch (Exception ex) { listener.error(IntegrationAssertions.toFlightRuntimeException(ex)); } } @Override public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) { Dataset h = datasets.get(descriptor); if (h == null) { throw CallStatus.NOT_FOUND .withDescription("Unknown descriptor: " + descriptor) .toRuntimeException(); } return h.getFlightInfo(location); } @Override public Runnable acceptPut( CallContext context, final FlightStream flightStream, final StreamListener ackStream) { return () -> { List batches = new ArrayList<>(); try { try (VectorSchemaRoot root = flightStream.getRoot()) { VectorUnloader unloader = new VectorUnloader(root); while (flightStream.next()) { ackStream.onNext(PutResult.metadata(flightStream.getLatestMetadata())); batches.add(unloader.getRecordBatch()); } // Closing the stream will release the dictionaries, take ownership final Dataset dataset = new Dataset( flightStream.getDescriptor(), flightStream.getSchema(), flightStream.takeDictionaryOwnership(), batches); batches.clear(); datasets.put(flightStream.getDescriptor(), dataset); } finally { AutoCloseables.close(batches); } } catch (Exception ex) { ackStream.onError(IntegrationAssertions.toFlightRuntimeException(ex)); } }; } @Override public void close() throws Exception { AutoCloseables.close(datasets.values()); datasets.clear(); } private static final class Dataset implements AutoCloseable { private final FlightDescriptor descriptor; private final Schema schema; private final DictionaryProvider dictionaryProvider; private final List batches; private Dataset( FlightDescriptor descriptor, Schema schema, DictionaryProvider dictionaryProvider, List batches) { this.descriptor = descriptor; this.schema = schema; this.dictionaryProvider = dictionaryProvider; this.batches = new ArrayList<>(batches); } public FlightInfo getFlightInfo(Location location) { ByteBuffer serializedDescriptor = descriptor.serialize(); byte[] descriptorBytes = new byte[serializedDescriptor.remaining()]; serializedDescriptor.get(descriptorBytes); final List endpoints = Collections.singletonList(new FlightEndpoint(new Ticket(descriptorBytes), location)); return new FlightInfo( messageFormatSchema(), descriptor, endpoints, batches.stream().mapToLong(ArrowRecordBatch::computeBodyLength).sum(), batches.stream().mapToInt(ArrowRecordBatch::getLength).sum()); } private Schema messageFormatSchema() { final Set dictionaryIdsUsed = new HashSet<>(); final List messageFormatFields = schema.getFields().stream() .map(f -> DictionaryUtility.toMessageFormat(f, dictionaryProvider, dictionaryIdsUsed)) .collect(Collectors.toList()); return new Schema(messageFormatFields, schema.getCustomMetadata()); } @Override public void close() throws Exception { AutoCloseables.close(batches); } public void streamTo(BufferAllocator allocator, ServerStreamListener listener) { try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { listener.start(root, dictionaryProvider); final VectorLoader loader = new VectorLoader(root); int counter = 0; for (ArrowRecordBatch batch : batches) { final byte[] rawMetadata = Integer.toString(counter).getBytes(StandardCharsets.UTF_8); final ArrowBuf metadata = allocator.buffer(rawMetadata.length); metadata.writeBytes(rawMetadata); loader.load(batch); // Transfers ownership of the buffer - do not free buffer ourselves listener.putNext(metadata); counter++; } listener.completed(); } } } } IntegrationTestClient.java000066400000000000000000000210321515451615200431150ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import java.io.File; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.util.List; import org.apache.arrow.flight.AsyncPutListener; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.PutResult; import org.apache.arrow.flight.Ticket; import org.apache.arrow.memory.ArrowBuf; 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.VectorUnloader; import org.apache.arrow.vector.ipc.JsonFileReader; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.Validator; import org.apache.commons.cli.CommandLine; import org.apache.commons.cli.CommandLineParser; import org.apache.commons.cli.DefaultParser; import org.apache.commons.cli.Options; import org.apache.commons.cli.ParseException; /** A Flight client for integration testing. */ class IntegrationTestClient { private static final org.slf4j.Logger LOGGER = org.slf4j.LoggerFactory.getLogger(IntegrationTestClient.class); private final Options options; private IntegrationTestClient() { options = new Options(); options.addOption("j", "json", true, "json file"); options.addOption("scenario", true, "The integration test scenario."); options.addOption("host", true, "The host to connect to."); options.addOption("port", true, "The port to connect to."); } public static void main(String[] args) { try { new IntegrationTestClient().run(args); } catch (ParseException e) { fatalError("Invalid parameters", e); } catch (IOException e) { fatalError("Error accessing files", e); } catch (Exception e) { fatalError("Unknown error", e); } } private static void fatalError(String message, Throwable e) { System.err.println(message); System.err.println(e.getMessage()); LOGGER.error(message, e); System.exit(1); } private void run(String[] args) throws Exception { final CommandLineParser parser = new DefaultParser(); final CommandLine cmd = parser.parse(options, args, false); final String host = cmd.getOptionValue("host", "localhost"); final int port = Integer.parseInt(cmd.getOptionValue("port", "31337")); final Location defaultLocation = Location.forGrpcInsecure(host, port); try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); final FlightClient client = FlightClient.builder(allocator, defaultLocation).build()) { if (cmd.hasOption("scenario")) { Scenarios.getScenario(cmd.getOptionValue("scenario")) .client(allocator, defaultLocation, client); } else { final String inputPath = cmd.getOptionValue("j"); testStream(allocator, client, inputPath); } } catch (InterruptedException e) { throw new RuntimeException(e); } } private static void testStream(BufferAllocator allocator, FlightClient client, String inputPath) throws IOException { // 1. Read data from JSON and upload to server. FlightDescriptor descriptor = FlightDescriptor.path(inputPath); try (JsonFileReader reader = new JsonFileReader(new File(inputPath), allocator); VectorSchemaRoot root = VectorSchemaRoot.create(reader.start(), allocator)) { FlightClient.ClientStreamListener stream = client.startPut( descriptor, root, reader, new AsyncPutListener() { int counter = 0; @Override public void onNext(PutResult val) { final byte[] metadataRaw = new byte[checkedCastToInt(val.getApplicationMetadata().readableBytes())]; val.getApplicationMetadata().readBytes(metadataRaw); final String metadata = new String(metadataRaw, StandardCharsets.UTF_8); if (!Integer.toString(counter).equals(metadata)) { throw new RuntimeException( String.format( "Invalid ACK from server. Expected '%d' but got '%s'.", counter, metadata)); } counter++; } }); int counter = 0; while (reader.read(root)) { final byte[] rawMetadata = Integer.toString(counter).getBytes(StandardCharsets.UTF_8); final ArrowBuf metadata = allocator.buffer(rawMetadata.length); metadata.writeBytes(rawMetadata); // Transfers ownership of the buffer, so do not release it ourselves stream.putNext(metadata); root.clear(); counter++; } stream.completed(); // Need to call this, or exceptions from the server get swallowed stream.getResult(); } // 2. Get the ticket for the data. FlightInfo info = client.getInfo(descriptor); List endpoints = info.getEndpoints(); if (endpoints.isEmpty()) { throw new RuntimeException("No endpoints returned from Flight server."); } for (FlightEndpoint endpoint : info.getEndpoints()) { // 3. Download the data from the server. List locations = endpoint.getLocations(); if (locations.isEmpty()) { // No locations provided, validate the server itself. testTicket(allocator, client, endpoint.getTicket(), inputPath); } else { // All locations should be equivalent, validate each one. for (Location location : locations) { try (FlightClient readClient = FlightClient.builder(allocator, location).build()) { testTicket(allocator, readClient, endpoint.getTicket(), inputPath); } catch (Exception e) { throw new RuntimeException(e); } } } } } private static void testTicket( BufferAllocator allocator, FlightClient readClient, Ticket ticket, String inputPath) { try (FlightStream stream = readClient.getStream(ticket); VectorSchemaRoot root = stream.getRoot(); VectorSchemaRoot downloadedRoot = VectorSchemaRoot.create(root.getSchema(), allocator); JsonFileReader reader = new JsonFileReader(new File(inputPath), allocator)) { VectorLoader loader = new VectorLoader(downloadedRoot); VectorUnloader unloader = new VectorUnloader(root); Schema jsonSchema = reader.start(); Validator.compareSchemas(root.getSchema(), jsonSchema); try (VectorSchemaRoot jsonRoot = VectorSchemaRoot.create(jsonSchema, allocator)) { while (stream.next()) { try (final ArrowRecordBatch arb = unloader.getRecordBatch()) { loader.load(arb); if (reader.read(jsonRoot)) { // 4. Validate the data. Validator.compareVectorSchemaRoot(jsonRoot, downloadedRoot); jsonRoot.clear(); } else { throw new RuntimeException("Flight stream has more batches than JSON"); } } } // Verify no more batches with data in JSON // NOTE: Currently the C++ Flight server skips empty batches at end of the stream if (reader.read(jsonRoot) && jsonRoot.getRowCount() > 0) { throw new RuntimeException("JSON has more batches with than Flight stream"); } } } catch (Exception e) { throw new RuntimeException(e); } } } IntegrationTestServer.java000066400000000000000000000073201515451615200431510ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.Location; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.commons.cli.CommandLine; import org.apache.commons.cli.CommandLineParser; import org.apache.commons.cli.DefaultParser; import org.apache.commons.cli.Options; import org.apache.commons.cli.ParseException; /** Flight server for integration testing. */ class IntegrationTestServer { private static final org.slf4j.Logger LOGGER = org.slf4j.LoggerFactory.getLogger(IntegrationTestServer.class); private final Options options; private IntegrationTestServer() { options = new Options(); options.addOption("port", true, "The port to serve on."); options.addOption("scenario", true, "The integration test scenario."); } private void run(String[] args) throws Exception { CommandLineParser parser = new DefaultParser(); CommandLine cmd = parser.parse(options, args, false); final int port = Integer.parseInt(cmd.getOptionValue("port", "31337")); final Location location = Location.forGrpcInsecure("localhost", port); final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); final FlightServer.Builder builder = FlightServer.builder().allocator(allocator).location(location); final FlightServer server; if (cmd.hasOption("scenario")) { final Scenario scenario = Scenarios.getScenario(cmd.getOptionValue("scenario")); scenario.buildServer(builder); server = builder.producer(scenario.producer(allocator, location)).build(); server.start(); } else { final IntegrationProducer producer = new IntegrationProducer(allocator, location); server = FlightServer.builder(allocator, location, producer).build().start(); producer.setLocation(Location.forGrpcInsecure("localhost", server.getPort())); } // Print out message for integration test script System.out.println("Server listening on localhost:" + server.getPort()); Runtime.getRuntime() .addShutdownHook( new Thread( () -> { try { System.out.println("\nExiting..."); AutoCloseables.close(server, allocator); } catch (Exception e) { e.printStackTrace(); } })); server.awaitTermination(); } public static void main(String[] args) { try { new IntegrationTestServer().run(args); } catch (ParseException e) { fatalError("Error parsing arguments", e); } catch (Exception e) { fatalError("Runtime error", e); } } private static void fatalError(String message, Throwable e) { System.err.println(message); System.err.println(e.getMessage()); LOGGER.error(message, e); System.exit(1); } } LocationReuseConnectionScenario.java000066400000000000000000000055071515451615200451240ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.nio.charset.StandardCharsets; import java.util.Collections; import java.util.List; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.NoOpFlightProducer; import org.apache.arrow.flight.Ticket; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.types.pojo.Schema; /** Test the 'arrow-flight-reuse-connection' scheme. */ public class LocationReuseConnectionScenario implements Scenario { @Override public FlightProducer producer(BufferAllocator allocator, Location location) throws Exception { return new ReuseConnectionProducer(); } @Override public void buildServer(FlightServer.Builder builder) throws Exception {} @Override public void client(BufferAllocator allocator, Location location, FlightClient client) throws Exception { final FlightInfo info = client.getInfo(FlightDescriptor.command("reuse".getBytes(StandardCharsets.UTF_8))); IntegrationAssertions.assertEquals(1, info.getEndpoints().size()); IntegrationAssertions.assertEquals(1, info.getEndpoints().get(0).getLocations().size()); Location actual = info.getEndpoints().get(0).getLocations().get(0); IntegrationAssertions.assertEquals(Location.reuseConnection().getUri(), actual.getUri()); } private static class ReuseConnectionProducer extends NoOpFlightProducer { @Override public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) { List endpoints = Collections.singletonList( new FlightEndpoint(new Ticket(new byte[0]), Location.reuseConnection())); return new FlightInfo( new Schema(Collections.emptyList()), descriptor, endpoints, /*bytes*/ -1, /*records*/ -1); } } } MiddlewareScenario.java000066400000000000000000000141771515451615200424100ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.Collections; import org.apache.arrow.flight.CallHeaders; import org.apache.arrow.flight.CallInfo; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightClientMiddleware; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightServerMiddleware; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.NoOpFlightProducer; import org.apache.arrow.flight.RequestContext; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.types.pojo.Schema; /** * Test an edge case in middleware: gRPC-Java consolidates headers and trailers if a call fails * immediately. On the gRPC implementation side, we need to watch for this, or else we'll have a * call with "no headers" if we only look for headers. */ final class MiddlewareScenario implements Scenario { private static final String HEADER = "x-middleware"; private static final String EXPECTED_HEADER_VALUE = "expected value"; private static final byte[] COMMAND_SUCCESS = "success".getBytes(StandardCharsets.UTF_8); @Override public FlightProducer producer(BufferAllocator allocator, Location location) { return new NoOpFlightProducer() { @Override public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) { if (descriptor.isCommand()) { if (Arrays.equals(COMMAND_SUCCESS, descriptor.getCommand())) { return new FlightInfo( new Schema(Collections.emptyList()), descriptor, Collections.emptyList(), -1, -1); } } throw CallStatus.UNIMPLEMENTED.toRuntimeException(); } }; } @Override public void buildServer(FlightServer.Builder builder) { builder.middleware( FlightServerMiddleware.Key.of("test"), new InjectingServerMiddleware.Factory()); } @Override public void client(BufferAllocator allocator, Location location, FlightClient ignored) throws Exception { final ExtractingClientMiddleware.Factory factory = new ExtractingClientMiddleware.Factory(); try (final FlightClient client = FlightClient.builder(allocator, location).intercept(factory).build()) { // Should fail immediately IntegrationAssertions.assertThrows( FlightRuntimeException.class, () -> client.getInfo(FlightDescriptor.command(new byte[0]))); if (!EXPECTED_HEADER_VALUE.equals(factory.extractedHeader)) { throw new AssertionError( "Expected to extract the header value '" + EXPECTED_HEADER_VALUE + "', but found: " + factory.extractedHeader); } // Should not fail factory.extractedHeader = ""; client.getInfo(FlightDescriptor.command(COMMAND_SUCCESS)); if (!EXPECTED_HEADER_VALUE.equals(factory.extractedHeader)) { throw new AssertionError( "Expected to extract the header value '" + EXPECTED_HEADER_VALUE + "', but found: " + factory.extractedHeader); } } } /** Middleware that inserts a constant value in outgoing requests. */ static class InjectingServerMiddleware implements FlightServerMiddleware { private final String headerValue; InjectingServerMiddleware(String incoming) { this.headerValue = incoming; } @Override public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) { outgoingHeaders.insert("x-middleware", headerValue); } @Override public void onCallCompleted(CallStatus status) {} @Override public void onCallErrored(Throwable err) {} /** The factory for the server middleware. */ static class Factory implements FlightServerMiddleware.Factory { @Override public InjectingServerMiddleware onCallStarted( CallInfo info, CallHeaders incomingHeaders, RequestContext context) { String incoming = incomingHeaders.get(HEADER); return new InjectingServerMiddleware(incoming == null ? "" : incoming); } } } /** Middleware that pulls a value out of incoming responses. */ static class ExtractingClientMiddleware implements FlightClientMiddleware { private final ExtractingClientMiddleware.Factory factory; public ExtractingClientMiddleware(ExtractingClientMiddleware.Factory factory) { this.factory = factory; } @Override public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) { outgoingHeaders.insert(HEADER, EXPECTED_HEADER_VALUE); } @Override public void onHeadersReceived(CallHeaders incomingHeaders) { this.factory.extractedHeader = incomingHeaders.get(HEADER); } @Override public void onCallCompleted(CallStatus status) {} /** The factory for the client middleware. */ static class Factory implements FlightClientMiddleware.Factory { String extractedHeader = null; @Override public FlightClientMiddleware onCallStarted(CallInfo info) { return new ExtractingClientMiddleware(this); } } } } OrderedScenario.java000066400000000000000000000147071515451615200417160ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.Objects; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.NoOpFlightProducer; import org.apache.arrow.flight.Ticket; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.message.IpcOption; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; /** Test the 'ordered' flag in FlightInfo. */ public class OrderedScenario implements Scenario { private static final Schema SCHEMA = new Schema( Collections.singletonList(Field.notNullable("number", Types.MinorType.INT.getType()))); private static final byte[] ORDERED_COMMAND = "ordered".getBytes(StandardCharsets.UTF_8); @Override public FlightProducer producer(BufferAllocator allocator, Location location) throws Exception { return new OrderedProducer(allocator); } @Override public void buildServer(FlightServer.Builder builder) throws Exception {} @Override public void client(BufferAllocator allocator, Location location, FlightClient client) throws Exception { final FlightInfo info = client.getInfo(FlightDescriptor.command(ORDERED_COMMAND)); IntegrationAssertions.assertTrue("ordered must be true", info.getOrdered()); IntegrationAssertions.assertEquals(3, info.getEndpoints().size()); int offset = 0; for (int multiplier : Arrays.asList(1, 10, 100)) { FlightEndpoint endpoint = info.getEndpoints().get(offset); IntegrationAssertions.assertTrue( "locations must be empty", endpoint.getLocations().isEmpty()); try (final FlightStream stream = client.getStream(endpoint.getTicket())) { IntegrationAssertions.assertEquals(SCHEMA, stream.getSchema()); IntegrationAssertions.assertTrue("stream must have a batch", stream.next()); IntVector number = (IntVector) stream.getRoot().getVector(0); IntegrationAssertions.assertEquals(3, stream.getRoot().getRowCount()); IntegrationAssertions.assertFalse("value must be non-null", number.isNull(0)); IntegrationAssertions.assertFalse("value must be non-null", number.isNull(1)); IntegrationAssertions.assertFalse("value must be non-null", number.isNull(2)); IntegrationAssertions.assertEquals(multiplier, number.get(0)); IntegrationAssertions.assertEquals(2 * multiplier, number.get(1)); IntegrationAssertions.assertEquals(3 * multiplier, number.get(2)); IntegrationAssertions.assertFalse("stream must have one batch", stream.next()); } offset++; } } private static class OrderedProducer extends NoOpFlightProducer { private static final byte[] TICKET_1 = "1".getBytes(StandardCharsets.UTF_8); private static final byte[] TICKET_2 = "2".getBytes(StandardCharsets.UTF_8); private static final byte[] TICKET_3 = "3".getBytes(StandardCharsets.UTF_8); private final BufferAllocator allocator; OrderedProducer(BufferAllocator allocator) { this.allocator = Objects.requireNonNull(allocator); } @Override public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { try (final VectorSchemaRoot root = VectorSchemaRoot.create(SCHEMA, allocator)) { IntVector number = (IntVector) root.getVector(0); if (Arrays.equals(ticket.getBytes(), TICKET_1)) { number.setSafe(0, 1); number.setSafe(1, 2); number.setSafe(2, 3); } else if (Arrays.equals(ticket.getBytes(), TICKET_2)) { number.setSafe(0, 10); number.setSafe(1, 20); number.setSafe(2, 30); } else if (Arrays.equals(ticket.getBytes(), TICKET_3)) { number.setSafe(0, 100); number.setSafe(1, 200); number.setSafe(2, 300); } else { listener.error( CallStatus.INVALID_ARGUMENT .withDescription( "Could not find flight: " + new String(ticket.getBytes(), StandardCharsets.UTF_8)) .toRuntimeException()); return; } root.setRowCount(3); listener.start(root); listener.putNext(); listener.completed(); } } @Override public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) { final boolean ordered = descriptor.isCommand() && Arrays.equals(descriptor.getCommand(), ORDERED_COMMAND); List endpoints; if (ordered) { endpoints = Arrays.asList( new FlightEndpoint(new Ticket(TICKET_1)), new FlightEndpoint(new Ticket(TICKET_2)), new FlightEndpoint(new Ticket(TICKET_3))); } else { endpoints = Arrays.asList( new FlightEndpoint(new Ticket(TICKET_1)), new FlightEndpoint(new Ticket(TICKET_3)), new FlightEndpoint(new Ticket(TICKET_2))); } return new FlightInfo( SCHEMA, descriptor, endpoints, /*bytes*/ -1, /*records*/ -1, ordered, IpcOption.DEFAULT); } } } PollFlightInfoProducer.java000066400000000000000000000046321515451615200432260ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.nio.charset.StandardCharsets; import java.time.Instant; import java.time.temporal.ChronoUnit; import java.util.Arrays; import java.util.Collections; import java.util.List; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.NoOpFlightProducer; import org.apache.arrow.flight.PollInfo; import org.apache.arrow.flight.Ticket; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; /** Test PollFlightInfo. */ class PollFlightInfoProducer extends NoOpFlightProducer { static final byte[] POLL_DESCRIPTOR = "poll".getBytes(StandardCharsets.UTF_8); @Override public PollInfo pollFlightInfo(CallContext context, FlightDescriptor descriptor) { Schema schema = new Schema( Collections.singletonList( Field.notNullable("number", Types.MinorType.UINT4.getType()))); List endpoints = Collections.singletonList( new FlightEndpoint(new Ticket("long-running query".getBytes(StandardCharsets.UTF_8)))); FlightInfo info = new FlightInfo(schema, descriptor, endpoints, -1, -1); if (descriptor.isCommand() && Arrays.equals(descriptor.getCommand(), POLL_DESCRIPTOR)) { return new PollInfo(info, null, 1.0, null); } else { return new PollInfo( info, FlightDescriptor.command(POLL_DESCRIPTOR), 0.1, Instant.now().plus(10, ChronoUnit.SECONDS)); } } } PollFlightInfoScenario.java000066400000000000000000000056651515451615200432150ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.nio.charset.StandardCharsets; import java.util.Optional; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.PollInfo; import org.apache.arrow.memory.BufferAllocator; /** Test PollFlightInfo. */ final class PollFlightInfoScenario implements Scenario { @Override public FlightProducer producer(BufferAllocator allocator, Location location) throws Exception { return new PollFlightInfoProducer(); } @Override public void buildServer(FlightServer.Builder builder) throws Exception {} @Override public void client(BufferAllocator allocator, Location location, FlightClient client) throws Exception { PollInfo info = client.pollInfo(FlightDescriptor.command("heavy query".getBytes(StandardCharsets.UTF_8))); IntegrationAssertions.assertNotNull(info.getFlightInfo()); Optional progress = info.getProgress(); IntegrationAssertions.assertTrue("progress is missing", progress.isPresent()); IntegrationAssertions.assertTrue( "progress is invalid", progress.get() >= 0.0 && progress.get() <= 1.0); IntegrationAssertions.assertTrue("expiration is missing", info.getExpirationTime().isPresent()); IntegrationAssertions.assertTrue( "descriptor is missing", info.getFlightDescriptor().isPresent()); info = client.pollInfo(info.getFlightDescriptor().get()); IntegrationAssertions.assertNotNull(info.getFlightInfo()); progress = info.getProgress(); IntegrationAssertions.assertTrue("progress is missing in finished query", progress.isPresent()); IntegrationAssertions.assertTrue( "progress isn't 1.0 in finished query", Math.abs(progress.get() - 1.0) < Math.ulp(1.0)); IntegrationAssertions.assertFalse( "expiration is set in finished query", info.getExpirationTime().isPresent()); IntegrationAssertions.assertFalse( "descriptor is set in finished query", info.getFlightDescriptor().isPresent()); } } Scenario.java000066400000000000000000000030371515451615200404030ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.Location; import org.apache.arrow.memory.BufferAllocator; /** A particular scenario in integration testing. */ interface Scenario { /** Construct the FlightProducer for a server in this scenario. */ FlightProducer producer(BufferAllocator allocator, Location location) throws Exception; /** Set any other server options. */ void buildServer(FlightServer.Builder builder) throws Exception; /** Run as the client in the scenario. */ void client(BufferAllocator allocator, Location location, FlightClient client) throws Exception; } Scenarios.java000066400000000000000000000102721515451615200405650ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.util.Map; import java.util.TreeMap; import java.util.concurrent.TimeUnit; import java.util.function.Supplier; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.Location; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; /** Scenarios for integration testing. */ final class Scenarios { private static Scenarios INSTANCE; private final Map> scenarios; private Scenarios() { scenarios = new TreeMap<>(); scenarios.put("auth:basic_proto", AuthBasicProtoScenario::new); scenarios.put( "expiration_time:cancel_flight_info", ExpirationTimeCancelFlightInfoScenario::new); scenarios.put( "expiration_time:renew_flight_endpoint", ExpirationTimeRenewFlightEndpointScenario::new); scenarios.put("expiration_time:do_get", ExpirationTimeDoGetScenario::new); scenarios.put("expiration_time:list_actions", ExpirationTimeListActionsScenario::new); scenarios.put("location:reuse_connection", LocationReuseConnectionScenario::new); scenarios.put("middleware", MiddlewareScenario::new); scenarios.put("ordered", OrderedScenario::new); scenarios.put("poll_flight_info", PollFlightInfoScenario::new); scenarios.put("flight_sql", FlightSqlScenario::new); scenarios.put("flight_sql:extension", FlightSqlExtensionScenario::new); scenarios.put("flight_sql:ingestion", FlightSqlIngestionScenario::new); scenarios.put("app_metadata_flight_info_endpoint", AppMetadataFlightInfoEndpointScenario::new); scenarios.put("session_options", SessionOptionsScenario::new); scenarios.put("do_exchange:echo", DoExchangeEchoScenario::new); } private static Scenarios getInstance() { if (INSTANCE == null) { INSTANCE = new Scenarios(); } return INSTANCE; } static Scenario getScenario(String scenario) { final Supplier ctor = getInstance().scenarios.get(scenario); if (ctor == null) { throw new IllegalArgumentException("Unknown integration test scenario: " + scenario); } return ctor.get(); } // Utility methods for implementing tests. public static void main(String[] args) { // Run scenarios one after the other final Location location = Location.forGrpcInsecure("localhost", 31337); for (final Map.Entry> entry : getInstance().scenarios.entrySet()) { System.out.println("Running test scenario: " + entry.getKey()); final Scenario scenario = entry.getValue().get(); try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) { final FlightServer.Builder builder = FlightServer.builder(allocator, location, scenario.producer(allocator, location)); scenario.buildServer(builder); try (final FlightServer server = builder.build()) { server.start(); try (final FlightClient client = FlightClient.builder(allocator, location).build()) { scenario.client(allocator, location, client); } server.shutdown(); server.awaitTermination(1, TimeUnit.SECONDS); System.out.println("Ran scenario " + entry.getKey()); } } catch (Exception e) { System.out.println("Exception while running scenario " + entry.getKey()); e.printStackTrace(); } } } } SessionOptionsProducer.java000066400000000000000000000107071515451615200433450ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.util.HashMap; import java.util.Map; import org.apache.arrow.flight.CloseSessionRequest; import org.apache.arrow.flight.CloseSessionResult; import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.flight.FlightServerMiddleware; import org.apache.arrow.flight.GetSessionOptionsRequest; import org.apache.arrow.flight.GetSessionOptionsResult; import org.apache.arrow.flight.ServerSessionMiddleware; import org.apache.arrow.flight.SessionOptionValue; import org.apache.arrow.flight.SessionOptionValueFactory; import org.apache.arrow.flight.SetSessionOptionsRequest; import org.apache.arrow.flight.SetSessionOptionsResult; import org.apache.arrow.flight.sql.NoOpFlightSqlProducer; /** * The server used for testing Sessions. * *

    SetSessionOptions(), GetSessionOptions(), and CloseSession() operate on a simple * SessionOptionValue store. */ final class SessionOptionsProducer extends NoOpFlightSqlProducer { private static final SessionOptionValue invalidOptionValue = SessionOptionValueFactory.makeSessionOptionValue("lol_invalid"); private final FlightServerMiddleware.Key sessionMiddlewareKey; SessionOptionsProducer(FlightServerMiddleware.Key sessionMiddlewareKey) { this.sessionMiddlewareKey = sessionMiddlewareKey; } @Override public void setSessionOptions( SetSessionOptionsRequest request, CallContext context, StreamListener listener) { Map errors = new HashMap(); ServerSessionMiddleware middleware = context.getMiddleware(sessionMiddlewareKey); ServerSessionMiddleware.Session session = middleware.getSession(); for (Map.Entry entry : request.getSessionOptions().entrySet()) { // Blacklisted option name if (entry.getKey().equals("lol_invalid")) { errors.put( entry.getKey(), new SetSessionOptionsResult.Error(SetSessionOptionsResult.ErrorValue.INVALID_NAME)); continue; } // Blacklisted option value // Recommend using a visitor to check polymorphic equality, but this check is easy if (entry.getValue().equals(invalidOptionValue)) { errors.put( entry.getKey(), new SetSessionOptionsResult.Error(SetSessionOptionsResult.ErrorValue.INVALID_VALUE)); continue; } // Business as usual: if (entry.getValue().isEmpty()) { session.eraseSessionOption(entry.getKey()); continue; } session.setSessionOption(entry.getKey(), entry.getValue()); } listener.onNext(new SetSessionOptionsResult(errors)); listener.onCompleted(); } @Override public void getSessionOptions( GetSessionOptionsRequest request, CallContext context, StreamListener listener) { ServerSessionMiddleware middleware = context.getMiddleware(sessionMiddlewareKey); final Map sessionOptions = middleware.getSession().getSessionOptions(); listener.onNext(new GetSessionOptionsResult(sessionOptions)); listener.onCompleted(); } @Override public void closeSession( CloseSessionRequest request, CallContext context, StreamListener listener) { ServerSessionMiddleware middleware = context.getMiddleware(sessionMiddlewareKey); try { middleware.closeSession(); } catch (FlightRuntimeException fre) { listener.onError(fre); return; } listener.onNext(new CloseSessionResult(CloseSessionResult.Status.CLOSED)); listener.onCompleted(); } } SessionOptionsScenario.java000066400000000000000000000140771515451615200433310ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import com.google.common.collect.ImmutableMap; import java.util.concurrent.atomic.AtomicInteger; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightServerMiddleware; import org.apache.arrow.flight.GetSessionOptionsRequest; import org.apache.arrow.flight.GetSessionOptionsResult; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.ServerSessionMiddleware; import org.apache.arrow.flight.SessionOptionValue; import org.apache.arrow.flight.SessionOptionValueFactory; import org.apache.arrow.flight.SetSessionOptionsRequest; import org.apache.arrow.flight.SetSessionOptionsResult; import org.apache.arrow.flight.client.ClientCookieMiddleware; import org.apache.arrow.flight.sql.FlightSqlClient; import org.apache.arrow.memory.BufferAllocator; /** Scenario to exercise Session Options functionality. */ final class SessionOptionsScenario implements Scenario { private final FlightServerMiddleware.Key key = FlightServerMiddleware.Key.of("sessionmiddleware"); @Override public FlightProducer producer(BufferAllocator allocator, Location location) throws Exception { return new SessionOptionsProducer(key); } @Override public void buildServer(FlightServer.Builder builder) { AtomicInteger counter = new AtomicInteger(1000); builder.middleware( key, new ServerSessionMiddleware.Factory(() -> String.valueOf(counter.getAndIncrement()))); } @Override public void client(BufferAllocator allocator, Location location, FlightClient ignored) throws Exception { final ClientCookieMiddleware.Factory factory = new ClientCookieMiddleware.Factory(); try (final FlightClient flightClient = FlightClient.builder(allocator, location).intercept(factory).build()) { final FlightSqlClient client = new FlightSqlClient(flightClient); // Set SetSessionOptionsRequest req1 = new SetSessionOptionsRequest( ImmutableMap.builder() .put("foolong", SessionOptionValueFactory.makeSessionOptionValue(123L)) .put("bardouble", SessionOptionValueFactory.makeSessionOptionValue(456.0)) .put( "lol_invalid", SessionOptionValueFactory.makeSessionOptionValue("this won't get set")) .put( "key_with_invalid_value", SessionOptionValueFactory.makeSessionOptionValue("lol_invalid")) .put( "big_ol_string_list", SessionOptionValueFactory.makeSessionOptionValue( new String[] {"a", "b", "sea", "dee", " ", " ", "geee", "(ã¥ï½¡â—•â€¿â€¿â—•ï½¡)ã¥"})) .build()); SetSessionOptionsResult res1 = client.setSessionOptions(req1); // Some errors IntegrationAssertions.assertEquals( ImmutableMap.builder() .put( "lol_invalid", new SetSessionOptionsResult.Error( SetSessionOptionsResult.ErrorValue.INVALID_NAME)) .put( "key_with_invalid_value", new SetSessionOptionsResult.Error( SetSessionOptionsResult.ErrorValue.INVALID_VALUE)) .build(), res1.getErrors()); // Some set, some omitted due to above errors GetSessionOptionsResult res2 = client.getSessionOptions(new GetSessionOptionsRequest()); IntegrationAssertions.assertEquals( ImmutableMap.builder() .put("foolong", SessionOptionValueFactory.makeSessionOptionValue(123L)) .put("bardouble", SessionOptionValueFactory.makeSessionOptionValue(456.0)) .put( "big_ol_string_list", SessionOptionValueFactory.makeSessionOptionValue( new String[] {"a", "b", "sea", "dee", " ", " ", "geee", "(ã¥ï½¡â—•â€¿â€¿â—•ï½¡)ã¥"})) .build(), res2.getSessionOptions()); // Update client.setSessionOptions( new SetSessionOptionsRequest( ImmutableMap.builder() // Delete .put("foolong", SessionOptionValueFactory.makeEmptySessionOptionValue()) // Update .put( "big_ol_string_list", SessionOptionValueFactory.makeSessionOptionValue( "a,b,sea,dee, , ,geee,(ã¥ï½¡â—•â€¿â€¿â—•ï½¡)ã¥")) .build())); GetSessionOptionsResult res4 = client.getSessionOptions(new GetSessionOptionsRequest()); IntegrationAssertions.assertEquals( ImmutableMap.builder() .put("bardouble", SessionOptionValueFactory.makeSessionOptionValue(456.0)) .put( "big_ol_string_list", SessionOptionValueFactory.makeSessionOptionValue( "a,b,sea,dee, , ,geee,(ã¥ï½¡â—•â€¿â€¿â—•ï½¡)ã¥")) .build(), res4.getSessionOptions()); } } } TestBufferAllocationListener.java000066400000000000000000000047201515451615200444250ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import java.util.ArrayList; import java.util.List; import org.apache.arrow.memory.AllocationListener; class TestBufferAllocationListener implements AllocationListener { static class Entry { StackTraceElement[] stackTrace; long size; boolean forAllocation; public Entry(StackTraceElement[] stackTrace, long size, boolean forAllocation) { this.stackTrace = stackTrace; this.size = size; this.forAllocation = forAllocation; } } List trail = new ArrayList<>(); public void onAllocation(long size) { trail.add(new Entry(Thread.currentThread().getStackTrace(), size, true)); } public void onRelease(long size) { trail.add(new Entry(Thread.currentThread().getStackTrace(), size, false)); } public void reThrowWithAddedAllocatorInfo(Exception e) { StringBuilder sb = new StringBuilder(); sb.append(e.getMessage()); sb.append("\n"); sb.append("[[Buffer allocation and release trail during the test execution: \n"); for (Entry trailEntry : trail) { sb.append( String.format( "%s: %d: %n%s", trailEntry.forAllocation ? "allocate" : "release", trailEntry.size, getStackTraceAsString(trailEntry.stackTrace))); } sb.append("]]"); throw new IllegalStateException(sb.toString(), e); } private String getStackTraceAsString(StackTraceElement[] elements) { StringBuilder sb = new StringBuilder(); for (int i = 1; i < elements.length; i++) { StackTraceElement s = elements[i]; sb.append("\t"); sb.append(s); sb.append("\n"); } return sb.toString(); } } apache-arrow-java-19.0.0/flight/flight-integration-tests/src/shade/000077500000000000000000000000001515451615200251645ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/shade/LICENSE.txt000066400000000000000000001627521515451615200270240ustar00rootroot00000000000000 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. apache-arrow-java-19.0.0/flight/flight-integration-tests/src/shade/NOTICE.txt000066400000000000000000000275171515451615200267220ustar00rootroot00000000000000Apache 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 apache-arrow-java-19.0.0/flight/flight-integration-tests/src/test/000077500000000000000000000000001515451615200250575ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/test/java/000077500000000000000000000000001515451615200260005ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/test/java/org/000077500000000000000000000000001515451615200265675ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/test/java/org/apache/000077500000000000000000000000001515451615200300105ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200311425ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/test/java/org/apache/arrow/flight/000077500000000000000000000000001515451615200324175ustar00rootroot00000000000000integration/000077500000000000000000000000001515451615200346635ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/test/java/org/apache/arrow/flighttests/000077500000000000000000000000001515451615200360255ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/test/java/org/apache/arrow/flight/integrationIntegrationTest.java000066400000000000000000000112331515451615200420130ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/test/java/org/apache/arrow/flight/integration/tests/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.integration.tests; import com.google.common.util.concurrent.ThreadFactoryBuilder; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.Location; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.Test; /** Run the integration test scenarios in-process. */ class IntegrationTest { @Test void authBasicProto() throws Exception { testScenario("auth:basic_proto"); } @Test void expirationTimeCancelFlightInfo() throws Exception { testScenario("expiration_time:cancel_flight_info"); } @Test void expirationTimeDoGet() throws Exception { testScenario("expiration_time:do_get"); } @Test void expirationTimeListActions() throws Exception { testScenario("expiration_time:list_actions"); } @Test void expirationTimeRenewFlightEndpoint() throws Exception { testScenario("expiration_time:renew_flight_endpoint"); } @Test void locationReuseConnection() throws Exception { testScenario("location:reuse_connection"); } @Test void middleware() throws Exception { testScenario("middleware"); } @Test void ordered() throws Exception { testScenario("ordered"); } @Test void pollFlightInfo() throws Exception { testScenario("poll_flight_info"); } @Test void flightSql() throws Exception { testScenario("flight_sql"); } @Test void flightSqlExtension() throws Exception { testScenario("flight_sql:extension"); } @Test void flightSqlIngestion() throws Exception { testScenario("flight_sql:ingestion"); } @Test void appMetadataFlightInfoEndpoint() throws Exception { testScenario("app_metadata_flight_info_endpoint"); } @Test void sessionOptions() throws Exception { testScenario("session_options"); } @Test void doExchangeEcho() throws Exception { testScenario("do_exchange:echo"); } void testScenario(String scenarioName) throws Exception { TestBufferAllocationListener listener = new TestBufferAllocationListener(); try (final BufferAllocator allocator = new RootAllocator(listener, Long.MAX_VALUE)) { final ExecutorService exec = Executors.newCachedThreadPool( new ThreadFactoryBuilder() .setNameFormat("integration-test-flight-server-executor-%d") .build()); final FlightServer.Builder builder = FlightServer.builder() .executor(exec) .allocator(allocator) .location(Location.forGrpcInsecure("0.0.0.0", 0)); final Scenario scenario = Scenarios.getScenario(scenarioName); scenario.buildServer(builder); builder.producer(scenario.producer(allocator, Location.forGrpcInsecure("0.0.0.0", 0))); try (final FlightServer server = builder.build()) { server.start(); final Location location = Location.forGrpcInsecure("localhost", server.getPort()); try (final FlightClient client = FlightClient.builder(allocator, location).build()) { scenario.client(allocator, location, client); } } // Shutdown the executor while allowing existing tasks to finish. // Without this wait, allocator.close() may get invoked earlier than an executor thread may // have finished freeing up resources // In that case, allocator.close() can throw an IllegalStateException for memory leak, leading // to flaky tests exec.shutdown(); final boolean unused = exec.awaitTermination(3, TimeUnit.SECONDS); } catch (IllegalStateException e) { // this could be due to Allocator detecting memory leak. Add allocation trail to help debug listener.reThrowWithAddedAllocatorInfo(e); } } } apache-arrow-java-19.0.0/flight/flight-integration-tests/src/test/resources/000077500000000000000000000000001515451615200270715ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-integration-tests/src/test/resources/logback.xml000066400000000000000000000025301515451615200312150ustar00rootroot00000000000000 %d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/000077500000000000000000000000001515451615200222535ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/jdbc-spotbugs-exclude.xml000066400000000000000000000033761515451615200272030ustar00rootroot00000000000000 apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/pom.xml000066400000000000000000000136731515451615200236020ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-flight 19.0.0 flight-sql-jdbc-core jar Arrow Flight SQL JDBC Driver Core Core implementation of JDBC driver based on Arrow Flight SQL. https://arrow.apache.org org.apache.arrow flight-core io.netty netty-transport-native-kqueue io.netty netty-transport-native-epoll io.grpc grpc-api io.grpc grpc-netty io.netty netty-transport org.apache.arrow arrow-memory-core org.apache.arrow arrow-memory-netty runtime org.apache.arrow arrow-vector ${arrow.vector.classifier} com.google.guava guava org.slf4j slf4j-api runtime com.google.protobuf protobuf-java org.hamcrest hamcrest 3.0 test commons-io commons-io 2.21.0 test org.mockito mockito-core test org.mockito mockito-junit-jupiter test com.squareup.okhttp3 mockwebserver3 5.3.2 test com.squareup.okhttp3 mockwebserver3-junit5 5.3.2 test com.squareup.okhttp3 okhttp-jvm 5.3.2 test com.squareup.okio okio-jvm 3.16.4 test io.netty netty-common org.apache.arrow flight-sql ${project.version} org.apache.calcite.avatica avatica 1.26.0 org.bouncycastle bcpkix-jdk18on 1.83 org.checkerframework checker-qual com.github.ben-manes.caffeine caffeine 3.2.3 com.nimbusds oauth2-oidc-sdk 11.20.1 true src/main/resources maven-surefire-plugin false ${project.basedir}/../../testing/data apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/000077500000000000000000000000001515451615200230425ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/000077500000000000000000000000001515451615200237665ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/000077500000000000000000000000001515451615200247075ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/000077500000000000000000000000001515451615200254765ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/000077500000000000000000000000001515451615200267175ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200300515ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/000077500000000000000000000000001515451615200313445ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/000077500000000000000000000000001515451615200322465ustar00rootroot00000000000000ArrowDatabaseMetadata.java000066400000000000000000001517431515451615200372250ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.sql.Types.BIGINT; import static java.sql.Types.BINARY; import static java.sql.Types.BIT; import static java.sql.Types.CHAR; import static java.sql.Types.DATE; import static java.sql.Types.DECIMAL; import static java.sql.Types.FLOAT; import static java.sql.Types.INTEGER; import static java.sql.Types.LONGNVARCHAR; import static java.sql.Types.LONGVARBINARY; import static java.sql.Types.NUMERIC; import static java.sql.Types.REAL; import static java.sql.Types.SMALLINT; import static java.sql.Types.TIMESTAMP; import static java.sql.Types.TINYINT; import static java.sql.Types.VARCHAR; import static org.apache.arrow.flight.sql.util.SqlInfoOptionsUtils.doesBitmaskTranslateToEnum; import com.google.protobuf.ProtocolMessageEnum; import java.io.ByteArrayInputStream; import java.io.IOException; import java.nio.channels.Channels; import java.nio.charset.Charset; import java.nio.charset.StandardCharsets; import java.sql.Connection; import java.sql.DatabaseMetaData; 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; import java.util.Map; import java.util.Optional; import java.util.concurrent.atomic.AtomicBoolean; import java.util.regex.Pattern; import java.util.stream.Collectors; import org.apache.arrow.driver.jdbc.utils.SqlTypes; import org.apache.arrow.driver.jdbc.utils.VectorSchemaRootTransformer; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.sql.FlightSqlColumnMetadata; import org.apache.arrow.flight.sql.FlightSqlProducer.Schemas; import org.apache.arrow.flight.sql.impl.FlightSql.SqlInfo; import org.apache.arrow.flight.sql.impl.FlightSql.SqlOuterJoinsSupportLevel; import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedElementActions; import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedGroupBy; import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedPositionedCommands; import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedResultSetType; import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedSubqueries; import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedUnions; import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportsConvert; import org.apache.arrow.flight.sql.impl.FlightSql.SqlTransactionIsolationLevel; import org.apache.arrow.flight.sql.impl.FlightSql.SupportedAnsi92SqlGrammarLevel; import org.apache.arrow.flight.sql.impl.FlightSql.SupportedSqlGrammar; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.IntVector; 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]; static final int NO_DECIMAL_DIGITS = 0; private static final int BASE10_RADIX = 10; static final int COLUMN_SIZE_BYTE = (int) Math.ceil((Byte.SIZE - 1) * Math.log(2) / Math.log(10)); static final int COLUMN_SIZE_SHORT = (int) Math.ceil((Short.SIZE - 1) * Math.log(2) / Math.log(10)); static final int COLUMN_SIZE_INT = (int) Math.ceil((Integer.SIZE - 1) * Math.log(2) / Math.log(10)); static final int COLUMN_SIZE_LONG = (int) Math.ceil((Long.SIZE - 1) * Math.log(2) / Math.log(10)); static final int COLUMN_SIZE_VARCHAR_AND_BINARY = 65536; static final int COLUMN_SIZE_DATE = "YYYY-MM-DD".length(); static final int COLUMN_SIZE_TIME = "HH:MM:ss".length(); static final int COLUMN_SIZE_TIME_MILLISECONDS = "HH:MM:ss.SSS".length(); static final int COLUMN_SIZE_TIME_MICROSECONDS = "HH:MM:ss.SSSSSS".length(); static final int COLUMN_SIZE_TIME_NANOSECONDS = "HH:MM:ss.SSSSSSSSS".length(); static final int COLUMN_SIZE_TIMESTAMP_SECONDS = COLUMN_SIZE_DATE + 1 + COLUMN_SIZE_TIME; static final int COLUMN_SIZE_TIMESTAMP_MILLISECONDS = COLUMN_SIZE_DATE + 1 + COLUMN_SIZE_TIME_MILLISECONDS; static final int COLUMN_SIZE_TIMESTAMP_MICROSECONDS = COLUMN_SIZE_DATE + 1 + COLUMN_SIZE_TIME_MICROSECONDS; static final int COLUMN_SIZE_TIMESTAMP_NANOSECONDS = COLUMN_SIZE_DATE + 1 + COLUMN_SIZE_TIME_NANOSECONDS; static final int DECIMAL_DIGITS_TIME_MILLISECONDS = 3; static final int DECIMAL_DIGITS_TIME_MICROSECONDS = 6; static final int DECIMAL_DIGITS_TIME_NANOSECONDS = 9; private static final Schema GET_COLUMNS_SCHEMA = new Schema( Arrays.asList( Field.nullable("TABLE_CAT", Types.MinorType.VARCHAR.getType()), Field.nullable("TABLE_SCHEM", Types.MinorType.VARCHAR.getType()), Field.notNullable("TABLE_NAME", Types.MinorType.VARCHAR.getType()), Field.notNullable("COLUMN_NAME", Types.MinorType.VARCHAR.getType()), Field.nullable("DATA_TYPE", Types.MinorType.INT.getType()), Field.nullable("TYPE_NAME", Types.MinorType.VARCHAR.getType()), Field.nullable("COLUMN_SIZE", Types.MinorType.INT.getType()), Field.nullable("BUFFER_LENGTH", Types.MinorType.INT.getType()), Field.nullable("DECIMAL_DIGITS", Types.MinorType.INT.getType()), Field.nullable("NUM_PREC_RADIX", Types.MinorType.INT.getType()), Field.notNullable("NULLABLE", Types.MinorType.INT.getType()), Field.nullable("REMARKS", Types.MinorType.VARCHAR.getType()), Field.nullable("COLUMN_DEF", Types.MinorType.VARCHAR.getType()), Field.nullable("SQL_DATA_TYPE", Types.MinorType.INT.getType()), Field.nullable("SQL_DATETIME_SUB", Types.MinorType.INT.getType()), Field.notNullable("CHAR_OCTET_LENGTH", Types.MinorType.INT.getType()), Field.notNullable("ORDINAL_POSITION", Types.MinorType.INT.getType()), Field.notNullable("IS_NULLABLE", Types.MinorType.VARCHAR.getType()), Field.nullable("SCOPE_CATALOG", Types.MinorType.VARCHAR.getType()), Field.nullable("SCOPE_SCHEMA", Types.MinorType.VARCHAR.getType()), Field.nullable("SCOPE_TABLE", Types.MinorType.VARCHAR.getType()), Field.nullable("SOURCE_DATA_TYPE", Types.MinorType.SMALLINT.getType()), Field.notNullable("IS_AUTOINCREMENT", Types.MinorType.VARCHAR.getType()), Field.notNullable("IS_GENERATEDCOLUMN", Types.MinorType.VARCHAR.getType()))); private final AtomicBoolean isCachePopulated = new AtomicBoolean(false); private final Map cachedSqlInfo = new EnumMap<>(SqlInfo.class); private static final Map sqlTypesToFlightEnumConvertTypes = new HashMap<>(); static { sqlTypesToFlightEnumConvertTypes.put(BIT, SqlSupportsConvert.SQL_CONVERT_BIT_VALUE); sqlTypesToFlightEnumConvertTypes.put(INTEGER, SqlSupportsConvert.SQL_CONVERT_INTEGER_VALUE); sqlTypesToFlightEnumConvertTypes.put(NUMERIC, SqlSupportsConvert.SQL_CONVERT_NUMERIC_VALUE); sqlTypesToFlightEnumConvertTypes.put(SMALLINT, SqlSupportsConvert.SQL_CONVERT_SMALLINT_VALUE); sqlTypesToFlightEnumConvertTypes.put(TINYINT, SqlSupportsConvert.SQL_CONVERT_TINYINT_VALUE); sqlTypesToFlightEnumConvertTypes.put(FLOAT, SqlSupportsConvert.SQL_CONVERT_FLOAT_VALUE); sqlTypesToFlightEnumConvertTypes.put(BIGINT, SqlSupportsConvert.SQL_CONVERT_BIGINT_VALUE); sqlTypesToFlightEnumConvertTypes.put(REAL, SqlSupportsConvert.SQL_CONVERT_REAL_VALUE); sqlTypesToFlightEnumConvertTypes.put(DECIMAL, SqlSupportsConvert.SQL_CONVERT_DECIMAL_VALUE); sqlTypesToFlightEnumConvertTypes.put(BINARY, SqlSupportsConvert.SQL_CONVERT_BINARY_VALUE); sqlTypesToFlightEnumConvertTypes.put( LONGVARBINARY, SqlSupportsConvert.SQL_CONVERT_LONGVARBINARY_VALUE); sqlTypesToFlightEnumConvertTypes.put(CHAR, SqlSupportsConvert.SQL_CONVERT_CHAR_VALUE); sqlTypesToFlightEnumConvertTypes.put(VARCHAR, SqlSupportsConvert.SQL_CONVERT_VARCHAR_VALUE); sqlTypesToFlightEnumConvertTypes.put( 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) { super(connection); } @Override public String getDatabaseProductName() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.FLIGHT_SQL_SERVER_NAME, String.class); } @Override public String getDatabaseProductVersion() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.FLIGHT_SQL_SERVER_VERSION, String.class); } @Override public String getIdentifierQuoteString() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_IDENTIFIER_QUOTE_CHAR, String.class); } @Override public boolean isReadOnly() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.FLIGHT_SQL_SERVER_READ_ONLY, Boolean.class); } @Override public String getSQLKeywords() throws SQLException { return convertListSqlInfoToString( getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_KEYWORDS, List.class)) .orElse(""); } @Override public String getNumericFunctions() throws SQLException { return convertListSqlInfoToString( getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_NUMERIC_FUNCTIONS, List.class)) .orElse(""); } @Override public String getStringFunctions() throws SQLException { return convertListSqlInfoToString( getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_STRING_FUNCTIONS, List.class)) .orElse(""); } @Override public String getSystemFunctions() throws SQLException { return convertListSqlInfoToString( getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_SYSTEM_FUNCTIONS, List.class)) .orElse(""); } @Override public String getTimeDateFunctions() throws SQLException { return convertListSqlInfoToString( getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_DATETIME_FUNCTIONS, List.class)) .orElse(""); } @Override public String getSearchStringEscape() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_SEARCH_STRING_ESCAPE, String.class); } @Override public String getExtraNameCharacters() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_EXTRA_NAME_CHARACTERS, String.class); } @Override public boolean supportsColumnAliasing() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_SUPPORTS_COLUMN_ALIASING, Boolean.class); } @Override public boolean nullPlusNonNullIsNull() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_NULL_PLUS_NULL_IS_NULL, Boolean.class); } @Override public boolean supportsConvert() throws SQLException { return !getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_SUPPORTS_CONVERT, Map.class).isEmpty(); } @Override public boolean supportsConvert(final int fromType, final int toType) throws SQLException { final Map> sqlSupportsConvert = getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_SUPPORTS_CONVERT, Map.class); if (!sqlTypesToFlightEnumConvertTypes.containsKey(fromType)) { return false; } final List list = sqlSupportsConvert.get(sqlTypesToFlightEnumConvertTypes.get(fromType)); return list != null && list.contains(sqlTypesToFlightEnumConvertTypes.get(toType)); } @Override public boolean supportsTableCorrelationNames() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTS_TABLE_CORRELATION_NAMES, Boolean.class); } @Override public boolean supportsDifferentTableCorrelationNames() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTS_DIFFERENT_TABLE_CORRELATION_NAMES, Boolean.class); } @Override public boolean supportsExpressionsInOrderBy() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTS_EXPRESSIONS_IN_ORDER_BY, Boolean.class); } @Override public boolean supportsOrderByUnrelated() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_SUPPORTS_ORDER_BY_UNRELATED, Boolean.class); } @Override public boolean supportsGroupBy() throws SQLException { final int bitmask = getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_SUPPORTED_GROUP_BY, Integer.class); return bitmask != 0; } @Override public boolean supportsGroupByUnrelated() throws SQLException { return getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTED_GROUP_BY, SqlSupportedGroupBy.SQL_GROUP_BY_UNRELATED); } @Override public boolean supportsLikeEscapeClause() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_SUPPORTS_LIKE_ESCAPE_CLAUSE, Boolean.class); } @Override public boolean supportsNonNullableColumns() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTS_NON_NULLABLE_COLUMNS, Boolean.class); } @Override public boolean supportsMinimumSQLGrammar() throws SQLException { return checkEnumLevel( Arrays.asList( getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTED_GRAMMAR, SupportedSqlGrammar.SQL_EXTENDED_GRAMMAR), getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTED_GRAMMAR, SupportedSqlGrammar.SQL_CORE_GRAMMAR), getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTED_GRAMMAR, SupportedSqlGrammar.SQL_MINIMUM_GRAMMAR))); } @Override public boolean supportsCoreSQLGrammar() throws SQLException { return checkEnumLevel( Arrays.asList( getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTED_GRAMMAR, SupportedSqlGrammar.SQL_EXTENDED_GRAMMAR), getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTED_GRAMMAR, SupportedSqlGrammar.SQL_CORE_GRAMMAR))); } @Override public boolean supportsExtendedSQLGrammar() throws SQLException { return getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTED_GRAMMAR, SupportedSqlGrammar.SQL_EXTENDED_GRAMMAR); } @Override public boolean supportsANSI92EntryLevelSQL() throws SQLException { return checkEnumLevel( Arrays.asList( getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_ANSI92_SUPPORTED_LEVEL, SupportedAnsi92SqlGrammarLevel.ANSI92_ENTRY_SQL), getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_ANSI92_SUPPORTED_LEVEL, SupportedAnsi92SqlGrammarLevel.ANSI92_INTERMEDIATE_SQL), getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_ANSI92_SUPPORTED_LEVEL, SupportedAnsi92SqlGrammarLevel.ANSI92_FULL_SQL))); } @Override public boolean supportsANSI92IntermediateSQL() throws SQLException { return checkEnumLevel( Arrays.asList( getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_ANSI92_SUPPORTED_LEVEL, SupportedAnsi92SqlGrammarLevel.ANSI92_ENTRY_SQL), getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_ANSI92_SUPPORTED_LEVEL, SupportedAnsi92SqlGrammarLevel.ANSI92_INTERMEDIATE_SQL))); } @Override public boolean supportsANSI92FullSQL() throws SQLException { return getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_ANSI92_SUPPORTED_LEVEL, SupportedAnsi92SqlGrammarLevel.ANSI92_FULL_SQL); } @Override public boolean supportsIntegrityEnhancementFacility() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTS_INTEGRITY_ENHANCEMENT_FACILITY, Boolean.class); } @Override public boolean supportsOuterJoins() throws SQLException { final int bitmask = getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_OUTER_JOINS_SUPPORT_LEVEL, Integer.class); return bitmask != 0; } @Override public boolean supportsFullOuterJoins() throws SQLException { return getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_OUTER_JOINS_SUPPORT_LEVEL, SqlOuterJoinsSupportLevel.SQL_FULL_OUTER_JOINS); } @Override public boolean supportsLimitedOuterJoins() throws SQLException { return getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_OUTER_JOINS_SUPPORT_LEVEL, SqlOuterJoinsSupportLevel.SQL_LIMITED_OUTER_JOINS); } @Override public String getSchemaTerm() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_SCHEMA_TERM, String.class); } @Override public String getProcedureTerm() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_PROCEDURE_TERM, String.class); } @Override public String getCatalogTerm() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_CATALOG_TERM, String.class); } @Override public boolean isCatalogAtStart() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_CATALOG_AT_START, Boolean.class); } @Override public boolean supportsSchemasInProcedureCalls() throws SQLException { return getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_SCHEMAS_SUPPORTED_ACTIONS, SqlSupportedElementActions.SQL_ELEMENT_IN_PROCEDURE_CALLS); } @Override public boolean supportsSchemasInIndexDefinitions() throws SQLException { return getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_SCHEMAS_SUPPORTED_ACTIONS, SqlSupportedElementActions.SQL_ELEMENT_IN_INDEX_DEFINITIONS); } @Override public boolean supportsSchemasInPrivilegeDefinitions() throws SQLException { return getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_SCHEMAS_SUPPORTED_ACTIONS, SqlSupportedElementActions.SQL_ELEMENT_IN_PRIVILEGE_DEFINITIONS); } @Override public boolean supportsCatalogsInIndexDefinitions() throws SQLException { return getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_CATALOGS_SUPPORTED_ACTIONS, SqlSupportedElementActions.SQL_ELEMENT_IN_INDEX_DEFINITIONS); } @Override public boolean supportsCatalogsInPrivilegeDefinitions() throws SQLException { return getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_CATALOGS_SUPPORTED_ACTIONS, SqlSupportedElementActions.SQL_ELEMENT_IN_PRIVILEGE_DEFINITIONS); } @Override public boolean supportsPositionedDelete() throws SQLException { return getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTED_POSITIONED_COMMANDS, SqlSupportedPositionedCommands.SQL_POSITIONED_DELETE); } @Override public boolean supportsPositionedUpdate() throws SQLException { return getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTED_POSITIONED_COMMANDS, SqlSupportedPositionedCommands.SQL_POSITIONED_UPDATE); } @Override public boolean supportsResultSetType(final int type) throws SQLException { final int bitmask = getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_SUPPORTED_RESULT_SET_TYPES, Integer.class); switch (type) { case ResultSet.TYPE_FORWARD_ONLY: return doesBitmaskTranslateToEnum( SqlSupportedResultSetType.SQL_RESULT_SET_TYPE_FORWARD_ONLY, bitmask); case ResultSet.TYPE_SCROLL_INSENSITIVE: return doesBitmaskTranslateToEnum( SqlSupportedResultSetType.SQL_RESULT_SET_TYPE_SCROLL_INSENSITIVE, bitmask); case ResultSet.TYPE_SCROLL_SENSITIVE: return doesBitmaskTranslateToEnum( SqlSupportedResultSetType.SQL_RESULT_SET_TYPE_SCROLL_SENSITIVE, bitmask); default: throw new SQLException( "Invalid result set type argument. The informed type is not defined in java.sql.ResultSet."); } } @Override public boolean supportsSelectForUpdate() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_SELECT_FOR_UPDATE_SUPPORTED, Boolean.class); } @Override public boolean supportsStoredProcedures() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_STORED_PROCEDURES_SUPPORTED, Boolean.class); } @Override public boolean supportsSubqueriesInComparisons() throws SQLException { return getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTED_SUBQUERIES, SqlSupportedSubqueries.SQL_SUBQUERIES_IN_COMPARISONS); } @Override public boolean supportsSubqueriesInExists() throws SQLException { return getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTED_SUBQUERIES, SqlSupportedSubqueries.SQL_SUBQUERIES_IN_EXISTS); } @Override public boolean supportsSubqueriesInIns() throws SQLException { return getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTED_SUBQUERIES, SqlSupportedSubqueries.SQL_SUBQUERIES_IN_INS); } @Override public boolean supportsSubqueriesInQuantifieds() throws SQLException { return getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTED_SUBQUERIES, SqlSupportedSubqueries.SQL_SUBQUERIES_IN_QUANTIFIEDS); } @Override public boolean supportsCorrelatedSubqueries() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty( SqlInfo.SQL_CORRELATED_SUBQUERIES_SUPPORTED, Boolean.class); } @Override public boolean supportsUnion() throws SQLException { final int bitmask = getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_SUPPORTED_UNIONS, Integer.class); return bitmask != 0; } @Override public boolean supportsUnionAll() throws SQLException { return getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTED_UNIONS, SqlSupportedUnions.SQL_UNION_ALL); } @Override public int getMaxBinaryLiteralLength() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_BINARY_LITERAL_LENGTH, Long.class) .intValue(); } @Override public int getMaxCharLiteralLength() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_CHAR_LITERAL_LENGTH, Long.class) .intValue(); } @Override public int getMaxColumnNameLength() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_COLUMN_NAME_LENGTH, Long.class) .intValue(); } @Override public int getMaxColumnsInGroupBy() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_COLUMNS_IN_GROUP_BY, Long.class) .intValue(); } @Override public int getMaxColumnsInIndex() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_COLUMNS_IN_INDEX, Long.class) .intValue(); } @Override public int getMaxColumnsInOrderBy() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_COLUMNS_IN_ORDER_BY, Long.class) .intValue(); } @Override public int getMaxColumnsInSelect() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_COLUMNS_IN_SELECT, Long.class) .intValue(); } @Override public int getMaxColumnsInTable() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_COLUMNS_IN_TABLE, Long.class) .intValue(); } @Override public int getMaxConnections() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_CONNECTIONS, Long.class).intValue(); } @Override public int getMaxCursorNameLength() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_CURSOR_NAME_LENGTH, Long.class) .intValue(); } @Override public int getMaxIndexLength() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_INDEX_LENGTH, Long.class).intValue(); } @Override public int getMaxSchemaNameLength() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_DB_SCHEMA_NAME_LENGTH, Long.class) .intValue(); } @Override public int getMaxProcedureNameLength() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_PROCEDURE_NAME_LENGTH, Long.class) .intValue(); } @Override public int getMaxCatalogNameLength() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_CATALOG_NAME_LENGTH, Long.class) .intValue(); } @Override public int getMaxRowSize() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_ROW_SIZE, Long.class).intValue(); } @Override public boolean doesMaxRowSizeIncludeBlobs() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_ROW_SIZE_INCLUDES_BLOBS, Boolean.class); } @Override public int getMaxStatementLength() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_STATEMENT_LENGTH, Long.class) .intValue(); } @Override public int getMaxStatements() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_STATEMENTS, Long.class).intValue(); } @Override public int getMaxTableNameLength() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_TABLE_NAME_LENGTH, Long.class) .intValue(); } @Override public int getMaxTablesInSelect() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_TABLES_IN_SELECT, Long.class) .intValue(); } @Override public int getMaxUserNameLength() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_MAX_USERNAME_LENGTH, Long.class).intValue(); } @Override public int getDefaultTransactionIsolation() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_DEFAULT_TRANSACTION_ISOLATION, Long.class) .intValue(); } @Override public boolean supportsTransactions() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_TRANSACTIONS_SUPPORTED, Boolean.class); } @Override public boolean supportsTransactionIsolationLevel(final int level) throws SQLException { final int bitmask = getSqlInfoAndCacheIfCacheIsEmpty( SqlInfo.SQL_SUPPORTED_TRANSACTIONS_ISOLATION_LEVELS, Integer.class); switch (level) { case Connection.TRANSACTION_NONE: return doesBitmaskTranslateToEnum( SqlTransactionIsolationLevel.SQL_TRANSACTION_NONE, bitmask); case Connection.TRANSACTION_READ_COMMITTED: return doesBitmaskTranslateToEnum( SqlTransactionIsolationLevel.SQL_TRANSACTION_READ_COMMITTED, bitmask); case Connection.TRANSACTION_READ_UNCOMMITTED: return doesBitmaskTranslateToEnum( SqlTransactionIsolationLevel.SQL_TRANSACTION_READ_UNCOMMITTED, bitmask); case Connection.TRANSACTION_REPEATABLE_READ: return doesBitmaskTranslateToEnum( SqlTransactionIsolationLevel.SQL_TRANSACTION_REPEATABLE_READ, bitmask); case Connection.TRANSACTION_SERIALIZABLE: return doesBitmaskTranslateToEnum( SqlTransactionIsolationLevel.SQL_TRANSACTION_SERIALIZABLE, bitmask); default: throw new SQLException( "Invalid transaction isolation level argument. The informed level is not defined in java.sql.Connection."); } } @Override public boolean dataDefinitionCausesTransactionCommit() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty( SqlInfo.SQL_DATA_DEFINITION_CAUSES_TRANSACTION_COMMIT, Boolean.class); } @Override public boolean dataDefinitionIgnoredInTransactions() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty( SqlInfo.SQL_DATA_DEFINITIONS_IN_TRANSACTIONS_IGNORED, Boolean.class); } @Override public boolean supportsBatchUpdates() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_BATCH_UPDATES_SUPPORTED, Boolean.class); } @Override public boolean supportsSavepoints() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_SAVEPOINTS_SUPPORTED, Boolean.class); } @Override public boolean supportsNamedParameters() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_NAMED_PARAMETERS_SUPPORTED, Boolean.class); } @Override public boolean locatorsUpdateCopy() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty(SqlInfo.SQL_LOCATORS_UPDATE_COPY, Boolean.class); } @Override public boolean supportsStoredFunctionsUsingCallSyntax() throws SQLException { return getSqlInfoAndCacheIfCacheIsEmpty( SqlInfo.SQL_STORED_FUNCTIONS_USING_CALL_SYNTAX_SUPPORTED, Boolean.class); } @Override public ArrowFlightConnection getConnection() throws SQLException { return (ArrowFlightConnection) super.getConnection(); } private T getSqlInfoAndCacheIfCacheIsEmpty( final SqlInfo sqlInfoCommand, final Class desiredType) throws SQLException { final ArrowFlightConnection connection = getConnection(); if (!isCachePopulated.get()) { // Lock-and-populate the cache. Only issue the call to getSqlInfo() once, // populate the cache, then mark it as populated. // Note that multiple callers from separate threads can see that the cache is not populated, // but only // one thread will try to populate the cache. Other threads will see the cache is non-empty // when acquiring // the lock on the cache and skip population. synchronized (cachedSqlInfo) { if (cachedSqlInfo.isEmpty()) { final FlightInfo sqlInfo = connection.getClientHandler().getSqlInfo(); try (final ResultSet resultSet = ArrowFlightJdbcFlightStreamResultSet.fromFlightInfo(connection, sqlInfo, null)) { while (resultSet.next()) { cachedSqlInfo.put( SqlInfo.forNumber((Integer) resultSet.getObject("info_name")), resultSet.getObject("value")); } } isCachePopulated.set(true); } } } 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) { if (sqlInfoList == null) { return Optional.empty(); } else { return Optional.of( sqlInfoList.stream().map(Object::toString).collect(Collectors.joining(", "))); } } private boolean getSqlInfoEnumOptionAndCacheIfCacheIsEmpty( final SqlInfo sqlInfoCommand, final ProtocolMessageEnum enumInstance) throws SQLException { final int bitmask = getSqlInfoAndCacheIfCacheIsEmpty(sqlInfoCommand, Integer.class); return doesBitmaskTranslateToEnum(enumInstance, bitmask); } private boolean checkEnumLevel(final List toCheck) { return toCheck.stream().anyMatch(e -> e); } @Override public ResultSet getCatalogs() throws SQLException { final ArrowFlightConnection connection = getConnection(); final FlightInfo flightInfoCatalogs = connection.getClientHandler().getCatalogs(); final BufferAllocator allocator = connection.getBufferAllocator(); final VectorSchemaRootTransformer transformer = new VectorSchemaRootTransformer.Builder(Schemas.GET_CATALOGS_SCHEMA, allocator) .renameFieldVector("catalog_name", "TABLE_CAT") .build(); return ArrowFlightJdbcFlightStreamResultSet.fromFlightInfo( connection, flightInfoCatalogs, transformer); } @Override public ResultSet getImportedKeys(final String catalog, final String schema, final String table) throws SQLException { final ArrowFlightConnection connection = getConnection(); final FlightInfo flightInfoImportedKeys = connection.getClientHandler().getImportedKeys(catalog, schema, table); final BufferAllocator allocator = connection.getBufferAllocator(); final VectorSchemaRootTransformer transformer = getForeignKeysTransformer(allocator); return ArrowFlightJdbcFlightStreamResultSet.fromFlightInfo( connection, flightInfoImportedKeys, transformer); } @Override public ResultSet getExportedKeys(final String catalog, final String schema, final String table) throws SQLException { final ArrowFlightConnection connection = getConnection(); final FlightInfo flightInfoExportedKeys = connection.getClientHandler().getExportedKeys(catalog, schema, table); final BufferAllocator allocator = connection.getBufferAllocator(); final VectorSchemaRootTransformer transformer = getForeignKeysTransformer(allocator); return ArrowFlightJdbcFlightStreamResultSet.fromFlightInfo( connection, flightInfoExportedKeys, transformer); } @Override public ResultSet getCrossReference( final String parentCatalog, final String parentSchema, final String parentTable, final String foreignCatalog, final String foreignSchema, final String foreignTable) throws SQLException { final ArrowFlightConnection connection = getConnection(); final FlightInfo flightInfoCrossReference = connection .getClientHandler() .getCrossReference( parentCatalog, parentSchema, parentTable, foreignCatalog, foreignSchema, foreignTable); final BufferAllocator allocator = connection.getBufferAllocator(); final VectorSchemaRootTransformer transformer = getForeignKeysTransformer(allocator); return ArrowFlightJdbcFlightStreamResultSet.fromFlightInfo( connection, flightInfoCrossReference, transformer); } /** * Transformer used on getImportedKeys, getExportedKeys and getCrossReference methods, since all * three share the same schema. */ private VectorSchemaRootTransformer getForeignKeysTransformer(final BufferAllocator allocator) { return new VectorSchemaRootTransformer.Builder(Schemas.GET_IMPORTED_KEYS_SCHEMA, allocator) .renameFieldVector("pk_catalog_name", "PKTABLE_CAT") .renameFieldVector("pk_db_schema_name", "PKTABLE_SCHEM") .renameFieldVector("pk_table_name", "PKTABLE_NAME") .renameFieldVector("pk_column_name", "PKCOLUMN_NAME") .renameFieldVector("fk_catalog_name", "FKTABLE_CAT") .renameFieldVector("fk_db_schema_name", "FKTABLE_SCHEM") .renameFieldVector("fk_table_name", "FKTABLE_NAME") .renameFieldVector("fk_column_name", "FKCOLUMN_NAME") .renameFieldVector("key_sequence", "KEY_SEQ") .renameFieldVector("fk_key_name", "FK_NAME") .renameFieldVector("pk_key_name", "PK_NAME") .renameFieldVector("update_rule", "UPDATE_RULE") .renameFieldVector("delete_rule", "DELETE_RULE") .addEmptyField("DEFERRABILITY", new ArrowType.Int(Byte.SIZE, false)) .build(); } @Override public ResultSet getSchemas(final String catalog, final String schemaPattern) throws SQLException { final ArrowFlightConnection connection = getConnection(); final FlightInfo flightInfoSchemas = connection.getClientHandler().getSchemas(catalog, schemaPattern); final BufferAllocator allocator = connection.getBufferAllocator(); final VectorSchemaRootTransformer transformer = new VectorSchemaRootTransformer.Builder(Schemas.GET_SCHEMAS_SCHEMA, allocator) .renameFieldVector("db_schema_name", "TABLE_SCHEM") .renameFieldVector("catalog_name", "TABLE_CATALOG") .build(); return ArrowFlightJdbcFlightStreamResultSet.fromFlightInfo( connection, flightInfoSchemas, transformer); } @Override public ResultSet getTableTypes() throws SQLException { final ArrowFlightConnection connection = getConnection(); final FlightInfo flightInfoTableTypes = connection.getClientHandler().getTableTypes(); final BufferAllocator allocator = connection.getBufferAllocator(); final VectorSchemaRootTransformer transformer = new VectorSchemaRootTransformer.Builder(Schemas.GET_TABLE_TYPES_SCHEMA, allocator) .renameFieldVector("table_type", "TABLE_TYPE") .build(); return ArrowFlightJdbcFlightStreamResultSet.fromFlightInfo( connection, flightInfoTableTypes, transformer); } @Override public ResultSet getTables( final String catalog, final String schemaPattern, final String tableNamePattern, final String[] types) throws SQLException { final ArrowFlightConnection connection = getConnection(); final List typesList = types == null ? null : Arrays.asList(types); final FlightInfo flightInfoTables = connection .getClientHandler() .getTables(catalog, schemaPattern, tableNamePattern, typesList, false); final BufferAllocator allocator = connection.getBufferAllocator(); final VectorSchemaRootTransformer transformer = new VectorSchemaRootTransformer.Builder(Schemas.GET_TABLES_SCHEMA_NO_SCHEMA, allocator) .renameFieldVector("catalog_name", "TABLE_CAT") .renameFieldVector("db_schema_name", "TABLE_SCHEM") .renameFieldVector("table_name", "TABLE_NAME") .renameFieldVector("table_type", "TABLE_TYPE") .addEmptyField("REMARKS", Types.MinorType.VARBINARY) .addEmptyField("TYPE_CAT", Types.MinorType.VARBINARY) .addEmptyField("TYPE_SCHEM", Types.MinorType.VARBINARY) .addEmptyField("TYPE_NAME", Types.MinorType.VARBINARY) .addEmptyField("SELF_REFERENCING_COL_NAME", Types.MinorType.VARBINARY) .addEmptyField("REF_GENERATION", Types.MinorType.VARBINARY) .build(); return ArrowFlightJdbcFlightStreamResultSet.fromFlightInfo( connection, flightInfoTables, transformer); } @Override public ResultSet getPrimaryKeys(final String catalog, final String schema, final String table) throws SQLException { final ArrowFlightConnection connection = getConnection(); final FlightInfo flightInfoPrimaryKeys = connection.getClientHandler().getPrimaryKeys(catalog, schema, table); final BufferAllocator allocator = connection.getBufferAllocator(); final VectorSchemaRootTransformer transformer = new VectorSchemaRootTransformer.Builder(Schemas.GET_PRIMARY_KEYS_SCHEMA, allocator) .renameFieldVector("catalog_name", "TABLE_CAT") .renameFieldVector("db_schema_name", "TABLE_SCHEM") .renameFieldVector("table_name", "TABLE_NAME") .renameFieldVector("column_name", "COLUMN_NAME") .renameFieldVector("key_sequence", "KEY_SEQ") .renameFieldVector("key_name", "PK_NAME") .build(); return ArrowFlightJdbcFlightStreamResultSet.fromFlightInfo( connection, flightInfoPrimaryKeys, transformer); } @Override public ResultSet getColumns( final String catalog, final String schemaPattern, final String tableNamePattern, final String columnNamePattern) throws SQLException { final ArrowFlightConnection connection = getConnection(); final FlightInfo flightInfoTables = connection .getClientHandler() .getTables(catalog, schemaPattern, tableNamePattern, null, true); final BufferAllocator allocator = connection.getBufferAllocator(); final Pattern columnNamePat = columnNamePattern != null ? Pattern.compile(sqlToRegexLike(columnNamePattern)) : null; return ArrowFlightJdbcFlightStreamResultSet.fromFlightInfo( connection, flightInfoTables, (originalRoot, transformedRoot) -> { int columnCounter = 0; if (transformedRoot == null) { transformedRoot = VectorSchemaRoot.create(GET_COLUMNS_SCHEMA, allocator); } final int originalRootRowCount = originalRoot.getRowCount(); final VarCharVector catalogNameVector = (VarCharVector) originalRoot.getVector("catalog_name"); final VarCharVector tableNameVector = (VarCharVector) originalRoot.getVector("table_name"); final VarCharVector schemaNameVector = (VarCharVector) originalRoot.getVector("db_schema_name"); final VarBinaryVector schemaVector = (VarBinaryVector) originalRoot.getVector("table_schema"); for (int i = 0; i < originalRootRowCount; i++) { final Text catalogName = catalogNameVector.getObject(i); final Text tableName = tableNameVector.getObject(i); final Text schemaName = schemaNameVector.getObject(i); final Schema currentSchema; try { currentSchema = MessageSerializer.deserializeSchema( new ReadChannel( Channels.newChannel(new ByteArrayInputStream(schemaVector.get(i))))); } catch (final IOException e) { throw new IOException( String.format("Failed to deserialize schema for table %s", tableName), e); } final List tableColumns = currentSchema.getFields(); columnCounter = setGetColumnsVectorSchemaRootFromFields( transformedRoot, columnCounter, tableColumns, catalogName, tableName, schemaName, columnNamePat); } transformedRoot.setRowCount(columnCounter); originalRoot.clear(); return transformedRoot; }); } private int setGetColumnsVectorSchemaRootFromFields( final VectorSchemaRoot currentRoot, int insertIndex, final List tableColumns, final Text catalogName, final Text tableName, final Text schemaName, final Pattern columnNamePattern) { int ordinalIndex = 1; final int tableColumnsSize = tableColumns.size(); final VarCharVector tableCatVector = (VarCharVector) currentRoot.getVector("TABLE_CAT"); final VarCharVector tableSchemVector = (VarCharVector) currentRoot.getVector("TABLE_SCHEM"); final VarCharVector tableNameVector = (VarCharVector) currentRoot.getVector("TABLE_NAME"); final VarCharVector columnNameVector = (VarCharVector) currentRoot.getVector("COLUMN_NAME"); final IntVector dataTypeVector = (IntVector) currentRoot.getVector("DATA_TYPE"); final VarCharVector typeNameVector = (VarCharVector) currentRoot.getVector("TYPE_NAME"); final IntVector columnSizeVector = (IntVector) currentRoot.getVector("COLUMN_SIZE"); final IntVector decimalDigitsVector = (IntVector) currentRoot.getVector("DECIMAL_DIGITS"); final IntVector numPrecRadixVector = (IntVector) currentRoot.getVector("NUM_PREC_RADIX"); final IntVector nullableVector = (IntVector) currentRoot.getVector("NULLABLE"); final IntVector ordinalPositionVector = (IntVector) currentRoot.getVector("ORDINAL_POSITION"); final VarCharVector isNullableVector = (VarCharVector) currentRoot.getVector("IS_NULLABLE"); final VarCharVector isAutoincrementVector = (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); final FlightSqlColumnMetadata columnMetadata = new FlightSqlColumnMetadata(field.getMetadata()); final String columnName = field.getName(); if (columnNamePattern != null && !columnNamePattern.matcher(columnName).matches()) { continue; } final ArrowType fieldType = field.getType(); if (catalogName != null) { tableCatVector.setSafe(insertIndex, catalogName); } if (schemaName != null) { tableSchemVector.setSafe(insertIndex, schemaName); } if (tableName != null) { tableNameVector.setSafe(insertIndex, tableName); } if (columnName != null) { columnNameVector.setSafe(insertIndex, columnName.getBytes(CHARSET)); } dataTypeVector.setSafe(insertIndex, SqlTypes.getSqlTypeIdFromArrowType(fieldType)); byte[] typeName = columnMetadata.getTypeName() != null ? columnMetadata.getTypeName().getBytes(CHARSET) : SqlTypes.getSqlTypeNameFromArrowType(fieldType).getBytes(CHARSET); typeNameVector.setSafe(insertIndex, typeName); // We aren't setting COLUMN_SIZE for ROWID SQL Types, as there's no such Arrow type. // We aren't setting COLUMN_SIZE nor DECIMAL_DIGITS for Float/Double as their precision and // scale are variable. if (fieldType instanceof ArrowType.Decimal) { numPrecRadixVector.setSafe(insertIndex, BASE10_RADIX); } else if (fieldType instanceof ArrowType.Int) { numPrecRadixVector.setSafe(insertIndex, BASE10_RADIX); } else if (fieldType instanceof ArrowType.FloatingPoint) { numPrecRadixVector.setSafe(insertIndex, BASE10_RADIX); } Integer decimalDigits = columnMetadata.getScale(); if (decimalDigits == null) { decimalDigits = getDecimalDigits(fieldType); } if (decimalDigits != null) { decimalDigitsVector.setSafe(insertIndex, decimalDigits); } Integer columnSize = columnMetadata.getPrecision(); if (columnSize == null) { columnSize = getColumnSize(fieldType); } if (columnSize != null) { columnSizeVector.setSafe(insertIndex, columnSize); } nullableVector.setSafe(insertIndex, field.isNullable() ? 1 : 0); isNullableVector.setSafe(insertIndex, booleanToYesOrNo(field.isNullable())); Boolean autoIncrement = columnMetadata.isAutoIncrement(); if (autoIncrement != null) { isAutoincrementVector.setSafe(insertIndex, booleanToYesOrNo(autoIncrement)); } else { 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); ordinalPositionVector.setSafe(insertIndex, ordinalIndex); insertIndex++; } return insertIndex; } private static byte[] booleanToYesOrNo(boolean autoIncrement) { return autoIncrement ? "YES".getBytes(CHARSET) : "NO".getBytes(CHARSET); } static Integer getDecimalDigits(final ArrowType fieldType) { // We aren't setting DECIMAL_DIGITS for Float/Double as their precision and scale are variable. if (fieldType instanceof ArrowType.Decimal) { final ArrowType.Decimal thisDecimal = (ArrowType.Decimal) fieldType; return thisDecimal.getScale(); } else if (fieldType instanceof ArrowType.Int) { return NO_DECIMAL_DIGITS; } else if (fieldType instanceof ArrowType.Timestamp) { switch (((ArrowType.Timestamp) fieldType).getUnit()) { case SECOND: return NO_DECIMAL_DIGITS; case MILLISECOND: return DECIMAL_DIGITS_TIME_MILLISECONDS; case MICROSECOND: return DECIMAL_DIGITS_TIME_MICROSECONDS; case NANOSECOND: return DECIMAL_DIGITS_TIME_NANOSECONDS; default: break; } } else if (fieldType instanceof ArrowType.Time) { switch (((ArrowType.Time) fieldType).getUnit()) { case SECOND: return NO_DECIMAL_DIGITS; case MILLISECOND: return DECIMAL_DIGITS_TIME_MILLISECONDS; case MICROSECOND: return DECIMAL_DIGITS_TIME_MICROSECONDS; case NANOSECOND: return DECIMAL_DIGITS_TIME_NANOSECONDS; default: break; } } else if (fieldType instanceof ArrowType.Date) { return NO_DECIMAL_DIGITS; } return null; } static Integer getColumnSize(final ArrowType fieldType) { // We aren't setting COLUMN_SIZE for ROWID SQL Types, as there's no such Arrow type. // We aren't setting COLUMN_SIZE nor DECIMAL_DIGITS for Float/Double as their precision and // scale are variable. if (fieldType instanceof ArrowType.Decimal) { final ArrowType.Decimal thisDecimal = (ArrowType.Decimal) fieldType; return thisDecimal.getPrecision(); } else if (fieldType instanceof ArrowType.Int) { final ArrowType.Int thisInt = (ArrowType.Int) fieldType; switch (thisInt.getBitWidth()) { case Byte.SIZE: return COLUMN_SIZE_BYTE; case Short.SIZE: return COLUMN_SIZE_SHORT; case Integer.SIZE: return COLUMN_SIZE_INT; case Long.SIZE: return COLUMN_SIZE_LONG; default: break; } } else if (fieldType instanceof ArrowType.Utf8 || fieldType instanceof ArrowType.Binary) { return COLUMN_SIZE_VARCHAR_AND_BINARY; } else if (fieldType instanceof ArrowType.Timestamp) { switch (((ArrowType.Timestamp) fieldType).getUnit()) { case SECOND: return COLUMN_SIZE_TIMESTAMP_SECONDS; case MILLISECOND: return COLUMN_SIZE_TIMESTAMP_MILLISECONDS; case MICROSECOND: return COLUMN_SIZE_TIMESTAMP_MICROSECONDS; case NANOSECOND: return COLUMN_SIZE_TIMESTAMP_NANOSECONDS; default: break; } } else if (fieldType instanceof ArrowType.Time) { switch (((ArrowType.Time) fieldType).getUnit()) { case SECOND: return COLUMN_SIZE_TIME; case MILLISECOND: return COLUMN_SIZE_TIME_MILLISECONDS; case MICROSECOND: return COLUMN_SIZE_TIME_MICROSECONDS; case NANOSECOND: return COLUMN_SIZE_TIME_NANOSECONDS; default: break; } } else if (fieldType instanceof ArrowType.Date) { return COLUMN_SIZE_DATE; } return null; } static String sqlToRegexLike(final String sqlPattern) { final int len = sqlPattern.length(); final StringBuilder javaPattern = new StringBuilder(len + len); for (int i = 0; i < len; i++) { final char currentChar = sqlPattern.charAt(i); if (JAVA_REGEX_SPECIALS.indexOf(currentChar) >= 0) { javaPattern.append('\\'); } switch (currentChar) { case '_': javaPattern.append('.'); break; case '%': javaPattern.append("."); javaPattern.append('*'); break; default: javaPattern.append(currentChar); break; } } return javaPattern.toString(); } } ArrowFlightConnection.java000066400000000000000000000226621515451615200373120ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty.replaceSemiColons; 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; import org.apache.arrow.util.AutoCloseables; 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 { private final BufferAllocator allocator; private final ArrowFlightSqlClientHandler clientHandler; private final ArrowFlightConnectionConfigImpl config; private ExecutorService executorService; private int metadataResultSetCount; private Map metadataResultSetMap = new HashMap<>(); /** * Creates a new {@link ArrowFlightConnection}. * * @param driver the {@link ArrowFlightJdbcDriver} to use. * @param factory the {@link AvaticaFactory} to use. * @param url the URL to use. * @param properties the {@link Properties} to use. * @param config the {@link ArrowFlightConnectionConfigImpl} to use. * @param allocator the {@link BufferAllocator} to use. * @param clientHandler the {@link ArrowFlightSqlClientHandler} to use. */ private ArrowFlightConnection( final ArrowFlightJdbcDriver driver, final AvaticaFactory factory, final String url, final Properties properties, final ArrowFlightConnectionConfigImpl config, final BufferAllocator allocator, final ArrowFlightSqlClientHandler clientHandler) { super(driver, factory, url, properties); 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; } /** * Creates a new {@link ArrowFlightConnection} to a {@link FlightClient}. * * @param driver the {@link ArrowFlightJdbcDriver} to use. * @param factory the {@link AvaticaFactory} to use. * @param url the URL to establish the connection to. * @param properties the {@link Properties} to use for this session. * @param allocator the {@link BufferAllocator} to use. * @return a new {@link ArrowFlightConnection}. * @throws SQLException on error. */ static ArrowFlightConnection createNewConnection( final ArrowFlightJdbcDriver driver, final AvaticaFactory factory, String url, final Properties properties, final BufferAllocator allocator) throws SQLException { url = replaceSemiColons(url); final ArrowFlightConnectionConfigImpl config = new ArrowFlightConnectionConfigImpl(properties); 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 DriverVersion driverVersion) throws SQLException { try { return new ArrowFlightSqlClientHandler.Builder() .withHost(config.getHost()) .withPort(config.getPort()) .withUsername(config.getUser()) .withPassword(config.getPassword()) .withTrustStorePath(config.getTrustStorePath()) .withTrustStorePassword(config.getTrustStorePassword()) .withSystemTrustStore(config.useSystemTrustStore()) .withTlsRootCertificates(config.getTlsRootCertificatesPath()) .withClientCertificate(config.getClientCertificatePath()) .withClientKey(config.getClientKeyPath()) .withBufferAllocator(allocator) .withEncryption(config.useEncryption()) .withDisableCertificateVerification(config.getDisableCertificateVerification()) .withToken(config.getToken()) .withCallOptions(config.toCallOption()) .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 { allocator.close(); } catch (final Exception allocatorCloseEx) { e.addSuppressed(allocatorCloseEx); } throw e; } } void reset() throws SQLException { // Clean up any open Statements try { AutoCloseables.close(statementMap.values()); } catch (final Exception e) { throw AvaticaConnection.HELPER.createException(e.getMessage(), e); } statementMap.clear(); // Reset Holdability this.setHoldability(this.metaData.getResultSetHoldability()); // Reset Meta ((ArrowFlightMetaImpl) this.meta).setDefaultConnectionProperties(); } /** * Gets the client {@link #clientHandler} backing this connection. * * @return the handler. */ ArrowFlightSqlClientHandler getClientHandler() { return clientHandler; } /** * Gets the {@link ExecutorService} of this connection. * * @return the {@link #executorService}. */ synchronized ExecutorService getExecutorService() { return executorService = executorService == null ? Executors.newFixedThreadPool( config.threadPoolSize(), new DefaultThreadFactory(getClass().getSimpleName())) : 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(); copy.putAll(info); return copy; } @Override public void close() throws SQLException { 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 { super.close(); } catch (final Exception e) { if (topLevelException == null) { topLevelException = e; } else { topLevelException.addSuppressed(e); } } if (topLevelException != null) { throw AvaticaConnection.HELPER.createException( topLevelException.getMessage(), topLevelException); } } BufferAllocator getBufferAllocator() { return allocator; } public ArrowFlightMetaImpl getMeta() { return (ArrowFlightMetaImpl) this.meta; } } ArrowFlightInfoStatement.java000066400000000000000000000025271515451615200377710ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.sql.SQLException; import java.sql.Statement; import org.apache.arrow.flight.FlightInfo; /** A {@link Statement} that deals with {@link FlightInfo}. */ public interface ArrowFlightInfoStatement extends Statement { @Override ArrowFlightConnection getConnection() throws SQLException; /** * Executes the query this {@link Statement} is holding. * * @return the {@link FlightInfo} for the results. * @throws SQLException on error. */ FlightInfo executeFlightInfoQuery() throws SQLException; } ArrowFlightJdbcArray.java000066400000000000000000000143361515451615200370530ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.sql.Array; import java.sql.ResultSet; import java.sql.SQLException; import java.sql.SQLFeatureNotSupportedException; import java.util.Arrays; import java.util.Map; import org.apache.arrow.driver.jdbc.accessor.impl.complex.AbstractArrowFlightJdbcListVectorAccessor; 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; import org.apache.arrow.vector.util.JsonStringArrayList; import org.apache.arrow.vector.util.TransferPair; /** * Implementation of {@link Array} using an underlying {@link FieldVector}. * * @see AbstractArrowFlightJdbcListVectorAccessor */ public class ArrowFlightJdbcArray implements Array { private final FieldVector dataVector; private final long startOffset; private final long valuesCount; /** * Instantiate an {@link Array} backed up by given {@link FieldVector}, limited by a start offset * and values count. * * @param dataVector underlying FieldVector, containing the Array items. * @param startOffset offset from FieldVector pointing to this Array's first value. * @param valuesCount how many items this Array contains. */ public ArrowFlightJdbcArray(FieldVector dataVector, long startOffset, long valuesCount) { this.dataVector = dataVector; this.startOffset = startOffset; this.valuesCount = valuesCount; } @Override public String getBaseTypeName() { final ArrowType arrowType = this.dataVector.getField().getType(); return SqlTypes.getSqlTypeNameFromArrowType(arrowType); } @Override public int getBaseType() { final ArrowType arrowType = this.dataVector.getField().getType(); return SqlTypes.getSqlTypeIdFromArrowType(arrowType); } @Override public Object getArray() throws SQLException { return getArray(null); } @Override public Object getArray(Map> map) throws SQLException { if (map != null) { throw new SQLFeatureNotSupportedException(); } return getArrayNoBoundCheck(this.dataVector, this.startOffset, this.valuesCount); } @Override public Object getArray(long index, int count) throws SQLException { return getArray(index, count, null); } private void checkBoundaries(long index, int count) { if (index < 0 || index + count > this.startOffset + this.valuesCount) { throw new ArrayIndexOutOfBoundsException(); } } private static Object getArrayNoBoundCheck(ValueVector dataVector, long start, long count) { Object[] result = new Object[LargeMemoryUtil.checkedCastToInt(count)]; for (int i = 0; i < count; i++) { result[i] = dataVector.getObject(LargeMemoryUtil.checkedCastToInt(start + i)); } return result; } @Override public Object getArray(long index, int count, Map> map) throws SQLException { if (map != null) { throw new SQLFeatureNotSupportedException(); } checkBoundaries(index, count); return getArrayNoBoundCheck( this.dataVector, LargeMemoryUtil.checkedCastToInt(this.startOffset + index), count); } @Override public ResultSet getResultSet() throws SQLException { return this.getResultSet(null); } @Override public ResultSet getResultSet(Map> map) throws SQLException { if (map != null) { throw new SQLFeatureNotSupportedException(); } return getResultSetNoBoundariesCheck(this.dataVector, this.startOffset, this.valuesCount); } @Override public ResultSet getResultSet(long index, int count) throws SQLException { return getResultSet(index, count, null); } 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(intStart, intCount); FieldVector valueVector = (FieldVector) transferPair.getTo(); VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.of(indexVector, valueVector); return ArrowFlightJdbcVectorSchemaRootResultSet.fromVectorSchemaRoot(vectorSchemaRoot); } @Override public ResultSet getResultSet(long index, int count, Map> map) throws SQLException { if (map != null) { throw new SQLFeatureNotSupportedException(); } checkBoundaries(index, count); return getResultSetNoBoundariesCheck( this.dataVector, LargeMemoryUtil.checkedCastToInt(this.startOffset + index), count); } @Override public void free() {} @Override public String toString() { JsonStringArrayList array = new JsonStringArrayList<>((int) this.valuesCount); try { array.addAll(Arrays.asList((Object[]) getArray())); } catch (SQLException e) { throw new RuntimeException(e); } return array.toString(); } } ArrowFlightJdbcConnectionPoolDataSource.java000066400000000000000000000106531515451615200426770ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.sql.SQLException; import java.util.Map; import java.util.Properties; import java.util.Queue; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; import javax.sql.ConnectionEvent; import javax.sql.ConnectionEventListener; import javax.sql.ConnectionPoolDataSource; import javax.sql.PooledConnection; import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl; /** {@link ConnectionPoolDataSource} implementation for Arrow Flight JDBC Driver. */ public class ArrowFlightJdbcConnectionPoolDataSource extends ArrowFlightJdbcDataSource implements ConnectionPoolDataSource, ConnectionEventListener, AutoCloseable { private final Map> pool = new ConcurrentHashMap<>(); /** * Instantiates a new DataSource. * * @param properties the properties * @param config the config. */ protected ArrowFlightJdbcConnectionPoolDataSource( final Properties properties, final ArrowFlightConnectionConfigImpl config) { super(properties, config); } /** * Creates a new {@link ArrowFlightJdbcConnectionPoolDataSource}. * * @param properties the properties. * @return a new data source. */ public static ArrowFlightJdbcConnectionPoolDataSource createNewDataSource( final Properties properties) { return new ArrowFlightJdbcConnectionPoolDataSource( properties, new ArrowFlightConnectionConfigImpl(properties)); } @Override public PooledConnection getPooledConnection() throws SQLException { final ArrowFlightConnectionConfigImpl config = getConfig(); return this.getPooledConnection(config.getUser(), config.getPassword()); } @Override public PooledConnection getPooledConnection(final String username, final String password) throws SQLException { final Properties properties = getProperties(username, password); Queue objectPool = pool.computeIfAbsent(properties, s -> new ConcurrentLinkedQueue<>()); ArrowFlightJdbcPooledConnection pooledConnection = objectPool.poll(); if (pooledConnection == null) { pooledConnection = createPooledConnection(new ArrowFlightConnectionConfigImpl(properties)); } else { pooledConnection.reset(); } return pooledConnection; } private ArrowFlightJdbcPooledConnection createPooledConnection( final ArrowFlightConnectionConfigImpl config) throws SQLException { ArrowFlightJdbcPooledConnection pooledConnection = new ArrowFlightJdbcPooledConnection(getConnection(config.getUser(), config.getPassword())); pooledConnection.addConnectionEventListener(this); return pooledConnection; } @Override public void connectionClosed(ConnectionEvent connectionEvent) { final ArrowFlightJdbcPooledConnection pooledConnection = (ArrowFlightJdbcPooledConnection) connectionEvent.getSource(); Queue connectionQueue = pool.get(pooledConnection.getProperties()); connectionQueue.add(pooledConnection); } @Override public void connectionErrorOccurred(ConnectionEvent connectionEvent) {} @Override public void close() throws Exception { SQLException lastException = null; for (Queue connections : this.pool.values()) { while (!connections.isEmpty()) { PooledConnection pooledConnection = connections.poll(); try { pooledConnection.close(); } catch (SQLException e) { lastException = e; } } } if (lastException != null) { throw lastException; } } } ArrowFlightJdbcCursor.java000066400000000000000000000063101515451615200372430ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.ArrayList; import java.util.Calendar; import java.util.List; import java.util.stream.Collectors; import java.util.stream.IntStream; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.calcite.avatica.ColumnMetaData; import org.apache.calcite.avatica.util.AbstractCursor; import org.apache.calcite.avatica.util.ArrayImpl; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** Arrow Flight Jdbc's Cursor class. */ public class ArrowFlightJdbcCursor extends AbstractCursor { private static final Logger LOGGER; private final VectorSchemaRoot root; private final int rowCount; private int currentRow = -1; static { LOGGER = LoggerFactory.getLogger(ArrowFlightJdbcCursor.class); } public ArrowFlightJdbcCursor(VectorSchemaRoot root) { this.root = root; rowCount = root.getRowCount(); } @Override public List createAccessors( List columns, Calendar localCalendar, ArrayImpl.Factory factory) { final List fieldVectors = root.getFieldVectors(); return IntStream.range(0, fieldVectors.size()) .mapToObj(root::getVector) .map(this::createAccessor) .collect(Collectors.toCollection(() -> new ArrayList<>(fieldVectors.size()))); } private Accessor createAccessor(FieldVector vector) { return ArrowFlightJdbcAccessorFactory.createAccessor( vector, this::getCurrentRow, (boolean wasNull) -> { // AbstractCursor creates a boolean array of length 1 to hold the wasNull value this.wasNull[0] = wasNull; }); } /** * ArrowFlightJdbcAccessors do not use {@link AbstractCursor.Getter}, as it would box primitive * types and cause performance issues. Each Accessor implementation works directly on Arrow * Vectors. */ @Override protected Getter createGetter(int column) { throw new UnsupportedOperationException("Not allowed."); } @Override public boolean next() { currentRow++; return currentRow < rowCount; } @Override public void close() { try { AutoCloseables.close(root); } catch (Exception e) { LOGGER.error(e.getMessage(), e); } } private int getCurrentRow() { return currentRow; } } ArrowFlightJdbcDataSource.java000066400000000000000000000103521515451615200400210ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty; import java.io.PrintWriter; import java.sql.SQLException; import java.sql.SQLFeatureNotSupportedException; import java.util.Properties; import java.util.logging.Logger; import javax.sql.DataSource; import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl; import org.apache.arrow.util.Preconditions; /** {@link DataSource} implementation for Arrow Flight JDBC Driver. */ public class ArrowFlightJdbcDataSource implements DataSource { private final Properties properties; private final ArrowFlightConnectionConfigImpl config; private PrintWriter logWriter; /** Instantiates a new DataSource. */ protected ArrowFlightJdbcDataSource( final Properties properties, final ArrowFlightConnectionConfigImpl config) { this.properties = Preconditions.checkNotNull(properties); this.config = Preconditions.checkNotNull(config); } /** * Gets the {@link #config} for this {@link ArrowFlightJdbcDataSource}. * * @return the {@link ArrowFlightConnectionConfigImpl}. */ protected final ArrowFlightConnectionConfigImpl getConfig() { return config; } /** * Gets a copy of the {@link #properties} for this {@link ArrowFlightJdbcDataSource} with the * provided {@code username} and {@code password}. * * @return the {@link Properties} for this data source. */ protected final Properties getProperties(final String username, final String password) { final Properties newProperties = new Properties(); newProperties.putAll(this.properties); if (username != null) { newProperties.replace(ArrowFlightConnectionProperty.USER.camelName(), username); } if (password != null) { newProperties.replace(ArrowFlightConnectionProperty.PASSWORD.camelName(), password); } return ArrowFlightJdbcDriver.lowerCasePropertyKeys(newProperties); } /** * Creates a new {@link ArrowFlightJdbcDataSource}. * * @param properties the properties. * @return a new data source. */ public static ArrowFlightJdbcDataSource createNewDataSource(final Properties properties) { return new ArrowFlightJdbcDataSource( properties, new ArrowFlightConnectionConfigImpl(properties)); } @Override public ArrowFlightConnection getConnection() throws SQLException { return getConnection(config.getUser(), config.getPassword()); } @Override public ArrowFlightConnection getConnection(final String username, final String password) throws SQLException { final Properties properties = getProperties(username, password); return new ArrowFlightJdbcDriver().connect(config.url(), properties); } @Override public T unwrap(Class aClass) throws SQLException { throw new SQLException("ArrowFlightJdbcDataSource is not a wrapper."); } @Override public boolean isWrapperFor(Class aClass) { return false; } @Override public PrintWriter getLogWriter() { return this.logWriter; } @Override public void setLogWriter(PrintWriter logWriter) { this.logWriter = logWriter; } @Override public void setLoginTimeout(int timeout) throws SQLException { throw new SQLFeatureNotSupportedException("Setting Login timeout is not supported."); } @Override public int getLoginTimeout() { return 0; } @Override public Logger getParentLogger() { return Logger.getLogger("ArrowFlightJdbc"); } } ArrowFlightJdbcDriver.java000066400000000000000000000243751515451615200372340ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty.replaceSemiColons; import java.io.BufferedReader; import java.io.IOException; import java.io.InputStream; import java.io.InputStreamReader; import java.io.Reader; import java.net.URI; import java.nio.charset.StandardCharsets; import java.sql.SQLException; import java.util.Map; import java.util.Objects; import java.util.Optional; import java.util.Properties; import java.util.logging.Logger; import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty; import org.apache.arrow.driver.jdbc.utils.UrlParser; import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.util.VisibleForTesting; import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.DriverVersion; import org.apache.calcite.avatica.Meta; import org.apache.calcite.avatica.UnregisteredDriver; /** JDBC driver for querying data from an Apache Arrow Flight server. */ public class ArrowFlightJdbcDriver extends UnregisteredDriver { private static final String CONNECT_STRING_PREFIX = "jdbc:arrow-flight-sql://"; private static final String CONNECT_STRING_PREFIX_DEPRECATED = "jdbc:arrow-flight://"; private static final String CONNECTION_STRING_EXPECTED = "jdbc:arrow-flight-sql://[host][:port][?param1=value&...]"; private static DriverVersion version; static { // Special code for supporting Java9 and higher. // Netty requires some extra properties to unlock some native memory management api // Setting this property if not already set externally // This has to be done before any netty class is being loaded final String key = "io.netty.tryReflectionSetAccessible"; final String tryReflectionSetAccessible = System.getProperty(key); if (tryReflectionSetAccessible == null) { System.setProperty(key, Boolean.TRUE.toString()); } new ArrowFlightJdbcDriver().register(); } @Override public Logger getParentLogger() { // Return the logger associated with the driver package ('org.apache.arrow.driver.jdbc') // When packaged in flight-sql-jdbc-driver, it will also apply to all shaded dependencies return Logger.getLogger(getClass().getPackage().getName()); } @Override public ArrowFlightConnection connect(final String url, final Properties info) throws SQLException { final Properties properties = new Properties(info); if (info != null) { properties.putAll(info); } if (url != null) { final Optional> maybeProperties = getUrlsArgs(url); if (!maybeProperties.isPresent()) { return null; } final Map propertiesFromUrl = maybeProperties.get(); properties.putAll(propertiesFromUrl); } try { return ArrowFlightConnection.createNewConnection( this, factory, url, lowerCasePropertyKeys(properties), new RootAllocator(Long.MAX_VALUE)); } catch (final FlightRuntimeException e) { throw new SQLException("Failed to connect.", e); } } @Override protected String getFactoryClassName(final JdbcVersion jdbcVersion) { return ArrowFlightJdbcFactory.class.getName(); } @Override @SuppressWarnings("StringSplitter") protected DriverVersion createDriverVersion() { if (version == null) { final InputStream flightProperties = this.getClass().getResourceAsStream("/properties/flight.properties"); if (flightProperties == null) { throw new RuntimeException( "Flight Properties not found. Ensure the JAR was built properly."); } try (final Reader reader = new BufferedReader(new InputStreamReader(flightProperties, StandardCharsets.UTF_8))) { final Properties properties = new Properties(); properties.load(reader); final String parentName = properties.getProperty("org.apache.arrow.flight.name"); final String parentVersion = properties.getProperty("org.apache.arrow.flight.version"); final String[] pVersion = parentVersion.split("\\."); final int parentMajorVersion = Integer.parseInt(pVersion[0]); final int parentMinorVersion = Integer.parseInt(pVersion[1]); final String childName = properties.getProperty("org.apache.arrow.flight.jdbc-driver.name"); final String childVersion = properties.getProperty("org.apache.arrow.flight.jdbc-driver.version"); final String[] cVersion = childVersion.split("\\."); final int childMajorVersion = Integer.parseInt(cVersion[0]); final int childMinorVersion = Integer.parseInt(cVersion[1]); version = new DriverVersion( childName, childVersion, parentName, parentVersion, true, childMajorVersion, childMinorVersion, parentMajorVersion, parentMinorVersion); } catch (final IOException e) { throw new RuntimeException("Failed to load driver version.", e); } } return version; } @Override public Meta createMeta(final AvaticaConnection connection) { return new ArrowFlightMetaImpl(connection); } @Override protected String getConnectStringPrefix() { return CONNECT_STRING_PREFIX; } @Override public boolean acceptsURL(final String url) { Preconditions.checkNotNull(url); return url.startsWith(CONNECT_STRING_PREFIX) || url.startsWith(CONNECT_STRING_PREFIX_DEPRECATED); } /** * Parses the provided url based on the format this driver accepts, retrieving arguments after the * {@link #CONNECT_STRING_PREFIX}. * *

    This method gets the args if the provided URL follows this pattern: {@code * jdbc:arrow-flight-sql://:[/?key1=val1&key2=val2&(...)]} * * * * * * * * * * * * * * * * * * * * * * * * * * * *
    GroupDefinitionValue
    ? — inaccessible{@link #getConnectStringPrefix} * the URL prefix accepted by this driver, i.e., * {@code "jdbc:arrow-flight-sql://"} *
    1IPv4 host name * first word after previous group and before "{@code :}" *
    2IPv4 port number * first number after previous group and before "{@code /?}" *
    3custom call parameters * all parameters provided after "{@code /?}" — must follow the * pattern: "{@code key=value}" with "{@code &}" separating a * parameter from another *
    * * @param url The url to parse. * @return the parsed arguments, or an empty optional if the driver does not handle this URL. * @throws SQLException If an error occurs while trying to parse the URL. */ @VisibleForTesting // ArrowFlightJdbcDriverTest Optional> getUrlsArgs(String url) throws SQLException { /* * * Perhaps this logic should be inside a utility class, separated from this * one, so as to better delegate responsibilities and concerns throughout * the code and increase maintainability. * * ===== * * Keep in mind that the URL must ALWAYS follow the pattern: * "jdbc:arrow-flight-sql://:[/?param1=value1¶m2=value2&(...)]." * */ final Properties resultMap = new Properties(); url = replaceSemiColons(url); if (!url.startsWith("jdbc:")) { throw new SQLException( "Connection string must start with 'jdbc:'. Expected format: " + CONNECTION_STRING_EXPECTED); } // It's necessary to use a string without "jdbc:" at the beginning to be parsed as a valid URL. url = url.substring(5); final URI uri; try { uri = URI.create(url); } catch (final IllegalArgumentException e) { throw new SQLException("Malformed/invalid URL!", e); } if (!Objects.equals(uri.getScheme(), "arrow-flight") && !Objects.equals(uri.getScheme(), "arrow-flight-sql")) { return Optional.empty(); } if (uri.getHost() == null) { throw new SQLException( "URL must have a host. Expected format: " + CONNECTION_STRING_EXPECTED); } else if (uri.getPort() < 0) { throw new SQLException( "URL must have a port. Expected format: " + CONNECTION_STRING_EXPECTED); } resultMap.put(ArrowFlightConnectionProperty.HOST.camelName(), uri.getHost()); // host resultMap.put(ArrowFlightConnectionProperty.PORT.camelName(), uri.getPort()); // port final String extraParams = uri.getRawQuery(); // optional params if (extraParams != null) { final Map keyValuePairs = UrlParser.parse(extraParams, "&"); resultMap.putAll(keyValuePairs); } return Optional.of(resultMap); } static Properties lowerCasePropertyKeys(final Properties properties) { final Properties resultProperty = new Properties(); properties.forEach((k, v) -> resultProperty.put(k.toString().toLowerCase(), v)); return resultProperty; } } ArrowFlightJdbcFactory.java000066400000000000000000000106601515451615200374000ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.sql.ResultSetMetaData; import java.sql.SQLException; import java.util.Properties; import java.util.TimeZone; import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler; import org.apache.arrow.memory.RootAllocator; import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.AvaticaFactory; import org.apache.calcite.avatica.AvaticaResultSetMetaData; import org.apache.calcite.avatica.AvaticaSpecificDatabaseMetaData; import org.apache.calcite.avatica.AvaticaStatement; import org.apache.calcite.avatica.Meta; import org.apache.calcite.avatica.QueryState; import org.apache.calcite.avatica.UnregisteredDriver; /** Factory for the Arrow Flight JDBC Driver. */ public class ArrowFlightJdbcFactory implements AvaticaFactory { private final int major; private final int minor; // This need to be public so Avatica can call this constructor public ArrowFlightJdbcFactory() { this(4, 1); } private ArrowFlightJdbcFactory(final int major, final int minor) { this.major = major; this.minor = minor; } @Override public AvaticaConnection newConnection( final UnregisteredDriver driver, final AvaticaFactory factory, final String url, final Properties info) throws SQLException { return ArrowFlightConnection.createNewConnection( (ArrowFlightJdbcDriver) driver, factory, url, info, new RootAllocator(Long.MAX_VALUE)); } @Override public AvaticaStatement newStatement( final AvaticaConnection connection, final Meta.StatementHandle handle, final int resultType, final int resultSetConcurrency, final int resultSetHoldability) { return new ArrowFlightStatement( (ArrowFlightConnection) connection, handle, resultType, resultSetConcurrency, resultSetHoldability); } @Override public ArrowFlightPreparedStatement newPreparedStatement( final AvaticaConnection connection, final Meta.StatementHandle statementHandle, final Meta.Signature signature, final int resultType, final int resultSetConcurrency, final int resultSetHoldability) throws SQLException { final ArrowFlightConnection flightConnection = (ArrowFlightConnection) connection; ArrowFlightSqlClientHandler.PreparedStatement preparedStatement = flightConnection.getMeta().getPreparedStatement(statementHandle); return ArrowFlightPreparedStatement.newPreparedStatement( flightConnection, preparedStatement, statementHandle, signature, resultType, resultSetConcurrency, resultSetHoldability); } @Override public ArrowFlightJdbcVectorSchemaRootResultSet newResultSet( final AvaticaStatement statement, final QueryState state, final Meta.Signature signature, final TimeZone timeZone, final Meta.Frame frame) throws SQLException { final ResultSetMetaData metaData = newResultSetMetaData(statement, signature); return new ArrowFlightJdbcFlightStreamResultSet( statement, state, signature, metaData, timeZone, frame); } @Override public AvaticaSpecificDatabaseMetaData newDatabaseMetaData(final AvaticaConnection connection) { return new ArrowDatabaseMetadata(connection); } @Override public ResultSetMetaData newResultSetMetaData( final AvaticaStatement avaticaStatement, final Meta.Signature signature) { return new AvaticaResultSetMetaData(avaticaStatement, null, signature); } @Override public int getJdbcMajorVersion() { return major; } @Override public int getJdbcMinorVersion() { return minor; } } ArrowFlightJdbcFlightStreamResultSet.java000066400000000000000000000222431515451615200422350ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.driver.jdbc.utils.FlightEndpointDataQueue.createNewQueue; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.util.Optional; import java.util.TimeZone; import java.util.concurrent.TimeUnit; import org.apache.arrow.driver.jdbc.client.CloseableEndpointStreamPair; import org.apache.arrow.driver.jdbc.utils.FlightEndpointDataQueue; import org.apache.arrow.driver.jdbc.utils.VectorSchemaRootTransformer; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightStream; 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.AvaticaResultSet; import org.apache.calcite.avatica.AvaticaResultSetMetaData; import org.apache.calcite.avatica.AvaticaStatement; import org.apache.calcite.avatica.Meta; import org.apache.calcite.avatica.QueryState; /** * {@link ResultSet} implementation for Arrow Flight used to access the results of multiple {@link * FlightStream} objects. */ public final class ArrowFlightJdbcFlightStreamResultSet extends ArrowFlightJdbcVectorSchemaRootResultSet { private final ArrowFlightConnection connection; private final FlightInfo flightInfo; private CloseableEndpointStreamPair currentEndpointData; private FlightEndpointDataQueue flightEndpointDataQueue; private VectorSchemaRootTransformer transformer; private VectorSchemaRoot currentVectorSchemaRoot; private Schema schema; private Integer id = null; // used for metadata result sets only /** Public constructor used by ArrowFlightJdbcFactory. */ ArrowFlightJdbcFlightStreamResultSet( final AvaticaStatement statement, final QueryState state, final Meta.Signature signature, final ResultSetMetaData resultSetMetaData, final TimeZone timeZone, final Meta.Frame firstFrame) throws SQLException { super(statement, state, signature, resultSetMetaData, timeZone, firstFrame); this.connection = (ArrowFlightConnection) statement.connection; this.flightInfo = ((ArrowFlightInfoStatement) statement).executeFlightInfoQuery(); } /** Private constructor for fromFlightInfo. */ private ArrowFlightJdbcFlightStreamResultSet( final ArrowFlightConnection connection, final QueryState state, final Meta.Signature signature, final ResultSetMetaData resultSetMetaData, final TimeZone timeZone, final Meta.Frame firstFrame, final FlightInfo flightInfo) throws SQLException { super(null, state, signature, resultSetMetaData, timeZone, firstFrame); this.connection = connection; this.flightInfo = flightInfo; this.id = connection.getNewMetadataResultSetId(this); } /** * Create a {@link ResultSet} which pulls data from given {@link FlightInfo}. This is used to * fetch result sets from DatabaseMetadata calls and skips the Avatica factory. * * @param connection The connection linked to the returned ResultSet. * @param flightInfo The FlightInfo from which data will be iterated by the returned ResultSet. * @param transformer Optional transformer for processing VectorSchemaRoot before access from * ResultSet * @return A ResultSet which pulls data from given FlightInfo. */ static ArrowFlightJdbcFlightStreamResultSet fromFlightInfo( final ArrowFlightConnection connection, final FlightInfo flightInfo, final VectorSchemaRootTransformer transformer) throws SQLException { // Similar to how org.apache.calcite.avatica.util.ArrayFactoryImpl does final TimeZone timeZone = TimeZone.getDefault(); final QueryState state = new QueryState(); final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null); final AvaticaResultSetMetaData resultSetMetaData = new AvaticaResultSetMetaData(null, null, signature); final ArrowFlightJdbcFlightStreamResultSet resultSet = new ArrowFlightJdbcFlightStreamResultSet( connection, state, signature, resultSetMetaData, timeZone, null, flightInfo); resultSet.transformer = transformer; resultSet.populateData(); return resultSet; } private void loadNewQueue() { Optional.ofNullable(flightEndpointDataQueue).ifPresent(AutoCloseables::closeNoChecked); flightEndpointDataQueue = createNewQueue(connection.getExecutorService()); } private void loadNewFlightStream() throws SQLException { if (currentEndpointData != null) { AutoCloseables.closeNoChecked(currentEndpointData); } this.currentEndpointData = getNextEndpointStream(true); } @Override protected AvaticaResultSet execute() throws SQLException { if (flightInfo != null) { schema = flightInfo.getSchemaOptional().orElse(null); populateData(); } return this; } private void populateData() throws SQLException { loadNewQueue(); flightEndpointDataQueue.enqueue(connection.getClientHandler().getStreams(flightInfo)); loadNewFlightStream(); // Ownership of the root will be passed onto the cursor. if (currentEndpointData != null) { populateDataForCurrentFlightStream(); } } private void populateDataForCurrentFlightStream() throws SQLException { final VectorSchemaRoot originalRoot = currentEndpointData.getStream().getRoot(); if (transformer != null) { try { currentVectorSchemaRoot = transformer.transform(originalRoot, currentVectorSchemaRoot); } catch (final Exception e) { throw new SQLException("Failed to transform VectorSchemaRoot.", e); } } else { currentVectorSchemaRoot = originalRoot; } populateData(currentVectorSchemaRoot, schema); } /** Expose appMetadata associated with the underlying FlightInfo for this ResultSet. */ public byte[] getAppMetadata() { return flightInfo.getAppMetadata(); } @Override public boolean next() throws SQLException { if (currentVectorSchemaRoot == null) { return false; } while (true) { final boolean hasNext = super.next(); final int maxRows = statement != null ? statement.getMaxRows() : 0; if (maxRows != 0 && this.getRow() > maxRows) { if (statement.isCloseOnCompletion()) { statement.close(); } return false; } if (hasNext) { return true; } if (currentEndpointData != null) { currentEndpointData.getStream().getRoot().clear(); if (currentEndpointData.getStream().next()) { populateDataForCurrentFlightStream(); continue; } flightEndpointDataQueue.enqueue(currentEndpointData); } currentEndpointData = getNextEndpointStream(false); if (currentEndpointData != null) { populateDataForCurrentFlightStream(); continue; } if (statement != null && statement.isCloseOnCompletion()) { statement.close(); } return false; } } @Override protected void cancel() { super.cancel(); final CloseableEndpointStreamPair currentEndpoint = this.currentEndpointData; if (currentEndpoint != null) { currentEndpoint.getStream().cancel("Cancel", null); } if (flightEndpointDataQueue != null) { try { flightEndpointDataQueue.close(); } catch (final Exception e) { throw new RuntimeException(e); } } } @Override public synchronized void close() { try { if (isClosed()) { return; } this.connection.onResultSetClose(id); if (flightEndpointDataQueue != null) { // flightStreamQueue should close currentFlightStream internally flightEndpointDataQueue.close(); } else if (currentEndpointData != null) { // close is only called for currentFlightStream if there's no queue currentEndpointData.close(); } } catch (final Exception e) { throw new RuntimeException(e); } finally { super.close(); } } private CloseableEndpointStreamPair getNextEndpointStream(final boolean canTimeout) throws SQLException { if (canTimeout) { final int statementTimeout = statement != null ? statement.getQueryTimeout() : 0; return statementTimeout != 0 ? flightEndpointDataQueue.next(statementTimeout, TimeUnit.SECONDS) : flightEndpointDataQueue.next(); } else { return flightEndpointDataQueue.next(); } } } ArrowFlightJdbcPooledConnection.java000066400000000000000000000064251515451615200412370ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.sql.Connection; import java.sql.SQLException; import java.util.Collections; import java.util.HashSet; import java.util.Properties; import java.util.Set; import javax.sql.ConnectionEvent; import javax.sql.ConnectionEventListener; import javax.sql.PooledConnection; import javax.sql.StatementEventListener; import org.apache.arrow.driver.jdbc.utils.ConnectionWrapper; /** {@link PooledConnection} implementation for Arrow Flight JDBC Driver. */ public class ArrowFlightJdbcPooledConnection implements PooledConnection { private final ArrowFlightConnection connection; private final Set eventListeners; private final Set statementEventListeners; private final class ConnectionHandle extends ConnectionWrapper { private boolean closed = false; public ConnectionHandle() { super(connection); } @Override public void close() throws SQLException { if (!closed) { closed = true; onConnectionClosed(); } } @Override public boolean isClosed() throws SQLException { return this.closed || super.isClosed(); } } ArrowFlightJdbcPooledConnection(ArrowFlightConnection connection) { this.connection = connection; this.eventListeners = Collections.synchronizedSet(new HashSet<>()); this.statementEventListeners = Collections.synchronizedSet(new HashSet<>()); } public Properties getProperties() { return connection.getClientInfo(); } @Override public Connection getConnection() throws SQLException { return new ConnectionHandle(); } @Override public void close() throws SQLException { this.connection.close(); } void reset() throws SQLException { this.connection.reset(); } @Override public void addConnectionEventListener(ConnectionEventListener listener) { eventListeners.add(listener); } @Override public void removeConnectionEventListener(ConnectionEventListener listener) { this.eventListeners.remove(listener); } @Override public void addStatementEventListener(StatementEventListener listener) { statementEventListeners.add(listener); } @Override public void removeStatementEventListener(StatementEventListener listener) { this.statementEventListeners.remove(listener); } private void onConnectionClosed() { ConnectionEvent connectionEvent = new ConnectionEvent(this); eventListeners.forEach(listener -> listener.connectionClosed(connectionEvent)); } } ArrowFlightJdbcTime.java000066400000000000000000000067611515451615200366760ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.calcite.avatica.util.DateTimeUtils.MILLIS_PER_DAY; import com.google.common.collect.ImmutableList; import java.sql.Time; import java.time.LocalTime; import java.time.temporal.ChronoField; import java.util.List; import java.util.Objects; import java.util.concurrent.TimeUnit; import org.apache.arrow.util.VisibleForTesting; /** * Wrapper class for Time objects to include the milliseconds part in ISO 8601 format in * this#toString. */ public class ArrowFlightJdbcTime extends Time { private static final List LEADING_ZEROES = ImmutableList.of("", "0", "00"); // Desired length of the millisecond portion should be 3 private static final int DESIRED_MILLIS_LENGTH = 3; // Millis of the date time object. private final int millisReprValue; /** * Constructs this object based on epoch millis. * * @param milliseconds milliseconds representing Time. */ public ArrowFlightJdbcTime(final long milliseconds) { super(milliseconds); millisReprValue = getMillisReprValue(milliseconds); } @VisibleForTesting ArrowFlightJdbcTime(final LocalTime time) { // Although the constructor is deprecated, this is the exact same code as // Time#valueOf(LocalTime) super(time.getHour(), time.getMinute(), time.getSecond()); millisReprValue = time.get(ChronoField.MILLI_OF_SECOND); } private int getMillisReprValue(long milliseconds) { // Extract the millisecond part from epoch nano day if (milliseconds >= MILLIS_PER_DAY) { // Convert to Epoch Day milliseconds %= MILLIS_PER_DAY; } else if (milliseconds < 0) { // LocalTime#ofNanoDay only accepts positive values milliseconds -= ((milliseconds / MILLIS_PER_DAY) - 1) * MILLIS_PER_DAY; } return LocalTime.ofNanoOfDay(TimeUnit.MILLISECONDS.toNanos(milliseconds)) .get(ChronoField.MILLI_OF_SECOND); } @Override public String toString() { final StringBuilder time = new StringBuilder().append(super.toString()); if (millisReprValue > 0) { final String millisString = Integer.toString(millisReprValue); // dot to separate the fractional seconds time.append("."); final int millisLength = millisString.length(); if (millisLength < DESIRED_MILLIS_LENGTH) { // add necessary leading zeroes time.append(LEADING_ZEROES.get(DESIRED_MILLIS_LENGTH - millisLength)); } time.append(millisString); } return time.toString(); } // Spotbugs requires these methods to be overridden @Override public boolean equals(Object obj) { return super.equals(obj); } @Override public int hashCode() { return Objects.hash(super.hashCode(), this.millisReprValue); } } ArrowFlightJdbcVectorSchemaRootResultSet.java000066400000000000000000000140041515451615200430670ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 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.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; /** {@link ResultSet} implementation used to access a {@link VectorSchemaRoot}. */ public class ArrowFlightJdbcVectorSchemaRootResultSet extends AvaticaResultSet { private static final Logger LOGGER = LoggerFactory.getLogger(ArrowFlightJdbcVectorSchemaRootResultSet.class); VectorSchemaRoot vectorSchemaRoot; ArrowFlightJdbcVectorSchemaRootResultSet( final AvaticaStatement statement, final QueryState state, final Signature signature, final ResultSetMetaData resultSetMetaData, final TimeZone timeZone, final Frame firstFrame) throws SQLException { super(statement, state, signature, resultSetMetaData, timeZone, firstFrame); } /** * Instantiate a ResultSet backed up by given VectorSchemaRoot. * * @param vectorSchemaRoot root from which the ResultSet will access. * @return a ResultSet which accesses the given VectorSchemaRoot */ public static ArrowFlightJdbcVectorSchemaRootResultSet fromVectorSchemaRoot( final VectorSchemaRoot vectorSchemaRoot) throws SQLException { // Similar to how org.apache.calcite.avatica.util.ArrayFactoryImpl does final TimeZone timeZone = TimeZone.getDefault(); final QueryState state = new QueryState(); final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null); final AvaticaResultSetMetaData resultSetMetaData = new AvaticaResultSetMetaData(null, null, signature); final ArrowFlightJdbcVectorSchemaRootResultSet resultSet = new ArrowFlightJdbcVectorSchemaRootResultSet( null, state, signature, resultSetMetaData, timeZone, null); resultSet.populateData(vectorSchemaRoot); return resultSet; } @Override protected AvaticaResultSet execute() throws SQLException { throw new RuntimeException("Can only execute with execute(VectorSchemaRoot)"); } void populateData(final VectorSchemaRoot vectorSchemaRoot) { populateData(vectorSchemaRoot, null); } void populateData(final VectorSchemaRoot vectorSchemaRoot, final Schema schema) { Schema currentSchema = schema == null ? vectorSchemaRoot.getSchema() : schema; final List columns = ConvertUtils.convertArrowFieldsToColumnMetaDataList(currentSchema.getFields()); signature.columns.clear(); signature.columns.addAll(columns); this.vectorSchemaRoot = vectorSchemaRoot; 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, localCalendar); } } @Override protected void cancel() { signature.columns.clear(); super.cancel(); try { AutoCloseables.close(vectorSchemaRoot); } catch (final Exception e) { throw new RuntimeException(e); } } @Override public void close() { final Set exceptions = new HashSet<>(); try { if (isClosed()) { return; } } catch (final SQLException e) { exceptions.add(e); } try { AutoCloseables.close(vectorSchemaRoot); } 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() .findAny() .ifPresent( e -> { throw new RuntimeException(e); }); } } ArrowFlightMetaImpl.java000066400000000000000000000265211515451615200367210ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.sql.Connection; import java.sql.SQLException; import java.sql.SQLTimeoutException; import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler.PreparedStatement; import org.apache.arrow.driver.jdbc.utils.AvaticaParameterBinder; import org.apache.arrow.driver.jdbc.utils.ConvertUtils; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.ColumnMetaData; import org.apache.calcite.avatica.MetaImpl; import org.apache.calcite.avatica.NoSuchStatementException; import org.apache.calcite.avatica.QueryState; import org.apache.calcite.avatica.remote.TypedValue; /** Metadata handler for Arrow Flight. */ public class ArrowFlightMetaImpl extends MetaImpl { private final Map statementHandlePreparedStatementMap; /** * Constructs a {@link MetaImpl} object specific for Arrow Flight. * * @param connection A {@link AvaticaConnection}. */ public ArrowFlightMetaImpl(final AvaticaConnection connection) { super(connection); this.statementHandlePreparedStatementMap = new ConcurrentHashMap<>(); setDefaultConnectionProperties(); } /** Construct a signature. */ static Signature newSignature(final String sql, Schema resultSetSchema, Schema parameterSchema) { List columnMetaData = resultSetSchema == null ? new ArrayList<>() : ConvertUtils.convertArrowFieldsToColumnMetaDataList(resultSetSchema.getFields()); List parameters = parameterSchema == null ? new ArrayList<>() : ConvertUtils.convertArrowFieldsToAvaticaParameters(parameterSchema.getFields()); StatementType 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); } @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 // this moment if (preparedStatement != null) { preparedStatement.close(); } } @Override public void commit(final ConnectionHandle connectionHandle) { // TODO Fill this stub. } @Override public ExecuteResult execute( final StatementHandle statementHandle, final List typedValues, final long maxRowCount) { Preconditions.checkArgument( connection.id.equals(statementHandle.connectionId), "Connection IDs are not consistent"); PreparedStatement preparedStatement = getPreparedStatement(statementHandle); if (preparedStatement == null) { throw new IllegalStateException("Prepared statement not found: " + statementHandle); } new AvaticaParameterBinder( preparedStatement, ((ArrowFlightConnection) connection).getBufferAllocator()) .bind(typedValues); if (statementHandle.signature == null || statementHandle.signature.statementType == StatementType.IS_DML) { // Update query long updatedCount = preparedStatement.executeUpdate(); return new ExecuteResult( Collections.singletonList( MetaResultSet.count(statementHandle.connectionId, statementHandle.id, updatedCount))); } else { // TODO Why is maxRowCount ignored? return new ExecuteResult( Collections.singletonList( MetaResultSet.create( statementHandle.connectionId, statementHandle.id, true, statementHandle.signature, null))); } } @Override public ExecuteResult execute( final StatementHandle statementHandle, final List typedValues, final int maxRowsInFirstFrame) { return execute(statementHandle, typedValues, (long) maxRowsInFirstFrame); } @Override public ExecuteBatchResult executeBatch( final StatementHandle statementHandle, final List> parameterValuesList) throws IllegalStateException { Preconditions.checkArgument( connection.id.equals(statementHandle.connectionId), "Connection IDs are not consistent"); PreparedStatement preparedStatement = getPreparedStatement(statementHandle); if (preparedStatement == null) { throw new IllegalStateException("Prepared statement not found: " + statementHandle); } final AvaticaParameterBinder binder = new AvaticaParameterBinder( preparedStatement, ((ArrowFlightConnection) connection).getBufferAllocator()); for (int i = 0; i < parameterValuesList.size(); i++) { binder.bind(parameterValuesList.get(i), i); } // Update query long[] updatedCounts = {preparedStatement.executeUpdate()}; return new ExecuteBatchResult(updatedCounts); } @Override public Frame fetch( final StatementHandle statementHandle, final long offset, final int fetchMaxRowCount) { /* * ArrowFlightMetaImpl does not use frames. * Instead, we have accessors that contain a VectorSchemaRoot with * the results. */ throw AvaticaConnection.HELPER.wrap( String.format("%s does not use frames.", this), AvaticaConnection.HELPER.unsupported()); } private PreparedStatement prepareForHandle(final String query, StatementHandle handle) { final PreparedStatement preparedStatement = ((ArrowFlightConnection) connection).getClientHandler().prepare(query); handle.signature = newSignature( query, preparedStatement.getDataSetSchema(), preparedStatement.getParameterSchema()); statementHandlePreparedStatementMap.put(new StatementHandleKey(handle), preparedStatement); return preparedStatement; } @Override public StatementHandle prepare( final ConnectionHandle connectionHandle, final String query, final long maxRowCount) { final StatementHandle handle = super.createStatement(connectionHandle); prepareForHandle(query, handle); return handle; } @Override public ExecuteResult prepareAndExecute( final StatementHandle statementHandle, final String query, final long maxRowCount, final PrepareCallback prepareCallback) throws NoSuchStatementException { return prepareAndExecute( statementHandle, query, maxRowCount, -1 /* Not used */, prepareCallback); } @Override public ExecuteResult prepareAndExecute( final StatementHandle handle, final String query, final long maxRowCount, final int maxRowsInFirstFrame, final PrepareCallback callback) throws NoSuchStatementException { try { PreparedStatement preparedStatement = prepareForHandle(query, handle); final StatementType statementType = preparedStatement.getType(); final long updateCount = statementType.equals(StatementType.UPDATE) ? preparedStatement.executeUpdate() : -1; synchronized (callback.getMonitor()) { callback.clear(); callback.assign(handle.signature, null, updateCount); } callback.execute(); final MetaResultSet metaResultSet = 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 // Exceptions. throw new RuntimeException(e); } catch (SQLException e) { throw new NoSuchStatementException(handle); } } @Override public ExecuteBatchResult prepareAndExecuteBatch( final StatementHandle statementHandle, final List queries) throws NoSuchStatementException { // TODO Fill this stub. return null; } @Override public void rollback(final ConnectionHandle connectionHandle) { // TODO Fill this stub. } @Override public boolean syncResults( final StatementHandle statementHandle, final QueryState queryState, final long offset) throws NoSuchStatementException { // TODO Fill this stub. 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 .setDirty(false) .setAutoCommit(true) .setReadOnly(true) .setCatalog(null) .setSchema(null) .setTransactionIsolation(Connection.TRANSACTION_NONE); } 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 // signature in // an UPDATE code path so we can't directly use StatementHandle as a map key. private static final class StatementHandleKey { public final String connectionId; public final int id; StatementHandleKey(StatementHandle statementHandle) { this.connectionId = statementHandle.connectionId; this.id = statementHandle.id; } @Override public boolean equals(Object o) { if (this == o) { return true; } if (o == null || getClass() != o.getClass()) { return false; } StatementHandleKey that = (StatementHandleKey) o; if (id != that.id) { return false; } return connectionId.equals(that.connectionId); } @Override public int hashCode() { int result = connectionId.hashCode(); result = 31 * result + id; return result; } } } ArrowFlightPreparedStatement.java000066400000000000000000000057671515451615200406510ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.sql.PreparedStatement; import java.sql.SQLException; import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.util.Preconditions; import org.apache.calcite.avatica.AvaticaPreparedStatement; import org.apache.calcite.avatica.Meta.Signature; import org.apache.calcite.avatica.Meta.StatementHandle; /** Arrow Flight JBCS's implementation {@link PreparedStatement}. */ public class ArrowFlightPreparedStatement extends AvaticaPreparedStatement implements ArrowFlightInfoStatement { private final ArrowFlightSqlClientHandler.PreparedStatement preparedStatement; private ArrowFlightPreparedStatement( final ArrowFlightConnection connection, final ArrowFlightSqlClientHandler.PreparedStatement preparedStatement, final StatementHandle handle, final Signature signature, final int resultSetType, final int resultSetConcurrency, final int resultSetHoldability) throws SQLException { super(connection, handle, signature, resultSetType, resultSetConcurrency, resultSetHoldability); this.preparedStatement = Preconditions.checkNotNull(preparedStatement); } static ArrowFlightPreparedStatement newPreparedStatement( final ArrowFlightConnection connection, final ArrowFlightSqlClientHandler.PreparedStatement preparedStmt, final StatementHandle statementHandle, final Signature signature, final int resultSetType, final int resultSetConcurrency, final int resultSetHoldability) throws SQLException { return new ArrowFlightPreparedStatement( connection, preparedStmt, statementHandle, signature, resultSetType, resultSetConcurrency, resultSetHoldability); } @Override public ArrowFlightConnection getConnection() throws SQLException { return (ArrowFlightConnection) super.getConnection(); } @Override public synchronized void close() throws SQLException { this.preparedStatement.close(); super.close(); } @Override public FlightInfo executeFlightInfoQuery() throws SQLException { return preparedStatement.executeQuery(); } } ArrowFlightStatement.java000066400000000000000000000046261515451615200371570ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.sql.SQLException; import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler.PreparedStatement; import org.apache.arrow.driver.jdbc.utils.ConvertUtils; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.calcite.avatica.AvaticaStatement; import org.apache.calcite.avatica.Meta; import org.apache.calcite.avatica.Meta.StatementHandle; /** A SQL statement for querying data from an Arrow Flight server. */ public class ArrowFlightStatement extends AvaticaStatement implements ArrowFlightInfoStatement { ArrowFlightStatement( final ArrowFlightConnection connection, final StatementHandle handle, final int resultSetType, final int resultSetConcurrency, final int resultSetHoldability) { super(connection, handle, resultSetType, resultSetConcurrency, resultSetHoldability); } @Override public ArrowFlightConnection getConnection() throws SQLException { return (ArrowFlightConnection) super.getConnection(); } @Override public FlightInfo executeFlightInfoQuery() throws SQLException { final PreparedStatement preparedStatement = getConnection().getMeta().getPreparedStatement(handle); final Meta.Signature signature = getSignature(); if (signature == null) { return null; } final Schema resultSetSchema = preparedStatement.getDataSetSchema(); signature.columns.addAll( ConvertUtils.convertArrowFieldsToColumnMetaDataList(resultSetSchema.getFields())); setSignature(signature); return preparedStatement.executeQuery(); } } accessor/000077500000000000000000000000001515451615200337715ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbcArrowFlightJdbcAccessor.java000066400000000000000000000156711515451615200413440ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.calcite.avatica.util.Cursor.Accessor; import java.io.InputStream; import java.io.Reader; import java.math.BigDecimal; import java.net.URL; import java.sql.Array; import java.sql.Blob; import java.sql.Clob; import java.sql.Date; import java.sql.NClob; import java.sql.Ref; import java.sql.SQLException; import java.sql.SQLXML; import java.sql.Struct; import java.sql.Time; import java.sql.Timestamp; import java.util.Calendar; import java.util.Map; import java.util.function.IntSupplier; /** Base Jdbc Accessor. */ public abstract class ArrowFlightJdbcAccessor implements Accessor { private final IntSupplier currentRowSupplier; // All the derived accessor classes should alter this as they encounter null Values protected boolean wasNull; protected ArrowFlightJdbcAccessorFactory.WasNullConsumer wasNullConsumer; protected ArrowFlightJdbcAccessor( final IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer wasNullConsumer) { this.currentRowSupplier = currentRowSupplier; this.wasNullConsumer = wasNullConsumer; } protected int getCurrentRow() { return currentRowSupplier.getAsInt(); } // It needs to be public so this method can be accessed when creating the complex types. public abstract Class getObjectClass(); @Override public boolean wasNull() { return wasNull; } @Override public String getString() throws SQLException { final Object object = getObject(); if (object == null) { return null; } return object.toString(); } @Override public boolean getBoolean() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public byte getByte() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public short getShort() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public int getInt() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public long getLong() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public float getFloat() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public double getDouble() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public BigDecimal getBigDecimal() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public BigDecimal getBigDecimal(final int i) throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public byte[] getBytes() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public InputStream getAsciiStream() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public InputStream getUnicodeStream() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public InputStream getBinaryStream() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public Object getObject() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public Reader getCharacterStream() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public Object getObject(final Map> map) throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public Ref getRef() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public Blob getBlob() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public Clob getClob() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public Array getArray() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public Struct getStruct() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public Date getDate(final Calendar calendar) throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public Time getTime(final Calendar calendar) throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public Timestamp getTimestamp(final Calendar calendar) throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public URL getURL() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public NClob getNClob() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public SQLXML getSQLXML() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public String getNString() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public Reader getNCharacterStream() throws SQLException { throw getOperationNotSupported(this.getClass()); } @Override public T getObject(final Class type) throws SQLException { final Object value; if (type == Byte.class) { value = getByte(); } else if (type == Short.class) { value = getShort(); } else if (type == Integer.class) { value = getInt(); } else if (type == Long.class) { value = getLong(); } else if (type == Float.class) { value = getFloat(); } else if (type == Double.class) { value = getDouble(); } else if (type == Boolean.class) { value = getBoolean(); } else if (type == BigDecimal.class) { value = getBigDecimal(); } else if (type == String.class) { value = getString(); } else if (type == byte[].class) { value = getBytes(); } else { value = getObject(); } return !type.isPrimitive() && wasNull ? null : type.cast(value); } private static SQLException getOperationNotSupported(final Class type) { return new SQLException(String.format("Operation not supported for type: %s.", type.getName())); } } ArrowFlightJdbcAccessorFactory.java000066400000000000000000000305111515451615200426620ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; 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; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeStampVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.complex.ArrowFlightJdbcDenseUnionVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.complex.ArrowFlightJdbcFixedSizeListVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.complex.ArrowFlightJdbcLargeListVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.complex.ArrowFlightJdbcListVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.complex.ArrowFlightJdbcMapVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.complex.ArrowFlightJdbcStructVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.complex.ArrowFlightJdbcUnionVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.numeric.ArrowFlightJdbcBaseIntVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.numeric.ArrowFlightJdbcBitVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.numeric.ArrowFlightJdbcDecimalVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.numeric.ArrowFlightJdbcFloat4VectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.numeric.ArrowFlightJdbcFloat8VectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.text.ArrowFlightJdbcVarCharVectorAccessor; 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.DurationVector; 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.IntervalDayVector; import org.apache.arrow.vector.IntervalMonthDayNanoVector; import org.apache.arrow.vector.IntervalYearVector; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.LargeVarCharVector; 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.TimeStampVector; 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.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; 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.UnionVector; /** Factory to instantiate the accessors. */ public class ArrowFlightJdbcAccessorFactory { /** * Create an accessor according to its type. * * @param vector an instance of an arrow vector. * @param getCurrentRow a supplier to check which row is being accessed. * @return an instance of one of the accessors. */ public static ArrowFlightJdbcAccessor createAccessor( ValueVector vector, IntSupplier getCurrentRow, WasNullConsumer setCursorWasNull) { if (vector instanceof UInt1Vector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (UInt1Vector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof UInt2Vector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (UInt2Vector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof UInt4Vector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (UInt4Vector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof UInt8Vector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (UInt8Vector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof TinyIntVector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (TinyIntVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof SmallIntVector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (SmallIntVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof IntVector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (IntVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof BigIntVector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (BigIntVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof Float4Vector) { return new ArrowFlightJdbcFloat4VectorAccessor( (Float4Vector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof Float8Vector) { return new ArrowFlightJdbcFloat8VectorAccessor( (Float8Vector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof BitVector) { return new ArrowFlightJdbcBitVectorAccessor( (BitVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof DecimalVector) { return new ArrowFlightJdbcDecimalVectorAccessor( (DecimalVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof Decimal256Vector) { return new ArrowFlightJdbcDecimalVectorAccessor( (Decimal256Vector) vector, getCurrentRow, setCursorWasNull); } 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); } else if (vector instanceof TimeStampVector) { return new ArrowFlightJdbcTimeStampVectorAccessor( (TimeStampVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof TimeNanoVector) { return new ArrowFlightJdbcTimeVectorAccessor( (TimeNanoVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof TimeMicroVector) { return new ArrowFlightJdbcTimeVectorAccessor( (TimeMicroVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof TimeMilliVector) { return new ArrowFlightJdbcTimeVectorAccessor( (TimeMilliVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof TimeSecVector) { return new ArrowFlightJdbcTimeVectorAccessor( (TimeSecVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof DateDayVector) { return new ArrowFlightJdbcDateVectorAccessor( ((DateDayVector) vector), getCurrentRow, setCursorWasNull); } else if (vector instanceof DateMilliVector) { return new ArrowFlightJdbcDateVectorAccessor( ((DateMilliVector) vector), getCurrentRow, setCursorWasNull); } else if (vector instanceof VarCharVector) { return new ArrowFlightJdbcVarCharVectorAccessor( (VarCharVector) vector, getCurrentRow, setCursorWasNull); } 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); } else if (vector instanceof IntervalDayVector) { return new ArrowFlightJdbcIntervalVectorAccessor( ((IntervalDayVector) vector), getCurrentRow, setCursorWasNull); } else if (vector instanceof IntervalYearVector) { return new ArrowFlightJdbcIntervalVectorAccessor( ((IntervalYearVector) vector), getCurrentRow, setCursorWasNull); } else if (vector instanceof IntervalMonthDayNanoVector) { return new ArrowFlightJdbcIntervalVectorAccessor( ((IntervalMonthDayNanoVector) vector), getCurrentRow, setCursorWasNull); } else if (vector instanceof StructVector) { return new ArrowFlightJdbcStructVectorAccessor( (StructVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof MapVector) { return new ArrowFlightJdbcMapVectorAccessor( (MapVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof ListVector) { return new ArrowFlightJdbcListVectorAccessor( (ListVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof LargeListVector) { return new ArrowFlightJdbcLargeListVectorAccessor( (LargeListVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof FixedSizeListVector) { return new ArrowFlightJdbcFixedSizeListVectorAccessor( (FixedSizeListVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof UnionVector) { return new ArrowFlightJdbcUnionVectorAccessor( (UnionVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof DenseUnionVector) { return new ArrowFlightJdbcDenseUnionVectorAccessor( (DenseUnionVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof NullVector || vector == null) { return new ArrowFlightJdbcNullVectorAccessor(setCursorWasNull); } throw new UnsupportedOperationException( "Unsupported vector type: " + vector.getClass().getName()); } /** Functional interface used to propagate that the value accessed was null or not. */ @FunctionalInterface public interface WasNullConsumer { void setWasNull(boolean wasNull); } } impl/000077500000000000000000000000001515451615200347325ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessorArrowFlightJdbcNullVectorAccessor.java000066400000000000000000000030641515451615200443140ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.vector.NullVector; /** Accessor for the Arrow type {@link NullVector}. */ public class ArrowFlightJdbcNullVectorAccessor extends ArrowFlightJdbcAccessor { public ArrowFlightJdbcNullVectorAccessor( ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(null, setCursorWasNull); } @Override public Class getObjectClass() { return Object.class; } @Override public boolean wasNull() { return true; } @Override public Object getObject() { this.wasNullConsumer.setWasNull(true); return null; } } binary/000077500000000000000000000000001515451615200362165ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/implArrowFlightJdbcBinaryVectorAccessor.java000066400000000000000000000103501515451615200461060ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.io.ByteArrayInputStream; import java.io.CharArrayReader; import java.io.InputStream; import java.io.Reader; import java.nio.charset.StandardCharsets; 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.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 * LargeVarBinaryVector}. */ public class ArrowFlightJdbcBinaryVectorAccessor extends ArrowFlightJdbcAccessor { private interface ByteArrayGetter { byte[] get(int index); } private final ByteArrayGetter getter; public ArrowFlightJdbcBinaryVectorAccessor( FixedSizeBinaryVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { this(vector::get, currentRowSupplier, setCursorWasNull); } public ArrowFlightJdbcBinaryVectorAccessor( VarBinaryVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { this(vector::get, currentRowSupplier, setCursorWasNull); } public ArrowFlightJdbcBinaryVectorAccessor( LargeVarBinaryVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { 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, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.getter = getter; } @Override public byte[] getBytes() { byte[] bytes = getter.get(getCurrentRow()); this.wasNull = bytes == null; this.wasNullConsumer.setWasNull(this.wasNull); return bytes; } @Override public Object getObject() { return this.getBytes(); } @Override public Class getObjectClass() { return byte[].class; } @Override public String getString() { byte[] bytes = this.getBytes(); if (bytes == null) { return null; } return new String(bytes, StandardCharsets.UTF_8); } @Override public InputStream getAsciiStream() { byte[] bytes = getBytes(); if (bytes == null) { return null; } return new ByteArrayInputStream(bytes); } @Override public InputStream getUnicodeStream() { byte[] bytes = getBytes(); if (bytes == null) { return null; } return new ByteArrayInputStream(bytes); } @Override public InputStream getBinaryStream() { byte[] bytes = getBytes(); if (bytes == null) { return null; } return new ByteArrayInputStream(bytes); } @Override public Reader getCharacterStream() { String string = getString(); if (string == null) { return null; } return new CharArrayReader(string.toCharArray()); } } ArrowFlightJdbcUuidVectorAccessor.java000066400000000000000000000055601515451615200455770ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } calendar/000077500000000000000000000000001515451615200365035ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/implArrowFlightJdbcDateVectorAccessor.java000066400000000000000000000123251515451615200460300ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.calendar; import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDateVectorGetter.Getter; import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDateVectorGetter.Holder; import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDateVectorGetter.createGetter; import static org.apache.arrow.driver.jdbc.utils.DateTimeUtils.getTimestampValue; import static org.apache.calcite.avatica.util.DateTimeUtils.MILLIS_PER_DAY; 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; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.driver.jdbc.utils.DateTimeUtils; import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DateMilliVector; import org.apache.arrow.vector.ValueVector; /** Accessor for the Arrow types: {@link DateDayVector} and {@link DateMilliVector}. */ public class ArrowFlightJdbcDateVectorAccessor extends ArrowFlightJdbcAccessor { private final Getter getter; private final TimeUnit timeUnit; private final Holder holder; /** * Instantiate an accessor for a {@link DateDayVector}. * * @param vector an instance of a DateDayVector. * @param currentRowSupplier the supplier to track the lines. * @param setCursorWasNull the consumer to set if value was null. */ public ArrowFlightJdbcDateVectorAccessor( DateDayVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.holder = new Holder(); this.getter = createGetter(vector); this.timeUnit = getTimeUnitForVector(vector); } /** * Instantiate an accessor for a {@link DateMilliVector}. * * @param vector an instance of a DateMilliVector. * @param currentRowSupplier the supplier to track the lines. */ public ArrowFlightJdbcDateVectorAccessor( DateMilliVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.holder = new Holder(); this.getter = createGetter(vector); this.timeUnit = getTimeUnitForVector(vector); } @Override public Class getObjectClass() { return Date.class; } @Override 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(); if (this.wasNull) { return null; } long value = holder.value; long milliseconds = this.timeUnit.toMillis(value); long millisWithCalendar = DateTimeUtils.applyCalendarOffset(milliseconds, calendar); return new Date(getTimestampValue(millisWithCalendar).getTime()); } private void fillHolder() { getter.get(getCurrentRow(), holder); this.wasNull = holder.isSet == 0; this.wasNullConsumer.setWasNull(this.wasNull); } @Override public Timestamp getTimestamp(Calendar calendar) { Date date = getDate(calendar); if (date == null) { return null; } return new Timestamp(date.getTime()); } @Override public String getString() { fillHolder(); if (wasNull) { return null; } long milliseconds = timeUnit.toMillis(holder.value); return unixDateToString((int) (milliseconds / MILLIS_PER_DAY)); } protected static TimeUnit getTimeUnitForVector(ValueVector vector) { if (vector instanceof DateDayVector) { return TimeUnit.DAYS; } else if (vector instanceof DateMilliVector) { return TimeUnit.MILLISECONDS; } throw new IllegalArgumentException("Invalid Arrow vector"); } private LocalDate getLocalDate() { return getDate(null).toLocalDate(); } } ArrowFlightJdbcDateVectorGetter.java000066400000000000000000000043601515451615200455200ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.calendar; import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DateMilliVector; import org.apache.arrow.vector.holders.NullableDateDayHolder; import org.apache.arrow.vector.holders.NullableDateMilliHolder; /** Auxiliary class used to unify data access on TimeStampVectors. */ final class ArrowFlightJdbcDateVectorGetter { private ArrowFlightJdbcDateVectorGetter() { // Prevent instantiation. } /** * Auxiliary class meant to unify Date*Vector#get implementations with different classes of * ValueHolders. */ static class Holder { int isSet; // Tells if value is set; 0 = not set, 1 = set long value; // Holds actual value in its respective timeunit } /** Functional interface used to unify Date*Vector#get implementations. */ @FunctionalInterface interface Getter { void get(int index, Holder holder); } static Getter createGetter(DateDayVector vector) { NullableDateDayHolder auxHolder = new NullableDateDayHolder(); return (index, holder) -> { vector.get(index, auxHolder); holder.isSet = auxHolder.isSet; holder.value = auxHolder.value; }; } static Getter createGetter(DateMilliVector vector) { NullableDateMilliHolder auxHolder = new NullableDateMilliHolder(); return (index, holder) -> { vector.get(index, auxHolder); holder.isSet = auxHolder.isSet; holder.value = auxHolder.value; }; } } ArrowFlightJdbcDurationVectorAccessor.java000066400000000000000000000035231515451615200467400ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.calendar; import java.time.Duration; 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.DurationVector; /** Accessor for the Arrow type {@link DurationVector}. */ public class ArrowFlightJdbcDurationVectorAccessor extends ArrowFlightJdbcAccessor { private final DurationVector vector; public ArrowFlightJdbcDurationVectorAccessor( DurationVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.vector = vector; } @Override public Class getObjectClass() { return Duration.class; } @Override public Object getObject() { Duration duration = vector.getObject(getCurrentRow()); this.wasNull = duration == null; this.wasNullConsumer.setWasNull(this.wasNull); return duration; } } ArrowFlightJdbcIntervalVectorAccessor.java000066400000000000000000000140051515451615200467340ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.calendar; import static org.apache.arrow.driver.jdbc.utils.IntervalStringUtils.formatIntervalDay; import static org.apache.arrow.driver.jdbc.utils.IntervalStringUtils.formatIntervalYear; import static org.apache.arrow.vector.util.DateUtility.yearsToMonths; import java.sql.SQLException; import java.time.Duration; import java.time.Period; 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.BaseFixedWidthVector; import org.apache.arrow.vector.IntervalDayVector; import org.apache.arrow.vector.IntervalMonthDayNanoVector; import org.apache.arrow.vector.IntervalYearVector; import org.apache.arrow.vector.PeriodDuration; import org.apache.arrow.vector.holders.NullableIntervalDayHolder; import org.apache.arrow.vector.holders.NullableIntervalMonthDayNanoHolder; import org.apache.arrow.vector.holders.NullableIntervalYearHolder; /** Accessor for the Arrow type {@link IntervalDayVector}. */ public class ArrowFlightJdbcIntervalVectorAccessor extends ArrowFlightJdbcAccessor { private final BaseFixedWidthVector vector; private final StringGetter stringGetter; private final Class objectClass; /** * Instantiate an accessor for a {@link IntervalDayVector}. * * @param vector an instance of a IntervalDayVector. * @param currentRowSupplier the supplier to track the rows. * @param setCursorWasNull the consumer to set if value was null. */ public ArrowFlightJdbcIntervalVectorAccessor( IntervalDayVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.vector = vector; stringGetter = (index) -> { final NullableIntervalDayHolder holder = new NullableIntervalDayHolder(); vector.get(index, holder); if (holder.isSet == 0) { return null; } else { final int days = holder.days; final int millis = holder.milliseconds; return formatIntervalDay(Duration.ofDays(days).plusMillis(millis)); } }; objectClass = java.time.Duration.class; } /** * Instantiate an accessor for a {@link IntervalYearVector}. * * @param vector an instance of a IntervalYearVector. * @param currentRowSupplier the supplier to track the rows. * @param setCursorWasNull the consumer to set if value was null. */ public ArrowFlightJdbcIntervalVectorAccessor( IntervalYearVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.vector = vector; stringGetter = (index) -> { final NullableIntervalYearHolder holder = new NullableIntervalYearHolder(); vector.get(index, holder); if (holder.isSet == 0) { return null; } else { final int interval = holder.value; final int years = (interval / yearsToMonths); final int months = (interval % yearsToMonths); return formatIntervalYear(Period.ofYears(years).plusMonths(months)); } }; objectClass = java.time.Period.class; } /** * Instantiate an accessor for a {@link IntervalMonthDayNanoVector}. * * @param vector an instance of a IntervalMonthDayNanoVector. * @param currentRowSupplier the supplier to track the rows. * @param setCursorWasNull the consumer to set if value was null. */ public ArrowFlightJdbcIntervalVectorAccessor( IntervalMonthDayNanoVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.vector = vector; stringGetter = (index) -> { final NullableIntervalMonthDayNanoHolder holder = new NullableIntervalMonthDayNanoHolder(); vector.get(index, holder); if (holder.isSet == 0) { return null; } else { final int months = holder.months; final int days = holder.days; final long nanos = holder.nanoseconds; final Period period = Period.ofMonths(months).plusDays(days); final Duration duration = Duration.ofNanos(nanos); return new PeriodDuration(period, duration).toISO8601IntervalString(); } }; objectClass = PeriodDuration.class; } @Override public Class getObjectClass() { return objectClass; } @Override public String getString() throws SQLException { String result = stringGetter.get(getCurrentRow()); wasNull = result == null; wasNullConsumer.setWasNull(wasNull); return result; } @Override public Object getObject() { Object object = vector.getObject(getCurrentRow()); wasNull = object == null; wasNullConsumer.setWasNull(wasNull); return object; } /** Functional interface used to unify Interval*Vector#getAsStringBuilder implementations. */ @FunctionalInterface interface StringGetter { String get(int index); } } ArrowFlightJdbcTimeStampVectorAccessor.java000066400000000000000000000227361515451615200470650ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.calendar; import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeStampVectorGetter.Getter; import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeStampVectorGetter.Holder; 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; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.vector.TimeStampVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.util.DateUtility; /** Accessor for the Arrow types extending from {@link TimeStampVector}. */ public class ArrowFlightJdbcTimeStampVectorAccessor extends ArrowFlightJdbcAccessor { private final TimeZone timeZone; private final Getter getter; 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 { LocalDateTime fromLong(long value); } /** Instantiate a ArrowFlightJdbcTimeStampVectorAccessor for given vector. */ public ArrowFlightJdbcTimeStampVectorAccessor( TimeStampVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); 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); } @Override 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; this.wasNullConsumer.setWasNull(this.wasNull); if (this.wasNull) { return null; } long value = holder.value; LocalDateTime localDateTime = this.longToLocalDateTime.fromLong(value); // 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 = localDateTime.minus( timeZone.getOffset(millis) - this.timeZone.getOffset(millis), ChronoUnit.MILLIS); } return localDateTime; } @Override public Date getDate(Calendar calendar) { LocalDateTime localDateTime = getLocalDateTime(calendar); if (localDateTime == null) { return null; } return new Date(getTimestampWithOffset(calendar, localDateTime).getTime()); } @Override public Time getTime(Calendar calendar) { LocalDateTime localDateTime = getLocalDateTime(calendar); if (localDateTime == null) { return null; } return new Time(getTimestampWithOffset(calendar, localDateTime).getTime()); } @Override public Timestamp getTimestamp(Calendar calendar) { LocalDateTime localDateTime = getLocalDateTime(calendar); if (localDateTime == null) { 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); } protected static TimeUnit getTimeUnitForVector(TimeStampVector vector) { ArrowType.Timestamp arrowType = (ArrowType.Timestamp) vector.getField().getFieldType().getType(); switch (arrowType.getUnit()) { case NANOSECOND: return TimeUnit.NANOSECONDS; case MICROSECOND: return TimeUnit.MICROSECONDS; case MILLISECOND: return TimeUnit.MILLISECONDS; case SECOND: return TimeUnit.SECONDS; default: throw new UnsupportedOperationException("Invalid Arrow time unit"); } } protected static LongToLocalDateTime getLongToLocalDateTimeForVector( TimeStampVector vector, TimeZone timeZone) { String timeZoneID = timeZone.getID(); ArrowType.Timestamp arrowType = (ArrowType.Timestamp) vector.getField().getFieldType().getType(); switch (arrowType.getUnit()) { case NANOSECOND: return nanoseconds -> DateUtility.getLocalDateTimeFromEpochNano(nanoseconds, timeZoneID); case MICROSECOND: return microseconds -> DateUtility.getLocalDateTimeFromEpochMicro(microseconds, timeZoneID); case MILLISECOND: return milliseconds -> DateUtility.getLocalDateTimeFromEpochMilli(milliseconds, timeZoneID); case SECOND: return seconds -> DateUtility.getLocalDateTimeFromEpochMilli( TimeUnit.SECONDS.toMillis(seconds), timeZoneID); default: throw new UnsupportedOperationException("Invalid Arrow time unit"); } } protected static TimeZone getTimeZoneForVector(TimeStampVector vector) { ArrowType.Timestamp arrowType = (ArrowType.Timestamp) vector.getField().getFieldType().getType(); 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; } } ArrowFlightJdbcTimeStampVectorGetter.java000066400000000000000000000136661515451615200465570ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.calendar; 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.TimeStampVector; import org.apache.arrow.vector.holders.NullableTimeStampMicroHolder; import org.apache.arrow.vector.holders.NullableTimeStampMicroTZHolder; import org.apache.arrow.vector.holders.NullableTimeStampMilliHolder; import org.apache.arrow.vector.holders.NullableTimeStampMilliTZHolder; import org.apache.arrow.vector.holders.NullableTimeStampNanoHolder; import org.apache.arrow.vector.holders.NullableTimeStampNanoTZHolder; import org.apache.arrow.vector.holders.NullableTimeStampSecHolder; import org.apache.arrow.vector.holders.NullableTimeStampSecTZHolder; /** Auxiliary class used to unify data access on TimeStampVectors. */ final class ArrowFlightJdbcTimeStampVectorGetter { private ArrowFlightJdbcTimeStampVectorGetter() { // Prevent instantiation. } /** * Auxiliary class meant to unify TimeStamp*Vector#get implementations with different classes of * ValueHolders. */ static class Holder { int isSet; // Tells if value is set; 0 = not set, 1 = set long value; // Holds actual value in its respective timeunit } /** Functional interface used to unify TimeStamp*Vector#get implementations. */ @FunctionalInterface interface Getter { void get(int index, Holder holder); } static Getter createGetter(TimeStampVector vector) { if (vector instanceof TimeStampNanoVector) { return createGetter((TimeStampNanoVector) vector); } else if (vector instanceof TimeStampNanoTZVector) { return createGetter((TimeStampNanoTZVector) vector); } else if (vector instanceof TimeStampMicroVector) { return createGetter((TimeStampMicroVector) vector); } else if (vector instanceof TimeStampMicroTZVector) { return createGetter((TimeStampMicroTZVector) vector); } else if (vector instanceof TimeStampMilliVector) { return createGetter((TimeStampMilliVector) vector); } else if (vector instanceof TimeStampMilliTZVector) { return createGetter((TimeStampMilliTZVector) vector); } else if (vector instanceof TimeStampSecVector) { return createGetter((TimeStampSecVector) vector); } else if (vector instanceof TimeStampSecTZVector) { return createGetter((TimeStampSecTZVector) vector); } throw new UnsupportedOperationException("Unsupported Timestamp vector type"); } private static Getter createGetter(TimeStampNanoVector vector) { NullableTimeStampNanoHolder auxHolder = new NullableTimeStampNanoHolder(); return (index, holder) -> { vector.get(index, auxHolder); holder.isSet = auxHolder.isSet; holder.value = auxHolder.value; }; } private static Getter createGetter(TimeStampNanoTZVector vector) { NullableTimeStampNanoTZHolder auxHolder = new NullableTimeStampNanoTZHolder(); return (index, holder) -> { vector.get(index, auxHolder); holder.isSet = auxHolder.isSet; holder.value = auxHolder.value; }; } private static Getter createGetter(TimeStampMicroVector vector) { NullableTimeStampMicroHolder auxHolder = new NullableTimeStampMicroHolder(); return (index, holder) -> { vector.get(index, auxHolder); holder.isSet = auxHolder.isSet; holder.value = auxHolder.value; }; } private static Getter createGetter(TimeStampMicroTZVector vector) { NullableTimeStampMicroTZHolder auxHolder = new NullableTimeStampMicroTZHolder(); return (index, holder) -> { vector.get(index, auxHolder); holder.isSet = auxHolder.isSet; holder.value = auxHolder.value; }; } private static Getter createGetter(TimeStampMilliVector vector) { NullableTimeStampMilliHolder auxHolder = new NullableTimeStampMilliHolder(); return (index, holder) -> { vector.get(index, auxHolder); holder.isSet = auxHolder.isSet; holder.value = auxHolder.value; }; } private static Getter createGetter(TimeStampMilliTZVector vector) { NullableTimeStampMilliTZHolder auxHolder = new NullableTimeStampMilliTZHolder(); return (index, holder) -> { vector.get(index, auxHolder); holder.isSet = auxHolder.isSet; holder.value = auxHolder.value; }; } private static Getter createGetter(TimeStampSecVector vector) { NullableTimeStampSecHolder auxHolder = new NullableTimeStampSecHolder(); return (index, holder) -> { vector.get(index, auxHolder); holder.isSet = auxHolder.isSet; holder.value = auxHolder.value; }; } private static Getter createGetter(TimeStampSecTZVector vector) { NullableTimeStampSecTZHolder auxHolder = new NullableTimeStampSecTZHolder(); return (index, holder) -> { vector.get(index, auxHolder); holder.isSet = auxHolder.isSet; holder.value = auxHolder.value; }; } } ArrowFlightJdbcTimeVectorAccessor.java000066400000000000000000000143261515451615200460540ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.calendar; import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeVectorGetter.Getter; 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; import org.apache.arrow.driver.jdbc.ArrowFlightJdbcTime; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.driver.jdbc.utils.DateTimeUtils; 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.ValueVector; /** * Accessor for the Arrow types: {@link TimeNanoVector}, {@link TimeMicroVector}, {@link * TimeMilliVector} and {@link TimeSecVector}. */ public class ArrowFlightJdbcTimeVectorAccessor extends ArrowFlightJdbcAccessor { private final Getter getter; private final TimeUnit timeUnit; private final Holder holder; /** * Instantiate an accessor for a {@link TimeNanoVector}. * * @param vector an instance of a TimeNanoVector. * @param currentRowSupplier the supplier to track the lines. * @param setCursorWasNull the consumer to set if value was null. */ public ArrowFlightJdbcTimeVectorAccessor( TimeNanoVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.holder = new Holder(); this.getter = createGetter(vector); this.timeUnit = getTimeUnitForVector(vector); } /** * Instantiate an accessor for a {@link TimeMicroVector}. * * @param vector an instance of a TimeMicroVector. * @param currentRowSupplier the supplier to track the lines. * @param setCursorWasNull the consumer to set if value was null. */ public ArrowFlightJdbcTimeVectorAccessor( TimeMicroVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.holder = new Holder(); this.getter = createGetter(vector); this.timeUnit = getTimeUnitForVector(vector); } /** * Instantiate an accessor for a {@link TimeMilliVector}. * * @param vector an instance of a TimeMilliVector. * @param currentRowSupplier the supplier to track the lines. */ public ArrowFlightJdbcTimeVectorAccessor( TimeMilliVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.holder = new Holder(); this.getter = createGetter(vector); this.timeUnit = getTimeUnitForVector(vector); } /** * Instantiate an accessor for a {@link TimeSecVector}. * * @param vector an instance of a TimeSecVector. * @param currentRowSupplier the supplier to track the lines. */ public ArrowFlightJdbcTimeVectorAccessor( TimeSecVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.holder = new Holder(); this.getter = createGetter(vector); this.timeUnit = getTimeUnitForVector(vector); } @Override public Class getObjectClass() { return Time.class; } @Override 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(); if (this.wasNull) { return null; } long value = holder.value; long milliseconds = this.timeUnit.toMillis(value); 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; this.wasNullConsumer.setWasNull(this.wasNull); } @Override public Timestamp getTimestamp(Calendar calendar) { Time time = getTime(calendar); if (time == null) { return null; } return new Timestamp(time.getTime()); } protected static TimeUnit getTimeUnitForVector(ValueVector vector) { if (vector instanceof TimeNanoVector) { return TimeUnit.NANOSECONDS; } else if (vector instanceof TimeMicroVector) { return TimeUnit.MICROSECONDS; } else if (vector instanceof TimeMilliVector) { return TimeUnit.MILLISECONDS; } else if (vector instanceof TimeSecVector) { return TimeUnit.SECONDS; } throw new IllegalArgumentException("Invalid Arrow vector"); } } ArrowFlightJdbcTimeVectorGetter.java000066400000000000000000000060051515451615200455370ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.calendar; 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.holders.NullableTimeMicroHolder; import org.apache.arrow.vector.holders.NullableTimeMilliHolder; import org.apache.arrow.vector.holders.NullableTimeNanoHolder; import org.apache.arrow.vector.holders.NullableTimeSecHolder; /** Auxiliary class used to unify data access on Time*Vectors. */ final class ArrowFlightJdbcTimeVectorGetter { private ArrowFlightJdbcTimeVectorGetter() { // Prevent instantiation. } /** * Auxiliary class meant to unify TimeStamp*Vector#get implementations with different classes of * ValueHolders. */ static class Holder { int isSet; // Tells if value is set; 0 = not set, 1 = set long value; // Holds actual value in its respective timeunit } /** Functional interface used to unify TimeStamp*Vector#get implementations. */ @FunctionalInterface interface Getter { void get(int index, Holder holder); } static Getter createGetter(TimeNanoVector vector) { NullableTimeNanoHolder auxHolder = new NullableTimeNanoHolder(); return (index, holder) -> { vector.get(index, auxHolder); holder.isSet = auxHolder.isSet; holder.value = auxHolder.value; }; } static Getter createGetter(TimeMicroVector vector) { NullableTimeMicroHolder auxHolder = new NullableTimeMicroHolder(); return (index, holder) -> { vector.get(index, auxHolder); holder.isSet = auxHolder.isSet; holder.value = auxHolder.value; }; } static Getter createGetter(TimeMilliVector vector) { NullableTimeMilliHolder auxHolder = new NullableTimeMilliHolder(); return (index, holder) -> { vector.get(index, auxHolder); holder.isSet = auxHolder.isSet; holder.value = auxHolder.value; }; } static Getter createGetter(TimeSecVector vector) { NullableTimeSecHolder auxHolder = new NullableTimeSecHolder(); return (index, holder) -> { vector.get(index, auxHolder); holder.isSet = auxHolder.isSet; holder.value = auxHolder.value; }; } } complex/000077500000000000000000000000001515451615200364015ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/implAbstractArrowFlightJdbcListVectorAccessor.java000066400000000000000000000050121515451615200474430ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex; import java.sql.Array; import java.util.List; import java.util.function.IntSupplier; import org.apache.arrow.driver.jdbc.ArrowFlightJdbcArray; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.LargeListVector; import org.apache.arrow.vector.complex.ListVector; /** * Base Accessor for the Arrow types {@link ListVector}, {@link LargeListVector} and {@link * FixedSizeListVector}. */ public abstract class AbstractArrowFlightJdbcListVectorAccessor extends ArrowFlightJdbcAccessor { protected AbstractArrowFlightJdbcListVectorAccessor( IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); } @Override public Class getObjectClass() { return List.class; } protected abstract long getStartOffset(int index); protected abstract long getEndOffset(int index); protected abstract FieldVector getDataVector(); protected abstract boolean isNull(int index); @Override public final Array getArray() { int index = getCurrentRow(); FieldVector dataVector = getDataVector(); this.wasNull = isNull(index); this.wasNullConsumer.setWasNull(this.wasNull); if (this.wasNull) { return null; } long startOffset = getStartOffset(index); long endOffset = getEndOffset(index); long valuesCount = endOffset - startOffset; return new ArrowFlightJdbcArray(dataVector, startOffset, valuesCount); } } AbstractArrowFlightJdbcUnionVectorAccessor.java000066400000000000000000000155721515451615200476340ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex; import java.io.InputStream; import java.io.Reader; import java.math.BigDecimal; import java.net.URL; import java.sql.Array; import java.sql.Blob; import java.sql.Clob; import java.sql.Date; import java.sql.NClob; import java.sql.Ref; import java.sql.SQLException; import java.sql.SQLXML; import java.sql.Struct; import java.sql.Time; import java.sql.Timestamp; import java.util.Calendar; import java.util.Map; 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.driver.jdbc.accessor.impl.ArrowFlightJdbcNullVectorAccessor; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.UnionVector; /** Base accessor for {@link UnionVector} and {@link DenseUnionVector}. */ public abstract class AbstractArrowFlightJdbcUnionVectorAccessor extends ArrowFlightJdbcAccessor { /** * Array of accessors for each type contained in UnionVector. Index corresponds to UnionVector and * DenseUnionVector typeIds which are both limited to 128. */ private final ArrowFlightJdbcAccessor[] accessors = new ArrowFlightJdbcAccessor[128]; private final ArrowFlightJdbcNullVectorAccessor nullAccessor = new ArrowFlightJdbcNullVectorAccessor((boolean wasNull) -> {}); protected AbstractArrowFlightJdbcUnionVectorAccessor( IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); } protected abstract ArrowFlightJdbcAccessor createAccessorForVector(ValueVector vector); protected abstract byte getCurrentTypeId(); protected abstract ValueVector getVectorByTypeId(byte typeId); /** * Returns an accessor for UnionVector child vector on current row. * * @return ArrowFlightJdbcAccessor for child vector on current row. */ protected ArrowFlightJdbcAccessor getAccessor() { // Get the typeId and child vector for the current row being accessed. byte typeId = this.getCurrentTypeId(); ValueVector vector = this.getVectorByTypeId(typeId); if (typeId < 0) { // typeId may be negative if the current row has no type defined. return this.nullAccessor; } // Ensure there is an accessor for given typeId if (this.accessors[typeId] == null) { this.accessors[typeId] = this.createAccessorForVector(vector); } return this.accessors[typeId]; } @Override public Class getObjectClass() { return getAccessor().getObjectClass(); } @Override public boolean wasNull() { return getAccessor().wasNull(); } @Override public String getString() throws SQLException { return getAccessor().getString(); } @Override public boolean getBoolean() throws SQLException { return getAccessor().getBoolean(); } @Override public byte getByte() throws SQLException { return getAccessor().getByte(); } @Override public short getShort() throws SQLException { return getAccessor().getShort(); } @Override public int getInt() throws SQLException { return getAccessor().getInt(); } @Override public long getLong() throws SQLException { return getAccessor().getLong(); } @Override public float getFloat() throws SQLException { return getAccessor().getFloat(); } @Override public double getDouble() throws SQLException { return getAccessor().getDouble(); } @Override public BigDecimal getBigDecimal() throws SQLException { return getAccessor().getBigDecimal(); } @Override public BigDecimal getBigDecimal(int i) throws SQLException { return getAccessor().getBigDecimal(i); } @Override public byte[] getBytes() throws SQLException { return getAccessor().getBytes(); } @Override public InputStream getAsciiStream() throws SQLException { return getAccessor().getAsciiStream(); } @Override public InputStream getUnicodeStream() throws SQLException { return getAccessor().getUnicodeStream(); } @Override public InputStream getBinaryStream() throws SQLException { return getAccessor().getBinaryStream(); } @Override public Object getObject() throws SQLException { return getAccessor().getObject(); } @Override public Reader getCharacterStream() throws SQLException { return getAccessor().getCharacterStream(); } @Override public Object getObject(Map> map) throws SQLException { return getAccessor().getObject(map); } @Override public Ref getRef() throws SQLException { return getAccessor().getRef(); } @Override public Blob getBlob() throws SQLException { return getAccessor().getBlob(); } @Override public Clob getClob() throws SQLException { return getAccessor().getClob(); } @Override public Array getArray() throws SQLException { return getAccessor().getArray(); } @Override public Struct getStruct() throws SQLException { return getAccessor().getStruct(); } @Override public Date getDate(Calendar calendar) throws SQLException { return getAccessor().getDate(calendar); } @Override public Time getTime(Calendar calendar) throws SQLException { return getAccessor().getTime(calendar); } @Override public Timestamp getTimestamp(Calendar calendar) throws SQLException { return getAccessor().getTimestamp(calendar); } @Override public URL getURL() throws SQLException { return getAccessor().getURL(); } @Override public NClob getNClob() throws SQLException { return getAccessor().getNClob(); } @Override public SQLXML getSQLXML() throws SQLException { return getAccessor().getSQLXML(); } @Override public String getNString() throws SQLException { return getAccessor().getNString(); } @Override public Reader getNCharacterStream() throws SQLException { return getAccessor().getNCharacterStream(); } @Override public T getObject(Class type) throws SQLException { return getAccessor().getObject(type); } } ArrowFlightJdbcDenseUnionVectorAccessor.java000066400000000000000000000045411515451615200471210ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex; 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.ValueVector; import org.apache.arrow.vector.complex.DenseUnionVector; /** Accessor for the Arrow type {@link DenseUnionVector}. */ public class ArrowFlightJdbcDenseUnionVectorAccessor extends AbstractArrowFlightJdbcUnionVectorAccessor { private final DenseUnionVector vector; /** * Instantiate an accessor for a {@link DenseUnionVector}. * * @param vector an instance of a DenseUnionVector. * @param currentRowSupplier the supplier to track the rows. * @param setCursorWasNull the consumer to set if value was null. */ public ArrowFlightJdbcDenseUnionVectorAccessor( DenseUnionVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.vector = vector; } @Override protected ArrowFlightJdbcAccessor createAccessorForVector(ValueVector vector) { return ArrowFlightJdbcAccessorFactory.createAccessor( vector, () -> this.vector.getOffset(this.getCurrentRow()), (boolean wasNull) -> {}); } @Override protected byte getCurrentTypeId() { int index = getCurrentRow(); return this.vector.getTypeId(index); } @Override protected ValueVector getVectorByTypeId(byte typeId) { return this.vector.getVectorByType(typeId); } } ArrowFlightJdbcFixedSizeListVectorAccessor.java000066400000000000000000000042601515451615200475760ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex; import java.util.List; import java.util.function.IntSupplier; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.complex.FixedSizeListVector; /** Accessor for the Arrow type {@link FixedSizeListVector}. */ public class ArrowFlightJdbcFixedSizeListVectorAccessor extends AbstractArrowFlightJdbcListVectorAccessor { private final FixedSizeListVector vector; public ArrowFlightJdbcFixedSizeListVectorAccessor( FixedSizeListVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.vector = vector; } @Override protected long getStartOffset(int index) { return (long) vector.getListSize() * index; } @Override protected long getEndOffset(int index) { return (long) vector.getListSize() * (index + 1); } @Override protected FieldVector getDataVector() { return vector.getDataVector(); } @Override protected boolean isNull(int index) { return vector.isNull(index); } @Override public Object getObject() { List object = vector.getObject(getCurrentRow()); this.wasNull = object == null; this.wasNullConsumer.setWasNull(this.wasNull); return object; } } ArrowFlightJdbcLargeListVectorAccessor.java000066400000000000000000000043541515451615200467420ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex; import java.util.List; import java.util.function.IntSupplier; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.complex.LargeListVector; /** Accessor for the Arrow type {@link LargeListVector}. */ public class ArrowFlightJdbcLargeListVectorAccessor extends AbstractArrowFlightJdbcListVectorAccessor { private final LargeListVector vector; public ArrowFlightJdbcLargeListVectorAccessor( LargeListVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.vector = vector; } @Override protected long getStartOffset(int index) { return vector.getOffsetBuffer().getLong((long) index * LargeListVector.OFFSET_WIDTH); } @Override protected long getEndOffset(int index) { return vector.getOffsetBuffer().getLong((long) (index + 1) * LargeListVector.OFFSET_WIDTH); } @Override protected FieldVector getDataVector() { return vector.getDataVector(); } @Override protected boolean isNull(int index) { return vector.isNull(index); } @Override public Object getObject() { List object = vector.getObject(getCurrentRow()); this.wasNull = object == null; this.wasNullConsumer.setWasNull(this.wasNull); return object; } } ArrowFlightJdbcListVectorAccessor.java000066400000000000000000000044521515451615200457660ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex; import java.util.List; import java.util.function.IntSupplier; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.complex.BaseRepeatedValueVector; import org.apache.arrow.vector.complex.ListVector; /** Accessor for the Arrow type {@link ListVector}. */ public class ArrowFlightJdbcListVectorAccessor extends AbstractArrowFlightJdbcListVectorAccessor { private final ListVector vector; public ArrowFlightJdbcListVectorAccessor( ListVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.vector = vector; } @Override protected long getStartOffset(int index) { return vector.getOffsetBuffer().getInt((long) index * BaseRepeatedValueVector.OFFSET_WIDTH); } @Override protected long getEndOffset(int index) { return vector .getOffsetBuffer() .getInt((long) (index + 1) * BaseRepeatedValueVector.OFFSET_WIDTH); } @Override protected FieldVector getDataVector() { return vector.getDataVector(); } @Override protected boolean isNull(int index) { return vector.isNull(index); } @Override public Object getObject() { List object = vector.getObject(getCurrentRow()); this.wasNull = object == null; this.wasNullConsumer.setWasNull(this.wasNull); return object; } } ArrowFlightJdbcMapVectorAccessor.java000066400000000000000000000054611515451615200455710ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex; import java.util.Map; import java.util.function.IntSupplier; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.complex.BaseRepeatedValueVector; import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.impl.UnionMapReader; import org.apache.arrow.vector.util.JsonStringHashMap; /** Accessor for the Arrow type {@link MapVector}. */ public class ArrowFlightJdbcMapVectorAccessor extends AbstractArrowFlightJdbcListVectorAccessor { private final MapVector vector; public ArrowFlightJdbcMapVectorAccessor( MapVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.vector = vector; } @Override public Class getObjectClass() { return Map.class; } @Override public Object getObject() { int index = getCurrentRow(); this.wasNull = vector.isNull(index); this.wasNullConsumer.setWasNull(this.wasNull); if (this.wasNull) { return null; } Map result = new JsonStringHashMap<>(); UnionMapReader reader = vector.getReader(); reader.setPosition(index); while (reader.next()) { Object key = reader.key().readObject(); Object value = reader.value().readObject(); result.put(key, value); } return result; } @Override protected long getStartOffset(int index) { return vector.getOffsetBuffer().getInt((long) index * BaseRepeatedValueVector.OFFSET_WIDTH); } @Override protected long getEndOffset(int index) { return vector .getOffsetBuffer() .getInt((long) (index + 1) * BaseRepeatedValueVector.OFFSET_WIDTH); } @Override protected boolean isNull(int index) { return vector.isNull(index); } @Override protected FieldVector getDataVector() { return vector.getDataVector(); } } ArrowFlightJdbcStructVectorAccessor.java000066400000000000000000000046151515451615200463400ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex; import java.sql.Struct; import java.util.List; import java.util.Map; import java.util.function.IntSupplier; import java.util.stream.Collectors; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.vector.complex.StructVector; import org.apache.calcite.avatica.util.StructImpl; /** Accessor for the Arrow type {@link StructVector}. */ public class ArrowFlightJdbcStructVectorAccessor extends ArrowFlightJdbcAccessor { private final StructVector vector; public ArrowFlightJdbcStructVectorAccessor( StructVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.vector = vector; } @Override public Class getObjectClass() { return Map.class; } @Override public Object getObject() { Map object = vector.getObject(getCurrentRow()); this.wasNull = object == null; this.wasNullConsumer.setWasNull(this.wasNull); return object; } @Override public Struct getStruct() { int currentRow = getCurrentRow(); this.wasNull = vector.isNull(currentRow); this.wasNullConsumer.setWasNull(this.wasNull); if (this.wasNull) { return null; } List attributes = vector.getChildrenFromFields().stream() .map(vector -> vector.getObject(currentRow)) .collect(Collectors.toList()); return new StructImpl(attributes); } } ArrowFlightJdbcUnionVectorAccessor.java000066400000000000000000000044411515451615200461410ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex; 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.ValueVector; import org.apache.arrow.vector.complex.UnionVector; /** Accessor for the Arrow type {@link UnionVector}. */ public class ArrowFlightJdbcUnionVectorAccessor extends AbstractArrowFlightJdbcUnionVectorAccessor { private final UnionVector vector; /** * Instantiate an accessor for a {@link UnionVector}. * * @param vector an instance of a UnionVector. * @param currentRowSupplier the supplier to track the rows. * @param setCursorWasNull the consumer to set if value was null. */ public ArrowFlightJdbcUnionVectorAccessor( UnionVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.vector = vector; } @Override protected ArrowFlightJdbcAccessor createAccessorForVector(ValueVector vector) { return ArrowFlightJdbcAccessorFactory.createAccessor( vector, this::getCurrentRow, (boolean wasNull) -> {}); } @Override protected byte getCurrentTypeId() { int index = getCurrentRow(); return (byte) this.vector.getTypeValue(index); } @Override protected ValueVector getVectorByTypeId(byte typeId) { return this.vector.getVectorByType(typeId); } } numeric/000077500000000000000000000000001515451615200363745ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/implArrowFlightJdbcBaseIntVectorAccessor.java000066400000000000000000000147531515451615200464000ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/numeric/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.numeric; import static org.apache.arrow.driver.jdbc.accessor.impl.numeric.ArrowFlightJdbcNumericGetter.Getter; import static org.apache.arrow.driver.jdbc.accessor.impl.numeric.ArrowFlightJdbcNumericGetter.createGetter; import java.math.BigDecimal; import java.math.RoundingMode; 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.driver.jdbc.accessor.impl.numeric.ArrowFlightJdbcNumericGetter.NumericHolder; import org.apache.arrow.vector.BaseIntVector; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.SmallIntVector; 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.types.Types.MinorType; /** * Accessor for the arrow types: TinyIntVector, SmallIntVector, IntVector, BigIntVector, * UInt1Vector, UInt2Vector, UInt4Vector and UInt8Vector. */ public class ArrowFlightJdbcBaseIntVectorAccessor extends ArrowFlightJdbcAccessor { private final MinorType type; private final boolean isUnsigned; private final Getter getter; private final NumericHolder holder; public ArrowFlightJdbcBaseIntVectorAccessor( UInt1Vector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { this(vector, currentRowSupplier, true, setCursorWasNull); } public ArrowFlightJdbcBaseIntVectorAccessor( UInt2Vector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { this(vector, currentRowSupplier, true, setCursorWasNull); } public ArrowFlightJdbcBaseIntVectorAccessor( UInt4Vector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { this(vector, currentRowSupplier, true, setCursorWasNull); } public ArrowFlightJdbcBaseIntVectorAccessor( UInt8Vector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { this(vector, currentRowSupplier, true, setCursorWasNull); } public ArrowFlightJdbcBaseIntVectorAccessor( TinyIntVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { this(vector, currentRowSupplier, false, setCursorWasNull); } public ArrowFlightJdbcBaseIntVectorAccessor( SmallIntVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { this(vector, currentRowSupplier, false, setCursorWasNull); } public ArrowFlightJdbcBaseIntVectorAccessor( IntVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { this(vector, currentRowSupplier, false, setCursorWasNull); } public ArrowFlightJdbcBaseIntVectorAccessor( BigIntVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { this(vector, currentRowSupplier, false, setCursorWasNull); } private ArrowFlightJdbcBaseIntVectorAccessor( BaseIntVector vector, IntSupplier currentRowSupplier, boolean isUnsigned, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.type = vector.getMinorType(); this.holder = new NumericHolder(); this.getter = createGetter(vector); this.isUnsigned = isUnsigned; } @Override public long getLong() { getter.get(getCurrentRow(), holder); this.wasNull = holder.isSet == 0; this.wasNullConsumer.setWasNull(this.wasNull); if (this.wasNull) { return 0; } return holder.value; } @Override public Class getObjectClass() { return Long.class; } @Override public String getString() { final long number = getLong(); if (this.wasNull) { return null; } else { return isUnsigned ? Long.toUnsignedString(number) : Long.toString(number); } } @Override public byte getByte() { return (byte) getLong(); } @Override public short getShort() { return (short) getLong(); } @Override public int getInt() { return (int) getLong(); } @Override public float getFloat() { return (float) getLong(); } @Override public double getDouble() { return (double) getLong(); } @Override public BigDecimal getBigDecimal() { final BigDecimal value = BigDecimal.valueOf(getLong()); return this.wasNull ? null : value; } @Override public BigDecimal getBigDecimal(int scale) { final BigDecimal value = BigDecimal.valueOf(this.getDouble()).setScale(scale, RoundingMode.HALF_UP); return this.wasNull ? null : value; } @Override public Number getObject() { final Number number; switch (type) { case TINYINT: case UINT1: number = getByte(); break; case SMALLINT: case UINT2: number = getShort(); break; case INT: case UINT4: number = getInt(); break; case BIGINT: case UINT8: number = getLong(); break; default: throw new IllegalStateException("No valid MinorType was provided."); } return wasNull ? null : number; } @Override public boolean getBoolean() { final long value = getLong(); return value != 0; } } ArrowFlightJdbcBitVectorAccessor.java000066400000000000000000000061321515451615200455610ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/numeric/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.numeric; import java.math.BigDecimal; 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.BitVector; import org.apache.arrow.vector.holders.NullableBitHolder; /** Accessor for the arrow {@link BitVector}. */ public class ArrowFlightJdbcBitVectorAccessor extends ArrowFlightJdbcAccessor { private final BitVector vector; private final NullableBitHolder holder; /** * Constructor for the BitVectorAccessor. * * @param vector an instance of a {@link BitVector}. * @param currentRowSupplier a supplier to check which row is being accessed. * @param setCursorWasNull the consumer to set if value was null. */ public ArrowFlightJdbcBitVectorAccessor( BitVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.vector = vector; this.holder = new NullableBitHolder(); } @Override public Class getObjectClass() { return Boolean.class; } @Override public String getString() { final boolean value = getBoolean(); return wasNull ? null : Boolean.toString(value); } @Override public boolean getBoolean() { return this.getLong() != 0; } @Override public byte getByte() { return (byte) this.getLong(); } @Override public short getShort() { return (short) this.getLong(); } @Override public int getInt() { return (int) this.getLong(); } @Override public long getLong() { vector.get(getCurrentRow(), holder); this.wasNull = holder.isSet == 0; this.wasNullConsumer.setWasNull(this.wasNull); if (this.wasNull) { return 0; } return holder.value; } @Override public float getFloat() { return this.getLong(); } @Override public double getDouble() { return this.getLong(); } @Override public BigDecimal getBigDecimal() { final long value = this.getLong(); return this.wasNull ? null : BigDecimal.valueOf(value); } @Override public Object getObject() { final boolean value = this.getBoolean(); return this.wasNull ? null : value; } } ArrowFlightJdbcDecimalVectorAccessor.java000066400000000000000000000074551515451615200464120ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/numeric/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.numeric; import java.math.BigDecimal; import java.math.RoundingMode; 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.Decimal256Vector; import org.apache.arrow.vector.DecimalVector; /** Accessor for {@link DecimalVector} and {@link Decimal256Vector}. */ public class ArrowFlightJdbcDecimalVectorAccessor extends ArrowFlightJdbcAccessor { private final Getter getter; /** Functional interface used to unify Decimal*Vector#getObject implementations. */ @FunctionalInterface interface Getter { BigDecimal getObject(int index); } public ArrowFlightJdbcDecimalVectorAccessor( DecimalVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.getter = vector::getObject; } public ArrowFlightJdbcDecimalVectorAccessor( Decimal256Vector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.getter = vector::getObject; } @Override public Class getObjectClass() { return BigDecimal.class; } @Override public BigDecimal getBigDecimal() { final BigDecimal value = getter.getObject(getCurrentRow()); this.wasNull = value == null; this.wasNullConsumer.setWasNull(this.wasNull); return value; } @Override public String getString() { final BigDecimal value = this.getBigDecimal(); return this.wasNull ? null : value.toString(); } @Override public boolean getBoolean() { final BigDecimal value = this.getBigDecimal(); return !this.wasNull && !value.equals(BigDecimal.ZERO); } @Override public byte getByte() { final BigDecimal value = this.getBigDecimal(); return this.wasNull ? 0 : value.byteValue(); } @Override public short getShort() { final BigDecimal value = this.getBigDecimal(); return this.wasNull ? 0 : value.shortValue(); } @Override public int getInt() { final BigDecimal value = this.getBigDecimal(); return this.wasNull ? 0 : value.intValue(); } @Override public long getLong() { final BigDecimal value = this.getBigDecimal(); return this.wasNull ? 0 : value.longValue(); } @Override public float getFloat() { final BigDecimal value = this.getBigDecimal(); return this.wasNull ? 0 : value.floatValue(); } @Override public double getDouble() { final BigDecimal value = this.getBigDecimal(); return this.wasNull ? 0 : value.doubleValue(); } @Override public BigDecimal getBigDecimal(int scale) { final BigDecimal value = this.getBigDecimal(); return this.wasNull ? null : value.setScale(scale, RoundingMode.HALF_UP); } @Override public Object getObject() { return this.getBigDecimal(); } } ArrowFlightJdbcFloat4VectorAccessor.java000066400000000000000000000072631515451615200462020ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/numeric/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.numeric; import java.math.BigDecimal; import java.math.RoundingMode; import java.sql.SQLException; 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.Float4Vector; import org.apache.arrow.vector.holders.NullableFloat4Holder; /** Accessor for the Float4Vector. */ public class ArrowFlightJdbcFloat4VectorAccessor extends ArrowFlightJdbcAccessor { private final Float4Vector vector; private final NullableFloat4Holder holder; /** * Instantiate a accessor for the {@link Float4Vector}. * * @param vector an instance of a Float4Vector. * @param currentRowSupplier the supplier to track the lines. * @param setCursorWasNull the consumer to set if value was null. */ public ArrowFlightJdbcFloat4VectorAccessor( Float4Vector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.holder = new NullableFloat4Holder(); this.vector = vector; } @Override public Class getObjectClass() { return Float.class; } @Override public String getString() { final float value = this.getFloat(); return this.wasNull ? null : Float.toString(value); } @Override public boolean getBoolean() { return this.getFloat() != 0.0; } @Override public byte getByte() { return (byte) this.getFloat(); } @Override public short getShort() { return (short) this.getFloat(); } @Override public int getInt() { return (int) this.getFloat(); } @Override public long getLong() { return (long) this.getFloat(); } @Override public float getFloat() { vector.get(getCurrentRow(), holder); this.wasNull = holder.isSet == 0; this.wasNullConsumer.setWasNull(this.wasNull); if (this.wasNull) { return 0; } return holder.value; } @Override public double getDouble() { return this.getFloat(); } @Override public BigDecimal getBigDecimal() throws SQLException { final float value = this.getFloat(); if (Float.isInfinite(value) || Float.isNaN(value)) { throw new SQLException("BigDecimal doesn't support Infinite/NaN."); } return this.wasNull ? null : BigDecimal.valueOf(value); } @Override public BigDecimal getBigDecimal(int scale) throws SQLException { final float value = this.getFloat(); if (Float.isInfinite(value) || Float.isNaN(value)) { throw new SQLException("BigDecimal doesn't support Infinite/NaN."); } return this.wasNull ? null : BigDecimal.valueOf(value).setScale(scale, RoundingMode.HALF_UP); } @Override public Object getObject() { final float value = this.getFloat(); return this.wasNull ? null : value; } } ArrowFlightJdbcFloat8VectorAccessor.java000066400000000000000000000073151515451615200462040ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/numeric/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.numeric; import java.math.BigDecimal; import java.math.RoundingMode; import java.sql.SQLException; 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.Float8Vector; import org.apache.arrow.vector.holders.NullableFloat8Holder; /** Accessor for the Float8Vector. */ public class ArrowFlightJdbcFloat8VectorAccessor extends ArrowFlightJdbcAccessor { private final Float8Vector vector; private final NullableFloat8Holder holder; /** * Instantiate a accessor for the {@link Float8Vector}. * * @param vector an instance of a Float8Vector. * @param currentRowSupplier the supplier to track the lines. * @param setCursorWasNull the consumer to set if value was null. */ public ArrowFlightJdbcFloat8VectorAccessor( Float8Vector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.holder = new NullableFloat8Holder(); this.vector = vector; } @Override public Class getObjectClass() { return Double.class; } @Override public double getDouble() { vector.get(getCurrentRow(), holder); this.wasNull = holder.isSet == 0; this.wasNullConsumer.setWasNull(this.wasNull); if (this.wasNull) { return 0; } return holder.value; } @Override public Object getObject() { final double value = this.getDouble(); return this.wasNull ? null : value; } @Override public String getString() { final double value = this.getDouble(); return this.wasNull ? null : Double.toString(value); } @Override public boolean getBoolean() { return this.getDouble() != 0.0; } @Override public byte getByte() { return (byte) this.getDouble(); } @Override public short getShort() { return (short) this.getDouble(); } @Override public int getInt() { return (int) this.getDouble(); } @Override public long getLong() { return (long) this.getDouble(); } @Override public float getFloat() { return (float) this.getDouble(); } @Override public BigDecimal getBigDecimal() throws SQLException { final double value = this.getDouble(); if (Double.isInfinite(value) || Double.isNaN(value)) { throw new SQLException("BigDecimal doesn't support Infinite/NaN."); } return this.wasNull ? null : BigDecimal.valueOf(value); } @Override public BigDecimal getBigDecimal(int scale) throws SQLException { final double value = this.getDouble(); if (Double.isInfinite(value) || Double.isNaN(value)) { throw new SQLException("BigDecimal doesn't support Infinite/NaN."); } return this.wasNull ? null : BigDecimal.valueOf(value).setScale(scale, RoundingMode.HALF_UP); } } ArrowFlightJdbcNumericGetter.java000066400000000000000000000157301515451615200447560ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/numeric/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.numeric; import org.apache.arrow.vector.BaseIntVector; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.SmallIntVector; 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.holders.NullableBigIntHolder; import org.apache.arrow.vector.holders.NullableIntHolder; import org.apache.arrow.vector.holders.NullableSmallIntHolder; import org.apache.arrow.vector.holders.NullableTinyIntHolder; import org.apache.arrow.vector.holders.NullableUInt1Holder; import org.apache.arrow.vector.holders.NullableUInt2Holder; import org.apache.arrow.vector.holders.NullableUInt4Holder; import org.apache.arrow.vector.holders.NullableUInt8Holder; /** A custom getter for values from the {@link BaseIntVector}. */ class ArrowFlightJdbcNumericGetter { /** A holder for values from the {@link BaseIntVector}. */ static class NumericHolder { int isSet; // Tells if value is set; 0 = not set, 1 = set long value; // Holds actual value } /** Functional interface for a getter to baseInt values. */ @FunctionalInterface interface Getter { void get(int index, NumericHolder holder); } /** * Main class that will check the type of the vector to create a specific getter. * * @param vector an instance of the {@link BaseIntVector} * @return a getter. */ static Getter createGetter(BaseIntVector vector) { if (vector instanceof UInt1Vector) { return createGetter((UInt1Vector) vector); } else if (vector instanceof UInt2Vector) { return createGetter((UInt2Vector) vector); } else if (vector instanceof UInt4Vector) { return createGetter((UInt4Vector) vector); } else if (vector instanceof UInt8Vector) { return createGetter((UInt8Vector) vector); } else if (vector instanceof TinyIntVector) { return createGetter((TinyIntVector) vector); } else if (vector instanceof SmallIntVector) { return createGetter((SmallIntVector) vector); } else if (vector instanceof IntVector) { return createGetter((IntVector) vector); } else if (vector instanceof BigIntVector) { return createGetter((BigIntVector) vector); } throw new UnsupportedOperationException("No valid IntVector was provided."); } /** * Create a specific getter for {@link UInt1Vector}. * * @param vector an instance of the {@link UInt1Vector} * @return a getter. */ private static Getter createGetter(UInt1Vector vector) { NullableUInt1Holder nullableUInt1Holder = new NullableUInt1Holder(); return (index, holder) -> { vector.get(index, nullableUInt1Holder); holder.isSet = nullableUInt1Holder.isSet; holder.value = nullableUInt1Holder.value; }; } /** * Create a specific getter for {@link UInt2Vector}. * * @param vector an instance of the {@link UInt2Vector} * @return a getter. */ private static Getter createGetter(UInt2Vector vector) { NullableUInt2Holder nullableUInt2Holder = new NullableUInt2Holder(); return (index, holder) -> { vector.get(index, nullableUInt2Holder); holder.isSet = nullableUInt2Holder.isSet; holder.value = nullableUInt2Holder.value; }; } /** * Create a specific getter for {@link UInt4Vector}. * * @param vector an instance of the {@link UInt4Vector} * @return a getter. */ private static Getter createGetter(UInt4Vector vector) { NullableUInt4Holder nullableUInt4Holder = new NullableUInt4Holder(); return (index, holder) -> { vector.get(index, nullableUInt4Holder); holder.isSet = nullableUInt4Holder.isSet; holder.value = nullableUInt4Holder.value; }; } /** * Create a specific getter for {@link UInt8Vector}. * * @param vector an instance of the {@link UInt8Vector} * @return a getter. */ private static Getter createGetter(UInt8Vector vector) { NullableUInt8Holder nullableUInt8Holder = new NullableUInt8Holder(); return (index, holder) -> { vector.get(index, nullableUInt8Holder); holder.isSet = nullableUInt8Holder.isSet; holder.value = nullableUInt8Holder.value; }; } /** * Create a specific getter for {@link TinyIntVector}. * * @param vector an instance of the {@link TinyIntVector} * @return a getter. */ private static Getter createGetter(TinyIntVector vector) { NullableTinyIntHolder nullableTinyIntHolder = new NullableTinyIntHolder(); return (index, holder) -> { vector.get(index, nullableTinyIntHolder); holder.isSet = nullableTinyIntHolder.isSet; holder.value = nullableTinyIntHolder.value; }; } /** * Create a specific getter for {@link SmallIntVector}. * * @param vector an instance of the {@link SmallIntVector} * @return a getter. */ private static Getter createGetter(SmallIntVector vector) { NullableSmallIntHolder nullableSmallIntHolder = new NullableSmallIntHolder(); return (index, holder) -> { vector.get(index, nullableSmallIntHolder); holder.isSet = nullableSmallIntHolder.isSet; holder.value = nullableSmallIntHolder.value; }; } /** * Create a specific getter for {@link IntVector}. * * @param vector an instance of the {@link IntVector} * @return a getter. */ private static Getter createGetter(IntVector vector) { NullableIntHolder nullableIntHolder = new NullableIntHolder(); return (index, holder) -> { vector.get(index, nullableIntHolder); holder.isSet = nullableIntHolder.isSet; holder.value = nullableIntHolder.value; }; } /** * Create a specific getter for {@link BigIntVector}. * * @param vector an instance of the {@link BigIntVector} * @return a getter. */ private static Getter createGetter(BigIntVector vector) { NullableBigIntHolder nullableBigIntHolder = new NullableBigIntHolder(); return (index, holder) -> { vector.get(index, nullableBigIntHolder); holder.isSet = nullableBigIntHolder.isSet; holder.value = nullableBigIntHolder.value; }; } } text/000077500000000000000000000000001515451615200357165ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/implArrowFlightJdbcVarCharVectorAccessor.java000066400000000000000000000165671515451615200457300ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/text/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.text; import static java.nio.charset.StandardCharsets.US_ASCII; import java.io.ByteArrayInputStream; import java.io.CharArrayReader; import java.io.InputStream; import java.io.Reader; import java.math.BigDecimal; import java.nio.charset.StandardCharsets; import java.sql.Date; import java.sql.SQLException; import java.sql.Time; import java.sql.Timestamp; import java.util.Calendar; 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.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}. */ public class ArrowFlightJdbcVarCharVectorAccessor extends ArrowFlightJdbcAccessor { /** Functional interface to help integrating VarCharVector and LargeVarCharVector. */ @FunctionalInterface interface Getter { byte[] get(int index); } private final Getter getter; public ArrowFlightJdbcVarCharVectorAccessor( VarCharVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { this(vector::get, currentRowSupplier, setCursorWasNull); } public ArrowFlightJdbcVarCharVectorAccessor( LargeVarCharVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { this(vector::get, currentRowSupplier, setCursorWasNull); } public ArrowFlightJdbcVarCharVectorAccessor( ViewVarCharVector vector, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { this(vector::get, currentRowSupplier, setCursorWasNull); } ArrowFlightJdbcVarCharVectorAccessor( Getter getter, IntSupplier currentRowSupplier, ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { super(currentRowSupplier, setCursorWasNull); this.getter = getter; } @Override public Class getObjectClass() { return String.class; } @Override public String getObject() { final byte[] bytes = getBytes(); return bytes == null ? null : new String(bytes, StandardCharsets.UTF_8); } @Override public String getString() { return getObject(); } @Override public byte[] getBytes() { final byte[] bytes = this.getter.get(getCurrentRow()); this.wasNull = bytes == null; this.wasNullConsumer.setWasNull(this.wasNull); return bytes; } @Override public boolean getBoolean() throws SQLException { String value = getString(); if (value == null || value.equalsIgnoreCase("false") || value.equals("0")) { return false; } else if (value.equalsIgnoreCase("true") || value.equals("1")) { return true; } else { throw new SQLException("It is not possible to convert this value to boolean: " + value); } } @Override public byte getByte() throws SQLException { try { return Byte.parseByte(this.getString()); } catch (Exception e) { throw new SQLException(e); } } @Override public short getShort() throws SQLException { try { return Short.parseShort(this.getString()); } catch (Exception e) { throw new SQLException(e); } } @Override public int getInt() throws SQLException { try { return Integer.parseInt(this.getString()); } catch (Exception e) { throw new SQLException(e); } } @Override public long getLong() throws SQLException { try { return Long.parseLong(this.getString()); } catch (Exception e) { throw new SQLException(e); } } @Override public float getFloat() throws SQLException { try { return Float.parseFloat(this.getString()); } catch (Exception e) { throw new SQLException(e); } } @Override public double getDouble() throws SQLException { try { return Double.parseDouble(this.getString()); } catch (Exception e) { throw new SQLException(e); } } @Override public BigDecimal getBigDecimal() throws SQLException { try { return new BigDecimal(this.getString()); } catch (Exception e) { throw new SQLException(e); } } @Override public BigDecimal getBigDecimal(int i) throws SQLException { try { return BigDecimal.valueOf(this.getLong(), i); } catch (Exception e) { throw new SQLException(e); } } @Override public InputStream getAsciiStream() { final String textValue = getString(); if (textValue == null) { return null; } // Already in UTF-8 return new ByteArrayInputStream(textValue.getBytes(US_ASCII)); } @Override public InputStream getUnicodeStream() { final byte[] value = getBytes(); if (value == null) { return null; } // Already in UTF-8 final Text textValue = new Text(value); return new ByteArrayInputStream(textValue.getBytes(), 0, (int) textValue.getLength()); } @Override public Reader getCharacterStream() { return new CharArrayReader(getString().toCharArray()); } @Override public Date getDate(Calendar calendar) throws SQLException { try { Date date = Date.valueOf(getString()); if (calendar == null) { return date; } // Use Calendar to apply time zone's offset long milliseconds = date.getTime(); return new Date(DateTimeUtils.applyCalendarOffset(milliseconds, calendar)); } catch (Exception e) { throw new SQLException(e); } } @Override public Time getTime(Calendar calendar) throws SQLException { try { Time time = Time.valueOf(getString()); if (calendar == null) { return time; } // Use Calendar to apply time zone's offset long milliseconds = time.getTime(); return new Time(DateTimeUtils.applyCalendarOffset(milliseconds, calendar)); } catch (Exception e) { throw new SQLException(e); } } @Override public Timestamp getTimestamp(Calendar calendar) throws SQLException { try { Timestamp timestamp = Timestamp.valueOf(getString()); if (calendar == null) { return timestamp; } // Use Calendar to apply time zone's offset long milliseconds = timestamp.getTime(); return new Timestamp(DateTimeUtils.applyCalendarOffset(milliseconds, calendar)); } catch (Exception e) { throw new SQLException(e); } } } client/000077500000000000000000000000001515451615200334455ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbcArrowFlightSqlClientHandler.java000066400000000000000000001216451515451615200416660ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 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; import java.util.HashSet; import java.util.List; 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.SessionOptionValueFactory; import org.apache.arrow.flight.SetSessionOptionsRequest; import org.apache.arrow.flight.SetSessionOptionsResult; import org.apache.arrow.flight.auth2.BearerCredentialWriter; import org.apache.arrow.flight.auth2.ClientBearerHeaderHandler; 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; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; 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; import org.slf4j.LoggerFactory; /** A {@link FlightSqlClient} handler. */ public final class ArrowFlightSqlClientHandler implements AutoCloseable { private static final Logger LOGGER = LoggerFactory.getLogger(ArrowFlightSqlClientHandler.class); // 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 @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; } /** * Creates a new {@link ArrowFlightSqlClientHandler} from the provided {@code client} and {@code * options}. * * @param client the {@link FlightClient} to manage under a {@link FlightSqlClient} wrapper. * @param options the {@link CallOption}s to persist in between subsequent client calls. * @return a new {@link ArrowFlightSqlClientHandler}. */ static ArrowFlightSqlClientHandler createNewHandler( final String cacheKey, final FlightClient client, final Builder builder, final Collection options, final Optional catalog, final @Nullable FlightClientCache flightClientCache) { final ArrowFlightSqlClientHandler handler = new ArrowFlightSqlClientHandler( cacheKey, new FlightSqlClient(client), builder, options, catalog, flightClientCache); handler.setSetCatalogInSessionIfPresent(); return handler; } /** * Gets the {@link #options} for the subsequent calls from this handler. * * @return the {@link CallOption}s. */ private CallOption[] getOptions() { return options.toArray(new CallOption[0]); } /** * Makes an RPC "getStream" request based on the provided {@link FlightInfo} object. Retrieves the * result of the query previously prepared with "getInfo." * * @param flightInfo The {@link FlightInfo} instance from which to fetch results. * @return a {@code FlightStream} of results. */ public List getStreams(final FlightInfo flightInfo) throws SQLException { final ArrayList endpoints = new ArrayList<>(flightInfo.getEndpoints().size()); 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. 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. // 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; 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 = new CloseableEndpointStreamPair( sqlClient.getStream(endpoint.getTicket(), getOptions()), null); break; } final Builder builderForEndpoint = new Builder(ArrowFlightSqlClientHandler.this.builder) .withHost(endpointUri.getHost()) .withPort(endpointUri.getPort()) .withEncryption(endpointUri.getScheme().equals(LocationSchemes.GRPC_TLS)) .withClientCache(flightClientCache) .withConnectTimeout(builder.connectTimeout); ArrowFlightSqlClientHandler endpointHandler = null; try { endpointHandler = builderForEndpoint.build(); stream = new CloseableEndpointStreamPair( endpointHandler.sqlClient.getStream( endpoint.getTicket(), endpointHandler.getOptions()), endpointHandler.sqlClient); // Make sure we actually get data from the server 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) { endpoints.add(stream); } else if (exceptions.isEmpty()) { // This should never happen... throw new IllegalStateException("Could not connect to endpoint and no errors occurred"); } else { Exception ex = exceptions.remove(0); while (!exceptions.isEmpty()) { ex.addSuppressed(exceptions.remove(exceptions.size() - 1)); } throw ex; } } } } catch (Exception outerException) { try { AutoCloseables.close(endpoints); } catch (Exception innerEx) { outerException.addSuppressed(innerEx); } if (outerException instanceof SQLException) { throw (SQLException) outerException; } throw new SQLException(outerException); } return endpoints; } /** * Makes an RPC "getInfo" request based on the provided {@code query} object. * * @param query The query. * @return a {@code FlightStream} of results. */ public FlightInfo getInfo(final String query) { return sqlClient.execute(query, getOptions()); } @Override public void close() throws SQLException { if (catalog.isPresent()) { 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 { /** * Executes this {@link PreparedStatement}. * * @return the {@link FlightInfo} representing the outcome of this query execution. * @throws SQLException on error. */ FlightInfo executeQuery() throws SQLException; /** * Executes a {@link StatementType#UPDATE} query. * * @return the number of rows affected. */ long executeUpdate(); /** * Gets the {@link StatementType} of this {@link PreparedStatement}. * * @return the Statement Type. */ StatementType getType(); /** * Gets the {@link Schema} of this {@link PreparedStatement}. * * @return {@link Schema}. */ Schema getDataSetSchema(); /** * Gets the {@link Schema} of the parameters for this {@link PreparedStatement}. * * @return {@link Schema}. */ Schema getParameterSchema(); void setParameters(VectorSchemaRoot parameters); @Override void close(); } /** A connection is created with catalog set as a session option. */ private void setSetCatalogInSessionIfPresent() { if (catalog.isPresent()) { 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()) { final Map errors = result.getErrors(); for (final Map.Entry error : errors.entrySet()) { LOGGER.warn(error.toString()); } 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); } } /** * Creates a new {@link PreparedStatement} for the given {@code query}. * * @param query the SQL query. * @return a new prepared statement. */ public PreparedStatement prepare(final String query) { final FlightSqlClient.PreparedStatement preparedStatement = sqlClient.prepare(query, getOptions()); return new PreparedStatement() { @Override public FlightInfo executeQuery() throws SQLException { return preparedStatement.execute(getOptions()); } @Override public long executeUpdate() { return preparedStatement.executeUpdate(getOptions()); } @Override public StatementType getType() { final Schema schema = preparedStatement.getResultSetSchema(); return schema.getFields().isEmpty() ? StatementType.UPDATE : StatementType.SELECT; } @Override public Schema getDataSetSchema() { return preparedStatement.getResultSetSchema(); } @Override public Schema getParameterSchema() { return preparedStatement.getParameterSchema(); } @Override public void setParameters(VectorSchemaRoot parameters) { preparedStatement.setParameters(parameters); } @Override public void close() { try { preparedStatement.close(getOptions()); } catch (FlightRuntimeException fre) { handleBenignCloseException(fre, "closing PreparedStatement"); } } }; } /** * Makes an RPC "getCatalogs" request. * * @return a {@code FlightStream} of results. */ public FlightInfo getCatalogs() { return sqlClient.getCatalogs(getOptions()); } /** * Makes an RPC "getImportedKeys" request based on the provided info. * * @param catalog The catalog name. Must match the catalog name as it is stored in the database. * Retrieves those without a catalog. Null means that the catalog name should not be used to * narrow the search. * @param schema The schema name. Must match the schema name as it is stored in the database. "" * retrieves those without a schema. Null means that the schema name should not be used to * narrow the search. * @param table The table name. Must match the table name as it is stored in the database. * @return a {@code FlightStream} of results. */ public FlightInfo getImportedKeys(final String catalog, final String schema, final String table) { return sqlClient.getImportedKeys(TableRef.of(catalog, schema, table), getOptions()); } /** * Makes an RPC "getExportedKeys" request based on the provided info. * * @param catalog The catalog name. Must match the catalog name as it is stored in the database. * Retrieves those without a catalog. Null means that the catalog name should not be used to * narrow the search. * @param schema The schema name. Must match the schema name as it is stored in the database. "" * retrieves those without a schema. Null means that the schema name should not be used to * narrow the search. * @param table The table name. Must match the table name as it is stored in the database. * @return a {@code FlightStream} of results. */ public FlightInfo getExportedKeys(final String catalog, final String schema, final String table) { return sqlClient.getExportedKeys(TableRef.of(catalog, schema, table), getOptions()); } /** * Makes an RPC "getSchemas" request based on the provided info. * * @param catalog The catalog name. Must match the catalog name as it is stored in the database. * Retrieves those without a catalog. Null means that the catalog name should not be used to * narrow the search. * @param schemaPattern The schema name pattern. Must match the schema name as it is stored in the * database. Null means that schema name should not be used to narrow down the search. * @return a {@code FlightStream} of results. */ public FlightInfo getSchemas(final String catalog, final String schemaPattern) { return sqlClient.getSchemas(catalog, schemaPattern, getOptions()); } /** * Makes an RPC "getTableTypes" request. * * @return a {@code FlightStream} of results. */ public FlightInfo getTableTypes() { return sqlClient.getTableTypes(getOptions()); } /** * Makes an RPC "getTables" request based on the provided info. * * @param catalog The catalog name. Must match the catalog name as it is stored in the database. * Retrieves those without a catalog. Null means that the catalog name should not be used to * narrow the search. * @param schemaPattern The schema name pattern. Must match the schema name as it is stored in the * database. "" retrieves those without a schema. Null means that the schema name should not * be used to narrow the search. * @param tableNamePattern The table name pattern. Must match the table name as it is stored in * the database. * @param types The list of table types, which must be from the list of table types to include. * Null returns all types. * @param includeSchema Whether to include schema. * @return a {@code FlightStream} of results. */ public FlightInfo getTables( final String catalog, final String schemaPattern, final String tableNamePattern, final List types, final boolean includeSchema) { return sqlClient.getTables( catalog, schemaPattern, tableNamePattern, types, includeSchema, getOptions()); } /** * Gets SQL info. * * @return the SQL info. */ public FlightInfo getSqlInfo(SqlInfo... info) { return sqlClient.getSqlInfo(info, getOptions()); } /** * Makes an RPC "getPrimaryKeys" request based on the provided info. * * @param catalog The catalog name; must match the catalog name as it is stored in the database. * "" retrieves those without a catalog. Null means that the catalog name should not be used * to narrow the search. * @param schema The schema name; must match the schema name as it is stored in the database. "" * retrieves those without a schema. Null means that the schema name should not be used to * narrow the search. * @param table The table name. Must match the table name as it is stored in the database. * @return a {@code FlightStream} of results. */ public FlightInfo getPrimaryKeys(final String catalog, final String schema, final String table) { return sqlClient.getPrimaryKeys(TableRef.of(catalog, schema, table), getOptions()); } /** * Makes an RPC "getCrossReference" request based on the provided info. * * @param pkCatalog The catalog name. Must match the catalog name as it is stored in the database. * Retrieves those without a catalog. Null means that the catalog name should not be used to * narrow the search. * @param pkSchema The schema name. Must match the schema name as it is stored in the database. "" * retrieves those without a schema. Null means that the schema name should not be used to * narrow the search. * @param pkTable The table name. Must match the table name as it is stored in the database. * @param fkCatalog The catalog name. Must match the catalog name as it is stored in the database. * Retrieves those without a catalog. Null means that the catalog name should not be used to * narrow the search. * @param fkSchema The schema name. Must match the schema name as it is stored in the database. "" * retrieves those without a schema. Null means that the schema name should not be used to * narrow the search. * @param fkTable The table name. Must match the table name as it is stored in the database. * @return a {@code FlightStream} of results. */ public FlightInfo getCrossReference( String pkCatalog, String pkSchema, String pkTable, String fkCatalog, String fkSchema, String fkTable) { return sqlClient.getCrossReference( TableRef.of(pkCatalog, pkSchema, pkTable), TableRef.of(fkCatalog, fkSchema, fkTable), getOptions()); } /** 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; private int port; @VisibleForTesting String username; @VisibleForTesting String password; @VisibleForTesting String trustStorePath; @VisibleForTesting String trustStorePassword; @VisibleForTesting String token; @VisibleForTesting boolean useEncryption = true; @VisibleForTesting boolean disableCertificateVerification; @VisibleForTesting boolean useSystemTrustStore = true; @VisibleForTesting String tlsRootCertificatesPath; @VisibleForTesting String clientCertificatePath; @VisibleForTesting String clientKeyPath; @VisibleForTesting private BufferAllocator allocator; @VisibleForTesting boolean retainCookies = true; @VisibleForTesting boolean retainAuth = true; @VisibleForTesting Optional catalog = Optional.empty(); @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 ClientIncomingAuthHeaderMiddleware.Factory authFactory = new ClientIncomingAuthHeaderMiddleware.Factory(new ClientBearerHeaderHandler()); @VisibleForTesting ClientCookieMiddleware.Factory cookieFactory = new ClientCookieMiddleware.Factory(); DriverVersion driverVersion; public Builder() {} /** * Copies the builder. * * @param original The builder to base this copy off of. */ @VisibleForTesting Builder(Builder original) { this.middlewareFactories.addAll(original.middlewareFactories); this.options.addAll(original.options); this.host = original.host; this.port = original.port; this.username = original.username; this.password = original.password; this.trustStorePath = original.trustStorePath; this.trustStorePassword = original.trustStorePassword; this.token = original.token; this.useEncryption = original.useEncryption; this.disableCertificateVerification = original.disableCertificateVerification; this.useSystemTrustStore = original.useSystemTrustStore; this.tlsRootCertificatesPath = original.tlsRootCertificatesPath; this.clientCertificatePath = original.clientCertificatePath; this.clientKeyPath = original.clientKeyPath; this.allocator = original.allocator; this.catalog = original.catalog; this.oauthConfig = original.oauthConfig; if (original.retainCookies) { this.cookieFactory = original.cookieFactory; } if (original.retainAuth) { this.authFactory = original.authFactory; } this.driverVersion = original.driverVersion; } /** * Sets the host for this handler. * * @param host the host. * @return this instance. */ public Builder withHost(final String host) { this.host = host; return this; } /** * Sets the port for this handler. * * @param port the port. * @return this instance. */ public Builder withPort(final int port) { this.port = port; return this; } /** * Sets the username for this handler. * * @param username the username. * @return this instance. */ public Builder withUsername(final String username) { this.username = username; return this; } /** * Sets the password for this handler. * * @param password the password. * @return this instance. */ public Builder withPassword(final String password) { this.password = password; return this; } /** * Sets the KeyStore path for this handler. * * @param trustStorePath the KeyStore path. * @return this instance. */ public Builder withTrustStorePath(final String trustStorePath) { this.trustStorePath = trustStorePath; return this; } /** * Sets the KeyStore password for this handler. * * @param trustStorePassword the KeyStore password. * @return this instance. */ public Builder withTrustStorePassword(final String trustStorePassword) { this.trustStorePassword = trustStorePassword; return this; } /** * Sets whether to use TLS encryption in this handler. * * @param useEncryption whether to use TLS encryption. * @return this instance. */ public Builder withEncryption(final boolean useEncryption) { this.useEncryption = useEncryption; return this; } /** * Sets whether to disable the certificate verification in this handler. * * @param disableCertificateVerification whether to disable certificate verification. * @return this instance. */ public Builder withDisableCertificateVerification( final boolean disableCertificateVerification) { this.disableCertificateVerification = disableCertificateVerification; return this; } /** * Sets whether to use the certificates from the operating system. * * @param useSystemTrustStore whether to use the system operating certificates. * @return this instance. */ public Builder withSystemTrustStore(final boolean useSystemTrustStore) { this.useSystemTrustStore = useSystemTrustStore; return this; } /** * Sets the TLS root certificate path as an alternative to using the System or other Trust * Store. The path must contain a valid PEM file. * * @param tlsRootCertificatesPath the TLS root certificate path (if TLS is required). * @return this instance. */ public Builder withTlsRootCertificates(final String tlsRootCertificatesPath) { this.tlsRootCertificatesPath = tlsRootCertificatesPath; return this; } /** * Sets the mTLS client certificate path (if mTLS is required). * * @param clientCertificatePath the mTLS client certificate path (if mTLS is required). * @return this instance. */ public Builder withClientCertificate(final String clientCertificatePath) { this.clientCertificatePath = clientCertificatePath; return this; } /** * Sets the mTLS client certificate private key path (if mTLS is required). * * @param clientKeyPath the mTLS client certificate private key path (if mTLS is required). * @return this instance. */ public Builder withClientKey(final String clientKeyPath) { this.clientKeyPath = clientKeyPath; return this; } /** * Sets the token used in the token authentication. * * @param token the token value. * @return this builder instance. */ public Builder withToken(final String token) { this.token = token; return this; } /** * Sets the {@link BufferAllocator} to use in this handler. * * @param allocator the allocator. * @return this instance. */ public Builder withBufferAllocator(final BufferAllocator allocator) { this.allocator = allocator.newChildAllocator("ArrowFlightSqlClientHandler", 0, allocator.getLimit()); return this; } /** * Indicates if cookies should be re-used by connections spawned for getStreams() calls. * * @param retainCookies The flag indicating if cookies should be re-used. * @return this builder instance. */ public Builder withRetainCookies(boolean retainCookies) { this.retainCookies = retainCookies; return this; } /** * Indicates if bearer tokens negotiated should be re-used by connections spawned for * getStreams() calls. * * @param retainAuth The flag indicating if auth tokens should be re-used. * @return this builder instance. */ public Builder withRetainAuth(boolean retainAuth) { this.retainAuth = retainAuth; return this; } /** * Adds the provided {@code factories} to the list of {@link #middlewareFactories} of this * handler. * * @param factories the factories to add. * @return this instance. */ public Builder withMiddlewareFactories(final FlightClientMiddleware.Factory... factories) { return withMiddlewareFactories(Arrays.asList(factories)); } /** * Adds the provided {@code factories} to the list of {@link #middlewareFactories} of this * handler. * * @param factories the factories to add. * @return this instance. */ public Builder withMiddlewareFactories( final Collection factories) { this.middlewareFactories.addAll(factories); return this; } /** * Adds the provided {@link CallOption}s to this handler. * * @param options the options * @return this instance. */ public Builder withCallOptions(final CallOption... options) { return withCallOptions(Arrays.asList(options)); } /** * Adds the provided {@link CallOption}s to this handler. * * @param options the options * @return this instance. */ public Builder withCallOptions(final Collection options) { this.options.addAll(options); return this; } /** * Sets the catalog for this handler if it is not null. * * @param catalog the catalog * @return this instance. */ public Builder withCatalog(@Nullable final String catalog) { this.catalog = Optional.ofNullable(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. * * @return a new client handler. * @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 // itself. Set buildTimeMiddlewareFactories = new HashSet<>(this.middlewareFactories); FlightClient client = null; boolean isUsingUserPasswordAuth = username != null && token == null; try { // Token should take priority since some apps pass in a username/password even // when a token // is provided if (isUsingUserPasswordAuth) { buildTimeMiddlewareFactories.add(authFactory); } final NettyClientBuilder clientBuilder = new NettyClientBuilder(); clientBuilder.allocator(allocator); String userAgent = String.format(USER_AGENT_TEMPLATE, DEFAULT_VERSION); if (driverVersion != null && driverVersion.versionString != null) { userAgent = String.format(USER_AGENT_TEMPLATE, driverVersion.versionString); } buildTimeMiddlewareFactories.add(new ClientCookieMiddleware.Factory()); buildTimeMiddlewareFactories.forEach(clientBuilder::intercept); if (useEncryption) { clientBuilder.useTls(); } Location location = getLocation(); clientBuilder.location(location); if (useEncryption) { if (disableCertificateVerification) { clientBuilder.verifyServer(false); } else { if (tlsRootCertificatesPath != null) { clientBuilder.trustedCertificates( ClientAuthenticationUtils.getTlsRootCertificatesStream(tlsRootCertificatesPath)); } else if (useSystemTrustStore) { clientBuilder.trustedCertificates( ClientAuthenticationUtils.getCertificateInputStreamFromSystem( trustStorePassword)); } else if (trustStorePath != null) { clientBuilder.trustedCertificates( ClientAuthenticationUtils.getCertificateStream( trustStorePath, trustStorePassword)); } } if (clientCertificatePath != null && clientKeyPath != null) { clientBuilder.clientCertificate( ClientAuthenticationUtils.getClientCertificateStream(clientCertificatePath), ClientAuthenticationUtils.getClientKeyStream(clientKeyPath)); } } 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<>(); // 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()); } else { // Otherwise do the handshake and get the token if possible. credentialOptions.add( ClientAuthenticationUtils.getAuthenticate( client, username, password, authFactory, options.toArray(new CallOption[0]))); } } else if (token != null) { credentialOptions.add( ClientAuthenticationUtils.getAuthenticate( client, new CredentialCallOption(new BearerCredentialWriter(token)), options.toArray(new CallOption[0]))); } return ArrowFlightSqlClientHandler.createNewHandler( getCacheKey(), client, this, credentialOptions, catalog, flightClientCache); } catch (final IllegalArgumentException | GeneralSecurityException | IOException | FlightRuntimeException e) { final SQLException originalException = new SQLException(e); if (client != null) { try { client.close(); } catch (final InterruptedException interruptedException) { originalException.addSuppressed(interruptedException); } } throw originalException; } } } } CloseableEndpointStreamPair.java000066400000000000000000000030731515451615200416750ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.sql.FlightSqlClient; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; /** Represents a connection to a {@link org.apache.arrow.flight.FlightEndpoint}. */ public class CloseableEndpointStreamPair implements AutoCloseable { private final FlightStream stream; private final FlightSqlClient client; public CloseableEndpointStreamPair(FlightStream stream, FlightSqlClient client) { this.stream = Preconditions.checkNotNull(stream); this.client = client; } public FlightStream getStream() { return stream; } @Override public void close() throws Exception { AutoCloseables.close(stream, client); } } oauth/000077500000000000000000000000001515451615200345655ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/clientAbstractOAuthTokenProvider.java000066400000000000000000000100201515451615200426410ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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(); } ClientCredentialsTokenProvider.java000066400000000000000000000047261515451615200435510ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } OAuthConfiguration.java000066400000000000000000000204761515451615200412110ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } } OAuthCredentialWriter.java000066400000000000000000000032701515451615200416420ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } } OAuthTokenException.java000066400000000000000000000023061515451615200413310ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } OAuthTokenProvider.java000066400000000000000000000024541515451615200411710ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; } OAuthTokenProviders.java000066400000000000000000000323301515451615200413500ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } } } TokenExchangeTokenProvider.java000066400000000000000000000056151515451615200426760ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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(); } } TokenInfo.java000066400000000000000000000033171515451615200373300ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } utils/000077500000000000000000000000001515451615200346055ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/clientClientAuthenticationUtils.java000066400000000000000000000264041515451615200426150ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.io.ByteArrayInputStream; import java.io.IOException; import java.io.InputStream; import java.io.StringWriter; import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.Paths; import java.security.GeneralSecurityException; import java.security.KeyStore; import java.security.KeyStoreException; import java.security.NoSuchAlgorithmException; import java.security.cert.Certificate; import java.security.cert.CertificateException; import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; import java.util.Enumeration; import java.util.List; import org.apache.arrow.flight.CallOption; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.auth2.BasicAuthCredentialWriter; import org.apache.arrow.flight.auth2.ClientIncomingAuthHeaderMiddleware; import org.apache.arrow.flight.grpc.CredentialCallOption; import org.apache.arrow.util.Preconditions; import org.apache.arrow.util.VisibleForTesting; import org.bouncycastle.openssl.jcajce.JcaPEMWriter; /** Utils for {@link FlightClientHandler} authentication. */ public final class ClientAuthenticationUtils { private ClientAuthenticationUtils() { // Prevent instantiation. } /** * Gets the {@link CredentialCallOption} for the provided authentication info. * * @param client the client. * @param credential the credential as CallOptions. * @param options the {@link CallOption}s to use. * @return the credential call option. */ public static CredentialCallOption getAuthenticate( final FlightClient client, final CredentialCallOption credential, final CallOption... options) { final List theseOptions = new ArrayList<>(); theseOptions.add(credential); theseOptions.addAll(Arrays.asList(options)); client.handshake(theseOptions.toArray(new CallOption[0])); return (CredentialCallOption) theseOptions.get(0); } /** * Gets the {@link CredentialCallOption} for the provided authentication info. * * @param client the client. * @param username the username. * @param password the password. * @param factory the {@link ClientIncomingAuthHeaderMiddleware.Factory} to use. * @param options the {@link CallOption}s to use. * @return the credential call option. */ public static CredentialCallOption getAuthenticate( final FlightClient client, final String username, final String password, final ClientIncomingAuthHeaderMiddleware.Factory factory, final CallOption... options) { return getAuthenticate( client, new CredentialCallOption(new BasicAuthCredentialWriter(username, password)), factory, options); } private static CredentialCallOption getAuthenticate( final FlightClient client, final CredentialCallOption token, final ClientIncomingAuthHeaderMiddleware.Factory factory, final CallOption... options) { final List theseOptions = new ArrayList<>(); theseOptions.add(token); theseOptions.addAll(Arrays.asList(options)); client.handshake(theseOptions.toArray(new CallOption[0])); return factory.getCredentialCallOption(); } @VisibleForTesting static KeyStore getKeyStoreInstance(String instance) throws KeyStoreException, CertificateException, IOException, NoSuchAlgorithmException { KeyStore keyStore = KeyStore.getInstance(instance); keyStore.load(null, null); return keyStore; } @VisibleForTesting static KeyStore getDefaultKeyStoreInstance(String password) throws KeyStoreException, CertificateException, NoSuchAlgorithmException, IOException { try (InputStream fileInputStream = getKeystoreInputStream()) { KeyStore keyStore = KeyStore.getInstance(KeyStore.getDefaultType()); keyStore.load(fileInputStream, password == null ? null : password.toCharArray()); return keyStore; } } static String getOperatingSystem() { return System.getProperty("os.name"); } /** * Check if the operating system running the software is Windows. * * @return whether is the windows system. */ public static boolean isWindows() { return getOperatingSystem().contains("Windows"); } /** * Check if the operating system running the software is Mac. * * @return whether is the mac system. */ public static boolean isMac() { return getOperatingSystem().contains("Mac"); } /** * It gets the trusted certificate based on the operating system and loads all the certificate * into a {@link InputStream}. * * @return An input stream with all the certificates. * @throws KeyStoreException if a key store could not be loaded. * @throws CertificateException if a certificate could not be found. * @throws IOException if it fails reading the file. */ public static InputStream getCertificateInputStreamFromSystem(String password) throws KeyStoreException, CertificateException, IOException, NoSuchAlgorithmException { List keyStoreList = new ArrayList<>(); if (isWindows()) { keyStoreList.add(getKeyStoreInstance("Windows-ROOT")); keyStoreList.add(getKeyStoreInstance("Windows-MY")); } else if (isMac()) { keyStoreList.add(getKeyStoreInstance("KeychainStore")); keyStoreList.add(getDefaultKeyStoreInstance(password)); } else { keyStoreList.add(getDefaultKeyStoreInstance(password)); } return getCertificatesInputStream(keyStoreList); } @VisibleForTesting static InputStream getKeystoreInputStream() throws IOException { Path path = Paths.get(System.getProperty("java.home"), "lib", "security", "cacerts"); if (Files.notExists(path)) { // for JDK8 path = Paths.get(System.getProperty("java.home"), "jre", "lib", "security", "cacerts"); } return Files.newInputStream(path); } @VisibleForTesting static void getCertificatesInputStream(KeyStore keyStore, JcaPEMWriter pemWriter) throws IOException, KeyStoreException { Enumeration aliases = keyStore.aliases(); while (aliases.hasMoreElements()) { String alias = aliases.nextElement(); if (keyStore.isCertificateEntry(alias)) { pemWriter.writeObject(keyStore.getCertificate(alias)); } } pemWriter.flush(); } @VisibleForTesting static InputStream getCertificatesInputStream(Collection keyStores) throws IOException, KeyStoreException { try (final StringWriter writer = new StringWriter(); final JcaPEMWriter pemWriter = new JcaPEMWriter(writer)) { for (KeyStore keyStore : keyStores) { getCertificatesInputStream(keyStore, pemWriter); } return new ByteArrayInputStream(writer.toString().getBytes(StandardCharsets.UTF_8)); } } /** * Generates an {@link InputStream} that contains certificates for a private key. * * @param keyStorePath The path of the KeyStore. * @param keyStorePass The password of the KeyStore. * @return a new {code InputStream} containing the certificates. * @throws GeneralSecurityException on error. * @throws IOException on error. */ public static InputStream getCertificateStream( final String keyStorePath, final String keyStorePass) throws GeneralSecurityException, IOException { Preconditions.checkNotNull(keyStorePath, "KeyStore path cannot be null!"); Preconditions.checkNotNull(keyStorePass, "KeyStorePass cannot be null!"); final KeyStore keyStore = KeyStore.getInstance(KeyStore.getDefaultType()); try (final InputStream keyStoreStream = Files.newInputStream(Paths.get(keyStorePath))) { keyStore.load(keyStoreStream, keyStorePass.toCharArray()); } return getSingleCertificateInputStream(keyStore); } /** * Generates an {@link InputStream} that contains certificates for path-based TLS Root * Certificates. * * @param tlsRootsCertificatesPath The path of the TLS Root Certificates. * @return a new {code InputStream} containing the certificates. * @throws GeneralSecurityException on error. * @throws IOException on error. */ public static InputStream getTlsRootCertificatesStream(final String tlsRootsCertificatesPath) throws GeneralSecurityException, IOException { Preconditions.checkNotNull( tlsRootsCertificatesPath, "TLS Root certificates path cannot be null!"); return Files.newInputStream(Paths.get(tlsRootsCertificatesPath)); } /** * Generates an {@link InputStream} that contains certificates for a path-based mTLS Client * Certificate. * * @param clientCertificatePath The path of the mTLS Client Certificate. * @return a new {code InputStream} containing the certificates. * @throws GeneralSecurityException on error. * @throws IOException on error. */ public static InputStream getClientCertificateStream(final String clientCertificatePath) throws GeneralSecurityException, IOException { Preconditions.checkNotNull(clientCertificatePath, "Client certificate path cannot be null!"); return Files.newInputStream(Paths.get(clientCertificatePath)); } /** * Generates an {@link InputStream} that contains certificates for a path-based mTLS Client Key. * * @param clientKeyPath The path of the mTLS Client Key. * @return a new {code InputStream} containing the certificates. * @throws GeneralSecurityException on error. * @throws IOException on error. */ public static InputStream getClientKeyStream(final String clientKeyPath) throws GeneralSecurityException, IOException { Preconditions.checkNotNull(clientKeyPath, "Client key path cannot be null!"); return Files.newInputStream(Paths.get(clientKeyPath)); } private static InputStream getSingleCertificateInputStream(KeyStore keyStore) throws KeyStoreException, IOException, CertificateException { final Enumeration aliases = keyStore.aliases(); while (aliases.hasMoreElements()) { final String alias = aliases.nextElement(); if (keyStore.isCertificateEntry(alias)) { return toInputStream(keyStore.getCertificate(alias)); } } throw new CertificateException("Keystore did not have a certificate."); } private static InputStream toInputStream(final Certificate certificate) throws IOException { try (final StringWriter writer = new StringWriter(); final JcaPEMWriter pemWriter = new JcaPEMWriter(writer)) { pemWriter.writeObject(certificate); pemWriter.flush(); return new ByteArrayInputStream(writer.toString().getBytes(StandardCharsets.UTF_8)); } } } FlightClientCache.java000066400000000000000000000041571515451615200407570ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 {} } FlightLocationQueue.java000066400000000000000000000045341515451615200413710ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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(); } } converter/000077500000000000000000000000001515451615200341765ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbcAvaticaParameterConverter.java000066400000000000000000000034331515451615200421450ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** * Interface for a class in charge of converting between AvaticaParameters and TypedValues and * Arrow. */ public interface AvaticaParameterConverter { /** * Bind a TypedValue to a FieldVector at the given index. * * @param vector FieldVector that the parameter should be bound to. * @param typedValue TypedValue to bind as a parameter. * @param index Vector index (0-indexed) that the TypedValue should be bound to. * @return Whether the value was set successfully. */ boolean bindParameter(FieldVector vector, TypedValue typedValue, int index); /** * Create an AvaticaParameter from the given Field. * * @param field Arrow Field to generate an AvaticaParameter from. */ AvaticaParameter createParameter(Field field); } impl/000077500000000000000000000000001515451615200351375ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converterBaseAvaticaParameterConverter.java000066400000000000000000000037701515451615200437050ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.driver.jdbc.converter.AvaticaParameterConverter; import org.apache.arrow.driver.jdbc.utils.SqlTypes; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.SqlType; /** * Base AvaticaParameterConverter with a generic createParameter method that can be used by most * Arrow types. */ abstract class BaseAvaticaParameterConverter implements AvaticaParameterConverter { protected AvaticaParameter createParameter(Field field, boolean signed) { final String name = field.getName(); final ArrowType arrowType = field.getType(); final String typeName = arrowType.toString(); final int precision = 0; // Would have to know about the actual number final int scale = 0; // According to https://www.postgresql.org/docs/current/datatype-numeric.html final int jdbcType = SqlTypes.getSqlTypeIdFromArrowType(arrowType); final String className = SqlType.valueOf(jdbcType).clazz.getCanonicalName(); return new AvaticaParameter(signed, precision, scale, jdbcType, typeName, className, name); } } BinaryAvaticaParameterConverter.java000066400000000000000000000033701515451615200442530ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for Binary Arrow types. */ public class BinaryAvaticaParameterConverter extends BaseAvaticaParameterConverter { public BinaryAvaticaParameterConverter(ArrowType.Binary type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { byte[] value = (byte[]) typedValue.toJdbc(null); if (vector instanceof VarBinaryVector) { ((VarBinaryVector) vector).setSafe(index, value); return true; } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } BinaryViewAvaticaParameterConverter.java000066400000000000000000000034241515451615200451060ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 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; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for BinaryView Arrow types. */ public class BinaryViewAvaticaParameterConverter extends BaseAvaticaParameterConverter { public BinaryViewAvaticaParameterConverter(ArrowType.BinaryView type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { byte[] value = (byte[]) typedValue.toJdbc(null); if (vector instanceof ViewVarBinaryVector) { ((ViewVarBinaryVector) vector).setSafe(index, value); return true; } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } BoolAvaticaParameterConverter.java000066400000000000000000000033451515451615200437240ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for Bool Arrow types. */ public class BoolAvaticaParameterConverter extends BaseAvaticaParameterConverter { public BoolAvaticaParameterConverter(ArrowType.Bool type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { boolean value = (boolean) typedValue.toLocal(); if (vector instanceof BitVector) { ((BitVector) vector).setSafe(index, value ? 1 : 0); return true; } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } DateAvaticaParameterConverter.java000066400000000000000000000036201515451615200437020ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DateMilliVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for Date Arrow types. */ public class DateAvaticaParameterConverter extends BaseAvaticaParameterConverter { public DateAvaticaParameterConverter(ArrowType.Date type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { int value = (int) typedValue.toLocal(); if (vector instanceof DateMilliVector) { ((DateMilliVector) vector).setSafe(index, value); return true; } else if (vector instanceof DateDayVector) { ((DateDayVector) vector).setSafe(index, value); return true; } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } DecimalAvaticaParameterConverter.java000066400000000000000000000034271515451615200443700ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.math.BigDecimal; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for Decimal Arrow types. */ public class DecimalAvaticaParameterConverter extends BaseAvaticaParameterConverter { public DecimalAvaticaParameterConverter(ArrowType.Decimal type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { BigDecimal value = (BigDecimal) typedValue.toLocal(); if (vector instanceof DecimalVector) { ((DecimalVector) vector).setSafe(index, value); return true; } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, true); } } DurationAvaticaParameterConverter.java000066400000000000000000000030351515451615200446120ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for Duration Arrow types. */ public class DurationAvaticaParameterConverter extends BaseAvaticaParameterConverter { public DurationAvaticaParameterConverter(ArrowType.Duration type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } FixedSizeBinaryAvaticaParameterConverter.java000066400000000000000000000034561515451615200460730ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.FixedSizeBinaryVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for FixedSizeBinary Arrow types. */ public class FixedSizeBinaryAvaticaParameterConverter extends BaseAvaticaParameterConverter { public FixedSizeBinaryAvaticaParameterConverter(ArrowType.FixedSizeBinary type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { byte[] value = (byte[]) typedValue.toJdbc(null); if (vector instanceof FixedSizeBinaryVector) { ((FixedSizeBinaryVector) vector).setSafe(index, value); return true; } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } FixedSizeListAvaticaParameterConverter.java000066400000000000000000000062751515451615200455640ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.List; import org.apache.arrow.driver.jdbc.utils.AvaticaParameterBinder; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for FixedSizeList Arrow types. */ public class FixedSizeListAvaticaParameterConverter extends BaseAvaticaParameterConverter { public FixedSizeListAvaticaParameterConverter(ArrowType.FixedSizeList type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { final List values = (List) typedValue.value; final int arraySize = values.size(); if (vector instanceof FixedSizeListVector) { FixedSizeListVector listVector = ((FixedSizeListVector) vector); FieldVector childVector = listVector.getDataVector(); int maxArraySize = listVector.getListSize(); if (arraySize != maxArraySize) { if (!childVector.getField().isNullable()) { throw new UnsupportedOperationException( "Each array must contain " + maxArraySize + " elements"); } else if (arraySize > maxArraySize) { throw new UnsupportedOperationException( "Each array must contain at most " + maxArraySize + " elements"); } } int startPos = listVector.startNewValue(index); for (int i = 0; i < arraySize; i++) { Object val = values.get(i); int childIndex = startPos + i; if (val == null) { if (childVector.getField().isNullable()) { childVector.setNull(childIndex); } else { throw new UnsupportedOperationException("Can't set null on non-nullable child list"); } } else { childVector .getField() .getType() .accept( new AvaticaParameterBinder.BinderVisitor( childVector, TypedValue.ofSerial(typedValue.componentType, val), childIndex)); } } listVector.setValueCount(index + 1); return true; } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } FloatingPointAvaticaParameterConverter.java000066400000000000000000000037041515451615200456050ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for FloatingPoint Arrow types. */ public class FloatingPointAvaticaParameterConverter extends BaseAvaticaParameterConverter { public FloatingPointAvaticaParameterConverter(ArrowType.FloatingPoint type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { Number value = (Number) typedValue.value; if (vector instanceof Float4Vector) { ((Float4Vector) vector).setSafe(index, value.floatValue()); return true; } else if (vector instanceof Float8Vector) { ((Float8Vector) vector).setSafe(index, value.doubleValue()); return true; } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, true); } } IntAvaticaParameterConverter.java000066400000000000000000000057761515451615200435750ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.SmallIntVector; 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.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for Int Arrow types. */ public class IntAvaticaParameterConverter extends BaseAvaticaParameterConverter { private final ArrowType.Int type; public IntAvaticaParameterConverter(ArrowType.Int type) { this.type = type; } @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { Number value = (Number) typedValue.value; if (vector instanceof TinyIntVector) { ((TinyIntVector) vector).setSafe(index, value.intValue()); return true; } else if (vector instanceof SmallIntVector) { ((SmallIntVector) vector).setSafe(index, value.intValue()); return true; } else if (vector instanceof IntVector) { ((IntVector) vector).setSafe(index, value.intValue()); return true; } else if (vector instanceof BigIntVector) { ((BigIntVector) vector).setSafe(index, value.longValue()); return true; } else if (vector instanceof UInt1Vector) { ((UInt1Vector) vector).setSafe(index, value.intValue()); return true; } else if (vector instanceof UInt2Vector) { ((UInt2Vector) vector).setSafe(index, value.intValue()); return true; } else if (vector instanceof UInt4Vector) { ((UInt4Vector) vector).setSafe(index, value.intValue()); return true; } else if (vector instanceof UInt8Vector) { ((UInt8Vector) vector).setSafe(index, value.longValue()); return true; } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, type.getIsSigned()); } } IntervalAvaticaParameterConverter.java000066400000000000000000000035211515451615200446110ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for Interval Arrow types. */ public class IntervalAvaticaParameterConverter extends BaseAvaticaParameterConverter { public IntervalAvaticaParameterConverter(ArrowType.Interval type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { // Object value = typedValue.toLocal(); // if (vector instanceof IntervalDayVector) { // ((IntervalDayVector) vector).setSafe(index, () value); // } else if (vector instanceof IntervalYearVector) { // ((IntervalYearVector) vector).setSafe(index, () value); // } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } LargeBinaryAvaticaParameterConverter.java000066400000000000000000000034331515451615200452260ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for LargeBinary Arrow types. */ public class LargeBinaryAvaticaParameterConverter extends BaseAvaticaParameterConverter { public LargeBinaryAvaticaParameterConverter(ArrowType.LargeBinary type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { byte[] value = (byte[]) typedValue.toJdbc(null); if (vector instanceof LargeVarBinaryVector) { ((LargeVarBinaryVector) vector).setSafe(index, value); return true; } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } LargeListAvaticaParameterConverter.java000066400000000000000000000054631515451615200447220ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.List; import org.apache.arrow.driver.jdbc.utils.AvaticaParameterBinder; import org.apache.arrow.memory.util.LargeMemoryUtil; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.complex.LargeListVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for LargeList Arrow types. */ public class LargeListAvaticaParameterConverter extends BaseAvaticaParameterConverter { public LargeListAvaticaParameterConverter(ArrowType.LargeList type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { final List values = (List) typedValue.value; if (vector instanceof LargeListVector) { LargeListVector listVector = ((LargeListVector) vector); FieldVector childVector = listVector.getDataVector(); long startPos = listVector.startNewValue(index); for (int i = 0; i < values.size(); i++) { Object val = values.get(i); int childIndex = LargeMemoryUtil.checkedCastToInt(startPos) + i; if (val == null) { if (childVector.getField().isNullable()) { childVector.setNull(childIndex); } else { throw new UnsupportedOperationException("Can't set null on non-nullable child list"); } } else { childVector .getField() .getType() .accept( new AvaticaParameterBinder.BinderVisitor( childVector, TypedValue.ofSerial(typedValue.componentType, val), childIndex)); } } listVector.endValue(index, values.size()); listVector.setValueCount(index + 1); return true; } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } LargeUtf8AvaticaParameterConverter.java000066400000000000000000000034761515451615200446370ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.LargeVarCharVector; 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; /** AvaticaParameterConverter for LargeUtf8 Arrow types. */ public class LargeUtf8AvaticaParameterConverter extends BaseAvaticaParameterConverter { public LargeUtf8AvaticaParameterConverter(ArrowType.LargeUtf8 type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { String value = (String) typedValue.toLocal(); if (vector instanceof LargeVarCharVector) { ((LargeVarCharVector) vector).setSafe(index, new Text(value)); return true; } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } ListAvaticaParameterConverter.java000066400000000000000000000052631515451615200437450ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.List; import org.apache.arrow.driver.jdbc.utils.AvaticaParameterBinder; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for List Arrow types. */ public class ListAvaticaParameterConverter extends BaseAvaticaParameterConverter { public ListAvaticaParameterConverter(ArrowType.List type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { final List values = (List) typedValue.value; if (vector instanceof ListVector) { ListVector listVector = ((ListVector) vector); FieldVector childVector = listVector.getDataVector(); int startPos = listVector.startNewValue(index); for (int i = 0; i < values.size(); i++) { Object val = values.get(i); int childIndex = startPos + i; if (val == null) { if (childVector.getField().isNullable()) { childVector.setNull(childIndex); } else { throw new UnsupportedOperationException("Can't set null on non-nullable child list"); } } else { childVector .getField() .getType() .accept( new AvaticaParameterBinder.BinderVisitor( childVector, TypedValue.ofSerial(typedValue.componentType, val), childIndex)); } } listVector.endValue(index, values.size()); listVector.setValueCount(index + 1); return true; } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } MapAvaticaParameterConverter.java000066400000000000000000000030111515451615200435340ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for Map Arrow types. */ public class MapAvaticaParameterConverter extends BaseAvaticaParameterConverter { public MapAvaticaParameterConverter(ArrowType.Map type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } NullAvaticaParameterConverter.java000066400000000000000000000034351515451615200437430ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.NullVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for Null Arrow types. */ public class NullAvaticaParameterConverter extends BaseAvaticaParameterConverter { public NullAvaticaParameterConverter(ArrowType.Null type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { Object value = typedValue.toLocal(); if (vector instanceof NullVector) { if (value != null) { throw new RuntimeException("Can't set non-null value on NullVector"); } vector.setNull(index); } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } StructAvaticaParameterConverter.java000066400000000000000000000030251515451615200443100ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for Struct Arrow types. */ public class StructAvaticaParameterConverter extends BaseAvaticaParameterConverter { public StructAvaticaParameterConverter(ArrowType.Struct type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } TimeAvaticaParameterConverter.java000066400000000000000000000043531515451615200437270ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.FieldVector; 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.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for Time Arrow types. */ public class TimeAvaticaParameterConverter extends BaseAvaticaParameterConverter { public TimeAvaticaParameterConverter(ArrowType.Time type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { int value = (int) typedValue.toLocal(); if (vector instanceof TimeMicroVector) { ((TimeMicroVector) vector).setSafe(index, value); return true; } else if (vector instanceof TimeMilliVector) { ((TimeMilliVector) vector).setSafe(index, value); return true; } else if (vector instanceof TimeNanoVector) { ((TimeNanoVector) vector).setSafe(index, value); return true; } else if (vector instanceof TimeSecVector) { ((TimeSecVector) vector).setSafe(index, value); return true; } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } TimestampAvaticaParameterConverter.java000066400000000000000000000061041515451615200447700ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.FieldVector; 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.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for Timestamp Arrow types. */ public class TimestampAvaticaParameterConverter extends BaseAvaticaParameterConverter { public TimestampAvaticaParameterConverter(ArrowType.Timestamp type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { long value = (long) typedValue.toLocal(); if (vector instanceof TimeStampSecVector) { ((TimeStampSecVector) vector).setSafe(index, value); return true; } else if (vector instanceof TimeStampMicroVector) { ((TimeStampMicroVector) vector).setSafe(index, value); return true; } else if (vector instanceof TimeStampMilliVector) { ((TimeStampMilliVector) vector).setSafe(index, value); return true; } else if (vector instanceof TimeStampNanoVector) { ((TimeStampNanoVector) vector).setSafe(index, value); return true; } else if (vector instanceof TimeStampSecTZVector) { ((TimeStampSecTZVector) vector).setSafe(index, value); return true; } else if (vector instanceof TimeStampMicroTZVector) { ((TimeStampMicroTZVector) vector).setSafe(index, value); return true; } else if (vector instanceof TimeStampMilliTZVector) { ((TimeStampMilliTZVector) vector).setSafe(index, value); return true; } else if (vector instanceof TimeStampNanoTZVector) { ((TimeStampNanoTZVector) vector).setSafe(index, value); return true; } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } UnionAvaticaParameterConverter.java000066400000000000000000000030211515451615200441100ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; /** AvaticaParameterConverter for Union Arrow types. */ public class UnionAvaticaParameterConverter extends BaseAvaticaParameterConverter { public UnionAvaticaParameterConverter(ArrowType.Union type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } Utf8AvaticaParameterConverter.java000066400000000000000000000034331515451615200436550ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VarCharVector; 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; /** AvaticaParameterConverter for Utf8 Arrow types. */ public class Utf8AvaticaParameterConverter extends BaseAvaticaParameterConverter { public Utf8AvaticaParameterConverter(ArrowType.Utf8 type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { String value = (String) typedValue.toLocal(); if (vector instanceof VarCharVector) { ((VarCharVector) vector).setSafe(index, new Text(value)); return true; } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } Utf8ViewAvaticaParameterConverter.java000066400000000000000000000034671515451615200445170ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 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; /** AvaticaParameterConverter for Utf8View Arrow types. */ public class Utf8ViewAvaticaParameterConverter extends BaseAvaticaParameterConverter { public Utf8ViewAvaticaParameterConverter(ArrowType.Utf8View type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { String value = (String) typedValue.toLocal(); if (vector instanceof ViewVarCharVector) { ((ViewVarCharVector) vector).setSafe(index, new Text(value)); return true; } return false; } @Override public AvaticaParameter createParameter(Field field) { return createParameter(field, false); } } UuidAvaticaParameterConverter.java000066400000000000000000000075451515451615200437450ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } utils/000077500000000000000000000000001515451615200333275ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbcArrowFlightConnectionConfigImpl.java000066400000000000000000000326741515451615200424260ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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; import org.apache.arrow.flight.HeaderCallOption; import org.apache.arrow.util.Preconditions; 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 { public ArrowFlightConnectionConfigImpl(final Properties properties) { super(properties); } /** * Gets the host. * * @return the host. */ public String getHost() { return ArrowFlightConnectionProperty.HOST.getString(properties); } /** * Gets the port. * * @return the port. */ public int getPort() { return ArrowFlightConnectionProperty.PORT.getInteger(properties); } /** * Gets the host. * * @return the host. */ public String getUser() { return ArrowFlightConnectionProperty.USER.getString(properties); } /** * Gets the host. * * @return the host. */ public String getPassword() { return ArrowFlightConnectionProperty.PASSWORD.getString(properties); } public String getToken() { return ArrowFlightConnectionProperty.TOKEN.getString(properties); } /** * Gets the KeyStore path. * * @return the path. */ public String getTrustStorePath() { return ArrowFlightConnectionProperty.TRUST_STORE.getString(properties); } /** * Gets the KeyStore password. * * @return the password. */ public String getTrustStorePassword() { return ArrowFlightConnectionProperty.TRUST_STORE_PASSWORD.getString(properties); } /** * Check if the JDBC should use the trusted store files from the operating system. * * @return whether to use system trusted store certificates. */ public boolean useSystemTrustStore() { return ArrowFlightConnectionProperty.USE_SYSTEM_TRUST_STORE.getBoolean(properties); } public String getTlsRootCertificatesPath() { return ArrowFlightConnectionProperty.TLS_ROOT_CERTS.getString(properties); } public String getClientCertificatePath() { return ArrowFlightConnectionProperty.CLIENT_CERTIFICATE.getString(properties); } public String getClientKeyPath() { return ArrowFlightConnectionProperty.CLIENT_KEY.getString(properties); } /** * Whether to use TLS encryption. * * @return whether to use TLS encryption. */ public boolean useEncryption() { return ArrowFlightConnectionProperty.USE_ENCRYPTION.getBoolean(properties); } public boolean getDisableCertificateVerification() { return ArrowFlightConnectionProperty.CERTIFICATE_VERIFICATION.getBoolean(properties); } /** * Gets the thread pool size. * * @return the thread pool size. */ public int threadPoolSize() { return ArrowFlightConnectionProperty.THREAD_POOL_SIZE.getInteger(properties); } /** * Indicates if sub-connections created for stream retrieval should reuse cookies from the main * connection. */ public boolean retainCookies() { return ArrowFlightConnectionProperty.RETAIN_COOKIES.getBoolean(properties); } /** * Indicates if sub-connections created for stream retrieval should reuse bearer tokens created * from the main connection. */ public boolean retainAuth() { return ArrowFlightConnectionProperty.RETAIN_AUTH.getBoolean(properties); } /** * The catalog to which a connection is made. * * @return the catalog. */ 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}. * * @return the call options. */ public CallOption toCallOption() { final CallHeaders headers = new FlightCallHeaders(); Map headerAttributes = getHeaderAttributes(); headerAttributes.forEach(headers::insert); return new HeaderCallOption(headers); } /** * Gets which properties should be added as headers. * * @return {@link Map} */ public Map getHeaderAttributes() { Map headers = new HashMap<>(); ArrowFlightConnectionProperty[] builtInProperties = ArrowFlightConnectionProperty.values(); properties.forEach( (key, val) -> { // For built-in properties before adding new headers if (Arrays.stream(builtInProperties) .noneMatch( builtInProperty -> builtInProperty.camelName.equalsIgnoreCase(key.toString()))) { headers.put(key.toString(), val.toString()); } }); 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), PORT("port", null, Type.NUMBER, true), USER("user", null, Type.STRING, false), PASSWORD("password", null, Type.STRING, false), USE_ENCRYPTION("useEncryption", true, Type.BOOLEAN, false), CERTIFICATE_VERIFICATION("disableCertificateVerification", false, Type.BOOLEAN, false), TRUST_STORE("trustStore", null, Type.STRING, false), TRUST_STORE_PASSWORD("trustStorePassword", null, Type.STRING, false), USE_SYSTEM_TRUST_STORE("useSystemTrustStore", true, Type.BOOLEAN, false), TLS_ROOT_CERTS("tlsRootCerts", null, Type.STRING, false), CLIENT_CERTIFICATE("clientCertificate", null, Type.STRING, false), CLIENT_KEY("clientKey", null, Type.STRING, false), THREAD_POOL_SIZE("threadPoolSize", 1, Type.NUMBER, false), 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), 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; private final Type type; private final boolean required; ArrowFlightConnectionProperty( final String camelName, final Object defaultValue, final Type type, final boolean required) { this.camelName = Preconditions.checkNotNull(camelName); this.defaultValue = defaultValue; this.type = Preconditions.checkNotNull(type); this.required = required; } /** * Gets the property. * * @param properties the properties from which to fetch this property. * @return the property. */ public Object get(final Properties properties) { Preconditions.checkNotNull(properties, "Properties cannot be null."); Object value = properties.get(camelName); if (value == null) { value = properties.get(camelName.toLowerCase()); } if (required) { if (value == null) { throw new IllegalStateException( String.format("Required property not provided: <%s>.", this)); } return value; } else { return value != null ? value : defaultValue; } } /** * Gets the property as Boolean. * * @param properties the properties from which to fetch this property. * @return the property. */ public Boolean getBoolean(final Properties properties) { final String valueFromProperties = String.valueOf(get(properties)); return valueFromProperties.equals("1") || valueFromProperties.equals("true"); } /** * Gets the property as Integer. * * @param properties the properties from which to fetch this property. * @return the property. */ public Integer getInteger(final Properties properties) { final String valueFromProperties = String.valueOf(get(properties)); return valueFromProperties.equals("null") ? null : Integer.parseInt(valueFromProperties); } /** * Gets the property as String. * * @param properties the properties from which to fetch this property. * @return the property. */ public String getString(final Properties properties) { return Objects.toString(get(properties), null); } @Override public String camelName() { return camelName; } @Override public Object defaultValue() { return defaultValue; } @Override public Type type() { return type; } @Override public PropEnv wrap(final Properties properties) { throw new UnsupportedOperationException("Operation unsupported."); } @Override public boolean required() { return required; } @Override public Class valueClass() { return type.defaultValueClass(); } /** * Replaces the semicolons in the URL to the proper format. * * @param url the current connection string * @return the formatted url */ public static String replaceSemiColons(String url) { if (url != null) { url = url.replaceFirst(";", "?"); url = url.replaceAll(";", "&"); } return url; } } } AvaticaParameterBinder.java000066400000000000000000000275501515451615200405400ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; 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; import org.apache.arrow.driver.jdbc.converter.impl.DurationAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.FixedSizeBinaryAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.FixedSizeListAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.FloatingPointAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.IntAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.IntervalAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.LargeBinaryAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.LargeListAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.LargeUtf8AvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.ListAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.MapAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.NullAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.StructAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.TimeAvaticaParameterConverter; 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 * the VectorSchemaRoot representing the PreparedStatement parameters. * *

    NOTE: Make sure to close the parameters VectorSchemaRoot once we're done with them. */ public class AvaticaParameterBinder { private final PreparedStatement preparedStatement; private final VectorSchemaRoot parameters; /** * Instantiate a new AvaticaParameterBinder. * * @param preparedStatement The PreparedStatement to bind parameters to. * @param bufferAllocator The BufferAllocator to use for allocating memory. */ public AvaticaParameterBinder( PreparedStatement preparedStatement, BufferAllocator bufferAllocator) { this.parameters = VectorSchemaRoot.create(preparedStatement.getParameterSchema(), bufferAllocator); this.preparedStatement = preparedStatement; } /** * Bind the given Avatica values to the prepared statement. * * @param typedValues The parameter values. */ public void bind(List typedValues) { bind(typedValues, 0); } /** * Bind the given Avatica values to the prepared statement at the given index. * * @param typedValues The parameter values. * @param index index for parameter. */ public void bind(List typedValues, int index) { if (preparedStatement.getParameterSchema().getFields().size() != typedValues.size()) { throw new IllegalStateException( String.format( "Prepared statement has %s parameters, but only received %s", preparedStatement.getParameterSchema().getFields().size(), typedValues.size())); } for (int i = 0; i < typedValues.size(); i++) { bind(parameters.getVector(i), typedValues.get(i), index); } if (!typedValues.isEmpty()) { parameters.setRowCount(index + 1); preparedStatement.setParameters(parameters); } } /** * Bind a TypedValue to the given index on the FieldVector. * * @param vector FieldVector to bind to. * @param typedValue TypedValue to bind to the vector. * @param index Vector index to bind the value at. */ private void bind(FieldVector vector, @Nullable TypedValue typedValue, int index) { try { if (typedValue == null || typedValue.value == null) { if (vector.getField().isNullable()) { vector.setNull(index); } else { throw new UnsupportedOperationException("Can't set null on non-nullable parameter"); } } else if (!vector .getField() .getType() .accept(new BinderVisitor(vector, typedValue, index))) { throw new UnsupportedOperationException( String.format("Binding to vector type %s is not yet supported", vector.getClass())); } } catch (ClassCastException e) { throw new UnsupportedOperationException( String.format( "Binding value of type %s is not yet supported for expected Arrow type %s", typedValue == null ? "null" : typedValue.type, vector.getField().getType())); } } /** * ArrowTypeVisitor that binds Avatica TypedValues to the given FieldVector at the specified * index. */ public static class BinderVisitor implements ArrowType.ArrowTypeVisitor { private final FieldVector vector; private final TypedValue typedValue; private final int index; /** * Instantiate a new BinderVisitor. * * @param vector FieldVector to bind values to. * @param value TypedValue to bind. * @param index Vector index (0-based) to bind the value to. */ public BinderVisitor(FieldVector vector, TypedValue value, int index) { this.vector = vector; this.typedValue = value; this.index = index; } @Override public Boolean visit(ArrowType.Null type) { return new NullAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.Struct type) { return new StructAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.List type) { return new ListAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.LargeList type) { return new LargeListAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.FixedSizeList type) { return new FixedSizeListAvaticaParameterConverter(type) .bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.Union type) { return new UnionAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.Map type) { return new MapAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.Int type) { return new IntAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.FloatingPoint type) { return new FloatingPointAvaticaParameterConverter(type) .bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.Utf8 type) { return new Utf8AvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.Utf8View type) { return new Utf8ViewAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.LargeUtf8 type) { return new LargeUtf8AvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.Binary type) { return new BinaryAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.BinaryView type) { return new BinaryViewAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.LargeBinary type) { return new LargeBinaryAvaticaParameterConverter(type) .bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.FixedSizeBinary type) { return new FixedSizeBinaryAvaticaParameterConverter(type) .bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.Bool type) { return new BoolAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.Decimal type) { return new DecimalAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.Date type) { return new DateAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.Time type) { return new TimeAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.Timestamp type) { return new TimestampAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.Interval type) { return new IntervalAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.Duration type) { return new DurationAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override public Boolean visit(ArrowType.ListView type) { throw new UnsupportedOperationException("Binding is not yet supported for type " + type); } @Override public Boolean visit(ArrowType.LargeListView type) { throw new UnsupportedOperationException("Binding is not yet supported for type " + type); } @Override 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); } } } ConnectionWrapper.java000066400000000000000000000234531515451615200376410ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import static com.google.common.base.Preconditions.checkNotNull; import java.sql.Array; import java.sql.Blob; import java.sql.CallableStatement; import java.sql.Clob; import java.sql.Connection; import java.sql.DatabaseMetaData; import java.sql.NClob; import java.sql.PreparedStatement; import java.sql.SQLClientInfoException; import java.sql.SQLException; import java.sql.SQLWarning; import java.sql.SQLXML; import java.sql.Savepoint; import java.sql.Statement; import java.sql.Struct; import java.util.Map; import java.util.Properties; import java.util.concurrent.Executor; import org.apache.arrow.driver.jdbc.ArrowFlightJdbcPooledConnection; /** * Auxiliary wrapper class for {@link Connection}, used on {@link ArrowFlightJdbcPooledConnection}. */ public class ConnectionWrapper implements Connection { private final Connection realConnection; public ConnectionWrapper(final Connection connection) { realConnection = checkNotNull(connection); } @Override public T unwrap(final Class type) { return type.cast(realConnection); } @Override public boolean isWrapperFor(final Class type) { return realConnection.getClass().isAssignableFrom(type); } @Override public Statement createStatement() throws SQLException { return realConnection.createStatement(); } @Override public PreparedStatement prepareStatement(final String sqlQuery) throws SQLException { return realConnection.prepareStatement(sqlQuery); } @Override public CallableStatement prepareCall(final String sqlQuery) throws SQLException { return realConnection.prepareCall(sqlQuery); } @Override public String nativeSQL(final String sqlStatement) throws SQLException { return realConnection.nativeSQL(sqlStatement); } @Override public void setAutoCommit(boolean autoCommit) throws SQLException { realConnection.setAutoCommit(autoCommit); } @Override public boolean getAutoCommit() throws SQLException { return realConnection.getAutoCommit(); } @Override public void commit() throws SQLException { realConnection.commit(); } @Override public void rollback() throws SQLException { realConnection.rollback(); } @Override public void close() throws SQLException { realConnection.close(); } @Override public boolean isClosed() throws SQLException { return realConnection.isClosed(); } @Override public DatabaseMetaData getMetaData() throws SQLException { return realConnection.getMetaData(); } @Override public void setReadOnly(final boolean readOnly) throws SQLException { realConnection.setReadOnly(readOnly); } @Override public boolean isReadOnly() throws SQLException { return realConnection.isReadOnly(); } @Override public void setCatalog(final String catalogName) throws SQLException { realConnection.setCatalog(catalogName); } @Override public String getCatalog() throws SQLException { return realConnection.getCatalog(); } @Override public void setTransactionIsolation(final int transactionIsolationId) throws SQLException { realConnection.setTransactionIsolation(transactionIsolationId); } @Override public int getTransactionIsolation() throws SQLException { return realConnection.getTransactionIsolation(); } @Override public SQLWarning getWarnings() throws SQLException { return realConnection.getWarnings(); } @Override public void clearWarnings() throws SQLException { realConnection.clearWarnings(); } @Override public Statement createStatement(final int resultSetTypeId, final int resultSetConcurrencyId) throws SQLException { return realConnection.createStatement(resultSetTypeId, resultSetConcurrencyId); } @Override public PreparedStatement prepareStatement( final String sqlQuery, final int resultSetTypeId, final int resultSetConcurrencyId) throws SQLException { return realConnection.prepareStatement(sqlQuery, resultSetTypeId, resultSetConcurrencyId); } @Override public CallableStatement prepareCall( final String query, final int resultSetTypeId, final int resultSetConcurrencyId) throws SQLException { return realConnection.prepareCall(query, resultSetTypeId, resultSetConcurrencyId); } @Override public Map> getTypeMap() throws SQLException { return realConnection.getTypeMap(); } @Override public void setTypeMap(final Map> typeNameToClass) throws SQLException { realConnection.setTypeMap(typeNameToClass); } @Override public void setHoldability(final int holdabilityId) throws SQLException { realConnection.setHoldability(holdabilityId); } @Override public int getHoldability() throws SQLException { return realConnection.getHoldability(); } @Override public Savepoint setSavepoint() throws SQLException { return realConnection.setSavepoint(); } @Override public Savepoint setSavepoint(final String savepointName) throws SQLException { return realConnection.setSavepoint(savepointName); } @Override public void rollback(final Savepoint savepoint) throws SQLException { realConnection.rollback(savepoint); } @Override public void releaseSavepoint(final Savepoint savepoint) throws SQLException { realConnection.releaseSavepoint(savepoint); } @Override public Statement createStatement( final int resultSetType, final int resultSetConcurrency, final int resultSetHoldability) throws SQLException { return realConnection.createStatement( resultSetType, resultSetConcurrency, resultSetHoldability); } @Override public PreparedStatement prepareStatement( final String sqlQuery, final int resultSetType, final int resultSetConcurrency, final int resultSetHoldability) throws SQLException { return realConnection.prepareStatement( sqlQuery, resultSetType, resultSetConcurrency, resultSetHoldability); } @Override public CallableStatement prepareCall( final String sqlQuery, final int resultSetType, final int resultSetConcurrency, final int resultSetHoldability) throws SQLException { return realConnection.prepareCall( sqlQuery, resultSetType, resultSetConcurrency, resultSetHoldability); } @Override public PreparedStatement prepareStatement(final String sqlQuery, final int autoGeneratedKeysId) throws SQLException { return realConnection.prepareStatement(sqlQuery, autoGeneratedKeysId); } @Override public PreparedStatement prepareStatement(final String sqlQuery, final int[] columnIndices) throws SQLException { return realConnection.prepareStatement(sqlQuery, columnIndices); } @Override public PreparedStatement prepareStatement(final String sqlQuery, final String[] columnNames) throws SQLException { return realConnection.prepareStatement(sqlQuery, columnNames); } @Override public Clob createClob() throws SQLException { return realConnection.createClob(); } @Override public Blob createBlob() throws SQLException { return realConnection.createBlob(); } @Override public NClob createNClob() throws SQLException { return realConnection.createNClob(); } @Override public SQLXML createSQLXML() throws SQLException { return realConnection.createSQLXML(); } @Override public boolean isValid(final int timeout) throws SQLException { return realConnection.isValid(timeout); } @Override public void setClientInfo(final String propertyName, final String propertyValue) throws SQLClientInfoException { realConnection.setClientInfo(propertyName, propertyValue); } @Override public void setClientInfo(final Properties properties) throws SQLClientInfoException { realConnection.setClientInfo(properties); } @Override public String getClientInfo(final String propertyName) throws SQLException { return realConnection.getClientInfo(propertyName); } @Override public Properties getClientInfo() throws SQLException { return realConnection.getClientInfo(); } @Override public Array createArrayOf(final String typeName, final Object[] elements) throws SQLException { return realConnection.createArrayOf(typeName, elements); } @Override public Struct createStruct(final String typeName, final Object[] attributes) throws SQLException { return realConnection.createStruct(typeName, attributes); } @Override public void setSchema(final String schemaName) throws SQLException { realConnection.setSchema(schemaName); } @Override public String getSchema() throws SQLException { return realConnection.getSchema(); } @Override public void abort(final Executor executor) throws SQLException { realConnection.abort(executor); } @Override public void setNetworkTimeout(final Executor executor, final int timeoutInMillis) throws SQLException { realConnection.setNetworkTimeout(executor, timeoutInMillis); } @Override public int getNetworkTimeout() throws SQLException { return realConnection.getNetworkTimeout(); } } ConvertUtils.java000066400000000000000000000270661515451615200366460ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import java.util.List; import java.util.Map; import java.util.stream.Collectors; import java.util.stream.Stream; 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; import org.apache.arrow.driver.jdbc.converter.impl.DurationAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.FixedSizeBinaryAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.FixedSizeListAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.FloatingPointAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.IntAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.IntervalAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.LargeBinaryAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.LargeListAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.LargeUtf8AvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.ListAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.MapAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.NullAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.StructAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.TimeAvaticaParameterConverter; 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.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; import org.apache.calcite.avatica.proto.Common; /** Convert objects between Arrow and Avatica. */ public final class ConvertUtils { private ConvertUtils() {} /** * Convert Fields To Column MetaData List functions. * * @param fields list of {@link Field}. * @return list of {@link ColumnMetaData}. */ public static List convertArrowFieldsToColumnMetaDataList( final List fields) { return Stream.iterate(0, Math::incrementExact) .limit(fields.size()) .map( index -> { final Field field = fields.get(index); final ArrowType fieldType = field.getType(); final Common.ColumnMetaData.Builder builder = Common.ColumnMetaData.newBuilder() .setOrdinal(index) .setColumnName(field.getName()) .setLabel(field.getName()); setOnColumnMetaDataBuilder(builder, field.getMetadata()); builder.setType( Common.AvaticaType.newBuilder() .setId(SqlTypes.getSqlTypeIdFromArrowType(fieldType)) .setName(SqlTypes.getSqlTypeNameFromArrowType(fieldType)) .build()); return ColumnMetaData.fromProto(builder.build()); }) .collect(Collectors.toList()); } /** * Set on Column MetaData Builder. * * @param builder {@link Common.ColumnMetaData.Builder} * @param metadataMap {@link Map} */ public static void setOnColumnMetaDataBuilder( final Common.ColumnMetaData.Builder builder, final Map metadataMap) { final FlightSqlColumnMetadata columnMetadata = new FlightSqlColumnMetadata(metadataMap); final String catalogName = columnMetadata.getCatalogName(); if (catalogName != null) { builder.setCatalogName(catalogName); } final String schemaName = columnMetadata.getSchemaName(); if (schemaName != null) { builder.setSchemaName(schemaName); } final String tableName = columnMetadata.getTableName(); if (tableName != null) { builder.setTableName(tableName); } final Integer precision = columnMetadata.getPrecision(); if (precision != null) { builder.setPrecision(precision); } final Integer scale = columnMetadata.getScale(); if (scale != null) { builder.setScale(scale); } final Boolean isAutoIncrement = columnMetadata.isAutoIncrement(); if (isAutoIncrement != null) { builder.setAutoIncrement(isAutoIncrement); } final Boolean caseSensitive = columnMetadata.isCaseSensitive(); if (caseSensitive != null) { builder.setCaseSensitive(caseSensitive); } final Boolean readOnly = columnMetadata.isReadOnly(); if (readOnly != null) { builder.setReadOnly(readOnly); } final Boolean searchable = columnMetadata.isSearchable(); if (searchable != null) { builder.setSearchable(searchable); } final String remarks = columnMetadata.getRemarks(); if (remarks != null) { builder.setLabel(remarks); } } /** * Convert Fields To Avatica Parameters. * * @param fields list of {@link Field}. * @return list of {@link AvaticaParameter}. */ public static List convertArrowFieldsToAvaticaParameters( final List fields) { return fields.stream() .map(field -> field.getType().accept(new ConverterVisitor(field))) .collect(Collectors.toList()); } private static class ConverterVisitor implements ArrowType.ArrowTypeVisitor { private final Field field; private ConverterVisitor(Field field) { this.field = field; } @Override public AvaticaParameter visit(ArrowType.Null type) { return new NullAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.Struct type) { return new StructAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.List type) { return new ListAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.LargeList type) { return new LargeListAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.FixedSizeList type) { return new FixedSizeListAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.Union type) { return new UnionAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.Map type) { return new MapAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.Int type) { return new IntAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.FloatingPoint type) { return new FloatingPointAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.Utf8 type) { return new Utf8AvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.Utf8View type) { return new Utf8ViewAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.LargeUtf8 type) { return new LargeUtf8AvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.Binary type) { return new BinaryAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.BinaryView type) { return new BinaryViewAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.LargeBinary type) { return new LargeBinaryAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.FixedSizeBinary type) { return new FixedSizeBinaryAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.Bool type) { return new BoolAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.Decimal type) { return new DecimalAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.Date type) { return new DateAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.Time type) { return new TimeAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.Timestamp type) { return new TimestampAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.Interval type) { return new IntervalAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.Duration type) { return new DurationAvaticaParameterConverter(type).createParameter(field); } @Override public AvaticaParameter visit(ArrowType.ListView type) { throw new UnsupportedOperationException( "AvaticaParameter not yet supported for type " + type); } @Override public AvaticaParameter visit(ArrowType.LargeListView type) { throw new UnsupportedOperationException( "AvaticaParameter not yet supported for type " + type); } @Override 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); } } } DateTimeUtils.java000066400000000000000000000046531515451615200367170ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import static org.apache.calcite.avatica.util.DateTimeUtils.MILLIS_PER_DAY; import java.sql.Timestamp; import java.time.LocalDate; import java.time.LocalDateTime; import java.time.LocalTime; import java.util.Calendar; import java.util.TimeZone; import java.util.concurrent.TimeUnit; /** Datetime utility functions. */ public class DateTimeUtils { private DateTimeUtils() { // Prevent instantiation. } /** Subtracts given Calendar's TimeZone offset from epoch milliseconds. */ public static long applyCalendarOffset(long milliseconds, Calendar calendar) { if (calendar == null) { calendar = Calendar.getInstance(); } final TimeZone tz = calendar.getTimeZone(); final TimeZone defaultTz = TimeZone.getDefault(); if (tz != defaultTz) { milliseconds -= tz.getOffset(milliseconds) - defaultTz.getOffset(milliseconds); } return milliseconds; } /** * Converts Epoch millis to a {@link Timestamp} object. * * @param millisWithCalendar the Timestamp in Epoch millis * @return a {@link Timestamp} object representing the given Epoch millis */ public static Timestamp getTimestampValue(long millisWithCalendar) { long milliseconds = millisWithCalendar; if (milliseconds < 0) { // LocalTime#ofNanoDay only accepts positive values milliseconds -= ((milliseconds / MILLIS_PER_DAY) - 1) * MILLIS_PER_DAY; } return Timestamp.valueOf( LocalDateTime.of( LocalDate.ofEpochDay(millisWithCalendar / MILLIS_PER_DAY), LocalTime.ofNanoOfDay(TimeUnit.MILLISECONDS.toNanos(milliseconds % MILLIS_PER_DAY)))); } } FlightEndpointDataQueue.java000066400000000000000000000223451515451615200407150ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import static java.lang.String.format; import static java.util.Collections.synchronizedSet; import static org.apache.arrow.util.Preconditions.checkNotNull; import static org.apache.arrow.util.Preconditions.checkState; import java.sql.SQLException; import java.sql.SQLTimeoutException; import java.util.Collection; import java.util.HashSet; import java.util.Set; import java.util.concurrent.CancellationException; import java.util.concurrent.CompletionService; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorCompletionService; import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import org.apache.arrow.driver.jdbc.client.CloseableEndpointStreamPair; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.flight.FlightStream; import org.apache.calcite.avatica.AvaticaConnection; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * Auxiliary class used to handle consuming of multiple {@link FlightStream}. * *

    The usage follows this routine: * *

      *
    1. Create a FlightStreamQueue; *
    2. Call enqueue(FlightStream) for all streams to be consumed; *
    3. Call next() to get a FlightStream that is ready to consume *
    4. Consume the given FlightStream and add it back to the queue - call * enqueue(FlightStream) *
    5. Repeat from (3) until next() returns null. *
    */ public class FlightEndpointDataQueue implements AutoCloseable { private static final Logger LOGGER = LoggerFactory.getLogger(FlightEndpointDataQueue.class); private final CompletionService completionService; private final Set> futures = synchronizedSet(new HashSet<>()); private final Set endpointsToClose = synchronizedSet(new HashSet<>()); private final AtomicBoolean closed = new AtomicBoolean(); /** Instantiate a new FlightStreamQueue. */ protected FlightEndpointDataQueue( final CompletionService executorService) { completionService = checkNotNull(executorService); } /** * Creates a new {@link FlightEndpointDataQueue} from the provided {@link ExecutorService}. * * @param service the service from which to create a new queue. * @return a new queue. */ public static FlightEndpointDataQueue createNewQueue(final ExecutorService service) { return new FlightEndpointDataQueue(new ExecutorCompletionService<>(service)); } /** * Gets whether this queue is closed. * * @return a boolean indicating whether this resource is closed. */ public boolean isClosed() { return closed.get(); } /** Auxiliary functional interface for getting ready-to-consume FlightStreams. */ @FunctionalInterface interface EndpointStreamSupplier { Future get() throws SQLException; } private CloseableEndpointStreamPair next(final EndpointStreamSupplier endpointStreamSupplier) throws SQLException { checkOpen(); while (!futures.isEmpty()) { final Future future = endpointStreamSupplier.get(); futures.remove(future); try { final CloseableEndpointStreamPair endpoint = future.get(); // Get the next FlightStream that has a root with content. if (endpoint != null) { return endpoint; } } catch (final ExecutionException e) { // Unwrap one layer final Throwable cause = e.getCause(); if (cause instanceof FlightRuntimeException) { throw (FlightRuntimeException) cause; } throw AvaticaConnection.HELPER.wrap(e.getMessage(), e); } catch (InterruptedException | CancellationException e) { throw AvaticaConnection.HELPER.wrap(e.getMessage(), e); } } return null; } /** * Blocking request with timeout to get the next ready FlightStream in queue. * * @param timeoutValue the amount of time to be waited * @param timeoutUnit the timeoutValue time unit * @return a FlightStream that is ready to consume or null if all FlightStreams are ended. */ public CloseableEndpointStreamPair next(final long timeoutValue, final TimeUnit timeoutUnit) throws SQLException { return next( () -> { try { final Future future = completionService.poll(timeoutValue, timeoutUnit); if (future != null) { return future; } } catch (final InterruptedException e) { throw new SQLTimeoutException("Query was interrupted", e); } throw new SQLTimeoutException( String.format("Query timed out after %d %s", timeoutValue, timeoutUnit)); }); } /** * Blocking request to get the next ready FlightStream in queue. * * @return a FlightStream that is ready to consume or null if all FlightStreams are ended. */ public CloseableEndpointStreamPair next() throws SQLException { return next( () -> { try { return completionService.take(); } catch (final InterruptedException e) { throw AvaticaConnection.HELPER.wrap(e.getMessage(), e); } }); } /** Checks if this queue is open. */ public synchronized void checkOpen() { checkState(!isClosed(), format("%s closed", this.getClass().getSimpleName())); } /** Readily adds given {@link FlightStream}s to the queue. */ public void enqueue(final Collection endpointRequests) { endpointRequests.forEach(this::enqueue); } /** Adds given {@link FlightStream} to the queue. */ public synchronized void enqueue(final CloseableEndpointStreamPair endpointRequest) { checkNotNull(endpointRequest); checkOpen(); endpointsToClose.add(endpointRequest); futures.add( completionService.submit( () -> { // `FlightStream#next` will block until new data can be read or stream is over. while (endpointRequest.getStream().next()) { if (endpointRequest.getStream().getRoot().getRowCount() > 0) { return endpointRequest; } } return null; })); } private static boolean isCallStatusCancelled(final Exception e) { return e.getCause() instanceof FlightRuntimeException && ((FlightRuntimeException) e.getCause()).status().code() == CallStatus.CANCELLED.code(); } @Override public synchronized void close() throws SQLException { if (isClosed()) { return; } final Set exceptions = new HashSet<>(); try { for (final CloseableEndpointStreamPair endpointToClose : endpointsToClose) { try { endpointToClose.getStream().cancel("Cancelling this FlightStream.", null); } catch (final Exception e) { final String errorMsg = "Failed to cancel a FlightStream."; LOGGER.error(errorMsg, e); exceptions.add(new SQLException(errorMsg, e)); } } futures.forEach( future -> { try { // TODO: Consider adding a hardcoded timeout? future.get(); } catch (final InterruptedException | ExecutionException e) { // Ignore if future is already cancelled if (!isCallStatusCancelled(e)) { final String errorMsg = "Failed consuming a future during close."; LOGGER.error(errorMsg, e); exceptions.add(new SQLException(errorMsg, e)); } } }); for (final CloseableEndpointStreamPair endpointToClose : endpointsToClose) { try { endpointToClose.close(); } catch (final Exception e) { final String errorMsg = "Failed to close a FlightStream."; LOGGER.error(errorMsg, e); exceptions.add(new SQLException(errorMsg, e)); } } } finally { endpointsToClose.clear(); futures.clear(); closed.set(true); } if (!exceptions.isEmpty()) { final SQLException sqlException = new SQLException("Failed to close streams."); exceptions.forEach(sqlException::setNextException); throw sqlException; } } } IntervalStringUtils.java000066400000000000000000000053261515451615200401740ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import java.time.Duration; import java.time.Period; import org.apache.arrow.vector.util.DateUtility; /** * Utility class to format periods similar to Oracle's representation of "INTERVAL * to *" data * type. */ public final class IntervalStringUtils { /** Constructor Method of class. */ private IntervalStringUtils() {} /** * Formats a period similar to Oracle INTERVAL YEAR TO MONTH data type
    * . For example, the string "+21-02" defines an interval of 21 years and 2 months. */ public static String formatIntervalYear(final Period p) { long months = p.toTotalMonths(); boolean neg = false; if (months < 0) { months = -months; neg = true; } final int years = (int) (months / DateUtility.yearsToMonths); months = months % DateUtility.yearsToMonths; return String.format("%c%03d-%02d", neg ? '-' : '+', years, months); } /** * Formats a period similar to Oracle INTERVAL DAY TO SECOND data type.
    * . For example, the string "-001 18:25:16.766" defines an interval of - 1 day 18 hours 25 * minutes 16 seconds and 766 milliseconds. */ public static String formatIntervalDay(final Duration d) { long millis = d.toMillis(); boolean neg = false; if (millis < 0) { millis = -millis; neg = true; } final int days = (int) (millis / DateUtility.daysToStandardMillis); millis = millis % DateUtility.daysToStandardMillis; final int hours = (int) (millis / DateUtility.hoursToMillis); millis = millis % DateUtility.hoursToMillis; final int minutes = (int) (millis / DateUtility.minutesToMillis); millis = millis % DateUtility.minutesToMillis; final int seconds = (int) (millis / DateUtility.secondsToMillis); millis = millis % DateUtility.secondsToMillis; return String.format( "%c%03d %02d:%02d:%02d.%03d", neg ? '-' : '+', days, hours, minutes, seconds, millis); } } SqlTypes.java000066400000000000000000000134001515451615200357540ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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 { typeIdToName.put(Types.BIT, "BIT"); typeIdToName.put(Types.TINYINT, "TINYINT"); typeIdToName.put(Types.SMALLINT, "SMALLINT"); typeIdToName.put(Types.INTEGER, "INTEGER"); typeIdToName.put(Types.BIGINT, "BIGINT"); typeIdToName.put(Types.FLOAT, "FLOAT"); typeIdToName.put(Types.REAL, "REAL"); typeIdToName.put(Types.DOUBLE, "DOUBLE"); typeIdToName.put(Types.NUMERIC, "NUMERIC"); typeIdToName.put(Types.DECIMAL, "DECIMAL"); typeIdToName.put(Types.CHAR, "CHAR"); typeIdToName.put(Types.VARCHAR, "VARCHAR"); typeIdToName.put(Types.LONGVARCHAR, "LONGVARCHAR"); typeIdToName.put(Types.DATE, "DATE"); typeIdToName.put(Types.TIME, "TIME"); typeIdToName.put(Types.TIMESTAMP, "TIMESTAMP"); typeIdToName.put(Types.BINARY, "BINARY"); typeIdToName.put(Types.VARBINARY, "VARBINARY"); typeIdToName.put(Types.LONGVARBINARY, "LONGVARBINARY"); typeIdToName.put(Types.NULL, "NULL"); typeIdToName.put(Types.OTHER, "OTHER"); typeIdToName.put(Types.JAVA_OBJECT, "JAVA_OBJECT"); typeIdToName.put(Types.DISTINCT, "DISTINCT"); typeIdToName.put(Types.STRUCT, "STRUCT"); typeIdToName.put(Types.ARRAY, "ARRAY"); typeIdToName.put(Types.BLOB, "BLOB"); typeIdToName.put(Types.CLOB, "CLOB"); typeIdToName.put(Types.REF, "REF"); typeIdToName.put(Types.DATALINK, "DATALINK"); typeIdToName.put(Types.BOOLEAN, "BOOLEAN"); typeIdToName.put(Types.ROWID, "ROWID"); typeIdToName.put(Types.NCHAR, "NCHAR"); typeIdToName.put(Types.NVARCHAR, "NVARCHAR"); typeIdToName.put(Types.LONGNVARCHAR, "LONGNVARCHAR"); typeIdToName.put(Types.NCLOB, "NCLOB"); typeIdToName.put(Types.SQLXML, "SQLXML"); typeIdToName.put(Types.REF_CURSOR, "REF_CURSOR"); typeIdToName.put(Types.TIME_WITH_TIMEZONE, "TIME_WITH_TIMEZONE"); typeIdToName.put(Types.TIMESTAMP_WITH_TIMEZONE, "TIMESTAMP_WITH_TIMEZONE"); } /** * Convert given {@link ArrowType} to its corresponding SQL type name. * * @param arrowType type to convert from * @return corresponding SQL type name. * @see java.sql.Types */ public static String getSqlTypeNameFromArrowType(ArrowType arrowType) { final int typeId = getSqlTypeIdFromArrowType(arrowType); return typeIdToName.get(typeId); } /** * Convert given {@link ArrowType} to its corresponding SQL type ID. * * @param arrowType type to convert from * @return corresponding SQL type ID. * @see java.sql.Types */ public static int getSqlTypeIdFromArrowType(ArrowType arrowType) { final ArrowType.ArrowTypeID typeID = arrowType.getTypeID(); switch (typeID) { case Int: final int bitWidth = ((ArrowType.Int) arrowType).getBitWidth(); switch (bitWidth) { case 8: return Types.TINYINT; case 16: return Types.SMALLINT; case 32: return Types.INTEGER; case 64: return Types.BIGINT; default: break; } 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; case Date: return Types.DATE; case Time: return Types.TIME; case 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: return Types.DECIMAL; case FloatingPoint: final FloatingPointPrecision floatingPointPrecision = ((ArrowType.FloatingPoint) arrowType).getPrecision(); switch (floatingPointPrecision) { case DOUBLE: return Types.DOUBLE; case SINGLE: return Types.FLOAT; default: break; } break; case List: case FixedSizeList: case LargeList: return Types.ARRAY; case Struct: case Duration: case Interval: case Map: case Union: return Types.JAVA_OBJECT; case NONE: case Null: return Types.NULL; default: break; } throw new IllegalArgumentException("Unsupported ArrowType " + arrowType); } } UrlParser.java000066400000000000000000000044271515451615200361200ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import java.io.UnsupportedEncodingException; import java.net.URLDecoder; import java.util.HashMap; import java.util.Map; /** URL Parser for extracting key values from a connection string. */ public final class UrlParser { private UrlParser() {} /** * Parse URL key value parameters. * *

    URL-decodes keys and values. * * @param url {@link String} * @return {@link Map} */ @SuppressWarnings("StringSplitter") public static Map parse(String url, String separator) { Map resultMap = new HashMap<>(); if (url != null) { String[] keyValues = url.split(separator); for (String keyValue : keyValues) { try { int separatorKey = keyValue.indexOf("="); // Find the first equal sign to split key and value. if (separatorKey != -1) { // Avoid crashes when not finding an equal sign in the property value. String key = keyValue.substring(0, separatorKey); key = URLDecoder.decode(key, "UTF-8"); String value = ""; if (!keyValue.endsWith("=")) { // Avoid crashes for empty values. value = keyValue.substring(separatorKey + 1); } value = URLDecoder.decode(value, "UTF-8"); resultMap.put(key, value); } } catch (UnsupportedEncodingException e) { throw new RuntimeException(e); } } } return resultMap; } } VectorSchemaRootTransformer.java000066400000000000000000000147111515451615200416500ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import java.util.ArrayList; import java.util.Collection; import java.util.List; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; 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; /** Converts Arrow's {@link VectorSchemaRoot} format to one JDBC would expect. */ @FunctionalInterface public interface VectorSchemaRootTransformer { VectorSchemaRoot transform(VectorSchemaRoot originalRoot, VectorSchemaRoot transformedRoot) throws Exception; /** Transformer's helper class; builds a new {@link VectorSchemaRoot}. */ class Builder { private final Schema schema; private final BufferAllocator bufferAllocator; private final List newFields = new ArrayList<>(); private final Collection tasks = new ArrayList<>(); /** * Constructor for the VectorSchemaRootTransformer's Builder. * * @param schema The Arrow schema. * @param bufferAllocator The BufferAllocator to use for allocating memory. */ public Builder(final Schema schema, final BufferAllocator bufferAllocator) { this.schema = schema; this.bufferAllocator = bufferAllocator.newChildAllocator( "VectorSchemaRootTransformer", 0, bufferAllocator.getLimit()); } /** * Add task to transform a vector to a new vector renaming it. This also adds * transformedVectorName to the transformed {@link VectorSchemaRoot} schema. * * @param originalVectorName Name of the original vector to be transformed. * @param transformedVectorName Name of the vector that is the result of the transformation. * @return a VectorSchemaRoot instance with a task to rename a field vector. */ public Builder renameFieldVector( final String originalVectorName, final String transformedVectorName) { tasks.add( (originalRoot, transformedRoot) -> { final FieldVector originalVector = originalRoot.getVector(originalVectorName); final FieldVector transformedVector = transformedRoot.getVector(transformedVectorName); final ArrowType originalType = originalVector.getField().getType(); final ArrowType transformedType = transformedVector.getField().getType(); if (!originalType.equals(transformedType)) { throw new IllegalArgumentException( String.format( "Cannot transfer vector with field type %s to %s", originalType, transformedType)); } if (originalVector instanceof BaseVariableWidthVector) { ((BaseVariableWidthVector) originalVector) .transferTo(((BaseVariableWidthVector) transformedVector)); } else if (originalVector instanceof BaseFixedWidthVector) { ((BaseFixedWidthVector) originalVector) .transferTo(((BaseFixedWidthVector) transformedVector)); } else { throw new IllegalStateException( String.format("Cannot transfer vector of type %s", originalVector.getClass())); } }); final Field originalField = schema.findField(originalVectorName); newFields.add( new Field( transformedVectorName, new FieldType( originalField.isNullable(), originalField.getType(), originalField.getDictionary(), originalField.getMetadata()), originalField.getChildren())); return this; } /** * Adds an empty field to the transformed {@link VectorSchemaRoot} schema. * * @param fieldName Name of the field to be added. * @param fieldType Type of the field to be added. * @return a VectorSchemaRoot instance with the current tasks. */ public Builder addEmptyField(final String fieldName, final Types.MinorType fieldType) { newFields.add(Field.nullable(fieldName, fieldType.getType())); return this; } /** * Adds an empty field to the transformed {@link VectorSchemaRoot} schema. * * @param fieldName Name of the field to be added. * @param fieldType Type of the field to be added. * @return a VectorSchemaRoot instance with the current tasks. */ public Builder addEmptyField(final String fieldName, final ArrowType fieldType) { newFields.add(Field.nullable(fieldName, fieldType)); return this; } /** * Build the {@link VectorSchemaRoot} with applied transformation tasks. * * @return The built {@link VectorSchemaRoot}. */ public VectorSchemaRootTransformer build() { return (originalRoot, transformedRoot) -> { if (transformedRoot == null) { transformedRoot = VectorSchemaRoot.create(new Schema(newFields), bufferAllocator); } for (final Task task : tasks) { task.run(originalRoot, transformedRoot); } transformedRoot.setRowCount(originalRoot.getRowCount()); originalRoot.clear(); return transformedRoot; }; } /** * Functional interface used to a task to transform a VectorSchemaRoot into a new * VectorSchemaRoot. */ @FunctionalInterface interface Task { void run(VectorSchemaRoot originalRoot, VectorSchemaRoot transformedRoot); } } } apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/resources/000077500000000000000000000000001515451615200260005ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/resources/META-INF/000077500000000000000000000000001515451615200271405ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/resources/META-INF/services/000077500000000000000000000000001515451615200307635ustar00rootroot00000000000000java.sql.Driver000066400000000000000000000014771515451615200336110ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/resources/META-INF/services# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. org.apache.arrow.driver.jdbc.ArrowFlightJdbcDriver apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/resources/properties/000077500000000000000000000000001515451615200301745ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/main/resources/properties/flight.properties000066400000000000000000000017761515451615200336020ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. org.apache.arrow.flight.jdbc-driver.name=${project.name} org.apache.arrow.flight.jdbc-driver.version=${project.version} org.apache.arrow.flight.name=${project.groupId}\:${project.artifactId} org.apache.arrow.flight.version=${project.version} apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/000077500000000000000000000000001515451615200240215ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/000077500000000000000000000000001515451615200247425ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/000077500000000000000000000000001515451615200255315ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/000077500000000000000000000000001515451615200267525ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200301045ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/000077500000000000000000000000001515451615200313775ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/000077500000000000000000000000001515451615200323015ustar00rootroot00000000000000ArrowDatabaseMetadataTest.java000066400000000000000000002317101515451615200401110ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 com.google.protobuf.ByteString.copyFrom; import static java.lang.String.format; import static java.sql.Types.BIGINT; import static java.sql.Types.BIT; import static java.sql.Types.INTEGER; import static java.sql.Types.JAVA_OBJECT; import static java.util.Collections.singletonList; 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.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; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; 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 com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.protobuf.Message; import java.sql.Connection; import java.sql.DatabaseMetaData; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.function.Consumer; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; 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; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetPrimaryKeys; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTableTypes; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTables; import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedSubqueries; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.UInt1Vector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; 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.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; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; /** Class containing the tests from the {@link ArrowDatabaseMetadata}. */ @SuppressWarnings("DoubleBraceInitialization") public class ArrowDatabaseMetadataTest { public static final boolean EXPECTED_MAX_ROW_SIZE_INCLUDES_BLOBS = false; private static final MockFlightSqlProducer FLIGHT_SQL_PRODUCER = new MockFlightSqlProducer(); private static final MockFlightSqlProducer FLIGHT_SQL_PRODUCER_EMPTY_SQLINFO = new MockFlightSqlProducer(); @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION = FlightServerTestExtension.createStandardTestExtension(FLIGHT_SQL_PRODUCER); @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_EMPTY_SQLINFO_TEST_RULE = FlightServerTestExtension.createStandardTestExtension(FLIGHT_SQL_PRODUCER_EMPTY_SQLINFO); private static final int ROW_COUNT = 10; private static final List> EXPECTED_GET_CATALOGS_RESULTS = range(0, ROW_COUNT) .mapToObj(i -> format("catalog #%d", i)) .map(Object.class::cast) .map(Collections::singletonList) .collect(toList()); private static final List> EXPECTED_GET_TABLE_TYPES_RESULTS = range(0, ROW_COUNT) .mapToObj(i -> format("table_type #%d", i)) .map(Object.class::cast) .map(Collections::singletonList) .collect(toList()); private static final List> EXPECTED_GET_TABLES_RESULTS = range(0, ROW_COUNT) .mapToObj( i -> new Object[] { format("catalog_name #%d", i), format("db_schema_name #%d", i), format("table_name #%d", i), format("table_type #%d", i), // TODO Add these fields to FlightSQL, as it's currently not possible to fetch // them. null, null, null, null, null, null }) .map(Arrays::asList) .collect(toList()); private static final List> EXPECTED_GET_SCHEMAS_RESULTS = range(0, ROW_COUNT) .mapToObj( i -> new Object[] {format("db_schema_name #%d", i), format("catalog_name #%d", i)}) .map(Arrays::asList) .collect(toList()); private static final List> EXPECTED_GET_EXPORTED_AND_IMPORTED_KEYS_RESULTS = range(0, ROW_COUNT) .mapToObj( i -> new Object[] { format("pk_catalog_name #%d", i), format("pk_db_schema_name #%d", i), format("pk_table_name #%d", i), format("pk_column_name #%d", i), format("fk_catalog_name #%d", i), format("fk_db_schema_name #%d", i), format("fk_table_name #%d", i), format("fk_column_name #%d", i), i, format("fk_key_name #%d", i), format("pk_key_name #%d", i), (byte) i, (byte) i, // TODO Add this field to FlightSQL, as it's currently not possible to fetch it. null }) .map(Arrays::asList) .collect(toList()); private static final List> EXPECTED_CROSS_REFERENCE_RESULTS = EXPECTED_GET_EXPORTED_AND_IMPORTED_KEYS_RESULTS; private static final List> EXPECTED_PRIMARY_KEYS_RESULTS = range(0, ROW_COUNT) .mapToObj( i -> new Object[] { format("catalog_name #%d", i), format("db_schema_name #%d", i), format("table_name #%d", i), format("column_name #%d", i), i, format("key_name #%d", i) }) .map(Arrays::asList) .collect(toList()); private static final List FIELDS_GET_IMPORTED_EXPORTED_KEYS = ImmutableList.of( "PKTABLE_CAT", "PKTABLE_SCHEM", "PKTABLE_NAME", "PKCOLUMN_NAME", "FKTABLE_CAT", "FKTABLE_SCHEM", "FKTABLE_NAME", "FKCOLUMN_NAME", "KEY_SEQ", "FK_NAME", "PK_NAME", "UPDATE_RULE", "DELETE_RULE", "DEFERRABILITY"); private static final List FIELDS_GET_CROSS_REFERENCE = FIELDS_GET_IMPORTED_EXPORTED_KEYS; private static final String TARGET_TABLE = "TARGET_TABLE"; private static final String TARGET_FOREIGN_TABLE = "FOREIGN_TABLE"; private static final String EXPECTED_DATABASE_PRODUCT_NAME = "Test Server Name"; private static final String EXPECTED_DATABASE_PRODUCT_VERSION = "v0.0.1-alpha"; private static final String EXPECTED_IDENTIFIER_QUOTE_STRING = "\""; private static final boolean EXPECTED_IS_READ_ONLY = true; private static final String EXPECTED_SQL_KEYWORDS = "ADD, ADD CONSTRAINT, ALTER, ALTER TABLE, ANY, USER, TABLE"; private static final String EXPECTED_NUMERIC_FUNCTIONS = "ABS(), ACOS(), ASIN(), ATAN(), CEIL(), CEILING(), COT()"; private static final String EXPECTED_STRING_FUNCTIONS = "ASCII, CHAR, CHARINDEX, CONCAT, CONCAT_WS, FORMAT, LEFT"; private static final String EXPECTED_SYSTEM_FUNCTIONS = "CAST, CONVERT, CHOOSE, ISNULL, IS_NUMERIC, IIF, TRY_CAST"; private static final String EXPECTED_TIME_DATE_FUNCTIONS = "GETDATE(), DATEPART(), DATEADD(), DATEDIFF()"; private static final String EXPECTED_SEARCH_STRING_ESCAPE = "\\"; private static final String EXPECTED_EXTRA_NAME_CHARACTERS = ""; private static final boolean EXPECTED_SUPPORTS_COLUMN_ALIASING = true; private static final boolean EXPECTED_NULL_PLUS_NULL_IS_NULL = true; private static final boolean EXPECTED_SQL_SUPPORTS_CONVERT = true; private static final boolean EXPECTED_INVALID_SQL_SUPPORTS_CONVERT = false; private static final boolean EXPECTED_SUPPORTS_TABLE_CORRELATION_NAMES = true; private static final boolean EXPECTED_SUPPORTS_DIFFERENT_TABLE_CORRELATION_NAMES = false; private static final boolean EXPECTED_EXPRESSIONS_IN_ORDER_BY = true; private static final boolean EXPECTED_SUPPORTS_ORDER_BY_UNRELATED = true; private static final boolean EXPECTED_SUPPORTS_GROUP_BY = true; private static final boolean EXPECTED_SUPPORTS_GROUP_BY_UNRELATED = true; private static final boolean EXPECTED_SUPPORTS_LIKE_ESCAPE_CLAUSE = true; private static final boolean EXPECTED_NON_NULLABLE_COLUMNS = true; private static final boolean EXPECTED_MINIMUM_SQL_GRAMMAR = true; private static final boolean EXPECTED_CORE_SQL_GRAMMAR = true; private static final boolean EXPECTED_EXTEND_SQL_GRAMMAR = false; private static final boolean EXPECTED_ANSI92_ENTRY_LEVEL_SQL = true; private static final boolean EXPECTED_ANSI92_INTERMEDIATE_SQL = true; private static final boolean EXPECTED_ANSI92_FULL_SQL = false; private static final String EXPECTED_SCHEMA_TERM = "schema"; private static final String EXPECTED_PROCEDURE_TERM = "procedure"; private static final String EXPECTED_CATALOG_TERM = "catalog"; private static final boolean EXPECTED_SUPPORTS_INTEGRITY_ENHANCEMENT_FACILITY = true; private static final boolean EXPECTED_SUPPORTS_OUTER_JOINS = true; private static final boolean EXPECTED_SUPPORTS_FULL_OUTER_JOINS = true; private static final boolean EXPECTED_SUPPORTS_LIMITED_JOINS = false; private static final boolean EXPECTED_CATALOG_AT_START = true; private static final boolean EXPECTED_SCHEMAS_IN_PROCEDURE_CALLS = true; private static final boolean EXPECTED_SCHEMAS_IN_INDEX_DEFINITIONS = true; private static final boolean EXPECTED_SCHEMAS_IN_PRIVILEGE_DEFINITIONS = false; private static final boolean EXPECTED_CATALOGS_IN_INDEX_DEFINITIONS = true; private static final boolean EXPECTED_CATALOGS_IN_PRIVILEGE_DEFINITIONS = false; private static final boolean EXPECTED_POSITIONED_DELETE = true; private static final boolean EXPECTED_POSITIONED_UPDATE = false; private static final boolean EXPECTED_TYPE_FORWARD_ONLY = true; private static final boolean EXPECTED_TYPE_SCROLL_INSENSITIVE = true; private static final boolean EXPECTED_TYPE_SCROLL_SENSITIVE = false; private static final boolean EXPECTED_SELECT_FOR_UPDATE_SUPPORTED = false; private static final boolean EXPECTED_STORED_PROCEDURES_SUPPORTED = false; private static final boolean EXPECTED_SUBQUERIES_IN_COMPARISON = true; private static final boolean EXPECTED_SUBQUERIES_IN_EXISTS = false; private static final boolean EXPECTED_SUBQUERIES_IN_INS = false; private static final boolean EXPECTED_SUBQUERIES_IN_QUANTIFIEDS = false; private static final SqlSupportedSubqueries[] EXPECTED_SUPPORTED_SUBQUERIES = new SqlSupportedSubqueries[] {SqlSupportedSubqueries.SQL_SUBQUERIES_IN_COMPARISONS}; private static final boolean EXPECTED_CORRELATED_SUBQUERIES_SUPPORTED = true; private static final boolean EXPECTED_SUPPORTS_UNION = true; private static final boolean EXPECTED_SUPPORTS_UNION_ALL = true; private static final int EXPECTED_MAX_BINARY_LITERAL_LENGTH = 0; private static final int EXPECTED_MAX_CHAR_LITERAL_LENGTH = 0; private static final int EXPECTED_MAX_COLUMN_NAME_LENGTH = 1024; private static final int EXPECTED_MAX_COLUMNS_IN_GROUP_BY = 0; private static final int EXPECTED_MAX_COLUMNS_IN_INDEX = 0; private static final int EXPECTED_MAX_COLUMNS_IN_ORDER_BY = 0; private static final int EXPECTED_MAX_COLUMNS_IN_SELECT = 0; private static final int EXPECTED_MAX_CONNECTIONS = 0; private static final int EXPECTED_MAX_CURSOR_NAME_LENGTH = 1024; private static final int EXPECTED_MAX_INDEX_LENGTH = 0; private static final int EXPECTED_SCHEMA_NAME_LENGTH = 1024; private static final int EXPECTED_MAX_PROCEDURE_NAME_LENGTH = 0; private static final int EXPECTED_MAX_CATALOG_NAME_LENGTH = 1024; private static final int EXPECTED_MAX_ROW_SIZE = 0; private static final int EXPECTED_MAX_STATEMENT_LENGTH = 0; private static final int EXPECTED_MAX_STATEMENTS = 0; private static final int EXPECTED_MAX_TABLE_NAME_LENGTH = 1024; private static final int EXPECTED_MAX_TABLES_IN_SELECT = 0; private static final int EXPECTED_MAX_USERNAME_LENGTH = 1024; private static final int EXPECTED_DEFAULT_TRANSACTION_ISOLATION = 0; private static final boolean EXPECTED_TRANSACTIONS_SUPPORTED = false; private static final boolean EXPECTED_TRANSACTION_NONE = false; private static final boolean EXPECTED_TRANSACTION_READ_UNCOMMITTED = false; private static final boolean EXPECTED_TRANSACTION_READ_COMMITTED = true; private static final boolean EXPECTED_TRANSACTION_REPEATABLE_READ = false; private static final boolean EXPECTED_TRANSACTION_SERIALIZABLE = true; private static final boolean EXPECTED_DATA_DEFINITION_CAUSES_TRANSACTION_COMMIT = true; private static final boolean EXPECTED_DATA_DEFINITIONS_IN_TRANSACTIONS_IGNORED = false; private static final boolean EXPECTED_BATCH_UPDATES_SUPPORTED = true; private static final boolean EXPECTED_SAVEPOINTS_SUPPORTED = false; private static final boolean EXPECTED_NAMED_PARAMETERS_SUPPORTED = false; private static final boolean EXPECTED_LOCATORS_UPDATE_COPY = true; private static final boolean EXPECTED_STORED_FUNCTIONS_USING_CALL_SYNTAX_SUPPORTED = false; private static final List> EXPECTED_GET_COLUMNS_RESULTS; private static Connection connection; static { 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); List expectedGetColumnsIsNullable = Arrays.asList("YES", "YES", "NO"); EXPECTED_GET_COLUMNS_RESULTS = range(0, ROW_COUNT * 3) .mapToObj( i -> new Object[] { format("catalog_name #%d", i / 3), format("db_schema_name #%d", i / 3), format("table_name%d", i / 3), format("column_%d", (i % 3) + 1), expectedGetColumnsDataTypes.get(i % 3), expectedGetColumnsTypeName.get(i % 3), expectedGetColumnsColumnSize.get(i % 3), null, expectedGetColumnsDecimalDigits.get(i % 3), expectedGetColumnsRadix.get(i % 3), !Objects.equals(expectedGetColumnsIsNullable.get(i % 3), "NO") ? 1 : 0, format("column description #%d", (i % 3) + 1), null, null, null, null, (i % 3) + 1, expectedGetColumnsIsNullable.get(i % 3), null, null, null, null, "", "" }) .map(Arrays::asList) .collect(toList()); } public final ResultSetTestUtils resultSetTestUtils = new ResultSetTestUtils(); @BeforeAll public static void setUpBeforeClass() throws SQLException { connection = FLIGHT_SERVER_TEST_EXTENSION.getConnection(false); final Message commandGetCatalogs = CommandGetCatalogs.getDefaultInstance(); final Consumer commandGetCatalogsResultProducer = listener -> { try (final BufferAllocator allocator = new RootAllocator(); final VectorSchemaRoot root = VectorSchemaRoot.create(Schemas.GET_CATALOGS_SCHEMA, allocator)) { final VarCharVector catalogName = (VarCharVector) root.getVector("catalog_name"); range(0, ROW_COUNT) .forEach(i -> catalogName.setSafe(i, new Text(format("catalog #%d", i)))); root.setRowCount(ROW_COUNT); listener.start(root); listener.putNext(); } catch (final Throwable throwable) { listener.error(throwable); } finally { listener.completed(); } }; FLIGHT_SQL_PRODUCER.addCatalogQuery(commandGetCatalogs, commandGetCatalogsResultProducer); 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, ROW_COUNT) .forEach(i -> tableType.setSafe(i, new Text(format("table_type #%d", i)))); root.setRowCount(ROW_COUNT); listener.start(root); listener.putNext(); } catch (final Throwable throwable) { listener.error(throwable); } finally { listener.completed(); } }; FLIGHT_SQL_PRODUCER.addCatalogQuery(commandGetTableTypes, commandGetTableTypesResultProducer); final Message commandGetTables = CommandGetTables.getDefaultInstance(); final Consumer commandGetTablesResultProducer = listener -> { try (final BufferAllocator allocator = new RootAllocator(); final VectorSchemaRoot root = VectorSchemaRoot.create(Schemas.GET_TABLES_SCHEMA_NO_SCHEMA, allocator)) { final VarCharVector catalogName = (VarCharVector) root.getVector("catalog_name"); final VarCharVector schemaName = (VarCharVector) root.getVector("db_schema_name"); final VarCharVector tableName = (VarCharVector) root.getVector("table_name"); final VarCharVector tableType = (VarCharVector) root.getVector("table_type"); range(0, ROW_COUNT) .peek(i -> catalogName.setSafe(i, new Text(format("catalog_name #%d", i)))) .peek(i -> schemaName.setSafe(i, new Text(format("db_schema_name #%d", i)))) .peek(i -> tableName.setSafe(i, new Text(format("table_name #%d", i)))) .forEach(i -> tableType.setSafe(i, new Text(format("table_type #%d", i)))); root.setRowCount(ROW_COUNT); listener.start(root); listener.putNext(); } catch (final Throwable throwable) { listener.error(throwable); } finally { listener.completed(); } }; FLIGHT_SQL_PRODUCER.addCatalogQuery(commandGetTables, commandGetTablesResultProducer); final Message commandGetTablesWithSchema = CommandGetTables.newBuilder().setIncludeSchema(true).build(); final Consumer commandGetTablesWithSchemaResultProducer = listener -> { 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(field1, field2, field3)))) .toByteArray(); final VarCharVector catalogName = (VarCharVector) root.getVector("catalog_name"); final VarCharVector schemaName = (VarCharVector) root.getVector("db_schema_name"); final VarCharVector tableName = (VarCharVector) root.getVector("table_name"); final VarCharVector tableType = (VarCharVector) root.getVector("table_type"); final VarBinaryVector tableSchema = (VarBinaryVector) root.getVector("table_schema"); range(0, ROW_COUNT) .peek(i -> catalogName.setSafe(i, new Text(format("catalog_name #%d", i)))) .peek(i -> schemaName.setSafe(i, new Text(format("db_schema_name #%d", i)))) .peek(i -> tableName.setSafe(i, new Text(format("table_name%d", i)))) .peek(i -> tableType.setSafe(i, new Text(format("table_type #%d", i)))) .forEach(i -> tableSchema.setSafe(i, filledTableSchemaBytes)); root.setRowCount(ROW_COUNT); listener.start(root); listener.putNext(); } catch (final Throwable throwable) { listener.error(throwable); } finally { listener.completed(); } }; FLIGHT_SQL_PRODUCER.addCatalogQuery( commandGetTablesWithSchema, commandGetTablesWithSchemaResultProducer); final Message commandGetDbSchemas = CommandGetDbSchemas.getDefaultInstance(); final Consumer commandGetSchemasResultProducer = listener -> { try (final BufferAllocator allocator = new RootAllocator(); final VectorSchemaRoot root = VectorSchemaRoot.create(Schemas.GET_SCHEMAS_SCHEMA, allocator)) { final VarCharVector catalogName = (VarCharVector) root.getVector("catalog_name"); final VarCharVector schemaName = (VarCharVector) root.getVector("db_schema_name"); range(0, ROW_COUNT) .peek(i -> catalogName.setSafe(i, new Text(format("catalog_name #%d", i)))) .forEach(i -> schemaName.setSafe(i, new Text(format("db_schema_name #%d", i)))); root.setRowCount(ROW_COUNT); listener.start(root); listener.putNext(); } catch (final Throwable throwable) { listener.error(throwable); } finally { listener.completed(); } }; FLIGHT_SQL_PRODUCER.addCatalogQuery(commandGetDbSchemas, commandGetSchemasResultProducer); final Message commandGetExportedKeys = CommandGetExportedKeys.newBuilder().setTable(TARGET_TABLE).build(); final Message commandGetImportedKeys = CommandGetImportedKeys.newBuilder().setTable(TARGET_TABLE).build(); final Message commandGetCrossReference = CommandGetCrossReference.newBuilder() .setPkTable(TARGET_TABLE) .setFkTable(TARGET_FOREIGN_TABLE) .build(); final Consumer commandGetExportedAndImportedKeysResultProducer = listener -> { try (final BufferAllocator allocator = new RootAllocator(); final VectorSchemaRoot root = VectorSchemaRoot.create(Schemas.GET_IMPORTED_KEYS_SCHEMA, allocator)) { final VarCharVector pkCatalogName = (VarCharVector) root.getVector("pk_catalog_name"); final VarCharVector pkSchemaName = (VarCharVector) root.getVector("pk_db_schema_name"); final VarCharVector pkTableName = (VarCharVector) root.getVector("pk_table_name"); final VarCharVector pkColumnName = (VarCharVector) root.getVector("pk_column_name"); final VarCharVector fkCatalogName = (VarCharVector) root.getVector("fk_catalog_name"); final VarCharVector fkSchemaName = (VarCharVector) root.getVector("fk_db_schema_name"); final VarCharVector fkTableName = (VarCharVector) root.getVector("fk_table_name"); final VarCharVector fkColumnName = (VarCharVector) root.getVector("fk_column_name"); final IntVector keySequence = (IntVector) root.getVector("key_sequence"); final VarCharVector fkKeyName = (VarCharVector) root.getVector("fk_key_name"); final VarCharVector pkKeyName = (VarCharVector) root.getVector("pk_key_name"); final UInt1Vector updateRule = (UInt1Vector) root.getVector("update_rule"); final UInt1Vector deleteRule = (UInt1Vector) root.getVector("delete_rule"); range(0, ROW_COUNT) .peek(i -> pkCatalogName.setSafe(i, new Text(format("pk_catalog_name #%d", i)))) .peek(i -> pkSchemaName.setSafe(i, new Text(format("pk_db_schema_name #%d", i)))) .peek(i -> pkTableName.setSafe(i, new Text(format("pk_table_name #%d", i)))) .peek(i -> pkColumnName.setSafe(i, new Text(format("pk_column_name #%d", i)))) .peek(i -> fkCatalogName.setSafe(i, new Text(format("fk_catalog_name #%d", i)))) .peek(i -> fkSchemaName.setSafe(i, new Text(format("fk_db_schema_name #%d", i)))) .peek(i -> fkTableName.setSafe(i, new Text(format("fk_table_name #%d", i)))) .peek(i -> fkColumnName.setSafe(i, new Text(format("fk_column_name #%d", i)))) .peek(i -> keySequence.setSafe(i, i)) .peek(i -> fkKeyName.setSafe(i, new Text(format("fk_key_name #%d", i)))) .peek(i -> pkKeyName.setSafe(i, new Text(format("pk_key_name #%d", i)))) .peek(i -> updateRule.setSafe(i, i)) .forEach(i -> deleteRule.setSafe(i, i)); root.setRowCount(ROW_COUNT); listener.start(root); listener.putNext(); } catch (final Throwable throwable) { listener.error(throwable); } finally { listener.completed(); } }; FLIGHT_SQL_PRODUCER.addCatalogQuery( commandGetExportedKeys, commandGetExportedAndImportedKeysResultProducer); FLIGHT_SQL_PRODUCER.addCatalogQuery( commandGetImportedKeys, commandGetExportedAndImportedKeysResultProducer); FLIGHT_SQL_PRODUCER.addCatalogQuery( commandGetCrossReference, commandGetExportedAndImportedKeysResultProducer); final Message commandGetPrimaryKeys = CommandGetPrimaryKeys.newBuilder().setTable(TARGET_TABLE).build(); final Consumer commandGetPrimaryKeysResultProducer = listener -> { try (final BufferAllocator allocator = new RootAllocator(); final VectorSchemaRoot root = VectorSchemaRoot.create(Schemas.GET_PRIMARY_KEYS_SCHEMA, allocator)) { final VarCharVector catalogName = (VarCharVector) root.getVector("catalog_name"); final VarCharVector schemaName = (VarCharVector) root.getVector("db_schema_name"); final VarCharVector tableName = (VarCharVector) root.getVector("table_name"); final VarCharVector columnName = (VarCharVector) root.getVector("column_name"); final IntVector keySequence = (IntVector) root.getVector("key_sequence"); final VarCharVector keyName = (VarCharVector) root.getVector("key_name"); range(0, ROW_COUNT) .peek(i -> catalogName.setSafe(i, new Text(format("catalog_name #%d", i)))) .peek(i -> schemaName.setSafe(i, new Text(format("db_schema_name #%d", i)))) .peek(i -> tableName.setSafe(i, new Text(format("table_name #%d", i)))) .peek(i -> columnName.setSafe(i, new Text(format("column_name #%d", i)))) .peek(i -> keySequence.setSafe(i, i)) .forEach(i -> keyName.setSafe(i, new Text(format("key_name #%d", i)))); root.setRowCount(ROW_COUNT); listener.start(root); listener.putNext(); } catch (final Throwable throwable) { listener.error(throwable); } finally { listener.completed(); } }; FLIGHT_SQL_PRODUCER.addCatalogQuery(commandGetPrimaryKeys, commandGetPrimaryKeysResultProducer); FLIGHT_SQL_PRODUCER .getSqlInfoBuilder() .withSqlOuterJoinSupportLevel(FlightSql.SqlOuterJoinsSupportLevel.SQL_FULL_OUTER_JOINS) .withFlightSqlServerName(EXPECTED_DATABASE_PRODUCT_NAME) .withFlightSqlServerVersion(EXPECTED_DATABASE_PRODUCT_VERSION) .withSqlIdentifierQuoteChar(EXPECTED_IDENTIFIER_QUOTE_STRING) .withFlightSqlServerReadOnly(EXPECTED_IS_READ_ONLY) .withSqlKeywords(EXPECTED_SQL_KEYWORDS.split("\\s*,\\s*")) .withSqlNumericFunctions(EXPECTED_NUMERIC_FUNCTIONS.split("\\s*,\\s*")) .withSqlStringFunctions(EXPECTED_STRING_FUNCTIONS.split("\\s*,\\s*")) .withSqlSystemFunctions(EXPECTED_SYSTEM_FUNCTIONS.split("\\s*,\\s*")) .withSqlDatetimeFunctions(EXPECTED_TIME_DATE_FUNCTIONS.split("\\s*,\\s*")) .withSqlSearchStringEscape(EXPECTED_SEARCH_STRING_ESCAPE) .withSqlExtraNameCharacters(EXPECTED_EXTRA_NAME_CHARACTERS) .withSqlSupportsColumnAliasing(EXPECTED_SUPPORTS_COLUMN_ALIASING) .withSqlNullPlusNullIsNull(EXPECTED_NULL_PLUS_NULL_IS_NULL) .withSqlSupportsConvert( ImmutableMap.of( SQL_CONVERT_BIT_VALUE, Arrays.asList(SQL_CONVERT_INTEGER_VALUE, SQL_CONVERT_BIGINT_VALUE))) .withSqlSupportsTableCorrelationNames(EXPECTED_SUPPORTS_TABLE_CORRELATION_NAMES) .withSqlSupportsDifferentTableCorrelationNames( EXPECTED_SUPPORTS_DIFFERENT_TABLE_CORRELATION_NAMES) .withSqlSupportsExpressionsInOrderBy(EXPECTED_EXPRESSIONS_IN_ORDER_BY) .withSqlSupportsOrderByUnrelated(EXPECTED_SUPPORTS_ORDER_BY_UNRELATED) .withSqlSupportedGroupBy(FlightSql.SqlSupportedGroupBy.SQL_GROUP_BY_UNRELATED) .withSqlSupportsLikeEscapeClause(EXPECTED_SUPPORTS_LIKE_ESCAPE_CLAUSE) .withSqlSupportsNonNullableColumns(EXPECTED_NON_NULLABLE_COLUMNS) .withSqlSupportedGrammar( FlightSql.SupportedSqlGrammar.SQL_CORE_GRAMMAR, FlightSql.SupportedSqlGrammar.SQL_MINIMUM_GRAMMAR) .withSqlAnsi92SupportedLevel( FlightSql.SupportedAnsi92SqlGrammarLevel.ANSI92_ENTRY_SQL, FlightSql.SupportedAnsi92SqlGrammarLevel.ANSI92_INTERMEDIATE_SQL) .withSqlSupportsIntegrityEnhancementFacility( EXPECTED_SUPPORTS_INTEGRITY_ENHANCEMENT_FACILITY) .withSqlSchemaTerm(EXPECTED_SCHEMA_TERM) .withSqlCatalogTerm(EXPECTED_CATALOG_TERM) .withSqlProcedureTerm(EXPECTED_PROCEDURE_TERM) .withSqlCatalogAtStart(EXPECTED_CATALOG_AT_START) .withSqlSchemasSupportedActions( FlightSql.SqlSupportedElementActions.SQL_ELEMENT_IN_PROCEDURE_CALLS, FlightSql.SqlSupportedElementActions.SQL_ELEMENT_IN_INDEX_DEFINITIONS) .withSqlCatalogsSupportedActions( FlightSql.SqlSupportedElementActions.SQL_ELEMENT_IN_INDEX_DEFINITIONS) .withSqlSupportedPositionedCommands( FlightSql.SqlSupportedPositionedCommands.SQL_POSITIONED_DELETE) .withSqlSelectForUpdateSupported(EXPECTED_SELECT_FOR_UPDATE_SUPPORTED) .withSqlStoredProceduresSupported(EXPECTED_STORED_PROCEDURES_SUPPORTED) .withSqlSubQueriesSupported(EXPECTED_SUPPORTED_SUBQUERIES) .withSqlCorrelatedSubqueriesSupported(EXPECTED_CORRELATED_SUBQUERIES_SUPPORTED) .withSqlSupportedUnions(FlightSql.SqlSupportedUnions.SQL_UNION_ALL) .withSqlMaxBinaryLiteralLength(EXPECTED_MAX_BINARY_LITERAL_LENGTH) .withSqlMaxCharLiteralLength(EXPECTED_MAX_CHAR_LITERAL_LENGTH) .withSqlMaxColumnNameLength(EXPECTED_MAX_COLUMN_NAME_LENGTH) .withSqlMaxColumnsInGroupBy(EXPECTED_MAX_COLUMNS_IN_GROUP_BY) .withSqlMaxColumnsInIndex(EXPECTED_MAX_COLUMNS_IN_INDEX) .withSqlMaxColumnsInOrderBy(EXPECTED_MAX_COLUMNS_IN_ORDER_BY) .withSqlMaxColumnsInSelect(EXPECTED_MAX_COLUMNS_IN_SELECT) .withSqlMaxConnections(EXPECTED_MAX_CONNECTIONS) .withSqlMaxCursorNameLength(EXPECTED_MAX_CURSOR_NAME_LENGTH) .withSqlMaxIndexLength(EXPECTED_MAX_INDEX_LENGTH) .withSqlDbSchemaNameLength(EXPECTED_SCHEMA_NAME_LENGTH) .withSqlMaxProcedureNameLength(EXPECTED_MAX_PROCEDURE_NAME_LENGTH) .withSqlMaxCatalogNameLength(EXPECTED_MAX_CATALOG_NAME_LENGTH) .withSqlMaxRowSize(EXPECTED_MAX_ROW_SIZE) .withSqlMaxRowSizeIncludesBlobs(EXPECTED_MAX_ROW_SIZE_INCLUDES_BLOBS) .withSqlMaxStatementLength(EXPECTED_MAX_STATEMENT_LENGTH) .withSqlMaxStatements(EXPECTED_MAX_STATEMENTS) .withSqlMaxTableNameLength(EXPECTED_MAX_TABLE_NAME_LENGTH) .withSqlMaxTablesInSelect(EXPECTED_MAX_TABLES_IN_SELECT) .withSqlMaxUsernameLength(EXPECTED_MAX_USERNAME_LENGTH) .withSqlDefaultTransactionIsolation(EXPECTED_DEFAULT_TRANSACTION_ISOLATION) .withSqlTransactionsSupported(EXPECTED_TRANSACTIONS_SUPPORTED) .withSqlSupportedTransactionsIsolationLevels( FlightSql.SqlTransactionIsolationLevel.SQL_TRANSACTION_SERIALIZABLE, FlightSql.SqlTransactionIsolationLevel.SQL_TRANSACTION_READ_COMMITTED) .withSqlDataDefinitionCausesTransactionCommit( EXPECTED_DATA_DEFINITION_CAUSES_TRANSACTION_COMMIT) .withSqlDataDefinitionsInTransactionsIgnored( EXPECTED_DATA_DEFINITIONS_IN_TRANSACTIONS_IGNORED) .withSqlSupportedResultSetTypes( FlightSql.SqlSupportedResultSetType.SQL_RESULT_SET_TYPE_FORWARD_ONLY, FlightSql.SqlSupportedResultSetType.SQL_RESULT_SET_TYPE_SCROLL_INSENSITIVE) .withSqlBatchUpdatesSupported(EXPECTED_BATCH_UPDATES_SUPPORTED) .withSqlSavepointsSupported(EXPECTED_SAVEPOINTS_SUPPORTED) .withSqlNamedParametersSupported(EXPECTED_NAMED_PARAMETERS_SUPPORTED) .withSqlLocatorsUpdateCopy(EXPECTED_LOCATORS_UPDATE_COPY) .withSqlStoredFunctionsUsingCallSyntaxSupported( EXPECTED_STORED_FUNCTIONS_USING_CALL_SYNTAX_SUPPORTED); } @AfterAll public static void tearDown() throws Exception { AutoCloseables.close(connection, FLIGHT_SQL_PRODUCER, FLIGHT_SQL_PRODUCER_EMPTY_SQLINFO); } @Test public void testGetCatalogsCanBeAccessedByIndices() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getCatalogs()) { resultSetTestUtils.testData(resultSet, EXPECTED_GET_CATALOGS_RESULTS); } } @Test public void testGetCatalogsCanBeAccessedByNames() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getCatalogs()) { resultSetTestUtils.testData( resultSet, singletonList("TABLE_CAT"), EXPECTED_GET_CATALOGS_RESULTS); } } @Test public void testTableTypesCanBeAccessedByIndices() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getTableTypes()) { resultSetTestUtils.testData(resultSet, EXPECTED_GET_TABLE_TYPES_RESULTS); } } @Test public void testTableTypesCanBeAccessedByNames() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getTableTypes()) { resultSetTestUtils.testData( resultSet, singletonList("TABLE_TYPE"), EXPECTED_GET_TABLE_TYPES_RESULTS); } } @Test public void testGetTablesCanBeAccessedByIndices() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getTables(null, null, null, null)) { resultSetTestUtils.testData(resultSet, EXPECTED_GET_TABLES_RESULTS); } } @Test public void testGetTablesCanBeAccessedByNames() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getTables(null, null, null, null)) { resultSetTestUtils.testData( resultSet, ImmutableList.of( "TABLE_CAT", "TABLE_SCHEM", "TABLE_NAME", "TABLE_TYPE", "REMARKS", "TYPE_CAT", "TYPE_SCHEM", "TYPE_NAME", "SELF_REFERENCING_COL_NAME", "REF_GENERATION"), EXPECTED_GET_TABLES_RESULTS); } } @Test public void testGetSchemasCanBeAccessedByIndices() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getSchemas()) { resultSetTestUtils.testData(resultSet, EXPECTED_GET_SCHEMAS_RESULTS); } } @Test public void testGetSchemasCanBeAccessedByNames() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getSchemas()) { resultSetTestUtils.testData( resultSet, ImmutableList.of("TABLE_SCHEM", "TABLE_CATALOG"), EXPECTED_GET_SCHEMAS_RESULTS); } } @Test public void testGetExportedKeysCanBeAccessedByIndices() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getExportedKeys(null, null, TARGET_TABLE)) { resultSetTestUtils.testData(resultSet, EXPECTED_GET_EXPORTED_AND_IMPORTED_KEYS_RESULTS); } } @Test public void testGetExportedKeysCanBeAccessedByNames() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getExportedKeys(null, null, TARGET_TABLE)) { resultSetTestUtils.testData( resultSet, FIELDS_GET_IMPORTED_EXPORTED_KEYS, EXPECTED_GET_EXPORTED_AND_IMPORTED_KEYS_RESULTS); } } @Test public void testGetImportedKeysCanBeAccessedByIndices() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getImportedKeys(null, null, TARGET_TABLE)) { resultSetTestUtils.testData(resultSet, EXPECTED_GET_EXPORTED_AND_IMPORTED_KEYS_RESULTS); } } @Test public void testGetImportedKeysCanBeAccessedByNames() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getImportedKeys(null, null, TARGET_TABLE)) { resultSetTestUtils.testData( resultSet, FIELDS_GET_IMPORTED_EXPORTED_KEYS, EXPECTED_GET_EXPORTED_AND_IMPORTED_KEYS_RESULTS); } } @Test public void testGetCrossReferenceCanBeAccessedByIndices() throws SQLException { try (final ResultSet resultSet = connection .getMetaData() .getCrossReference(null, null, TARGET_TABLE, null, null, TARGET_FOREIGN_TABLE)) { resultSetTestUtils.testData(resultSet, EXPECTED_CROSS_REFERENCE_RESULTS); } } @Test public void testGetGetCrossReferenceCanBeAccessedByNames() throws SQLException { try (final ResultSet resultSet = connection .getMetaData() .getCrossReference(null, null, TARGET_TABLE, null, null, TARGET_FOREIGN_TABLE)) { resultSetTestUtils.testData( resultSet, FIELDS_GET_CROSS_REFERENCE, EXPECTED_CROSS_REFERENCE_RESULTS); } } @Test public void testPrimaryKeysCanBeAccessedByIndices() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getPrimaryKeys(null, null, TARGET_TABLE)) { resultSetTestUtils.testData(resultSet, EXPECTED_PRIMARY_KEYS_RESULTS); } } @Test public void testPrimaryKeysCanBeAccessedByNames() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getPrimaryKeys(null, null, TARGET_TABLE)) { resultSetTestUtils.testData( resultSet, ImmutableList.of( "TABLE_CAT", "TABLE_SCHEM", "TABLE_NAME", "COLUMN_NAME", "KEY_SEQ", "PK_NAME"), EXPECTED_PRIMARY_KEYS_RESULTS); } } @Test public void testGetColumnsCanBeAccessedByIndices() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getColumns(null, null, null, null)) { resultSetTestUtils.testData(resultSet, EXPECTED_GET_COLUMNS_RESULTS); } } @Test public void testGetColumnsCanByIndicesFilteringColumnNames() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getColumns(null, null, null, "column_1")) { resultSetTestUtils.testData( resultSet, EXPECTED_GET_COLUMNS_RESULTS.stream() .filter(insideList -> Objects.equals(insideList.get(3), "column_1")) .collect(toList())); } } @Test public void testGetSqlInfo() throws SQLException { final DatabaseMetaData metaData = connection.getMetaData(); assertThat(metaData.getDatabaseProductName(), is(EXPECTED_DATABASE_PRODUCT_NAME)); assertThat(metaData.getDatabaseProductVersion(), is(EXPECTED_DATABASE_PRODUCT_VERSION)); assertThat(metaData.getIdentifierQuoteString(), is(EXPECTED_IDENTIFIER_QUOTE_STRING)); assertThat(metaData.isReadOnly(), is(EXPECTED_IS_READ_ONLY)); assertThat(metaData.getSQLKeywords(), is(EXPECTED_SQL_KEYWORDS)); assertThat(metaData.getNumericFunctions(), is(EXPECTED_NUMERIC_FUNCTIONS)); assertThat(metaData.getStringFunctions(), is(EXPECTED_STRING_FUNCTIONS)); assertThat(metaData.getSystemFunctions(), is(EXPECTED_SYSTEM_FUNCTIONS)); assertThat(metaData.getTimeDateFunctions(), is(EXPECTED_TIME_DATE_FUNCTIONS)); assertThat(metaData.getSearchStringEscape(), is(EXPECTED_SEARCH_STRING_ESCAPE)); assertThat(metaData.getExtraNameCharacters(), is(EXPECTED_EXTRA_NAME_CHARACTERS)); assertThat(metaData.supportsConvert(), is(EXPECTED_SQL_SUPPORTS_CONVERT)); assertThat(metaData.supportsConvert(BIT, INTEGER), is(EXPECTED_SQL_SUPPORTS_CONVERT)); assertThat(metaData.supportsConvert(BIT, BIGINT), is(EXPECTED_SQL_SUPPORTS_CONVERT)); assertThat( metaData.supportsConvert(BIGINT, INTEGER), is(EXPECTED_INVALID_SQL_SUPPORTS_CONVERT)); assertThat( metaData.supportsConvert(JAVA_OBJECT, INTEGER), is(EXPECTED_INVALID_SQL_SUPPORTS_CONVERT)); assertThat( metaData.supportsTableCorrelationNames(), is(EXPECTED_SUPPORTS_TABLE_CORRELATION_NAMES)); assertThat(metaData.supportsExpressionsInOrderBy(), is(EXPECTED_EXPRESSIONS_IN_ORDER_BY)); assertThat(metaData.supportsOrderByUnrelated(), is(EXPECTED_SUPPORTS_ORDER_BY_UNRELATED)); assertThat(metaData.supportsGroupBy(), is(EXPECTED_SUPPORTS_GROUP_BY)); assertThat(metaData.supportsGroupByUnrelated(), is(EXPECTED_SUPPORTS_GROUP_BY_UNRELATED)); assertThat(metaData.supportsLikeEscapeClause(), is(EXPECTED_SUPPORTS_LIKE_ESCAPE_CLAUSE)); assertThat(metaData.supportsNonNullableColumns(), is(EXPECTED_NON_NULLABLE_COLUMNS)); assertThat(metaData.supportsMinimumSQLGrammar(), is(EXPECTED_MINIMUM_SQL_GRAMMAR)); assertThat(metaData.supportsCoreSQLGrammar(), is(EXPECTED_CORE_SQL_GRAMMAR)); assertThat(metaData.supportsExtendedSQLGrammar(), is(EXPECTED_EXTEND_SQL_GRAMMAR)); assertThat(metaData.supportsANSI92EntryLevelSQL(), is(EXPECTED_ANSI92_ENTRY_LEVEL_SQL)); assertThat(metaData.supportsANSI92IntermediateSQL(), is(EXPECTED_ANSI92_INTERMEDIATE_SQL)); assertThat(metaData.supportsANSI92FullSQL(), is(EXPECTED_ANSI92_FULL_SQL)); assertThat(metaData.supportsOuterJoins(), is(EXPECTED_SUPPORTS_OUTER_JOINS)); assertThat(metaData.supportsFullOuterJoins(), is(EXPECTED_SUPPORTS_FULL_OUTER_JOINS)); assertThat(metaData.supportsLimitedOuterJoins(), is(EXPECTED_SUPPORTS_LIMITED_JOINS)); assertThat(metaData.getSchemaTerm(), is(EXPECTED_SCHEMA_TERM)); assertThat(metaData.getProcedureTerm(), is(EXPECTED_PROCEDURE_TERM)); assertThat(metaData.getCatalogTerm(), is(EXPECTED_CATALOG_TERM)); assertThat(metaData.isCatalogAtStart(), is(EXPECTED_CATALOG_AT_START)); assertThat(metaData.supportsSchemasInProcedureCalls(), is(EXPECTED_SCHEMAS_IN_PROCEDURE_CALLS)); assertThat( metaData.supportsSchemasInIndexDefinitions(), is(EXPECTED_SCHEMAS_IN_INDEX_DEFINITIONS)); assertThat( metaData.supportsCatalogsInIndexDefinitions(), is(EXPECTED_CATALOGS_IN_INDEX_DEFINITIONS)); assertThat(metaData.supportsPositionedDelete(), is(EXPECTED_POSITIONED_DELETE)); assertThat(metaData.supportsPositionedUpdate(), is(EXPECTED_POSITIONED_UPDATE)); assertThat( metaData.supportsResultSetType(ResultSet.TYPE_FORWARD_ONLY), is(EXPECTED_TYPE_FORWARD_ONLY)); assertThat(metaData.supportsSelectForUpdate(), is(EXPECTED_SELECT_FOR_UPDATE_SUPPORTED)); assertThat(metaData.supportsStoredProcedures(), is(EXPECTED_STORED_PROCEDURES_SUPPORTED)); assertThat(metaData.supportsSubqueriesInComparisons(), is(EXPECTED_SUBQUERIES_IN_COMPARISON)); assertThat(metaData.supportsSubqueriesInExists(), is(EXPECTED_SUBQUERIES_IN_EXISTS)); assertThat(metaData.supportsSubqueriesInIns(), is(EXPECTED_SUBQUERIES_IN_INS)); assertThat(metaData.supportsSubqueriesInQuantifieds(), is(EXPECTED_SUBQUERIES_IN_QUANTIFIEDS)); assertThat( metaData.supportsCorrelatedSubqueries(), is(EXPECTED_CORRELATED_SUBQUERIES_SUPPORTED)); assertThat(metaData.supportsUnion(), is(EXPECTED_SUPPORTS_UNION)); assertThat(metaData.supportsUnionAll(), is(EXPECTED_SUPPORTS_UNION_ALL)); assertThat(metaData.getMaxBinaryLiteralLength(), is(EXPECTED_MAX_BINARY_LITERAL_LENGTH)); assertThat(metaData.getMaxCharLiteralLength(), is(EXPECTED_MAX_CHAR_LITERAL_LENGTH)); assertThat(metaData.getMaxColumnsInGroupBy(), is(EXPECTED_MAX_COLUMNS_IN_GROUP_BY)); assertThat(metaData.getMaxColumnsInIndex(), is(EXPECTED_MAX_COLUMNS_IN_INDEX)); assertThat(metaData.getMaxColumnsInOrderBy(), is(EXPECTED_MAX_COLUMNS_IN_ORDER_BY)); assertThat(metaData.getMaxColumnsInSelect(), is(EXPECTED_MAX_COLUMNS_IN_SELECT)); assertThat(metaData.getMaxConnections(), is(EXPECTED_MAX_CONNECTIONS)); assertThat(metaData.getMaxCursorNameLength(), is(EXPECTED_MAX_CURSOR_NAME_LENGTH)); assertThat(metaData.getMaxIndexLength(), is(EXPECTED_MAX_INDEX_LENGTH)); assertThat(metaData.getMaxSchemaNameLength(), is(EXPECTED_SCHEMA_NAME_LENGTH)); assertThat(metaData.getMaxProcedureNameLength(), is(EXPECTED_MAX_PROCEDURE_NAME_LENGTH)); assertThat(metaData.getMaxCatalogNameLength(), is(EXPECTED_MAX_CATALOG_NAME_LENGTH)); assertThat(metaData.getMaxRowSize(), is(EXPECTED_MAX_ROW_SIZE)); assertThat(metaData.doesMaxRowSizeIncludeBlobs(), is(EXPECTED_MAX_ROW_SIZE_INCLUDES_BLOBS)); assertThat(metaData.getMaxStatementLength(), is(EXPECTED_MAX_STATEMENT_LENGTH)); assertThat(metaData.getMaxStatements(), is(EXPECTED_MAX_STATEMENTS)); assertThat(metaData.getMaxTableNameLength(), is(EXPECTED_MAX_TABLE_NAME_LENGTH)); assertThat(metaData.getMaxTablesInSelect(), is(EXPECTED_MAX_TABLES_IN_SELECT)); assertThat(metaData.getMaxUserNameLength(), is(EXPECTED_MAX_USERNAME_LENGTH)); assertThat( metaData.getDefaultTransactionIsolation(), is(EXPECTED_DEFAULT_TRANSACTION_ISOLATION)); assertThat(metaData.supportsTransactions(), is(EXPECTED_TRANSACTIONS_SUPPORTED)); assertThat(metaData.supportsBatchUpdates(), is(EXPECTED_BATCH_UPDATES_SUPPORTED)); assertThat(metaData.supportsSavepoints(), is(EXPECTED_SAVEPOINTS_SUPPORTED)); assertThat(metaData.supportsNamedParameters(), is(EXPECTED_NAMED_PARAMETERS_SUPPORTED)); assertThat(metaData.locatorsUpdateCopy(), is(EXPECTED_LOCATORS_UPDATE_COPY)); assertThat( metaData.supportsResultSetType(ResultSet.TYPE_SCROLL_INSENSITIVE), is(EXPECTED_TYPE_SCROLL_INSENSITIVE)); assertThat( metaData.supportsResultSetType(ResultSet.TYPE_SCROLL_SENSITIVE), is(EXPECTED_TYPE_SCROLL_SENSITIVE)); assertThat( metaData.supportsSchemasInPrivilegeDefinitions(), is(EXPECTED_SCHEMAS_IN_PRIVILEGE_DEFINITIONS)); assertThat( metaData.supportsCatalogsInPrivilegeDefinitions(), is(EXPECTED_CATALOGS_IN_PRIVILEGE_DEFINITIONS)); assertThat( metaData.supportsTransactionIsolationLevel(Connection.TRANSACTION_NONE), is(EXPECTED_TRANSACTION_NONE)); assertThat( metaData.supportsTransactionIsolationLevel(Connection.TRANSACTION_READ_COMMITTED), is(EXPECTED_TRANSACTION_READ_COMMITTED)); assertThat( metaData.supportsTransactionIsolationLevel(Connection.TRANSACTION_READ_UNCOMMITTED), is(EXPECTED_TRANSACTION_READ_UNCOMMITTED)); assertThat( metaData.supportsTransactionIsolationLevel(Connection.TRANSACTION_REPEATABLE_READ), is(EXPECTED_TRANSACTION_REPEATABLE_READ)); assertThat( metaData.supportsTransactionIsolationLevel(Connection.TRANSACTION_SERIALIZABLE), is(EXPECTED_TRANSACTION_SERIALIZABLE)); assertThat( metaData.dataDefinitionCausesTransactionCommit(), is(EXPECTED_DATA_DEFINITION_CAUSES_TRANSACTION_COMMIT)); assertThat( metaData.dataDefinitionIgnoredInTransactions(), is(EXPECTED_DATA_DEFINITIONS_IN_TRANSACTIONS_IGNORED)); assertThat( metaData.supportsStoredFunctionsUsingCallSyntax(), is(EXPECTED_STORED_FUNCTIONS_USING_CALL_SYNTAX_SUPPORTED)); assertThat( metaData.supportsIntegrityEnhancementFacility(), is(EXPECTED_SUPPORTS_INTEGRITY_ENHANCEMENT_FACILITY)); assertThat( metaData.supportsDifferentTableCorrelationNames(), is(EXPECTED_SUPPORTS_DIFFERENT_TABLE_CORRELATION_NAMES)); ThrowableAssertionUtils.simpleAssertThrowableClass( SQLException.class, () -> metaData.supportsTransactionIsolationLevel(Connection.TRANSACTION_SERIALIZABLE + 1)); ThrowableAssertionUtils.simpleAssertThrowableClass( SQLException.class, () -> metaData.supportsResultSetType(ResultSet.HOLD_CURSORS_OVER_COMMIT)); } @Test public void testGetColumnsCanBeAccessedByNames() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getColumns(null, null, null, null)) { resultSetTestUtils.testData( resultSet, ImmutableList.of( "TABLE_CAT", "TABLE_SCHEM", "TABLE_NAME", "COLUMN_NAME", "DATA_TYPE", "TYPE_NAME", "COLUMN_SIZE", "BUFFER_LENGTH", "DECIMAL_DIGITS", "NUM_PREC_RADIX", "NULLABLE", "REMARKS", "COLUMN_DEF", "SQL_DATA_TYPE", "SQL_DATETIME_SUB", "CHAR_OCTET_LENGTH", "ORDINAL_POSITION", "IS_NULLABLE", "SCOPE_CATALOG", "SCOPE_SCHEMA", "SCOPE_TABLE", "SOURCE_DATA_TYPE", "IS_AUTOINCREMENT", "IS_GENERATEDCOLUMN"), EXPECTED_GET_COLUMNS_RESULTS); } } @Test public void testGetProcedures() throws SQLException { try (final ResultSet resultSet = connection.getMetaData().getProcedures(null, null, null)) { // Maps ordinal index to column name according to JDBC documentation final Map expectedGetProceduresSchema = new HashMap() { { put(1, "PROCEDURE_CAT"); put(2, "PROCEDURE_SCHEM"); put(3, "PROCEDURE_NAME"); put(4, "FUTURE_USE1"); put(5, "FUTURE_USE2"); put(6, "FUTURE_USE3"); put(7, "REMARKS"); put(8, "PROCEDURE_TYPE"); put(9, "SPECIFIC_NAME"); } }; testEmptyResultSet(resultSet, expectedGetProceduresSchema); } } @Test public void testGetProcedureColumns() throws SQLException { try (ResultSet resultSet = connection.getMetaData().getProcedureColumns(null, null, null, null)) { // Maps ordinal index to column name according to JDBC documentation final Map expectedGetProcedureColumnsSchema = new HashMap() { { put(1, "PROCEDURE_CAT"); put(2, "PROCEDURE_SCHEM"); put(3, "PROCEDURE_NAME"); put(4, "COLUMN_NAME"); put(5, "COLUMN_TYPE"); put(6, "DATA_TYPE"); put(7, "TYPE_NAME"); put(8, "PRECISION"); put(9, "LENGTH"); put(10, "SCALE"); put(11, "RADIX"); put(12, "NULLABLE"); put(13, "REMARKS"); put(14, "COLUMN_DEF"); put(15, "SQL_DATA_TYPE"); put(16, "SQL_DATETIME_SUB"); put(17, "CHAR_OCTET_LENGTH"); put(18, "ORDINAL_POSITION"); put(19, "IS_NULLABLE"); put(20, "SPECIFIC_NAME"); } }; testEmptyResultSet(resultSet, expectedGetProcedureColumnsSchema); } } @Test public void testGetColumnPrivileges() throws SQLException { try (ResultSet resultSet = connection.getMetaData().getColumnPrivileges(null, null, null, null)) { // Maps ordinal index to column name according to JDBC documentation final Map expectedGetColumnPrivilegesSchema = new HashMap() { { put(1, "TABLE_CAT"); put(2, "TABLE_SCHEM"); put(3, "TABLE_NAME"); put(4, "COLUMN_NAME"); put(5, "GRANTOR"); put(6, "GRANTEE"); put(7, "PRIVILEGE"); put(8, "IS_GRANTABLE"); } }; testEmptyResultSet(resultSet, expectedGetColumnPrivilegesSchema); } } @Test public void testGetTablePrivileges() throws SQLException { try (ResultSet resultSet = connection.getMetaData().getTablePrivileges(null, null, null)) { // Maps ordinal index to column name according to JDBC documentation final Map expectedGetTablePrivilegesSchema = new HashMap() { { put(1, "TABLE_CAT"); put(2, "TABLE_SCHEM"); put(3, "TABLE_NAME"); put(4, "GRANTOR"); put(5, "GRANTEE"); put(6, "PRIVILEGE"); put(7, "IS_GRANTABLE"); } }; testEmptyResultSet(resultSet, expectedGetTablePrivilegesSchema); } } @Test public void testGetBestRowIdentifier() throws SQLException { try (ResultSet resultSet = connection.getMetaData().getBestRowIdentifier(null, null, null, 0, true)) { // Maps ordinal index to column name according to JDBC documentation final Map expectedGetBestRowIdentifierSchema = new HashMap() { { put(1, "SCOPE"); put(2, "COLUMN_NAME"); put(3, "DATA_TYPE"); put(4, "TYPE_NAME"); put(5, "COLUMN_SIZE"); put(6, "BUFFER_LENGTH"); put(7, "DECIMAL_DIGITS"); put(8, "PSEUDO_COLUMN"); } }; testEmptyResultSet(resultSet, expectedGetBestRowIdentifierSchema); } } @Test public void testGetVersionColumns() throws SQLException { try (ResultSet resultSet = connection.getMetaData().getVersionColumns(null, null, null)) { // Maps ordinal index to column name according to JDBC documentation final Map expectedGetVersionColumnsSchema = new HashMap() { { put(1, "SCOPE"); put(2, "COLUMN_NAME"); put(3, "DATA_TYPE"); put(4, "TYPE_NAME"); put(5, "COLUMN_SIZE"); put(6, "BUFFER_LENGTH"); put(7, "DECIMAL_DIGITS"); put(8, "PSEUDO_COLUMN"); } }; testEmptyResultSet(resultSet, expectedGetVersionColumnsSchema); } } @Test public void testGetTypeInfo() throws SQLException { try (ResultSet resultSet = connection.getMetaData().getTypeInfo()) { // Maps ordinal index to column name according to JDBC documentation final Map expectedGetTypeInfoSchema = new HashMap() { { put(1, "TYPE_NAME"); put(2, "DATA_TYPE"); put(3, "PRECISION"); put(4, "LITERAL_PREFIX"); put(5, "LITERAL_SUFFIX"); put(6, "CREATE_PARAMS"); put(7, "NULLABLE"); put(8, "CASE_SENSITIVE"); put(9, "SEARCHABLE"); put(10, "UNSIGNED_ATTRIBUTE"); put(11, "FIXED_PREC_SCALE"); put(12, "AUTO_INCREMENT"); put(13, "LOCAL_TYPE_NAME"); put(14, "MINIMUM_SCALE"); put(15, "MAXIMUM_SCALE"); put(16, "SQL_DATA_TYPE"); put(17, "SQL_DATETIME_SUB"); put(18, "NUM_PREC_RADIX"); } }; testEmptyResultSet(resultSet, expectedGetTypeInfoSchema); } } @Test public void testGetIndexInfo() throws SQLException { try (ResultSet resultSet = connection.getMetaData().getIndexInfo(null, null, null, false, true)) { // Maps ordinal index to column name according to JDBC documentation final Map expectedGetIndexInfoSchema = new HashMap() { { put(1, "TABLE_CAT"); put(2, "TABLE_SCHEM"); put(3, "TABLE_NAME"); put(4, "NON_UNIQUE"); put(5, "INDEX_QUALIFIER"); put(6, "INDEX_NAME"); put(7, "TYPE"); put(8, "ORDINAL_POSITION"); put(9, "COLUMN_NAME"); put(10, "ASC_OR_DESC"); put(11, "CARDINALITY"); put(12, "PAGES"); put(13, "FILTER_CONDITION"); } }; testEmptyResultSet(resultSet, expectedGetIndexInfoSchema); } } @Test public void testGetUDTs() throws SQLException { try (ResultSet resultSet = connection.getMetaData().getUDTs(null, null, null, null)) { // Maps ordinal index to column name according to JDBC documentation final Map expectedGetUDTsSchema = new HashMap() { { put(1, "TYPE_CAT"); put(2, "TYPE_SCHEM"); put(3, "TYPE_NAME"); put(4, "CLASS_NAME"); put(5, "DATA_TYPE"); put(6, "REMARKS"); put(7, "BASE_TYPE"); } }; testEmptyResultSet(resultSet, expectedGetUDTsSchema); } } @Test public void testGetSuperTypes() throws SQLException { try (ResultSet resultSet = connection.getMetaData().getSuperTypes(null, null, null)) { // Maps ordinal index to column name according to JDBC documentation final Map expectedGetSuperTypesSchema = new HashMap() { { put(1, "TYPE_CAT"); put(2, "TYPE_SCHEM"); put(3, "TYPE_NAME"); put(4, "SUPERTYPE_CAT"); put(5, "SUPERTYPE_SCHEM"); put(6, "SUPERTYPE_NAME"); } }; testEmptyResultSet(resultSet, expectedGetSuperTypesSchema); } } @Test public void testGetSuperTables() throws SQLException { try (ResultSet resultSet = connection.getMetaData().getSuperTables(null, null, null)) { // Maps ordinal index to column name according to JDBC documentation final Map expectedGetSuperTablesSchema = new HashMap() { { put(1, "TABLE_CAT"); put(2, "TABLE_SCHEM"); put(3, "TABLE_NAME"); put(4, "SUPERTABLE_NAME"); } }; testEmptyResultSet(resultSet, expectedGetSuperTablesSchema); } } @Test public void testGetAttributes() throws SQLException { try (ResultSet resultSet = connection.getMetaData().getAttributes(null, null, null, null)) { // Maps ordinal index to column name according to JDBC documentation final Map expectedGetAttributesSchema = new HashMap() { { put(1, "TYPE_CAT"); put(2, "TYPE_SCHEM"); put(3, "TYPE_NAME"); put(4, "ATTR_NAME"); put(5, "DATA_TYPE"); put(6, "ATTR_TYPE_NAME"); put(7, "ATTR_SIZE"); put(8, "DECIMAL_DIGITS"); put(9, "NUM_PREC_RADIX"); put(10, "NULLABLE"); put(11, "REMARKS"); put(12, "ATTR_DEF"); put(13, "SQL_DATA_TYPE"); put(14, "SQL_DATETIME_SUB"); put(15, "CHAR_OCTET_LENGTH"); put(16, "ORDINAL_POSITION"); put(17, "IS_NULLABLE"); put(18, "SCOPE_CATALOG"); put(19, "SCOPE_SCHEMA"); put(20, "SCOPE_TABLE"); put(21, "SOURCE_DATA_TYPE"); } }; testEmptyResultSet(resultSet, expectedGetAttributesSchema); } } @Test public void testGetClientInfoProperties() throws SQLException { try (ResultSet resultSet = connection.getMetaData().getClientInfoProperties()) { // Maps ordinal index to column name according to JDBC documentation final Map expectedGetClientInfoPropertiesSchema = new HashMap() { { put(1, "NAME"); put(2, "MAX_LEN"); put(3, "DEFAULT_VALUE"); put(4, "DESCRIPTION"); } }; testEmptyResultSet(resultSet, expectedGetClientInfoPropertiesSchema); } } @Test public void testGetFunctions() throws SQLException { try (ResultSet resultSet = connection.getMetaData().getFunctions(null, null, null)) { // Maps ordinal index to column name according to JDBC documentation final Map expectedGetFunctionsSchema = new HashMap() { { put(1, "FUNCTION_CAT"); put(2, "FUNCTION_SCHEM"); put(3, "FUNCTION_NAME"); put(4, "REMARKS"); put(5, "FUNCTION_TYPE"); put(6, "SPECIFIC_NAME"); } }; testEmptyResultSet(resultSet, expectedGetFunctionsSchema); } } @Test public void testGetFunctionColumns() throws SQLException { try (ResultSet resultSet = connection.getMetaData().getFunctionColumns(null, null, null, null)) { // Maps ordinal index to column name according to JDBC documentation final Map expectedGetFunctionColumnsSchema = new HashMap() { { put(1, "FUNCTION_CAT"); put(2, "FUNCTION_SCHEM"); put(3, "FUNCTION_NAME"); put(4, "COLUMN_NAME"); put(5, "COLUMN_TYPE"); put(6, "DATA_TYPE"); put(7, "TYPE_NAME"); put(8, "PRECISION"); put(9, "LENGTH"); put(10, "SCALE"); put(11, "RADIX"); put(12, "NULLABLE"); put(13, "REMARKS"); put(14, "CHAR_OCTET_LENGTH"); put(15, "ORDINAL_POSITION"); put(16, "IS_NULLABLE"); put(17, "SPECIFIC_NAME"); } }; testEmptyResultSet(resultSet, expectedGetFunctionColumnsSchema); } } @Test public void testGetPseudoColumns() throws SQLException { try (ResultSet resultSet = connection.getMetaData().getPseudoColumns(null, null, null, null)) { // Maps ordinal index to column name according to JDBC documentation final Map expectedGetPseudoColumnsSchema = new HashMap() { { put(1, "TABLE_CAT"); put(2, "TABLE_SCHEM"); put(3, "TABLE_NAME"); put(4, "COLUMN_NAME"); put(5, "DATA_TYPE"); put(6, "COLUMN_SIZE"); put(7, "DECIMAL_DIGITS"); put(8, "NUM_PREC_RADIX"); put(9, "COLUMN_USAGE"); put(10, "REMARKS"); put(11, "CHAR_OCTET_LENGTH"); put(12, "IS_NULLABLE"); } }; testEmptyResultSet(resultSet, expectedGetPseudoColumnsSchema); } } private void testEmptyResultSet( final ResultSet resultSet, final Map expectedResultSetSchema) throws SQLException { assertFalse(resultSet.next()); final ResultSetMetaData resultSetMetaData = resultSet.getMetaData(); for (final Map.Entry entry : expectedResultSetSchema.entrySet()) { assertEquals(entry.getValue(), resultSetMetaData.getColumnLabel(entry.getKey())); } } @Test public void testGetColumnSize() { assertEquals( Integer.valueOf(ArrowDatabaseMetadata.COLUMN_SIZE_BYTE), ArrowDatabaseMetadata.getColumnSize(new ArrowType.Int(Byte.SIZE, true))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.COLUMN_SIZE_SHORT), ArrowDatabaseMetadata.getColumnSize(new ArrowType.Int(Short.SIZE, true))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.COLUMN_SIZE_INT), ArrowDatabaseMetadata.getColumnSize(new ArrowType.Int(Integer.SIZE, true))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.COLUMN_SIZE_LONG), ArrowDatabaseMetadata.getColumnSize(new ArrowType.Int(Long.SIZE, true))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.COLUMN_SIZE_VARCHAR_AND_BINARY), ArrowDatabaseMetadata.getColumnSize(new ArrowType.Utf8())); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.COLUMN_SIZE_VARCHAR_AND_BINARY), ArrowDatabaseMetadata.getColumnSize(new ArrowType.Binary())); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.COLUMN_SIZE_TIMESTAMP_SECONDS), ArrowDatabaseMetadata.getColumnSize(new ArrowType.Timestamp(TimeUnit.SECOND, null))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.COLUMN_SIZE_TIMESTAMP_MILLISECONDS), ArrowDatabaseMetadata.getColumnSize(new ArrowType.Timestamp(TimeUnit.MILLISECOND, null))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.COLUMN_SIZE_TIMESTAMP_MICROSECONDS), ArrowDatabaseMetadata.getColumnSize(new ArrowType.Timestamp(TimeUnit.MICROSECOND, null))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.COLUMN_SIZE_TIMESTAMP_NANOSECONDS), ArrowDatabaseMetadata.getColumnSize(new ArrowType.Timestamp(TimeUnit.NANOSECOND, null))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.COLUMN_SIZE_TIME), ArrowDatabaseMetadata.getColumnSize(new ArrowType.Time(TimeUnit.SECOND, Integer.SIZE))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.COLUMN_SIZE_TIME_MILLISECONDS), ArrowDatabaseMetadata.getColumnSize( new ArrowType.Time(TimeUnit.MILLISECOND, Integer.SIZE))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.COLUMN_SIZE_TIME_MICROSECONDS), ArrowDatabaseMetadata.getColumnSize( new ArrowType.Time(TimeUnit.MICROSECOND, Integer.SIZE))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.COLUMN_SIZE_TIME_NANOSECONDS), ArrowDatabaseMetadata.getColumnSize(new ArrowType.Time(TimeUnit.NANOSECOND, Integer.SIZE))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.COLUMN_SIZE_DATE), ArrowDatabaseMetadata.getColumnSize(new ArrowType.Date(DateUnit.DAY))); assertNull( ArrowDatabaseMetadata.getColumnSize( new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE))); } @Test public void testGetDecimalDigits() { assertEquals( Integer.valueOf(ArrowDatabaseMetadata.NO_DECIMAL_DIGITS), ArrowDatabaseMetadata.getDecimalDigits(new ArrowType.Int(Byte.SIZE, true))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.NO_DECIMAL_DIGITS), ArrowDatabaseMetadata.getDecimalDigits(new ArrowType.Timestamp(TimeUnit.SECOND, null))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.DECIMAL_DIGITS_TIME_MILLISECONDS), ArrowDatabaseMetadata.getDecimalDigits( new ArrowType.Timestamp(TimeUnit.MILLISECOND, null))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.DECIMAL_DIGITS_TIME_MICROSECONDS), ArrowDatabaseMetadata.getDecimalDigits( new ArrowType.Timestamp(TimeUnit.MICROSECOND, null))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.DECIMAL_DIGITS_TIME_NANOSECONDS), ArrowDatabaseMetadata.getDecimalDigits(new ArrowType.Timestamp(TimeUnit.NANOSECOND, null))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.NO_DECIMAL_DIGITS), ArrowDatabaseMetadata.getDecimalDigits(new ArrowType.Time(TimeUnit.SECOND, Integer.SIZE))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.DECIMAL_DIGITS_TIME_MILLISECONDS), ArrowDatabaseMetadata.getDecimalDigits( new ArrowType.Time(TimeUnit.MILLISECOND, Integer.SIZE))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.DECIMAL_DIGITS_TIME_MICROSECONDS), ArrowDatabaseMetadata.getDecimalDigits( new ArrowType.Time(TimeUnit.MICROSECOND, Integer.SIZE))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.DECIMAL_DIGITS_TIME_NANOSECONDS), ArrowDatabaseMetadata.getDecimalDigits( new ArrowType.Time(TimeUnit.NANOSECOND, Integer.SIZE))); assertEquals( Integer.valueOf(ArrowDatabaseMetadata.NO_DECIMAL_DIGITS), ArrowDatabaseMetadata.getDecimalDigits(new ArrowType.Date(DateUnit.DAY))); assertNull(ArrowDatabaseMetadata.getDecimalDigits(new ArrowType.Utf8())); } @Test public void testSqlToRegexLike() { assertEquals(".*", ArrowDatabaseMetadata.sqlToRegexLike("%")); assertEquals(".", ArrowDatabaseMetadata.sqlToRegexLike("_")); assertEquals("\\*", ArrowDatabaseMetadata.sqlToRegexLike("*")); assertEquals("T\\*E.S.*T", ArrowDatabaseMetadata.sqlToRegexLike("T*E_S%T")); } @Test 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)); } } } ArrowFlightJdbcArrayTest.java000066400000000000000000000153221515451615200377420ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import java.sql.ResultSet; import java.sql.SQLException; import java.sql.SQLFeatureNotSupportedException; import java.sql.Types; import java.util.Arrays; import java.util.HashMap; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.util.JsonStringArrayList; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.junit.jupiter.api.extension.RegisterExtension; import org.mockito.junit.jupiter.MockitoExtension; @ExtendWith(MockitoExtension.class) public class ArrowFlightJdbcArrayTest { @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); IntVector dataVector; @BeforeEach public void setup() { dataVector = rootAllocatorTestExtension.createIntVector(); } @AfterEach public void tearDown() { this.dataVector.close(); } @Test public void testShouldGetBaseTypeNameReturnCorrectTypeName() { ArrowFlightJdbcArray arrowFlightJdbcArray = new ArrowFlightJdbcArray(dataVector, 0, dataVector.getValueCount()); assertEquals("INTEGER", arrowFlightJdbcArray.getBaseTypeName()); } @Test public void testShouldGetBaseTypeReturnCorrectType() { ArrowFlightJdbcArray arrowFlightJdbcArray = new ArrowFlightJdbcArray(dataVector, 0, dataVector.getValueCount()); assertEquals(Types.INTEGER, arrowFlightJdbcArray.getBaseType()); } @Test public void testShouldGetArrayReturnValidArray() throws SQLException { ArrowFlightJdbcArray arrowFlightJdbcArray = new ArrowFlightJdbcArray(dataVector, 0, dataVector.getValueCount()); Object[] array = (Object[]) arrowFlightJdbcArray.getArray(); Object[] expected = new Object[dataVector.getValueCount()]; for (int i = 0; i < expected.length; i++) { expected[i] = dataVector.getObject(i); } assertArrayEquals(array, expected); } @Test public void testShouldGetArrayReturnValidArrayWithOffsets() throws SQLException { ArrowFlightJdbcArray arrowFlightJdbcArray = new ArrowFlightJdbcArray(dataVector, 0, dataVector.getValueCount()); Object[] array = (Object[]) arrowFlightJdbcArray.getArray(1, 5); Object[] expected = new Object[5]; for (int i = 0; i < expected.length; i++) { expected[i] = dataVector.getObject(i + 1); } assertArrayEquals(array, expected); } @Test public void testShouldGetArrayWithOffsetsThrowArrayIndexOutOfBoundsException() throws SQLException { ArrowFlightJdbcArray arrowFlightJdbcArray = new ArrowFlightJdbcArray(dataVector, 0, dataVector.getValueCount()); assertThrows( ArrayIndexOutOfBoundsException.class, () -> arrowFlightJdbcArray.getArray(0, dataVector.getValueCount() + 1)); } @Test public void testShouldGetArrayWithMapNotBeSupported() throws SQLException { ArrowFlightJdbcArray arrowFlightJdbcArray = new ArrowFlightJdbcArray(dataVector, 0, dataVector.getValueCount()); HashMap> map = new HashMap<>(); assertThrows(SQLFeatureNotSupportedException.class, () -> arrowFlightJdbcArray.getArray(map)); } @Test public void testShouldGetArrayWithOffsetsAndMapNotBeSupported() throws SQLException { ArrowFlightJdbcArray arrowFlightJdbcArray = new ArrowFlightJdbcArray(dataVector, 0, dataVector.getValueCount()); HashMap> map = new HashMap<>(); assertThrows( SQLFeatureNotSupportedException.class, () -> arrowFlightJdbcArray.getArray(0, 5, map)); } @Test public void testShouldGetResultSetReturnValidResultSet() throws SQLException { ArrowFlightJdbcArray arrowFlightJdbcArray = new ArrowFlightJdbcArray(dataVector, 0, dataVector.getValueCount()); try (ResultSet resultSet = arrowFlightJdbcArray.getResultSet()) { int count = 0; while (resultSet.next()) { assertEquals((Object) resultSet.getInt(2), dataVector.getObject(count)); count++; } } } @Test public void testShouldGetResultSetReturnValidResultSetWithOffsets() throws SQLException { ArrowFlightJdbcArray arrowFlightJdbcArray = new ArrowFlightJdbcArray(dataVector, 0, dataVector.getValueCount()); try (ResultSet resultSet = arrowFlightJdbcArray.getResultSet(3, 5)) { int count = 0; while (resultSet.next()) { assertEquals((Object) resultSet.getInt(2), dataVector.getObject(count + 3)); count++; } assertEquals(5, count); } } @Test public void testToString() throws SQLException { ArrowFlightJdbcArray arrowFlightJdbcArray = new ArrowFlightJdbcArray(dataVector, 0, dataVector.getValueCount()); JsonStringArrayList array = new JsonStringArrayList<>(); array.addAll(Arrays.asList((Object[]) arrowFlightJdbcArray.getArray())); assertEquals(array.toString(), arrowFlightJdbcArray.toString()); } @Test public void testShouldGetResultSetWithMapNotBeSupported() throws SQLException { ArrowFlightJdbcArray arrowFlightJdbcArray = new ArrowFlightJdbcArray(dataVector, 0, dataVector.getValueCount()); HashMap> map = new HashMap<>(); assertThrows( SQLFeatureNotSupportedException.class, () -> arrowFlightJdbcArray.getResultSet(map)); } @Test public void testShouldGetResultSetWithOffsetsAndMapNotBeSupported() throws SQLException { ArrowFlightJdbcArray arrowFlightJdbcArray = new ArrowFlightJdbcArray(dataVector, 0, dataVector.getValueCount()); HashMap> map = new HashMap<>(); assertThrows( SQLFeatureNotSupportedException.class, () -> arrowFlightJdbcArray.getResultSet(0, 5, map)); } } ArrowFlightJdbcConnectionCookieTest.java000066400000000000000000000040461515451615200421160ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertNull; import java.sql.Connection; import java.sql.SQLException; import java.sql.Statement; import org.apache.arrow.driver.jdbc.utils.CoreMockedSqlProducers; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; public class ArrowFlightJdbcConnectionCookieTest { @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION = FlightServerTestExtension.createStandardTestExtension( CoreMockedSqlProducers.getLegacyProducer()); @Test public void testCookies() throws SQLException { try (Connection connection = FLIGHT_SERVER_TEST_EXTENSION.getConnection(false); Statement statement = connection.createStatement()) { // Expect client didn't receive cookies before any operation 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.getInterceptorFactory().getCookie()); } } } ArrowFlightJdbcConnectionPoolDataSourceTest.java000066400000000000000000000126761515451615200436010ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; import java.sql.Connection; import javax.sql.PooledConnection; import org.apache.arrow.driver.jdbc.authentication.UserPasswordAuthentication; import org.apache.arrow.driver.jdbc.utils.ConnectionWrapper; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; 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; public class ArrowFlightJdbcConnectionPoolDataSourceTest { @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION; private static final MockFlightSqlProducer PRODUCER = new MockFlightSqlProducer(); static { UserPasswordAuthentication authentication = new UserPasswordAuthentication.Builder() .user("user1", "pass1") .user("user2", "pass2") .build(); FLIGHT_SERVER_TEST_EXTENSION = new FlightServerTestExtension.Builder() .authentication(authentication) .producer(PRODUCER) .build(); } private ArrowFlightJdbcConnectionPoolDataSource dataSource; @BeforeEach public void setUp() { dataSource = FLIGHT_SERVER_TEST_EXTENSION.createConnectionPoolDataSource(false); } @AfterEach public void tearDown() throws Exception { dataSource.close(); } @Test public void testShouldInnerConnectionIsClosedReturnCorrectly() throws Exception { PooledConnection pooledConnection = dataSource.getPooledConnection(); Connection connection = pooledConnection.getConnection(); assertFalse(connection.isClosed()); connection.close(); assertTrue(connection.isClosed()); } @Test public void testShouldInnerConnectionShouldIgnoreDoubleClose() throws Exception { PooledConnection pooledConnection = dataSource.getPooledConnection(); Connection connection = pooledConnection.getConnection(); assertFalse(connection.isClosed()); connection.close(); assertTrue(connection.isClosed()); } @Test public void testShouldInnerConnectionIsClosedReturnTrueIfPooledConnectionCloses() throws Exception { PooledConnection pooledConnection = dataSource.getPooledConnection(); Connection connection = pooledConnection.getConnection(); assertFalse(connection.isClosed()); pooledConnection.close(); assertTrue(connection.isClosed()); } @Test public void testShouldReuseConnectionsOnPool() throws Exception { PooledConnection pooledConnection = dataSource.getPooledConnection("user1", "pass1"); ConnectionWrapper connection = ((ConnectionWrapper) pooledConnection.getConnection()); assertFalse(connection.isClosed()); connection.close(); assertTrue(connection.isClosed()); assertFalse(connection.unwrap(ArrowFlightConnection.class).isClosed()); PooledConnection pooledConnection2 = dataSource.getPooledConnection("user1", "pass1"); ConnectionWrapper connection2 = ((ConnectionWrapper) pooledConnection2.getConnection()); assertFalse(connection2.isClosed()); connection2.close(); assertTrue(connection2.isClosed()); assertFalse(connection2.unwrap(ArrowFlightConnection.class).isClosed()); assertSame(pooledConnection, pooledConnection2); assertNotSame(connection, connection2); assertSame( connection.unwrap(ArrowFlightConnection.class), connection2.unwrap(ArrowFlightConnection.class)); } @Test public void testShouldNotMixConnectionsForDifferentUsers() throws Exception { PooledConnection pooledConnection = dataSource.getPooledConnection("user1", "pass1"); ConnectionWrapper connection = ((ConnectionWrapper) pooledConnection.getConnection()); assertFalse(connection.isClosed()); connection.close(); assertTrue(connection.isClosed()); assertFalse(connection.unwrap(ArrowFlightConnection.class).isClosed()); PooledConnection pooledConnection2 = dataSource.getPooledConnection("user2", "pass2"); ConnectionWrapper connection2 = ((ConnectionWrapper) pooledConnection2.getConnection()); assertFalse(connection2.isClosed()); connection2.close(); assertTrue(connection2.isClosed()); assertFalse(connection2.unwrap(ArrowFlightConnection.class).isClosed()); assertNotSame(pooledConnection, pooledConnection2); assertNotSame(connection, connection2); assertNotSame( connection.unwrap(ArrowFlightConnection.class), connection2.unwrap(ArrowFlightConnection.class)); } } ArrowFlightJdbcCursorTest.java000066400000000000000000000222031515451615200401350ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertTrue; import com.google.common.collect.ImmutableList; import java.sql.SQLException; import java.util.ArrayList; import java.util.List; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.IntervalDayVector; import org.apache.arrow.vector.TimeMilliVector; import org.apache.arrow.vector.TimeStampMilliVector; import org.apache.arrow.vector.UInt4Vector; 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.LargeListVector; 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.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.IntervalUnit; 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.FieldType; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.calcite.avatica.util.Cursor; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; /** Tests for {@link ArrowFlightJdbcCursor}. */ public class ArrowFlightJdbcCursorTest { ArrowFlightJdbcCursor cursor; BufferAllocator allocator; @AfterEach public void cleanUp() { allocator.close(); cursor.close(); } @Test public void testBinaryVectorNullTrue() throws SQLException { final VectorSchemaRoot root = getVectorSchemaRoot("Binary", new ArrowType.Binary(), null); ((VarBinaryVector) root.getVector("Binary")).setNull(0); testCursorWasNull(root); } @Test public void testDateVectorNullTrue() throws SQLException { final VectorSchemaRoot root = getVectorSchemaRoot("Date", new ArrowType.Date(DateUnit.DAY), null); ((DateDayVector) root.getVector("Date")).setNull(0); testCursorWasNull(root); } @Test public void testDurationVectorNullTrue() throws SQLException { final VectorSchemaRoot root = getVectorSchemaRoot("Duration", new ArrowType.Duration(TimeUnit.MILLISECOND), null); ((DurationVector) root.getVector("Duration")).setNull(0); testCursorWasNull(root); } @Test public void testDateInternalNullTrue() throws SQLException { final VectorSchemaRoot root = getVectorSchemaRoot("Interval", new ArrowType.Interval(IntervalUnit.DAY_TIME), null); ((IntervalDayVector) root.getVector("Interval")).setNull(0); testCursorWasNull(root); } @Test public void testTimeStampVectorNullTrue() throws SQLException { final VectorSchemaRoot root = getVectorSchemaRoot("TimeStamp", new ArrowType.Timestamp(TimeUnit.MILLISECOND, null), null); ((TimeStampMilliVector) root.getVector("TimeStamp")).setNull(0); testCursorWasNull(root); } @Test public void testTimeVectorNullTrue() throws SQLException { final VectorSchemaRoot root = getVectorSchemaRoot("Time", new ArrowType.Time(TimeUnit.MILLISECOND, 32), null); ((TimeMilliVector) root.getVector("Time")).setNull(0); testCursorWasNull(root); } @Test public void testFixedSizeListVectorNullTrue() throws SQLException { List fieldList = new ArrayList<>(); fieldList.add(new Field("Null", new FieldType(true, new ArrowType.Null(), null), null)); final VectorSchemaRoot root = getVectorSchemaRoot("FixedSizeList", new ArrowType.FixedSizeList(10), fieldList); ((FixedSizeListVector) root.getVector("FixedSizeList")).setNull(0); testCursorWasNull(root); } @Test public void testLargeListVectorNullTrue() throws SQLException { List fieldList = new ArrayList<>(); fieldList.add(new Field("Null", new FieldType(true, new ArrowType.Null(), null), null)); final VectorSchemaRoot root = getVectorSchemaRoot("LargeList", new ArrowType.LargeList(), fieldList); ((LargeListVector) root.getVector("LargeList")).setNull(0); testCursorWasNull(root); } @Test public void testListVectorNullTrue() throws SQLException { List fieldList = new ArrayList<>(); fieldList.add(new Field("Null", new FieldType(true, new ArrowType.Null(), null), null)); final VectorSchemaRoot root = getVectorSchemaRoot("List", new ArrowType.List(), fieldList); ((ListVector) root.getVector("List")).setNull(0); testCursorWasNull(root); } @Test public void testMapVectorNullTrue() throws SQLException { List structChildren = new ArrayList<>(); structChildren.add(new Field("Key", new FieldType(false, new ArrowType.Utf8(), null), null)); structChildren.add(new Field("Value", new FieldType(false, new ArrowType.Utf8(), null), null)); List fieldList = new ArrayList<>(); fieldList.add( new Field("Struct", new FieldType(false, new ArrowType.Struct(), null), structChildren)); final VectorSchemaRoot root = getVectorSchemaRoot("Map", new ArrowType.Map(false), fieldList); ((MapVector) root.getVector("Map")).setNull(0); testCursorWasNull(root); } @Test public void testStructVectorNullTrue() throws SQLException { final VectorSchemaRoot root = getVectorSchemaRoot("Struct", new ArrowType.Struct(), null); ((StructVector) root.getVector("Struct")).setNull(0); testCursorWasNull(root); } @Test public void testBaseIntVectorNullTrue() throws SQLException { final VectorSchemaRoot root = getVectorSchemaRoot("BaseInt", new ArrowType.Int(32, false), null); ((UInt4Vector) root.getVector("BaseInt")).setNull(0); testCursorWasNull(root); } @Test public void testBitVectorNullTrue() throws SQLException { final VectorSchemaRoot root = getVectorSchemaRoot("Bit", new ArrowType.Bool(), null); ((BitVector) root.getVector("Bit")).setNull(0); testCursorWasNull(root); } @Test public void testDecimalVectorNullTrue() throws SQLException { final VectorSchemaRoot root = getVectorSchemaRoot("Decimal", new ArrowType.Decimal(2, 2, 128), null); ((DecimalVector) root.getVector("Decimal")).setNull(0); testCursorWasNull(root); } @Test public void testFloat4VectorNullTrue() throws SQLException { final VectorSchemaRoot root = getVectorSchemaRoot( "Float4", new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE), null); ((Float4Vector) root.getVector("Float4")).setNull(0); testCursorWasNull(root); } @Test public void testFloat8VectorNullTrue() throws SQLException { final VectorSchemaRoot root = getVectorSchemaRoot( "Float8", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE), null); ((Float8Vector) root.getVector("Float8")).setNull(0); testCursorWasNull(root); } @Test public void testVarCharVectorNullTrue() throws SQLException { final VectorSchemaRoot root = getVectorSchemaRoot("VarChar", new ArrowType.Utf8(), null); ((VarCharVector) root.getVector("VarChar")).setNull(0); testCursorWasNull(root); } @Test public void testNullVectorNullTrue() throws SQLException { final VectorSchemaRoot root = getVectorSchemaRoot("Null", new ArrowType.Null(), null); testCursorWasNull(root); } private VectorSchemaRoot getVectorSchemaRoot( String name, ArrowType arrowType, List children) { final Schema schema = new Schema( ImmutableList.of(new Field(name, new FieldType(true, arrowType, null), children))); allocator = new RootAllocator(Long.MAX_VALUE); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); root.allocateNew(); return root; } private void testCursorWasNull(VectorSchemaRoot root) throws SQLException { root.setRowCount(1); cursor = new ArrowFlightJdbcCursor(root); cursor.next(); List accessorList = cursor.createAccessors(null, null, null); accessorList.get(0).getObject(); assertTrue(cursor.wasNull()); root.close(); } } ArrowFlightJdbcDriverTest.java000066400000000000000000000422411515451615200401170ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.sql.Connection; import java.sql.Driver; import java.sql.DriverManager; import java.sql.SQLException; import java.util.Collection; import java.util.Map; import java.util.Properties; import org.apache.arrow.driver.jdbc.authentication.UserPasswordAuthentication; import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; 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 ArrowFlightJdbcDriver}. */ public class ArrowFlightJdbcDriverTest { @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION; private static final MockFlightSqlProducer PRODUCER = new MockFlightSqlProducer(); static { UserPasswordAuthentication authentication = new UserPasswordAuthentication.Builder() .user("user1", "pass1") .user("user2", "pass2") .build(); FLIGHT_SERVER_TEST_EXTENSION = new FlightServerTestExtension.Builder() .authentication(authentication) .producer(PRODUCER) .build(); } private BufferAllocator allocator; private ArrowFlightJdbcConnectionPoolDataSource dataSource; @BeforeEach public void setUp() throws Exception { allocator = new RootAllocator(Long.MAX_VALUE); dataSource = FLIGHT_SERVER_TEST_EXTENSION.createConnectionPoolDataSource(); } @AfterEach public void tearDown() throws Exception { Collection childAllocators = allocator.getChildAllocators(); AutoCloseables.close(childAllocators.toArray(new AutoCloseable[0])); AutoCloseables.close(dataSource, allocator); } /** * Tests whether the {@link ArrowFlightJdbcDriver} is registered in the {@link DriverManager}. * * @throws SQLException If an error occurs. (This is not supposed to happen.) */ @Test public void testDriverIsRegisteredInDriverManager() throws Exception { assertTrue( DriverManager.getDriver("jdbc:arrow-flight://localhost:32010") instanceof ArrowFlightJdbcDriver); assertTrue( DriverManager.getDriver("jdbc:arrow-flight-sql://localhost:32010") instanceof ArrowFlightJdbcDriver); } /** * Tests whether the {@link ArrowFlightJdbcDriver} returns null when provided with an unsupported * URL prefix. */ @Test public void testShouldDeclineUrlWithUnsupportedPrefix() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); assertNull( driver.connect( "jdbc:mysql://localhost:32010", dataSource.getProperties("flight", "flight123"))); } /** * Tests whether the {@link ArrowFlightJdbcDriver} can establish a successful connection to the * Arrow Flight client. * * @throws Exception If the connection fails to be established. */ @Test public void testShouldConnectWhenProvidedWithValidUrl() throws Exception { // Get the Arrow Flight JDBC driver by providing a URL with a valid prefix. final Driver driver = new ArrowFlightJdbcDriver(); try (Connection connection = driver.connect( "jdbc:arrow-flight://" + dataSource.getConfig().getHost() + ":" + dataSource.getConfig().getPort() + "?" + "useEncryption=false", dataSource.getProperties( dataSource.getConfig().getUser(), dataSource.getConfig().getPassword()))) { assertTrue(connection.isValid(300)); } try (Connection connection = driver.connect( "jdbc:arrow-flight-sql://" + dataSource.getConfig().getHost() + ":" + dataSource.getConfig().getPort() + "?" + "useEncryption=false", dataSource.getProperties( dataSource.getConfig().getUser(), dataSource.getConfig().getPassword()))) { assertTrue(connection.isValid(300)); } } @Test public void testConnectWithInsensitiveCasePropertyKeys() throws Exception { // Get the Arrow Flight JDBC driver by providing a URL with insensitive case property keys. final Driver driver = new ArrowFlightJdbcDriver(); try (Connection connection = driver.connect( "jdbc:arrow-flight://" + dataSource.getConfig().getHost() + ":" + dataSource.getConfig().getPort() + "?" + "UseEncryptIon=false", dataSource.getProperties( dataSource.getConfig().getUser(), dataSource.getConfig().getPassword()))) { assertTrue(connection.isValid(300)); } try (Connection connection = driver.connect( "jdbc:arrow-flight-sql://" + dataSource.getConfig().getHost() + ":" + dataSource.getConfig().getPort() + "?" + "UseEncryptIon=false", dataSource.getProperties( dataSource.getConfig().getUser(), dataSource.getConfig().getPassword()))) { assertTrue(connection.isValid(300)); } } @Test public void testConnectWithInsensitiveCasePropertyKeys2() throws Exception { // Get the Arrow Flight JDBC driver by providing a property object with insensitive case keys. final Driver driver = new ArrowFlightJdbcDriver(); Properties properties = dataSource.getProperties( dataSource.getConfig().getUser(), dataSource.getConfig().getPassword()); properties.put("UseEncryptIon", "false"); try (Connection connection = driver.connect( "jdbc:arrow-flight://" + dataSource.getConfig().getHost() + ":" + dataSource.getConfig().getPort(), properties)) { assertTrue(connection.isValid(300)); } try (Connection connection = driver.connect( "jdbc:arrow-flight-sql://" + dataSource.getConfig().getHost() + ":" + dataSource.getConfig().getPort(), properties)) { assertTrue(connection.isValid(300)); } } /** * 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. * * @throws SQLException If an error occurs. */ @Test public void testShouldThrowExceptionWhenAttemptingToConnectToMalformedUrl() throws SQLException { final Driver driver = new ArrowFlightJdbcDriver(); final String malformedUri = "yes:??/chainsaw.i=T333"; assertThrows( SQLException.class, () -> driver.connect(malformedUri, dataSource.getProperties("flight", "flight123"))); } /** * Tests whether an exception is thrown upon attempting to connect to a malformed URI. * * @throws SQLException If an error occurs. */ @Test public void testShouldThrowExceptionWhenAttemptingToConnectToUrlNoPrefix() throws SQLException { final Driver driver = new ArrowFlightJdbcDriver(); final String malformedUri = "localhost:32010"; assertThrows( SQLException.class, () -> driver.connect( malformedUri, dataSource.getProperties( dataSource.getConfig().getUser(), dataSource.getConfig().getPassword()))); } /** Tests whether an exception is thrown upon attempting to connect to a malformed URI. */ @Test public void testShouldThrowExceptionWhenAttemptingToConnectToUrlNoPort() { final Driver driver = new ArrowFlightJdbcDriver(); SQLException e = assertThrows( SQLException.class, () -> { Properties properties = dataSource.getProperties( dataSource.getConfig().getUser(), dataSource.getConfig().getPassword()); Connection conn = driver.connect("jdbc:arrow-flight://localhost", properties); conn.close(); }); assertTrue(e.getMessage().contains("URL must have a port")); e = assertThrows( SQLException.class, () -> { Properties properties = dataSource.getProperties( dataSource.getConfig().getUser(), dataSource.getConfig().getPassword()); Connection conn = driver.connect("jdbc:arrow-flight-sql://localhost", properties); conn.close(); }); assertTrue(e.getMessage().contains("URL must have a port")); } /** Tests whether an exception is thrown upon attempting to connect to a malformed URI. */ @Test public void testShouldThrowExceptionWhenAttemptingToConnectToUrlNoHost() { final Driver driver = new ArrowFlightJdbcDriver(); SQLException e = assertThrows( SQLException.class, () -> { Properties properties = dataSource.getProperties( dataSource.getConfig().getUser(), dataSource.getConfig().getPassword()); Connection conn = driver.connect("jdbc:arrow-flight://32010:localhost", properties); conn.close(); }); assertTrue(e.getMessage().contains("URL must have a host")); e = assertThrows( SQLException.class, () -> { Properties properties = dataSource.getProperties( dataSource.getConfig().getUser(), dataSource.getConfig().getPassword()); Connection conn = driver.connect("jdbc:arrow-flight-sql://32010:localhost", properties); conn.close(); }); assertTrue(e.getMessage().contains("URL must have a host")); } /** * Tests whether {@link ArrowFlightJdbcDriver#getUrlsArgs} returns the correct URL parameters. * * @throws Exception If an error occurs. */ @Test public void testDriverUrlParsingMechanismShouldReturnTheDesiredArgsFromUrl() throws Exception { final ArrowFlightJdbcDriver driver = new ArrowFlightJdbcDriver(); final Map parsedArgs = driver .getUrlsArgs("jdbc:arrow-flight-sql://localhost:2222/?key1=value1&key2=value2&a=b") .orElseThrow(() -> new RuntimeException("URL was rejected")); // Check size == the amount of args provided (scheme not included) assertEquals(5, parsedArgs.size()); // Check host == the provided host assertEquals(parsedArgs.get(ArrowFlightConnectionProperty.HOST.camelName()), "localhost"); // Check port == the provided port assertEquals(parsedArgs.get(ArrowFlightConnectionProperty.PORT.camelName()), 2222); // Check all other non-default arguments assertEquals(parsedArgs.get("key1"), "value1"); assertEquals(parsedArgs.get("key2"), "value2"); assertEquals(parsedArgs.get("a"), "b"); } @Test public void testDriverUrlParsingMechanismShouldReturnTheDesiredArgsFromUrlWithSemicolon() throws Exception { final ArrowFlightJdbcDriver driver = new ArrowFlightJdbcDriver(); final Map parsedArgs = driver .getUrlsArgs("jdbc:arrow-flight-sql://localhost:2222/;key1=value1;key2=value2;a=b") .orElseThrow(() -> new RuntimeException("URL was rejected")); // Check size == the amount of args provided (scheme not included) assertEquals(5, parsedArgs.size()); // Check host == the provided host assertEquals(parsedArgs.get(ArrowFlightConnectionProperty.HOST.camelName()), "localhost"); // Check port == the provided port assertEquals(parsedArgs.get(ArrowFlightConnectionProperty.PORT.camelName()), 2222); // Check all other non-default arguments assertEquals(parsedArgs.get("key1"), "value1"); assertEquals(parsedArgs.get("key2"), "value2"); assertEquals(parsedArgs.get("a"), "b"); } @Test public void testDriverUrlParsingMechanismShouldReturnTheDesiredArgsFromUrlWithOneSemicolon() throws Exception { final ArrowFlightJdbcDriver driver = new ArrowFlightJdbcDriver(); final Map parsedArgs = driver .getUrlsArgs("jdbc:arrow-flight-sql://localhost:2222/;key1=value1") .orElseThrow(() -> new RuntimeException("URL was rejected")); // Check size == the amount of args provided (scheme not included) assertEquals(3, parsedArgs.size()); // Check host == the provided host assertEquals(parsedArgs.get(ArrowFlightConnectionProperty.HOST.camelName()), "localhost"); // Check port == the provided port assertEquals(parsedArgs.get(ArrowFlightConnectionProperty.PORT.camelName()), 2222); // Check all other non-default arguments assertEquals(parsedArgs.get("key1"), "value1"); } @Test public void testDriverUrlParsingMechanismShouldReturnEmptyOptionalForUnknownScheme() throws SQLException { final ArrowFlightJdbcDriver driver = new ArrowFlightJdbcDriver(); assertFalse(driver.getUrlsArgs("jdbc:malformed-url-flight://localhost:2222").isPresent()); } /** * Tests whether {@code ArrowFlightJdbcDriverTest#getUrlsArgs} returns the correct URL parameters * when the host is an IP Address. * * @throws Exception If an error occurs. */ @Test public void testDriverUrlParsingMechanismShouldWorkWithIPAddress() throws Exception { final ArrowFlightJdbcDriver driver = new ArrowFlightJdbcDriver(); final Map parsedArgs = driver .getUrlsArgs("jdbc:arrow-flight-sql://0.0.0.0:2222") .orElseThrow(() -> new RuntimeException("URL was rejected")); // Check size == the amount of args provided (scheme not included) assertEquals(2, parsedArgs.size()); // Check host == the provided host assertEquals(parsedArgs.get(ArrowFlightConnectionProperty.HOST.camelName()), "0.0.0.0"); // Check port == the provided port assertEquals(parsedArgs.get(ArrowFlightConnectionProperty.PORT.camelName()), 2222); } /** * Tests whether {@code ArrowFlightJdbcDriverTest#getUrlsArgs} escape especial characters and * returns the correct URL parameters when the especial character '&' is embedded in the query * parameters values. * * @throws Exception If an error occurs. */ @Test public void testDriverUrlParsingMechanismShouldWorkWithEmbeddedEspecialCharacter() throws Exception { final ArrowFlightJdbcDriver driver = new ArrowFlightJdbcDriver(); final Map parsedArgs = driver .getUrlsArgs( "jdbc:arrow-flight-sql://0.0.0.0:2222?test1=test1value&test2%26continue=test2value&test3=test3value") .orElseThrow(() -> new RuntimeException("URL was rejected")); // Check size == the amount of args provided (scheme not included) assertEquals(5, parsedArgs.size()); // Check host == the provided host assertEquals(parsedArgs.get(ArrowFlightConnectionProperty.HOST.camelName()), "0.0.0.0"); // Check port == the provided port assertEquals(parsedArgs.get(ArrowFlightConnectionProperty.PORT.camelName()), 2222); // Check all other non-default arguments assertEquals(parsedArgs.get("test1"), "test1value"); assertEquals(parsedArgs.get("test2&continue"), "test2value"); assertEquals(parsedArgs.get("test3"), "test3value"); } } ArrowFlightJdbcFactoryTest.java000066400000000000000000000070611515451615200402740ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ImmutableMap; import java.lang.reflect.Constructor; import java.sql.Connection; import java.util.Properties; import org.apache.arrow.driver.jdbc.authentication.UserPasswordAuthentication; import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.calcite.avatica.UnregisteredDriver; 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 ArrowFlightJdbcDriver}. */ public class ArrowFlightJdbcFactoryTest { @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION; private static final MockFlightSqlProducer PRODUCER = new MockFlightSqlProducer(); static { UserPasswordAuthentication authentication = new UserPasswordAuthentication.Builder() .user("user1", "pass1") .user("user2", "pass2") .build(); FLIGHT_SERVER_TEST_EXTENSION = new FlightServerTestExtension.Builder() .authentication(authentication) .producer(PRODUCER) .build(); } private BufferAllocator allocator; private ArrowFlightJdbcConnectionPoolDataSource dataSource; @BeforeEach public void setUp() throws Exception { allocator = new RootAllocator(Long.MAX_VALUE); dataSource = FLIGHT_SERVER_TEST_EXTENSION.createConnectionPoolDataSource(); } @AfterEach public void tearDown() throws Exception { AutoCloseables.close(dataSource, allocator); } @Test public void testShouldBeAbleToEstablishAConnectionSuccessfully() throws Exception { UnregisteredDriver driver = new ArrowFlightJdbcDriver(); Constructor constructor = ArrowFlightJdbcFactory.class.getConstructor(); constructor.setAccessible(true); ArrowFlightJdbcFactory factory = constructor.newInstance(); final Properties properties = new Properties(); properties.putAll( ImmutableMap.of( ArrowFlightConnectionProperty.HOST.camelName(), "localhost", ArrowFlightConnectionProperty.PORT.camelName(), 32010, ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false)); try (Connection connection = factory.newConnection( driver, constructor.newInstance(), "jdbc:arrow-flight-sql://localhost:32010", properties)) { assert connection.isValid(300); } } } ArrowFlightJdbcTimeTest.java000066400000000000000000000060351515451615200375630ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.hamcrest.CoreMatchers.endsWith; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; import java.time.LocalTime; import java.util.concurrent.TimeUnit; import org.junit.jupiter.api.Test; public class ArrowFlightJdbcTimeTest { final int hour = 5; final int minute = 6; final int second = 7; @Test public void testPrintingMillisNoLeadingZeroes() { // testing the regular case where the precision of the millisecond is 3 LocalTime dateTime = LocalTime.of(hour, minute, second, (int) TimeUnit.MILLISECONDS.toNanos(999)); ArrowFlightJdbcTime time = new ArrowFlightJdbcTime(dateTime); assertThat(time.toString(), endsWith(".999")); assertThat(time.getHours(), is(hour)); assertThat(time.getMinutes(), is(minute)); assertThat(time.getSeconds(), is(second)); } @Test public void testPrintingMillisOneLeadingZeroes() { // test case where one leading zero needs to be added LocalTime dateTime = LocalTime.of(hour, minute, second, (int) TimeUnit.MILLISECONDS.toNanos(99)); ArrowFlightJdbcTime time = new ArrowFlightJdbcTime(dateTime); assertThat(time.toString(), endsWith(".099")); assertThat(time.getHours(), is(hour)); assertThat(time.getMinutes(), is(minute)); assertThat(time.getSeconds(), is(second)); } @Test public void testPrintingMillisTwoLeadingZeroes() { // test case where two leading zeroes needs to be added LocalTime dateTime = LocalTime.of(hour, minute, second, (int) TimeUnit.MILLISECONDS.toNanos(1)); ArrowFlightJdbcTime time = new ArrowFlightJdbcTime(dateTime); assertThat(time.toString(), endsWith(".001")); assertThat(time.getHours(), is(hour)); assertThat(time.getMinutes(), is(minute)); assertThat(time.getSeconds(), is(second)); } @Test public void testEquality() { // tests #equals and #hashCode for coverage checks LocalTime dateTime = LocalTime.of(hour, minute, second, (int) TimeUnit.MILLISECONDS.toNanos(1)); ArrowFlightJdbcTime time1 = new ArrowFlightJdbcTime(dateTime); ArrowFlightJdbcTime time2 = new ArrowFlightJdbcTime(dateTime); assertThat(time1, is(time2)); assertThat(time1.hashCode(), is(time2.hashCode())); } } ArrowFlightPreparedStatementTest.java000066400000000000000000000241621515451615200415320ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.hamcrest.CoreMatchers.equalTo; 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; import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.SQLException; import java.util.Arrays; import java.util.Collections; import java.util.List; import org.apache.arrow.driver.jdbc.utils.CoreMockedSqlProducers; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; import org.apache.arrow.flight.sql.FlightSqlUtils; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorSchemaRoot; 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; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; public class ArrowFlightPreparedStatementTest { public static final MockFlightSqlProducer PRODUCER = CoreMockedSqlProducers.getLegacyProducer(); @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION = FlightServerTestExtension.createStandardTestExtension(PRODUCER); private static Connection connection; @BeforeAll public static void setup() throws SQLException { connection = FLIGHT_SERVER_TEST_EXTENSION.getConnection(false); } @AfterAll public static void tearDown() throws SQLException { connection.close(); } @BeforeEach public void before() { PRODUCER.clearActionTypeCounter(); } @Test public void testSimpleQueryNoParameterBinding() throws SQLException { final String query = CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD; try (final PreparedStatement preparedStatement = connection.prepareStatement(query); final ResultSet resultSet = preparedStatement.executeQuery()) { CoreMockedSqlProducers.assertLegacyRegularSqlResultSet(resultSet); } } @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 testQueryWithParameterBinding() throws SQLException { final String query = "Fake query with parameters"; final Schema schema = new Schema(Collections.singletonList(Field.nullable("", Types.MinorType.INT.getType()))); final Schema parameterSchema = new Schema( Arrays.asList( Field.nullable("", ArrowType.Utf8.INSTANCE), new Field( "", FieldType.nullable(ArrowType.List.INSTANCE), Collections.singletonList(Field.nullable("", Types.MinorType.INT.getType()))))); final List> expected = Collections.singletonList(Arrays.asList(new Text("foo"), new Integer[] {1, 2, null})); PRODUCER.addSelectQuery( query, schema, Collections.singletonList( listener -> { try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { ((IntVector) root.getVector(0)).setSafe(0, 10); root.setRowCount(1); listener.start(root); listener.putNext(); } catch (final Throwable throwable) { listener.error(throwable); } finally { listener.completed(); } })); PRODUCER.addExpectedParameters(query, parameterSchema, expected); try (final PreparedStatement preparedStatement = connection.prepareStatement(query)) { preparedStatement.setString(1, "foo"); preparedStatement.setArray( 2, connection.createArrayOf("INTEGER", new Integer[] {1, 2, null})); try (final ResultSet resultSet = preparedStatement.executeQuery()) { resultSet.next(); assert true; } } } @Test @Disabled("https://github.com/apache/arrow/issues/34741: flaky test") public void testPreparedStatementExecutionOnce() throws SQLException { final PreparedStatement statement = connection.prepareStatement(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD); // Expect that there is one entry in the map -- {prepared statement action type, invocation // count}. assertEquals(PRODUCER.getActionTypeCounter().size(), 1); // Expect that the prepared statement was executed exactly once. assertEquals( PRODUCER .getActionTypeCounter() .get(FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_STATEMENT.getType()), 1); statement.close(); } @Test public void testReturnColumnCount() throws SQLException { final String query = CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD; try (final PreparedStatement psmt = connection.prepareStatement(query)) { assertAll( "Column count is as expected", () -> assertThat("ID", equalTo(psmt.getMetaData().getColumnName(1))), () -> assertThat("Name", equalTo(psmt.getMetaData().getColumnName(2))), () -> assertThat("Age", equalTo(psmt.getMetaData().getColumnName(3))), () -> assertThat("Salary", equalTo(psmt.getMetaData().getColumnName(4))), () -> assertThat("Hire Date", equalTo(psmt.getMetaData().getColumnName(5))), () -> assertThat("Last Sale", equalTo(psmt.getMetaData().getColumnName(6))), () -> assertThat(6, equalTo(psmt.getMetaData().getColumnCount()))); } } @Test public void testUpdateQuery() throws SQLException { String query = "Fake update"; PRODUCER.addUpdateQuery(query, /*updatedRows*/ 42); try (final PreparedStatement stmt = connection.prepareStatement(query)) { int updated = stmt.executeUpdate(); assertEquals(42, updated); } } @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 testUpdateQueryWithParameters() throws SQLException { String query = "Fake update with parameters"; PRODUCER.addUpdateQuery(query, /*updatedRows*/ 42); PRODUCER.addExpectedParameters( query, new Schema(Collections.singletonList(Field.nullable("", ArrowType.Utf8.INSTANCE))), Collections.singletonList( Collections.singletonList(new Text("foo".getBytes(StandardCharsets.UTF_8))))); try (final PreparedStatement stmt = connection.prepareStatement(query)) { // TODO: make sure this is validated on the server too stmt.setString(1, "foo"); int updated = stmt.executeUpdate(); assertEquals(42, updated); } } @Test public void testUpdateQueryWithBatchedParameters() throws SQLException { String query = "Fake update with batched parameters"; Schema parameterSchema = new Schema( Arrays.asList( Field.nullable("", ArrowType.Utf8.INSTANCE), new Field( "", FieldType.nullable(ArrowType.List.INSTANCE), Collections.singletonList(Field.nullable("", Types.MinorType.INT.getType()))))); List> expected = Arrays.asList( Arrays.asList(new Text("foo"), new Integer[] {1, 2, null}), Arrays.asList(new Text("bar"), new Integer[] {0, -1, 100000})); PRODUCER.addUpdateQuery(query, /*updatedRows*/ 42); PRODUCER.addExpectedParameters(query, parameterSchema, expected); try (final PreparedStatement stmt = connection.prepareStatement(query)) { // TODO: make sure this is validated on the server too stmt.setString(1, "foo"); stmt.setArray(2, connection.createArrayOf("INTEGER", new Integer[] {1, 2, null})); stmt.addBatch(); stmt.setString(1, "bar"); stmt.setArray(2, connection.createArrayOf("INTEGER", new Integer[] {0, -1, 100000})); stmt.addBatch(); int[] updated = stmt.executeBatch(); assertEquals(42, updated[0]); } } } ArrowFlightStatementExecuteTest.java000066400000000000000000000160641515451615200413740ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.hamcrest.CoreMatchers.allOf; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.CoreMatchers.not; import static org.hamcrest.CoreMatchers.nullValue; import static org.hamcrest.MatcherAssert.assertThat; import java.sql.Connection; import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Statement; import java.util.Collections; import java.util.HashSet; import java.util.Set; import java.util.stream.Collectors; import java.util.stream.IntStream; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.UInt1Vector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.calcite.avatica.AvaticaUtils; 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; /** Tests for {@link ArrowFlightStatement#execute}. */ public class ArrowFlightStatementExecuteTest { private static final String SAMPLE_QUERY_CMD = "SELECT * FROM this_test"; private static final int SAMPLE_QUERY_ROWS = Byte.MAX_VALUE; private static final String VECTOR_NAME = "Unsigned Byte"; private static final Schema SAMPLE_QUERY_SCHEMA = new Schema(Collections.singletonList(Field.nullable(VECTOR_NAME, MinorType.UINT1.getType()))); private static final String SAMPLE_UPDATE_QUERY = "UPDATE this_table SET this_field = that_field FROM this_test WHERE this_condition"; private static final long SAMPLE_UPDATE_COUNT = 100L; 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 MockFlightSqlProducer PRODUCER = new MockFlightSqlProducer(); @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION = FlightServerTestExtension.createStandardTestExtension(PRODUCER); private Connection connection; private Statement statement; @BeforeAll public static void setUpBeforeClass() { PRODUCER.addSelectQuery( SAMPLE_QUERY_CMD, 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(); } })); PRODUCER.addUpdateQuery(SAMPLE_UPDATE_QUERY, SAMPLE_UPDATE_COUNT); PRODUCER.addUpdateQuery(SAMPLE_LARGE_UPDATE_QUERY, SAMPLE_LARGE_UPDATE_COUNT); } @BeforeEach public void setUp() throws SQLException { connection = FLIGHT_SERVER_TEST_EXTENSION.getConnection(false); statement = connection.createStatement(); } @AfterEach public void tearDown() throws Exception { AutoCloseables.close(statement, connection); } @AfterAll public static void tearDownAfterClass() throws Exception { AutoCloseables.close(PRODUCER); } @Test public void testExecuteShouldRunSelectQuery() throws SQLException { assertThat(statement.execute(SAMPLE_QUERY_CMD), is(true)); // Means this is a SELECT query. 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 testExecuteShouldRunUpdateQueryForSmallUpdate() throws SQLException { assertThat(statement.execute(SAMPLE_UPDATE_QUERY), is(false)); // Means this is an UPDATE query. assertThat( (long) statement.getUpdateCount(), is(allOf(equalTo(statement.getLargeUpdateCount()), equalTo(SAMPLE_UPDATE_COUNT)))); assertThat(statement.getResultSet(), is(nullValue())); } @Test public void testExecuteShouldRunUpdateQueryForLargeUpdate() throws SQLException { assertThat(statement.execute(SAMPLE_LARGE_UPDATE_QUERY), 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())); } @Test public void testUpdateCountShouldStartOnZero() throws SQLException { assertThat( (long) statement.getUpdateCount(), is(allOf(equalTo(statement.getLargeUpdateCount()), equalTo(0L)))); assertThat(statement.getResultSet(), is(nullValue())); } } ArrowFlightStatementExecuteUpdateTest.java000066400000000000000000000205701515451615200425340ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.lang.String.format; import static org.hamcrest.CoreMatchers.allOf; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.instanceOf; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.CoreMatchers.not; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.assertThrows; import java.sql.Connection; import java.sql.SQLException; import java.sql.SQLFeatureNotSupportedException; import java.sql.Statement; import java.util.Collections; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; 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.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.calcite.avatica.AvaticaUtils; 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; /** Tests for {@link ArrowFlightStatement#executeUpdate}. */ public class ArrowFlightStatementExecuteUpdateTest { private static final String UPDATE_SAMPLE_QUERY = "UPDATE sample_table SET sample_col = sample_val WHERE sample_condition"; private static final int UPDATE_SAMPLE_QUERY_AFFECTED_COLS = 10; private static final String LARGE_UPDATE_SAMPLE_QUERY = "UPDATE large_sample_table SET large_sample_col = large_sample_val WHERE large_sample_condition"; private static final long LARGE_UPDATE_SAMPLE_QUERY_AFFECTED_COLS = (long) Integer.MAX_VALUE + 1; private static final String REGULAR_QUERY_SAMPLE = "SELECT * FROM NOT_UPDATE_QUERY"; private static final Schema REGULAR_QUERY_SCHEMA = new Schema( Collections.singletonList(Field.nullable("placeholder", MinorType.VARCHAR.getType()))); private static final MockFlightSqlProducer PRODUCER = new MockFlightSqlProducer(); @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION = FlightServerTestExtension.createStandardTestExtension(PRODUCER); public Connection connection; public Statement statement; @BeforeAll public static void setUpBeforeClass() { PRODUCER.addUpdateQuery(UPDATE_SAMPLE_QUERY, UPDATE_SAMPLE_QUERY_AFFECTED_COLS); PRODUCER.addUpdateQuery(LARGE_UPDATE_SAMPLE_QUERY, LARGE_UPDATE_SAMPLE_QUERY_AFFECTED_COLS); PRODUCER.addSelectQuery( REGULAR_QUERY_SAMPLE, REGULAR_QUERY_SCHEMA, Collections.singletonList( listener -> { try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); final VectorSchemaRoot root = VectorSchemaRoot.create(REGULAR_QUERY_SCHEMA, allocator)) { listener.start(root); listener.putNext(); } catch (final Throwable throwable) { listener.error(throwable); } finally { listener.completed(); } })); } @BeforeEach public void setUp() throws SQLException { connection = FLIGHT_SERVER_TEST_EXTENSION.getConnection(false); statement = connection.createStatement(); } @AfterEach public void tearDown() throws Exception { AutoCloseables.close(statement, connection); } @AfterAll public static void tearDownAfterClass() throws Exception { AutoCloseables.close(PRODUCER); } @Test public void testExecuteUpdateShouldReturnNumColsAffectedForNumRowsFittingInt() throws SQLException { assertThat(statement.executeUpdate(UPDATE_SAMPLE_QUERY), is(UPDATE_SAMPLE_QUERY_AFFECTED_COLS)); } @Test public void testExecuteUpdateShouldReturnSaturatedNumColsAffectedIfDoesNotFitInInt() throws SQLException { final long result = statement.executeUpdate(LARGE_UPDATE_SAMPLE_QUERY); final long expectedRowCountRaw = LARGE_UPDATE_SAMPLE_QUERY_AFFECTED_COLS; assertThat( result, is( allOf( not(equalTo(expectedRowCountRaw)), equalTo( (long) AvaticaUtils.toSaturatedInt( expectedRowCountRaw))))); // Because of long-to-integer overflow. } @Test public void testExecuteLargeUpdateShouldReturnNumColsAffected() throws SQLException { assertThat( statement.executeLargeUpdate(LARGE_UPDATE_SAMPLE_QUERY), is(LARGE_UPDATE_SAMPLE_QUERY_AFFECTED_COLS)); } @Test // TODO Implement `Statement#executeUpdate(String, int)` public void testExecuteUpdateUnsupportedWithDriverFlag() throws SQLException { assertThrows( SQLFeatureNotSupportedException.class, () -> { assertThat( statement.executeUpdate(UPDATE_SAMPLE_QUERY, Statement.NO_GENERATED_KEYS), is(UPDATE_SAMPLE_QUERY_AFFECTED_COLS)); }); } @Test // TODO Implement `Statement#executeUpdate(String, int[])` public void testExecuteUpdateUnsupportedWithArrayOfInts() throws SQLException { assertThrows( SQLFeatureNotSupportedException.class, () -> { assertThat( statement.executeUpdate(UPDATE_SAMPLE_QUERY, new int[0]), is(UPDATE_SAMPLE_QUERY_AFFECTED_COLS)); }); } @Test // TODO Implement `Statement#executeUpdate(String, String[])` public void testExecuteUpdateUnsupportedWithArraysOfStrings() throws SQLException { assertThrows( SQLFeatureNotSupportedException.class, () -> { assertThat( statement.executeUpdate(UPDATE_SAMPLE_QUERY, new String[0]), is(UPDATE_SAMPLE_QUERY_AFFECTED_COLS)); }); } @Test public void testExecuteShouldExecuteUpdateQueryAutomatically() throws SQLException { assertThat( statement.execute(UPDATE_SAMPLE_QUERY), is(false)); // Meaning there was an update query. assertThat( statement.execute(REGULAR_QUERY_SAMPLE), is(true)); // Meaning there was a select query. } @Test public void testShouldFailToPrepareStatementForNullQuery() { int count = 0; try { assertThat(statement.execute(null), is(false)); } catch (final SQLException e) { count++; assertThat(e.getCause(), is(instanceOf(NullPointerException.class))); } assertThat(count, is(1)); } @Test public void testShouldFailToPrepareStatementForClosedStatement() throws SQLException { statement.close(); assertThat(statement.isClosed(), is(true)); int count = 0; try { statement.execute(UPDATE_SAMPLE_QUERY); } catch (final SQLException e) { count++; assertThat(e.getMessage(), is("Statement closed")); } assertThat(count, is(1)); } @Test public void testShouldFailToPrepareStatementForBadStatement() { final String badQuery = "BAD INVALID STATEMENT"; int count = 0; try { statement.execute(badQuery); } catch (final SQLException e) { count++; /* * The error message is up to whatever implementation of `FlightSqlProducer` * the driver is communicating with. However, for the purpose of this test, * we simply throw an `IllegalArgumentException` for queries not registered * in our `MockFlightSqlProducer`. */ assertThat( e.getMessage(), is(format("Error while executing SQL \"%s\": Query not found", badQuery))); } assertThat(count, is(1)); } } ConnectionMutualTlsTest.java000066400000000000000000000424261515451615200377070ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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 java.io.File; import java.net.URLEncoder; import java.sql.Connection; import java.sql.Driver; import java.sql.DriverManager; import java.sql.SQLException; import java.util.Properties; 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.FlightSqlTestCertificates; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; 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 encrypted connections. */ public class ConnectionMutualTlsTest { @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION; private static final String tlsRootCertsPath; private static final String clientMTlsCertPath; private static final String badClientMTlsCertPath; private static final String clientMTlsKeyPath; private static final String badClientMTlsKeyPath; private static final MockFlightSqlProducer PRODUCER = new MockFlightSqlProducer(); private static final String userTest = "user1"; private static final String passTest = "pass1"; static { final FlightSqlTestCertificates.CertKeyPair certKey = FlightSqlTestCertificates.exampleTlsCerts().get(0); tlsRootCertsPath = certKey.cert.getPath(); final File serverMTlsCACert = FlightSqlTestCertificates.exampleCACert(); final FlightSqlTestCertificates.CertKeyPair clientMTlsCertKey = FlightSqlTestCertificates.exampleTlsCerts().get(1); clientMTlsCertPath = clientMTlsCertKey.cert.getPath(); clientMTlsKeyPath = clientMTlsCertKey.key.getPath(); badClientMTlsCertPath = clientMTlsCertPath + ".bad"; badClientMTlsKeyPath = clientMTlsKeyPath + ".bad"; UserPasswordAuthentication authentication = new UserPasswordAuthentication.Builder().user(userTest, passTest).build(); FLIGHT_SERVER_TEST_EXTENSION = new FlightServerTestExtension.Builder() .authentication(authentication) .useEncryption(certKey.cert, certKey.key) .useMTlsClientVerification(serverMTlsCACert) .producer(PRODUCER) .build(); } private BufferAllocator allocator; @BeforeEach public void setUp() throws Exception { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void tearDown() throws Exception { allocator.getChildAllocators().forEach(BufferAllocator::close); AutoCloseables.close(allocator); } /** * Try to instantiate an encrypted FlightClient. * * @throws Exception on error. */ @Test public void testGetEncryptedClientAuthenticated() throws Exception { try (ArrowFlightSqlClientHandler client = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) .withPort(FLIGHT_SERVER_TEST_EXTENSION.getPort()) .withUsername(userTest) .withPassword(passTest) .withTlsRootCertificates(tlsRootCertsPath) .withClientCertificate(clientMTlsCertPath) .withClientKey(clientMTlsKeyPath) .withBufferAllocator(allocator) .withEncryption(true) .build()) { assertNotNull(client); } } /** * Try to instantiate an encrypted FlightClient providing a bad mTLS Cert Path. It's expected to * receive the SQLException. */ @Test public void testGetEncryptedClientWithBadMTlsCertPath() { assertThrows( SQLException.class, () -> { try (ArrowFlightSqlClientHandler handler = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) .withPort(FLIGHT_SERVER_TEST_EXTENSION.getPort()) .withUsername(userTest) .withPassword(passTest) .withTlsRootCertificates(tlsRootCertsPath) .withClientCertificate(badClientMTlsCertPath) .withClientKey(clientMTlsKeyPath) .withBufferAllocator(allocator) .withEncryption(true) .build()) { fail(); } }); } /** * Try to instantiate an encrypted FlightClient providing a bad mTLS Key Path. It's expected to * receive the SQLException. */ @Test public void testGetEncryptedClientWithBadMTlsKeyPath() { assertThrows( SQLException.class, () -> { try (ArrowFlightSqlClientHandler handler = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) .withPort(FLIGHT_SERVER_TEST_EXTENSION.getPort()) .withUsername(userTest) .withPassword(passTest) .withTlsRootCertificates(tlsRootCertsPath) .withClientCertificate(clientMTlsCertPath) .withClientKey(badClientMTlsKeyPath) .withBufferAllocator(allocator) .withEncryption(true) .build()) { fail(); } }); } /** * Try to instantiate an encrypted FlightClient without credentials. * * @throws Exception on error. */ @Test public void testGetNonAuthenticatedEncryptedClientNoAuth() throws Exception { try (ArrowFlightSqlClientHandler client = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) .withTlsRootCertificates(tlsRootCertsPath) .withClientCertificate(clientMTlsCertPath) .withClientKey(clientMTlsKeyPath) .withBufferAllocator(allocator) .withEncryption(true) .build()) { assertNotNull(client); } } /** * Check if an encrypted connection can be established successfully when the provided valid * credentials and a valid TLS Root Certs path. * * @throws Exception on error. */ @Test public void testGetEncryptedConnectionWithValidCredentialsAndTlsRootsPath() throws Exception { 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(ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), tlsRootCertsPath); properties.put( ArrowFlightConnectionProperty.CLIENT_CERTIFICATE.camelName(), clientMTlsCertPath); properties.put(ArrowFlightConnectionProperty.CLIENT_KEY.camelName(), clientMTlsKeyPath); final ArrowFlightJdbcDataSource dataSource = ArrowFlightJdbcDataSource.createNewDataSource(properties); try (final Connection connection = dataSource.getConnection()) { assertTrue(connection.isValid(300)); } } /** * Check if an encrypted connection can be established successfully when not providing * authentication. * * @throws Exception on error. */ @Test public void testGetNonAuthenticatedEncryptedConnection() 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.USE_ENCRYPTION.camelName(), true); properties.put(ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), tlsRootCertsPath); properties.put( ArrowFlightConnectionProperty.CLIENT_CERTIFICATE.camelName(), clientMTlsCertPath); properties.put(ArrowFlightConnectionProperty.CLIENT_KEY.camelName(), clientMTlsKeyPath); final ArrowFlightJdbcDataSource dataSource = ArrowFlightJdbcDataSource.createNewDataSource(properties); try (final Connection connection = dataSource.getConnection()) { assertTrue(connection.isValid(300)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using just a connection url. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueCorrectCastUrlWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); final String jdbcUrl = String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s" + "&useEncryption=true&%s=%s&%s=%s&%s=%s", FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest, ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), URLEncoder.encode(tlsRootCertsPath, "UTF-8"), ArrowFlightConnectionProperty.CLIENT_CERTIFICATE.camelName(), URLEncoder.encode(clientMTlsCertPath, "UTF-8"), ArrowFlightConnectionProperty.CLIENT_KEY.camelName(), URLEncoder.encode(clientMTlsKeyPath, "UTF-8")); try (Connection connection = DriverManager.getConnection(jdbcUrl)) { assertTrue(connection.isValid(0)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using a connection url and properties with String K-V pairs. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.setProperty(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.setProperty(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.setProperty( ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), tlsRootCertsPath); properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "true"); properties.setProperty( ArrowFlightConnectionProperty.CLIENT_CERTIFICATE.camelName(), clientMTlsCertPath); properties.setProperty(ArrowFlightConnectionProperty.CLIENT_KEY.camelName(), clientMTlsKeyPath); final String jdbcUrl = String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()); try (Connection connection = DriverManager.getConnection(jdbcUrl, properties)) { assertTrue(connection.isValid(0)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using a connection url and properties with Object K-V pairs. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueCorrectCastUrlAndPropertiesUsingPutWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), true); properties.put(ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), tlsRootCertsPath); properties.put( ArrowFlightConnectionProperty.CLIENT_CERTIFICATE.camelName(), clientMTlsCertPath); properties.put(ArrowFlightConnectionProperty.CLIENT_KEY.camelName(), clientMTlsKeyPath); final String jdbcUrl = String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()); try (Connection connection = DriverManager.getConnection(jdbcUrl, properties)) { assertTrue(connection.isValid(0)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using just a connection url and using 0 and 1 as ssl values. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueIntegerCorrectCastUrlWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); final String jdbcUrl = String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s" + "&useEncryption=1&useSystemTrustStore=0&%s=%s&%s=%s&%s=%s", FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest, ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), URLEncoder.encode(tlsRootCertsPath, "UTF-8"), ArrowFlightConnectionProperty.CLIENT_CERTIFICATE.camelName(), URLEncoder.encode(clientMTlsCertPath, "UTF-8"), ArrowFlightConnectionProperty.CLIENT_KEY.camelName(), URLEncoder.encode(clientMTlsKeyPath, "UTF-8")); try (Connection connection = DriverManager.getConnection(jdbcUrl)) { assertTrue(connection.isValid(0)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using a connection url and properties with String K-V pairs and using 0 and 1 as * ssl values. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueIntegerCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.setProperty(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.setProperty(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.setProperty( ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), tlsRootCertsPath); properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "1"); properties.setProperty( ArrowFlightConnectionProperty.CLIENT_CERTIFICATE.camelName(), clientMTlsCertPath); properties.setProperty(ArrowFlightConnectionProperty.CLIENT_KEY.camelName(), clientMTlsKeyPath); final String jdbcUrl = String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()); try (Connection connection = DriverManager.getConnection(jdbcUrl, properties)) { assertTrue(connection.isValid(0)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using a connection url and properties with Object K-V pairs and using 0 and 1 as * ssl values. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueIntegerCorrectCastUrlAndPropertiesUsingPutWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), 1); properties.put(ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), tlsRootCertsPath); properties.put( ArrowFlightConnectionProperty.CLIENT_CERTIFICATE.camelName(), clientMTlsCertPath); properties.put(ArrowFlightConnectionProperty.CLIENT_KEY.camelName(), clientMTlsKeyPath); final String jdbcUrl = String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()); try (Connection connection = DriverManager.getConnection(jdbcUrl, properties)) { assertTrue(connection.isValid(0)); } } } ConnectionTest.java000066400000000000000000000711071515451615200360320ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 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; import org.junit.jupiter.api.extension.RegisterExtension; /** Tests for {@link Connection}. */ public class ConnectionTest { @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION; private static final MockFlightSqlProducer PRODUCER = new MockFlightSqlProducer(); private static final String userTest = "user1"; private static final String passTest = "pass1"; static { UserPasswordAuthentication authentication = new UserPasswordAuthentication.Builder().user(userTest, passTest).build(); FLIGHT_SERVER_TEST_EXTENSION = new FlightServerTestExtension.Builder() .authentication(authentication) .producer(PRODUCER) .build(); } private BufferAllocator allocator; @BeforeEach public void setUp() throws Exception { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void tearDown() throws Exception { allocator.getChildAllocators().forEach(BufferAllocator::close); AutoCloseables.close(allocator); } /** * Checks if an unencrypted connection can be established successfully when the provided valid * credentials. * * @throws SQLException on error. */ @Test public void testUnencryptedConnectionShouldOpenSuccessfullyWhenProvidedValidCredentials() throws Exception { 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); try (Connection connection = DriverManager.getConnection( "jdbc:arrow-flight-sql://" + FLIGHT_SERVER_TEST_EXTENSION.getHost() + ":" + FLIGHT_SERVER_TEST_EXTENSION.getPort(), properties)) { assertTrue(connection.isValid(300)); } } /** * Checks if a token is provided it takes precedence over username/pass. In this case, the * connection should fail if a token is passed in. */ @Test public void testTokenOverridesUsernameAndPasswordAuth() { 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(ArrowFlightConnectionProperty.TOKEN.camelName(), "token"); properties.put("useEncryption", false); SQLException e = assertThrows( SQLException.class, () -> { try (Connection conn = DriverManager.getConnection( "jdbc:arrow-flight-sql://" + FLIGHT_SERVER_TEST_EXTENSION.getHost() + ":" + FLIGHT_SERVER_TEST_EXTENSION.getPort(), properties)) { fail(); } }); assertTrue(e.getMessage().contains("UNAUTHENTICATED")); } /** * Checks if the exception SQLException is thrown when trying to establish a connection without a * host. * * @throws SQLException on error. */ @Test public void testUnencryptedConnectionWithEmptyHost() throws Exception { final Properties properties = new Properties(); properties.put("user", userTest); properties.put("password", passTest); final String invalidUrl = "jdbc:arrow-flight-sql://"; assertThrows( SQLException.class, () -> { try (Connection conn = DriverManager.getConnection(invalidUrl, properties)) { fail("Expected SQLException."); } }); } /** * Try to instantiate a basic FlightClient. * * @throws URISyntaxException on error. */ @Test public void testGetBasicClientAuthenticatedShouldOpenConnection() throws Exception { try (ArrowFlightSqlClientHandler client = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) .withPort(FLIGHT_SERVER_TEST_EXTENSION.getPort()) .withEncryption(false) .withUsername(userTest) .withPassword(passTest) .withBufferAllocator(allocator) .build()) { assertNotNull(client); } } /** * Checks if the exception IllegalArgumentException is thrown when trying to establish an * unencrypted connection providing with an invalid port. * * @throws SQLException on error. */ @Test public void testUnencryptedConnectionProvidingInvalidPort() throws Exception { final Properties properties = new Properties(); properties.put(ArrowFlightConnectionProperty.HOST.camelName(), "localhost"); properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false); final String invalidUrl = "jdbc:arrow-flight-sql://" + FLIGHT_SERVER_TEST_EXTENSION.getHost() + ":" + 65537; assertThrows( SQLException.class, () -> { try (Connection conn = DriverManager.getConnection(invalidUrl, properties)) { fail("Expected SQLException"); } }); } /** * Try to instantiate a basic FlightClient. * * @throws URISyntaxException on error. */ @Test public void testGetBasicClientNoAuthShouldOpenConnection() throws Exception { try (ArrowFlightSqlClientHandler client = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) .withBufferAllocator(allocator) .withEncryption(false) .build()) { assertNotNull(client); } } /** * Checks if an unencrypted connection can be established successfully when not providing * credentials. * * @throws SQLException on error. */ @Test public void testUnencryptedConnectionShouldOpenSuccessfullyWithoutAuthentication() throws Exception { final Properties properties = new Properties(); properties.put(ArrowFlightConnectionProperty.HOST.camelName(), "localhost"); properties.put( ArrowFlightConnectionProperty.PORT.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getPort()); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false); try (Connection connection = DriverManager.getConnection("jdbc:arrow-flight-sql://localhost:32010", properties)) { assertTrue(connection.isValid(300)); } } /** * Check if an unencrypted connection throws an exception when provided with invalid credentials. * * @throws SQLException The exception expected to be thrown. */ @Test public void testUnencryptedConnectionShouldThrowExceptionWhenProvidedWithInvalidCredentials() throws Exception { final Properties properties = new Properties(); properties.put(ArrowFlightConnectionProperty.HOST.camelName(), "localhost"); properties.put(ArrowFlightConnectionProperty.USER.camelName(), "invalidUser"); properties.put( ArrowFlightConnectionProperty.PORT.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getPort()); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false); properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), "invalidPassword"); assertThrows( SQLException.class, () -> { try (Connection ignored = DriverManager.getConnection("jdbc:arrow-flight-sql://localhost:32010", properties)) { fail(); } }); } /** * Check if an non-encrypted connection can be established successfully when connecting through * the DriverManager using just a connection url. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyFalseCorrectCastUrlWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=false", FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest))) { assertTrue(connection.isValid(0)); } } /** * Check if an non-encrypted connection can be established successfully when connecting through * the DriverManager using a connection url and properties with String K-V pairs. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyFalseCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.setProperty(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.setProperty(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "false"); try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), properties)) { assertTrue(connection.isValid(0)); } } /** * Check if an non-encrypted connection can be established successfully when connecting through * the DriverManager using a connection url and properties with Object K-V pairs. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyFalseCorrectCastUrlAndPropertiesUsingPutWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false); try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), properties)) { assertTrue(connection.isValid(0)); } } /** * Check if an non-encrypted connection can be established successfully when connecting through * the DriverManager using just a connection url and using 0 and 1 as ssl values. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyFalseIntegerCorrectCastUrlWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=0", FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest))) { assertTrue(connection.isValid(0)); } } /** * Check if an non-encrypted connection can be established successfully when connecting through * the DriverManager using a connection url and properties with String K-V pairs and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyFalseIntegerCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.setProperty(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.setProperty(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "0"); try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), properties)) { assertTrue(connection.isValid(0)); } } /** * Check if an non-encrypted connection can be established successfully when connecting through * the DriverManager using a connection url and properties with Object K-V pairs and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyFalseIntegerCorrectCastUrlAndPropertiesUsingPutWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), 0); try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), properties)) { assertTrue(connection.isValid(0)); } } /** * Check if an non-encrypted connection can be established successfully when connecting through * the DriverManager using just a connection url. * * @throws Exception on error. */ @Test public void testThreadPoolSizeConnectionPropertyCorrectCastUrlWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&threadPoolSize=1&useEncryption=%s", FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest, false))) { assertTrue(connection.isValid(0)); } } /** * Check if an non-encrypted connection can be established successfully when connecting through * the DriverManager using a connection url and properties with String K-V pairs and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test public void testThreadPoolSizeConnectionPropertyCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.setProperty(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.setProperty(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.setProperty(ArrowFlightConnectionProperty.THREAD_POOL_SIZE.camelName(), "1"); properties.put("useEncryption", false); try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), properties)) { assertTrue(connection.isValid(0)); } } /** * Check if an non-encrypted connection can be established successfully when connecting through * the DriverManager using a connection url and properties with Object K-V pairs and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test public void testThreadPoolSizeConnectionPropertyCorrectCastUrlAndPropertiesUsingPutWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put(ArrowFlightConnectionProperty.THREAD_POOL_SIZE.camelName(), 1); properties.put("useEncryption", false); try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), properties)) { assertTrue(connection.isValid(0)); } } /** * Check if an non-encrypted connection can be established successfully when connecting through * the DriverManager using just a connection url. * * @throws Exception on error. */ @Test public void testPasswordConnectionPropertyIntegerCorrectCastUrlWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s&useEncryption=%s", FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest, false))) { assertTrue(connection.isValid(0)); } } /** * Check if an non-encrypted connection can be established successfully when connecting through * the DriverManager using a connection url and properties with String K-V pairs and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test public void testPasswordConnectionPropertyIntegerCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.setProperty(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.setProperty(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put("useEncryption", false); try (Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), properties)) { assertTrue(connection.isValid(0)); } } /** * Check if an non-encrypted connection can be established successfully when connecting through * the DriverManager using a connection url and properties with Object K-V pairs and using 0 and 1 * as ssl values. * * @throws Exception on error. */ @Test public void testPasswordConnectionPropertyIntegerCorrectCastUrlAndPropertiesUsingPutWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); 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( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), properties)) { 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()); } } } ConnectionTlsRootCertsTest.java000066400000000000000000000331061515451615200403570ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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 java.net.URLEncoder; import java.sql.Connection; import java.sql.Driver; import java.sql.DriverManager; import java.sql.SQLException; import java.util.Properties; 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.FlightSqlTestCertificates; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; 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 encrypted connections. */ public class ConnectionTlsRootCertsTest { @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION; private static final String tlsRootCertsPath; private static final String badTlsRootCertsPath; private static final MockFlightSqlProducer PRODUCER = new MockFlightSqlProducer(); private static final String userTest = "user1"; private static final String passTest = "pass1"; static { final FlightSqlTestCertificates.CertKeyPair certKey = FlightSqlTestCertificates.exampleTlsCerts().get(0); tlsRootCertsPath = certKey.cert.getPath(); badTlsRootCertsPath = certKey.cert.getPath() + ".bad"; UserPasswordAuthentication authentication = new UserPasswordAuthentication.Builder().user(userTest, passTest).build(); FLIGHT_SERVER_TEST_EXTENSION = new FlightServerTestExtension.Builder() .authentication(authentication) .useEncryption(certKey.cert, certKey.key) .producer(PRODUCER) .build(); } private BufferAllocator allocator; @BeforeEach public void setUp() throws Exception { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void tearDown() throws Exception { allocator.getChildAllocators().forEach(BufferAllocator::close); AutoCloseables.close(allocator); } /** * Try to instantiate an encrypted FlightClient. * * @throws Exception on error. */ @Test public void testGetEncryptedClientAuthenticated() throws Exception { try (ArrowFlightSqlClientHandler client = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) .withPort(FLIGHT_SERVER_TEST_EXTENSION.getPort()) .withUsername(userTest) .withPassword(passTest) .withTlsRootCertificates(tlsRootCertsPath) .withBufferAllocator(allocator) .withEncryption(true) .build()) { assertNotNull(client); } } /** * Try to instantiate an encrypted FlightClient providing a bad TLS Root Certs Path. It's expected * to receive the SQLException. */ @Test public void testGetEncryptedClientWithNoCertificateOnKeyStore() { assertThrows( SQLException.class, () -> { try (ArrowFlightSqlClientHandler handler = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) .withTlsRootCertificates(badTlsRootCertsPath) .withBufferAllocator(allocator) .withEncryption(true) .build()) { fail(); } }); } /** * Try to instantiate an encrypted FlightClient without credentials. * * @throws Exception on error. */ @Test public void testGetNonAuthenticatedEncryptedClientNoAuth() throws Exception { try (ArrowFlightSqlClientHandler client = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) .withTlsRootCertificates(tlsRootCertsPath) .withBufferAllocator(allocator) .withEncryption(true) .build()) { assertNotNull(client); } } /** * Check if an encrypted connection can be established successfully when the provided valid * credentials and a valid TLS Root Certs path. * * @throws Exception on error. */ @Test public void testGetEncryptedConnectionWithValidCredentialsAndTlsRootsPath() throws Exception { 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(ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), tlsRootCertsPath); final ArrowFlightJdbcDataSource dataSource = ArrowFlightJdbcDataSource.createNewDataSource(properties); try (final Connection connection = dataSource.getConnection()) { assertTrue(connection.isValid(300)); } } /** * Check if an encrypted connection can be established successfully when not providing * authentication. * * @throws Exception on error. */ @Test public void testGetNonAuthenticatedEncryptedConnection() 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.USE_ENCRYPTION.camelName(), true); properties.put(ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), tlsRootCertsPath); final ArrowFlightJdbcDataSource dataSource = ArrowFlightJdbcDataSource.createNewDataSource(properties); try (final Connection connection = dataSource.getConnection()) { assertTrue(connection.isValid(300)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using just a connection url. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueCorrectCastUrlWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s" + "&useEncryption=true&%s=%s", FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest, ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), URLEncoder.encode(tlsRootCertsPath, "UTF-8")))) { assertTrue(connection.isValid(0)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using a connection url and properties with String K-V pairs. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.setProperty(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.setProperty(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.setProperty( ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), tlsRootCertsPath); properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "true"); try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), properties)) { assertTrue(connection.isValid(0)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using a connection url and properties with Object K-V pairs. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueCorrectCastUrlAndPropertiesUsingPutWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), true); properties.put(ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), tlsRootCertsPath); try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), properties)) { assertTrue(connection.isValid(0)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using just a connection url and using 0 and 1 as ssl values. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueIntegerCorrectCastUrlWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s" + "&useEncryption=1&useSystemTrustStore=0&%s=%s", FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest, ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), URLEncoder.encode(tlsRootCertsPath, "UTF-8")))) { assertTrue(connection.isValid(0)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using a connection url and properties with String K-V pairs and using 0 and 1 as * ssl values. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueIntegerCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.setProperty(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.setProperty(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.setProperty( ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), tlsRootCertsPath); properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "1"); try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), properties)) { assertTrue(connection.isValid(0)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using a connection url and properties with Object K-V pairs and using 0 and 1 as * ssl values. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueIntegerCorrectCastUrlAndPropertiesUsingPutWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), 1); properties.put(ArrowFlightConnectionProperty.TLS_ROOT_CERTS.camelName(), tlsRootCertsPath); try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), properties)) { assertTrue(connection.isValid(0)); } } } ConnectionTlsTest.java000066400000000000000000000451471515451615200365220ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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 java.net.URLEncoder; import java.nio.file.Paths; import java.sql.Connection; import java.sql.Driver; import java.sql.DriverManager; import java.sql.SQLException; import java.util.Properties; 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.FlightSqlTestCertificates; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; 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 encrypted connections. */ public class ConnectionTlsTest { @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION; private static final MockFlightSqlProducer PRODUCER = new MockFlightSqlProducer(); private static final String userTest = "user1"; private static final String passTest = "pass1"; static { final FlightSqlTestCertificates.CertKeyPair certKey = FlightSqlTestCertificates.exampleTlsCerts().get(0); UserPasswordAuthentication authentication = new UserPasswordAuthentication.Builder().user(userTest, passTest).build(); FLIGHT_SERVER_TEST_EXTENSION = new FlightServerTestExtension.Builder() .authentication(authentication) .useEncryption(certKey.cert, certKey.key) .producer(PRODUCER) .build(); } private String trustStorePath; private String noCertificateKeyStorePath; private final String trustStorePass = "flight"; private BufferAllocator allocator; @BeforeEach public void setUp() throws Exception { trustStorePath = Paths.get(Preconditions.checkNotNull(getClass().getResource("/keys/keyStore.jks")).toURI()) .toString(); noCertificateKeyStorePath = Paths.get( Preconditions.checkNotNull(getClass().getResource("/keys/noCertificate.jks")) .toURI()) .toString(); allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void tearDown() throws Exception { allocator.getChildAllocators().forEach(BufferAllocator::close); AutoCloseables.close(allocator); } /** * Try to instantiate an encrypted FlightClient. * * @throws Exception on error. */ @Test public void testGetEncryptedClientAuthenticatedWithDisableCertVerification() throws Exception { try (ArrowFlightSqlClientHandler client = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) .withPort(FLIGHT_SERVER_TEST_EXTENSION.getPort()) .withUsername(userTest) .withPassword(passTest) .withDisableCertificateVerification(true) .withBufferAllocator(allocator) .withEncryption(true) .build()) { assertNotNull(client); } } /** * Try to instantiate an encrypted FlightClient. * * @throws Exception on error. */ @Test public void testGetEncryptedClientAuthenticated() throws Exception { try (ArrowFlightSqlClientHandler client = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) .withPort(FLIGHT_SERVER_TEST_EXTENSION.getPort()) .withSystemTrustStore(false) .withUsername(userTest) .withPassword(passTest) .withTrustStorePath(trustStorePath) .withTrustStorePassword(trustStorePass) .withBufferAllocator(allocator) .withEncryption(true) .build()) { assertNotNull(client); } } /** * Try to instantiate an encrypted FlightClient providing a keystore without certificate. It's * expected to receive the SQLException. * * @throws Exception on error. */ @Test public void testGetEncryptedClientWithNoCertificateOnKeyStore() throws Exception { final String noCertificateKeyStorePassword = "flight1"; assertThrows( SQLException.class, () -> { try (ArrowFlightSqlClientHandler ignored = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) .withTrustStorePath(noCertificateKeyStorePath) .withTrustStorePassword(noCertificateKeyStorePassword) .withSystemTrustStore(false) .withBufferAllocator(allocator) .withEncryption(true) .build()) { fail(); } }); } /** * Try to instantiate an encrypted FlightClient without credentials. * * @throws Exception on error. */ @Test public void testGetNonAuthenticatedEncryptedClientNoAuth() throws Exception { try (ArrowFlightSqlClientHandler client = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) .withSystemTrustStore(false) .withTrustStorePath(trustStorePath) .withTrustStorePassword(trustStorePass) .withBufferAllocator(allocator) .withEncryption(true) .build()) { assertNotNull(client); } } /** * Try to instantiate an encrypted FlightClient with an invalid password to the keystore file. * It's expected to receive the SQLException. * * @throws Exception on error. */ @Test public void testGetEncryptedClientWithKeyStoreBadPasswordAndNoAuth() throws Exception { String keyStoreBadPassword = "badPassword"; assertThrows( SQLException.class, () -> { try (ArrowFlightSqlClientHandler ignored = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) .withSystemTrustStore(false) .withTrustStorePath(trustStorePath) .withTrustStorePassword(keyStoreBadPassword) .withBufferAllocator(allocator) .withEncryption(true) .build()) { fail(); } }); } /** * Check if an encrypted connection can be established successfully when the provided valid * credentials and a valid Keystore. * * @throws Exception on error. */ @Test public void testGetEncryptedConnectionWithValidCredentialsAndKeyStore() throws Exception { 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(ArrowFlightConnectionProperty.TRUST_STORE.camelName(), trustStorePath); properties.put(ArrowFlightConnectionProperty.USE_SYSTEM_TRUST_STORE.camelName(), false); properties.put(ArrowFlightConnectionProperty.TRUST_STORE_PASSWORD.camelName(), trustStorePass); final ArrowFlightJdbcDataSource dataSource = ArrowFlightJdbcDataSource.createNewDataSource(properties); try (final Connection connection = dataSource.getConnection()) { assertTrue(connection.isValid(300)); } } /** * Check if the SQLException is thrown when trying to establish an encrypted connection providing * valid credentials but invalid password to the Keystore. * * @throws SQLException on error. */ @Test public void testGetAuthenticatedEncryptedConnectionWithKeyStoreBadPassword() 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(), true); properties.put(ArrowFlightConnectionProperty.TRUST_STORE.camelName(), trustStorePath); properties.put(ArrowFlightConnectionProperty.TRUST_STORE_PASSWORD.camelName(), "badpassword"); final ArrowFlightJdbcDataSource dataSource = ArrowFlightJdbcDataSource.createNewDataSource(properties); assertThrows( SQLException.class, () -> { try (final Connection ignored = dataSource.getConnection()) { fail(); } }); } /** * Check if an encrypted connection can be established successfully when not providing * authentication. * * @throws Exception on error. */ @Test public void testGetNonAuthenticatedEncryptedConnection() 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.USE_ENCRYPTION.camelName(), true); properties.put(ArrowFlightConnectionProperty.USE_SYSTEM_TRUST_STORE.camelName(), false); properties.put(ArrowFlightConnectionProperty.TRUST_STORE.camelName(), trustStorePath); properties.put(ArrowFlightConnectionProperty.TRUST_STORE_PASSWORD.camelName(), trustStorePass); final ArrowFlightJdbcDataSource dataSource = ArrowFlightJdbcDataSource.createNewDataSource(properties); try (final Connection connection = dataSource.getConnection()) { assertTrue(connection.isValid(300)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using just a connection url. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueCorrectCastUrlWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s" + "&useEncryption=true&useSystemTrustStore=false&%s=%s&%s=%s", FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest, ArrowFlightConnectionProperty.TRUST_STORE.camelName(), URLEncoder.encode(trustStorePath, "UTF-8"), ArrowFlightConnectionProperty.TRUST_STORE_PASSWORD.camelName(), URLEncoder.encode(trustStorePass, "UTF-8")))) { assertTrue(connection.isValid(0)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using a connection url and properties with String K-V pairs. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.setProperty(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.setProperty(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.setProperty(ArrowFlightConnectionProperty.TRUST_STORE.camelName(), trustStorePath); properties.setProperty( ArrowFlightConnectionProperty.TRUST_STORE_PASSWORD.camelName(), trustStorePass); properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "true"); properties.setProperty( ArrowFlightConnectionProperty.USE_SYSTEM_TRUST_STORE.camelName(), "false"); try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), properties)) { assertTrue(connection.isValid(0)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using a connection url and properties with Object K-V pairs. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueCorrectCastUrlAndPropertiesUsingPutWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), true); properties.put(ArrowFlightConnectionProperty.USE_SYSTEM_TRUST_STORE.camelName(), false); properties.put(ArrowFlightConnectionProperty.TRUST_STORE.camelName(), trustStorePath); properties.put(ArrowFlightConnectionProperty.TRUST_STORE_PASSWORD.camelName(), trustStorePass); try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), properties)) { assertTrue(connection.isValid(0)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using just a connection url and using 0 and 1 as ssl values. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueIntegerCorrectCastUrlWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s?user=%s&password=%s" + "&useEncryption=1&useSystemTrustStore=0&%s=%s&%s=%s", FLIGHT_SERVER_TEST_EXTENSION.getPort(), userTest, passTest, ArrowFlightConnectionProperty.TRUST_STORE.camelName(), URLEncoder.encode(trustStorePath, "UTF-8"), ArrowFlightConnectionProperty.TRUST_STORE_PASSWORD.camelName(), URLEncoder.encode(trustStorePass, "UTF-8")))) { assertTrue(connection.isValid(0)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using a connection url and properties with String K-V pairs and using 0 and 1 as * ssl values. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueIntegerCorrectCastUrlAndPropertiesUsingSetPropertyWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.setProperty(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.setProperty(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.setProperty(ArrowFlightConnectionProperty.TRUST_STORE.camelName(), trustStorePath); properties.setProperty( ArrowFlightConnectionProperty.TRUST_STORE_PASSWORD.camelName(), trustStorePass); properties.setProperty(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), "1"); properties.setProperty(ArrowFlightConnectionProperty.USE_SYSTEM_TRUST_STORE.camelName(), "0"); try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), properties)) { assertTrue(connection.isValid(0)); } } /** * Check if an encrypted connection can be established successfully when connecting through the * DriverManager using a connection url and properties with Object K-V pairs and using 0 and 1 as * ssl values. * * @throws Exception on error. */ @Test public void testTLSConnectionPropertyTrueIntegerCorrectCastUrlAndPropertiesUsingPutWithDriverManager() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); DriverManager.registerDriver(driver); Properties properties = new Properties(); properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest); properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), 1); properties.put(ArrowFlightConnectionProperty.USE_SYSTEM_TRUST_STORE.camelName(), 0); properties.put(ArrowFlightConnectionProperty.TRUST_STORE.camelName(), trustStorePath); properties.put(ArrowFlightConnectionProperty.TRUST_STORE_PASSWORD.camelName(), trustStorePass); try (final Connection connection = DriverManager.getConnection( String.format( "jdbc:arrow-flight-sql://localhost:%s", FLIGHT_SERVER_TEST_EXTENSION.getPort()), properties)) { assertTrue(connection.isValid(0)); } } } FlightServerTestExtension.java000066400000000000000000000301161515451615200402270ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.apache.arrow.driver.jdbc.utils.FlightSqlTestCertificates.CertKeyPair; import java.io.File; import java.io.IOException; import java.lang.reflect.Method; import java.sql.Connection; 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; import org.apache.arrow.driver.jdbc.authentication.UserPasswordAuthentication; import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl; 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; import org.apache.arrow.flight.RequestContext; import org.apache.arrow.flight.sql.FlightSqlProducer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; import org.junit.jupiter.api.extension.AfterAllCallback; import org.junit.jupiter.api.extension.BeforeAllCallback; import org.junit.jupiter.api.extension.ExtensionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * Utility class for unit tests that need to instantiate a {@link FlightServer} and interact with * it. */ public class FlightServerTestExtension implements BeforeAllCallback, AfterAllCallback, AutoCloseable { public static final String DEFAULT_USER = "flight-test-user"; public static final String DEFAULT_PASSWORD = "flight-test-password"; private static final Logger LOGGER = LoggerFactory.getLogger(FlightServerTestExtension.class); private final Properties properties; private final ArrowFlightConnectionConfigImpl config; private final BufferAllocator allocator; private final FlightSqlProducer producer; private final Authentication authentication; private final CertKeyPair certKeyPair; private final File mTlsCACert; private final InterceptorMiddleware.Factory interceptorFactory = new InterceptorMiddleware.Factory(); private FlightServerTestExtension( final Properties properties, final ArrowFlightConnectionConfigImpl config, final BufferAllocator allocator, final FlightSqlProducer producer, final Authentication authentication, final CertKeyPair certKeyPair, final File mTlsCACert) { this.properties = Preconditions.checkNotNull(properties); this.config = Preconditions.checkNotNull(config); this.allocator = Preconditions.checkNotNull(allocator); this.producer = Preconditions.checkNotNull(producer); this.authentication = authentication; this.certKeyPair = certKeyPair; this.mTlsCACert = mTlsCACert; } /** * Create a {@link FlightServerTestExtension} with standard values such as: user, password, * localhost. * * @param producer the producer used to create the FlightServerTestExtension. * @return the FlightServerTestExtension. */ public static FlightServerTestExtension createStandardTestExtension( final FlightSqlProducer producer) { UserPasswordAuthentication authentication = new UserPasswordAuthentication.Builder().user(DEFAULT_USER, DEFAULT_PASSWORD).build(); return new Builder().authentication(authentication).producer(producer).build(); } ArrowFlightJdbcDataSource createDataSource() { return ArrowFlightJdbcDataSource.createNewDataSource(properties); } public ArrowFlightJdbcConnectionPoolDataSource createConnectionPoolDataSource() { return ArrowFlightJdbcConnectionPoolDataSource.createNewDataSource(properties); } public ArrowFlightJdbcConnectionPoolDataSource createConnectionPoolDataSource( boolean useEncryption) { setUseEncryption(useEncryption); return ArrowFlightJdbcConnectionPoolDataSource.createNewDataSource(properties); } public Connection getConnection(boolean useEncryption, String token) throws SQLException { properties.put("token", token); return getConnection(useEncryption); } public Connection getConnection(boolean useEncryption) throws SQLException { setUseEncryption(useEncryption); 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 InterceptorMiddleware.Factory getInterceptorFactory() { return interceptorFactory; } @FunctionalInterface public interface CheckedFunction { R apply(T t) throws IOException; } private FlightServer initiateServer(Location location) throws IOException { FlightServer.Builder builder = FlightServer.builder(allocator, location, producer) .headerAuthenticator(authentication.authenticate()) .middleware(FlightServerMiddleware.Key.of("KEY"), interceptorFactory); if (certKeyPair != null) { builder.useTls(certKeyPair.cert, certKeyPair.key); } if (mTlsCACert != null) { builder.useMTlsClientVerification(mTlsCACert); } return builder.build(); } @Override public void beforeAll(ExtensionContext context) throws Exception { try { FlightServer flightServer = getStartServer(this::initiateServer, 3); properties.put("port", flightServer.getPort()); LOGGER.info("Started " + FlightServer.class.getName() + " as " + flightServer); context.getStore(ExtensionContext.Namespace.GLOBAL).put("flightServer", flightServer); } catch (Exception e) { LOGGER.error("Failed to start FlightServer", e); throw e; } } @Override public void afterAll(ExtensionContext context) throws Exception { FlightServer flightServer = context.getStore(ExtensionContext.Namespace.GLOBAL).get("flightServer", FlightServer.class); if (flightServer != null) { flightServer.close(); } close(); } private FlightServer getStartServer( CheckedFunction newServerFromLocation, int retries) throws IOException { final Deque exceptions = new ArrayDeque<>(); for (; retries > 0; retries--) { final FlightServer server = newServerFromLocation.apply(Location.forGrpcInsecure("localhost", 0)); try { Method start = server.getClass().getMethod("start"); start.setAccessible(true); start.invoke(server); return server; } catch (ReflectiveOperationException e) { exceptions.add(e); } } exceptions.forEach(e -> LOGGER.error("Failed to start FlightServer", e)); throw new IOException(exceptions.pop().getCause()); } /** * Sets a port to be used. * * @return the port value. */ public int getPort() { return config.getPort(); } /** * Sets a host to be used. * * @return the host value. */ public String getHost() { return config.getHost(); } @Override public void close() throws Exception { allocator.getChildAllocators().forEach(BufferAllocator::close); AutoCloseables.close(allocator); } /** Builder for {@link FlightServerTestExtension}. */ public static final class Builder { private final Properties properties; private FlightSqlProducer producer; private Authentication authentication; private CertKeyPair certKeyPair; private File mTlsCACert; public Builder() { this.properties = new Properties(); this.properties.put("host", "localhost"); } /** * Sets the producer that will be used in the server rule. * * @param producer the flight sql producer. * @return the Builder. */ public Builder producer(final FlightSqlProducer producer) { this.producer = producer; return this; } /** * Sets the type of the authentication that will be used in the server rules. There are two * types of authentication: {@link UserPasswordAuthentication} and {@link TokenAuthentication}. * * @param authentication the type of authentication. * @return the Builder. */ public Builder authentication(final Authentication authentication) { this.authentication = authentication; return this; } /** * Enable TLS on the server. * * @param certChain The certificate chain to use. * @param key The private key to use. * @return the Builder. */ public Builder useEncryption(final File certChain, final File key) { certKeyPair = new CertKeyPair(certChain, key); return this; } /** * Enable Client Verification via mTLS on the server. * * @param mTlsCACert The CA certificate to use for client verification. * @return the Builder. */ public Builder useMTlsClientVerification(final File mTlsCACert) { this.mTlsCACert = mTlsCACert; return this; } /** * Builds the {@link FlightServerTestExtension} using the provided values. * * @return a {@link FlightServerTestExtension}. */ public FlightServerTestExtension build() { authentication.populateProperties(properties); return new FlightServerTestExtension( properties, new ArrowFlightConnectionConfigImpl(properties), new RootAllocator(Long.MAX_VALUE), producer, authentication, certKeyPair, mTlsCACert); } } /** * A middleware to handle with the cookies in the server. It is used to test if cookies are being * sent properly. */ static class InterceptorMiddleware implements FlightServerMiddleware { private final Factory factory; public InterceptorMiddleware(Factory factory) { this.factory = factory; } @Override public void onBeforeSendingHeaders(CallHeaders callHeaders) { if (!factory.receivedCookieHeader) { callHeaders.insert("Set-Cookie", "k=v"); } } @Override public void onCallCompleted(CallStatus callStatus) {} @Override public void onCallErrored(Throwable throwable) {} /** A factory for the MiddlewareCookie. */ static class Factory implements FlightServerMiddleware.Factory { private final Map receivedCallHeaders = new HashMap<>(); private boolean receivedCookieHeader = false; private String cookie; @Override public InterceptorMiddleware onCallStarted( CallInfo callInfo, CallHeaders callHeaders, RequestContext requestContext) { cookie = callHeaders.get("Cookie"); receivedCookieHeader = null != cookie; 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); } } } } OAuthIntegrationTest.java000066400000000000000000000467121515451615200371630ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } ResultSetMetadataTest.java000066400000000000000000000170111515451615200373200ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.notNullValue; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.assertThrows; import java.sql.Connection; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Statement; import java.sql.Types; import org.apache.arrow.driver.jdbc.utils.CoreMockedSqlProducers; import org.hamcrest.CoreMatchers; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; public class ResultSetMetadataTest { private static ResultSetMetaData metadata; private static Connection connection; @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION = FlightServerTestExtension.createStandardTestExtension( CoreMockedSqlProducers.getLegacyProducer()); @BeforeAll public static void setup() throws SQLException { connection = FLIGHT_SERVER_TEST_EXTENSION.getConnection(false); try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.LEGACY_METADATA_SQL_CMD)) { metadata = resultSet.getMetaData(); } } @AfterAll public static void teardown() throws SQLException { connection.close(); } /** Test if {@link ResultSetMetaData} object is not null. */ @Test public void testShouldGetResultSetMetadata() { assertThat(metadata, CoreMatchers.is(notNullValue())); } /** * Test if {@link ResultSetMetaData#getColumnCount()} returns the correct values. * * @throws SQLException in case of error. */ @Test public void testShouldGetColumnCount() throws SQLException { final int columnCount = metadata.getColumnCount(); assert columnCount == 3; } /** * Test if {@link ResultSetMetaData#getColumnTypeName(int)} returns the correct type name for each * column. * * @throws SQLException in case of error. */ @Test public void testShouldGetColumnTypesName() throws SQLException { final String firstColumn = metadata.getColumnTypeName(1); final String secondColumn = metadata.getColumnTypeName(2); final String thirdColumn = metadata.getColumnTypeName(3); assertThat(firstColumn, equalTo("BIGINT")); assertThat(secondColumn, equalTo("VARCHAR")); assertThat(thirdColumn, equalTo("FLOAT")); } /** * Test if {@link ResultSetMetaData#getColumnTypeName(int)} passing an column index that does not * exist. */ @Test public void testShouldGetColumnTypesNameFromOutOfBoundIndex() { assertThrows( IndexOutOfBoundsException.class, () -> { metadata.getColumnTypeName(4); }); } /** * Test if {@link ResultSetMetaData#getColumnName(int)} returns the correct name for each column. * * @throws SQLException in case of error. */ @Test public void testShouldGetColumnNames() throws SQLException { final String firstColumn = metadata.getColumnName(1); final String secondColumn = metadata.getColumnName(2); final String thirdColumn = metadata.getColumnName(3); assertThat(firstColumn, equalTo("integer0")); assertThat(secondColumn, equalTo("string1")); assertThat(thirdColumn, equalTo("float2")); } /** * Test {@link ResultSetMetaData#getColumnTypeName(int)} passing an column index that does not * exist. */ @Test public void testShouldGetColumnNameFromOutOfBoundIndex() { assertThrows(IndexOutOfBoundsException.class, () -> metadata.getColumnName(4)); } /** * Test if {@link ResultSetMetaData#getColumnType(int)}returns the correct values. * * @throws SQLException in case of error. */ @Test public void testShouldGetColumnType() throws SQLException { final int firstColumn = metadata.getColumnType(1); final int secondColumn = metadata.getColumnType(2); final int thirdColumn = metadata.getColumnType(3); assertThat(firstColumn, equalTo(Types.BIGINT)); assertThat(secondColumn, equalTo(Types.VARCHAR)); assertThat(thirdColumn, equalTo(Types.FLOAT)); } @Test public void testShouldGetPrecision() throws SQLException { assertThat(metadata.getPrecision(1), equalTo(10)); assertThat(metadata.getPrecision(2), equalTo(65535)); assertThat(metadata.getPrecision(3), equalTo(15)); } @Test public void testShouldGetScale() throws SQLException { assertThat(metadata.getScale(1), equalTo(0)); assertThat(metadata.getScale(2), equalTo(0)); assertThat(metadata.getScale(3), equalTo(20)); } @Test public void testShouldGetCatalogName() throws SQLException { assertThat(metadata.getCatalogName(1), equalTo("CATALOG_NAME_1")); assertThat(metadata.getCatalogName(2), equalTo("CATALOG_NAME_2")); assertThat(metadata.getCatalogName(3), equalTo("CATALOG_NAME_3")); } @Test public void testShouldGetSchemaName() throws SQLException { assertThat(metadata.getSchemaName(1), equalTo("SCHEMA_NAME_1")); assertThat(metadata.getSchemaName(2), equalTo("SCHEMA_NAME_2")); assertThat(metadata.getSchemaName(3), equalTo("SCHEMA_NAME_3")); } @Test public void testShouldGetTableName() throws SQLException { assertThat(metadata.getTableName(1), equalTo("TABLE_NAME_1")); assertThat(metadata.getTableName(2), equalTo("TABLE_NAME_2")); assertThat(metadata.getTableName(3), equalTo("TABLE_NAME_3")); } @Test public void testShouldIsAutoIncrement() throws SQLException { assertThat(metadata.isAutoIncrement(1), equalTo(true)); assertThat(metadata.isAutoIncrement(2), equalTo(false)); assertThat(metadata.isAutoIncrement(3), equalTo(false)); } @Test public void testShouldIsCaseSensitive() throws SQLException { assertThat(metadata.isCaseSensitive(1), equalTo(false)); assertThat(metadata.isCaseSensitive(2), equalTo(true)); assertThat(metadata.isCaseSensitive(3), equalTo(false)); } @Test public void testShouldIsReadonly() throws SQLException { assertThat(metadata.isReadOnly(1), equalTo(true)); assertThat(metadata.isReadOnly(2), equalTo(false)); assertThat(metadata.isReadOnly(3), equalTo(false)); } @Test public void testShouldIsSearchable() throws SQLException { assertThat(metadata.isSearchable(1), equalTo(true)); assertThat(metadata.isSearchable(2), equalTo(true)); assertThat(metadata.isSearchable(3), equalTo(true)); } /** * Test if {@link ResultSetMetaData#getColumnTypeName(int)} passing an column index that does not * exist. */ @Test public void testShouldGetColumnTypesFromOutOfBoundIndex() { assertThrows(IndexOutOfBoundsException.class, () -> metadata.getColumnType(4)); } } ResultSetTest.java000066400000000000000000001147371515451615200356740ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.lang.String.format; import static java.util.Collections.synchronizedSet; import static org.apache.arrow.flight.Location.forGrpcInsecure; 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.*; 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; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.HashSet; 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; import org.apache.arrow.driver.jdbc.utils.PartitionedFlightSqlProducer; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightStatusCode; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.Ticket; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorSchemaRoot; 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.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; import org.junit.jupiter.api.extension.RegisterExtension; public class ResultSetTest { private static final Random RANDOM = new Random(10); @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION = FlightServerTestExtension.createStandardTestExtension( CoreMockedSqlProducers.getLegacyProducer()); private static Connection connection; @BeforeAll public static void setup() throws SQLException { connection = FLIGHT_SERVER_TEST_EXTENSION.getConnection(false); } @AfterAll public static void tearDown() throws SQLException { connection.close(); } private static void resultSetNextUntilDone(ResultSet resultSet) throws SQLException { while (resultSet.next()) { // TODO: implement resultSet.last() // Pass to the next until resultSet is done } } /** * Tests whether the {@link ArrowFlightJdbcDriver} can run a query successfully. * * @throws Exception If the connection fails to be established. */ @Test public void testShouldRunSelectQuery() throws Exception { try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { CoreMockedSqlProducers.assertLegacyRegularSqlResultSet(resultSet); } } @Test public void testShouldExecuteQueryNotBlockIfClosedBeforeEnd() throws Exception { try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { for (int i = 0; i < 7500; i++) { assertTrue(resultSet.next()); } } } /** * Tests whether the {@link ArrowFlightJdbcDriver} query only returns only the amount of value set * by {@link org.apache.calcite.avatica.AvaticaStatement#setMaxRows(int)}. * * @throws Exception If the connection fails to be established. */ @Test public void testShouldRunSelectQuerySettingMaxRowLimit() throws Exception { try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { final int maxRowsLimit = 3; statement.setMaxRows(maxRowsLimit); assertThat(statement.getMaxRows(), is(maxRowsLimit)); int count = 0; int columns = 6; for (; resultSet.next(); count++) { for (int column = 1; column <= columns; column++) { resultSet.getObject(column); } assertThat("Test Name #" + count, is(resultSet.getString(2))); } assertThat(maxRowsLimit, is(count)); } } /** * Tests whether the {@link ArrowFlightJdbcDriver} fails upon attempting to run an invalid query. * * @throws Exception If the connection fails to be established. */ @Test public void testShouldThrowExceptionUponAttemptingToExecuteAnInvalidSelectQuery() { assertThrows( SQLException.class, () -> { try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery("SELECT * FROM SHOULD-FAIL")) { fail(); } }); } /** * Tests whether the {@link ArrowFlightJdbcDriver} query only returns only the amount of value set * by {@link org.apache.calcite.avatica.AvaticaStatement#setLargeMaxRows(long)} (int)}. * * @throws Exception If the connection fails to be established. */ @Test public void testShouldRunSelectQuerySettingLargeMaxRowLimit() throws Exception { try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { final long maxRowsLimit = 3; statement.setLargeMaxRows(maxRowsLimit); assertThat(statement.getLargeMaxRows(), is(maxRowsLimit)); int count = 0; int columns = resultSet.getMetaData().getColumnCount(); for (; resultSet.next(); count++) { for (int column = 1; column <= columns; column++) { resultSet.getObject(column); } assertEquals("Test Name #" + count, resultSet.getString(2)); } assertEquals(maxRowsLimit, count); } } @Test public void testColumnCountShouldRemainConsistentForResultSetThroughoutEntireDuration() throws SQLException { final Set counts = new HashSet<>(); try (final Statement statement = connection.createStatement(); final ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { while (resultSet.next()) { counts.add(resultSet.getMetaData().getColumnCount()); } } assertThat(counts, is(ImmutableSet.of(6))); } /** * Tests whether the {@link ArrowFlightJdbcDriver} close the statement after complete ResultSet * when call {@link org.apache.calcite.avatica.AvaticaStatement#closeOnCompletion()}. * * @throws Exception If the connection fails to be established. */ @Test public void testShouldCloseStatementWhenIsCloseOnCompletion() throws Exception { try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { statement.closeOnCompletion(); resultSetNextUntilDone(resultSet); assertThat(statement.isClosed(), is(true)); } } /** * Tests whether the {@link ArrowFlightJdbcDriver} close the statement after complete ResultSet * with max rows limit when call {@link * org.apache.calcite.avatica.AvaticaStatement#closeOnCompletion()}. * * @throws Exception If the connection fails to be established. */ @Test public void testShouldCloseStatementWhenIsCloseOnCompletionWithMaxRowsLimit() throws Exception { try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { final long maxRowsLimit = 3; statement.setLargeMaxRows(maxRowsLimit); statement.closeOnCompletion(); resultSetNextUntilDone(resultSet); assertThat(statement.isClosed(), is(true)); } } /** * Tests whether the {@link ArrowFlightJdbcDriver} not close the statement after complete * ResultSet with max rows limit when call {@link * org.apache.calcite.avatica.AvaticaStatement#closeOnCompletion()}. * * @throws Exception If the connection fails to be established. */ @Test public void testShouldNotCloseStatementWhenIsNotCloseOnCompletionWithMaxRowsLimit() throws Exception { try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { final long maxRowsLimit = 3; statement.setLargeMaxRows(maxRowsLimit); assertThat(statement.isClosed(), is(false)); resultSetNextUntilDone(resultSet); assertThat(resultSet.isClosed(), is(false)); assertThat(resultSet, is(instanceOf(ArrowFlightJdbcFlightStreamResultSet.class))); } } @Test public void testShouldCancelQueryUponCancelAfterQueryingResultSet() throws SQLException { try (final Statement statement = connection.createStatement(); final ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { final int column = RANDOM.nextInt(resultSet.getMetaData().getColumnCount()) + 1; assertThat(resultSet.isClosed(), is(false)); assertThat(resultSet.next(), is(true)); assertDoesNotThrow(() -> resultSet.getObject(column)); statement.cancel(); // Should reset `ResultSet`; keep both `ResultSet` and `Connection` open. assertThat(statement.isClosed(), is(false)); assertThat(resultSet.isClosed(), is(false)); assertThat(resultSet.getMetaData().getColumnCount(), is(0)); } } @Test public void testShouldInterruptFlightStreamsIfQueryIsCancelledMidQuerying() throws SQLException, InterruptedException { try (final Statement statement = connection.createStatement()) { final CountDownLatch latch = new CountDownLatch(1); final Set exceptions = synchronizedSet(new HashSet<>(1)); final Thread thread = new Thread( () -> { try (final ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { final int cachedColumnCount = resultSet.getMetaData().getColumnCount(); Thread.sleep(300); while (resultSet.next()) { resultSet.getObject(RANDOM.nextInt(cachedColumnCount) + 1); } } catch (final SQLException | InterruptedException e) { exceptions.add(e); } finally { latch.countDown(); } }); thread.setName("Test Case: interrupt query execution before first retrieval"); thread.start(); statement.cancel(); thread.join(); assertThat( exceptions.stream() .map(Exception::getMessage) .map(StringBuilder::new) .reduce(StringBuilder::append) .orElseThrow(IllegalArgumentException::new) .toString(), is("Statement canceled")); } } @Test public void testShouldInterruptFlightStreamsIfQueryIsCancelledMidProcessingForTimeConsumingQueries() throws SQLException, InterruptedException { final String query = CoreMockedSqlProducers.LEGACY_CANCELLATION_SQL_CMD; try (final Statement statement = connection.createStatement()) { final Set exceptions = synchronizedSet(new HashSet<>(1)); final Thread thread = new Thread( () -> { try (final ResultSet ignored = statement.executeQuery(query)) { fail(); } catch (final SQLException e) { exceptions.add(e); } }); thread.setName("Test Case: interrupt query execution mid-process"); thread.setPriority(Thread.MAX_PRIORITY); thread.start(); Thread.sleep(5000); // Let the other thread attempt to retrieve results. statement.cancel(); thread.join(); assertThat( exceptions.stream() .map(Exception::getMessage) .map(StringBuilder::new) .reduce(StringBuilder::append) .orElseThrow(IllegalStateException::new) .toString(), anyOf( is(format("Error while executing SQL \"%s\": Query canceled", query)), allOf( containsString(format("Error while executing SQL \"%s\"", query)), anyOf(containsString("CANCELLED"), containsString("Cancelling"))))); } } @Test public void testShouldInterruptFlightStreamsIfQueryTimeoutIsOver() throws SQLException { final String query = CoreMockedSqlProducers.LEGACY_CANCELLATION_SQL_CMD; final int timeoutValue = 2; final String timeoutUnit = "SECONDS"; try (final Statement statement = connection.createStatement()) { statement.setQueryTimeout(timeoutValue); final Set exceptions = new HashSet<>(1); try { statement.executeQuery(query); } catch (final Exception e) { exceptions.add(e); } final Throwable comparisonCause = exceptions.stream().findFirst().orElseThrow(RuntimeException::new).getCause().getCause(); assertThat(comparisonCause, is(instanceOf(SQLTimeoutException.class))); assertThat( comparisonCause.getMessage(), is(format("Query timed out after %d %s", timeoutValue, timeoutUnit))); } } @Test public void testFlightStreamsQueryShouldNotTimeout() throws SQLException { final String query = CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD; final int timeoutValue = 5; try (Statement statement = connection.createStatement()) { statement.setQueryTimeout(timeoutValue); try (ResultSet resultSet = statement.executeQuery(query)) { CoreMockedSqlProducers.assertLegacyRegularSqlResultSet(resultSet); } } } @Test public void testPartitionedFlightServer() throws Exception { // Arrange final Schema schema = new Schema( Arrays.asList(Field.nullablePrimitive("int_column", new ArrowType.Int(32, true)))); try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); VectorSchemaRoot firstPartition = VectorSchemaRoot.create(schema, allocator); VectorSchemaRoot secondPartition = VectorSchemaRoot.create(schema, allocator)) { firstPartition.setRowCount(1); ((IntVector) firstPartition.getVector(0)).set(0, 1); secondPartition.setRowCount(1); ((IntVector) secondPartition.getVector(0)).set(0, 2); // Construct the data-only nodes first. FlightProducer firstProducer = new PartitionedFlightSqlProducer.DataOnlyFlightSqlProducer( new Ticket("first".getBytes(StandardCharsets.UTF_8)), firstPartition); FlightProducer secondProducer = new PartitionedFlightSqlProducer.DataOnlyFlightSqlProducer( new Ticket("second".getBytes(StandardCharsets.UTF_8)), secondPartition); final FlightServer.Builder firstBuilder = FlightServer.builder(allocator, forGrpcInsecure("localhost", 0), firstProducer); final FlightServer.Builder secondBuilder = FlightServer.builder(allocator, forGrpcInsecure("localhost", 0), secondProducer); // Run the data-only nodes so that we can get the Locations they are running at. try (FlightServer firstServer = firstBuilder.build(); FlightServer secondServer = secondBuilder.build()) { firstServer.start(); secondServer.start(); final FlightEndpoint firstEndpoint = new FlightEndpoint( new Ticket("first".getBytes(StandardCharsets.UTF_8)), firstServer.getLocation()); final FlightEndpoint secondEndpoint = new FlightEndpoint( new Ticket("second".getBytes(StandardCharsets.UTF_8)), secondServer.getLocation()); // Finally start the root node. try (final PartitionedFlightSqlProducer rootProducer = new PartitionedFlightSqlProducer(schema, firstEndpoint, secondEndpoint); 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())); Statement newStatement = newConnection.createStatement(); // Act ResultSet result = newStatement.executeQuery("Select partitioned_data")) { List resultData = new ArrayList<>(); while (result.next()) { resultData.add(result.getInt(1)); } // Assert assertEquals( firstPartition.getRowCount() + secondPartition.getRowCount(), resultData.size()); assertTrue(resultData.contains(((IntVector) firstPartition.getVector(0)).get(0))); assertTrue(resultData.contains(((IntVector) secondPartition.getVector(0)).get(0))); } } } } @Test public void testPartitionedFlightServerIgnoreFailure() throws Exception { final Schema schema = new Schema( Collections.singletonList( Field.nullablePrimitive("int_column", new ArrowType.Int(32, true)))); try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { final FlightEndpoint firstEndpoint = new FlightEndpoint( new Ticket("first".getBytes(StandardCharsets.UTF_8)), Location.forGrpcInsecure("127.0.0.2", 1234), Location.forGrpcInsecure("127.0.0.3", 1234)); try (final PartitionedFlightSqlProducer rootProducer = new PartitionedFlightSqlProducer(schema, firstEndpoint); 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())); Statement newStatement = newConnection.createStatement()) { final SQLException e = assertThrows( SQLException.class, () -> { ResultSet result = newStatement.executeQuery("Select partitioned_data"); while (result.next()) {} }); final Throwable cause = e.getCause(); assertTrue(cause instanceof FlightRuntimeException); final FlightRuntimeException fre = (FlightRuntimeException) cause; assertEquals(FlightStatusCode.UNAVAILABLE, fre.status().code()); } } } @Test public void testPartitionedFlightServerAllFailure() throws Exception { // Arrange final Schema schema = new Schema( Collections.singletonList( Field.nullablePrimitive("int_column", new ArrowType.Int(32, true)))); try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); VectorSchemaRoot firstPartition = VectorSchemaRoot.create(schema, allocator)) { firstPartition.setRowCount(1); ((IntVector) firstPartition.getVector(0)).set(0, 1); // Construct the data-only nodes first. FlightProducer firstProducer = new PartitionedFlightSqlProducer.DataOnlyFlightSqlProducer( new Ticket("first".getBytes(StandardCharsets.UTF_8)), firstPartition); final FlightServer.Builder firstBuilder = FlightServer.builder(allocator, forGrpcInsecure("localhost", 0), firstProducer); // Run the data-only nodes so that we can get the Locations they are running at. try (FlightServer firstServer = firstBuilder.build()) { firstServer.start(); final Location badLocation = Location.forGrpcInsecure("127.0.0.2", 1234); final FlightEndpoint firstEndpoint = new FlightEndpoint( new Ticket("first".getBytes(StandardCharsets.UTF_8)), badLocation, firstServer.getLocation()); // Finally start the root node. try (final PartitionedFlightSqlProducer rootProducer = new PartitionedFlightSqlProducer(schema, firstEndpoint); 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())); Statement newStatement = newConnection.createStatement(); // Act ResultSet result = newStatement.executeQuery("Select partitioned_data")) { List resultData = new ArrayList<>(); while (result.next()) { resultData.add(result.getInt(1)); } // Assert assertEquals(firstPartition.getRowCount(), resultData.size()); assertTrue(resultData.contains(((IntVector) firstPartition.getVector(0)).get(0))); } } } } @Test public void testFallbackFlightServer() 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())); Statement newStatement = newConnection.createStatement(); ResultSet result = newStatement.executeQuery("fallback")) { 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))); } } } @Test public void testFallbackSecondFlightServer() 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())); Statement newStatement = newConnection.createStatement(); ResultSet result = newStatement.executeQuery("fallback with error")) { 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))); } } } @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(); ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_WITH_EMPTY_SQL_CMD)) { long rowCount = 0; while (resultSet.next()) { ++rowCount; } assertEquals(2, rowCount); } } @Test public void testResultSetAppMetadata() throws Exception { try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { assertArrayEquals( ((ArrowFlightJdbcFlightStreamResultSet) resultSet).getAppMetadata(), "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)); } } } TimestampResultSetTest.java000066400000000000000000000157511515451615200375540ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } } TokenAuthenticationTest.java000066400000000000000000000045101515451615200377050ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; import java.sql.Connection; import java.sql.SQLException; import org.apache.arrow.driver.jdbc.authentication.TokenAuthentication; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; import org.apache.arrow.util.AutoCloseables; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; public class TokenAuthenticationTest { 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("1234").build()) .producer(FLIGHT_SQL_PRODUCER) .build(); } @AfterAll public static void tearDownAfterClass() { AutoCloseables.closeNoChecked(FLIGHT_SQL_PRODUCER); } @Test public void connectUsingTokenAuthenticationShouldFail() throws SQLException { assertThrows( SQLException.class, () -> FLIGHT_SERVER_TEST_EXTENSION.getConnection(false, "invalid")); } @Test public void connectUsingTokenAuthenticationShouldSuccess() throws SQLException { try (Connection connection = FLIGHT_SERVER_TEST_EXTENSION.getConnection(false, "1234")) { assertFalse(connection.isClosed()); } } } accessor/000077500000000000000000000000001515451615200340245ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbcArrowFlightJdbcAccessorFactoryTest.java000066400000000000000000000475211515451615200435660ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; 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; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeStampVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.complex.ArrowFlightJdbcDenseUnionVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.complex.ArrowFlightJdbcFixedSizeListVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.complex.ArrowFlightJdbcLargeListVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.complex.ArrowFlightJdbcListVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.complex.ArrowFlightJdbcMapVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.complex.ArrowFlightJdbcStructVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.complex.ArrowFlightJdbcUnionVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.numeric.ArrowFlightJdbcBaseIntVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.numeric.ArrowFlightJdbcBitVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.numeric.ArrowFlightJdbcDecimalVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.numeric.ArrowFlightJdbcFloat4VectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.numeric.ArrowFlightJdbcFloat8VectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.text.ArrowFlightJdbcVarCharVectorAccessor; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.DurationVector; import org.apache.arrow.vector.IntervalDayVector; import org.apache.arrow.vector.IntervalMonthDayNanoVector; import org.apache.arrow.vector.IntervalYearVector; 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; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.FieldType; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; public class ArrowFlightJdbcAccessorFactoryTest { public static final IntSupplier GET_CURRENT_ROW = () -> 0; @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); @Test public void createAccessorForUInt1Vector() { try (ValueVector valueVector = rootAllocatorTestExtension.createUInt1Vector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcBaseIntVectorAccessor); } } @Test public void createAccessorForUInt2Vector() { try (ValueVector valueVector = rootAllocatorTestExtension.createUInt2Vector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcBaseIntVectorAccessor); } } @Test public void createAccessorForUInt4Vector() { try (ValueVector valueVector = rootAllocatorTestExtension.createUInt4Vector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcBaseIntVectorAccessor); } } @Test public void createAccessorForUInt8Vector() { try (ValueVector valueVector = rootAllocatorTestExtension.createUInt8Vector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcBaseIntVectorAccessor); } } @Test public void createAccessorForTinyIntVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createTinyIntVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcBaseIntVectorAccessor); } } @Test public void createAccessorForSmallIntVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createSmallIntVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcBaseIntVectorAccessor); } } @Test public void createAccessorForIntVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createIntVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcBaseIntVectorAccessor); } } @Test public void createAccessorForBigIntVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createBigIntVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcBaseIntVectorAccessor); } } @Test public void createAccessorForFloat4Vector() { try (ValueVector valueVector = rootAllocatorTestExtension.createFloat4Vector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcFloat4VectorAccessor); } } @Test public void createAccessorForFloat8Vector() { try (ValueVector valueVector = rootAllocatorTestExtension.createFloat8Vector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcFloat8VectorAccessor); } } @Test public void createAccessorForBitVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createBitVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcBitVectorAccessor); } } @Test public void createAccessorForDecimalVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createDecimalVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcDecimalVectorAccessor); } } @Test public void createAccessorForDecimal256Vector() { try (ValueVector valueVector = rootAllocatorTestExtension.createDecimal256Vector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcDecimalVectorAccessor); } } @Test public void createAccessorForVarBinaryVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createVarBinaryVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcBinaryVectorAccessor); } } @Test public void createAccessorForLargeVarBinaryVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createLargeVarBinaryVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcBinaryVectorAccessor); } } @Test public void createAccessorForFixedSizeBinaryVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createFixedSizeBinaryVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcBinaryVectorAccessor); } } @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()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcTimeStampVectorAccessor); } } @Test public void createAccessorForTimeNanoVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createTimeNanoVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcTimeVectorAccessor); } } @Test public void createAccessorForTimeMicroVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createTimeMicroVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcTimeVectorAccessor); } } @Test public void createAccessorForTimeMilliVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createTimeMilliVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcTimeVectorAccessor); } } @Test public void createAccessorForTimeSecVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createTimeSecVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcTimeVectorAccessor); } } @Test public void createAccessorForDateDayVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createDateDayVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcDateVectorAccessor); } } @Test public void createAccessorForDateMilliVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createDateMilliVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcDateVectorAccessor); } } @Test public void createAccessorForVarCharVector() { try (ValueVector valueVector = new VarCharVector("", rootAllocatorTestExtension.getRootAllocator())) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcVarCharVectorAccessor); } } @Test public void createAccessorForLargeVarCharVector() { try (ValueVector valueVector = new LargeVarCharVector("", rootAllocatorTestExtension.getRootAllocator())) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcVarCharVectorAccessor); } } @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 = new DurationVector( "", new FieldType(true, new ArrowType.Duration(TimeUnit.MILLISECOND), null), rootAllocatorTestExtension.getRootAllocator())) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcDurationVectorAccessor); } } @Test public void createAccessorForIntervalDayVector() { try (ValueVector valueVector = new IntervalDayVector("", rootAllocatorTestExtension.getRootAllocator())) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcIntervalVectorAccessor); } } @Test public void createAccessorForIntervalYearVector() { try (ValueVector valueVector = new IntervalYearVector("", rootAllocatorTestExtension.getRootAllocator())) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcIntervalVectorAccessor); } } @Test public void createAccessorForIntervalMonthDayNanoVector() { try (ValueVector valueVector = new IntervalMonthDayNanoVector("", rootAllocatorTestExtension.getRootAllocator())) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcIntervalVectorAccessor); } } @Test public void createAccessorForUnionVector() { try (ValueVector valueVector = new UnionVector("", rootAllocatorTestExtension.getRootAllocator(), null, null)) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcUnionVectorAccessor); } } @Test public void createAccessorForDenseUnionVector() { try (ValueVector valueVector = new DenseUnionVector("", rootAllocatorTestExtension.getRootAllocator(), null, null)) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcDenseUnionVectorAccessor); } } @Test public void createAccessorForStructVector() { try (ValueVector valueVector = StructVector.empty("", rootAllocatorTestExtension.getRootAllocator())) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcStructVectorAccessor); } } @Test public void createAccessorForListVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createListVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcListVectorAccessor); } } @Test public void createAccessorForLargeListVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createLargeListVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcLargeListVectorAccessor); } } @Test public void createAccessorForFixedSizeListVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createFixedSizeListVector()) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); assertTrue(accessor instanceof ArrowFlightJdbcFixedSizeListVectorAccessor); } } @Test public void createAccessorForMapVector() { try (ValueVector valueVector = MapVector.empty("", rootAllocatorTestExtension.getRootAllocator(), true)) { ArrowFlightJdbcAccessor accessor = ArrowFlightJdbcAccessorFactory.createAccessor( valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); 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); } } } ArrowFlightJdbcAccessorTest.java000066400000000000000000000272211515451615200422310ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.junit.jupiter.api.Assertions.assertArrayEquals; 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.math.BigDecimal; import java.nio.charset.StandardCharsets; import java.sql.SQLException; import java.util.HashMap; import java.util.Map; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; @ExtendWith(MockitoExtension.class) public class ArrowFlightJdbcAccessorTest { static class MockedArrowFlightJdbcAccessor extends ArrowFlightJdbcAccessor { protected MockedArrowFlightJdbcAccessor() { super(() -> 0, (boolean wasNull) -> {}); } @Override public Class getObjectClass() { return Long.class; } } @Mock MockedArrowFlightJdbcAccessor accessor; @Test public void testShouldGetObjectWithByteClassReturnGetByte() throws SQLException { byte expected = Byte.MAX_VALUE; when(accessor.getByte()).thenReturn(expected); when(accessor.getObject(Byte.class)).thenCallRealMethod(); assertEquals(accessor.getObject(Byte.class), (Object) expected); verify(accessor).getByte(); } @Test public void testShouldGetObjectWithShortClassReturnGetShort() throws SQLException { short expected = Short.MAX_VALUE; when(accessor.getShort()).thenReturn(expected); when(accessor.getObject(Short.class)).thenCallRealMethod(); assertEquals(accessor.getObject(Short.class), (Object) expected); verify(accessor).getShort(); } @Test public void testShouldGetObjectWithIntegerClassReturnGetInt() throws SQLException { int expected = Integer.MAX_VALUE; when(accessor.getInt()).thenReturn(expected); when(accessor.getObject(Integer.class)).thenCallRealMethod(); assertEquals(accessor.getObject(Integer.class), (Object) expected); verify(accessor).getInt(); } @Test public void testShouldGetObjectWithLongClassReturnGetLong() throws SQLException { long expected = Long.MAX_VALUE; when(accessor.getLong()).thenReturn(expected); when(accessor.getObject(Long.class)).thenCallRealMethod(); assertEquals(accessor.getObject(Long.class), (Object) expected); verify(accessor).getLong(); } @Test public void testShouldGetObjectWithFloatClassReturnGetFloat() throws SQLException { float expected = Float.MAX_VALUE; when(accessor.getFloat()).thenReturn(expected); when(accessor.getObject(Float.class)).thenCallRealMethod(); assertEquals(accessor.getObject(Float.class), (Object) expected); verify(accessor).getFloat(); } @Test public void testShouldGetObjectWithDoubleClassReturnGetDouble() throws SQLException { double expected = Double.MAX_VALUE; when(accessor.getDouble()).thenReturn(expected); when(accessor.getObject(Double.class)).thenCallRealMethod(); assertEquals(accessor.getObject(Double.class), (Object) expected); verify(accessor).getDouble(); } @Test public void testShouldGetObjectWithBooleanClassReturnGetBoolean() throws SQLException { when(accessor.getBoolean()).thenReturn(true); when(accessor.getObject(Boolean.class)).thenCallRealMethod(); assertEquals(true, accessor.getObject(Boolean.class)); verify(accessor).getBoolean(); } @Test public void testShouldGetObjectWithBigDecimalClassReturnGetBigDecimal() throws SQLException { BigDecimal expected = BigDecimal.TEN; when(accessor.getBigDecimal()).thenReturn(expected); when(accessor.getObject(BigDecimal.class)).thenCallRealMethod(); assertEquals(expected, accessor.getObject(BigDecimal.class)); verify(accessor).getBigDecimal(); } @Test public void testShouldGetObjectWithStringClassReturnGetString() throws SQLException { String expected = "STRING_VALUE"; when(accessor.getString()).thenReturn(expected); when(accessor.getObject(String.class)).thenCallRealMethod(); assertEquals(expected, accessor.getObject(String.class)); verify(accessor).getString(); } @Test public void testShouldGetObjectWithByteArrayClassReturnGetBytes() throws SQLException { byte[] expected = "STRING_VALUE".getBytes(StandardCharsets.UTF_8); when(accessor.getBytes()).thenReturn(expected); when(accessor.getObject(byte[].class)).thenCallRealMethod(); assertArrayEquals(accessor.getObject(byte[].class), expected); verify(accessor).getBytes(); } @Test public void testShouldGetObjectWithObjectClassReturnGetObject() throws SQLException { Object expected = new Object(); when(accessor.getObject()).thenReturn(expected); when(accessor.getObject(Object.class)).thenCallRealMethod(); assertEquals(expected, accessor.getObject(Object.class)); verify(accessor).getObject(); } @Test public void testShouldGetObjectWithAccessorsObjectClassReturnGetObject() throws SQLException { Class objectClass = Long.class; when(accessor.getObject(objectClass)).thenCallRealMethod(); accessor.getObject(objectClass); verify(accessor).getObject(objectClass); } @Test public void testShouldFailToGetBoolean() throws SQLException { when(accessor.getBoolean()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getBoolean()); } @Test public void testShouldFailToGetByte() throws SQLException { when(accessor.getByte()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getByte()); } @Test public void testShouldFailToGetShort() throws SQLException { when(accessor.getShort()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getShort()); } @Test public void testShouldFailToGetInt() throws SQLException { when(accessor.getInt()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getInt()); } @Test public void testShouldFailToGetLong() throws SQLException { when(accessor.getLong()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getLong()); } @Test public void testShouldFailToGetFloat() throws SQLException { when(accessor.getFloat()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getFloat()); } @Test public void testShouldFailToGetDouble() throws SQLException { when(accessor.getDouble()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getDouble()); } @Test public void testShouldFailToGetBigDecimal() throws SQLException { when(accessor.getBigDecimal()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getBigDecimal()); } @Test public void testShouldFailToGetBytes() throws SQLException { when(accessor.getBytes()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getBytes()); } @Test public void testShouldFailToGetAsciiStream() throws SQLException { when(accessor.getAsciiStream()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getAsciiStream()); } @Test public void testShouldFailToGetUnicodeStream() throws SQLException { when(accessor.getUnicodeStream()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getUnicodeStream()); } @Test public void testShouldFailToGetBinaryStream() throws SQLException { when(accessor.getBinaryStream()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getBinaryStream()); } @Test public void testShouldFailToGetObject() throws SQLException { when(accessor.getObject()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getObject()); } @Test public void testShouldFailToGetObjectMap() throws SQLException { Map> map = new HashMap<>(); when(accessor.getObject(map)).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getObject(map)); } @Test public void testShouldFailToGetCharacterStream() throws SQLException { when(accessor.getCharacterStream()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getCharacterStream()); } @Test public void testShouldFailToGetRef() throws SQLException { when(accessor.getRef()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getRef()); } @Test public void testShouldFailToGetBlob() throws SQLException { when(accessor.getBlob()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getBlob()); } @Test public void testShouldFailToGetClob() throws SQLException { when(accessor.getClob()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getClob()); } @Test public void testShouldFailToGetArray() throws SQLException { when(accessor.getArray()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getArray()); } @Test public void testShouldFailToGetStruct() throws SQLException { when(accessor.getStruct()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getStruct()); } @Test public void testShouldFailToGetURL() throws SQLException { when(accessor.getURL()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getURL()); } @Test public void testShouldFailToGetNClob() throws SQLException { when(accessor.getNClob()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getNClob()); } @Test public void testShouldFailToGetSQLXML() throws SQLException { when(accessor.getSQLXML()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getSQLXML()); } @Test public void testShouldFailToGetNString() throws SQLException { when(accessor.getNString()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getNString()); } @Test public void testShouldFailToGetNCharacterStream() throws SQLException { when(accessor.getNCharacterStream()).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getNCharacterStream()); } @Test public void testShouldFailToGetDate() throws SQLException { when(accessor.getDate(null)).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getDate(null)); } @Test public void testShouldFailToGetTime() throws SQLException { when(accessor.getTime(null)).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getTime(null)); } @Test public void testShouldFailToGetTimestamp() throws SQLException { when(accessor.getTimestamp(null)).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getTimestamp(null)); } @Test public void testShouldFailToGetBigDecimalWithValue() throws SQLException { when(accessor.getBigDecimal(0)).thenCallRealMethod(); assertThrows(SQLException.class, () -> accessor.getBigDecimal(0)); } } impl/000077500000000000000000000000001515451615200347655ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessorArrowFlightJdbcNullVectorAccessorTest.java000066400000000000000000000025401515451615200452050ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; import org.junit.jupiter.api.Test; public class ArrowFlightJdbcNullVectorAccessorTest { ArrowFlightJdbcNullVectorAccessor accessor = new ArrowFlightJdbcNullVectorAccessor((boolean wasNull) -> {}); @Test public void testShouldWasNullReturnTrue() { assertTrue(accessor.wasNull()); } @Test public void testShouldGetObjectReturnNull() { assertNull(accessor.getObject()); } } binary/000077500000000000000000000000001515451615200362515ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/implArrowFlightJdbcBinaryVectorAccessorTest.java000066400000000000000000000245541515451615200470140ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.nio.charset.StandardCharsets.US_ASCII; import static java.nio.charset.StandardCharsets.UTF_8; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; import java.io.InputStream; import java.io.Reader; import java.util.function.Supplier; import java.util.stream.Stream; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.FixedSizeBinaryVector; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.commons.io.IOUtils; import org.hamcrest.CoreMatchers; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; public class ArrowFlightJdbcBinaryVectorAccessorTest { @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private ValueVector vector; private final AccessorTestUtils.AccessorSupplier accessorSupplier = (vector, getCurrentRow) -> { ArrowFlightJdbcAccessorFactory.WasNullConsumer noOpWasNullConsumer = (boolean wasNull) -> {}; if (vector instanceof VarBinaryVector) { return new ArrowFlightJdbcBinaryVectorAccessor( ((VarBinaryVector) vector), getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof LargeVarBinaryVector) { return new ArrowFlightJdbcBinaryVectorAccessor( ((LargeVarBinaryVector) vector), getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof FixedSizeBinaryVector) { return new ArrowFlightJdbcBinaryVectorAccessor( ((FixedSizeBinaryVector) vector), getCurrentRow, noOpWasNullConsumer); } return null; }; private final AccessorTestUtils.AccessorIterator accessorIterator = new AccessorTestUtils.AccessorIterator<>(accessorSupplier); public static Stream data() { return Stream.of( Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createVarBinaryVector(), "VarBinaryVector"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createLargeVarBinaryVector(), "LargeVarBinaryVector"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createFixedSizeBinaryVector(), "FixedSizeBinaryVector")); } public void setup(Supplier vectorSupplier) { this.vector = vectorSupplier.get(); } @AfterEach public void tearDown() { this.vector.close(); } @ParameterizedTest @MethodSource("data") public void testShouldGetStringReturnExpectedString(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcBinaryVectorAccessor::getString, (accessor) -> is(new String(accessor.getBytes(), UTF_8))); } @ParameterizedTest @MethodSource("data") public void testShouldGetStringReturnNull(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); vector.reset(); vector.setValueCount(5); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcBinaryVectorAccessor::getString, CoreMatchers.nullValue()); } @ParameterizedTest @MethodSource("data") public void testShouldGetBytesReturnExpectedByteArray(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcBinaryVectorAccessor::getBytes, (accessor, currentRow) -> { if (vector instanceof VarBinaryVector) { return is(((VarBinaryVector) vector).get(currentRow)); } else if (vector instanceof LargeVarBinaryVector) { return is(((LargeVarBinaryVector) vector).get(currentRow)); } else if (vector instanceof FixedSizeBinaryVector) { return is(((FixedSizeBinaryVector) vector).get(currentRow)); } return null; }); } @ParameterizedTest @MethodSource("data") public void testShouldGetBytesReturnNull(Supplier vectorSupplier) { setup(vectorSupplier); vector.reset(); vector.setValueCount(5); ArrowFlightJdbcBinaryVectorAccessor accessor = accessorSupplier.supply(vector, () -> 0); assertThat(accessor.getBytes(), CoreMatchers.equalTo(null)); assertThat(accessor.wasNull(), is(true)); } @ParameterizedTest @MethodSource("data") public void testShouldGetObjectReturnAsGetBytes(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcBinaryVectorAccessor::getObject, (accessor) -> is(accessor.getBytes())); } @ParameterizedTest @MethodSource("data") public void testShouldGetObjectReturnNull(Supplier vectorSupplier) { setup(vectorSupplier); vector.reset(); vector.setValueCount(5); ArrowFlightJdbcBinaryVectorAccessor accessor = accessorSupplier.supply(vector, () -> 0); assertThat(accessor.getObject(), CoreMatchers.equalTo(null)); assertThat(accessor.wasNull(), is(true)); } @ParameterizedTest @MethodSource("data") public void testShouldGetUnicodeStreamReturnCorrectInputStream( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.iterate( vector, (accessor, currentRow) -> { InputStream inputStream = accessor.getUnicodeStream(); String actualString = IOUtils.toString(inputStream, UTF_8); assertThat(accessor.wasNull(), is(false)); assertThat(actualString, is(accessor.getString())); }); } @ParameterizedTest @MethodSource("data") public void testShouldGetUnicodeStreamReturnNull(Supplier vectorSupplier) { setup(vectorSupplier); vector.reset(); vector.setValueCount(5); ArrowFlightJdbcBinaryVectorAccessor accessor = accessorSupplier.supply(vector, () -> 0); assertThat(accessor.getUnicodeStream(), CoreMatchers.equalTo(null)); assertThat(accessor.wasNull(), is(true)); } @ParameterizedTest @MethodSource("data") public void testShouldGetAsciiStreamReturnCorrectInputStream(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.iterate( vector, (accessor, currentRow) -> { InputStream inputStream = accessor.getAsciiStream(); String actualString = IOUtils.toString(inputStream, US_ASCII); assertThat(accessor.wasNull(), is(false)); assertThat(actualString, is(accessor.getString())); }); } @ParameterizedTest @MethodSource("data") public void testShouldGetAsciiStreamReturnNull(Supplier vectorSupplier) { setup(vectorSupplier); vector.reset(); vector.setValueCount(5); ArrowFlightJdbcBinaryVectorAccessor accessor = accessorSupplier.supply(vector, () -> 0); assertThat(accessor.getAsciiStream(), CoreMatchers.equalTo(null)); assertThat(accessor.wasNull(), is(true)); } @ParameterizedTest @MethodSource("data") public void testShouldGetBinaryStreamReturnCurrentInputStream( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.iterate( vector, (accessor, currentRow) -> { InputStream inputStream = accessor.getBinaryStream(); String actualString = IOUtils.toString(inputStream, UTF_8); assertThat(accessor.wasNull(), is(false)); assertThat(actualString, is(accessor.getString())); }); } @ParameterizedTest @MethodSource("data") public void testShouldGetBinaryStreamReturnNull(Supplier vectorSupplier) { setup(vectorSupplier); vector.reset(); vector.setValueCount(5); ArrowFlightJdbcBinaryVectorAccessor accessor = accessorSupplier.supply(vector, () -> 0); assertThat(accessor.getBinaryStream(), CoreMatchers.equalTo(null)); assertThat(accessor.wasNull(), is(true)); } @ParameterizedTest @MethodSource("data") public void testShouldGetCharacterStreamReturnCorrectReader(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.iterate( vector, (accessor, currentRow) -> { Reader characterStream = accessor.getCharacterStream(); String actualString = IOUtils.toString(characterStream); assertThat(accessor.wasNull(), is(false)); assertThat(actualString, is(accessor.getString())); }); } @ParameterizedTest @MethodSource("data") public void testShouldGetCharacterStreamReturnNull(Supplier vectorSupplier) { setup(vectorSupplier); vector.reset(); vector.setValueCount(5); ArrowFlightJdbcBinaryVectorAccessor accessor = accessorSupplier.supply(vector, () -> 0); assertThat(accessor.getCharacterStream(), CoreMatchers.equalTo(null)); assertThat(accessor.wasNull(), is(true)); } } ArrowFlightJdbcUuidVectorAccessorTest.java000066400000000000000000000143551515451615200464740ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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)); } } calendar/000077500000000000000000000000001515451615200365365ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/implArrowFlightJdbcDateVectorAccessorTest.java000066400000000000000000000252301515451615200467220ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.calendar; import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDateVectorAccessor.getTimeUnitForVector; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.CoreMatchers.not; import static org.hamcrest.MatcherAssert.assertThat; import java.sql.Date; import java.sql.Timestamp; import java.time.LocalDateTime; import java.util.Calendar; import java.util.TimeZone; import java.util.concurrent.TimeUnit; import java.util.function.Supplier; import java.util.stream.Stream; import org.apache.arrow.driver.jdbc.accessor.impl.text.ArrowFlightJdbcVarCharVectorAccessor; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DateMilliVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.util.Text; import org.hamcrest.CoreMatchers; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; public class ArrowFlightJdbcDateVectorAccessorTest { public static final String AMERICA_VANCOUVER = "America/Vancouver"; @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private BaseFixedWidthVector vector; private final AccessorTestUtils.AccessorSupplier accessorSupplier = (vector, getCurrentRow) -> { if (vector instanceof DateDayVector) { return new ArrowFlightJdbcDateVectorAccessor( (DateDayVector) vector, getCurrentRow, (boolean wasNull) -> {}); } else if (vector instanceof DateMilliVector) { return new ArrowFlightJdbcDateVectorAccessor( (DateMilliVector) vector, getCurrentRow, (boolean wasNull) -> {}); } return null; }; private final AccessorTestUtils.AccessorIterator accessorIterator = new AccessorTestUtils.AccessorIterator<>(accessorSupplier); public static Stream data() { return Stream.of( Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createDateDayVector(), "DateDayVector"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createDateMilliVector(), "DateMilliVector")); } public void setup(Supplier vectorSupplier) { this.vector = vectorSupplier.get(); } @AfterEach public void tearDown() { this.vector.close(); } @ParameterizedTest @MethodSource("data") public void testShouldGetTimestampReturnValidTimestampWithoutCalendar( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, accessor -> accessor.getTimestamp(null), (accessor, currentRow) -> is(getTimestampForVector(currentRow))); } @ParameterizedTest @MethodSource("data") public void testShouldGetObjectWithDateClassReturnValidDateWithoutCalendar( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, accessor -> accessor.getObject(Date.class), (accessor, currentRow) -> is(new Date(getTimestampForVector(currentRow).getTime()))); } @ParameterizedTest @MethodSource("data") public void testShouldGetTimestampReturnValidTimestampWithCalendar( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); TimeZone timeZone = TimeZone.getTimeZone(AMERICA_VANCOUVER); Calendar calendar = Calendar.getInstance(timeZone); accessorIterator.iterate( vector, (accessor, currentRow) -> { final Timestamp resultWithoutCalendar = accessor.getTimestamp(null); final Timestamp result = accessor.getTimestamp(calendar); long offset = timeZone.getOffset(resultWithoutCalendar.getTime()); assertThat(resultWithoutCalendar.getTime() - result.getTime(), is(offset)); assertThat(accessor.wasNull(), is(false)); }); } @ParameterizedTest @MethodSource("data") public void testShouldGetTimestampReturnNull(Supplier vectorSupplier) { setup(vectorSupplier); vector.setNull(0); ArrowFlightJdbcDateVectorAccessor accessor = accessorSupplier.supply(vector, () -> 0); assertThat(accessor.getTimestamp(null), CoreMatchers.equalTo(null)); assertThat(accessor.wasNull(), is(true)); } @ParameterizedTest @MethodSource("data") public void testShouldGetDateReturnValidDateWithoutCalendar( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, accessor -> accessor.getDate(null), (accessor, currentRow) -> is(new Date(getTimestampForVector(currentRow).getTime()))); } @ParameterizedTest @MethodSource("data") public void testShouldGetDateReturnValidDateWithCalendar( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); TimeZone timeZone = TimeZone.getTimeZone(AMERICA_VANCOUVER); Calendar calendar = Calendar.getInstance(timeZone); accessorIterator.iterate( vector, (accessor, currentRow) -> { final Date resultWithoutCalendar = accessor.getDate(null); final Date result = accessor.getDate(calendar); long offset = timeZone.getOffset(resultWithoutCalendar.getTime()); assertThat(resultWithoutCalendar.getTime() - result.getTime(), is(offset)); assertThat(accessor.wasNull(), is(false)); }); } @ParameterizedTest @MethodSource("data") public void testShouldGetDateReturnNull(Supplier vectorSupplier) { setup(vectorSupplier); vector.setNull(0); ArrowFlightJdbcDateVectorAccessor accessor = accessorSupplier.supply(vector, () -> 0); assertThat(accessor.getDate(null), CoreMatchers.equalTo(null)); assertThat(accessor.wasNull(), is(true)); } private Timestamp getTimestampForVector(int currentRow) { Object object = vector.getObject(currentRow); Timestamp expectedTimestamp = null; if (object instanceof LocalDateTime) { expectedTimestamp = Timestamp.valueOf((LocalDateTime) object); } else if (object instanceof Number) { long value = ((Number) object).longValue(); TimeUnit timeUnit = getTimeUnitForVector(vector); long millis = timeUnit.toMillis(value); expectedTimestamp = new Timestamp(millis); } return expectedTimestamp; } @ParameterizedTest @MethodSource("data") public void testShouldGetObjectClass(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcDateVectorAccessor::getObjectClass, equalTo(Date.class)); } @ParameterizedTest @MethodSource("data") public void testShouldGetStringBeConsistentWithVarCharAccessorWithoutCalendar( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); assertGetStringIsConsistentWithVarCharAccessor(null); } @ParameterizedTest @MethodSource("data") public void testShouldGetStringBeConsistentWithVarCharAccessorWithCalendar( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); Calendar calendar = Calendar.getInstance(TimeZone.getTimeZone(AMERICA_VANCOUVER)); assertGetStringIsConsistentWithVarCharAccessor(calendar); } @ParameterizedTest @MethodSource("data") public void testValidateGetStringTimeZoneConsistency( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.iterate( vector, (accessor, currentRow) -> { final TimeZone defaultTz = TimeZone.getDefault(); try { final String string = accessor.getString(); // Should always be UTC as no calendar is provided // Validate with UTC Date date = accessor.getDate(null); TimeZone.setDefault(TimeZone.getTimeZone("UTC")); assertThat(date.toString(), is(string)); // Validate with different TZ TimeZone.setDefault(TimeZone.getTimeZone(AMERICA_VANCOUVER)); assertThat(date.toString(), not(string)); assertThat(accessor.wasNull(), is(false)); } finally { // Set default Tz back TimeZone.setDefault(defaultTz); } }); } private void assertGetStringIsConsistentWithVarCharAccessor(Calendar calendar) throws Exception { try (VarCharVector varCharVector = new VarCharVector("", rootAllocatorTestExtension.getRootAllocator())) { varCharVector.allocateNew(1); ArrowFlightJdbcVarCharVectorAccessor varCharVectorAccessor = new ArrowFlightJdbcVarCharVectorAccessor(varCharVector, () -> 0, (boolean wasNull) -> {}); accessorIterator.iterate( vector, (accessor, currentRow) -> { final String string = accessor.getString(); varCharVector.set(0, new Text(string)); varCharVector.setValueCount(1); Date dateFromVarChar = varCharVectorAccessor.getDate(calendar); Date date = accessor.getDate(calendar); assertThat(date, is(dateFromVarChar)); assertThat(accessor.wasNull(), is(false)); }); } } } ArrowFlightJdbcDurationVectorAccessorTest.java000066400000000000000000000101131515451615200476240ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.calendar; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import java.time.Duration; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.DurationVector; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.FieldType; 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; public class ArrowFlightJdbcDurationVectorAccessorTest { @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private DurationVector vector; private final AccessorTestUtils.AccessorSupplier accessorSupplier = (vector, getCurrentRow) -> new ArrowFlightJdbcDurationVectorAccessor( (DurationVector) vector, getCurrentRow, (boolean wasNull) -> {}); private final AccessorTestUtils.AccessorIterator accessorIterator = new AccessorTestUtils.AccessorIterator<>(accessorSupplier); @BeforeEach public void setup() { FieldType fieldType = new FieldType(true, new ArrowType.Duration(TimeUnit.MILLISECOND), null); this.vector = new DurationVector("", fieldType, rootAllocatorTestExtension.getRootAllocator()); int valueCount = 10; this.vector.setValueCount(valueCount); for (int i = 0; i < valueCount; i++) { this.vector.set(i, java.util.concurrent.TimeUnit.DAYS.toMillis(i + 1)); } } @AfterEach public void tearDown() { this.vector.close(); } @Test public void getObject() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcDurationVectorAccessor::getObject, (accessor, currentRow) -> is(Duration.ofDays(currentRow + 1))); } @Test public void getObjectForNull() throws Exception { int valueCount = vector.getValueCount(); for (int i = 0; i < valueCount; i++) { vector.setNull(i); } accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcDurationVectorAccessor::getObject, (accessor, currentRow) -> equalTo(null)); } @Test public void getString() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcAccessor::getString, (accessor, currentRow) -> is(Duration.ofDays(currentRow + 1).toString())); } @Test public void getStringForNull() throws Exception { int valueCount = vector.getValueCount(); for (int i = 0; i < valueCount; i++) { vector.setNull(i); } accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcAccessor::getString, (accessor, currentRow) -> equalTo(null)); } @Test public void testShouldGetObjectClass() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcAccessor::getObjectClass, (accessor, currentRow) -> equalTo(Duration.class)); } } ArrowFlightJdbcIntervalVectorAccessorTest.java000066400000000000000000000345021515451615200476330ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.calendar; import static org.apache.arrow.driver.jdbc.utils.IntervalStringUtils.formatIntervalDay; import static org.apache.arrow.driver.jdbc.utils.IntervalStringUtils.formatIntervalYear; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import static org.junit.jupiter.api.Assertions.assertEquals; import java.time.Duration; import java.time.Period; import java.time.format.DateTimeParseException; import java.util.function.Supplier; import java.util.stream.Stream; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.IntervalDayVector; import org.apache.arrow.vector.IntervalMonthDayNanoVector; import org.apache.arrow.vector.IntervalYearVector; import org.apache.arrow.vector.PeriodDuration; import org.apache.arrow.vector.ValueVector; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; public class ArrowFlightJdbcIntervalVectorAccessorTest { @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private ValueVector vector; private final AccessorTestUtils.AccessorSupplier accessorSupplier = (vector, getCurrentRow) -> { ArrowFlightJdbcAccessorFactory.WasNullConsumer noOpWasNullConsumer = (boolean wasNull) -> {}; if (vector instanceof IntervalDayVector) { return new ArrowFlightJdbcIntervalVectorAccessor( (IntervalDayVector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof IntervalYearVector) { return new ArrowFlightJdbcIntervalVectorAccessor( (IntervalYearVector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof IntervalMonthDayNanoVector) { return new ArrowFlightJdbcIntervalVectorAccessor( (IntervalMonthDayNanoVector) vector, getCurrentRow, noOpWasNullConsumer); } return null; }; final AccessorTestUtils.AccessorIterator accessorIterator = new AccessorTestUtils.AccessorIterator<>(accessorSupplier); public static Stream data() { return Stream.of( Arguments.of( (Supplier) () -> { IntervalDayVector vector = new IntervalDayVector("", rootAllocatorTestExtension.getRootAllocator()); int valueCount = 10; vector.setValueCount(valueCount); for (int i = 0; i < valueCount; i++) { vector.set(i, i + 1, (i + 1) * 1000); } return vector; }, "IntervalDayVector"), Arguments.of( (Supplier) () -> { IntervalYearVector vector = new IntervalYearVector("", rootAllocatorTestExtension.getRootAllocator()); int valueCount = 10; vector.setValueCount(valueCount); for (int i = 0; i < valueCount; i++) { vector.set(i, i + 1); } return vector; }, "IntervalYearVector"), Arguments.of( (Supplier) () -> { IntervalMonthDayNanoVector vector = new IntervalMonthDayNanoVector( "", rootAllocatorTestExtension.getRootAllocator()); int valueCount = 10; vector.setValueCount(valueCount); for (int i = 0; i < valueCount; i++) { vector.set(i, i + 1, (i + 1) * 10, (i + 1) * 100); } return vector; }, "IntervalMonthDayNanoVector")); } public void setup(Supplier vectorSupplier) { this.vector = vectorSupplier.get(); } @AfterEach public void tearDown() { if (this.vector != null) { this.vector.close(); } } @ParameterizedTest @MethodSource("data") public void testShouldGetObjectReturnValidObject(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcIntervalVectorAccessor::getObject, (accessor, currentRow) -> is(getExpectedObject(vector, currentRow))); } @ParameterizedTest @MethodSource("data") public void testShouldGetObjectPassingObjectClassAsParameterReturnValidObject( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); Class objectClass = getExpectedObjectClassForVector(vector); accessorIterator.assertAccessorGetter( vector, accessor -> accessor.getObject(objectClass), (accessor, currentRow) -> is(getExpectedObject(vector, currentRow))); } @ParameterizedTest @MethodSource("data") public void testShouldGetObjectReturnNull(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); setAllNullOnVector(vector); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcIntervalVectorAccessor::getObject, (accessor, currentRow) -> equalTo(null)); } private String getStringOnVector(ValueVector vector, int index) { Object object = getExpectedObject(vector, index); if (object == null) { return null; } else if (vector instanceof IntervalDayVector) { return formatIntervalDay(Duration.parse(object.toString())); } else if (vector instanceof IntervalYearVector) { return formatIntervalYear(Period.parse(object.toString())); } else if (vector instanceof IntervalMonthDayNanoVector) { String iso8601IntervalString = ((PeriodDuration) object).toISO8601IntervalString(); String[] periodAndDuration = iso8601IntervalString.split("T"); if (periodAndDuration.length == 1) { // If there is no 'T', then either Period or Duration is zero, and the other one will // successfully parse it String periodOrDuration = periodAndDuration[0]; try { return new PeriodDuration(Period.parse(periodOrDuration), Duration.ZERO) .toISO8601IntervalString(); } catch (DateTimeParseException e) { return new PeriodDuration(Period.ZERO, Duration.parse(periodOrDuration)) .toISO8601IntervalString(); } } else { // If there is a 'T', both Period and Duration are non-zero, and we just need to prepend the // 'PT' to the // duration for both to parse successfully Period parse = Period.parse(periodAndDuration[0]); Duration duration = Duration.parse("PT" + periodAndDuration[1]); return new PeriodDuration(parse, duration).toISO8601IntervalString(); } } return null; } @ParameterizedTest @MethodSource("data") public void testShouldGetIntervalYear(Supplier vectorSupplier) { setup(vectorSupplier); assertEquals("-002-00", formatIntervalYear(Period.parse("P-2Y"))); assertEquals("-001-01", formatIntervalYear(Period.parse("P-1Y-1M"))); assertEquals("-001-02", formatIntervalYear(Period.parse("P-1Y-2M"))); assertEquals("-002-03", formatIntervalYear(Period.parse("P-2Y-3M"))); assertEquals("-002-04", formatIntervalYear(Period.parse("P-2Y-4M"))); assertEquals("-011-01", formatIntervalYear(Period.parse("P-11Y-1M"))); assertEquals("+002-00", formatIntervalYear(Period.parse("P+2Y"))); assertEquals("+001-01", formatIntervalYear(Period.parse("P+1Y1M"))); assertEquals("+001-02", formatIntervalYear(Period.parse("P+1Y2M"))); assertEquals("+002-03", formatIntervalYear(Period.parse("P+2Y3M"))); assertEquals("+002-04", formatIntervalYear(Period.parse("P+2Y4M"))); assertEquals("+011-01", formatIntervalYear(Period.parse("P+11Y1M"))); } @ParameterizedTest @MethodSource("data") public void testShouldGetIntervalDay(Supplier vectorSupplier) { setup(vectorSupplier); assertEquals("-001 00:00:00.000", formatIntervalDay(Duration.parse("PT-24H"))); assertEquals("+001 00:00:00.000", formatIntervalDay(Duration.parse("PT+24H"))); assertEquals("-000 01:00:00.000", formatIntervalDay(Duration.parse("PT-1H"))); // "JDK-8054978: java.time.Duration.parse() fails for negative duration with 0 seconds and // nanos" not fixed on JDK8 // assertEquals("-000 01:00:00.001", formatIntervalDay(Duration.parse("PT-1H-0M-00.001S"))); assertEquals("-000 01:00:00.001", formatIntervalDay(Duration.ofHours(-1).minusMillis(1))); assertEquals("-000 01:01:01.000", formatIntervalDay(Duration.parse("PT-1H-1M-1S"))); assertEquals("-000 02:02:02.002", formatIntervalDay(Duration.parse("PT-2H-2M-02.002S"))); assertEquals("-000 23:59:59.999", formatIntervalDay(Duration.parse("PT-23H-59M-59.999S"))); // "JDK-8054978: java.time.Duration.parse() fails for negative duration with 0 seconds and // nanos" not fixed on JDK8 // assertEquals("-000 11:59:00.100", formatIntervalDay(Duration.parse("PT-11H-59M-00.100S"))); assertEquals( "-000 11:59:00.100", formatIntervalDay(Duration.ofHours(-11).minusMinutes(59).minusMillis(100))); assertEquals("-000 05:02:03.000", formatIntervalDay(Duration.parse("PT-5H-2M-3S"))); assertEquals("-000 22:22:22.222", formatIntervalDay(Duration.parse("PT-22H-22M-22.222S"))); assertEquals("+000 01:00:00.000", formatIntervalDay(Duration.parse("PT+1H"))); assertEquals("+000 01:00:00.001", formatIntervalDay(Duration.parse("PT+1H0M00.001S"))); assertEquals("+000 01:01:01.000", formatIntervalDay(Duration.parse("PT+1H1M1S"))); assertEquals("+000 02:02:02.002", formatIntervalDay(Duration.parse("PT+2H2M02.002S"))); assertEquals("+000 23:59:59.999", formatIntervalDay(Duration.parse("PT+23H59M59.999S"))); assertEquals("+000 11:59:00.100", formatIntervalDay(Duration.parse("PT+11H59M00.100S"))); assertEquals("+000 05:02:03.000", formatIntervalDay(Duration.parse("PT+5H2M3S"))); assertEquals("+000 22:22:22.222", formatIntervalDay(Duration.parse("PT+22H22M22.222S"))); } @ParameterizedTest @MethodSource("data") public void testIntervalDayWithJodaPeriodObject(Supplier vectorSupplier) { setup(vectorSupplier); assertEquals("+1567 00:00:00.000", formatIntervalDay(Duration.ofDays(1567))); assertEquals("-1567 00:00:00.000", formatIntervalDay(Duration.ofDays(-1567))); } @ParameterizedTest @MethodSource("data") public void testShouldGetStringReturnCorrectString(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcIntervalVectorAccessor::getString, (accessor, currentRow) -> is(getStringOnVector(vector, currentRow))); } @ParameterizedTest @MethodSource("data") public void testShouldGetStringReturnNull(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); setAllNullOnVector(vector); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcIntervalVectorAccessor::getString, (accessor, currentRow) -> equalTo(null)); } @ParameterizedTest @MethodSource("data") public void testShouldGetObjectClassReturnCorrectClass(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); Class expectedObjectClass = getExpectedObjectClassForVector(vector); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcIntervalVectorAccessor::getObjectClass, (accessor, currentRow) -> equalTo(expectedObjectClass)); } private Class getExpectedObjectClassForVector(ValueVector vector) { if (vector instanceof IntervalDayVector) { return Duration.class; } else if (vector instanceof IntervalYearVector) { return Period.class; } else if (vector instanceof IntervalMonthDayNanoVector) { return PeriodDuration.class; } return null; } private void setAllNullOnVector(ValueVector vector) { int valueCount = vector.getValueCount(); if (vector instanceof IntervalDayVector) { for (int i = 0; i < valueCount; i++) { ((IntervalDayVector) vector).setNull(i); } } else if (vector instanceof IntervalYearVector) { for (int i = 0; i < valueCount; i++) { ((IntervalYearVector) vector).setNull(i); } } else if (vector instanceof IntervalMonthDayNanoVector) { for (int i = 0; i < valueCount; i++) { ((IntervalMonthDayNanoVector) vector).setNull(i); } } } private Object getExpectedObject(ValueVector vector, int currentRow) { if (vector instanceof IntervalDayVector) { return Duration.ofDays(currentRow + 1).plusMillis((currentRow + 1) * 1000L); } else if (vector instanceof IntervalYearVector) { return Period.ofMonths(currentRow + 1); } else if (vector instanceof IntervalMonthDayNanoVector) { Period period = Period.ofMonths(currentRow + 1).plusDays((currentRow + 1) * 10L); Duration duration = Duration.ofNanos((currentRow + 1) * 100L); return new PeriodDuration(period, duration); } return null; } } ArrowFlightJdbcTimeStampVectorAccessorTest.java000066400000000000000000000463531515451615200477610ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.calendar; 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; import java.util.stream.Stream; import org.apache.arrow.driver.jdbc.accessor.impl.text.ArrowFlightJdbcVarCharVectorAccessor; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.TimeStampMicroVector; import org.apache.arrow.vector.TimeStampMilliVector; import org.apache.arrow.vector.TimeStampNanoVector; import org.apache.arrow.vector.TimeStampSecVector; import org.apache.arrow.vector.TimeStampVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.util.Text; import org.hamcrest.CoreMatchers; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Assumptions; import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; public class ArrowFlightJdbcTimeStampVectorAccessorTest { public static final String AMERICA_VANCOUVER = "America/Vancouver"; public static final String ASIA_BANGKOK = "Asia/Bangkok"; public static final String AMERICA_SAO_PAULO = "America/Sao_Paulo"; @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private TimeStampVector vector; private final AccessorTestUtils.AccessorSupplier accessorSupplier = (vector, getCurrentRow) -> new ArrowFlightJdbcTimeStampVectorAccessor( (TimeStampVector) vector, getCurrentRow, (boolean wasNull) -> {}); private final AccessorTestUtils.AccessorIterator accessorIterator = new AccessorTestUtils.AccessorIterator<>(accessorSupplier); public static Stream data() { return Stream.of( Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeStampNanoVector(), "TimeStampNanoVector", null), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeStampNanoTZVector("UTC"), "TimeStampNanoTZVector", "UTC"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeStampNanoTZVector(AMERICA_VANCOUVER), "TimeStampNanoTZVector", AMERICA_VANCOUVER), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeStampNanoTZVector(ASIA_BANGKOK), "TimeStampNanoTZVector", ASIA_BANGKOK), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeStampMicroVector(), "TimeStampMicroVector", null), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeStampMicroTZVector("UTC"), "TimeStampMicroTZVector", "UTC"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeStampMicroTZVector(AMERICA_VANCOUVER), "TimeStampMicroTZVector", AMERICA_VANCOUVER), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeStampMicroTZVector(ASIA_BANGKOK), "TimeStampMicroTZVector", ASIA_BANGKOK), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeStampMilliVector(), "TimeStampMilliVector", null), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeStampMilliTZVector("UTC"), "TimeStampMilliTZVector", "UTC"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeStampMilliTZVector(AMERICA_VANCOUVER), "TimeStampMilliTZVector", AMERICA_VANCOUVER), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeStampMilliTZVector(ASIA_BANGKOK), "TimeStampMilliTZVector", ASIA_BANGKOK), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeStampSecVector(), "TimeStampSecVector", null), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeStampSecTZVector("UTC"), "TimeStampSecTZVector", "UTC"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeStampSecTZVector(AMERICA_VANCOUVER), "TimeStampSecTZVector", AMERICA_VANCOUVER), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeStampSecTZVector(ASIA_BANGKOK), "TimeStampSecTZVector", ASIA_BANGKOK)); } public void setup(Supplier vectorSupplier) { this.vector = vectorSupplier.get(); } @AfterEach public void tearDown() { this.vector.close(); } @ParameterizedTest @MethodSource("data") public void testShouldGetTimestampReturnValidTimestampWithoutCalendar( Supplier vectorSupplier, String vectorType, String timeZone) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, accessor -> accessor.getTimestamp(null), (accessor, currentRow) -> is(getTimestampForVector(currentRow, timeZone))); } @ParameterizedTest @MethodSource("data") public void testShouldGetTimestampReturnValidTimestampWithCalendar( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); TimeZone timeZone = TimeZone.getTimeZone(AMERICA_SAO_PAULO); Calendar calendar = Calendar.getInstance(timeZone); TimeZone timeZoneForVector = getTimeZoneForVector(vector); accessorIterator.iterate( vector, (accessor, currentRow) -> { final Timestamp resultWithoutCalendar = accessor.getTimestamp(null); final Timestamp result = accessor.getTimestamp(calendar); long offset = (long) timeZone.getOffset(resultWithoutCalendar.getTime()) - timeZoneForVector.getOffset(resultWithoutCalendar.getTime()); assertThat(resultWithoutCalendar.getTime() - result.getTime(), is(offset)); assertThat(accessor.wasNull(), is(false)); }); } @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) { setup(vectorSupplier); vector.setNull(0); ArrowFlightJdbcTimeStampVectorAccessor accessor = accessorSupplier.supply(vector, () -> 0); assertThat(accessor.getTimestamp(null), CoreMatchers.equalTo(null)); assertThat(accessor.wasNull(), is(true)); } @ParameterizedTest @MethodSource("data") public void testShouldGetDateReturnValidDateWithoutCalendar( Supplier vectorSupplier, String vectorType, String timeZone) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, accessor -> accessor.getDate(null), (accessor, currentRow) -> is(new Date(getTimestampForVector(currentRow, timeZone).getTime()))); } @ParameterizedTest @MethodSource("data") public void testShouldGetDateReturnValidDateWithCalendar(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); TimeZone timeZone = TimeZone.getTimeZone(AMERICA_SAO_PAULO); Calendar calendar = Calendar.getInstance(timeZone); TimeZone timeZoneForVector = getTimeZoneForVector(vector); accessorIterator.iterate( vector, (accessor, currentRow) -> { final Date resultWithoutCalendar = accessor.getDate(null); final Date result = accessor.getDate(calendar); long offset = (long) timeZone.getOffset(resultWithoutCalendar.getTime()) - timeZoneForVector.getOffset(resultWithoutCalendar.getTime()); assertThat(resultWithoutCalendar.getTime() - result.getTime(), is(offset)); assertThat(accessor.wasNull(), is(false)); }); } @ParameterizedTest @MethodSource("data") public void testShouldGetDateReturnNull(Supplier vectorSupplier) { setup(vectorSupplier); vector.setNull(0); ArrowFlightJdbcTimeStampVectorAccessor accessor = accessorSupplier.supply(vector, () -> 0); assertThat(accessor.getDate(null), CoreMatchers.equalTo(null)); assertThat(accessor.wasNull(), is(true)); } @ParameterizedTest @MethodSource("data") public void testShouldGetTimeReturnValidTimeWithoutCalendar( Supplier vectorSupplier, String vectorType, String timeZone) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, accessor -> accessor.getTime(null), (accessor, currentRow) -> is(new Time(getTimestampForVector(currentRow, timeZone).getTime()))); } @ParameterizedTest @MethodSource("data") public void testShouldGetTimeReturnValidTimeWithCalendar(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); TimeZone timeZone = TimeZone.getTimeZone(AMERICA_SAO_PAULO); Calendar calendar = Calendar.getInstance(timeZone); TimeZone timeZoneForVector = getTimeZoneForVector(vector); accessorIterator.iterate( vector, (accessor, currentRow) -> { final Time resultWithoutCalendar = accessor.getTime(null); final Time result = accessor.getTime(calendar); long offset = (long) timeZone.getOffset(resultWithoutCalendar.getTime()) - timeZoneForVector.getOffset(resultWithoutCalendar.getTime()); assertThat(resultWithoutCalendar.getTime() - result.getTime(), is(offset)); assertThat(accessor.wasNull(), is(false)); }); } @ParameterizedTest @MethodSource("data") public void testShouldGetTimeReturnNull(Supplier vectorSupplier) { setup(vectorSupplier); vector.setNull(0); ArrowFlightJdbcTimeStampVectorAccessor accessor = accessorSupplier.supply(vector, () -> 0); assertThat(accessor.getTime(null), CoreMatchers.equalTo(null)); assertThat(accessor.wasNull(), is(true)); } private Timestamp getTimestampForVector(int currentRow, String timeZone) { Object object = vector.getObject(currentRow); Timestamp expectedTimestamp = null; if (object instanceof LocalDateTime) { expectedTimestamp = Timestamp.valueOf((LocalDateTime) object); } else if (object instanceof Long) { TimeUnit timeUnit = getTimeUnitForVector(vector); long millis = timeUnit.toMillis((Long) object); long offset = TimeZone.getTimeZone(timeZone).getOffset(millis); expectedTimestamp = new Timestamp(millis + offset); } 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 { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcTimeStampVectorAccessor::getObjectClass, equalTo(Timestamp.class)); } @ParameterizedTest @MethodSource("data") public void testShouldGetStringBeConsistentWithVarCharAccessorWithoutCalendar( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); assertGetStringIsConsistentWithVarCharAccessor(null); } @ParameterizedTest @MethodSource("data") public void testShouldGetStringBeConsistentWithVarCharAccessorWithCalendar( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); // Ignore for TimeStamp vectors with TZ, as VarChar accessor won't consider their TZ Assumptions.assumeTrue( vector instanceof TimeStampNanoVector || vector instanceof TimeStampMicroVector || vector instanceof TimeStampMilliVector || vector instanceof TimeStampSecVector); Calendar calendar = Calendar.getInstance(TimeZone.getTimeZone(AMERICA_VANCOUVER)); assertGetStringIsConsistentWithVarCharAccessor(calendar); } private void assertGetStringIsConsistentWithVarCharAccessor(Calendar calendar) throws Exception { try (VarCharVector varCharVector = new VarCharVector("", rootAllocatorTestExtension.getRootAllocator())) { varCharVector.allocateNew(1); ArrowFlightJdbcVarCharVectorAccessor varCharVectorAccessor = new ArrowFlightJdbcVarCharVectorAccessor(varCharVector, () -> 0, (boolean wasNull) -> {}); accessorIterator.iterate( vector, (accessor, currentRow) -> { final String string = accessor.getString(); varCharVector.set(0, new Text(string)); varCharVector.setValueCount(1); Timestamp timestampFromVarChar = varCharVectorAccessor.getTimestamp(calendar); Timestamp timestamp = accessor.getTimestamp(calendar); assertThat(timestamp, is(timestampFromVarChar)); assertThat(accessor.wasNull(), is(false)); }); } } } ArrowFlightJdbcTimeVectorAccessorTest.java000066400000000000000000000264341515451615200467520ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.calendar; import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeVectorAccessor.getTimeUnitForVector; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.CoreMatchers.not; import static org.hamcrest.MatcherAssert.assertThat; import java.sql.Time; import java.sql.Timestamp; import java.time.LocalDateTime; import java.util.Calendar; import java.util.TimeZone; import java.util.concurrent.TimeUnit; import java.util.function.Supplier; import java.util.stream.Stream; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.driver.jdbc.accessor.impl.text.ArrowFlightJdbcVarCharVectorAccessor; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.BaseFixedWidthVector; 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.VarCharVector; import org.apache.arrow.vector.util.Text; import org.hamcrest.CoreMatchers; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; public class ArrowFlightJdbcTimeVectorAccessorTest { public static final String AMERICA_VANCOUVER = "America/Vancouver"; @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private BaseFixedWidthVector vector; private final AccessorTestUtils.AccessorSupplier accessorSupplier = (vector, getCurrentRow) -> { ArrowFlightJdbcAccessorFactory.WasNullConsumer noOpWasNullConsumer = (boolean wasNull) -> {}; if (vector instanceof TimeNanoVector) { return new ArrowFlightJdbcTimeVectorAccessor( (TimeNanoVector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof TimeMicroVector) { return new ArrowFlightJdbcTimeVectorAccessor( (TimeMicroVector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof TimeMilliVector) { return new ArrowFlightJdbcTimeVectorAccessor( (TimeMilliVector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof TimeSecVector) { return new ArrowFlightJdbcTimeVectorAccessor( (TimeSecVector) vector, getCurrentRow, noOpWasNullConsumer); } return null; }; private final AccessorTestUtils.AccessorIterator accessorIterator = new AccessorTestUtils.AccessorIterator<>(accessorSupplier); public static Stream data() { return Stream.of( Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeNanoVector(), "TimeNanoVector"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeMicroVector(), "TimeMicroVector"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeMilliVector(), "TimeMilliVector"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTimeSecVector(), "TimeSecVector")); } public void setup(Supplier vectorSupplier) { this.vector = vectorSupplier.get(); } @AfterEach public void tearDown() { this.vector.close(); } @ParameterizedTest @MethodSource("data") public void testShouldGetTimestampReturnValidTimestampWithoutCalendar( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, accessor -> accessor.getTimestamp(null), (accessor, currentRow) -> is(getTimestampForVector(currentRow))); } @ParameterizedTest @MethodSource("data") public void testShouldGetTimestampReturnValidTimestampWithCalendar( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); TimeZone timeZone = TimeZone.getTimeZone(AMERICA_VANCOUVER); Calendar calendar = Calendar.getInstance(timeZone); accessorIterator.iterate( vector, (accessor, currentRow) -> { final Timestamp resultWithoutCalendar = accessor.getTimestamp(null); final Timestamp result = accessor.getTimestamp(calendar); long offset = timeZone.getOffset(resultWithoutCalendar.getTime()); assertThat(resultWithoutCalendar.getTime() - result.getTime(), is(offset)); assertThat(accessor.wasNull(), is(false)); }); } @ParameterizedTest @MethodSource("data") public void testShouldGetTimestampReturnNull(Supplier vectorSupplier) { setup(vectorSupplier); vector.setNull(0); ArrowFlightJdbcTimeVectorAccessor accessor = accessorSupplier.supply(vector, () -> 0); assertThat(accessor.getTimestamp(null), CoreMatchers.equalTo(null)); assertThat(accessor.wasNull(), is(true)); } @ParameterizedTest @MethodSource("data") public void testShouldGetTimeReturnValidTimeWithoutCalendar( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, accessor -> accessor.getTime(null), (accessor, currentRow) -> { Timestamp expectedTimestamp = getTimestampForVector(currentRow); return is(new Time(expectedTimestamp.getTime())); }); } @ParameterizedTest @MethodSource("data") public void testShouldGetTimeReturnValidTimeWithCalendar( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); TimeZone timeZone = TimeZone.getTimeZone(AMERICA_VANCOUVER); Calendar calendar = Calendar.getInstance(timeZone); accessorIterator.iterate( vector, (accessor, currentRow) -> { final Time resultWithoutCalendar = accessor.getTime(null); final Time result = accessor.getTime(calendar); long offset = timeZone.getOffset(resultWithoutCalendar.getTime()); assertThat(resultWithoutCalendar.getTime() - result.getTime(), is(offset)); assertThat(accessor.wasNull(), is(false)); }); } @ParameterizedTest @MethodSource("data") public void testShouldGetTimeReturnNull(Supplier vectorSupplier) { setup(vectorSupplier); vector.setNull(0); ArrowFlightJdbcTimeVectorAccessor accessor = accessorSupplier.supply(vector, () -> 0); assertThat(accessor.getTime(null), CoreMatchers.equalTo(null)); assertThat(accessor.wasNull(), is(true)); } private Timestamp getTimestampForVector(int currentRow) { Object object = vector.getObject(currentRow); Timestamp expectedTimestamp = null; if (object instanceof LocalDateTime) { expectedTimestamp = Timestamp.valueOf((LocalDateTime) object); } else if (object instanceof Number) { long value = ((Number) object).longValue(); TimeUnit timeUnit = getTimeUnitForVector(vector); long millis = timeUnit.toMillis(value); expectedTimestamp = new Timestamp(millis); } return expectedTimestamp; } @ParameterizedTest @MethodSource("data") public void testShouldGetObjectClass(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcTimeVectorAccessor::getObjectClass, equalTo(Time.class)); } @ParameterizedTest @MethodSource("data") public void testShouldGetStringBeConsistentWithVarCharAccessorWithoutCalendar( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); assertGetStringIsConsistentWithVarCharAccessor(null); } @ParameterizedTest @MethodSource("data") public void testShouldGetStringBeConsistentWithVarCharAccessorWithCalendar( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); Calendar calendar = Calendar.getInstance(TimeZone.getTimeZone(AMERICA_VANCOUVER)); assertGetStringIsConsistentWithVarCharAccessor(calendar); } @ParameterizedTest @MethodSource("data") public void testValidateGetStringTimeZoneConsistency( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.iterate( vector, (accessor, currentRow) -> { final TimeZone defaultTz = TimeZone.getDefault(); try { final String string = accessor.getString(); // Should always be UTC as no calendar is provided // Validate with UTC Time time = accessor.getTime(null); TimeZone.setDefault(TimeZone.getTimeZone("UTC")); assertThat(time.toString(), is(string)); // Validate with different TZ TimeZone.setDefault(TimeZone.getTimeZone(AMERICA_VANCOUVER)); assertThat(time.toString(), not(string)); assertThat(accessor.wasNull(), is(false)); } finally { // Set default Tz back TimeZone.setDefault(defaultTz); } }); } private void assertGetStringIsConsistentWithVarCharAccessor(Calendar calendar) throws Exception { try (VarCharVector varCharVector = new VarCharVector("", rootAllocatorTestExtension.getRootAllocator())) { varCharVector.allocateNew(1); ArrowFlightJdbcVarCharVectorAccessor varCharVectorAccessor = new ArrowFlightJdbcVarCharVectorAccessor(varCharVector, () -> 0, (boolean wasNull) -> {}); accessorIterator.iterate( vector, (accessor, currentRow) -> { final String string = accessor.getString(); varCharVector.set(0, new Text(string)); varCharVector.setValueCount(1); Time timeFromVarChar = varCharVectorAccessor.getTime(calendar); Time time = accessor.getTime(calendar); assertThat(time, is(timeFromVarChar)); assertThat(accessor.wasNull(), is(false)); }); } } } complex/000077500000000000000000000000001515451615200364345ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/implAbstractArrowFlightJdbcListAccessorTest.java000066400000000000000000000166351515451615200471700ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; import java.sql.Array; import java.sql.ResultSet; import java.util.Arrays; import java.util.List; import java.util.function.Supplier; import java.util.stream.Stream; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.LargeListVector; import org.apache.arrow.vector.complex.ListVector; import org.hamcrest.CoreMatchers; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; public class AbstractArrowFlightJdbcListAccessorTest { @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private ValueVector vector; private final AccessorTestUtils.AccessorSupplier accessorSupplier = (vector, getCurrentRow) -> { ArrowFlightJdbcAccessorFactory.WasNullConsumer noOpWasNullConsumer = (boolean wasNull) -> {}; if (vector instanceof ListVector) { return new ArrowFlightJdbcListVectorAccessor( (ListVector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof LargeListVector) { return new ArrowFlightJdbcLargeListVectorAccessor( (LargeListVector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof FixedSizeListVector) { return new ArrowFlightJdbcFixedSizeListVectorAccessor( (FixedSizeListVector) vector, getCurrentRow, noOpWasNullConsumer); } return null; }; final AccessorTestUtils.AccessorIterator accessorIterator = new AccessorTestUtils.AccessorIterator<>(accessorSupplier); public static Stream data() { return Stream.of( Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createListVector(), "ListVector"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createLargeListVector(), "LargeListVector"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createFixedSizeListVector(), "FixedSizeListVector")); } public void setup(Supplier vectorSupplier) { this.vector = vectorSupplier.get(); } @AfterEach public void tearDown() { this.vector.close(); } @ParameterizedTest @MethodSource("data") public void testShouldGetObjectClassReturnCorrectClass(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, AbstractArrowFlightJdbcListVectorAccessor::getObjectClass, (accessor, currentRow) -> equalTo(List.class)); } @ParameterizedTest @MethodSource("data") public void testShouldGetObjectReturnValidList(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, AbstractArrowFlightJdbcListVectorAccessor::getObject, (accessor, currentRow) -> equalTo(Arrays.asList(0, currentRow, currentRow * 2, currentRow * 3, currentRow * 4))); } @ParameterizedTest @MethodSource("data") public void testShouldGetObjectReturnNull(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); vector.clear(); vector.allocateNewSafe(); vector.setValueCount(5); accessorIterator.assertAccessorGetter( vector, AbstractArrowFlightJdbcListVectorAccessor::getObject, (accessor, currentRow) -> CoreMatchers.nullValue()); } @ParameterizedTest @MethodSource("data") public void testShouldGetArrayReturnValidArray(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.iterate( vector, (accessor, currentRow) -> { Array array = accessor.getArray(); assert array != null; Object[] arrayObject = (Object[]) array.getArray(); assertThat( arrayObject, equalTo( new Object[] {0, currentRow, currentRow * 2, currentRow * 3, currentRow * 4})); }); } @ParameterizedTest @MethodSource("data") public void testShouldGetArrayReturnNull(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); vector.clear(); vector.allocateNewSafe(); vector.setValueCount(5); accessorIterator.assertAccessorGetter( vector, AbstractArrowFlightJdbcListVectorAccessor::getArray, CoreMatchers.nullValue()); } @ParameterizedTest @MethodSource("data") public void testShouldGetArrayReturnValidArrayPassingOffsets(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.iterate( vector, (accessor, currentRow) -> { Array array = accessor.getArray(); assert array != null; Object[] arrayObject = (Object[]) array.getArray(1, 3); assertThat( arrayObject, equalTo(new Object[] {currentRow, currentRow * 2, currentRow * 3})); }); } @ParameterizedTest @MethodSource("data") public void testShouldGetArrayGetResultSetReturnValidResultSet( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.iterate( vector, (accessor, currentRow) -> { Array array = accessor.getArray(); assert array != null; try (ResultSet rs = array.getResultSet()) { int count = 0; while (rs.next()) { // 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++; } assertThat(count, equalTo(5)); } }); } } AbstractArrowFlightJdbcUnionVectorAccessorTest.java000066400000000000000000000164671515451615200505330ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import java.sql.SQLException; import java.util.Calendar; import java.util.Map; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.ArrowFlightJdbcNullVectorAccessor; import org.apache.arrow.vector.NullVector; import org.apache.arrow.vector.ValueVector; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.Spy; import org.mockito.junit.jupiter.MockitoExtension; @ExtendWith(MockitoExtension.class) public class AbstractArrowFlightJdbcUnionVectorAccessorTest { @Mock ArrowFlightJdbcAccessor innerAccessor; @Spy AbstractArrowFlightJdbcUnionVectorAccessorMock accessor; @BeforeEach public void setup() { when(accessor.getAccessor()).thenReturn(innerAccessor); } @Test public void testGetNCharacterStreamUsesSpecificAccessor() throws SQLException { accessor.getNCharacterStream(); verify(innerAccessor).getNCharacterStream(); } @Test public void testGetNStringUsesSpecificAccessor() throws SQLException { accessor.getNString(); verify(innerAccessor).getNString(); } @Test public void testGetSQLXMLUsesSpecificAccessor() throws SQLException { accessor.getSQLXML(); verify(innerAccessor).getSQLXML(); } @Test public void testGetNClobUsesSpecificAccessor() throws SQLException { accessor.getNClob(); verify(innerAccessor).getNClob(); } @Test public void testGetURLUsesSpecificAccessor() throws SQLException { accessor.getURL(); verify(innerAccessor).getURL(); } @Test public void testGetStructUsesSpecificAccessor() throws SQLException { accessor.getStruct(); verify(innerAccessor).getStruct(); } @Test public void testGetArrayUsesSpecificAccessor() throws SQLException { accessor.getArray(); verify(innerAccessor).getArray(); } @Test public void testGetClobUsesSpecificAccessor() throws SQLException { accessor.getClob(); verify(innerAccessor).getClob(); } @Test public void testGetBlobUsesSpecificAccessor() throws SQLException { accessor.getBlob(); verify(innerAccessor).getBlob(); } @Test public void testGetRefUsesSpecificAccessor() throws SQLException { accessor.getRef(); verify(innerAccessor).getRef(); } @Test public void testGetCharacterStreamUsesSpecificAccessor() throws SQLException { accessor.getCharacterStream(); verify(innerAccessor).getCharacterStream(); } @Test public void testGetBinaryStreamUsesSpecificAccessor() throws SQLException { accessor.getBinaryStream(); verify(innerAccessor).getBinaryStream(); } @Test public void testGetUnicodeStreamUsesSpecificAccessor() throws SQLException { accessor.getUnicodeStream(); verify(innerAccessor).getUnicodeStream(); } @Test public void testGetAsciiStreamUsesSpecificAccessor() throws SQLException { accessor.getAsciiStream(); verify(innerAccessor).getAsciiStream(); } @Test public void testGetBytesUsesSpecificAccessor() throws SQLException { accessor.getBytes(); verify(innerAccessor).getBytes(); } @Test public void testGetBigDecimalUsesSpecificAccessor() throws SQLException { accessor.getBigDecimal(); verify(innerAccessor).getBigDecimal(); } @Test public void testGetDoubleUsesSpecificAccessor() throws SQLException { accessor.getDouble(); verify(innerAccessor).getDouble(); } @Test public void testGetFloatUsesSpecificAccessor() throws SQLException { accessor.getFloat(); verify(innerAccessor).getFloat(); } @Test public void testGetLongUsesSpecificAccessor() throws SQLException { accessor.getLong(); verify(innerAccessor).getLong(); } @Test public void testGetIntUsesSpecificAccessor() throws SQLException { accessor.getInt(); verify(innerAccessor).getInt(); } @Test public void testGetShortUsesSpecificAccessor() throws SQLException { accessor.getShort(); verify(innerAccessor).getShort(); } @Test public void testGetByteUsesSpecificAccessor() throws SQLException { accessor.getByte(); verify(innerAccessor).getByte(); } @Test public void testGetBooleanUsesSpecificAccessor() throws SQLException { accessor.getBoolean(); verify(innerAccessor).getBoolean(); } @Test public void testGetStringUsesSpecificAccessor() throws SQLException { accessor.getString(); verify(innerAccessor).getString(); } @Test public void testGetObjectClassUsesSpecificAccessor() { accessor.getObjectClass(); verify(innerAccessor).getObjectClass(); } @Test public void testGetObjectWithClassUsesSpecificAccessor() throws SQLException { accessor.getObject(Object.class); verify(innerAccessor).getObject(Object.class); } @Test public void testGetTimestampUsesSpecificAccessor() throws SQLException { Calendar calendar = Calendar.getInstance(); accessor.getTimestamp(calendar); verify(innerAccessor).getTimestamp(calendar); } @Test public void testGetTimeUsesSpecificAccessor() throws SQLException { Calendar calendar = Calendar.getInstance(); accessor.getTime(calendar); verify(innerAccessor).getTime(calendar); } @Test public void testGetDateUsesSpecificAccessor() throws SQLException { Calendar calendar = Calendar.getInstance(); accessor.getDate(calendar); verify(innerAccessor).getDate(calendar); } @Test public void testGetObjectUsesSpecificAccessor() throws SQLException { Map> map = mock(Map.class); accessor.getObject(map); verify(innerAccessor).getObject(map); } @Test public void testGetBigDecimalWithScaleUsesSpecificAccessor() throws SQLException { accessor.getBigDecimal(2); verify(innerAccessor).getBigDecimal(2); } private static class AbstractArrowFlightJdbcUnionVectorAccessorMock extends AbstractArrowFlightJdbcUnionVectorAccessor { protected AbstractArrowFlightJdbcUnionVectorAccessorMock() { super(() -> 0, (boolean wasNull) -> {}); } @Override protected ArrowFlightJdbcAccessor createAccessorForVector(ValueVector vector) { return new ArrowFlightJdbcNullVectorAccessor((boolean wasNull) -> {}); } @Override protected byte getCurrentTypeId() { return 0; } @Override protected ValueVector getVectorByTypeId(byte typeId) { return new NullVector(); } } } ArrowFlightJdbcDenseUnionVectorAccessorTest.java000066400000000000000000000111061515451615200500070ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; import java.sql.Timestamp; import java.util.Arrays; import java.util.List; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.holders.NullableBigIntHolder; import org.apache.arrow.vector.holders.NullableFloat8Holder; import org.apache.arrow.vector.holders.NullableTimeStampMilliHolder; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.Field; 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; public class ArrowFlightJdbcDenseUnionVectorAccessorTest { @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private DenseUnionVector vector; private final AccessorTestUtils.AccessorSupplier accessorSupplier = (vector, getCurrentRow) -> new ArrowFlightJdbcDenseUnionVectorAccessor( (DenseUnionVector) vector, getCurrentRow, (boolean wasNull) -> { // No Operation }); private final AccessorTestUtils.AccessorIterator accessorIterator = new AccessorTestUtils.AccessorIterator<>(accessorSupplier); @BeforeEach public void setup() throws Exception { this.vector = DenseUnionVector.empty("", rootAllocatorTestExtension.getRootAllocator()); this.vector.allocateNew(); // write some data byte bigIntTypeId = this.vector.registerNewTypeId(Field.nullable("", Types.MinorType.BIGINT.getType())); byte float8TypeId = this.vector.registerNewTypeId(Field.nullable("", Types.MinorType.FLOAT8.getType())); byte timestampMilliTypeId = this.vector.registerNewTypeId(Field.nullable("", Types.MinorType.TIMESTAMPMILLI.getType())); NullableBigIntHolder nullableBigIntHolder = new NullableBigIntHolder(); nullableBigIntHolder.isSet = 1; nullableBigIntHolder.value = Long.MAX_VALUE; this.vector.setTypeId(0, bigIntTypeId); this.vector.setSafe(0, nullableBigIntHolder); NullableFloat8Holder nullableFloat4Holder = new NullableFloat8Holder(); nullableFloat4Holder.isSet = 1; nullableFloat4Holder.value = Math.PI; this.vector.setTypeId(1, float8TypeId); this.vector.setSafe(1, nullableFloat4Holder); NullableTimeStampMilliHolder nullableTimeStampMilliHolder = new NullableTimeStampMilliHolder(); nullableTimeStampMilliHolder.isSet = 1; nullableTimeStampMilliHolder.value = 1625702400000L; this.vector.setTypeId(2, timestampMilliTypeId); this.vector.setSafe(2, nullableTimeStampMilliHolder); nullableBigIntHolder.isSet = 0; this.vector.setTypeId(3, bigIntTypeId); this.vector.setSafe(3, nullableBigIntHolder); this.vector.setValueCount(5); } @AfterEach public void tearDown() { this.vector.close(); } @Test public void getObject() throws Exception { List result = accessorIterator.toList(vector); List expected = Arrays.asList(Long.MAX_VALUE, Math.PI, new Timestamp(1625702400000L), null, null); assertThat(result, is(expected)); } @Test public void getObjectForNull() throws Exception { vector.reset(); vector.setValueCount(5); accessorIterator.assertAccessorGetter( vector, AbstractArrowFlightJdbcUnionVectorAccessor::getObject, equalTo(null)); } } ArrowFlightJdbcMapVectorAccessorTest.java000066400000000000000000000165511515451615200464660ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex; 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.assertTrue; import java.sql.Array; import java.sql.ResultSet; import java.sql.SQLException; import java.util.Map; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.impl.UnionMapWriter; import org.apache.arrow.vector.util.JsonStringHashMap; 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; public class ArrowFlightJdbcMapVectorAccessorTest { @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private MapVector vector; @BeforeEach public void setup() { vector = MapVector.empty("", rootAllocatorTestExtension.getRootAllocator(), false); UnionMapWriter writer = vector.getWriter(); writer.allocate(); writer.setPosition(0); // optional writer.startMap(); writer.startEntry(); writer.key().integer().writeInt(1); writer.value().integer().writeInt(11); writer.endEntry(); writer.startEntry(); writer.key().integer().writeInt(2); writer.value().integer().writeInt(22); writer.endEntry(); writer.startEntry(); writer.key().integer().writeInt(3); writer.value().integer().writeInt(33); writer.endEntry(); writer.endMap(); writer.setPosition(1); writer.startMap(); writer.startEntry(); writer.key().integer().writeInt(2); writer.endEntry(); writer.endMap(); writer.setPosition(2); writer.startMap(); writer.startEntry(); writer.key().integer().writeInt(0); writer.value().integer().writeInt(2000); writer.endEntry(); writer.startEntry(); writer.key().integer().writeInt(1); writer.value().integer().writeInt(2001); writer.endEntry(); writer.startEntry(); writer.key().integer().writeInt(2); writer.value().integer().writeInt(2002); writer.endEntry(); writer.startEntry(); writer.key().integer().writeInt(3); writer.value().integer().writeInt(2003); writer.endEntry(); writer.endMap(); writer.setValueCount(3); } @AfterEach public void tearDown() { vector.close(); } @Test public void testShouldGetObjectReturnValidMap() { AccessorTestUtils.Cursor cursor = new AccessorTestUtils.Cursor(vector.getValueCount()); ArrowFlightJdbcMapVectorAccessor accessor = new ArrowFlightJdbcMapVectorAccessor( vector, cursor::getCurrentRow, (boolean wasNull) -> {}); Map expected = new JsonStringHashMap<>(); expected.put(1, 11); expected.put(2, 22); expected.put(3, 33); assertEquals(expected, accessor.getObject()); assertFalse(accessor.wasNull()); cursor.next(); expected = new JsonStringHashMap<>(); expected.put(2, null); assertEquals(expected, accessor.getObject()); assertFalse(accessor.wasNull()); cursor.next(); expected = new JsonStringHashMap<>(); expected.put(0, 2000); expected.put(1, 2001); expected.put(2, 2002); expected.put(3, 2003); assertEquals(expected, accessor.getObject()); assertFalse(accessor.wasNull()); } @Test public void testShouldGetObjectReturnNull() { vector.setNull(0); ArrowFlightJdbcMapVectorAccessor accessor = new ArrowFlightJdbcMapVectorAccessor(vector, () -> 0, (boolean wasNull) -> {}); assertNull(accessor.getObject()); assertTrue(accessor.wasNull()); } @Test public void testShouldGetArrayReturnValidArray() throws SQLException { AccessorTestUtils.Cursor cursor = new AccessorTestUtils.Cursor(vector.getValueCount()); ArrowFlightJdbcMapVectorAccessor accessor = new ArrowFlightJdbcMapVectorAccessor( vector, cursor::getCurrentRow, (boolean wasNull) -> {}); Array array = accessor.getArray(); assertNotNull(array); assertFalse(accessor.wasNull()); try (ResultSet resultSet = array.getResultSet()) { assertTrue(resultSet.next()); Map entry = resultSet.getObject(2, Map.class); assertEquals(1, entry.get("key")); assertEquals(11, entry.get("value")); assertTrue(resultSet.next()); entry = resultSet.getObject(2, Map.class); assertEquals(2, entry.get("key")); assertEquals(22, entry.get("value")); assertTrue(resultSet.next()); entry = resultSet.getObject(2, Map.class); assertEquals(3, entry.get("key")); assertEquals(33, entry.get("value")); assertFalse(resultSet.next()); } cursor.next(); array = accessor.getArray(); assertNotNull(array); assertFalse(accessor.wasNull()); try (ResultSet resultSet = array.getResultSet()) { assertTrue(resultSet.next()); Map entry = resultSet.getObject(2, Map.class); assertEquals(2, entry.get("key")); assertNull(entry.get("value")); assertFalse(resultSet.next()); } cursor.next(); array = accessor.getArray(); assertNotNull(array); assertFalse(accessor.wasNull()); try (ResultSet resultSet = array.getResultSet()) { assertTrue(resultSet.next()); Map entry = resultSet.getObject(2, Map.class); assertEquals(0, entry.get("key")); assertEquals(2000, entry.get("value")); assertTrue(resultSet.next()); entry = resultSet.getObject(2, Map.class); assertEquals(1, entry.get("key")); assertEquals(2001, entry.get("value")); assertTrue(resultSet.next()); entry = resultSet.getObject(2, Map.class); assertEquals(2, entry.get("key")); assertEquals(2002, entry.get("value")); assertTrue(resultSet.next()); entry = resultSet.getObject(2, Map.class); assertEquals(3, entry.get("key")); assertEquals(2003, entry.get("value")); assertFalse(resultSet.next()); } } @Test public void testShouldGetArrayReturnNull() { vector.setNull(0); ((StructVector) vector.getDataVector()).setNull(0); ArrowFlightJdbcMapVectorAccessor accessor = new ArrowFlightJdbcMapVectorAccessor(vector, () -> 0, (boolean wasNull) -> {}); assertNull(accessor.getArray()); assertTrue(accessor.wasNull()); } } ArrowFlightJdbcStructVectorAccessorTest.java000066400000000000000000000171741515451615200472370ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.nullValue; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.assertEquals; import java.sql.SQLException; import java.sql.Struct; import java.util.HashMap; import java.util.Map; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.complex.impl.UnionListWriter; import org.apache.arrow.vector.holders.NullableBitHolder; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.JsonStringArrayList; import org.apache.arrow.vector.util.JsonStringHashMap; 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; public class ArrowFlightJdbcStructVectorAccessorTest { @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private StructVector vector; private final AccessorTestUtils.AccessorSupplier accessorSupplier = (vector, getCurrentRow) -> new ArrowFlightJdbcStructVectorAccessor( (StructVector) vector, getCurrentRow, (boolean wasNull) -> {}); private final AccessorTestUtils.AccessorIterator accessorIterator = new AccessorTestUtils.AccessorIterator<>(accessorSupplier); @BeforeEach public void setUp() throws Exception { Map metadata = new HashMap<>(); metadata.put("k1", "v1"); FieldType type = new FieldType(true, ArrowType.Struct.INSTANCE, null, metadata); vector = new StructVector("", rootAllocatorTestExtension.getRootAllocator(), type, null); vector.allocateNew(); IntVector intVector = vector.addOrGet("int", FieldType.nullable(Types.MinorType.INT.getType()), IntVector.class); Float8Vector float8Vector = vector.addOrGet( "float8", FieldType.nullable(Types.MinorType.FLOAT8.getType()), Float8Vector.class); intVector.setSafe(0, 100); float8Vector.setSafe(0, 100.05); vector.setIndexDefined(0); intVector.setSafe(1, 200); float8Vector.setSafe(1, 200.1); vector.setIndexDefined(1); vector.setValueCount(2); } @AfterEach public void tearDown() throws Exception { vector.close(); } @Test public void testShouldGetObjectClassReturnMapClass() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcStructVectorAccessor::getObjectClass, (accessor, currentRow) -> equalTo(Map.class)); } @Test public void testShouldGetObjectReturnValidMap() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcStructVectorAccessor::getObject, (accessor, currentRow) -> { Map expected = new HashMap<>(); expected.put("int", 100 * (currentRow + 1)); expected.put("float8", 100.05 * (currentRow + 1)); return equalTo(expected); }); } @Test public void testShouldGetObjectReturnNull() throws Exception { vector.setNull(0); vector.setNull(1); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcStructVectorAccessor::getObject, (accessor, currentRow) -> nullValue()); } @Test public void testShouldGetStructReturnValidStruct() throws Exception { accessorIterator.iterate( vector, (accessor, currentRow) -> { Struct struct = accessor.getStruct(); assert struct != null; Object[] expected = new Object[] {100 * (currentRow + 1), 100.05 * (currentRow + 1)}; assertThat(struct.getAttributes(), equalTo(expected)); }); } @Test public void testShouldGetStructReturnNull() throws Exception { vector.setNull(0); vector.setNull(1); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcStructVectorAccessor::getStruct, (accessor, currentRow) -> nullValue()); } @Test public void testShouldGetObjectWorkWithNestedComplexData() throws SQLException { try (StructVector rootVector = StructVector.empty("", rootAllocatorTestExtension.getRootAllocator())) { StructVector structVector = rootVector.addOrGetStruct("struct"); FieldType intFieldType = FieldType.nullable(Types.MinorType.INT.getType()); IntVector intVector = structVector.addOrGet("int", intFieldType, IntVector.class); FieldType float8FieldType = FieldType.nullable(Types.MinorType.FLOAT8.getType()); Float8Vector float8Vector = structVector.addOrGet("float8", float8FieldType, Float8Vector.class); ListVector listVector = rootVector.addOrGetList("list"); UnionListWriter listWriter = listVector.getWriter(); listWriter.allocate(); UnionVector unionVector = rootVector.addOrGetUnion("union"); intVector.setSafe(0, 100); intVector.setValueCount(1); float8Vector.setSafe(0, 100.05); float8Vector.setValueCount(1); structVector.setIndexDefined(0); listWriter.setPosition(0); listWriter.startList(); listWriter.bigInt().writeBigInt(Long.MAX_VALUE); listWriter.bigInt().writeBigInt(Long.MIN_VALUE); listWriter.endList(); listVector.setValueCount(1); unionVector.setType(0, Types.MinorType.BIT); NullableBitHolder holder = new NullableBitHolder(); holder.isSet = 1; holder.value = 1; unionVector.setSafe(0, holder); unionVector.setValueCount(1); rootVector.setIndexDefined(0); rootVector.setValueCount(1); Map expected = new JsonStringHashMap<>(); Map nestedStruct = new JsonStringHashMap<>(); nestedStruct.put("int", 100); nestedStruct.put("float8", 100.05); expected.put("struct", nestedStruct); JsonStringArrayList nestedList = new JsonStringArrayList<>(); nestedList.add(Long.MAX_VALUE); nestedList.add(Long.MIN_VALUE); expected.put("list", nestedList); expected.put("union", true); ArrowFlightJdbcStructVectorAccessor accessor = new ArrowFlightJdbcStructVectorAccessor(rootVector, () -> 0, (boolean wasNull) -> {}); assertEquals(expected, accessor.getObject()); assertEquals(expected.toString(), accessor.getString()); } } } ArrowFlightJdbcUnionVectorAccessorTest.java000066400000000000000000000100331515451615200470260ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; import java.sql.Timestamp; import java.util.Arrays; import java.util.List; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.holders.NullableBigIntHolder; import org.apache.arrow.vector.holders.NullableFloat8Holder; import org.apache.arrow.vector.holders.NullableTimeStampMilliHolder; import org.apache.arrow.vector.types.Types; 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; public class ArrowFlightJdbcUnionVectorAccessorTest { @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private UnionVector vector; private final AccessorTestUtils.AccessorSupplier accessorSupplier = (vector, getCurrentRow) -> new ArrowFlightJdbcUnionVectorAccessor( (UnionVector) vector, getCurrentRow, (boolean wasNull) -> {}); private final AccessorTestUtils.AccessorIterator accessorIterator = new AccessorTestUtils.AccessorIterator<>(accessorSupplier); @BeforeEach public void setup() { this.vector = UnionVector.empty("", rootAllocatorTestExtension.getRootAllocator()); this.vector.allocateNew(); NullableBigIntHolder nullableBigIntHolder = new NullableBigIntHolder(); nullableBigIntHolder.isSet = 1; nullableBigIntHolder.value = Long.MAX_VALUE; this.vector.setType(0, Types.MinorType.BIGINT); this.vector.setSafe(0, nullableBigIntHolder); NullableFloat8Holder nullableFloat4Holder = new NullableFloat8Holder(); nullableFloat4Holder.isSet = 1; nullableFloat4Holder.value = Math.PI; this.vector.setType(1, Types.MinorType.FLOAT8); this.vector.setSafe(1, nullableFloat4Holder); NullableTimeStampMilliHolder nullableTimeStampMilliHolder = new NullableTimeStampMilliHolder(); nullableTimeStampMilliHolder.isSet = 1; nullableTimeStampMilliHolder.value = 1625702400000L; this.vector.setType(2, Types.MinorType.TIMESTAMPMILLI); this.vector.setSafe(2, nullableTimeStampMilliHolder); nullableBigIntHolder.isSet = 0; this.vector.setType(3, Types.MinorType.BIGINT); this.vector.setSafe(3, nullableBigIntHolder); this.vector.setValueCount(5); } @AfterEach public void tearDown() { this.vector.close(); } @Test public void getObject() throws Exception { List result = accessorIterator.toList(vector); List expected = Arrays.asList(Long.MAX_VALUE, Math.PI, new Timestamp(1625702400000L), null, null); assertThat(result, is(expected)); } @Test public void getObjectForNull() throws Exception { vector.reset(); vector.setValueCount(5); accessorIterator.assertAccessorGetter( vector, AbstractArrowFlightJdbcUnionVectorAccessor::getObject, equalTo(null)); } } numeric/000077500000000000000000000000001515451615200364275ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/implArrowFlightJdbcBaseIntVectorAccessorTest.java000066400000000000000000000201451515451615200472630ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/numeric/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.numeric; import static org.hamcrest.CoreMatchers.equalTo; import java.util.function.Supplier; import java.util.stream.Stream; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.BaseIntVector; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.SmallIntVector; 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.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; public class ArrowFlightJdbcBaseIntVectorAccessorTest { @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private BaseIntVector vector; private final AccessorTestUtils.AccessorSupplier accessorSupplier = (vector, getCurrentRow) -> { ArrowFlightJdbcAccessorFactory.WasNullConsumer noOpWasNullConsumer = (boolean wasNull) -> {}; if (vector instanceof UInt1Vector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (UInt1Vector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof UInt2Vector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (UInt2Vector) vector, getCurrentRow, noOpWasNullConsumer); } else { if (vector instanceof UInt4Vector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (UInt4Vector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof UInt8Vector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (UInt8Vector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof TinyIntVector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (TinyIntVector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof SmallIntVector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (SmallIntVector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof IntVector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (IntVector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof BigIntVector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (BigIntVector) vector, getCurrentRow, noOpWasNullConsumer); } } throw new UnsupportedOperationException(); }; private final AccessorTestUtils.AccessorIterator accessorIterator = new AccessorTestUtils.AccessorIterator<>(accessorSupplier); public static Stream data() { return Stream.of( Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createIntVector(), "IntVector"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createSmallIntVector(), "SmallIntVector"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createTinyIntVector(), "TinyIntVector"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createBigIntVector(), "BigIntVector"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createUInt1Vector(), "UInt1Vector"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createUInt2Vector(), "UInt2Vector"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createUInt4Vector(), "UInt4Vector"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createUInt8Vector(), "UInt8Vector")); } public void setup(Supplier vectorSupplier) { this.vector = vectorSupplier.get(); } @AfterEach public void tearDown() { this.vector.close(); } @ParameterizedTest @MethodSource("data") public void testShouldConvertToByteMethodFromBaseIntVector(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcBaseIntVectorAccessor::getByte, (accessor, currentRow) -> equalTo((byte) accessor.getLong())); } @ParameterizedTest @MethodSource("data") public void testShouldConvertToShortMethodFromBaseIntVector( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcBaseIntVectorAccessor::getShort, (accessor, currentRow) -> equalTo((short) accessor.getLong())); } @ParameterizedTest @MethodSource("data") public void testShouldConvertToIntegerMethodFromBaseIntVector( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcBaseIntVectorAccessor::getInt, (accessor, currentRow) -> equalTo((int) accessor.getLong())); } @ParameterizedTest @MethodSource("data") public void testShouldConvertToFloatMethodFromBaseIntVector( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcBaseIntVectorAccessor::getFloat, (accessor, currentRow) -> equalTo((float) accessor.getLong())); } @ParameterizedTest @MethodSource("data") public void testShouldConvertToDoubleMethodFromBaseIntVector( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcBaseIntVectorAccessor::getDouble, (accessor, currentRow) -> equalTo((double) accessor.getLong())); } @ParameterizedTest @MethodSource("data") public void testShouldConvertToBooleanMethodFromBaseIntVector( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcBaseIntVectorAccessor::getBoolean, (accessor, currentRow) -> equalTo(accessor.getLong() != 0L)); } @ParameterizedTest @MethodSource("data") public void testShouldGetObjectClass(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcBaseIntVectorAccessor::getObjectClass, equalTo(Long.class)); } } ArrowFlightJdbcBaseIntVectorAccessorUnitTest.java000066400000000000000000000217721515451615200501320ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/numeric/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.numeric; import static org.hamcrest.CoreMatchers.equalTo; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.SmallIntVector; 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.hamcrest.CoreMatchers; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.junit.jupiter.api.extension.RegisterExtension; import org.mockito.junit.jupiter.MockitoExtension; @ExtendWith(MockitoExtension.class) public class ArrowFlightJdbcBaseIntVectorAccessorUnitTest { @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private static UInt4Vector int4Vector; private static UInt8Vector int8Vector; private static IntVector intVectorWithNull; private static TinyIntVector tinyIntVector; private static SmallIntVector smallIntVector; private static IntVector intVector; private static BigIntVector bigIntVector; private final AccessorTestUtils.AccessorSupplier accessorSupplier = (vector, getCurrentRow) -> { ArrowFlightJdbcAccessorFactory.WasNullConsumer noOpWasNullConsumer = (boolean wasNull) -> {}; if (vector instanceof UInt1Vector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (UInt1Vector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof UInt2Vector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (UInt2Vector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof UInt4Vector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (UInt4Vector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof UInt8Vector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (UInt8Vector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof TinyIntVector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (TinyIntVector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof SmallIntVector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (SmallIntVector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof IntVector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (IntVector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof BigIntVector) { return new ArrowFlightJdbcBaseIntVectorAccessor( (BigIntVector) vector, getCurrentRow, noOpWasNullConsumer); } return null; }; private final AccessorTestUtils.AccessorIterator accessorIterator = new AccessorTestUtils.AccessorIterator<>(accessorSupplier); @BeforeAll public static void setup() { int4Vector = new UInt4Vector("ID", rootAllocatorTestExtension.getRootAllocator()); int4Vector.setSafe(0, 0x80000001); int4Vector.setValueCount(1); int8Vector = new UInt8Vector("ID", rootAllocatorTestExtension.getRootAllocator()); int8Vector.setSafe(0, 0xFFFFFFFFFFFFFFFFL); int8Vector.setValueCount(1); intVectorWithNull = new IntVector("ID", rootAllocatorTestExtension.getRootAllocator()); intVectorWithNull.setNull(0); intVectorWithNull.setValueCount(1); tinyIntVector = new TinyIntVector("ID", rootAllocatorTestExtension.getRootAllocator()); tinyIntVector.setSafe(0, 0xAA); tinyIntVector.setValueCount(1); smallIntVector = new SmallIntVector("ID", rootAllocatorTestExtension.getRootAllocator()); smallIntVector.setSafe(0, 0xAABB); smallIntVector.setValueCount(1); intVector = new IntVector("ID", rootAllocatorTestExtension.getRootAllocator()); intVector.setSafe(0, 0xAABBCCDD); intVector.setValueCount(1); bigIntVector = new BigIntVector("ID", rootAllocatorTestExtension.getRootAllocator()); bigIntVector.setSafe(0, 0xAABBCCDDEEFFAABBL); bigIntVector.setValueCount(1); } @AfterAll public static void tearDown() throws Exception { AutoCloseables.close( bigIntVector, intVector, smallIntVector, tinyIntVector, int4Vector, int8Vector, intVectorWithNull, rootAllocatorTestExtension); } @Test public void testShouldGetStringFromUnsignedValue() throws Exception { accessorIterator.assertAccessorGetter( int8Vector, ArrowFlightJdbcBaseIntVectorAccessor::getString, equalTo("18446744073709551615")); } @Test public void testShouldGetBytesFromIntVectorThrowsSqlException() throws Exception { accessorIterator.assertAccessorGetterThrowingException( intVector, ArrowFlightJdbcBaseIntVectorAccessor::getBytes); } @Test public void testShouldGetStringFromIntVectorWithNull() throws Exception { accessorIterator.assertAccessorGetter( intVectorWithNull, ArrowFlightJdbcBaseIntVectorAccessor::getString, CoreMatchers.nullValue()); } @Test public void testShouldGetObjectFromInt() throws Exception { accessorIterator.assertAccessorGetter( intVector, ArrowFlightJdbcBaseIntVectorAccessor::getObject, equalTo(0xAABBCCDD)); } @Test public void testShouldGetObjectFromTinyInt() throws Exception { accessorIterator.assertAccessorGetter( tinyIntVector, ArrowFlightJdbcBaseIntVectorAccessor::getObject, equalTo((byte) 0xAA)); } @Test public void testShouldGetObjectFromSmallInt() throws Exception { accessorIterator.assertAccessorGetter( smallIntVector, ArrowFlightJdbcBaseIntVectorAccessor::getObject, equalTo((short) 0xAABB)); } @Test public void testShouldGetObjectFromBigInt() throws Exception { accessorIterator.assertAccessorGetter( bigIntVector, ArrowFlightJdbcBaseIntVectorAccessor::getObject, equalTo(0xAABBCCDDEEFFAABBL)); } @Test public void testShouldGetObjectFromUnsignedInt() throws Exception { accessorIterator.assertAccessorGetter( int4Vector, ArrowFlightJdbcBaseIntVectorAccessor::getObject, equalTo(0x80000001)); } @Test public void testShouldGetObjectFromIntVectorWithNull() throws Exception { accessorIterator.assertAccessorGetter( intVectorWithNull, ArrowFlightJdbcBaseIntVectorAccessor::getObject, CoreMatchers.nullValue()); } @Test public void testShouldGetBigDecimalFromIntVectorWithNull() throws Exception { accessorIterator.assertAccessorGetter( intVectorWithNull, ArrowFlightJdbcBaseIntVectorAccessor::getBigDecimal, CoreMatchers.nullValue()); } @Test public void testShouldGetBigDecimalWithScaleFromIntVectorWithNull() throws Exception { accessorIterator.assertAccessorGetter( intVectorWithNull, accessor -> accessor.getBigDecimal(2), CoreMatchers.nullValue()); } @Test public void testShouldGetBytesFromSmallVectorThrowsSqlException() throws Exception { accessorIterator.assertAccessorGetterThrowingException( smallIntVector, ArrowFlightJdbcBaseIntVectorAccessor::getBytes); } @Test public void testShouldGetBytesFromTinyIntVectorThrowsSqlException() throws Exception { accessorIterator.assertAccessorGetterThrowingException( tinyIntVector, ArrowFlightJdbcBaseIntVectorAccessor::getBytes); } @Test public void testShouldGetBytesFromBigIntVectorThrowsSqlException() throws Exception { accessorIterator.assertAccessorGetterThrowingException( bigIntVector, ArrowFlightJdbcBaseIntVectorAccessor::getBytes); } } ArrowFlightJdbcBitVectorAccessorTest.java000066400000000000000000000124251515451615200464560ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/numeric/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.numeric; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import java.math.BigDecimal; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils.AccessorIterator; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils.CheckedFunction; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.BitVector; 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; public class ArrowFlightJdbcBitVectorAccessorTest { @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private final AccessorTestUtils.AccessorSupplier accessorSupplier = (vector, getCurrentRow) -> new ArrowFlightJdbcBitVectorAccessor( (BitVector) vector, getCurrentRow, (boolean wasNull) -> {}); private final AccessorIterator accessorIterator = new AccessorIterator<>(accessorSupplier); private BitVector vector; private BitVector vectorWithNull; private boolean[] arrayToAssert; @BeforeEach public void setup() { this.arrayToAssert = new boolean[] {false, true}; this.vector = rootAllocatorTestExtension.createBitVector(); this.vectorWithNull = rootAllocatorTestExtension.createBitVectorForNullTests(); } @AfterEach public void tearDown() { this.vector.close(); this.vectorWithNull.close(); } private void iterate( final CheckedFunction function, final T result, final T resultIfFalse, final BitVector vector) throws Exception { accessorIterator.assertAccessorGetter( vector, function, (accessor, currentRow) -> is(arrayToAssert[currentRow] ? result : resultIfFalse)); } @Test public void testShouldGetBooleanMethodFromBitVector() throws Exception { iterate(ArrowFlightJdbcBitVectorAccessor::getBoolean, true, false, vector); } @Test public void testShouldGetByteMethodFromBitVector() throws Exception { iterate(ArrowFlightJdbcBitVectorAccessor::getByte, (byte) 1, (byte) 0, vector); } @Test public void testShouldGetShortMethodFromBitVector() throws Exception { iterate(ArrowFlightJdbcBitVectorAccessor::getShort, (short) 1, (short) 0, vector); } @Test public void testShouldGetIntMethodFromBitVector() throws Exception { iterate(ArrowFlightJdbcBitVectorAccessor::getInt, 1, 0, vector); } @Test public void testShouldGetLongMethodFromBitVector() throws Exception { iterate(ArrowFlightJdbcBitVectorAccessor::getLong, (long) 1, (long) 0, vector); } @Test public void testShouldGetFloatMethodFromBitVector() throws Exception { iterate(ArrowFlightJdbcBitVectorAccessor::getFloat, (float) 1, (float) 0, vector); } @Test public void testShouldGetDoubleMethodFromBitVector() throws Exception { iterate(ArrowFlightJdbcBitVectorAccessor::getDouble, (double) 1, (double) 0, vector); } @Test public void testShouldGetBigDecimalMethodFromBitVector() throws Exception { iterate( ArrowFlightJdbcBitVectorAccessor::getBigDecimal, BigDecimal.ONE, BigDecimal.ZERO, vector); } @Test public void testShouldGetBigDecimalMethodFromBitVectorFromNull() throws Exception { iterate(ArrowFlightJdbcBitVectorAccessor::getBigDecimal, null, null, vectorWithNull); } @Test public void testShouldGetObjectMethodFromBitVector() throws Exception { iterate(ArrowFlightJdbcBitVectorAccessor::getObject, true, false, vector); } @Test public void testShouldGetObjectMethodFromBitVectorFromNull() throws Exception { iterate(ArrowFlightJdbcBitVectorAccessor::getObject, null, null, vectorWithNull); } @Test public void testShouldGetStringMethodFromBitVector() throws Exception { iterate(ArrowFlightJdbcBitVectorAccessor::getString, "true", "false", vector); } @Test public void testShouldGetStringMethodFromBitVectorFromNull() throws Exception { iterate(ArrowFlightJdbcBitVectorAccessor::getString, null, null, vectorWithNull); } @Test public void testShouldGetObjectClass() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcBitVectorAccessor::getObjectClass, equalTo(Boolean.class)); } } ArrowFlightJdbcDecimalVectorAccessorTest.java000066400000000000000000000273641515451615200473060ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/numeric/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.numeric; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import java.math.BigDecimal; import java.util.function.Supplier; import java.util.stream.Stream; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.Decimal256Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.ValueVector; import org.hamcrest.CoreMatchers; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; public class ArrowFlightJdbcDecimalVectorAccessorTest { @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private ValueVector vector; private ValueVector vectorWithNull; private final AccessorTestUtils.AccessorSupplier accessorSupplier = (vector, getCurrentRow) -> { ArrowFlightJdbcAccessorFactory.WasNullConsumer noOpWasNullConsumer = (boolean wasNull) -> {}; if (vector instanceof DecimalVector) { return new ArrowFlightJdbcDecimalVectorAccessor( (DecimalVector) vector, getCurrentRow, noOpWasNullConsumer); } else if (vector instanceof Decimal256Vector) { return new ArrowFlightJdbcDecimalVectorAccessor( (Decimal256Vector) vector, getCurrentRow, noOpWasNullConsumer); } return null; }; private final AccessorTestUtils.AccessorIterator accessorIterator = new AccessorTestUtils.AccessorIterator<>(accessorSupplier); public static Stream data() { return Stream.of( Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createDecimalVector(), "DecimalVector"), Arguments.of( (Supplier) () -> rootAllocatorTestExtension.createDecimal256Vector(), "Decimal256Vector")); } public void setup(Supplier vectorSupplier) { this.vector = vectorSupplier.get(); this.vectorWithNull = vectorSupplier.get(); this.vectorWithNull.clear(); this.vectorWithNull.setValueCount(5); } @AfterEach public void tearDown() { this.vector.close(); this.vectorWithNull.close(); } @ParameterizedTest @MethodSource("data") public void testShouldGetBigDecimalFromDecimalVector(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcDecimalVectorAccessor::getBigDecimal, (accessor, currentRow) -> CoreMatchers.notNullValue()); } @ParameterizedTest @MethodSource("data") public void testShouldGetDoubleMethodFromDecimalVector(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcDecimalVectorAccessor::getDouble, (accessor, currentRow) -> equalTo(accessor.getBigDecimal().doubleValue())); } @ParameterizedTest @MethodSource("data") public void testShouldGetFloatMethodFromDecimalVector(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcDecimalVectorAccessor::getFloat, (accessor, currentRow) -> equalTo(accessor.getBigDecimal().floatValue())); } @ParameterizedTest @MethodSource("data") public void testShouldGetLongMethodFromDecimalVector(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcDecimalVectorAccessor::getLong, (accessor, currentRow) -> equalTo(accessor.getBigDecimal().longValue())); } @ParameterizedTest @MethodSource("data") public void testShouldGetIntMethodFromDecimalVector(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcDecimalVectorAccessor::getInt, (accessor, currentRow) -> equalTo(accessor.getBigDecimal().intValue())); } @ParameterizedTest @MethodSource("data") public void testShouldGetShortMethodFromDecimalVector(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcDecimalVectorAccessor::getShort, (accessor, currentRow) -> equalTo(accessor.getBigDecimal().shortValue())); } @ParameterizedTest @MethodSource("data") public void testShouldGetByteMethodFromDecimalVector(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcDecimalVectorAccessor::getByte, (accessor, currentRow) -> equalTo(accessor.getBigDecimal().byteValue())); } @ParameterizedTest @MethodSource("data") public void testShouldGetStringMethodFromDecimalVector(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcDecimalVectorAccessor::getString, (accessor, currentRow) -> equalTo(accessor.getBigDecimal().toString())); } @ParameterizedTest @MethodSource("data") public void testShouldGetBooleanMethodFromDecimalVector(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcDecimalVectorAccessor::getBoolean, (accessor, currentRow) -> equalTo(!accessor.getBigDecimal().equals(BigDecimal.ZERO))); } @ParameterizedTest @MethodSource("data") public void testShouldGetObjectMethodFromDecimalVector(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcDecimalVectorAccessor::getObject, (accessor, currentRow) -> equalTo(accessor.getBigDecimal())); } @ParameterizedTest @MethodSource("data") public void testShouldGetObjectClass(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcDecimalVectorAccessor::getObjectClass, (accessor, currentRow) -> equalTo(BigDecimal.class)); } @ParameterizedTest @MethodSource("data") public void testShouldGetBigDecimalMethodFromDecimalVectorWithNull( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vectorWithNull, ArrowFlightJdbcDecimalVectorAccessor::getBigDecimal, (accessor, currentRow) -> CoreMatchers.nullValue()); } @ParameterizedTest @MethodSource("data") public void testShouldGetObjectMethodFromDecimalVectorWithNull( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vectorWithNull, ArrowFlightJdbcDecimalVectorAccessor::getObject, (accessor, currentRow) -> CoreMatchers.nullValue()); } @ParameterizedTest @MethodSource("data") public void testShouldGetStringMethodFromDecimalVectorWithNull( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vectorWithNull, ArrowFlightJdbcDecimalVectorAccessor::getString, (accessor, currentRow) -> CoreMatchers.nullValue()); } @ParameterizedTest @MethodSource("data") public void testShouldGetByteMethodFromDecimalVectorWithNull(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vectorWithNull, ArrowFlightJdbcDecimalVectorAccessor::getByte, (accessor, currentRow) -> is((byte) 0)); } @ParameterizedTest @MethodSource("data") public void testShouldGetShortMethodFromDecimalVectorWithNull( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vectorWithNull, ArrowFlightJdbcDecimalVectorAccessor::getShort, (accessor, currentRow) -> is((short) 0)); } @ParameterizedTest @MethodSource("data") public void testShouldGetIntMethodFromDecimalVectorWithNull(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vectorWithNull, ArrowFlightJdbcDecimalVectorAccessor::getInt, (accessor, currentRow) -> is(0)); } @ParameterizedTest @MethodSource("data") public void testShouldGetLongMethodFromDecimalVectorWithNull(Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vectorWithNull, ArrowFlightJdbcDecimalVectorAccessor::getLong, (accessor, currentRow) -> is((long) 0)); } @ParameterizedTest @MethodSource("data") public void testShouldGetFloatMethodFromDecimalVectorWithNull( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vectorWithNull, ArrowFlightJdbcDecimalVectorAccessor::getFloat, (accessor, currentRow) -> is(0.0f)); } @ParameterizedTest @MethodSource("data") public void testShouldGetDoubleMethodFromDecimalVectorWithNull( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vectorWithNull, ArrowFlightJdbcDecimalVectorAccessor::getDouble, (accessor, currentRow) -> is(0.0D)); } @ParameterizedTest @MethodSource("data") public void testShouldGetBooleanMethodFromDecimalVectorWithNull( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vectorWithNull, ArrowFlightJdbcDecimalVectorAccessor::getBoolean, (accessor, currentRow) -> is(false)); } @ParameterizedTest @MethodSource("data") public void testShouldGetBigDecimalWithScaleMethodFromDecimalVectorWithNull( Supplier vectorSupplier) throws Exception { setup(vectorSupplier); accessorIterator.assertAccessorGetter( vectorWithNull, accessor -> accessor.getBigDecimal(2), (accessor, currentRow) -> CoreMatchers.nullValue()); } } ArrowFlightJdbcFloat4VectorAccessorTest.java000066400000000000000000000171521515451615200470730ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/numeric/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.numeric; 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.math.BigDecimal; import java.math.RoundingMode; import java.sql.SQLException; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.Float4Vector; import org.hamcrest.CoreMatchers; 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; public class ArrowFlightJdbcFloat4VectorAccessorTest { @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private Float4Vector vector; private final AccessorTestUtils.AccessorSupplier accessorSupplier = (vector, getCurrentRow) -> new ArrowFlightJdbcFloat4VectorAccessor( (Float4Vector) vector, getCurrentRow, (boolean wasNull) -> {}); private final AccessorTestUtils.AccessorIterator accessorIterator = new AccessorTestUtils.AccessorIterator<>(accessorSupplier); @BeforeEach public void setup() { this.vector = rootAllocatorTestExtension.createFloat4Vector(); } @AfterEach public void tearDown() { this.vector.close(); } @Test public void testShouldGetFloatMethodFromFloat4Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat4VectorAccessor::getFloat, (accessor, currentRow) -> is(vector.get(currentRow))); } @Test public void testShouldGetObjectMethodFromFloat4Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat4VectorAccessor::getObject, (accessor) -> is(accessor.getFloat())); } @Test public void testShouldGetStringMethodFromFloat4Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat4VectorAccessor::getString, accessor -> is(Float.toString(accessor.getFloat()))); } @Test public void testShouldGetStringMethodFromFloat4VectorWithNull() throws Exception { try (final Float4Vector float4Vector = new Float4Vector("ID", rootAllocatorTestExtension.getRootAllocator())) { float4Vector.setNull(0); float4Vector.setValueCount(1); accessorIterator.assertAccessorGetter( float4Vector, ArrowFlightJdbcFloat4VectorAccessor::getString, CoreMatchers.nullValue()); } } @Test public void testShouldGetFloatMethodFromFloat4VectorWithNull() throws Exception { try (final Float4Vector float4Vector = new Float4Vector("ID", rootAllocatorTestExtension.getRootAllocator())) { float4Vector.setNull(0); float4Vector.setValueCount(1); accessorIterator.assertAccessorGetter( float4Vector, ArrowFlightJdbcFloat4VectorAccessor::getFloat, is(0.0f)); } } @Test public void testShouldGetBigDecimalMethodFromFloat4VectorWithNull() throws Exception { try (final Float4Vector float4Vector = new Float4Vector("ID", rootAllocatorTestExtension.getRootAllocator())) { float4Vector.setNull(0); float4Vector.setValueCount(1); accessorIterator.assertAccessorGetter( float4Vector, ArrowFlightJdbcFloat4VectorAccessor::getBigDecimal, CoreMatchers.nullValue()); } } @Test public void testShouldGetObjectMethodFromFloat4VectorWithNull() throws Exception { try (final Float4Vector float4Vector = new Float4Vector("ID", rootAllocatorTestExtension.getRootAllocator())) { float4Vector.setNull(0); float4Vector.setValueCount(1); accessorIterator.assertAccessorGetter( float4Vector, ArrowFlightJdbcFloat4VectorAccessor::getObject, CoreMatchers.nullValue()); } } @Test public void testShouldGetBooleanMethodFromFloat4Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat4VectorAccessor::getBoolean, accessor -> is(accessor.getFloat() != 0.0f)); } @Test public void testShouldGetByteMethodFromFloat4Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat4VectorAccessor::getByte, accessor -> is((byte) accessor.getFloat())); } @Test public void testShouldGetShortMethodFromFloat4Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat4VectorAccessor::getShort, accessor -> is((short) accessor.getFloat())); } @Test public void testShouldGetIntMethodFromFloat4Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat4VectorAccessor::getInt, accessor -> is((int) accessor.getFloat())); } @Test public void testShouldGetLongMethodFromFloat4Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat4VectorAccessor::getLong, accessor -> is((long) accessor.getFloat())); } @Test public void testShouldGetDoubleMethodFromFloat4Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat4VectorAccessor::getDouble, accessor -> is((double) accessor.getFloat())); } @Test public void testShouldGetBigDecimalMethodFromFloat4Vector() throws Exception { accessorIterator.iterate( vector, (accessor, currentRow) -> { float value = accessor.getFloat(); if (Float.isInfinite(value) || Float.isNaN(value)) { assertThrows(SQLException.class, accessor::getBigDecimal); } else { assertThat(accessor.getBigDecimal(), is(BigDecimal.valueOf(value))); } }); } @Test public void testShouldGetBigDecimalWithScaleMethodFromFloat4Vector() throws Exception { accessorIterator.iterate( vector, (accessor, currentRow) -> { float value = accessor.getFloat(); if (Float.isInfinite(value) || Float.isNaN(value)) { assertThrows(SQLException.class, () -> accessor.getBigDecimal(9)); } else { assertThat( accessor.getBigDecimal(9), is(BigDecimal.valueOf(value).setScale(9, RoundingMode.HALF_UP))); } }); } @Test public void testShouldGetObjectClass() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat4VectorAccessor::getObjectClass, accessor -> equalTo(Float.class)); } } ArrowFlightJdbcFloat8VectorAccessorTest.java000066400000000000000000000157601515451615200471020ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/numeric/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.numeric; 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.math.BigDecimal; import java.math.RoundingMode; import java.sql.SQLException; import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.vector.Float8Vector; import org.hamcrest.CoreMatchers; 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; public class ArrowFlightJdbcFloat8VectorAccessorTest { @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private Float8Vector vector; private Float8Vector vectorWithNull; private final AccessorTestUtils.AccessorSupplier accessorSupplier = (vector, getCurrentRow) -> new ArrowFlightJdbcFloat8VectorAccessor( (Float8Vector) vector, getCurrentRow, (boolean wasNull) -> {}); private final AccessorTestUtils.AccessorIterator accessorIterator = new AccessorTestUtils.AccessorIterator<>(accessorSupplier); @BeforeEach public void setup() { this.vector = rootAllocatorTestExtension.createFloat8Vector(); this.vectorWithNull = rootAllocatorTestExtension.createFloat8VectorForNullTests(); } @AfterEach public void tearDown() { this.vector.close(); this.vectorWithNull.close(); } @Test public void testShouldGetDoubleMethodFromFloat8Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat8VectorAccessor::getDouble, (accessor, currentRow) -> is(vector.getValueAsDouble(currentRow))); } @Test public void testShouldGetObjectMethodFromFloat8Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat8VectorAccessor::getObject, (accessor) -> is(accessor.getDouble())); } @Test public void testShouldGetStringMethodFromFloat8Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat8VectorAccessor::getString, (accessor) -> is(Double.toString(accessor.getDouble()))); } @Test public void testShouldGetBooleanMethodFromFloat8Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat8VectorAccessor::getBoolean, (accessor) -> is(accessor.getDouble() != 0.0)); } @Test public void testShouldGetByteMethodFromFloat8Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat8VectorAccessor::getByte, (accessor) -> is((byte) accessor.getDouble())); } @Test public void testShouldGetShortMethodFromFloat8Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat8VectorAccessor::getShort, (accessor) -> is((short) accessor.getDouble())); } @Test public void testShouldGetIntMethodFromFloat8Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat8VectorAccessor::getInt, (accessor) -> is((int) accessor.getDouble())); } @Test public void testShouldGetLongMethodFromFloat8Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat8VectorAccessor::getLong, (accessor) -> is((long) accessor.getDouble())); } @Test public void testShouldGetFloatMethodFromFloat8Vector() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat8VectorAccessor::getFloat, (accessor) -> is((float) accessor.getDouble())); } @Test public void testShouldGetBigDecimalMethodFromFloat8Vector() throws Exception { accessorIterator.iterate( vector, (accessor, currentRow) -> { double value = accessor.getDouble(); if (Double.isInfinite(value) || Double.isNaN(value)) { assertThrows(SQLException.class, accessor::getBigDecimal); } else { assertThat(accessor.getBigDecimal(), is(BigDecimal.valueOf(value))); } }); } @Test public void testShouldGetObjectClass() throws Exception { accessorIterator.assertAccessorGetter( vector, ArrowFlightJdbcFloat8VectorAccessor::getObjectClass, equalTo(Double.class)); } @Test public void testShouldGetStringMethodFromFloat8VectorWithNull() throws Exception { accessorIterator.assertAccessorGetter( vectorWithNull, ArrowFlightJdbcFloat8VectorAccessor::getString, CoreMatchers.nullValue()); } @Test public void testShouldGetFloatMethodFromFloat8VectorWithNull() throws Exception { accessorIterator.assertAccessorGetter( vectorWithNull, ArrowFlightJdbcFloat8VectorAccessor::getFloat, is(0.0f)); } @Test public void testShouldGetBigDecimalMethodFromFloat8VectorWithNull() throws Exception { accessorIterator.assertAccessorGetter( vectorWithNull, ArrowFlightJdbcFloat8VectorAccessor::getBigDecimal, CoreMatchers.nullValue()); } @Test public void testShouldGetBigDecimalWithScaleMethodFromFloat4Vector() throws Exception { accessorIterator.iterate( vector, (accessor, currentRow) -> { double value = accessor.getDouble(); if (Double.isInfinite(value) || Double.isNaN(value)) { assertThrows(SQLException.class, () -> accessor.getBigDecimal(9)); } else { assertThat( accessor.getBigDecimal(9), is(BigDecimal.valueOf(value).setScale(9, RoundingMode.HALF_UP))); } }); } @Test public void testShouldGetObjectMethodFromFloat8VectorWithNull() throws Exception { accessorIterator.assertAccessorGetter( vectorWithNull, ArrowFlightJdbcFloat8VectorAccessor::getObject, CoreMatchers.nullValue()); } } text/000077500000000000000000000000001515451615200357515ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/implArrowFlightJdbcVarCharVectorAccessorTest.java000066400000000000000000000601441515451615200466110ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/text/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.text; import static java.nio.charset.StandardCharsets.US_ASCII; import static java.nio.charset.StandardCharsets.UTF_8; import static org.apache.commons.io.IOUtils.toByteArray; import static org.apache.commons.io.IOUtils.toCharArray; import static org.hamcrest.CoreMatchers.equalTo; 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; import java.io.InputStream; import java.io.Reader; import java.math.BigDecimal; import java.sql.Date; import java.sql.SQLException; import java.sql.Time; import java.sql.Timestamp; import java.text.SimpleDateFormat; import java.util.Calendar; import java.util.TimeZone; import java.util.function.IntSupplier; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDateVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeStampVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeVectorAccessor; import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; import org.apache.arrow.driver.jdbc.utils.ThrowableAssertionUtils; 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; import org.junit.jupiter.api.extension.ExtendWith; import org.junit.jupiter.api.extension.RegisterExtension; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; @ExtendWith(MockitoExtension.class) public class ArrowFlightJdbcVarCharVectorAccessorTest { private ArrowFlightJdbcVarCharVectorAccessor accessor; private final SimpleDateFormat dateTimeFormat = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSXXX"); private final SimpleDateFormat timeFormat = new SimpleDateFormat("HH:mm:ss.SSSXXX"); @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); @Mock private ArrowFlightJdbcVarCharVectorAccessor.Getter getter; @BeforeEach public void setUp() { IntSupplier currentRowSupplier = () -> 0; accessor = new ArrowFlightJdbcVarCharVectorAccessor( getter, currentRowSupplier, (boolean wasNull) -> {}); } @Test public void testShouldGetStringFromNullReturnNull() { when(getter.get(0)).thenReturn(null); final String result = accessor.getString(); assertThat(result, equalTo(null)); } @Test public void testShouldGetStringReturnValidString() { Text value = new Text("Value for Test."); when(getter.get(0)).thenReturn(value.copyBytes()); final String result = accessor.getString(); assertThat(result, instanceOf(String.class)); assertThat(result, equalTo(value.toString())); } @Test public void testShouldGetObjectReturnValidString() { Text value = new Text("Value for Test."); when(getter.get(0)).thenReturn(value.copyBytes()); final String result = accessor.getObject(); assertThat(result, instanceOf(String.class)); assertThat(result, equalTo(value.toString())); } @Test public void testShouldGetByteThrowsExceptionForNonNumericValue() { Text value = new Text("Invalid value for byte."); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getByte()); } @Test public void testShouldGetByteThrowsExceptionForOutOfRangePositiveValue() { Text value = new Text("128"); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getByte()); } @Test public void testShouldGetByteThrowsExceptionForOutOfRangeNegativeValue() { Text value = new Text("-129"); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getByte()); } @Test public void testShouldGetByteReturnValidPositiveByte() throws Exception { Text value = new Text("127"); when(getter.get(0)).thenReturn(value.copyBytes()); byte result = accessor.getByte(); assertThat(result, instanceOf(Byte.class)); assertThat(result, equalTo((byte) 127)); } @Test public void testShouldGetByteReturnValidNegativeByte() throws Exception { Text value = new Text("-128"); when(getter.get(0)).thenReturn(value.copyBytes()); byte result = accessor.getByte(); assertThat(result, instanceOf(Byte.class)); assertThat(result, equalTo((byte) -128)); } @Test public void testShouldGetShortThrowsExceptionForNonNumericValue() { Text value = new Text("Invalid value for short."); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getShort()); } @Test public void testShouldGetShortThrowsExceptionForOutOfRangePositiveValue() { Text value = new Text("32768"); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getShort()); } @Test public void testShouldGetShortThrowsExceptionForOutOfRangeNegativeValue() { Text value = new Text("-32769"); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getShort()); } @Test public void testShouldGetShortReturnValidPositiveShort() throws Exception { Text value = new Text("32767"); when(getter.get(0)).thenReturn(value.copyBytes()); short result = accessor.getShort(); assertThat(result, instanceOf(Short.class)); assertThat(result, equalTo((short) 32767)); } @Test public void testShouldGetShortReturnValidNegativeShort() throws Exception { Text value = new Text("-32768"); when(getter.get(0)).thenReturn(value.copyBytes()); short result = accessor.getShort(); assertThat(result, instanceOf(Short.class)); assertThat(result, equalTo((short) -32768)); } @Test public void testShouldGetIntThrowsExceptionForNonNumericValue() { Text value = new Text("Invalid value for int."); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getInt()); } @Test public void testShouldGetIntThrowsExceptionForOutOfRangePositiveValue() { Text value = new Text("2147483648"); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getInt()); } @Test public void testShouldGetIntThrowsExceptionForOutOfRangeNegativeValue() { Text value = new Text("-2147483649"); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getInt()); } @Test public void testShouldGetIntReturnValidPositiveInteger() throws Exception { Text value = new Text("2147483647"); when(getter.get(0)).thenReturn(value.copyBytes()); int result = accessor.getInt(); assertThat(result, instanceOf(Integer.class)); assertThat(result, equalTo(2147483647)); } @Test public void testShouldGetIntReturnValidNegativeInteger() throws Exception { Text value = new Text("-2147483648"); when(getter.get(0)).thenReturn(value.copyBytes()); int result = accessor.getInt(); assertThat(result, instanceOf(Integer.class)); assertThat(result, equalTo(-2147483648)); } @Test public void testShouldGetLongThrowsExceptionForNonNumericValue() { Text value = new Text("Invalid value for long."); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getLong()); } @Test public void testShouldGetLongThrowsExceptionForOutOfRangePositiveValue() { Text value = new Text("9223372036854775808"); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getLong()); } @Test public void testShouldGetLongThrowsExceptionForOutOfRangeNegativeValue() { Text value = new Text("-9223372036854775809"); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getLong()); } @Test public void testShouldGetLongReturnValidPositiveLong() throws Exception { Text value = new Text("9223372036854775807"); when(getter.get(0)).thenReturn(value.copyBytes()); long result = accessor.getLong(); assertThat(result, instanceOf(Long.class)); assertThat(result, equalTo(9223372036854775807L)); } @Test public void testShouldGetLongReturnValidNegativeLong() throws Exception { Text value = new Text("-9223372036854775808"); when(getter.get(0)).thenReturn(value.copyBytes()); long result = accessor.getLong(); assertThat(result, instanceOf(Long.class)); assertThat(result, equalTo(-9223372036854775808L)); } @Test public void testShouldBigDecimalWithParametersThrowsExceptionForNonNumericValue() { Text value = new Text("Invalid value for BigDecimal."); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getBigDecimal(1)); } @Test public void testShouldGetBigDecimalThrowsExceptionForNonNumericValue() { Text value = new Text("Invalid value for BigDecimal."); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getBigDecimal()); } @Test public void testShouldGetBigDecimalReturnValidPositiveBigDecimal() throws Exception { Text value = new Text("9223372036854775807000.999"); when(getter.get(0)).thenReturn(value.copyBytes()); BigDecimal result = accessor.getBigDecimal(); assertThat(result, instanceOf(BigDecimal.class)); assertThat(result, equalTo(new BigDecimal("9223372036854775807000.999"))); } @Test public void testShouldGetBigDecimalReturnValidNegativeBigDecimal() throws Exception { Text value = new Text("-9223372036854775807000.999"); when(getter.get(0)).thenReturn(value.copyBytes()); BigDecimal result = accessor.getBigDecimal(); assertThat(result, instanceOf(BigDecimal.class)); assertThat(result, equalTo(new BigDecimal("-9223372036854775807000.999"))); } @Test public void testShouldGetDoubleThrowsExceptionForNonNumericValue() { Text value = new Text("Invalid value for double."); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getDouble()); } @Test public void testShouldGetDoubleReturnValidPositiveDouble() throws Exception { Text value = new Text("1.7976931348623157E308D"); when(getter.get(0)).thenReturn(value.copyBytes()); double result = accessor.getDouble(); assertThat(result, instanceOf(Double.class)); assertThat(result, equalTo(1.7976931348623157E308D)); } @Test public void testShouldGetDoubleReturnValidNegativeDouble() throws Exception { Text value = new Text("-1.7976931348623157E308D"); when(getter.get(0)).thenReturn(value.copyBytes()); double result = accessor.getDouble(); assertThat(result, instanceOf(Double.class)); assertThat(result, equalTo(-1.7976931348623157E308D)); } @Test public void testShouldGetDoubleWorkWithPositiveInfinity() throws Exception { Text value = new Text("Infinity"); when(getter.get(0)).thenReturn(value.copyBytes()); double result = accessor.getDouble(); assertThat(result, instanceOf(Double.class)); assertThat(result, equalTo(Double.POSITIVE_INFINITY)); } @Test public void testShouldGetDoubleWorkWithNegativeInfinity() throws Exception { Text value = new Text("-Infinity"); when(getter.get(0)).thenReturn(value.copyBytes()); double result = accessor.getDouble(); assertThat(result, instanceOf(Double.class)); assertThat(result, equalTo(Double.NEGATIVE_INFINITY)); } @Test public void testShouldGetDoubleWorkWithNaN() throws Exception { Text value = new Text("NaN"); when(getter.get(0)).thenReturn(value.copyBytes()); double result = accessor.getDouble(); assertThat(result, instanceOf(Double.class)); assertThat(result, equalTo(Double.NaN)); } @Test public void testShouldGetFloatThrowsExceptionForNonNumericValue() { Text value = new Text("Invalid value for float."); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getFloat()); } @Test public void testShouldGetFloatReturnValidPositiveFloat() throws Exception { Text value = new Text("3.4028235E38F"); when(getter.get(0)).thenReturn(value.copyBytes()); float result = accessor.getFloat(); assertThat(result, instanceOf(Float.class)); assertThat(result, equalTo(3.4028235E38F)); } @Test public void testShouldGetFloatReturnValidNegativeFloat() throws Exception { Text value = new Text("-3.4028235E38F"); when(getter.get(0)).thenReturn(value.copyBytes()); float result = accessor.getFloat(); assertThat(result, instanceOf(Float.class)); assertThat(result, equalTo(-3.4028235E38F)); } @Test public void testShouldGetFloatWorkWithPositiveInfinity() throws Exception { Text value = new Text("Infinity"); when(getter.get(0)).thenReturn(value.copyBytes()); float result = accessor.getFloat(); assertThat(result, instanceOf(Float.class)); assertThat(result, equalTo(Float.POSITIVE_INFINITY)); } @Test public void testShouldGetFloatWorkWithNegativeInfinity() throws Exception { Text value = new Text("-Infinity"); when(getter.get(0)).thenReturn(value.copyBytes()); float result = accessor.getFloat(); assertThat(result, instanceOf(Float.class)); assertThat(result, equalTo(Float.NEGATIVE_INFINITY)); } @Test public void testShouldGetFloatWorkWithNaN() throws Exception { Text value = new Text("NaN"); when(getter.get(0)).thenReturn(value.copyBytes()); float result = accessor.getFloat(); assertThat(result, instanceOf(Float.class)); assertThat(result, equalTo(Float.NaN)); } @Test public void testShouldGetDateThrowsExceptionForNonDateValue() { Text value = new Text("Invalid value for date."); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getDate(null)); } @Test public void testShouldGetDateReturnValidDateWithoutCalendar() throws Exception { Text value = new Text("2021-07-02"); when(getter.get(0)).thenReturn(value.copyBytes()); Date result = accessor.getDate(null); assertThat(result, instanceOf(Date.class)); Calendar calendar = Calendar.getInstance(); calendar.setTime(result); assertThat(dateTimeFormat.format(calendar.getTime()), equalTo("2021-07-02T00:00:00.000Z")); } @Test public void testShouldGetDateReturnValidDateWithCalendar() throws Exception { Text value = new Text("2021-07-02"); when(getter.get(0)).thenReturn(value.copyBytes()); Calendar calendar = Calendar.getInstance(TimeZone.getTimeZone("America/Sao_Paulo")); Date result = accessor.getDate(calendar); calendar = Calendar.getInstance(TimeZone.getTimeZone("Etc/UTC")); calendar.setTime(result); assertThat(dateTimeFormat.format(calendar.getTime()), equalTo("2021-07-02T03:00:00.000Z")); } @Test public void testShouldGetTimeThrowsExceptionForNonTimeValue() { Text value = new Text("Invalid value for time."); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getTime(null)); } @Test public void testShouldGetTimeReturnValidDateWithoutCalendar() throws Exception { Text value = new Text("02:30:00"); when(getter.get(0)).thenReturn(value.copyBytes()); Time result = accessor.getTime(null); Calendar calendar = Calendar.getInstance(); calendar.setTime(result); assertThat(timeFormat.format(calendar.getTime()), equalTo("02:30:00.000Z")); } @Test public void testShouldGetTimeReturnValidDateWithCalendar() throws Exception { Text value = new Text("02:30:00"); when(getter.get(0)).thenReturn(value.copyBytes()); Calendar calendar = Calendar.getInstance(TimeZone.getTimeZone("America/Sao_Paulo")); Time result = accessor.getTime(calendar); calendar = Calendar.getInstance(TimeZone.getTimeZone("Etc/UTC")); calendar.setTime(result); assertThat(timeFormat.format(calendar.getTime()), equalTo("05:30:00.000Z")); } @Test public void testShouldGetTimestampThrowsExceptionForNonTimeValue() { Text value = new Text("Invalid value for timestamp."); when(getter.get(0)).thenReturn(value.copyBytes()); assertThrows(SQLException.class, () -> accessor.getTimestamp(null)); } @Test public void testShouldGetTimestampReturnValidDateWithoutCalendar() throws Exception { Text value = new Text("2021-07-02 02:30:00.000"); when(getter.get(0)).thenReturn(value.copyBytes()); Timestamp result = accessor.getTimestamp(null); Calendar calendar = Calendar.getInstance(); calendar.setTime(result); assertThat(dateTimeFormat.format(calendar.getTime()), equalTo("2021-07-02T02:30:00.000Z")); } @Test public void testShouldGetTimestampReturnValidDateWithCalendar() throws Exception { Text value = new Text("2021-07-02 02:30:00.000"); when(getter.get(0)).thenReturn(value.copyBytes()); Calendar calendar = Calendar.getInstance(TimeZone.getTimeZone("America/Sao_Paulo")); Timestamp result = accessor.getTimestamp(calendar); calendar = Calendar.getInstance(TimeZone.getTimeZone("Etc/UTC")); calendar.setTime(result); assertThat(dateTimeFormat.format(calendar.getTime()), equalTo("2021-07-02T05:30:00.000Z")); } private void assertGetBoolean(Text value, boolean expectedResult) throws SQLException { when(getter.get(0)).thenReturn(value == null ? null : value.copyBytes()); boolean result = accessor.getBoolean(); assertThat(result, equalTo(expectedResult)); } private void assertGetBooleanForSQLException(Text value) { when(getter.get(0)).thenReturn(value == null ? null : value.copyBytes()); ThrowableAssertionUtils.simpleAssertThrowableClass( SQLException.class, () -> accessor.getBoolean()); } @Test public void testShouldGetBooleanThrowsSQLExceptionForInvalidValue() { assertGetBooleanForSQLException(new Text("anything")); } @Test public void testShouldGetBooleanThrowsSQLExceptionForEmpty() { assertGetBooleanForSQLException(new Text("")); } @Test public void testShouldGetBooleanReturnFalseFor0() throws Exception { assertGetBoolean(new Text("0"), false); } @Test public void testShouldGetBooleanReturnFalseForFalseString() throws Exception { assertGetBoolean(new Text("false"), false); } @Test public void testShouldGetBooleanReturnFalseForNull() throws Exception { assertGetBoolean(null, false); } @Test public void testShouldGetBytesReturnValidByteArray() { Text value = new Text("Value for Test."); when(getter.get(0)).thenReturn(value.copyBytes()); final byte[] result = accessor.getBytes(); assertThat(result, instanceOf(byte[].class)); assertThat(result, equalTo(value.toString().getBytes(UTF_8))); } @Test public void testShouldGetUnicodeStreamReturnValidInputStream() throws Exception { Text value = new Text("Value for Test."); when(getter.get(0)).thenReturn(value.copyBytes()); try (final InputStream result = accessor.getUnicodeStream()) { byte[] resultBytes = toByteArray(result); assertThat(new String(resultBytes, UTF_8), equalTo(value.toString())); } } @Test public void testShouldGetAsciiStreamReturnValidInputStream() throws Exception { Text valueText = new Text("Value for Test."); byte[] valueAscii = valueText.toString().getBytes(US_ASCII); when(getter.get(0)).thenReturn(valueText.copyBytes()); try (final InputStream result = accessor.getAsciiStream()) { byte[] resultBytes = toByteArray(result); assertArrayEquals(valueAscii, resultBytes); } } @Test public void testShouldGetCharacterStreamReturnValidReader() throws Exception { Text value = new Text("Value for Test."); when(getter.get(0)).thenReturn(value.copyBytes()); try (Reader result = accessor.getCharacterStream()) { char[] resultChars = toCharArray(result); assertThat(new String(resultChars), equalTo(value.toString())); } } @Test public void testShouldGetTimeStampBeConsistentWithTimeStampAccessor() throws Exception { try (TimeStampVector timeStampVector = rootAllocatorTestExtension.createTimeStampMilliVector()) { ArrowFlightJdbcTimeStampVectorAccessor timeStampVectorAccessor = new ArrowFlightJdbcTimeStampVectorAccessor( timeStampVector, () -> 0, (boolean wasNull) -> {}); Text value = new Text(timeStampVectorAccessor.getString()); when(getter.get(0)).thenReturn(value.copyBytes()); Timestamp timestamp = accessor.getTimestamp(null); assertThat(timestamp, equalTo(timeStampVectorAccessor.getTimestamp(null))); } } @Test public void testShouldGetTimeBeConsistentWithTimeAccessor() throws Exception { try (TimeMilliVector timeVector = rootAllocatorTestExtension.createTimeMilliVector()) { ArrowFlightJdbcTimeVectorAccessor timeVectorAccessor = new ArrowFlightJdbcTimeVectorAccessor(timeVector, () -> 0, (boolean wasNull) -> {}); Text value = new Text(timeVectorAccessor.getString()); when(getter.get(0)).thenReturn(value.copyBytes()); Time time = accessor.getTime(null); assertThat(time, equalTo(timeVectorAccessor.getTime(null))); } } @Test public void testShouldGetDateBeConsistentWithDateAccessor() throws Exception { try (DateMilliVector dateVector = rootAllocatorTestExtension.createDateMilliVector()) { ArrowFlightJdbcDateVectorAccessor dateVectorAccessor = new ArrowFlightJdbcDateVectorAccessor(dateVector, () -> 0, (boolean wasNull) -> {}); Text value = new Text(dateVectorAccessor.getString()); when(getter.get(0)).thenReturn(value.copyBytes()); Date date = accessor.getDate(null); assertThat(date, equalTo(dateVectorAccessor.getDate(null))); } } @Test 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()); } } } authentication/000077500000000000000000000000001515451615200352415ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbcAuthentication.java000066400000000000000000000025531515451615200410700ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/authentication/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.authentication; import java.util.Properties; import org.apache.arrow.flight.auth2.CallHeaderAuthenticator; public interface Authentication { /** * Create a {@link CallHeaderAuthenticator} which is used to authenticate the connection. * * @return a CallHeaderAuthenticator. */ CallHeaderAuthenticator authenticate(); /** * Uses the validCredentials variable and populate the Properties object. * * @param properties the Properties object that will be populated. */ void populateProperties(Properties properties); } TokenAuthentication.java000066400000000000000000000051001515451615200420600ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/authentication/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.authentication; import java.util.ArrayList; import java.util.List; import java.util.Properties; import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl; import org.apache.arrow.flight.CallHeaders; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.auth2.CallHeaderAuthenticator; public class TokenAuthentication implements Authentication { private final List validCredentials; public TokenAuthentication(List validCredentials) { this.validCredentials = validCredentials; } @Override public CallHeaderAuthenticator authenticate() { return new CallHeaderAuthenticator() { @Override public AuthResult authenticate(CallHeaders incomingHeaders) { String authorization = incomingHeaders.get("authorization"); if (!validCredentials.contains(authorization)) { throw CallStatus.UNAUTHENTICATED .withDescription("Invalid credentials.") .toRuntimeException(); } return new AuthResult() { @Override public String getPeerIdentity() { return authorization; } }; } }; } @Override public void populateProperties(Properties properties) { this.validCredentials.forEach( value -> properties.put( ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty.TOKEN.camelName(), value)); } public static final class Builder { private final List tokenList = new ArrayList<>(); public TokenAuthentication.Builder token(String token) { tokenList.add("Bearer " + token); return this; } public TokenAuthentication build() { return new TokenAuthentication(tokenList); } } } UserPasswordAuthentication.java000066400000000000000000000055201515451615200434470ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/authentication/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.authentication; import java.util.HashMap; import java.util.Map; import java.util.Properties; import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.auth2.BasicCallHeaderAuthenticator; import org.apache.arrow.flight.auth2.CallHeaderAuthenticator; import org.apache.arrow.flight.auth2.GeneratedBearerTokenAuthenticator; public class UserPasswordAuthentication implements Authentication { private final Map validCredentials; public UserPasswordAuthentication(Map validCredentials) { this.validCredentials = validCredentials; } private String getCredentials(String key) { return validCredentials.getOrDefault(key, null); } @Override public CallHeaderAuthenticator authenticate() { return new GeneratedBearerTokenAuthenticator( new BasicCallHeaderAuthenticator( (username, password) -> { if (validCredentials.containsKey(username) && getCredentials(username).equals(password)) { return () -> username; } throw CallStatus.UNAUTHENTICATED .withDescription("Invalid credentials.") .toRuntimeException(); })); } @Override public void populateProperties(Properties properties) { validCredentials.forEach( (key, value) -> { properties.put( ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty.USER.camelName(), key); properties.put( ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty.PASSWORD.camelName(), value); }); } public static class Builder { Map credentials = new HashMap<>(); public Builder user(String username, String password) { credentials.put(username, password); return this; } public UserPasswordAuthentication build() { return new UserPasswordAuthentication(credentials); } } } client/000077500000000000000000000000001515451615200335005ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbcArrowFlightSqlClientHandlerBuilderTest.java000066400000000000000000000153531515451615200440660ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.Optional; import org.apache.arrow.driver.jdbc.FlightServerTestExtension; import org.apache.arrow.driver.jdbc.utils.CoreMockedSqlProducers; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; /** Test the behavior of ArrowFlightSqlClientHandler.Builder. */ public class ArrowFlightSqlClientHandlerBuilderTest { @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION = FlightServerTestExtension.createStandardTestExtension( CoreMockedSqlProducers.getLegacyProducer()); private static BufferAllocator allocator; @BeforeAll public static void setup() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterAll public static void tearDown() { allocator.close(); } @Test public void testRetainCookiesOnAuthOff() throws Exception { // Arrange final ArrowFlightSqlClientHandler.Builder rootBuilder = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) .withPort(FLIGHT_SERVER_TEST_EXTENSION.getPort()) .withBufferAllocator(allocator) .withUsername(FlightServerTestExtension.DEFAULT_USER) .withPassword(FlightServerTestExtension.DEFAULT_PASSWORD) .withEncryption(false) .withRetainCookies(true) .withRetainAuth(false); try (ArrowFlightSqlClientHandler rootHandler = rootBuilder.build()) { // Act final ArrowFlightSqlClientHandler.Builder testBuilder = new ArrowFlightSqlClientHandler.Builder(rootBuilder); // Assert assertSame(rootBuilder.cookieFactory, testBuilder.cookieFactory); assertNotSame(rootBuilder.authFactory, testBuilder.authFactory); } } @Test public void testRetainCookiesOffAuthOff() throws Exception { // Arrange final ArrowFlightSqlClientHandler.Builder rootBuilder = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) .withPort(FLIGHT_SERVER_TEST_EXTENSION.getPort()) .withBufferAllocator(allocator) .withUsername(FlightServerTestExtension.DEFAULT_USER) .withPassword(FlightServerTestExtension.DEFAULT_PASSWORD) .withEncryption(false) .withRetainCookies(false) .withRetainAuth(false); try (ArrowFlightSqlClientHandler rootHandler = rootBuilder.build()) { // Act final ArrowFlightSqlClientHandler.Builder testBuilder = new ArrowFlightSqlClientHandler.Builder(rootBuilder); // Assert assertNotSame(rootBuilder.cookieFactory, testBuilder.cookieFactory); assertNotSame(rootBuilder.authFactory, testBuilder.authFactory); } } @Test public void testRetainCookiesOnAuthOn() throws Exception { // Arrange final ArrowFlightSqlClientHandler.Builder rootBuilder = new ArrowFlightSqlClientHandler.Builder() .withHost(FLIGHT_SERVER_TEST_EXTENSION.getHost()) .withPort(FLIGHT_SERVER_TEST_EXTENSION.getPort()) .withBufferAllocator(allocator) .withUsername(FlightServerTestExtension.DEFAULT_USER) .withPassword(FlightServerTestExtension.DEFAULT_PASSWORD) .withEncryption(false) .withRetainCookies(true) .withRetainAuth(true); try (ArrowFlightSqlClientHandler rootHandler = rootBuilder.build()) { // Act final ArrowFlightSqlClientHandler.Builder testBuilder = new ArrowFlightSqlClientHandler.Builder(rootBuilder); // Assert assertSame(rootBuilder.cookieFactory, testBuilder.cookieFactory); assertSame(rootBuilder.authFactory, testBuilder.authFactory); } } @Test public void testDefaults() { final ArrowFlightSqlClientHandler.Builder builder = new ArrowFlightSqlClientHandler.Builder(); // Validate all non-mandatory fields against defaults in ArrowFlightConnectionProperty. assertNull(builder.username); assertNull(builder.password); assertTrue(builder.useEncryption); assertFalse(builder.disableCertificateVerification); assertNull(builder.trustStorePath); assertNull(builder.trustStorePassword); assertTrue(builder.useSystemTrustStore); assertNull(builder.token); assertTrue(builder.retainAuth); assertTrue(builder.retainCookies); assertNull(builder.tlsRootCertificatesPath); assertNull(builder.clientCertificatePath); assertNull(builder.clientKeyPath); assertEquals(Optional.empty(), builder.catalog); assertNull(builder.flightClientCache); assertNull(builder.connectTimeout); assertNull(builder.driverVersion); } @Test public void testCatalog() { ArrowFlightSqlClientHandler.Builder rootBuilder = new ArrowFlightSqlClientHandler.Builder(); rootBuilder.withCatalog(null); assertFalse(rootBuilder.catalog.isPresent()); rootBuilder.withCatalog(""); assertTrue(rootBuilder.catalog.isPresent()); rootBuilder.withCatalog(" "); assertTrue(rootBuilder.catalog.isPresent()); final String noSpaces = "noSpaces"; rootBuilder.withCatalog(noSpaces); assertTrue(rootBuilder.catalog.isPresent()); assertEquals(noSpaces, rootBuilder.catalog.get()); final String nameWithSpaces = " spaces "; rootBuilder.withCatalog(nameWithSpaces); assertTrue(rootBuilder.catalog.isPresent()); assertEquals(nameWithSpaces, rootBuilder.catalog.get()); } } ArrowFlightSqlClientHandlerTest.java000066400000000000000000000073221515451615200425540ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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}, }; } } oauth/000077500000000000000000000000001515451615200346205ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/clientOAuthConfigurationTest.java000066400000000000000000000300121515451615200420670ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } OAuthCredentialWriterTest.java000066400000000000000000000071631515451615200425420ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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)); } } utils/000077500000000000000000000000001515451615200346405ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/clientClientAuthenticationUtilsTest.java000066400000000000000000000224071515451615200435070ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.mockito.Mockito.mock; import java.io.IOException; import java.io.InputStream; import java.lang.reflect.Method; import java.security.KeyStore; import java.security.KeyStoreException; import java.security.NoSuchAlgorithmException; import java.security.cert.Certificate; import java.security.cert.CertificateException; import java.util.Arrays; import java.util.Collections; import java.util.Enumeration; import org.bouncycastle.openssl.jcajce.JcaPEMWriter; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.MockedStatic; import org.mockito.Mockito; import org.mockito.junit.jupiter.MockitoExtension; @ExtendWith(MockitoExtension.class) public class ClientAuthenticationUtilsTest { @Mock KeyStore keyStoreMock; @Test public void testGetCertificatesInputStream() throws IOException, KeyStoreException { JcaPEMWriter pemWriterMock = mock(JcaPEMWriter.class); Certificate certificateMock = mock(Certificate.class); Enumeration alias = Collections.enumeration(Arrays.asList("test1", "test2")); Mockito.when(keyStoreMock.aliases()).thenReturn(alias); Mockito.when(keyStoreMock.isCertificateEntry("test1")).thenReturn(true); Mockito.when(keyStoreMock.getCertificate("test1")).thenReturn(certificateMock); ClientAuthenticationUtils.getCertificatesInputStream(keyStoreMock, pemWriterMock); Mockito.verify(pemWriterMock).writeObject(certificateMock); Mockito.verify(pemWriterMock).flush(); } @Test public void testGetKeyStoreInstance() throws IOException, KeyStoreException, CertificateException, NoSuchAlgorithmException { try (MockedStatic keyStoreMockedStatic = Mockito.mockStatic(KeyStore.class)) { keyStoreMockedStatic .when(() -> ClientAuthenticationUtils.getKeyStoreInstance(Mockito.any())) .thenReturn(keyStoreMock); KeyStore receiveKeyStore = ClientAuthenticationUtils.getKeyStoreInstance("test1"); Mockito.verify(keyStoreMock).load(null, null); assertEquals(receiveKeyStore, keyStoreMock); } } @Test public void testGetDefaultKeyStoreInstancePassword() throws IOException, KeyStoreException, CertificateException, NoSuchAlgorithmException { try (MockedStatic keyStoreMockedStatic = Mockito.mockStatic(KeyStore.class)) { keyStoreMockedStatic .when(() -> ClientAuthenticationUtils.getDefaultKeyStoreInstance("changeit")) .thenReturn(keyStoreMock); KeyStore receiveKeyStore = ClientAuthenticationUtils.getDefaultKeyStoreInstance("changeit"); assertEquals(receiveKeyStore, keyStoreMock); } } @Test public void testGetDefaultKeyStoreInstanceNoPassword() throws IOException, KeyStoreException, CertificateException, NoSuchAlgorithmException { try (MockedStatic keyStoreMockedStatic = Mockito.mockStatic(KeyStore.class)) { keyStoreMockedStatic .when(() -> ClientAuthenticationUtils.getDefaultKeyStoreInstance(null)) .thenReturn(keyStoreMock); KeyStore receiveKeyStore = ClientAuthenticationUtils.getDefaultKeyStoreInstance(null); assertEquals(receiveKeyStore, keyStoreMock); } } @Test public void testGetCertificateInputStreamFromMacSystem() throws IOException, KeyStoreException, CertificateException, NoSuchAlgorithmException { InputStream mock = mock(InputStream.class); try (MockedStatic keyStoreMockedStatic = createKeyStoreStaticMock(); MockedStatic clientAuthenticationUtilsMockedStatic = createClientAuthenticationUtilsStaticMock()) { setOperatingSystemMock(clientAuthenticationUtilsMockedStatic, false, true); keyStoreMockedStatic .when(() -> ClientAuthenticationUtils.getKeyStoreInstance("KeychainStore")) .thenReturn(keyStoreMock); keyStoreMockedStatic .when(() -> ClientAuthenticationUtils.getDefaultKeyStoreInstance("changeit")) .thenReturn(keyStoreMock); clientAuthenticationUtilsMockedStatic .when(ClientAuthenticationUtils::getKeystoreInputStream) .thenCallRealMethod(); keyStoreMockedStatic.when(KeyStore::getDefaultType).thenCallRealMethod(); keyStoreMockedStatic .when(() -> ClientAuthenticationUtils.getCertificatesInputStream(Mockito.any())) .thenReturn(mock); InputStream inputStream = ClientAuthenticationUtils.getCertificateInputStreamFromSystem("changeit"); assertEquals(inputStream, mock); } } @Test public void testGetCertificateInputStreamFromWindowsSystem() throws IOException, KeyStoreException, CertificateException, NoSuchAlgorithmException { InputStream mock = mock(InputStream.class); try (MockedStatic keyStoreMockedStatic = createKeyStoreStaticMock(); MockedStatic clientAuthenticationUtilsMockedStatic = createClientAuthenticationUtilsStaticMock()) { setOperatingSystemMock(clientAuthenticationUtilsMockedStatic, true, false); keyStoreMockedStatic .when(() -> ClientAuthenticationUtils.getKeyStoreInstance("Windows-ROOT")) .thenReturn(keyStoreMock); keyStoreMockedStatic .when(() -> ClientAuthenticationUtils.getKeyStoreInstance("Windows-MY")) .thenReturn(keyStoreMock); keyStoreMockedStatic .when(() -> ClientAuthenticationUtils.getCertificatesInputStream(Mockito.any())) .thenReturn(mock); InputStream inputStream = ClientAuthenticationUtils.getCertificateInputStreamFromSystem("test"); assertEquals(inputStream, mock); } } @Test public void testGetCertificateInputStreamFromLinuxSystem() throws IOException, KeyStoreException, CertificateException, NoSuchAlgorithmException { InputStream mock = mock(InputStream.class); try (MockedStatic keyStoreMockedStatic = createKeyStoreStaticMock(); MockedStatic clientAuthenticationUtilsMockedStatic = createClientAuthenticationUtilsStaticMock()) { setOperatingSystemMock(clientAuthenticationUtilsMockedStatic, false, false); keyStoreMockedStatic .when(() -> ClientAuthenticationUtils.getCertificatesInputStream(Mockito.any())) .thenReturn(mock); keyStoreMockedStatic .when(() -> ClientAuthenticationUtils.getDefaultKeyStoreInstance(Mockito.any())) .thenReturn(keyStoreMock); clientAuthenticationUtilsMockedStatic .when(ClientAuthenticationUtils::getKeystoreInputStream) .thenCallRealMethod(); keyStoreMockedStatic.when(KeyStore::getDefaultType).thenCallRealMethod(); InputStream inputStream = ClientAuthenticationUtils.getCertificateInputStreamFromSystem("changeit"); assertEquals(inputStream, mock); inputStream = ClientAuthenticationUtils.getCertificateInputStreamFromSystem(null); assertEquals(inputStream, mock); } } private MockedStatic createKeyStoreStaticMock() { return Mockito.mockStatic( KeyStore.class, invocationOnMock -> { Method method = invocationOnMock.getMethod(); if (method.getName().equals("getInstance")) { return invocationOnMock.callRealMethod(); } return method.invoke(invocationOnMock.getMock(), invocationOnMock.getArguments()); }); } private MockedStatic createClientAuthenticationUtilsStaticMock() { return Mockito.mockStatic( ClientAuthenticationUtils.class, invocationOnMock -> { Method method = invocationOnMock.getMethod(); if (method.getName().equals("getCertificateInputStreamFromSystem")) { return invocationOnMock.callRealMethod(); } return method.invoke(invocationOnMock.getMock(), invocationOnMock.getArguments()); }); } private void setOperatingSystemMock( MockedStatic clientAuthenticationUtilsMockedStatic, boolean isWindows, boolean isMac) { clientAuthenticationUtilsMockedStatic.when(ClientAuthenticationUtils::isMac).thenReturn(isMac); assertEquals(ClientAuthenticationUtils.isMac(), isMac); clientAuthenticationUtilsMockedStatic .when(ClientAuthenticationUtils::isWindows) .thenReturn(isWindows); assertEquals(ClientAuthenticationUtils.isWindows(), isWindows); } } FlightClientCacheTest.java000066400000000000000000000037001515451615200416430ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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())); } } FlightLocationQueueTest.java000066400000000000000000000047641515451615200422710ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } converter/000077500000000000000000000000001515451615200342315ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbcimpl/000077500000000000000000000000001515451615200351725ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/converterUuidAvaticaParameterConverterTest.java000066400000000000000000000135161515451615200446330ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/converter/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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())); } } utils/000077500000000000000000000000001515451615200333625ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbcAccessorTestUtils.java000066400000000000000000000115461515451615200376570ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; import java.sql.SQLException; import java.util.ArrayList; import java.util.List; import java.util.function.Consumer; import java.util.function.Function; import java.util.function.IntSupplier; import java.util.function.Supplier; import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor; import org.apache.arrow.vector.ValueVector; import org.hamcrest.Matcher; public class AccessorTestUtils { @FunctionalInterface public interface CheckedFunction { R apply(T t) throws SQLException; } public interface AccessorSupplier { T supply(ValueVector vector, IntSupplier getCurrentRow); } public interface AccessorConsumer { void accept(T accessor, int currentRow) throws Exception; } public interface MatcherGetter { Matcher get(T accessor, int currentRow); } public static class Cursor { int currentRow = 0; int limit; public Cursor(int limit) { this.limit = limit; } public void next() { currentRow++; } boolean hasNext() { return currentRow < limit; } public int getCurrentRow() { return currentRow; } } public static class AccessorIterator { private final AccessorSupplier accessorSupplier; public AccessorIterator(AccessorSupplier accessorSupplier) { this.accessorSupplier = accessorSupplier; } public void iterate(ValueVector vector, AccessorConsumer accessorConsumer) throws Exception { int valueCount = vector.getValueCount(); if (valueCount == 0) { throw new IllegalArgumentException("Vector is empty"); } Cursor cursor = new Cursor(valueCount); T accessor = accessorSupplier.supply(vector, cursor::getCurrentRow); while (cursor.hasNext()) { accessorConsumer.accept(accessor, cursor.getCurrentRow()); cursor.next(); } } public void iterate(ValueVector vector, Consumer accessorConsumer) throws Exception { iterate(vector, (accessor, currentRow) -> accessorConsumer.accept(accessor)); } public List toList(ValueVector vector) throws Exception { List result = new ArrayList<>(); iterate(vector, (accessor, currentRow) -> result.add(accessor.getObject())); return result; } public void assertAccessorGetter( ValueVector vector, CheckedFunction getter, MatcherGetter matcherGetter) throws Exception { iterate( vector, (accessor, currentRow) -> { R object = getter.apply(accessor); boolean wasNull = accessor.wasNull(); assertThat(object, matcherGetter.get(accessor, currentRow)); assertThat(wasNull, is(accessor.getObject() == null)); }); } public void assertAccessorGetterThrowingException( ValueVector vector, CheckedFunction getter) throws Exception { iterate( vector, (accessor, currentRow) -> ThrowableAssertionUtils.simpleAssertThrowableClass( SQLException.class, () -> getter.apply(accessor))); } public void assertAccessorGetter( ValueVector vector, CheckedFunction getter, Function> matcherGetter) throws Exception { assertAccessorGetter(vector, getter, (accessor, currentRow) -> matcherGetter.apply(accessor)); } public void assertAccessorGetter( ValueVector vector, CheckedFunction getter, Supplier> matcherGetter) throws Exception { assertAccessorGetter(vector, getter, (accessor, currentRow) -> matcherGetter.get()); } public void assertAccessorGetter( ValueVector vector, CheckedFunction getter, Matcher matcher) throws Exception { assertAccessorGetter(vector, getter, (accessor, currentRow) -> matcher); } } } ArrowFlightConnectionConfigImplTest.java000066400000000000000000000133461515451615200433140ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; 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; import java.util.stream.Stream; import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; public final class ArrowFlightConnectionConfigImplTest { private static final Random RANDOM = new Random(12L); private Properties properties; private ArrowFlightConnectionConfigImpl arrowFlightConnectionConfig; public ArrowFlightConnectionProperty property; public Object value; public Function arrowFlightConnectionConfigFunction; @BeforeEach public void setUp() { properties = new Properties(); arrowFlightConnectionConfig = new ArrowFlightConnectionConfigImpl(properties); } @ParameterizedTest @MethodSource("provideParameters") public void testGetProperty( ArrowFlightConnectionProperty property, Object value, Object expected, Function configFunction) { properties.put(property.camelName(), value); arrowFlightConnectionConfigFunction = configFunction; 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, 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, useEncryption, useEncryption, (Function) ArrowFlightConnectionConfigImpl::useEncryption), Arguments.of( THREAD_POOL_SIZE, 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::useClientCache)); } } ArrowFlightConnectionPropertyTest.java000066400000000000000000000071641515451615200431120ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import static org.apache.arrow.util.AutoCloseables.close; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assumptions.assumeTrue; import static org.mockito.MockitoAnnotations.openMocks; import java.util.ArrayList; import java.util.List; import java.util.Properties; import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.extension.ExtendWith; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; @ExtendWith(MockitoExtension.class) public final class ArrowFlightConnectionPropertyTest { @Mock public Properties properties; private AutoCloseable mockitoResource; public ArrowFlightConnectionProperty arrowFlightConnectionProperty; @BeforeEach public void setUp() { mockitoResource = openMocks(this); } @AfterEach public void tearDown() throws Exception { close(mockitoResource); } @ParameterizedTest @MethodSource("provideParameters") public void testWrapIsUnsupported(ArrowFlightConnectionProperty property) { this.arrowFlightConnectionProperty = property; ThrowableAssertionUtils.simpleAssertThrowableClass( UnsupportedOperationException.class, () -> arrowFlightConnectionProperty.wrap(properties)); } @ParameterizedTest @MethodSource("provideParameters") public void testRequiredPropertyThrows(ArrowFlightConnectionProperty property) { this.arrowFlightConnectionProperty = property; assumeTrue(arrowFlightConnectionProperty.required()); ThrowableAssertionUtils.simpleAssertThrowableClass( IllegalStateException.class, () -> arrowFlightConnectionProperty.get(new Properties())); } @ParameterizedTest @MethodSource("provideParameters") public void testOptionalPropertyReturnsDefault(ArrowFlightConnectionProperty property) { this.arrowFlightConnectionProperty = property; assumeTrue(!arrowFlightConnectionProperty.required()); assertEquals( arrowFlightConnectionProperty.defaultValue(), arrowFlightConnectionProperty.get(new Properties())); } public static List provideParameters() { final ArrowFlightConnectionProperty[] arrowFlightConnectionProperties = ArrowFlightConnectionProperty.values(); final List parameters = new ArrayList<>(arrowFlightConnectionProperties.length); for (final ArrowFlightConnectionProperty arrowFlightConnectionProperty : arrowFlightConnectionProperties) { parameters.add(Arguments.of(arrowFlightConnectionProperty)); } return parameters; } } ConnectionWrapperTest.java000066400000000000000000000416711515451615200405360ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import static java.lang.String.format; import static java.util.stream.IntStream.range; import static org.hamcrest.CoreMatchers.allOf; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.CoreMatchers.nullValue; import static org.hamcrest.CoreMatchers.sameInstance; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import java.sql.Connection; import java.sql.ResultSet; import java.sql.SQLClientInfoException; import java.sql.SQLException; import java.sql.Statement; import java.util.Arrays; import java.util.Random; import org.apache.arrow.driver.jdbc.ArrowFlightConnection; import org.apache.arrow.util.AutoCloseables; import org.apache.calcite.avatica.AvaticaConnection; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; @ExtendWith(MockitoExtension.class) public final class ConnectionWrapperTest { private static final String SCHEMA_NAME = "SCHEMA"; private static final String PLACEHOLDER_QUERY = "SELECT * FROM DOES_NOT_MATTER"; private static final int[] COLUMN_INDICES = range(0, 10).toArray(); private static final String[] COLUMN_NAMES = Arrays.stream(COLUMN_INDICES).mapToObj(i -> format("col%d", i)).toArray(String[]::new); private static final String TYPE_NAME = "TYPE_NAME"; private static final String SAVEPOINT_NAME = "SAVEPOINT"; private static final String CLIENT_INFO = "CLIENT_INFO"; private static final int RESULT_SET_TYPE = ResultSet.TYPE_FORWARD_ONLY; private static final int RESULT_SET_CONCURRENCY = ResultSet.CONCUR_READ_ONLY; private static final int RESULT_SET_HOLDABILITY = ResultSet.HOLD_CURSORS_OVER_COMMIT; private static final int GENERATED_KEYS = Statement.NO_GENERATED_KEYS; private static final Random RANDOM = new Random(Long.MAX_VALUE); private static final int TIMEOUT = RANDOM.nextInt(Integer.MAX_VALUE); @Mock public AvaticaConnection underlyingConnection; private ConnectionWrapper connectionWrapper; @BeforeEach public void setUp() { connectionWrapper = new ConnectionWrapper(underlyingConnection); } @AfterEach public void tearDown() throws Exception { AutoCloseables.close(connectionWrapper, underlyingConnection); } @Test public void testUnwrappingUnderlyingConnectionShouldReturnUnderlyingConnection() { assertThat( assertDoesNotThrow(() -> connectionWrapper.unwrap(Object.class)), is(sameInstance(underlyingConnection))); assertThat( assertDoesNotThrow(() -> connectionWrapper.unwrap(Connection.class)), is(sameInstance(underlyingConnection))); assertThat( assertDoesNotThrow(() -> connectionWrapper.unwrap(AvaticaConnection.class)), is(sameInstance(underlyingConnection))); ThrowableAssertionUtils.simpleAssertThrowableClass( ClassCastException.class, () -> connectionWrapper.unwrap(ArrowFlightConnection.class)); ThrowableAssertionUtils.simpleAssertThrowableClass( ClassCastException.class, () -> connectionWrapper.unwrap(ConnectionWrapper.class)); } @Test public void testCreateStatementShouldCreateStatementFromUnderlyingConnection() throws SQLException { assertThat( connectionWrapper.createStatement(), is(sameInstance(verify(underlyingConnection, times(1)).createStatement()))); assertThat( connectionWrapper.createStatement( RESULT_SET_TYPE, RESULT_SET_CONCURRENCY, RESULT_SET_HOLDABILITY), is( verify(underlyingConnection, times(1)) .createStatement(RESULT_SET_TYPE, RESULT_SET_CONCURRENCY, RESULT_SET_HOLDABILITY))); assertThat( connectionWrapper.createStatement(RESULT_SET_TYPE, RESULT_SET_CONCURRENCY), is( verify(underlyingConnection, times(1)) .createStatement(RESULT_SET_TYPE, RESULT_SET_CONCURRENCY))); } @Test public void testPrepareStatementShouldPrepareStatementFromUnderlyingConnection() throws SQLException { assertThat( connectionWrapper.prepareStatement(PLACEHOLDER_QUERY), is( sameInstance( verify(underlyingConnection, times(1)).prepareStatement(PLACEHOLDER_QUERY)))); assertThat( connectionWrapper.prepareStatement(PLACEHOLDER_QUERY, COLUMN_INDICES), is( allOf( sameInstance( verify(underlyingConnection, times(1)) .prepareStatement(PLACEHOLDER_QUERY, COLUMN_INDICES)), nullValue()))); assertThat( connectionWrapper.prepareStatement(PLACEHOLDER_QUERY, COLUMN_NAMES), is( allOf( sameInstance( verify(underlyingConnection, times(1)) .prepareStatement(PLACEHOLDER_QUERY, COLUMN_NAMES)), nullValue()))); assertThat( connectionWrapper.prepareStatement( PLACEHOLDER_QUERY, RESULT_SET_TYPE, RESULT_SET_CONCURRENCY), is( allOf( sameInstance( verify(underlyingConnection, times(1)) .prepareStatement( PLACEHOLDER_QUERY, RESULT_SET_TYPE, RESULT_SET_CONCURRENCY)), nullValue()))); assertThat( connectionWrapper.prepareStatement(PLACEHOLDER_QUERY, GENERATED_KEYS), is( allOf( sameInstance( verify(underlyingConnection, times(1)) .prepareStatement(PLACEHOLDER_QUERY, GENERATED_KEYS)), nullValue()))); } @Test public void testPrepareCallShouldPrepareCallFromUnderlyingConnection() throws SQLException { assertThat( connectionWrapper.prepareCall(PLACEHOLDER_QUERY), is(sameInstance(verify(underlyingConnection, times(1)).prepareCall(PLACEHOLDER_QUERY)))); assertThat( connectionWrapper.prepareCall(PLACEHOLDER_QUERY, RESULT_SET_TYPE, RESULT_SET_CONCURRENCY), is( verify(underlyingConnection, times(1)) .prepareCall(PLACEHOLDER_QUERY, RESULT_SET_TYPE, RESULT_SET_CONCURRENCY))); } @Test public void testNativeSqlShouldGetNativeSqlFromUnderlyingConnection() throws SQLException { assertThat( connectionWrapper.nativeSQL(PLACEHOLDER_QUERY), is(sameInstance(verify(underlyingConnection, times(1)).nativeSQL(PLACEHOLDER_QUERY)))); } @Test public void testSetAutoCommitShouldSetAutoCommitInUnderlyingConnection() throws SQLException { connectionWrapper.setAutoCommit(true); verify(underlyingConnection, times(1)).setAutoCommit(true); connectionWrapper.setAutoCommit(false); verify(underlyingConnection, times(1)).setAutoCommit(false); } @Test public void testGetAutoCommitShouldGetAutoCommitFromUnderlyingConnection() throws SQLException { assertThat( connectionWrapper.getAutoCommit(), is(verify(underlyingConnection, times(1)).getAutoCommit())); } @Test public void testCommitShouldCommitToUnderlyingConnection() throws SQLException { connectionWrapper.commit(); verify(underlyingConnection, times(1)).commit(); } @Test public void testRollbackShouldRollbackFromUnderlyingConnection() throws SQLException { connectionWrapper.rollback(); verify(underlyingConnection, times(1)).rollback(); } @Test public void testCloseShouldCloseUnderlyingConnection() throws SQLException { connectionWrapper.close(); verify(underlyingConnection, times(1)).close(); } @Test public void testIsClosedShouldGetStatusFromUnderlyingConnection() throws SQLException { assertThat(connectionWrapper.isClosed(), is(verify(underlyingConnection, times(1)).isClosed())); } @Test public void testGetMetadataShouldGetMetadataFromUnderlyingConnection() throws SQLException { assertThat( connectionWrapper.getMetaData(), is(verify(underlyingConnection, times(1)).getMetaData())); } @Test public void testSetReadOnlyShouldSetUnderlyingConnectionAsReadOnly() throws SQLException { connectionWrapper.setReadOnly(false); verify(underlyingConnection, times(1)).setReadOnly(false); connectionWrapper.setReadOnly(true); verify(underlyingConnection, times(1)).setReadOnly(true); } @Test public void testSetIsReadOnlyShouldGetStatusFromUnderlyingConnection() throws SQLException { assertThat(connectionWrapper.isReadOnly(), is(verify(underlyingConnection).isReadOnly())); } @Test public void testSetCatalogShouldSetCatalogInUnderlyingConnection() throws SQLException { final String catalog = "CATALOG"; connectionWrapper.setCatalog(catalog); verify(underlyingConnection, times(1)).setCatalog(catalog); } @Test public void testGetCatalogShouldGetCatalogFromUnderlyingConnection() throws SQLException { assertThat( connectionWrapper.getCatalog(), is(allOf(sameInstance(verify(underlyingConnection, times(1)).getCatalog()), nullValue()))); } @Test public void setTransactionIsolationShouldSetUnderlyingTransactionIsolation() throws SQLException { final int transactionIsolation = Connection.TRANSACTION_NONE; connectionWrapper.setTransactionIsolation(Connection.TRANSACTION_NONE); verify(underlyingConnection, times(1)).setTransactionIsolation(transactionIsolation); } @Test public void getTransactionIsolationShouldGetUnderlyingConnectionIsolation() throws SQLException { assertThat( connectionWrapper.getTransactionIsolation(), is(equalTo(verify(underlyingConnection, times(1)).getTransactionIsolation()))); } @Test public void getWarningShouldGetWarningsFromUnderlyingConnection() throws SQLException { assertThat( connectionWrapper.getWarnings(), is(allOf(sameInstance(verify(underlyingConnection, times(1)).getWarnings()), nullValue()))); } @Test public void testClearWarningShouldClearWarningsFromUnderlyingConnection() throws SQLException { connectionWrapper.clearWarnings(); verify(underlyingConnection, times(1)).clearWarnings(); } @Test public void getTypeMapShouldGetTypeMapFromUnderlyingConnection() throws SQLException { when(underlyingConnection.getTypeMap()).thenReturn(null); assertThat( connectionWrapper.getTypeMap(), is(verify(underlyingConnection, times(1)).getTypeMap())); } @Test public void testSetTypeMapShouldSetTypeMapFromUnderlyingConnection() throws SQLException { connectionWrapper.setTypeMap(null); verify(underlyingConnection, times(1)).setTypeMap(null); } @Test public void testSetHoldabilityShouldSetUnderlyingConnection() throws SQLException { connectionWrapper.setHoldability(RESULT_SET_HOLDABILITY); verify(underlyingConnection, times(1)).setHoldability(RESULT_SET_HOLDABILITY); } @Test public void testGetHoldabilityShouldGetHoldabilityFromUnderlyingConnection() throws SQLException { assertThat( connectionWrapper.getHoldability(), is(equalTo(verify(underlyingConnection, times(1)).getHoldability()))); } @Test public void testSetSavepointShouldSetSavepointInUnderlyingConnection() throws SQLException { assertThat( connectionWrapper.setSavepoint(), is( allOf( sameInstance(verify(underlyingConnection, times(1)).setSavepoint()), nullValue()))); assertThat( connectionWrapper.setSavepoint(SAVEPOINT_NAME), is(sameInstance(verify(underlyingConnection, times(1)).setSavepoint(SAVEPOINT_NAME)))); } @Test public void testRollbackShouldRollbackInUnderlyingConnection() throws SQLException { connectionWrapper.rollback(null); verify(underlyingConnection, times(1)).rollback(null); } @Test public void testReleaseSavepointShouldReleaseSavepointFromUnderlyingConnection() throws SQLException { connectionWrapper.releaseSavepoint(null); verify(underlyingConnection, times(1)).releaseSavepoint(null); } @Test public void testCreateClobShouldCreateClobFromUnderlyingConnection() throws SQLException { assertThat( connectionWrapper.createClob(), is(allOf(sameInstance(verify(underlyingConnection, times(1)).createClob()), nullValue()))); } @Test public void testCreateBlobShouldCreateBlobFromUnderlyingConnection() throws SQLException { assertThat( connectionWrapper.createBlob(), is(allOf(sameInstance(verify(underlyingConnection, times(1)).createBlob()), nullValue()))); } @Test public void testCreateNClobShouldCreateNClobFromUnderlyingConnection() throws SQLException { assertThat( connectionWrapper.createNClob(), is(allOf(sameInstance(verify(underlyingConnection, times(1)).createNClob()), nullValue()))); } @Test public void testCreateSQLXMLShouldCreateSQLXMLFromUnderlyingConnection() throws SQLException { assertThat( connectionWrapper.createSQLXML(), is( allOf( sameInstance(verify(underlyingConnection, times(1)).createSQLXML()), nullValue()))); } @Test public void testIsValidShouldReturnWhetherUnderlyingConnectionIsValid() throws SQLException { assertThat( connectionWrapper.isValid(TIMEOUT), is(verify(underlyingConnection, times(1)).isValid(TIMEOUT))); } @Test public void testSetClientInfoShouldSetClientInfoInUnderlyingConnection() throws SQLClientInfoException { connectionWrapper.setClientInfo(null); verify(underlyingConnection, times(1)).setClientInfo(null); } @Test public void testGetClientInfoShouldGetClientInfoFromUnderlyingConnection() throws SQLException { assertThat( connectionWrapper.getClientInfo(CLIENT_INFO), is( allOf( sameInstance(verify(underlyingConnection, times(1)).getClientInfo(CLIENT_INFO)), nullValue()))); assertThat( connectionWrapper.getClientInfo(), is( allOf( sameInstance(verify(underlyingConnection, times(1)).getClientInfo()), nullValue()))); } @Test public void testCreateArrayOfShouldCreateArrayFromUnderlyingConnection() throws SQLException { final Object[] elements = range(0, 100).boxed().toArray(); assertThat( connectionWrapper.createArrayOf(TYPE_NAME, elements), is( allOf( sameInstance( verify(underlyingConnection, times(1)).createArrayOf(TYPE_NAME, elements)), nullValue()))); } @Test public void testCreateStructShouldCreateStructFromUnderlyingConnection() throws SQLException { final Object[] attributes = range(0, 120).boxed().toArray(); assertThat( connectionWrapper.createStruct(TYPE_NAME, attributes), is( allOf( sameInstance( verify(underlyingConnection, times(1)).createStruct(TYPE_NAME, attributes)), nullValue()))); } @Test public void testSetSchemaShouldSetSchemaInUnderlyingConnection() throws SQLException { connectionWrapper.setSchema(SCHEMA_NAME); verify(underlyingConnection, times(1)).setSchema(SCHEMA_NAME); } @Test public void testGetSchemaShouldGetSchemaFromUnderlyingConnection() throws SQLException { assertThat( connectionWrapper.getSchema(), is(allOf(sameInstance(verify(underlyingConnection, times(1)).getSchema()), nullValue()))); } @Test public void testAbortShouldAbortUnderlyingConnection() throws SQLException { connectionWrapper.abort(null); verify(underlyingConnection, times(1)).abort(null); } @Test public void testSetNetworkTimeoutShouldSetNetworkTimeoutInUnderlyingConnection() throws SQLException { connectionWrapper.setNetworkTimeout(null, TIMEOUT); verify(underlyingConnection, times(1)).setNetworkTimeout(null, TIMEOUT); } @Test public void testGetNetworkTimeoutShouldGetNetworkTimeoutFromUnderlyingConnection() throws SQLException { assertThat( connectionWrapper.getNetworkTimeout(), is(equalTo(verify(underlyingConnection, times(1)).getNetworkTimeout()))); } } ConvertUtilsTest.java000066400000000000000000000163431515451615200375350ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; import com.google.common.collect.ImmutableList; import java.util.List; import org.apache.arrow.flight.sql.FlightSqlColumnMetadata; 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.calcite.avatica.ColumnMetaData; import org.apache.calcite.avatica.proto.Common; import org.junit.jupiter.api.Test; public class ConvertUtilsTest { @Test public void testShouldSetOnColumnMetaDataBuilder() { final Common.ColumnMetaData.Builder builder = Common.ColumnMetaData.newBuilder(); final FlightSqlColumnMetadata expectedColumnMetaData = new FlightSqlColumnMetadata.Builder() .catalogName("catalog1") .schemaName("schema1") .tableName("table1") .isAutoIncrement(true) .isCaseSensitive(true) .isReadOnly(true) .isSearchable(true) .precision(20) .scale(10) .remarks("test column") .build(); ConvertUtils.setOnColumnMetaDataBuilder(builder, expectedColumnMetaData.getMetadataMap()); assertBuilder(builder, expectedColumnMetaData); } @Test public void testShouldConvertArrowFieldsToColumnMetaDataList() { final List listField = ImmutableList.of( new Field( "col1", new FieldType( true, ArrowType.Utf8.INSTANCE, null, new FlightSqlColumnMetadata.Builder() .catalogName("catalog1") .schemaName("schema1") .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 = ImmutableList.of( ColumnMetaData.fromProto( Common.ColumnMetaData.newBuilder() .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 = ConvertUtils.convertArrowFieldsToColumnMetaDataList(listField); assertColumnMetaData(expectedColumnMetaData, actualColumnMetaData); } private void assertColumnMetaData( final List expected, final List actual) { assertThat(expected.size(), equalTo(actual.size())); int size = expected.size(); for (int i = 0; i < size; i++) { final ColumnMetaData expectedColumnMetaData = expected.get(i); final ColumnMetaData actualColumnMetaData = actual.get(i); 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)); assertThat(expectedColumnMetaData.scale, equalTo(actualColumnMetaData.scale)); assertThat(expectedColumnMetaData.caseSensitive, equalTo(actualColumnMetaData.caseSensitive)); assertThat(expectedColumnMetaData.searchable, equalTo(actualColumnMetaData.searchable)); } } private void assertBuilder( final Common.ColumnMetaData.Builder builder, final FlightSqlColumnMetadata flightSqlColumnMetaData) { final Integer precision = flightSqlColumnMetaData.getPrecision(); final Integer scale = flightSqlColumnMetaData.getScale(); assertThat(flightSqlColumnMetaData.getCatalogName(), equalTo(builder.getCatalogName())); assertThat(flightSqlColumnMetaData.getSchemaName(), equalTo(builder.getSchemaName())); assertThat(flightSqlColumnMetaData.getTableName(), equalTo(builder.getTableName())); assertThat(flightSqlColumnMetaData.isAutoIncrement(), equalTo(builder.getAutoIncrement())); assertThat(flightSqlColumnMetaData.isCaseSensitive(), equalTo(builder.getCaseSensitive())); assertThat(flightSqlColumnMetaData.isSearchable(), equalTo(builder.getSearchable())); 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())); } } CoreMockedSqlProducers.java000066400000000000000000000477351515451615200406270ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import static java.lang.String.format; 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.assertAll; import com.google.common.collect.ImmutableList; import java.sql.Date; import java.sql.ResultSet; 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; import org.apache.arrow.flight.sql.FlightSqlColumnMetadata; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BigIntVector; 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; 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.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. public final class CoreMockedSqlProducers { public static final String LEGACY_REGULAR_SQL_CMD = "SELECT * FROM TEST"; public static final String LEGACY_METADATA_SQL_CMD = "SELECT * FROM METADATA"; 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. } /** * Gets the {@link MockFlightSqlProducer} for legacy tests and backward compatibility. * * @return a new producer. */ public static MockFlightSqlProducer getLegacyProducer() { final MockFlightSqlProducer producer = new MockFlightSqlProducer(); addLegacyRegularSqlCmdSupport(producer); 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( ImmutableList.of( new Field("ID", new FieldType(true, new ArrowType.Int(64, true), null), null))); final List> resultProducers = new ArrayList<>(); Consumer dataRoot = listener -> { try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); final VectorSchemaRoot root = VectorSchemaRoot.create(querySchema, allocator)) { root.allocateNew(); root.setRowCount(0); listener.start(root); listener.putNext(); // empty root ((BigIntVector) root.getVector("ID")).setSafe(0, 100L); root.setRowCount(1); listener.putNext(); // data root root.clear(); root.setRowCount(0); listener.putNext(); // empty root ((BigIntVector) root.getVector("ID")).setSafe(0, 100L); root.setRowCount(1); listener.putNext(); // data root } finally { listener.completed(); } }; resultProducers.add(dataRoot); producer.addSelectQuery(LEGACY_REGULAR_WITH_EMPTY_SQL_CMD, querySchema, resultProducers); } private static void addLegacyRegularSqlCmdSupport(final MockFlightSqlProducer producer) { final Schema querySchema = new Schema( ImmutableList.of( new Field("ID", new FieldType(true, new ArrowType.Int(64, true), null), null), new Field("Name", new FieldType(true, new ArrowType.Utf8(), null), null), new Field("Age", new FieldType(true, new ArrowType.Int(32, false), null), null), new Field( "Salary", new FieldType( true, new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE), null), null), new Field( "Hire Date", new FieldType(true, new ArrowType.Date(DateUnit.DAY), null), null), new Field( "Last Sale", new FieldType(true, new ArrowType.Timestamp(TimeUnit.MILLISECOND, null), null), null))); final List> resultProducers = new ArrayList<>(); IntStream.range(0, 10) .forEach( page -> { resultProducers.add( listener -> { final int rowsPerPage = 5000; try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); final VectorSchemaRoot root = VectorSchemaRoot.create(querySchema, allocator)) { root.allocateNew(); listener.start(root); int batchSize = 500; int indexOnBatch = 0; int resultsOffset = page * rowsPerPage; for (int i = 0; i < rowsPerPage; i++) { ((BigIntVector) root.getVector("ID")) .setSafe( indexOnBatch, (long) Integer.MAX_VALUE + 1 + i + resultsOffset); ((VarCharVector) root.getVector("Name")) .setSafe(indexOnBatch, new Text("Test Name #" + (resultsOffset + i))); ((UInt4Vector) root.getVector("Age")) .setSafe(indexOnBatch, (int) Short.MAX_VALUE + 1 + i + resultsOffset); ((Float8Vector) root.getVector("Salary")) .setSafe( indexOnBatch, Math.scalb((double) (i + resultsOffset) / 2, i + resultsOffset)); ((DateDayVector) root.getVector("Hire Date")) .setSafe(indexOnBatch, i + resultsOffset); ((TimeStampMilliVector) root.getVector("Last Sale")) .setSafe(indexOnBatch, Long.MAX_VALUE - i - resultsOffset); indexOnBatch++; if (indexOnBatch == batchSize) { root.setRowCount(indexOnBatch); if (listener.isCancelled()) { return; } listener.putNext(); root.allocateNew(); indexOnBatch = 0; } } if (listener.isCancelled()) { return; } root.setRowCount(indexOnBatch); listener.putNext(); } finally { listener.completed(); } }); }); producer.addSelectQuery(LEGACY_REGULAR_SQL_CMD, querySchema, resultProducers); } private static void addLegacyMetadataSqlCmdSupport(final MockFlightSqlProducer producer) { final Schema metadataSchema = new Schema( ImmutableList.of( new Field( "integer0", new FieldType( true, new ArrowType.Int(64, true), null, new FlightSqlColumnMetadata.Builder() .catalogName("CATALOG_NAME_1") .schemaName("SCHEMA_NAME_1") .tableName("TABLE_NAME_1") .typeName("TYPE_NAME_1") .precision(10) .scale(0) .isAutoIncrement(true) .isCaseSensitive(false) .isReadOnly(true) .isSearchable(true) .build() .getMetadataMap()), null), new Field( "string1", new FieldType( true, new ArrowType.Utf8(), null, new FlightSqlColumnMetadata.Builder() .catalogName("CATALOG_NAME_2") .schemaName("SCHEMA_NAME_2") .tableName("TABLE_NAME_2") .typeName("TYPE_NAME_2") .precision(65535) .scale(0) .isAutoIncrement(false) .isCaseSensitive(true) .isReadOnly(false) .isSearchable(true) .build() .getMetadataMap()), null), new Field( "float2", new FieldType( true, new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE), null, new FlightSqlColumnMetadata.Builder() .catalogName("CATALOG_NAME_3") .schemaName("SCHEMA_NAME_3") .tableName("TABLE_NAME_3") .typeName("TYPE_NAME_3") .precision(15) .scale(20) .isAutoIncrement(false) .isCaseSensitive(false) .isReadOnly(false) .isSearchable(true) .build() .getMetadataMap()), null))); final Consumer formula = listener -> { try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); final VectorSchemaRoot root = VectorSchemaRoot.create(metadataSchema, allocator)) { root.allocateNew(); ((BigIntVector) root.getVector("integer0")).setSafe(0, 1); ((VarCharVector) root.getVector("string1")).setSafe(0, new Text("teste")); ((Float4Vector) root.getVector("float2")).setSafe(0, (float) 4.1); root.setRowCount(1); listener.start(root); listener.putNext(); } finally { listener.completed(); } }; producer.addSelectQuery( LEGACY_METADATA_SQL_CMD, metadataSchema, Collections.singletonList(formula)); } private static void addLegacyCancellationSqlCmdSupport(final MockFlightSqlProducer producer) { producer.addSelectQuery( LEGACY_CANCELLATION_SQL_CMD, new Schema( Collections.singletonList( new Field( "integer0", new FieldType(true, new ArrowType.Int(64, true), null), null))), Collections.singletonList( listener -> { // Should keep hanging until canceled. })); } /** * Asserts that the values in the provided {@link ResultSet} are expected for the legacy {@link * MockFlightSqlProducer}. * * @param resultSet the result set. * @throws SQLException on error. */ public static void assertLegacyRegularSqlResultSet(final ResultSet resultSet) throws SQLException { final int expectedRowCount = 50_000; final long[] expectedIds = new long[expectedRowCount]; final List expectedNames = new ArrayList<>(expectedRowCount); final int[] expectedAges = new int[expectedRowCount]; final double[] expectedSalaries = new double[expectedRowCount]; final List expectedHireDates = new ArrayList<>(expectedRowCount); final List expectedLastSales = new ArrayList<>(expectedRowCount); final long[] actualIds = new long[expectedRowCount]; final List actualNames = new ArrayList<>(expectedRowCount); final int[] actualAges = new int[expectedRowCount]; final double[] actualSalaries = new double[expectedRowCount]; final List actualHireDates = new ArrayList<>(expectedRowCount); final List actualLastSales = new ArrayList<>(expectedRowCount); int actualRowCount = 0; for (; resultSet.next(); actualRowCount++) { expectedIds[actualRowCount] = (long) Integer.MAX_VALUE + 1 + actualRowCount; expectedNames.add(format("Test Name #%d", actualRowCount)); expectedAges[actualRowCount] = (int) Short.MAX_VALUE + 1 + actualRowCount; expectedSalaries[actualRowCount] = Math.scalb((double) actualRowCount / 2, actualRowCount); expectedHireDates.add(new Date(86_400_000L * actualRowCount)); expectedLastSales.add(new Timestamp(Long.MAX_VALUE - actualRowCount)); actualIds[actualRowCount] = (long) resultSet.getObject(1); actualNames.add((String) resultSet.getObject(2)); actualAges[actualRowCount] = (int) resultSet.getObject(3); actualSalaries[actualRowCount] = (double) resultSet.getObject(4); actualHireDates.add((Date) resultSet.getObject(5)); actualLastSales.add((Timestamp) resultSet.getObject(6)); } final int finalActualRowCount = actualRowCount; assertAll( "ResultSet values are as expected", () -> assertThat(finalActualRowCount, is(equalTo(expectedRowCount))), () -> assertThat(actualIds, is(expectedIds)), () -> assertThat(actualNames, is(expectedNames)), () -> assertThat(actualAges, is(expectedAges)), () -> assertThat(actualSalaries, is(expectedSalaries)), () -> assertThat(actualHireDates, is(expectedHireDates)), () -> assertThat(actualLastSales, is(expectedLastSales))); } } DateTimeUtilsTest.java000066400000000000000000000066611515451615200376130ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; import java.sql.Timestamp; import java.time.Instant; import java.util.Calendar; import java.util.TimeZone; import org.junit.jupiter.api.Test; public class DateTimeUtilsTest { private final TimeZone defaultTimezone = TimeZone.getTimeZone("UTC"); private final TimeZone alternateTimezone = TimeZone.getTimeZone("America/Vancouver"); private final long positiveEpochMilli = 959817600000L; // 2000-06-01 00:00:00 UTC private final long negativeEpochMilli = -618105600000L; // 1950-06-01 00:00:00 UTC @Test public void testShouldGetOffsetWithSameTimeZone() { final TimeZone currentTimezone = TimeZone.getDefault(); final long epochMillis = positiveEpochMilli; final long offset = defaultTimezone.getOffset(epochMillis); TimeZone.setDefault(defaultTimezone); try { // Trying to guarantee timezone returns to its original value final long expected = epochMillis + offset; final long actual = DateTimeUtils.applyCalendarOffset(epochMillis, Calendar.getInstance(defaultTimezone)); assertThat(actual, is(expected)); } finally { // Reset Timezone TimeZone.setDefault(currentTimezone); } } @Test public void testShouldGetOffsetWithDifferentTimeZone() { final TimeZone currentTimezone = TimeZone.getDefault(); final long epochMillis = negativeEpochMilli; final long offset = alternateTimezone.getOffset(epochMillis); TimeZone.setDefault(alternateTimezone); try { // Trying to guarantee timezone returns to its original value final long expectedEpochMillis = epochMillis + offset; final long actualEpochMillis = DateTimeUtils.applyCalendarOffset(epochMillis, Calendar.getInstance(defaultTimezone)); assertThat(actualEpochMillis, is(expectedEpochMillis)); } finally { // Reset Timezone TimeZone.setDefault(currentTimezone); } } @Test public void testShouldGetTimestampPositive() { long epochMilli = positiveEpochMilli; final Instant instant = Instant.ofEpochMilli(epochMilli); final Timestamp expected = Timestamp.from(instant); final Timestamp actual = DateTimeUtils.getTimestampValue(epochMilli); assertThat(expected, is(actual)); } @Test public void testShouldGetTimestampNegative() { final long epochMilli = negativeEpochMilli; final Instant instant = Instant.ofEpochMilli(epochMilli); final Timestamp expected = Timestamp.from(instant); final Timestamp actual = DateTimeUtils.getTimestampValue(epochMilli); assertThat(expected, is(actual)); } } FallbackFlightSqlProducer.java000066400000000000000000000115411515451615200412500ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import com.google.protobuf.Any; import com.google.protobuf.ByteString; import com.google.protobuf.Message; import java.util.Collections; import java.util.List; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.Result; import org.apache.arrow.flight.Ticket; import org.apache.arrow.flight.sql.BasicFlightSqlProducer; import org.apache.arrow.flight.sql.impl.FlightSql; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; public class FallbackFlightSqlProducer extends BasicFlightSqlProducer { private final VectorSchemaRoot data; public FallbackFlightSqlProducer(VectorSchemaRoot resultData) { this.data = resultData; } @Override protected List determineEndpoints( T request, FlightDescriptor flightDescriptor, Schema schema) { return Collections.emptyList(); } @Override public void createPreparedStatement( FlightSql.ActionCreatePreparedStatementRequest request, CallContext context, StreamListener listener) { final FlightSql.ActionCreatePreparedStatementResult.Builder resultBuilder = FlightSql.ActionCreatePreparedStatementResult.newBuilder() .setPreparedStatementHandle(request.getQueryBytes()); final ByteString datasetSchemaBytes = ByteString.copyFrom(data.getSchema().serializeAsMessage()); resultBuilder.setDatasetSchema(datasetSchemaBytes); listener.onNext(new Result(Any.pack(resultBuilder.build()).toByteArray())); listener.onCompleted(); } @Override public FlightInfo getFlightInfoStatement( FlightSql.CommandStatementQuery command, CallContext context, FlightDescriptor descriptor) { return getFlightInfo(descriptor, command.getQuery()); } @Override public FlightInfo getFlightInfoPreparedStatement( FlightSql.CommandPreparedStatementQuery command, CallContext context, FlightDescriptor descriptor) { return getFlightInfo(descriptor, command.getPreparedStatementHandle().toStringUtf8()); } @Override public void getStreamStatement( FlightSql.TicketStatementQuery ticket, CallContext context, ServerStreamListener listener) { listener.start(data); listener.putNext(); listener.completed(); } @Override public void closePreparedStatement( FlightSql.ActionClosePreparedStatementRequest request, CallContext context, StreamListener listener) { listener.onCompleted(); } private FlightInfo getFlightInfo(FlightDescriptor descriptor, String query) { final List endpoints; final Ticket ticket = new Ticket(Any.pack(FlightSql.TicketStatementQuery.getDefaultInstance()).toByteArray()); if (query.equals("fallback")) { endpoints = Collections.singletonList( FlightEndpoint.builder(ticket, Location.reuseConnection()).build()); } else if (query.equals("fallback with error")) { endpoints = Collections.singletonList( FlightEndpoint.builder( ticket, 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(); } return FlightInfo.builder(data.getSchema(), descriptor, endpoints).build(); } } FlightEndpointDataQueueTest.java000066400000000000000000000054641515451615200416130ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; 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.assertDoesNotThrow; import static org.mockito.Mockito.mock; import java.util.concurrent.CompletionService; import org.apache.arrow.driver.jdbc.client.CloseableEndpointStreamPair; import org.junit.jupiter.api.BeforeEach; 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 FlightEndpointDataQueue}. */ @ExtendWith(MockitoExtension.class) public class FlightEndpointDataQueueTest { @Mock private CompletionService mockedService; private FlightEndpointDataQueue queue; @BeforeEach public void setUp() { queue = new FlightEndpointDataQueue(mockedService); } @Test public void testNextShouldRetrieveNullIfEmpty() throws Exception { assertThat(queue.next(), is(nullValue())); } @Test public void testNextShouldThrowExceptionUponClose() throws Exception { queue.close(); ThrowableAssertionUtils.simpleAssertThrowableClass( IllegalStateException.class, () -> queue.next()); } @Test public void testEnqueueShouldThrowExceptionUponClose() throws Exception { queue.close(); ThrowableAssertionUtils.simpleAssertThrowableClass( IllegalStateException.class, () -> queue.enqueue(mock(CloseableEndpointStreamPair.class))); } @Test public void testCheckOpen() throws Exception { assertDoesNotThrow( () -> { queue.checkOpen(); return true; }); queue.close(); ThrowableAssertionUtils.simpleAssertThrowableClass( IllegalStateException.class, () -> queue.checkOpen()); } @Test public void testShouldCloseQueue() throws Exception { assertThat(queue.isClosed(), is(false)); queue.close(); assertThat(queue.isClosed(), is(true)); } } FlightSqlTestCertificates.java000066400000000000000000000054551515451615200413210ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import java.io.File; import java.nio.file.Path; import java.nio.file.Paths; import java.util.Arrays; import java.util.List; import java.util.Objects; /** Utility class for unit tests that need to reference the certificate params. */ public class FlightSqlTestCertificates { public static final String TEST_DATA_ENV_VAR = "ARROW_TEST_DATA"; public static final String TEST_DATA_PROPERTY = "arrow.test.dataRoot"; static Path getTestDataRoot() { String path = System.getenv(TEST_DATA_ENV_VAR); if (path == null) { path = System.getProperty(TEST_DATA_PROPERTY); } return Paths.get( Objects.requireNonNull( path, String.format( "Could not find test data path. Set the environment variable %s or the JVM property %s.", TEST_DATA_ENV_VAR, TEST_DATA_PROPERTY))); } /** * Get the Path from the Files to be used in the encrypted test of Flight. * * @return the Path from the Files with certificates and keys. */ static Path getFlightTestDataRoot() { return getTestDataRoot().resolve("flight"); } /** * Create File object with the CA certificate. * * @return A File containing the Root CA certificate. */ public static File exampleCACert() { final Path root = getFlightTestDataRoot(); return root.resolve("root-ca.pem").toFile(); } /** * Create CertKeyPair object with the certificates and keys. * * @return A list with CertKeyPair. */ public static List exampleTlsCerts() { final Path root = getFlightTestDataRoot(); return Arrays.asList( new CertKeyPair(root.resolve("cert0.pem").toFile(), root.resolve("cert0.pkcs1").toFile()), new CertKeyPair(root.resolve("cert1.pem").toFile(), root.resolve("cert1.pkcs1").toFile())); } public static class CertKeyPair { public final File cert; public final File key; public CertKeyPair(File cert, File key) { this.cert = cert; this.key = key; } } } MockFlightSqlProducer.java000066400000000000000000000662161515451615200404530ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import static com.google.protobuf.Any.pack; import static com.google.protobuf.ByteString.copyFrom; import static java.lang.String.format; import static java.util.UUID.randomUUID; import static java.util.stream.Collectors.toList; import com.google.protobuf.Any; import com.google.protobuf.ByteString; import com.google.protobuf.Message; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.channels.Channels; import java.nio.charset.StandardCharsets; import java.util.AbstractMap.SimpleImmutableEntry; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Map.Entry; import java.util.Objects; import java.util.UUID; import java.util.function.BiConsumer; import java.util.function.Consumer; import java.util.stream.IntStream; import org.apache.arrow.flight.Action; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.Criteria; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightStream; 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; import org.apache.arrow.flight.sql.impl.FlightSql; import org.apache.arrow.flight.sql.impl.FlightSql.ActionClosePreparedStatementRequest; import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementRequest; import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementResult; 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; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetPrimaryKeys; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetSqlInfo; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTableTypes; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTables; import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementQuery; import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementUpdate; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementQuery; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementUpdate; import org.apache.arrow.flight.sql.impl.FlightSql.DoPutUpdateResult; import org.apache.arrow.flight.sql.impl.FlightSql.TicketStatementQuery; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.WriteChannel; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.JsonStringArrayList; import org.apache.calcite.avatica.Meta.StatementType; /** An ad-hoc {@link FlightSqlProducer} for tests. */ public final class MockFlightSqlProducer implements FlightSqlProducer { private final Map>> queryResults = new HashMap<>(); private final Map> selectResultProviders = new HashMap<>(); private final Map preparedStatements = new HashMap<>(); private final Map> catalogQueriesResults = new HashMap<>(); private final Map>> updateResultProviders = new HashMap<>(); private final SqlInfoBuilder sqlInfoBuilder = new SqlInfoBuilder(); private final Map parameterSchemas = new HashMap<>(); private final Map>> expectedParameterValues = new HashMap<>(); private final Map actionTypeCounter = new HashMap<>(); private static FlightInfo getFlightInfoExportedAndImportedKeys( final Message message, final FlightDescriptor descriptor) { return getFlightInfo(message, Schemas.GET_IMPORTED_KEYS_SCHEMA, descriptor); } private static FlightInfo getFlightInfo( final Message message, final Schema schema, final FlightDescriptor descriptor) { return new FlightInfo( schema, descriptor, Collections.singletonList(new FlightEndpoint(new Ticket(Any.pack(message).toByteArray()))), -1, -1); } public static ByteBuffer serializeSchema(final Schema schema) { final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); try { MessageSerializer.serialize(new WriteChannel(Channels.newChannel(outputStream)), schema); return ByteBuffer.wrap(outputStream.toByteArray()); } catch (final IOException e) { throw new RuntimeException("Failed to serialize schema", e); } } /** * Registers a new {@link StatementType#SELECT} SQL query. * * @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. */ public void addSelectQuery( final String sqlCommand, final Schema schema, final List> resultProviders) { final int providers = resultProviders.size(); final List uuids = IntStream.range(0, providers) .mapToObj(index -> new UUID(sqlCommand.hashCode(), Integer.hashCode(index))) .collect(toList()); queryResults.put(sqlCommand, new SimpleImmutableEntry<>(schema, uuids)); IntStream.range(0, providers) .forEach( index -> this.selectResultProviders.put(uuids.get(index), resultProviders.get(index))); } /** * Registers a new {@link StatementType#UPDATE} SQL query. * * @param sqlCommand the SQL command. * @param updatedRows the number of rows affected. */ public void addUpdateQuery(final String sqlCommand, final long updatedRows) { addUpdateQuery( sqlCommand, (flightStream, putResultStreamListener) -> { final DoPutUpdateResult result = DoPutUpdateResult.newBuilder().setRecordCount(updatedRows).build(); try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); final ArrowBuf buffer = allocator.buffer(result.getSerializedSize())) { buffer.writeBytes(result.toByteArray()); putResultStreamListener.onNext(PutResult.metadata(buffer)); } catch (final Throwable throwable) { putResultStreamListener.onError(throwable); } finally { putResultStreamListener.onCompleted(); } }); } /** * Adds a catalog query to the results. * * @param message the {@link Message} corresponding to the catalog query request type to register. * @param resultsProvider the results provider. */ public void addCatalogQuery( final Message message, final Consumer resultsProvider) { catalogQueriesResults.put(message, resultsProvider); } /** * Registers a new {@link StatementType#UPDATE} SQL query. * * @param sqlCommand the SQL command. * @param resultsProvider consumer for producing update results. */ void addUpdateQuery( final String sqlCommand, final BiConsumer> resultsProvider) { Preconditions.checkState( updateResultProviders.putIfAbsent(sqlCommand, resultsProvider) == null, format("Attempted to overwrite preexisting query: <%s>.", sqlCommand)); } /** Registers parameters expected to be provided with a prepared statement. */ public void addExpectedParameters( String query, Schema parameterSchema, List> expectedValues) { parameterSchemas.put(query, parameterSchema); expectedParameterValues.put(query, expectedValues); } @Override public void createPreparedStatement( final ActionCreatePreparedStatementRequest request, final CallContext callContext, final StreamListener listener) { try { final ByteString preparedStatementHandle = copyFrom(randomUUID().toString().getBytes(StandardCharsets.UTF_8)); final String query = request.getQuery(); final ActionCreatePreparedStatementResult.Builder resultBuilder = ActionCreatePreparedStatementResult.newBuilder() .setPreparedStatementHandle(preparedStatementHandle); final Entry> entry = queryResults.get(query); if (entry != null) { preparedStatements.put(preparedStatementHandle, query); final Schema datasetSchema = entry.getKey(); final ByteString datasetSchemaBytes = ByteString.copyFrom(serializeSchema(datasetSchema)); resultBuilder.setDatasetSchema(datasetSchemaBytes); } else if (updateResultProviders.containsKey(query)) { preparedStatements.put(preparedStatementHandle, query); } else { listener.onError( CallStatus.INVALID_ARGUMENT.withDescription("Query not found").toRuntimeException()); return; } final Schema parameterSchema = parameterSchemas.get(query); if (parameterSchema != null) { final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); MessageSerializer.serialize( new WriteChannel(Channels.newChannel(outputStream)), parameterSchema); resultBuilder.setParameterSchema(ByteString.copyFrom(outputStream.toByteArray())); } listener.onNext(new Result(pack(resultBuilder.build()).toByteArray())); } catch (final Throwable t) { listener.onError(t); } finally { listener.onCompleted(); } } @Override public void closePreparedStatement( final ActionClosePreparedStatementRequest actionClosePreparedStatementRequest, final CallContext callContext, final StreamListener streamListener) { // TODO Implement this method. streamListener.onCompleted(); } @Override public FlightInfo getFlightInfoStatement( final CommandStatementQuery commandStatementQuery, final CallContext callContext, final FlightDescriptor flightDescriptor) { final String query = commandStatementQuery.getQuery(); final Entry> queryInfo = Preconditions.checkNotNull( queryResults.get(query), format("Query not registered: <%s>.", query)); final List endpoints = queryInfo.getValue().stream() .map(TicketConversionUtils::getTicketBytesFromUuid) .map(TicketConversionUtils::getTicketStatementQueryFromHandle) .map(TicketConversionUtils::getEndpointFromMessage) .collect(toList()); return FlightInfo.builder(queryInfo.getKey(), flightDescriptor, endpoints) .setAppMetadata("foo".getBytes(StandardCharsets.UTF_8)) .build(); } @Override public FlightInfo getFlightInfoPreparedStatement( final CommandPreparedStatementQuery commandPreparedStatementQuery, final CallContext callContext, final FlightDescriptor flightDescriptor) { final ByteString preparedStatementHandle = commandPreparedStatementQuery.getPreparedStatementHandle(); final String query = Preconditions.checkNotNull( preparedStatements.get(preparedStatementHandle), format("No query registered under handle: <%s>.", preparedStatementHandle)); final Entry> queryInfo = Preconditions.checkNotNull( queryResults.get(query), format("Query not registered: <%s>.", query)); final List endpoints = queryInfo.getValue().stream() .map(TicketConversionUtils::getTicketBytesFromUuid) .map(TicketConversionUtils::getCommandPreparedStatementQueryFromHandle) .map(TicketConversionUtils::getEndpointFromMessage) .collect(toList()); return FlightInfo.builder(queryInfo.getKey(), flightDescriptor, endpoints) .setAppMetadata("foo".getBytes(StandardCharsets.UTF_8)) .build(); } @Override public SchemaResult getSchemaStatement( final CommandStatementQuery commandStatementQuery, final CallContext callContext, final FlightDescriptor flightDescriptor) { final String query = commandStatementQuery.getQuery(); final Entry> queryInfo = Preconditions.checkNotNull( queryResults.get(query), format("Query not registered: <%s>.", query)); return new SchemaResult(queryInfo.getKey()); } @Override public void getStreamStatement( final TicketStatementQuery ticketStatementQuery, final CallContext callContext, final ServerStreamListener serverStreamListener) { final UUID uuid = UUID.fromString(ticketStatementQuery.getStatementHandle().toStringUtf8()); Preconditions.checkNotNull( selectResultProviders.get(uuid), "No consumer was registered for the specified UUID: <%s>.", uuid) .accept(serverStreamListener); } @Override public void getStreamPreparedStatement( final CommandPreparedStatementQuery commandPreparedStatementQuery, final CallContext callContext, final ServerStreamListener serverStreamListener) { final UUID uuid = UUID.fromString(commandPreparedStatementQuery.getPreparedStatementHandle().toStringUtf8()); Preconditions.checkNotNull( selectResultProviders.get(uuid), "No consumer was registered for the specified UUID: <%s>.", uuid) .accept(serverStreamListener); } @Override public Runnable acceptPutStatement( final CommandStatementUpdate commandStatementUpdate, final CallContext callContext, final FlightStream flightStream, final StreamListener streamListener) { return () -> { final String query = commandStatementUpdate.getQuery(); final BiConsumer> resultProvider = Preconditions.checkNotNull( updateResultProviders.get(query), format("No consumer found for query: <%s>.", query)); resultProvider.accept(flightStream, streamListener); }; } private boolean validateParameters( String query, FlightStream flightStream, StreamListener streamListener) { final List> expectedValues = expectedParameterValues.get(query); if (expectedValues != null) { int index = 0; while (flightStream.next()) { final VectorSchemaRoot root = flightStream.getRoot(); for (int i = 0; i < root.getRowCount(); i++) { if (index >= expectedValues.size()) { streamListener.onError( CallStatus.INVALID_ARGUMENT .withDescription("More parameter rows provided than expected") .toRuntimeException()); return true; } List expectedRow = expectedValues.get(index++); if (root.getFieldVectors().size() != expectedRow.size()) { streamListener.onError( CallStatus.INVALID_ARGUMENT .withDescription("Parameter count mismatch") .toRuntimeException()); return true; } for (int paramIndex = 0; paramIndex < expectedRow.size(); paramIndex++) { Object expected = expectedRow.get(paramIndex); Object actual = root.getVector(paramIndex).getObject(i); boolean matches; if (expected.getClass().isArray()) { matches = Arrays.equals((Object[]) expected, ((JsonStringArrayList) actual).toArray()); } else { matches = Objects.equals(expected, actual); } if (!matches) { streamListener.onError( CallStatus.INVALID_ARGUMENT .withDescription( String.format( "Parameter mismatch. Expected: %s Actual: %s", expected, actual)) .toRuntimeException()); return true; } } } } if (index < expectedValues.size()) { streamListener.onError( CallStatus.INVALID_ARGUMENT .withDescription("Fewer parameter rows provided than expected") .toRuntimeException()); return true; } } return false; } @Override public Runnable acceptPutPreparedStatementUpdate( final CommandPreparedStatementUpdate commandPreparedStatementUpdate, final CallContext callContext, final FlightStream flightStream, final StreamListener streamListener) { final ByteString handle = commandPreparedStatementUpdate.getPreparedStatementHandle(); final String query = Preconditions.checkNotNull( preparedStatements.get(handle), format("No query registered under handle: <%s>.", handle)); if (validateParameters(query, flightStream, streamListener)) { return () -> {}; } return acceptPutStatement( CommandStatementUpdate.newBuilder().setQuery(query).build(), callContext, flightStream, streamListener); } @Override public Runnable acceptPutPreparedStatementQuery( final CommandPreparedStatementQuery commandPreparedStatementQuery, final CallContext callContext, final FlightStream flightStream, final StreamListener streamListener) { final ByteString handle = commandPreparedStatementQuery.getPreparedStatementHandle(); final String query = Preconditions.checkNotNull( preparedStatements.get(handle), format("No query registered under handle: <%s>.", handle)); if (validateParameters(query, flightStream, streamListener)) { return () -> {}; } return streamListener::onCompleted; } @Override public FlightInfo getFlightInfoSqlInfo( final CommandGetSqlInfo commandGetSqlInfo, final CallContext callContext, final FlightDescriptor flightDescriptor) { return getFlightInfo(commandGetSqlInfo, Schemas.GET_SQL_INFO_SCHEMA, flightDescriptor); } @Override public void getStreamSqlInfo( final CommandGetSqlInfo commandGetSqlInfo, final CallContext callContext, final ServerStreamListener serverStreamListener) { sqlInfoBuilder.send(commandGetSqlInfo.getInfoList(), serverStreamListener); } @Override public FlightInfo getFlightInfoTypeInfo( FlightSql.CommandGetXdbcTypeInfo request, CallContext context, FlightDescriptor descriptor) { // TODO Implement this return null; } @Override public void getStreamTypeInfo( FlightSql.CommandGetXdbcTypeInfo request, CallContext context, ServerStreamListener listener) { // TODO Implement this } @Override public FlightInfo getFlightInfoCatalogs( final CommandGetCatalogs commandGetCatalogs, final CallContext callContext, final FlightDescriptor flightDescriptor) { return getFlightInfo(commandGetCatalogs, Schemas.GET_CATALOGS_SCHEMA, flightDescriptor); } @Override public void getStreamCatalogs( final CallContext callContext, final ServerStreamListener serverStreamListener) { final CommandGetCatalogs command = CommandGetCatalogs.getDefaultInstance(); getStreamCatalogFunctions(command, serverStreamListener); } @Override public FlightInfo getFlightInfoSchemas( final CommandGetDbSchemas commandGetSchemas, final CallContext callContext, final FlightDescriptor flightDescriptor) { return getFlightInfo(commandGetSchemas, Schemas.GET_SCHEMAS_SCHEMA, flightDescriptor); } @Override public void getStreamSchemas( final CommandGetDbSchemas commandGetSchemas, final CallContext callContext, final ServerStreamListener serverStreamListener) { getStreamCatalogFunctions(commandGetSchemas, serverStreamListener); } @Override public FlightInfo getFlightInfoTables( final CommandGetTables commandGetTables, final CallContext callContext, final FlightDescriptor flightDescriptor) { return getFlightInfo(commandGetTables, Schemas.GET_TABLES_SCHEMA_NO_SCHEMA, flightDescriptor); } @Override public void getStreamTables( final CommandGetTables commandGetTables, final CallContext callContext, final ServerStreamListener serverStreamListener) { getStreamCatalogFunctions(commandGetTables, serverStreamListener); } @Override public FlightInfo getFlightInfoTableTypes( final CommandGetTableTypes commandGetTableTypes, final CallContext callContext, final FlightDescriptor flightDescriptor) { return getFlightInfo(commandGetTableTypes, Schemas.GET_TABLE_TYPES_SCHEMA, flightDescriptor); } @Override public void getStreamTableTypes( final CallContext callContext, final ServerStreamListener serverStreamListener) { final CommandGetTableTypes command = CommandGetTableTypes.getDefaultInstance(); getStreamCatalogFunctions(command, serverStreamListener); } @Override public FlightInfo getFlightInfoPrimaryKeys( final CommandGetPrimaryKeys commandGetPrimaryKeys, final CallContext callContext, final FlightDescriptor flightDescriptor) { return getFlightInfo(commandGetPrimaryKeys, Schemas.GET_PRIMARY_KEYS_SCHEMA, flightDescriptor); } @Override public void getStreamPrimaryKeys( final CommandGetPrimaryKeys commandGetPrimaryKeys, final CallContext callContext, final ServerStreamListener serverStreamListener) { getStreamCatalogFunctions(commandGetPrimaryKeys, serverStreamListener); } @Override public FlightInfo getFlightInfoExportedKeys( final CommandGetExportedKeys commandGetExportedKeys, final CallContext callContext, final FlightDescriptor flightDescriptor) { return getFlightInfoExportedAndImportedKeys(commandGetExportedKeys, flightDescriptor); } @Override public FlightInfo getFlightInfoImportedKeys( final CommandGetImportedKeys commandGetImportedKeys, final CallContext callContext, final FlightDescriptor flightDescriptor) { return getFlightInfoExportedAndImportedKeys(commandGetImportedKeys, flightDescriptor); } @Override public FlightInfo getFlightInfoCrossReference( final CommandGetCrossReference commandGetCrossReference, final CallContext callContext, final FlightDescriptor flightDescriptor) { return getFlightInfoExportedAndImportedKeys(commandGetCrossReference, flightDescriptor); } @Override public void getStreamExportedKeys( final CommandGetExportedKeys commandGetExportedKeys, final CallContext callContext, final ServerStreamListener serverStreamListener) { getStreamCatalogFunctions(commandGetExportedKeys, serverStreamListener); } @Override public void getStreamImportedKeys( final CommandGetImportedKeys commandGetImportedKeys, final CallContext callContext, final ServerStreamListener serverStreamListener) { getStreamCatalogFunctions(commandGetImportedKeys, serverStreamListener); } @Override public void getStreamCrossReference( final CommandGetCrossReference commandGetCrossReference, final CallContext callContext, final ServerStreamListener serverStreamListener) { getStreamCatalogFunctions(commandGetCrossReference, serverStreamListener); } @Override public void close() { // TODO No-op. } @Override public void listFlights( final CallContext callContext, final Criteria criteria, final StreamListener streamListener) { // TODO Implement this method. throw CallStatus.UNIMPLEMENTED.toRuntimeException(); } @Override public void doAction(CallContext context, Action action, StreamListener listener) { FlightSqlProducer.super.doAction(context, action, listener); actionTypeCounter.put( action.getType(), actionTypeCounter.getOrDefault(action.getType(), 0) + 1); } /** * Clear the `actionTypeCounter` map and restore to its default state. Intended to be used in * tests. */ public void clearActionTypeCounter() { actionTypeCounter.clear(); } public Map getActionTypeCounter() { return actionTypeCounter; } private void getStreamCatalogFunctions( final Message ticket, final ServerStreamListener serverStreamListener) { Preconditions.checkNotNull( catalogQueriesResults.get(ticket), format("Query not registered for ticket: <%s>", ticket)) .accept(serverStreamListener); } 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. } private static ByteString getTicketBytesFromUuid(final UUID uuid) { return ByteString.copyFromUtf8(uuid.toString()); } private static TicketStatementQuery getTicketStatementQueryFromHandle(final ByteString handle) { return TicketStatementQuery.newBuilder().setStatementHandle(handle).build(); } private static CommandPreparedStatementQuery getCommandPreparedStatementQueryFromHandle( final ByteString handle) { return CommandPreparedStatementQuery.newBuilder().setPreparedStatementHandle(handle).build(); } private static FlightEndpoint getEndpointFromMessage(final Message message) { return new FlightEndpoint(new Ticket(Any.pack(message).toByteArray())); } } } PartitionedFlightSqlProducer.java000066400000000000000000000102231515451615200420270ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import static com.google.protobuf.Any.pack; import com.google.protobuf.ByteString; import com.google.protobuf.Message; import java.util.Arrays; import java.util.List; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.NoOpFlightProducer; import org.apache.arrow.flight.Result; import org.apache.arrow.flight.Ticket; import org.apache.arrow.flight.sql.BasicFlightSqlProducer; import org.apache.arrow.flight.sql.impl.FlightSql; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; public class PartitionedFlightSqlProducer extends BasicFlightSqlProducer { /** A minimal FlightProducer intended to just serve data when given the correct Ticket. */ public static class DataOnlyFlightSqlProducer extends NoOpFlightProducer { private final Ticket ticket; private final VectorSchemaRoot data; public DataOnlyFlightSqlProducer(Ticket ticket, VectorSchemaRoot data) { this.ticket = ticket; this.data = data; } @Override public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { if (!Arrays.equals(ticket.getBytes(), this.ticket.getBytes())) { listener.error( CallStatus.INVALID_ARGUMENT.withDescription("Illegal ticket.").toRuntimeException()); return; } listener.start(data); listener.putNext(); listener.completed(); } } private final List endpoints; private final Schema schema; public PartitionedFlightSqlProducer(Schema schema, FlightEndpoint... endpoints) { this.schema = schema; this.endpoints = Arrays.asList(endpoints); } @Override protected List determineEndpoints( T request, FlightDescriptor flightDescriptor, Schema schema) { return endpoints; } @Override public void createPreparedStatement( FlightSql.ActionCreatePreparedStatementRequest request, CallContext context, StreamListener listener) { final FlightSql.ActionCreatePreparedStatementResult.Builder resultBuilder = FlightSql.ActionCreatePreparedStatementResult.newBuilder() .setPreparedStatementHandle(ByteString.EMPTY); final ByteString datasetSchemaBytes = ByteString.copyFrom(schema.serializeAsMessage()); resultBuilder.setDatasetSchema(datasetSchemaBytes); listener.onNext(new Result(pack(resultBuilder.build()).toByteArray())); listener.onCompleted(); } @Override public FlightInfo getFlightInfoStatement( FlightSql.CommandStatementQuery command, CallContext context, FlightDescriptor descriptor) { return FlightInfo.builder(schema, descriptor, endpoints).build(); } @Override public FlightInfo getFlightInfoPreparedStatement( FlightSql.CommandPreparedStatementQuery command, CallContext context, FlightDescriptor descriptor) { return FlightInfo.builder(schema, descriptor, endpoints).build(); } @Override public void closePreparedStatement( FlightSql.ActionClosePreparedStatementRequest request, CallContext context, StreamListener listener) { listener.onCompleted(); } // Note -- getStream() is intentionally not implemented. } ResultSetTestUtils.java000066400000000000000000000124461515451615200400470ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import static java.util.stream.IntStream.range; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; import java.sql.ResultSet; import java.sql.SQLException; import java.util.ArrayList; import java.util.List; import java.util.function.Function; /** * Utility class for testing that require asserting that the values in a {@link ResultSet} are * expected. */ public final class ResultSetTestUtils { /** * Checks that the values (rows and columns) in the provided {@link ResultSet} are expected. * * @param resultSet the {@code ResultSet} to assert. * @param expectedResults the rows and columns representing the only values the {@code resultSet} * is expected to have. * @param the type to be found in the expected results for the {@code resultSet}. * @throws SQLException if querying the {@code ResultSet} fails at some point unexpectedly. */ public static void testData(final ResultSet resultSet, final List> expectedResults) throws SQLException { testData( resultSet, range(1, resultSet.getMetaData().getColumnCount() + 1).toArray(), expectedResults); } /** * Checks that the values (rows and columns) in the provided {@link ResultSet} are expected. * * @param resultSet the {@code ResultSet} to assert. * @param columnNames the column names to fetch in the {@code ResultSet} for comparison. * @param expectedResults the rows and columns representing the only values the {@code resultSet} * is expected to have. * @param the type to be found in the expected results for the {@code resultSet}. * @throws SQLException if querying the {@code ResultSet} fails at some point unexpectedly. */ @SuppressWarnings("unchecked") public static void testData( final ResultSet resultSet, final List columnNames, final List> expectedResults) throws SQLException { testData( resultSet, data -> { final List columns = new ArrayList<>(); for (final String columnName : columnNames) { try { columns.add((T) resultSet.getObject(columnName)); } catch (final SQLException e) { throw new RuntimeException(e); } } return columns; }, expectedResults); } /** * Checks that the values (rows and columns) in the provided {@link ResultSet} are expected. * * @param resultSet the {@code ResultSet} to assert. * @param columnIndices the column indices to fetch in the {@code ResultSet} for comparison. * @param expectedResults the rows and columns representing the only values the {@code resultSet} * is expected to have. * @param the type to be found in the expected results for the {@code resultSet}. * @throws SQLException if querying the {@code ResultSet} fails at some point unexpectedly. */ @SuppressWarnings("unchecked") public static void testData( final ResultSet resultSet, final int[] columnIndices, final List> expectedResults) throws SQLException { testData( resultSet, data -> { final List columns = new ArrayList<>(); for (final int columnIndex : columnIndices) { try { columns.add((T) resultSet.getObject(columnIndex)); } catch (final SQLException e) { throw new RuntimeException(e); } } return columns; }, expectedResults); } /** * Checks that the values (rows and columns) in the provided {@link ResultSet} are expected. * * @param resultSet the {@code ResultSet} to assert. * @param dataConsumer the column indices to fetch in the {@code ResultSet} for comparison. * @param expectedResults the rows and columns representing the only values the {@code resultSet} * is expected to have. * @param the type to be found in the expected results for the {@code resultSet}. * @throws SQLException if querying the {@code ResultSet} fails at some point unexpectedly. */ public static void testData( final ResultSet resultSet, final Function> dataConsumer, final List> expectedResults) throws SQLException { final List> actualResults = new ArrayList<>(); while (resultSet.next()) { actualResults.add(dataConsumer.apply(resultSet)); } assertThat(actualResults, is(expectedResults)); } } RootAllocatorTestExtension.java000066400000000000000000000632441515451615200415570ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; 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; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; 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.FixedSizeBinaryVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.LargeVarBinaryVector; 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.UuidVector; import org.apache.arrow.vector.VarBinaryVector; 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.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; public class RootAllocatorTestExtension implements BeforeAllCallback, AfterAllCallback, AutoCloseable { public static final byte MAX_VALUE = Byte.MAX_VALUE; private final BufferAllocator rootAllocator = new RootAllocator(); private final Random random = new Random(10); @Override public void beforeAll(ExtensionContext context) { // no-op } @Override public void afterAll(ExtensionContext context) { try { close(); } catch (Exception e) { throw new RuntimeException(e); } } public BufferAllocator getRootAllocator() { return rootAllocator; } @Override public void close() throws Exception { this.rootAllocator.getChildAllocators().forEach(BufferAllocator::close); AutoCloseables.close(this.rootAllocator); } /** * Create a Float8Vector to be used in the accessor tests. * * @return Float8Vector */ public Float8Vector createFloat8Vector() { double[] doubleVectorValues = new double[] { 0, 1, -1, Byte.MIN_VALUE, Byte.MAX_VALUE, Short.MIN_VALUE, Short.MAX_VALUE, Integer.MIN_VALUE, Integer.MAX_VALUE, Long.MIN_VALUE, Long.MAX_VALUE, Float.MAX_VALUE, -Float.MAX_VALUE, Float.NEGATIVE_INFINITY, Float.POSITIVE_INFINITY, Float.MIN_VALUE, -Float.MIN_VALUE, Double.MAX_VALUE, -Double.MAX_VALUE, Double.NEGATIVE_INFINITY, Double.POSITIVE_INFINITY, Double.MIN_VALUE, -Double.MIN_VALUE, }; Float8Vector result = new Float8Vector("", this.getRootAllocator()); result.setValueCount(MAX_VALUE); for (int i = 0; i < MAX_VALUE; i++) { if (i < doubleVectorValues.length) { result.setSafe(i, doubleVectorValues[i]); } else { result.setSafe(i, random.nextDouble()); } } return result; } public Float8Vector createFloat8VectorForNullTests() { final Float8Vector float8Vector = new Float8Vector("ID", this.getRootAllocator()); float8Vector.allocateNew(1); float8Vector.setNull(0); float8Vector.setValueCount(1); return float8Vector; } /** * Create a Float4Vector to be used in the accessor tests. * * @return Float4Vector */ public Float4Vector createFloat4Vector() { float[] floatVectorValues = new float[] { 0, 1, -1, Byte.MIN_VALUE, Byte.MAX_VALUE, Short.MIN_VALUE, Short.MAX_VALUE, Integer.MIN_VALUE, Integer.MAX_VALUE, Long.MIN_VALUE, Long.MAX_VALUE, Float.MAX_VALUE, -Float.MAX_VALUE, Float.NEGATIVE_INFINITY, Float.POSITIVE_INFINITY, Float.MIN_VALUE, -Float.MIN_VALUE, }; Float4Vector result = new Float4Vector("", this.getRootAllocator()); result.setValueCount(MAX_VALUE); for (int i = 0; i < MAX_VALUE; i++) { if (i < floatVectorValues.length) { result.setSafe(i, floatVectorValues[i]); } else { result.setSafe(i, random.nextFloat()); } } return result; } /** * Create a IntVector to be used in the accessor tests. * * @return IntVector */ public IntVector createIntVector() { int[] intVectorValues = new int[] { 0, 1, -1, Byte.MIN_VALUE, Byte.MAX_VALUE, Short.MIN_VALUE, Short.MAX_VALUE, Integer.MIN_VALUE, Integer.MAX_VALUE, }; IntVector result = new IntVector("", this.getRootAllocator()); result.setValueCount(MAX_VALUE); for (int i = 0; i < MAX_VALUE; i++) { if (i < intVectorValues.length) { result.setSafe(i, intVectorValues[i]); } else { result.setSafe(i, random.nextInt()); } } return result; } /** * Create a SmallIntVector to be used in the accessor tests. * * @return SmallIntVector */ public SmallIntVector createSmallIntVector() { short[] smallIntVectorValues = new short[] { 0, 1, -1, Byte.MIN_VALUE, Byte.MAX_VALUE, Short.MIN_VALUE, Short.MAX_VALUE, }; SmallIntVector result = new SmallIntVector("", this.getRootAllocator()); result.setValueCount(MAX_VALUE); for (int i = 0; i < MAX_VALUE; i++) { if (i < smallIntVectorValues.length) { result.setSafe(i, smallIntVectorValues[i]); } else { result.setSafe(i, random.nextInt(Short.MAX_VALUE)); } } return result; } /** * Create a TinyIntVector to be used in the accessor tests. * * @return TinyIntVector */ public TinyIntVector createTinyIntVector() { byte[] byteVectorValues = new byte[] { 0, 1, -1, Byte.MIN_VALUE, Byte.MAX_VALUE, }; TinyIntVector result = new TinyIntVector("", this.getRootAllocator()); result.setValueCount(MAX_VALUE); for (int i = 0; i < MAX_VALUE; i++) { if (i < byteVectorValues.length) { result.setSafe(i, byteVectorValues[i]); } else { result.setSafe(i, random.nextInt(Byte.MAX_VALUE)); } } return result; } /** * Create a BigIntVector to be used in the accessor tests. * * @return BigIntVector */ public BigIntVector createBigIntVector() { long[] longVectorValues = new long[] { 0, 1, -1, Byte.MIN_VALUE, Byte.MAX_VALUE, Short.MIN_VALUE, Short.MAX_VALUE, Integer.MIN_VALUE, Integer.MAX_VALUE, Long.MIN_VALUE, Long.MAX_VALUE, }; BigIntVector result = new BigIntVector("", this.getRootAllocator()); result.setValueCount(MAX_VALUE); for (int i = 0; i < MAX_VALUE; i++) { if (i < longVectorValues.length) { result.setSafe(i, longVectorValues[i]); } else { result.setSafe(i, random.nextLong()); } } return result; } /** * Create a UInt1Vector to be used in the accessor tests. * * @return UInt1Vector */ public UInt1Vector createUInt1Vector() { short[] uInt1VectorValues = new short[] { 0, 1, -1, Byte.MIN_VALUE, Byte.MAX_VALUE, }; UInt1Vector result = new UInt1Vector("", this.getRootAllocator()); result.setValueCount(MAX_VALUE); for (int i = 0; i < MAX_VALUE; i++) { if (i < uInt1VectorValues.length) { result.setSafe(i, uInt1VectorValues[i]); } else { result.setSafe(i, random.nextInt(0x100)); } } return result; } /** * Create a UInt2Vector to be used in the accessor tests. * * @return UInt2Vector */ public UInt2Vector createUInt2Vector() { int[] uInt2VectorValues = new int[] { 0, 1, -1, Byte.MIN_VALUE, Byte.MAX_VALUE, Short.MIN_VALUE, Short.MAX_VALUE, }; UInt2Vector result = new UInt2Vector("", this.getRootAllocator()); result.setValueCount(MAX_VALUE); for (int i = 0; i < MAX_VALUE; i++) { if (i < uInt2VectorValues.length) { result.setSafe(i, uInt2VectorValues[i]); } else { result.setSafe(i, random.nextInt(0x10000)); } } return result; } /** * Create a UInt4Vector to be used in the accessor tests. * * @return UInt4Vector */ public UInt4Vector createUInt4Vector() { int[] uInt4VectorValues = new int[] { 0, 1, -1, Byte.MIN_VALUE, Byte.MAX_VALUE, Short.MIN_VALUE, Short.MAX_VALUE, Integer.MIN_VALUE, Integer.MAX_VALUE }; UInt4Vector result = new UInt4Vector("", this.getRootAllocator()); result.setValueCount(MAX_VALUE); for (int i = 0; i < MAX_VALUE; i++) { if (i < uInt4VectorValues.length) { result.setSafe(i, uInt4VectorValues[i]); } else { result.setSafe(i, random.nextInt(Integer.MAX_VALUE)); } } return result; } /** * Create a UInt8Vector to be used in the accessor tests. * * @return UInt8Vector */ public UInt8Vector createUInt8Vector() { long[] uInt8VectorValues = new long[] { 0, 1, -1, Byte.MIN_VALUE, Byte.MAX_VALUE, Short.MIN_VALUE, Short.MAX_VALUE, Integer.MIN_VALUE, Integer.MAX_VALUE, Long.MIN_VALUE, Long.MAX_VALUE }; UInt8Vector result = new UInt8Vector("", this.getRootAllocator()); result.setValueCount(MAX_VALUE); for (int i = 0; i < MAX_VALUE; i++) { if (i < uInt8VectorValues.length) { result.setSafe(i, uInt8VectorValues[i]); } else { result.setSafe(i, random.nextLong()); } } return result; } /** * Create a VarBinaryVector to be used in the accessor tests. * * @return VarBinaryVector */ public VarBinaryVector createVarBinaryVector() { return createVarBinaryVector(""); } /** * Create a VarBinaryVector to be used in the accessor tests. * * @return VarBinaryVector */ public VarBinaryVector createVarBinaryVector(final String fieldName) { VarBinaryVector valueVector = new VarBinaryVector(fieldName, this.getRootAllocator()); valueVector.allocateNew(3); valueVector.setSafe(0, (fieldName + "__BINARY_DATA_0001").getBytes(StandardCharsets.UTF_8)); valueVector.setSafe(1, (fieldName + "__BINARY_DATA_0002").getBytes(StandardCharsets.UTF_8)); valueVector.setSafe(2, (fieldName + "__BINARY_DATA_0003").getBytes(StandardCharsets.UTF_8)); valueVector.setValueCount(3); return valueVector; } /** * Create a LargeVarBinaryVector to be used in the accessor tests. * * @return LargeVarBinaryVector */ public LargeVarBinaryVector createLargeVarBinaryVector() { LargeVarBinaryVector valueVector = new LargeVarBinaryVector("", this.getRootAllocator()); valueVector.allocateNew(3); valueVector.setSafe(0, "BINARY_DATA_0001".getBytes(StandardCharsets.UTF_8)); valueVector.setSafe(1, "BINARY_DATA_0002".getBytes(StandardCharsets.UTF_8)); valueVector.setSafe(2, "BINARY_DATA_0003".getBytes(StandardCharsets.UTF_8)); valueVector.setValueCount(3); return valueVector; } /** * Create a FixedSizeBinaryVector to be used in the accessor tests. * * @return FixedSizeBinaryVector */ public FixedSizeBinaryVector createFixedSizeBinaryVector() { FixedSizeBinaryVector valueVector = new FixedSizeBinaryVector("", this.getRootAllocator(), 16); valueVector.allocateNew(3); valueVector.setSafe(0, "BINARY_DATA_0001".getBytes(StandardCharsets.UTF_8)); valueVector.setSafe(1, "BINARY_DATA_0002".getBytes(StandardCharsets.UTF_8)); valueVector.setSafe(2, "BINARY_DATA_0003".getBytes(StandardCharsets.UTF_8)); valueVector.setValueCount(3); return valueVector; } /** * Create a UInt8Vector to be used in the accessor tests. * * @return UInt8Vector */ public DecimalVector createDecimalVector() { BigDecimal[] bigDecimalValues = new BigDecimal[] { new BigDecimal(0), new BigDecimal(1), new BigDecimal(-1), new BigDecimal(Byte.MIN_VALUE), new BigDecimal(Byte.MAX_VALUE), new BigDecimal(-Short.MAX_VALUE), new BigDecimal(Short.MIN_VALUE), new BigDecimal(Integer.MIN_VALUE), new BigDecimal(Integer.MAX_VALUE), new BigDecimal(Long.MIN_VALUE), new BigDecimal(-Long.MAX_VALUE), new BigDecimal("170141183460469231731687303715884105727") }; DecimalVector result = new DecimalVector("ID", this.getRootAllocator(), 39, 0); result.setValueCount(MAX_VALUE); for (int i = 0; i < MAX_VALUE; i++) { if (i < bigDecimalValues.length) { result.setSafe(i, bigDecimalValues[i]); } else { result.setSafe(i, random.nextLong()); } } return result; } /** * Create a UInt8Vector to be used in the accessor tests. * * @return UInt8Vector */ public Decimal256Vector createDecimal256Vector() { BigDecimal[] bigDecimalValues = new BigDecimal[] { new BigDecimal(0), new BigDecimal(1), new BigDecimal(-1), new BigDecimal(Byte.MIN_VALUE), new BigDecimal(Byte.MAX_VALUE), new BigDecimal(-Short.MAX_VALUE), new BigDecimal(Short.MIN_VALUE), new BigDecimal(Integer.MIN_VALUE), new BigDecimal(Integer.MAX_VALUE), new BigDecimal(Long.MIN_VALUE), new BigDecimal(-Long.MAX_VALUE), new BigDecimal("170141183460469231731687303715884105727"), new BigDecimal("17014118346046923173168234157303715884105727"), new BigDecimal("1701411834604692317316823415265417303715884105727"), new BigDecimal("-17014118346046923173168234152654115451237303715884105727"), new BigDecimal("-17014118346046923173168234152654115451231545157303715884105727"), new BigDecimal("1701411834604692315815656534152654115451231545157303715884105727"), new BigDecimal("30560141183460469231581565634152654115451231545157303715884105727"), new BigDecimal( "57896044618658097711785492504343953926634992332820282019728792003956564819967"), new BigDecimal( "-56896044618658097711785492504343953926634992332820282019728792003956564819967") }; Decimal256Vector result = new Decimal256Vector("ID", this.getRootAllocator(), 77, 0); result.setValueCount(MAX_VALUE); for (int i = 0; i < MAX_VALUE; i++) { if (i < bigDecimalValues.length) { result.setSafe(i, bigDecimalValues[i]); } else { result.setSafe(i, random.nextLong()); } } return result; } public TimeStampNanoVector createTimeStampNanoVector() { TimeStampNanoVector valueVector = new TimeStampNanoVector("", this.getRootAllocator()); valueVector.allocateNew(2); valueVector.setSafe(0, TimeUnit.MILLISECONDS.toNanos(1625702400000L)); valueVector.setSafe(1, TimeUnit.MILLISECONDS.toNanos(1625788800000L)); valueVector.setValueCount(2); return valueVector; } public TimeStampNanoTZVector createTimeStampNanoTZVector(String timeZone) { TimeStampNanoTZVector valueVector = new TimeStampNanoTZVector("", this.getRootAllocator(), timeZone); valueVector.allocateNew(2); valueVector.setSafe(0, TimeUnit.MILLISECONDS.toNanos(1625702400000L)); valueVector.setSafe(1, TimeUnit.MILLISECONDS.toNanos(1625788800000L)); valueVector.setValueCount(2); return valueVector; } public TimeStampMicroVector createTimeStampMicroVector() { TimeStampMicroVector valueVector = new TimeStampMicroVector("", this.getRootAllocator()); valueVector.allocateNew(2); valueVector.setSafe(0, TimeUnit.MILLISECONDS.toMicros(1625702400000L)); valueVector.setSafe(1, TimeUnit.MILLISECONDS.toMicros(1625788800000L)); valueVector.setValueCount(2); return valueVector; } public TimeStampMicroTZVector createTimeStampMicroTZVector(String timeZone) { TimeStampMicroTZVector valueVector = new TimeStampMicroTZVector("", this.getRootAllocator(), timeZone); valueVector.allocateNew(2); valueVector.setSafe(0, TimeUnit.MILLISECONDS.toMicros(1625702400000L)); valueVector.setSafe(1, TimeUnit.MILLISECONDS.toMicros(1625788800000L)); valueVector.setValueCount(2); return valueVector; } public TimeStampMilliVector createTimeStampMilliVector() { TimeStampMilliVector valueVector = new TimeStampMilliVector("", this.getRootAllocator()); valueVector.allocateNew(2); valueVector.setSafe(0, 1625702400000L); valueVector.setSafe(1, 1625788800000L); valueVector.setValueCount(2); return valueVector; } public TimeStampMilliTZVector createTimeStampMilliTZVector(String timeZone) { TimeStampMilliTZVector valueVector = new TimeStampMilliTZVector("", this.getRootAllocator(), timeZone); valueVector.allocateNew(2); valueVector.setSafe(0, 1625702400000L); valueVector.setSafe(1, 1625788800000L); valueVector.setValueCount(2); return valueVector; } public TimeStampSecVector createTimeStampSecVector() { TimeStampSecVector valueVector = new TimeStampSecVector("", this.getRootAllocator()); valueVector.allocateNew(2); valueVector.setSafe(0, TimeUnit.MILLISECONDS.toSeconds(1625702400000L)); valueVector.setSafe(1, TimeUnit.MILLISECONDS.toSeconds(1625788800000L)); valueVector.setValueCount(2); return valueVector; } public TimeStampSecTZVector createTimeStampSecTZVector(String timeZone) { TimeStampSecTZVector valueVector = new TimeStampSecTZVector("", this.getRootAllocator(), timeZone); valueVector.allocateNew(2); valueVector.setSafe(0, TimeUnit.MILLISECONDS.toSeconds(1625702400000L)); valueVector.setSafe(1, TimeUnit.MILLISECONDS.toSeconds(1625788800000L)); valueVector.setValueCount(2); return valueVector; } public BitVector createBitVector() { BitVector valueVector = new BitVector("Value", this.getRootAllocator()); valueVector.allocateNew(2); valueVector.setSafe(0, 0); valueVector.setSafe(1, 1); valueVector.setValueCount(2); return valueVector; } public BitVector createBitVectorForNullTests() { final BitVector bitVector = new BitVector("ID", this.getRootAllocator()); bitVector.allocateNew(2); bitVector.setNull(0); bitVector.setValueCount(1); return bitVector; } public TimeNanoVector createTimeNanoVector() { TimeNanoVector valueVector = new TimeNanoVector("", this.getRootAllocator()); valueVector.allocateNew(5); valueVector.setSafe(0, 0); valueVector.setSafe(1, 1_000_000_000L); // 1 second valueVector.setSafe(2, 60 * 1_000_000_000L); // 1 minute valueVector.setSafe(3, 60 * 60 * 1_000_000_000L); // 1 hour valueVector.setSafe(4, (24 * 60 * 60 - 1) * 1_000_000_000L); // 23:59:59 valueVector.setValueCount(5); return valueVector; } public TimeMicroVector createTimeMicroVector() { TimeMicroVector valueVector = new TimeMicroVector("", this.getRootAllocator()); valueVector.allocateNew(5); valueVector.setSafe(0, 0); valueVector.setSafe(1, 1_000_000L); // 1 second valueVector.setSafe(2, 60 * 1_000_000L); // 1 minute valueVector.setSafe(3, 60 * 60 * 1_000_000L); // 1 hour valueVector.setSafe(4, (24 * 60 * 60 - 1) * 1_000_000L); // 23:59:59 valueVector.setValueCount(5); return valueVector; } public TimeMilliVector createTimeMilliVector() { TimeMilliVector valueVector = new TimeMilliVector("", this.getRootAllocator()); valueVector.allocateNew(5); valueVector.setSafe(0, 0); valueVector.setSafe(1, 1_000); // 1 second valueVector.setSafe(2, 60 * 1_000); // 1 minute valueVector.setSafe(3, 60 * 60 * 1_000); // 1 hour valueVector.setSafe(4, (24 * 60 * 60 - 1) * 1_000); // 23:59:59 valueVector.setValueCount(5); return valueVector; } public TimeSecVector createTimeSecVector() { TimeSecVector valueVector = new TimeSecVector("", this.getRootAllocator()); valueVector.allocateNew(5); valueVector.setSafe(0, 0); valueVector.setSafe(1, 1); // 1 second valueVector.setSafe(2, 60); // 1 minute valueVector.setSafe(3, 60 * 60); // 1 hour valueVector.setSafe(4, (24 * 60 * 60 - 1)); // 23:59:59 valueVector.setValueCount(5); return valueVector; } public DateDayVector createDateDayVector() { DateDayVector valueVector = new DateDayVector("", this.getRootAllocator()); valueVector.allocateNew(2); valueVector.setSafe(0, (int) TimeUnit.MILLISECONDS.toDays(1625702400000L)); valueVector.setSafe(1, (int) TimeUnit.MILLISECONDS.toDays(1625788800000L)); valueVector.setValueCount(2); return valueVector; } public DateMilliVector createDateMilliVector() { DateMilliVector valueVector = new DateMilliVector("", this.getRootAllocator()); valueVector.allocateNew(2); valueVector.setSafe(0, 1625702400000L); valueVector.setSafe(1, 1625788800000L); valueVector.setValueCount(2); return valueVector; } public ListVector createListVector() { return createListVector(""); } public ListVector createListVector(String fieldName) { ListVector valueVector = ListVector.empty(fieldName, this.getRootAllocator()); valueVector.setInitialCapacity(MAX_VALUE); UnionListWriter writer = valueVector.getWriter(); IntStream range = IntStream.range(0, MAX_VALUE); range.forEach( row -> { writer.startList(); writer.setPosition(row); IntStream.range(0, 5).map(j -> j * row).forEach(writer::writeInt); writer.setValueCount(5); writer.endList(); }); valueVector.setValueCount(MAX_VALUE); return valueVector; } public LargeListVector createLargeListVector() { LargeListVector valueVector = LargeListVector.empty("", this.getRootAllocator()); valueVector.setInitialCapacity(MAX_VALUE); UnionLargeListWriter writer = valueVector.getWriter(); IntStream range = IntStream.range(0, MAX_VALUE); range.forEach( row -> { writer.startList(); writer.setPosition(row); IntStream.range(0, 5).map(j -> j * row).forEach(writer::writeInt); writer.setValueCount(5); writer.endList(); }); valueVector.setValueCount(MAX_VALUE); return valueVector; } public FixedSizeListVector createFixedSizeListVector() { FixedSizeListVector valueVector = FixedSizeListVector.empty("", 5, this.getRootAllocator()); valueVector.setInitialCapacity(MAX_VALUE); UnionFixedSizeListWriter writer = valueVector.getWriter(); IntStream range = IntStream.range(0, MAX_VALUE); range.forEach( row -> { writer.startList(); writer.setPosition(row); IntStream.range(0, 5).map(j -> j * row).forEach(writer::writeInt); writer.setValueCount(5); writer.endList(); }); valueVector.setValueCount(MAX_VALUE); 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; } } SqlTypesTest.java000066400000000000000000000166511515451615200366620ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import static org.apache.arrow.driver.jdbc.utils.SqlTypes.getSqlTypeIdFromArrowType; import static org.apache.arrow.driver.jdbc.utils.SqlTypes.getSqlTypeNameFromArrowType; 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; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.UnionMode; import org.apache.arrow.vector.types.pojo.ArrowType; import org.junit.jupiter.api.Test; public class SqlTypesTest { @Test public void testGetSqlTypeIdFromArrowType() { assertEquals(Types.TINYINT, getSqlTypeIdFromArrowType(new ArrowType.Int(8, true))); assertEquals(Types.SMALLINT, getSqlTypeIdFromArrowType(new ArrowType.Int(16, true))); assertEquals(Types.INTEGER, getSqlTypeIdFromArrowType(new ArrowType.Int(32, true))); assertEquals(Types.BIGINT, getSqlTypeIdFromArrowType(new ArrowType.Int(64, true))); 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())); assertEquals(Types.DECIMAL, getSqlTypeIdFromArrowType(new ArrowType.Decimal(0, 0, 64))); assertEquals( Types.DOUBLE, getSqlTypeIdFromArrowType(new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE))); assertEquals( Types.FLOAT, getSqlTypeIdFromArrowType(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE))); assertEquals(Types.ARRAY, getSqlTypeIdFromArrowType(new ArrowType.List())); assertEquals(Types.ARRAY, getSqlTypeIdFromArrowType(new ArrowType.LargeList())); assertEquals(Types.ARRAY, getSqlTypeIdFromArrowType(new ArrowType.FixedSizeList(10))); assertEquals(Types.JAVA_OBJECT, getSqlTypeIdFromArrowType(new ArrowType.Struct())); assertEquals( Types.JAVA_OBJECT, getSqlTypeIdFromArrowType(new ArrowType.Duration(TimeUnit.MILLISECOND))); assertEquals( Types.JAVA_OBJECT, getSqlTypeIdFromArrowType(new ArrowType.Interval(IntervalUnit.DAY_TIME))); assertEquals( Types.JAVA_OBJECT, getSqlTypeIdFromArrowType(new ArrowType.Union(UnionMode.Dense, null))); assertEquals(Types.JAVA_OBJECT, getSqlTypeIdFromArrowType(new ArrowType.Map(true))); assertEquals(Types.NULL, getSqlTypeIdFromArrowType(new ArrowType.Null())); assertEquals(Types.OTHER, getSqlTypeIdFromArrowType(UuidType.INSTANCE)); } @Test public void testGetSqlTypeNameFromArrowType() { assertEquals("TINYINT", getSqlTypeNameFromArrowType(new ArrowType.Int(8, true))); assertEquals("SMALLINT", getSqlTypeNameFromArrowType(new ArrowType.Int(16, true))); assertEquals("INTEGER", getSqlTypeNameFromArrowType(new ArrowType.Int(32, true))); assertEquals("BIGINT", getSqlTypeNameFromArrowType(new ArrowType.Int(64, true))); 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())); assertEquals("DECIMAL", getSqlTypeNameFromArrowType(new ArrowType.Decimal(0, 0, 64))); assertEquals( "DOUBLE", getSqlTypeNameFromArrowType(new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE))); assertEquals( "FLOAT", getSqlTypeNameFromArrowType(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE))); assertEquals("ARRAY", getSqlTypeNameFromArrowType(new ArrowType.List())); assertEquals("ARRAY", getSqlTypeNameFromArrowType(new ArrowType.LargeList())); assertEquals("ARRAY", getSqlTypeNameFromArrowType(new ArrowType.FixedSizeList(10))); assertEquals("JAVA_OBJECT", getSqlTypeNameFromArrowType(new ArrowType.Struct())); assertEquals( "JAVA_OBJECT", getSqlTypeNameFromArrowType(new ArrowType.Duration(TimeUnit.MILLISECOND))); assertEquals( "JAVA_OBJECT", getSqlTypeNameFromArrowType(new ArrowType.Interval(IntervalUnit.DAY_TIME))); assertEquals( "JAVA_OBJECT", getSqlTypeNameFromArrowType(new ArrowType.Union(UnionMode.Dense, null))); assertEquals("JAVA_OBJECT", getSqlTypeNameFromArrowType(new ArrowType.Map(true))); assertEquals("NULL", getSqlTypeNameFromArrowType(new ArrowType.Null())); assertEquals("OTHER", getSqlTypeNameFromArrowType(UuidType.INSTANCE)); } } ThrowableAssertionUtils.java000066400000000000000000000043141515451615200410670ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; /** * Utility class to avoid upgrading JUnit to version >= 4.13 and keep using code to assert a {@link * Throwable}. This should be removed as soon as we can use the proper assertThrows/checkThrows. */ public class ThrowableAssertionUtils { private ThrowableAssertionUtils() {} public static void simpleAssertThrowableClass( final Class expectedThrowable, final ThrowingRunnable runnable) { try { runnable.run(); } catch (Throwable actualThrown) { if (expectedThrowable.isInstance(actualThrown)) { return; } else { final String mismatchMessage = String.format( "unexpected exception type thrown;\nexpected: %s\nactual: %s", formatClass(expectedThrowable), formatClass(actualThrown.getClass())); throw new AssertionError(mismatchMessage, actualThrown); } } final String notThrownMessage = String.format( "expected %s to be thrown, but nothing was thrown", formatClass(expectedThrowable)); throw new AssertionError(notThrownMessage); } private static String formatClass(final Class value) { // Fallback for anonymous inner classes final String className = value.getCanonicalName(); return className == null ? value.getName() : className; } public interface ThrowingRunnable { void run() throws Throwable; } } UrlParserTest.java000066400000000000000000000033131515451615200370040ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNull; import java.util.Map; import org.junit.jupiter.api.Test; class UrlParserTest { @Test void parse() { final Map parsed = UrlParser.parse("foo=bar&123=456", "&"); assertEquals(parsed.get("foo"), "bar"); assertEquals(parsed.get("123"), "456"); } @Test void parseEscaped() { final Map parsed = UrlParser.parse("foo=bar%26&%26123=456", "&"); assertEquals(parsed.get("foo"), "bar&"); assertEquals(parsed.get("&123"), "456"); } @Test void parseEmpty() { final Map parsed = UrlParser.parse("a=&b&foo=bar&123=456", "&"); assertEquals(parsed.get("a"), ""); assertNull(parsed.get("b")); assertEquals(parsed.get("foo"), "bar"); assertEquals(parsed.get("123"), "456"); } } VectorSchemaRootTransformerTest.java000066400000000000000000000132561515451615200425460ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.utils; 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.assertSame; import com.google.common.collect.ImmutableList; import java.util.List; import java.util.stream.Collectors; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; 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; import org.junit.jupiter.api.extension.RegisterExtension; public class VectorSchemaRootTransformerTest { @RegisterExtension public static RootAllocatorTestExtension rootAllocatorTestExtension = new RootAllocatorTestExtension(); private final BufferAllocator rootAllocator = rootAllocatorTestExtension.getRootAllocator(); @Test public void testTransformerBuilderWorksCorrectly() throws Exception { final VarBinaryVector field1 = rootAllocatorTestExtension.createVarBinaryVector("FIELD_1"); final VarBinaryVector field2 = rootAllocatorTestExtension.createVarBinaryVector("FIELD_2"); final VarBinaryVector field3 = rootAllocatorTestExtension.createVarBinaryVector("FIELD_3"); try (final VectorSchemaRoot originalRoot = VectorSchemaRoot.of(field1, field2, field3); final VectorSchemaRoot clonedRoot = cloneVectorSchemaRoot(originalRoot)) { final VectorSchemaRootTransformer.Builder builder = new VectorSchemaRootTransformer.Builder(originalRoot.getSchema(), rootAllocator); builder.renameFieldVector("FIELD_3", "FIELD_3_RENAMED"); builder.addEmptyField("EMPTY_FIELD", new ArrowType.Bool()); builder.renameFieldVector("FIELD_2", "FIELD_2_RENAMED"); builder.renameFieldVector("FIELD_1", "FIELD_1_RENAMED"); final VectorSchemaRootTransformer transformer = builder.build(); final Schema transformedSchema = new Schema( ImmutableList.of( Field.nullable("FIELD_3_RENAMED", new ArrowType.Binary()), Field.nullable("EMPTY_FIELD", new ArrowType.Bool()), Field.nullable("FIELD_2_RENAMED", new ArrowType.Binary()), Field.nullable("FIELD_1_RENAMED", new ArrowType.Binary()))); try (final VectorSchemaRoot transformedRoot = createVectorSchemaRoot(transformedSchema)) { assertSame(transformedRoot, transformer.transform(clonedRoot, transformedRoot)); assertEquals(transformedSchema, transformedRoot.getSchema()); final int rowCount = originalRoot.getRowCount(); assertEquals(rowCount, transformedRoot.getRowCount()); final VarBinaryVector originalField1 = (VarBinaryVector) originalRoot.getVector("FIELD_1"); final VarBinaryVector originalField2 = (VarBinaryVector) originalRoot.getVector("FIELD_2"); final VarBinaryVector originalField3 = (VarBinaryVector) originalRoot.getVector("FIELD_3"); final VarBinaryVector transformedField1 = (VarBinaryVector) transformedRoot.getVector("FIELD_1_RENAMED"); final VarBinaryVector transformedField2 = (VarBinaryVector) transformedRoot.getVector("FIELD_2_RENAMED"); final VarBinaryVector transformedField3 = (VarBinaryVector) transformedRoot.getVector("FIELD_3_RENAMED"); final FieldVector emptyField = transformedRoot.getVector("EMPTY_FIELD"); for (int i = 0; i < rowCount; i++) { assertArrayEquals(originalField1.getObject(i), transformedField1.getObject(i)); assertArrayEquals(originalField2.getObject(i), transformedField2.getObject(i)); assertArrayEquals(originalField3.getObject(i), transformedField3.getObject(i)); assertNull(emptyField.getObject(i)); } } } } private VectorSchemaRoot cloneVectorSchemaRoot(final VectorSchemaRoot originalRoot) { final VectorUnloader vectorUnloader = new VectorUnloader(originalRoot); try (final ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch()) { final VectorSchemaRoot clonedRoot = createVectorSchemaRoot(originalRoot.getSchema()); final VectorLoader vectorLoader = new VectorLoader(clonedRoot); vectorLoader.load(recordBatch); return clonedRoot; } } private VectorSchemaRoot createVectorSchemaRoot(final Schema schema) { final List fieldVectors = schema.getFields().stream() .map(field -> field.createVector(rootAllocator)) .collect(Collectors.toList()); return new VectorSchemaRoot(fieldVectors); } } apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/resources/000077500000000000000000000000001515451615200260335ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/resources/keys/000077500000000000000000000000001515451615200270065ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/resources/keys/keyStore.jks000066400000000000000000000030011515451615200313160ustar00rootroot00000000000000þíþí!cn=test,o=myorg\, inc.,st=ca,c=usyã*ŠX.509§0‚£0‚‹  âžòqzÝ 0  *†H†÷  0?1 0 UUS1 0 UCA10U  MyOrg, Inc.1 0 Utest0 190618181645Z 461103181645Z0?1 0 UUS1 0 UCA10U  MyOrg, Inc.1 0 Utest0‚"0  *†H†÷ ‚0‚ ‚£xÁóJ¯¶v¦:ØëÿôP2£ð3²Æ;õ³€a‘‘­6äX€: w\Fí‚ö Š{ngSq‡iÚ6H{ÛJ†>N*n#¤”8†óЛ۰}-n6ˆêÊ\é᤹‹/ýÆ Å¦SξÇ :oÆwŸ§<ê¨é[%ª’Œì†k&‹#Äê¨ñžò¹èÁ}(md¯¦K‚‡²0½gKßÿ/룡0ž0Ui‰BØé2xQ59%h1B7m “0oU#h0f€i‰BØé2xQ59%h1B7m “¡C¤A0?1 0 UUS1 0 UCA10U  MyOrg, Inc.1 0 Utest‚ âžòqzÝ 0 U0ÿ0  *†H†÷  ‚djÄÈÈ;Å¥p†¬—NqáuähÞÊI&=ßù.zÙû9øa}J8å”ÔiÖ[ð-ÝåGCyCbõ¬ªà‰´‘ Ý´L2_¹’„p{!ü¹f gÕ_*jtÏzLF†àY ²c‘Ï6Ä@N~ž«23«?¿½~¹;¸-aÄ+Q‚d·¦MOÔ¢ð>EGu5è3#XTïÉ|ƒÕiS”¤ÆËÔ˜ó‰î‹×•¸ïÎE ÿÕY.ª«ä¤!'81yxP yaNo=QÝÛ¼ªÒ.d;ˆÌŸnò¦†M1*1j8‘ / 'Oë‚_Ñô«Ä"¨^Ï  óB` ƒkùÍÖØ#4 i¶°cN}M©lh:‹ú1d%˜‡¸‹0’Õ‹½Ì Å–4ŠÌA #b£´{Ùí os>“¯[€~øfÙý4¸Ä>8ÿkî¹+"2XŽù`ª°Xz¿Hèw^<gbIŸ•ê‡IìÑ‚e¦€ —¾y(ò]•@³ ûw À5¥†ƒ@ínÖ¿nò<Цå9%®‚¸FIùÕ¬OÕ»º sà_Gsˆÿfû¿ÆC¤³`Ì0€ ÔƽcI³ Ú›ã}&Õ‚ÅÛFj¦íÖ'ævLÚ˜ŸH£6Ó>îñ:Õý¸y×ñŠLþX¬-VÎQKÒD‰;É á,t6Íd.‹Jgø8¹è0FÂ+£WÒèŸai¯wapache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/resources/keys/noCertificate.jks000066400000000000000000000047611515451615200323060ustar00rootroot000000000000000‚ í0‚ ¦ *†H†÷  ‚ —‚ “0‚ 0‚k *†H†÷  ‚\‚X0‚T0‚P *†H†÷   ‚û0‚÷0) *†H†÷  0‹W Cîù'šoŒofR7æÐÊòwÃP‚Èà1O-Í=]oîq…Ü×î2søŽVÔqDì=¯ µÝ<¾´l¯¥Çxgk“‘yY{{ñÉ‹å¾Oã‘¢t¸èµ˜prfUÉ›7™ˆNL“D®0\æú{'þÁ×AOŒ…mÒÀÔp;•ó_ €Ñšãêæ¼58qkgšp½Öº€KñäaéœÒâwî’gˆ±À’ Æ6î Æbz½SJ2lÑZu‚:r9—“ÌJß\Ožú§ü{Ìé2%>çŠZ\§?áÂÅÃ(RC~‰U+î®UXÍ*µú:BCNÐì7ª/÷áŸÑ¢ÌþM£ßÆg[‚xè,5`‚ãÀŽ`eÔt‡ÞÞ/ö3Ý>êó™Äá¹ë¼:cxÕÒáæ–ƒ~>}ÙC®j6˜¹¾ Z}û>ÿ¦AlÅÇ b*…ªýán:?\@-ydê‚ãôö–uîJnæl­ìæj-/]M…äUN{Äi%xDHʬê¦}ªºtÆ]†•j݉+z !ÒåŸ~êˆ^i§6÷nÔj uýb³ (ZxĨªÕ½çâóÄÔQ0,šbª€×g!7“>ÈWùXß ¦î’ëŒhܯN¼†¼çÞ׳y'ú… æ"ö(¢'µW-ÇuÀ³#[)FšÌPë@¨Æ)ºnááûÀ6Î2K£I*|b€ q ÷Qì\IÜN8Ar¨M1;öÃ¥:»óKžwǜ1a } ù*¿æ¢kuA•uó‡`ÛÙ3?)ò¢ØáSîÓÊ õ¶qû¹xã^k¢@› ›7'åÈPOª°°Íã¥Ujã‹Ý«ãè_)÷Ú¯yÓ7(fúOÂôLŠíÏ¿/”hÀîRIÊÎc®Ýå?—¬ÓÆ;Ú–Ôü¢ÆÅ;¥ô!ø´M߉›&kÁ;W’_¬–P¥[ù"E¥E½É…¦:o¼¿ñÔo`vè´1‹*º„ù8FäÞß/Y£Wz ’ÍPÞëfxJtº}}¶/U°…è7hYZ[íõeŸ¸5ÿk¤å—a•¶ì£p”»ü9[_½S:‘ùÉ•Ÿ.Ȭ¤±³çå@mæ¸Ì)T3-´V þ@ší+h û[`¹ón›ð=ÏËsÒ<ªVƒv›˜ÿm+OÏÄN(·"‚|MWæÐvÂ1 °_0^=ž „ìiŠ˜d÷cÊ­XÇE—†„ÕÅnm=³—7úò0µZîÞõP;¯”E£+‚9šÈ`÷\°]%Ïd­Í’hb?°»Ù…ì ‡4ñ £ì¥ @ÜŠ¦ßÉÂ0¿¿VüeL021B0 *†H†÷  1invalid0! *†H†÷  1Time 16236820030700‚ *†H†÷  ‚ 0‚ 0‚ *†H†÷ 0) *†H†÷  0Æ\‚ü ö}u®îµÖ48 ûø'²ÃP€‚Èc7ün-˜@°ßfŒÅ'÷Wâë2(eÒ«·û D°i}K@\®êÿ*?œÀ ¡ØÄPI:“MÝ5pÇ&,MË!öwlG=m»‹bÕS–ÙjK¨ûÆJã’\ç“åHõêàŒq³T«͎tcW'²Yt^ª]jŸµ\.-ªÒ°\Š¦³²Òkt[nà zmN:â¶&0C[·1; EE&§gøçÃ(èÁÝ'ÈŽ£b*AÁª© 'ƒ.º£×Fï ¨r¯¨ZÙzÕóêJwßHK6ß×ÛWÍ)ƒŽu3ã3>kƒÌT§£—l…”Ë'¥ø_\xj "©ržöׄþ¡3— ž^zŸ½wþûkG×hºýA½Þã;+wŽOs<É-yºé~à w#n±ÕÊÍ1⮌€È©7ºô¤­¨j¡?ª`ôi£\Êdè±9ÛYR·/¸h¢½(®}Æh©ûúôD«»,mèÐ"LoÓgo,ûÊ ÂA}WnÄ&E+ñçÁÄ:æ²}jó”„E[¯;MáSBì—'_…ßÌ] °çùƒÌÕoDŠoN6xMC]mœIðÞVî. ~¹º¾Ÿ¡$­Ÿ¨8ÅâlC5–0al³áªà3„=ø^<¬£¾ÜÕ°Œ½-ˆ™ÊŠX?ò-ÚîK–â)Rß³‘€}¸éìˆÆ ´dð¡nûÄG¹NGì®5sRQENÊÕäÍ| Æj@×aè-öšÊ®*ØÞÍø¹¢œÊîs¤~êÞ׿/äÐpŽeÀcóW–Њ?Å5„±È@e¤á.6E½†ß"¸ôÿ›PÉ˶EÇê0;ýR­-1f|ͱ ‘Í Jˆ¦eU¸‚ “Z?só©ÈËTä‡ñ%Ñ÷`üñZIÌ󥓢¸¦{äCŽw}b¡šÄó@8}þ“ cõ*­(“ØiE2¨öÛ©:=ÿNëÑ5ŵúA~8™Úî›’ñl¶j_Nƒ±"Zæ žLB»Xc-ÍЭ•œ,0>0!0 +c]$‚’®ð¹ ™0—™£ü÷°ç¢7_OŽ9<-IxhßœÛÜéɈ apache-arrow-java-19.0.0/flight/flight-sql-jdbc-core/src/test/resources/logback.xml000066400000000000000000000025361515451615200301650ustar00rootroot00000000000000 %d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n apache-arrow-java-19.0.0/flight/flight-sql-jdbc-driver/000077500000000000000000000000001515451615200226165ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-driver/pom.xml000066400000000000000000000201441515451615200241340ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-flight 19.0.0 flight-sql-jdbc-driver jar Arrow Flight SQL JDBC Driver A JDBC driver based on Arrow Flight SQL. https://arrow.apache.org org.apache.arrow flight-sql-jdbc-core runtime org.slf4j slf4j-jdk14 runtime com.google.guava guava runtime org.apache.maven.plugins maven-failsafe-plugin default-it integration-test verify org.apache.maven.plugins maven-shade-plugin shade package false true false *:* org.checkerframework:checker-qual org.codehaus.mojo:animal-sniffer-annotations javax.annotation:javax.annotation-api com.google.android:annotations com.google.errorprone:error_prone_annotations com.google.code.findbugs:jsr305 com.google.j2objc:j2objc-annotations com. org.apache.arrow.driver.jdbc.shaded.com. com.sun.** org. org.apache.arrow.driver.jdbc.shaded.org. org.apache.arrow.driver.jdbc.** org.apache.arrow.flight.name org.apache.arrow.flight.version org.apache.arrow.flight.jdbc-driver.name org.apache.arrow.flight.jdbc-driver.version io. org.apache.arrow.driver.jdbc.shaded.io. net. org.apache.arrow.driver.jdbc.shaded.net. mozilla. org.apache.arrow.driver.jdbc.shaded.mozilla. META-INF.native.libnetty_ META-INF.native.liborg_apache_arrow_driver_jdbc_shaded_netty_ META-INF.native.netty_ META-INF.native.org_apache_arrow_driver_jdbc_shaded_netty_ META-INF/LICENSE.txt src/shade/LICENSE.txt META-INF/NOTICE.txt src/shade/NOTICE.txt org.apache.arrow:arrow-vector codegen/** org.apache.calcite.avatica:* META-INF/services/java.sql.Driver *:* **/*.SF **/*.RSA **/*.DSA META-INF/native-image/ META-INF/proguard/ META-INF/versions/ **/*.proto **/module-info.class LICENSE.txt NOTICE.txt META-INF/*LICENSE* META-INF/*NOTICE* META-INF/license/* META-INF/licenses/**/* apache-arrow-java-19.0.0/flight/flight-sql-jdbc-driver/src/000077500000000000000000000000001515451615200234055ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-driver/src/shade/000077500000000000000000000000001515451615200244715ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-driver/src/shade/LICENSE.txt000066400000000000000000000452471515451615200263300ustar00rootroot00000000000000 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.** apache-arrow-java-19.0.0/flight/flight-sql-jdbc-driver/src/shade/NOTICE.txt000066400000000000000000000275051515451615200262240ustar00rootroot00000000000000Apache 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 apache-arrow-java-19.0.0/flight/flight-sql-jdbc-driver/src/test/000077500000000000000000000000001515451615200243645ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-driver/src/test/java/000077500000000000000000000000001515451615200253055ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-driver/src/test/java/org/000077500000000000000000000000001515451615200260745ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-driver/src/test/java/org/apache/000077500000000000000000000000001515451615200273155ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200304475ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/000077500000000000000000000000001515451615200317425ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/000077500000000000000000000000001515451615200326445ustar00rootroot00000000000000ITDriverJarValidation.java000066400000000000000000000145061515451615200375760ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertNotNull; import com.google.common.collect.ImmutableSet; import java.io.File; import java.io.IOException; import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Method; import java.net.JarURLConnection; import java.net.URISyntaxException; import java.net.URL; import java.net.URLClassLoader; import java.util.Set; import java.util.concurrent.TimeUnit; import java.util.function.Predicate; import java.util.jar.JarEntry; import java.util.jar.JarFile; import java.util.stream.Stream; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; import org.junit.jupiter.api.function.Executable; /** * Check the content of the JDBC driver jar * *

    After shading everything should be either under org.apache.arrow.driver.jdbc. package */ public class ITDriverJarValidation { /** * Use this property to provide path to the JDBC driver jar. Can be used to run the test from an * IDE */ public static final String JDBC_DRIVER_PATH_OVERRIDE = System.getProperty("arrow-flight-jdbc-driver.jar.override"); /** List of allowed prefixes a jar entry may match. */ public static final Set ALLOWED_PREFIXES = ImmutableSet.of( "org/apache/arrow/driver/jdbc/", // Driver code "META-INF/maven/", // Maven metadata (useful for security scanner "META-INF/services/", // ServiceLoader implementations "META-INF/license/", "META-INF/licenses/", // Prefixes for native libraries "META-INF/native/liborg_apache_arrow_driver_jdbc_shaded_", "META-INF/native/org_apache_arrow_driver_jdbc_shaded_"); /** 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", "META-INF/DEPENDENCIES", "META-INF/FastDoubleParser-LICENSE", "META-INF/FastDoubleParser-NOTICE", "META-INF/LICENSE", "META-INF/LICENSE.txt", "META-INF/NOTICE", "META-INF/NOTICE.txt", "META-INF/thirdparty-LICENSE", "META-INF/bigint-LICENSE"); // This method is designed to work with Maven failsafe plugin and expects the // JDBC driver jar to be present in the test classpath (instead of the individual classes) private static File getJdbcJarFile() throws IOException { // Check if an override has been set if (JDBC_DRIVER_PATH_OVERRIDE != null) { return new File(JDBC_DRIVER_PATH_OVERRIDE); } // Check classpath to find the driver jar (without loading the class) URL driverClassURL = ITDriverJarValidation.class .getClassLoader() .getResource("org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.class"); assertNotNull(driverClassURL, "Driver class was not detected in the classpath"); assertEquals( "jar", driverClassURL.getProtocol(), "Driver class was not found inside a jar file"); // Return the enclosing jar file JarURLConnection connection = (JarURLConnection) driverClassURL.openConnection(); try { return new File(connection.getJarFileURL().toURI()); } catch (URISyntaxException e) { throw new IOException(e); } } /** Validate the content of the jar to enforce all 3rd party dependencies have been shaded. */ @Test @Timeout(value = 2, unit = TimeUnit.MINUTES) public void validateShadedJar() throws IOException { try (JarFile jar = new JarFile(getJdbcJarFile())) { Stream executables = jar.stream() .filter(Predicate.not(JarEntry::isDirectory)) .map( entry -> { return () -> checkEntryAllowed(entry.getName()); }); Assertions.assertAll(executables); } } /** Check that relocated netty code can also load matching native library. */ @Test @Timeout(value = 2, unit = TimeUnit.MINUTES) public void checkNettyOpenSslNativeLoader() throws Throwable { try (URLClassLoader driverClassLoader = new URLClassLoader(new URL[] {getJdbcJarFile().toURI().toURL()}, null)) { Class openSslClass = driverClassLoader.loadClass( "org.apache.arrow.driver.jdbc.shaded.io.netty.handler.ssl.OpenSsl"); Method method = openSslClass.getDeclaredMethod("ensureAvailability"); try { method.invoke(null); } catch (InvocationTargetException e) { throw e.getCause(); } } } /** * Check if a jar entry is allowed. * *

    A jar entry is allowed if either it is part of the allowed files or it matches one of the * allowed prefixes * * @param name the jar entry name * @throws AssertionError if the entry is not allowed */ private void checkEntryAllowed(String name) { // Check if there's a matching file entry first if (ALLOWED_FILES.contains(name)) { return; } for (String prefix : ALLOWED_PREFIXES) { if (name.startsWith(prefix)) { return; } } throw new AssertionError("'" + name + "' is not an allowed jar entry"); } } apache-arrow-java-19.0.0/flight/flight-sql/000077500000000000000000000000001515451615200204255ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/pom.xml000066400000000000000000000107551515451615200217520ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-flight 19.0.0 flight-sql jar Arrow Flight SQL Contains utility classes to expose Flight SQL semantics for clients and servers over Arrow Flight 1 --add-reads=org.apache.arrow.flight.sql=org.slf4j --add-reads=org.apache.arrow.flight.core=ALL-UNNAMED --add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED org.apache.arrow flight-core io.netty netty-transport-native-unix-common io.netty netty-transport-native-kqueue io.netty netty-transport-native-epoll org.apache.arrow arrow-memory-core org.immutables value-annotations org.apache.arrow arrow-jdbc test com.google.guava guava com.google.protobuf protobuf-java org.apache.arrow arrow-vector ${arrow.vector.classifier} org.slf4j slf4j-api test org.apache.derby derby 10.15.2.0 test org.apache.commons commons-dbcp2 2.14.0 test commons-logging commons-logging org.apache.commons commons-pool2 2.13.1 test org.apache.commons commons-text 1.15.0 test commons-cli commons-cli 1.11.0 true org.assertj assertj-core test apache-arrow-java-19.0.0/flight/flight-sql/src/000077500000000000000000000000001515451615200212145ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/000077500000000000000000000000001515451615200221405ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/000077500000000000000000000000001515451615200230615ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/module-info.java000066400000000000000000000022701515451615200261430ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ module org.apache.arrow.flight.sql { exports org.apache.arrow.flight.sql; exports org.apache.arrow.flight.sql.example; exports org.apache.arrow.flight.sql.util; requires com.google.common; requires com.google.protobuf; requires java.sql; requires org.apache.arrow.flight.core; requires org.apache.arrow.memory.core; requires org.apache.arrow.vector; requires org.apache.commons.cli; } apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/000077500000000000000000000000001515451615200236505ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/000077500000000000000000000000001515451615200250715ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200262235ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/000077500000000000000000000000001515451615200275005ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/000077500000000000000000000000001515451615200302775ustar00rootroot00000000000000BasicFlightSqlProducer.java000066400000000000000000000106421515451615200354310ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql; import com.google.protobuf.Message; import java.util.List; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.sql.impl.FlightSql; import org.apache.arrow.vector.types.pojo.Schema; /** A {@link FlightSqlProducer} that implements getting FlightInfo for each metadata request. */ public abstract class BasicFlightSqlProducer extends NoOpFlightSqlProducer { @Override public FlightInfo getFlightInfoSqlInfo( FlightSql.CommandGetSqlInfo request, CallContext context, FlightDescriptor descriptor) { return generateFlightInfo(request, descriptor, Schemas.GET_SQL_INFO_SCHEMA); } @Override public FlightInfo getFlightInfoTypeInfo( FlightSql.CommandGetXdbcTypeInfo request, CallContext context, FlightDescriptor descriptor) { return generateFlightInfo(request, descriptor, Schemas.GET_TYPE_INFO_SCHEMA); } @Override public FlightInfo getFlightInfoCatalogs( FlightSql.CommandGetCatalogs request, CallContext context, FlightDescriptor descriptor) { return generateFlightInfo(request, descriptor, Schemas.GET_CATALOGS_SCHEMA); } @Override public FlightInfo getFlightInfoSchemas( FlightSql.CommandGetDbSchemas request, CallContext context, FlightDescriptor descriptor) { return generateFlightInfo(request, descriptor, Schemas.GET_SCHEMAS_SCHEMA); } @Override public FlightInfo getFlightInfoTables( FlightSql.CommandGetTables request, CallContext context, FlightDescriptor descriptor) { if (request.getIncludeSchema()) { return generateFlightInfo(request, descriptor, Schemas.GET_TABLES_SCHEMA); } return generateFlightInfo(request, descriptor, Schemas.GET_TABLES_SCHEMA_NO_SCHEMA); } @Override public FlightInfo getFlightInfoTableTypes( FlightSql.CommandGetTableTypes request, CallContext context, FlightDescriptor descriptor) { return generateFlightInfo(request, descriptor, Schemas.GET_TABLE_TYPES_SCHEMA); } @Override public FlightInfo getFlightInfoPrimaryKeys( FlightSql.CommandGetPrimaryKeys request, CallContext context, FlightDescriptor descriptor) { return generateFlightInfo(request, descriptor, Schemas.GET_PRIMARY_KEYS_SCHEMA); } @Override public FlightInfo getFlightInfoExportedKeys( FlightSql.CommandGetExportedKeys request, CallContext context, FlightDescriptor descriptor) { return generateFlightInfo(request, descriptor, Schemas.GET_EXPORTED_KEYS_SCHEMA); } @Override public FlightInfo getFlightInfoImportedKeys( FlightSql.CommandGetImportedKeys request, CallContext context, FlightDescriptor descriptor) { return generateFlightInfo(request, descriptor, Schemas.GET_IMPORTED_KEYS_SCHEMA); } @Override public FlightInfo getFlightInfoCrossReference( FlightSql.CommandGetCrossReference request, CallContext context, FlightDescriptor descriptor) { return generateFlightInfo(request, descriptor, Schemas.GET_CROSS_REFERENCE_SCHEMA); } /** * Return a list of FlightEndpoints for the given request and FlightDescriptor. This method should * validate that the request is supported by this FlightSqlProducer. */ protected abstract List determineEndpoints( T request, FlightDescriptor flightDescriptor, Schema schema); protected FlightInfo generateFlightInfo( T request, FlightDescriptor descriptor, Schema schema) { final List endpoints = determineEndpoints(request, descriptor, schema); return new FlightInfo(schema, descriptor, endpoints, -1, -1); } } CancelListener.java000066400000000000000000000033251515451615200337610ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql; import com.google.protobuf.Any; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.Result; import org.apache.arrow.flight.sql.impl.FlightSql; /** Typed StreamListener for cancelQuery. */ @SuppressWarnings("deprecation") class CancelListener implements FlightProducer.StreamListener { private final FlightProducer.StreamListener listener; CancelListener(FlightProducer.StreamListener listener) { this.listener = listener; } @Override public void onNext(CancelResult val) { FlightSql.ActionCancelQueryResult result = FlightSql.ActionCancelQueryResult.newBuilder().setResult(val.toProtocol()).build(); listener.onNext(new Result(Any.pack(result).toByteArray())); } @Override public void onError(Throwable t) { listener.onError(t); } @Override public void onCompleted() { listener.onCompleted(); } } CancelResult.java000066400000000000000000000032151515451615200334500ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql; import org.apache.arrow.flight.sql.impl.FlightSql; /** * The result of cancelling a query. * * @deprecated Prefer {@link org.apache.arrow.flight.CancelStatus}. */ @Deprecated public enum CancelResult { UNSPECIFIED, CANCELLED, CANCELLING, NOT_CANCELLABLE, ; FlightSql.ActionCancelQueryResult.CancelResult toProtocol() { switch (this) { default: case UNSPECIFIED: return FlightSql.ActionCancelQueryResult.CancelResult.CANCEL_RESULT_UNSPECIFIED; case CANCELLED: return FlightSql.ActionCancelQueryResult.CancelResult.CANCEL_RESULT_CANCELLED; case CANCELLING: return FlightSql.ActionCancelQueryResult.CancelResult.CANCEL_RESULT_CANCELLING; case NOT_CANCELLABLE: return FlightSql.ActionCancelQueryResult.CancelResult.CANCEL_RESULT_NOT_CANCELLABLE; } } } CancelStatusListener.java000066400000000000000000000031401515451615200351600ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql; import org.apache.arrow.flight.CancelFlightInfoResult; import org.apache.arrow.flight.CancelStatus; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.Result; /** Typed StreamListener for cancelFlightInfo. */ class CancelStatusListener implements FlightProducer.StreamListener { private final FlightProducer.StreamListener listener; CancelStatusListener(FlightProducer.StreamListener listener) { this.listener = listener; } @Override public void onNext(CancelStatus val) { listener.onNext(new Result(new CancelFlightInfoResult(val).serialize().array())); } @Override public void onError(Throwable t) { listener.onError(t); } @Override public void onCompleted() { listener.onCompleted(); } } CloseSessionResultListener.java000066400000000000000000000030621515451615200364020ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql; import org.apache.arrow.flight.CloseSessionResult; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.Result; /** Typed StreamListener for closeSession. */ public class CloseSessionResultListener implements FlightProducer.StreamListener { private final FlightProducer.StreamListener listener; CloseSessionResultListener(FlightProducer.StreamListener listener) { this.listener = listener; } @Override public void onNext(CloseSessionResult val) { listener.onNext(new Result(val.serialize().array())); } @Override public void onError(Throwable t) { listener.onError(t); } @Override public void onCompleted() { listener.onCompleted(); } } FlightEndpointListener.java000066400000000000000000000030411515451615200355050ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.Result; /** Typed StreamListener for renewFlightEndpoint. */ public class FlightEndpointListener implements FlightProducer.StreamListener { private final FlightProducer.StreamListener listener; FlightEndpointListener(FlightProducer.StreamListener listener) { this.listener = listener; } @Override public void onNext(FlightEndpoint val) { listener.onNext(new Result(val.serialize().array())); } @Override public void onError(Throwable t) { listener.onError(t); } @Override public void onCompleted() { listener.onCompleted(); } } FlightSqlClient.java000066400000000000000000001664321515451615200341330ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql; import static java.util.Objects.isNull; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionBeginSavepointRequest; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionBeginSavepointResult; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionBeginTransactionRequest; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionBeginTransactionResult; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionCancelQueryRequest; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionCancelQueryResult; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionClosePreparedStatementRequest; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementRequest; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedSubstraitPlanRequest; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionEndSavepointRequest; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionEndTransactionRequest; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCatalogs; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCrossReference; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetDbSchemas; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetExportedKeys; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetImportedKeys; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetPrimaryKeys; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetSqlInfo; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTableTypes; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTables; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetXdbcTypeInfo; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementUpdate; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementQuery; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementSubstraitPlan; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementUpdate; import static org.apache.arrow.flight.sql.impl.FlightSql.DoPutUpdateResult; import static org.apache.arrow.flight.sql.impl.FlightSql.SqlInfo; import com.google.protobuf.Any; import com.google.protobuf.ByteString; import com.google.protobuf.InvalidProtocolBufferException; import java.io.ByteArrayInputStream; import java.io.IOException; import java.nio.channels.Channels; import java.util.Arrays; import java.util.Collections; import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.concurrent.ExecutionException; import java.util.function.Consumer; import java.util.stream.Collectors; import org.apache.arrow.flight.Action; import org.apache.arrow.flight.CallOption; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.CancelFlightInfoRequest; import org.apache.arrow.flight.CancelFlightInfoResult; import org.apache.arrow.flight.CloseSessionRequest; import org.apache.arrow.flight.CloseSessionResult; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.GetSessionOptionsRequest; import org.apache.arrow.flight.GetSessionOptionsResult; import org.apache.arrow.flight.PutResult; import org.apache.arrow.flight.RenewFlightEndpointRequest; import org.apache.arrow.flight.Result; import org.apache.arrow.flight.SchemaResult; import org.apache.arrow.flight.SetSessionOptionsRequest; import org.apache.arrow.flight.SetSessionOptionsResult; import org.apache.arrow.flight.SyncPutListener; import org.apache.arrow.flight.Ticket; import org.apache.arrow.flight.sql.impl.FlightSql; import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementResult; import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementQuery; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions; import org.apache.arrow.flight.sql.util.TableRef; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.VectorSchemaRoot; 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; /** Flight client with Flight SQL semantics. */ public class FlightSqlClient implements AutoCloseable { private final FlightClient client; public FlightSqlClient(final FlightClient client) { this.client = Objects.requireNonNull(client, "Client cannot be null!"); } /** * Execute a query on the server. * * @param query The query to execute. * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo execute(final String query, final CallOption... options) { return execute(query, /*transaction*/ null, options); } /** * Execute a query on the server. * * @param query The query to execute. * @param transaction The transaction that this query is part of. * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo execute( final String query, Transaction transaction, final CallOption... options) { final CommandStatementQuery.Builder builder = CommandStatementQuery.newBuilder().setQuery(query); if (transaction != null) { builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId())); } final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); return client.getInfo(descriptor, options); } /** * Execute a Substrait plan on the server. * * @param plan The Substrait plan to execute. * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo executeSubstrait(SubstraitPlan plan, CallOption... options) { return executeSubstrait(plan, /*transaction*/ null, options); } /** * Execute a Substrait plan on the server. * * @param plan The Substrait plan to execute. * @param transaction The transaction that this query is part of. * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo executeSubstrait( SubstraitPlan plan, Transaction transaction, CallOption... options) { final CommandStatementSubstraitPlan.Builder builder = CommandStatementSubstraitPlan.newBuilder(); builder .getPlanBuilder() .setPlan(ByteString.copyFrom(plan.getPlan())) .setVersion(plan.getVersion()); if (transaction != null) { builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId())); } final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); return client.getInfo(descriptor, options); } /** Get the schema of the result set of a query. */ public SchemaResult getExecuteSchema( String query, Transaction transaction, CallOption... options) { final CommandStatementQuery.Builder builder = CommandStatementQuery.newBuilder(); builder.setQuery(query); if (transaction != null) { builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId())); } final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); return client.getSchema(descriptor, options); } /** Get the schema of the result set of a query. */ public SchemaResult getExecuteSchema(String query, CallOption... options) { return getExecuteSchema(query, /*transaction*/ null, options); } /** Get the schema of the result set of a Substrait plan. */ public SchemaResult getExecuteSubstraitSchema( SubstraitPlan plan, Transaction transaction, final CallOption... options) { final CommandStatementSubstraitPlan.Builder builder = CommandStatementSubstraitPlan.newBuilder(); builder .getPlanBuilder() .setPlan(ByteString.copyFrom(plan.getPlan())) .setVersion(plan.getVersion()); if (transaction != null) { builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId())); } final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); return client.getSchema(descriptor, options); } /** Get the schema of the result set of a Substrait plan. */ public SchemaResult getExecuteSubstraitSchema( SubstraitPlan substraitPlan, final CallOption... options) { return getExecuteSubstraitSchema(substraitPlan, /*transaction*/ null, options); } /** * Execute a bulk ingest on the server. * * @param data data to be ingested * @param ingestOptions options for the ingest request. * @param options RPC-layer hints for this call. * @return the number of rows affected. */ public long executeIngest( final VectorSchemaRoot data, final ExecuteIngestOptions ingestOptions, final CallOption... options) { return executeIngest(data, ingestOptions, /*transaction*/ null, options); } /** * Execute a bulk ingest on the server. * * @param dataReader data stream to be ingested * @param ingestOptions options for the ingest request. * @param options RPC-layer hints for this call. * @return the number of rows affected. */ public long executeIngest( final ArrowReader dataReader, final ExecuteIngestOptions ingestOptions, final CallOption... options) { return executeIngest(dataReader, ingestOptions, /*transaction*/ null, options); } /** * Execute a bulk ingest on the server. * * @param data data to be ingested * @param ingestOptions options for the ingest request. * @param transaction The transaction that this ingest request is part of. * @param options RPC-layer hints for this call. * @return the number of rows affected. */ public long executeIngest( final VectorSchemaRoot data, final ExecuteIngestOptions ingestOptions, Transaction transaction, final CallOption... options) { return executeIngest( data, ingestOptions, transaction, FlightClient.ClientStreamListener::putNext, options); } /** * Execute a bulk ingest on the server. * * @param dataReader data stream to be ingested * @param ingestOptions options for the ingest request. * @param transaction The transaction that this ingest request is part of. * @param options RPC-layer hints for this call. * @return the number of rows affected. */ public long executeIngest( final ArrowReader dataReader, final ExecuteIngestOptions ingestOptions, Transaction transaction, final CallOption... options) { try { return executeIngest( dataReader.getVectorSchemaRoot(), ingestOptions, transaction, listener -> { while (true) { try { if (!dataReader.loadNextBatch()) { break; } } catch (IOException e) { throw CallStatus.UNKNOWN.withCause(e).toRuntimeException(); } listener.putNext(); } }, options); } catch (IOException e) { throw CallStatus.UNKNOWN.withCause(e).toRuntimeException(); } } private long executeIngest( final VectorSchemaRoot data, final ExecuteIngestOptions ingestOptions, final Transaction transaction, final Consumer dataPutter, final CallOption... options) { try { final CommandStatementIngest.Builder builder = CommandStatementIngest.newBuilder(); if (transaction != null) { builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId())); } ingestOptions.updateCommandBuilder(builder); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); try (final SyncPutListener putListener = new SyncPutListener()) { final FlightClient.ClientStreamListener listener = client.startPut(descriptor, data, putListener, options); dataPutter.accept(listener); listener.completed(); listener.getResult(); try (final PutResult result = putListener.read()) { final DoPutUpdateResult doPutUpdateResult = DoPutUpdateResult.parseFrom(result.getApplicationMetadata().nioBuffer()); return doPutUpdateResult.getRecordCount(); } } } catch (final InterruptedException e) { throw CallStatus.CANCELLED.withCause(e).toRuntimeException(); } catch (final ExecutionException e) { throw CallStatus.CANCELLED.withCause(e.getCause()).toRuntimeException(); } catch (final InvalidProtocolBufferException e) { throw CallStatus.INTERNAL.withCause(e).toRuntimeException(); } } /** * Execute an update query on the server. * * @param query The query to execute. * @param options RPC-layer hints for this call. * @return the number of rows affected. */ public long executeUpdate(final String query, final CallOption... options) { return executeUpdate(query, /*transaction*/ null, options); } /** * Execute an update query on the server. * * @param query The query to execute. * @param transaction The transaction that this query is part of. * @param options RPC-layer hints for this call. * @return the number of rows affected. */ public long executeUpdate( final String query, Transaction transaction, final CallOption... options) { final CommandStatementUpdate.Builder builder = CommandStatementUpdate.newBuilder().setQuery(query); if (transaction != null) { builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId())); } final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); try (final SyncPutListener putListener = new SyncPutListener()) { final FlightClient.ClientStreamListener listener = client.startPut(descriptor, VectorSchemaRoot.of(), putListener, options); try (final PutResult result = putListener.read()) { final DoPutUpdateResult doPutUpdateResult = DoPutUpdateResult.parseFrom(result.getApplicationMetadata().nioBuffer()); return doPutUpdateResult.getRecordCount(); } finally { listener.getResult(); } } catch (final InterruptedException e) { throw CallStatus.CANCELLED.withCause(e).toRuntimeException(); } catch (final ExecutionException e) { throw CallStatus.CANCELLED.withCause(e.getCause()).toRuntimeException(); } catch (final InvalidProtocolBufferException e) { throw CallStatus.INTERNAL.withCause(e).toRuntimeException(); } } /** * Execute an update query on the server. * * @param plan The Substrait plan to execute. * @param options RPC-layer hints for this call. * @return the number of rows affected. */ public long executeSubstraitUpdate(SubstraitPlan plan, CallOption... options) { return executeSubstraitUpdate(plan, /*transaction*/ null, options); } /** * Execute an update query on the server. * * @param plan The Substrait plan to execute. * @param transaction The transaction that this query is part of. * @param options RPC-layer hints for this call. * @return the number of rows affected. */ public long executeSubstraitUpdate( SubstraitPlan plan, Transaction transaction, CallOption... options) { final CommandStatementSubstraitPlan.Builder builder = CommandStatementSubstraitPlan.newBuilder(); builder .getPlanBuilder() .setPlan(ByteString.copyFrom(plan.getPlan())) .setVersion(plan.getVersion()); if (transaction != null) { builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId())); } final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); try (final SyncPutListener putListener = new SyncPutListener()) { final FlightClient.ClientStreamListener listener = client.startPut(descriptor, VectorSchemaRoot.of(), putListener, options); try (final PutResult result = putListener.read()) { final DoPutUpdateResult doPutUpdateResult = DoPutUpdateResult.parseFrom(result.getApplicationMetadata().nioBuffer()); return doPutUpdateResult.getRecordCount(); } finally { listener.getResult(); } } catch (final InterruptedException e) { throw CallStatus.CANCELLED.withCause(e).toRuntimeException(); } catch (final ExecutionException e) { throw CallStatus.CANCELLED.withCause(e.getCause()).toRuntimeException(); } catch (final InvalidProtocolBufferException e) { throw CallStatus.INTERNAL.withCause(e).toRuntimeException(); } } /** * Request a list of catalogs. * * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo getCatalogs(final CallOption... options) { final CommandGetCatalogs.Builder builder = CommandGetCatalogs.newBuilder(); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); return client.getInfo(descriptor, options); } /** * Get the schema of {@link #getCatalogs(CallOption...)} from the server. * *

    Should be identical to {@link FlightSqlProducer.Schemas#GET_CATALOGS_SCHEMA}. */ public SchemaResult getCatalogsSchema(final CallOption... options) { final CommandGetCatalogs command = CommandGetCatalogs.getDefaultInstance(); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray()); return client.getSchema(descriptor, options); } /** * Request a list of schemas. * * @param catalog The catalog. * @param dbSchemaFilterPattern The schema filter pattern. * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo getSchemas( final String catalog, final String dbSchemaFilterPattern, final CallOption... options) { final CommandGetDbSchemas.Builder builder = CommandGetDbSchemas.newBuilder(); if (catalog != null) { builder.setCatalog(catalog); } if (dbSchemaFilterPattern != null) { builder.setDbSchemaFilterPattern(dbSchemaFilterPattern); } final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); return client.getInfo(descriptor, options); } /** * Get the schema of {@link #getSchemas(String, String, CallOption...)} from the server. * *

    Should be identical to {@link FlightSqlProducer.Schemas#GET_SCHEMAS_SCHEMA}. */ public SchemaResult getSchemasSchema(final CallOption... options) { final CommandGetDbSchemas command = CommandGetDbSchemas.getDefaultInstance(); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray()); return client.getSchema(descriptor, options); } /** * Get schema for a stream. * * @param descriptor The descriptor for the stream. * @param options RPC-layer hints for this call. */ public SchemaResult getSchema(FlightDescriptor descriptor, CallOption... options) { return client.getSchema(descriptor, options); } /** * Retrieve a stream from the server. * * @param ticket The ticket granting access to the data stream. * @param options RPC-layer hints for this call. */ public FlightStream getStream(Ticket ticket, CallOption... options) { return client.getStream(ticket, options); } /** * Request a set of Flight SQL metadata. * * @param info The set of metadata to retrieve. None to retrieve all metadata. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo getSqlInfo(final SqlInfo... info) { return getSqlInfo(info, new CallOption[0]); } /** * Request a set of Flight SQL metadata. * * @param info The set of metadata to retrieve. None to retrieve all metadata. * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo getSqlInfo(final SqlInfo[] info, final CallOption... options) { final int[] infoNumbers = Arrays.stream(info).mapToInt(SqlInfo::getNumber).toArray(); return getSqlInfo(infoNumbers, options); } /** * Request a set of Flight SQL metadata. Use this method if you would like to retrieve custom * metadata, where the custom metadata key values start from 10_000. * * @param info The set of metadata to retrieve. None to retrieve all metadata. * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo getSqlInfo(final int[] info, final CallOption... options) { return getSqlInfo(Arrays.stream(info).boxed().collect(Collectors.toList()), options); } /** * Request a set of Flight SQL metadata. Use this method if you would like to retrieve custom * metadata, where the custom metadata key values start from 10_000. * * @param info The set of metadata to retrieve. None to retrieve all metadata. * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo getSqlInfo(final Iterable info, final CallOption... options) { final CommandGetSqlInfo.Builder builder = CommandGetSqlInfo.newBuilder(); builder.addAllInfo(info); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); return client.getInfo(descriptor, options); } /** * Get the schema of {@link #getSqlInfo(SqlInfo...)} from the server. * *

    Should be identical to {@link FlightSqlProducer.Schemas#GET_SQL_INFO_SCHEMA}. */ public SchemaResult getSqlInfoSchema(final CallOption... options) { final CommandGetSqlInfo command = CommandGetSqlInfo.getDefaultInstance(); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray()); return client.getSchema(descriptor, options); } /** * Request the information about the data types supported related to a filter data type. * * @param dataType the data type to be used as filter. * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo getXdbcTypeInfo(final int dataType, final CallOption... options) { final CommandGetXdbcTypeInfo.Builder builder = CommandGetXdbcTypeInfo.newBuilder(); builder.setDataType(dataType); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); return client.getInfo(descriptor, options); } /** * Request the information about all the data types supported. * * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo getXdbcTypeInfo(final CallOption... options) { final CommandGetXdbcTypeInfo.Builder builder = CommandGetXdbcTypeInfo.newBuilder(); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); return client.getInfo(descriptor, options); } /** * Get the schema of {@link #getXdbcTypeInfo(CallOption...)} from the server. * *

    Should be identical to {@link FlightSqlProducer.Schemas#GET_TYPE_INFO_SCHEMA}. */ public SchemaResult getXdbcTypeInfoSchema(final CallOption... options) { final CommandGetXdbcTypeInfo command = CommandGetXdbcTypeInfo.getDefaultInstance(); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray()); return client.getSchema(descriptor, options); } /** * Request a list of tables. * * @param catalog The catalog. * @param dbSchemaFilterPattern The schema filter pattern. * @param tableFilterPattern The table filter pattern. * @param tableTypes The table types to include. * @param includeSchema True to include the schema upon return, false to not include the schema. * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo getTables( final String catalog, final String dbSchemaFilterPattern, final String tableFilterPattern, final List tableTypes, final boolean includeSchema, final CallOption... options) { final CommandGetTables.Builder builder = CommandGetTables.newBuilder(); if (catalog != null) { builder.setCatalog(catalog); } if (dbSchemaFilterPattern != null) { builder.setDbSchemaFilterPattern(dbSchemaFilterPattern); } if (tableFilterPattern != null) { builder.setTableNameFilterPattern(tableFilterPattern); } if (tableTypes != null) { builder.addAllTableTypes(tableTypes); } builder.setIncludeSchema(includeSchema); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); return client.getInfo(descriptor, options); } /** * Get the schema of {@link #getTables(String, String, String, List, boolean, CallOption...)} from * the server. * *

    Should be identical to {@link FlightSqlProducer.Schemas#GET_TABLES_SCHEMA} or {@link * FlightSqlProducer.Schemas#GET_TABLES_SCHEMA_NO_SCHEMA}. */ public SchemaResult getTablesSchema(boolean includeSchema, final CallOption... options) { final CommandGetTables command = CommandGetTables.newBuilder().setIncludeSchema(includeSchema).build(); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray()); return client.getSchema(descriptor, options); } /** * Request the primary keys for a table. * * @param tableRef An object which hold info about catalog, dbSchema and table. * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo getPrimaryKeys(final TableRef tableRef, final CallOption... options) { final CommandGetPrimaryKeys.Builder builder = CommandGetPrimaryKeys.newBuilder(); if (tableRef.getCatalog() != null) { builder.setCatalog(tableRef.getCatalog()); } if (tableRef.getDbSchema() != null) { builder.setDbSchema(tableRef.getDbSchema()); } Objects.requireNonNull(tableRef.getTable()); builder.setTable(tableRef.getTable()); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); return client.getInfo(descriptor, options); } /** * Get the schema of {@link #getPrimaryKeys(TableRef, CallOption...)} from the server. * *

    Should be identical to {@link FlightSqlProducer.Schemas#GET_PRIMARY_KEYS_SCHEMA}. */ public SchemaResult getPrimaryKeysSchema(final CallOption... options) { final CommandGetPrimaryKeys command = CommandGetPrimaryKeys.getDefaultInstance(); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray()); return client.getSchema(descriptor, options); } /** * Retrieves a description about the foreign key columns that reference the primary key columns of * the given table. * * @param tableRef An object which hold info about catalog, dbSchema and table. * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo getExportedKeys(final TableRef tableRef, final CallOption... options) { Objects.requireNonNull(tableRef.getTable(), "Table cannot be null."); final CommandGetExportedKeys.Builder builder = CommandGetExportedKeys.newBuilder(); if (tableRef.getCatalog() != null) { builder.setCatalog(tableRef.getCatalog()); } if (tableRef.getDbSchema() != null) { builder.setDbSchema(tableRef.getDbSchema()); } Objects.requireNonNull(tableRef.getTable()); builder.setTable(tableRef.getTable()); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); return client.getInfo(descriptor, options); } /** * Get the schema of {@link #getExportedKeys(TableRef, CallOption...)} from the server. * *

    Should be identical to {@link FlightSqlProducer.Schemas#GET_EXPORTED_KEYS_SCHEMA}. */ public SchemaResult getExportedKeysSchema(final CallOption... options) { final CommandGetExportedKeys command = CommandGetExportedKeys.getDefaultInstance(); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray()); return client.getSchema(descriptor, options); } /** * Retrieves the foreign key columns for the given table. * * @param tableRef An object which hold info about catalog, dbSchema and table. * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo getImportedKeys(final TableRef tableRef, final CallOption... options) { Objects.requireNonNull(tableRef.getTable(), "Table cannot be null."); final CommandGetImportedKeys.Builder builder = CommandGetImportedKeys.newBuilder(); if (tableRef.getCatalog() != null) { builder.setCatalog(tableRef.getCatalog()); } if (tableRef.getDbSchema() != null) { builder.setDbSchema(tableRef.getDbSchema()); } Objects.requireNonNull(tableRef.getTable()); builder.setTable(tableRef.getTable()); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); return client.getInfo(descriptor, options); } /** * Get the schema of {@link #getImportedKeys(TableRef, CallOption...)} from the server. * *

    Should be identical to {@link FlightSqlProducer.Schemas#GET_IMPORTED_KEYS_SCHEMA}. */ public SchemaResult getImportedKeysSchema(final CallOption... options) { final CommandGetImportedKeys command = CommandGetImportedKeys.getDefaultInstance(); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray()); return client.getSchema(descriptor, options); } /** * Retrieves a description of the foreign key columns that reference the given table's primary key * columns (the foreign keys exported by a table). * * @param pkTableRef An object which hold info about catalog, dbSchema and table from a primary * table. * @param fkTableRef An object which hold info about catalog, dbSchema and table from a foreign * table. * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo getCrossReference( final TableRef pkTableRef, final TableRef fkTableRef, final CallOption... options) { Objects.requireNonNull(pkTableRef.getTable(), "Parent Table cannot be null."); Objects.requireNonNull(fkTableRef.getTable(), "Foreign Table cannot be null."); final CommandGetCrossReference.Builder builder = CommandGetCrossReference.newBuilder(); if (pkTableRef.getCatalog() != null) { builder.setPkCatalog(pkTableRef.getCatalog()); } if (pkTableRef.getDbSchema() != null) { builder.setPkDbSchema(pkTableRef.getDbSchema()); } if (fkTableRef.getCatalog() != null) { builder.setFkCatalog(fkTableRef.getCatalog()); } if (fkTableRef.getDbSchema() != null) { builder.setFkDbSchema(fkTableRef.getDbSchema()); } builder.setPkTable(pkTableRef.getTable()); builder.setFkTable(fkTableRef.getTable()); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); return client.getInfo(descriptor, options); } /** * Get the schema of {@link #getCrossReference(TableRef, TableRef, CallOption...)} from the * server. * *

    Should be identical to {@link FlightSqlProducer.Schemas#GET_CROSS_REFERENCE_SCHEMA}. */ public SchemaResult getCrossReferenceSchema(final CallOption... options) { final CommandGetCrossReference command = CommandGetCrossReference.getDefaultInstance(); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray()); return client.getSchema(descriptor, options); } /** * Request a list of table types. * * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo getTableTypes(final CallOption... options) { final CommandGetTableTypes.Builder builder = CommandGetTableTypes.newBuilder(); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(builder.build()).toByteArray()); return client.getInfo(descriptor, options); } /** * Get the schema of {@link #getTableTypes(CallOption...)} from the server. * *

    Should be identical to {@link FlightSqlProducer.Schemas#GET_TABLE_TYPES_SCHEMA}. */ public SchemaResult getTableTypesSchema(final CallOption... options) { final CommandGetTableTypes command = CommandGetTableTypes.getDefaultInstance(); final FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(command).toByteArray()); return client.getSchema(descriptor, options); } /** * Create a prepared statement for a SQL query on the server. * * @param query The query to prepare. * @param options RPC-layer hints for this call. * @return The representation of the prepared statement which exists on the server. */ public PreparedStatement prepare(String query, CallOption... options) { return prepare(query, /*transaction*/ null, options); } /** * Create a prepared statement for a SQL query on the server. * * @param query The query to prepare. * @param transaction The transaction that this query is part of. * @param options RPC-layer hints for this call. * @return The representation of the prepared statement which exists on the server. */ public PreparedStatement prepare(String query, Transaction transaction, CallOption... options) { ActionCreatePreparedStatementRequest.Builder builder = ActionCreatePreparedStatementRequest.newBuilder().setQuery(query); if (transaction != null) { builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId())); } return new PreparedStatement( client, new Action( FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_STATEMENT.getType(), Any.pack(builder.build()).toByteArray()), options); } /** * Create a prepared statement for a Substrait plan on the server. * * @param plan The query to prepare. * @param options RPC-layer hints for this call. * @return The representation of the prepared statement which exists on the server. */ public PreparedStatement prepare(SubstraitPlan plan, CallOption... options) { return prepare(plan, /*transaction*/ null, options); } /** * Create a prepared statement for a Substrait plan on the server. * * @param plan The query to prepare. * @param transaction The transaction that this query is part of. * @param options RPC-layer hints for this call. * @return The representation of the prepared statement which exists on the server. */ public PreparedStatement prepare( SubstraitPlan plan, Transaction transaction, CallOption... options) { ActionCreatePreparedSubstraitPlanRequest.Builder builder = ActionCreatePreparedSubstraitPlanRequest.newBuilder(); builder .getPlanBuilder() .setPlan(ByteString.copyFrom(plan.getPlan())) .setVersion(plan.getVersion()); if (transaction != null) { builder.setTransactionId(ByteString.copyFrom(transaction.getTransactionId())); } return new PreparedStatement( client, new Action( FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_SUBSTRAIT_PLAN.getType(), Any.pack(builder.build()).toByteArray()), options); } /** Begin a transaction. */ public Transaction beginTransaction(CallOption... options) { final Action action = new Action( FlightSqlUtils.FLIGHT_SQL_BEGIN_TRANSACTION.getType(), Any.pack(ActionBeginTransactionRequest.getDefaultInstance()).toByteArray()); final Iterator preparedStatementResults = client.doAction(action, options); final ActionBeginTransactionResult result = FlightSqlUtils.unpackAndParseOrThrow( preparedStatementResults.next().getBody(), ActionBeginTransactionResult.class); preparedStatementResults.forEachRemaining((ignored) -> {}); if (result.getTransactionId().isEmpty()) { throw CallStatus.INTERNAL .withDescription("Server returned an empty transaction ID") .toRuntimeException(); } return new Transaction(result.getTransactionId().toByteArray()); } /** Create a savepoint within a transaction. */ public Savepoint beginSavepoint(Transaction transaction, String name, CallOption... options) { Preconditions.checkArgument( transaction.getTransactionId().length != 0, "Transaction must be initialized"); ActionBeginSavepointRequest request = ActionBeginSavepointRequest.newBuilder() .setTransactionId(ByteString.copyFrom(transaction.getTransactionId())) .setName(name) .build(); final Action action = new Action( FlightSqlUtils.FLIGHT_SQL_BEGIN_SAVEPOINT.getType(), Any.pack(request).toByteArray()); final Iterator preparedStatementResults = client.doAction(action, options); final ActionBeginSavepointResult result = FlightSqlUtils.unpackAndParseOrThrow( preparedStatementResults.next().getBody(), ActionBeginSavepointResult.class); preparedStatementResults.forEachRemaining((ignored) -> {}); if (result.getSavepointId().isEmpty()) { throw CallStatus.INTERNAL .withDescription("Server returned an empty transaction ID") .toRuntimeException(); } return new Savepoint(result.getSavepointId().toByteArray()); } /** Commit a transaction. */ public void commit(Transaction transaction, CallOption... options) { Preconditions.checkArgument( transaction.getTransactionId().length != 0, "Transaction must be initialized"); ActionEndTransactionRequest request = ActionEndTransactionRequest.newBuilder() .setTransactionId(ByteString.copyFrom(transaction.getTransactionId())) .setActionValue( ActionEndTransactionRequest.EndTransaction.END_TRANSACTION_COMMIT.getNumber()) .build(); final Action action = new Action( FlightSqlUtils.FLIGHT_SQL_END_TRANSACTION.getType(), Any.pack(request).toByteArray()); final Iterator preparedStatementResults = client.doAction(action, options); preparedStatementResults.forEachRemaining((ignored) -> {}); } /** Release a savepoint. */ public void release(Savepoint savepoint, CallOption... options) { Preconditions.checkArgument( savepoint.getSavepointId().length != 0, "Savepoint must be initialized"); ActionEndSavepointRequest request = ActionEndSavepointRequest.newBuilder() .setSavepointId(ByteString.copyFrom(savepoint.getSavepointId())) .setActionValue( ActionEndSavepointRequest.EndSavepoint.END_SAVEPOINT_RELEASE.getNumber()) .build(); final Action action = new Action( FlightSqlUtils.FLIGHT_SQL_END_SAVEPOINT.getType(), Any.pack(request).toByteArray()); final Iterator preparedStatementResults = client.doAction(action, options); preparedStatementResults.forEachRemaining((ignored) -> {}); } /** Rollback a transaction. */ public void rollback(Transaction transaction, CallOption... options) { Preconditions.checkArgument( transaction.getTransactionId().length != 0, "Transaction must be initialized"); ActionEndTransactionRequest request = ActionEndTransactionRequest.newBuilder() .setTransactionId(ByteString.copyFrom(transaction.getTransactionId())) .setActionValue( ActionEndTransactionRequest.EndTransaction.END_TRANSACTION_ROLLBACK.getNumber()) .build(); final Action action = new Action( FlightSqlUtils.FLIGHT_SQL_END_TRANSACTION.getType(), Any.pack(request).toByteArray()); final Iterator preparedStatementResults = client.doAction(action, options); preparedStatementResults.forEachRemaining((ignored) -> {}); } /** Rollback to a savepoint. */ public void rollback(Savepoint savepoint, CallOption... options) { Preconditions.checkArgument( savepoint.getSavepointId().length != 0, "Savepoint must be initialized"); ActionEndSavepointRequest request = ActionEndSavepointRequest.newBuilder() .setSavepointId(ByteString.copyFrom(savepoint.getSavepointId())) .setActionValue( ActionEndSavepointRequest.EndSavepoint.END_SAVEPOINT_RELEASE.getNumber()) .build(); final Action action = new Action( FlightSqlUtils.FLIGHT_SQL_END_SAVEPOINT.getType(), Any.pack(request).toByteArray()); final Iterator preparedStatementResults = client.doAction(action, options); preparedStatementResults.forEachRemaining((ignored) -> {}); } /** * Cancel execution of a distributed query. * * @param request The query to cancel. * @param options Call options. * @return The server response. */ public CancelFlightInfoResult cancelFlightInfo( CancelFlightInfoRequest request, CallOption... options) { return client.cancelFlightInfo(request, options); } /** * Explicitly cancel a running query. * *

    This lets a single client explicitly cancel work, no matter how many clients are * involved/whether the query is distributed or not, given server support. The * transaction/statement is not rolled back; it is the application's job to commit or rollback as * appropriate. This only indicates the client no longer wishes to read the remainder of the query * results or continue submitting data. * * @deprecated Prefer {@link #cancelFlightInfo}. */ @Deprecated public CancelResult cancelQuery(FlightInfo info, CallOption... options) { ActionCancelQueryRequest request = ActionCancelQueryRequest.newBuilder() .setInfo(ByteString.copyFrom(info.serialize())) .build(); final Action action = new Action( FlightSqlUtils.FLIGHT_SQL_CANCEL_QUERY.getType(), Any.pack(request).toByteArray()); final Iterator preparedStatementResults = client.doAction(action, options); final ActionCancelQueryResult result = FlightSqlUtils.unpackAndParseOrThrow( preparedStatementResults.next().getBody(), ActionCancelQueryResult.class); preparedStatementResults.forEachRemaining((ignored) -> {}); switch (result.getResult()) { case CANCEL_RESULT_UNSPECIFIED: return CancelResult.UNSPECIFIED; case CANCEL_RESULT_CANCELLED: return CancelResult.CANCELLED; case CANCEL_RESULT_CANCELLING: return CancelResult.CANCELLING; case CANCEL_RESULT_NOT_CANCELLABLE: return CancelResult.NOT_CANCELLABLE; case UNRECOGNIZED: default: throw CallStatus.INTERNAL .withDescription("Unknown result: " + result.getResult()) .toRuntimeException(); } } /** * Request the server to extend the lifetime of a query result set. * * @param request The result set partition. * @param options Call options. * @return The new endpoint with an updated expiration time. */ public FlightEndpoint renewFlightEndpoint( RenewFlightEndpointRequest request, CallOption... options) { return client.renewFlightEndpoint(request, options); } public SetSessionOptionsResult setSessionOptions( SetSessionOptionsRequest request, CallOption... options) { return client.setSessionOptions(request, options); } public GetSessionOptionsResult getSessionOptions( GetSessionOptionsRequest request, CallOption... options) { return client.getSessionOptions(request, options); } public CloseSessionResult closeSession(CloseSessionRequest request, CallOption... options) { return client.closeSession(request, options); } @Override public void close() throws Exception { AutoCloseables.close(client); } /** Class to encapsulate Flight SQL bulk ingest request options. * */ public static class ExecuteIngestOptions { private final String table; private final TableDefinitionOptions tableDefinitionOptions; private final boolean useTemporaryTable; private final String catalog; private final String schema; private final Map options; /** * Constructor. * * @param table The table to load data into. * @param tableDefinitionOptions The behavior for handling the table definition. * @param catalog The catalog of the destination table to load data into. If null, a * backend-specific default may be used. * @param schema The schema of the destination table to load data into. If null, a * backend-specific default may be used. * @param options Backend-specific options. Can be null if there are no options to be set. */ public ExecuteIngestOptions( String table, TableDefinitionOptions tableDefinitionOptions, String catalog, String schema, Map options) { this(table, tableDefinitionOptions, false, catalog, schema, options); } /** * Constructor. * * @param table The table to load data into. * @param tableDefinitionOptions The behavior for handling the table definition. * @param useTemporaryTable Use a temporary table for bulk ingestion. Temporary table may get * placed in a backend-specific schema and/or catalog and gets dropped at the end of the * session. If backend does not support ingesting using a temporary table or an explicit * choice of schema or catalog is incompatible with the server's namespacing decision, an * error is returned as part of {@link #executeIngest} request. * @param catalog The catalog of the destination table to load data into. If null, a * backend-specific default may be used. * @param schema The schema of the destination table to load data into. If null, a * backend-specific default may be used. * @param options Backend-specific options. Can be null if there are no options to be set. */ public ExecuteIngestOptions( String table, TableDefinitionOptions tableDefinitionOptions, boolean useTemporaryTable, String catalog, String schema, Map options) { this.table = table; this.tableDefinitionOptions = tableDefinitionOptions; this.useTemporaryTable = useTemporaryTable; this.catalog = catalog; this.schema = schema; this.options = options; } protected void updateCommandBuilder(CommandStatementIngest.Builder builder) { builder.setTable(table); builder.setTableDefinitionOptions(tableDefinitionOptions); builder.setTemporary(useTemporaryTable); if (!isNull(catalog)) { builder.setCatalog(catalog); } if (!isNull(schema)) { builder.setSchema(schema); } if (!isNull(options)) { builder.putAllOptions(options); } } } /** Helper class to encapsulate Flight SQL prepared statement logic. */ public static class PreparedStatement implements AutoCloseable { private final FlightClient client; private final ActionCreatePreparedStatementResult preparedStatementResult; private VectorSchemaRoot parameterBindingRoot; private boolean isClosed; private Schema resultSetSchema; private Schema parameterSchema; PreparedStatement(FlightClient client, Action action, CallOption... options) { this.client = client; final Iterator preparedStatementResults = client.doAction(action, options); preparedStatementResult = FlightSqlUtils.unpackAndParseOrThrow( preparedStatementResults.next().getBody(), ActionCreatePreparedStatementResult.class); isClosed = false; } /** * Set the {@link #parameterBindingRoot} containing the parameter binding from a {@link * PreparedStatement} operation. * * @param parameterBindingRoot a {@code VectorSchemaRoot} object containing the values to be * used in the {@code PreparedStatement} setters. */ public void setParameters(final VectorSchemaRoot parameterBindingRoot) { if (parameterBindingRoot == this.parameterBindingRoot) { // Nothing to do if we're attempting to set the same parameters again. return; } clearParameters(); this.parameterBindingRoot = parameterBindingRoot; } /** * Closes the {@link #parameterBindingRoot}, which contains the parameter binding from a {@link * PreparedStatement} operation, releasing its resources. */ public void clearParameters() { if (parameterBindingRoot != null) { parameterBindingRoot.close(); } } /** * Returns the Schema of the resultset. * * @return the Schema of the resultset. */ public Schema getResultSetSchema() { if (resultSetSchema == null) { final ByteString bytes = preparedStatementResult.getDatasetSchema(); resultSetSchema = deserializeSchema(bytes); } return resultSetSchema; } /** * Returns the Schema of the parameters. * * @return the Schema of the parameters. */ public Schema getParameterSchema() { if (parameterSchema == null) { final ByteString bytes = preparedStatementResult.getParameterSchema(); parameterSchema = deserializeSchema(bytes); } return parameterSchema; } /** Get the schema of the result set (should be identical to {@link #getResultSetSchema()}). */ public SchemaResult fetchSchema(CallOption... options) { checkOpen(); final FlightDescriptor descriptor = FlightDescriptor.command( Any.pack( CommandPreparedStatementQuery.newBuilder() .setPreparedStatementHandle( preparedStatementResult.getPreparedStatementHandle()) .build()) .toByteArray()); return client.getSchema(descriptor, options); } private Schema deserializeSchema(final ByteString bytes) { try { return bytes.isEmpty() ? new Schema(Collections.emptyList()) : MessageSerializer.deserializeSchema( new ReadChannel( Channels.newChannel(new ByteArrayInputStream(bytes.toByteArray())))); } catch (final IOException e) { throw new RuntimeException("Failed to deserialize schema", e); } } /** * Executes the prepared statement query on the server. * * @param options RPC-layer hints for this call. * @return a FlightInfo object representing the stream(s) to fetch. */ public FlightInfo execute(final CallOption... options) { checkOpen(); FlightDescriptor descriptor = FlightDescriptor.command( Any.pack( CommandPreparedStatementQuery.newBuilder() .setPreparedStatementHandle( preparedStatementResult.getPreparedStatementHandle()) .build()) .toByteArray()); if (parameterBindingRoot != null && parameterBindingRoot.getRowCount() > 0) { try (final SyncPutListener putListener = putParameters(descriptor, options)) { if (getParameterSchema().getFields().size() > 0 && parameterBindingRoot != null && parameterBindingRoot.getRowCount() > 0) { final PutResult read = putListener.read(); if (read != null) { try (final ArrowBuf metadata = read.getApplicationMetadata()) { final FlightSql.DoPutPreparedStatementResult doPutPreparedStatementResult = FlightSql.DoPutPreparedStatementResult.parseFrom(metadata.nioBuffer()); descriptor = FlightDescriptor.command( Any.pack( CommandPreparedStatementQuery.newBuilder() .setPreparedStatementHandle( doPutPreparedStatementResult.getPreparedStatementHandle()) .build()) .toByteArray()); } } } } catch (final InterruptedException e) { throw CallStatus.CANCELLED.withCause(e).toRuntimeException(); } catch (final ExecutionException e) { throw CallStatus.CANCELLED.withCause(e.getCause()).toRuntimeException(); } catch (final InvalidProtocolBufferException e) { throw CallStatus.INTERNAL.withCause(e).toRuntimeException(); } } return client.getInfo(descriptor, options); } private SyncPutListener putParameters(FlightDescriptor descriptor, CallOption... options) { final SyncPutListener putListener = new SyncPutListener(); FlightClient.ClientStreamListener listener = client.startPut(descriptor, parameterBindingRoot, putListener, options); listener.putNext(); listener.completed(); listener.getResult(); return putListener; } /** * Checks whether this client is open. * * @throws IllegalStateException if client is closed. */ protected final void checkOpen() { Preconditions.checkState(!isClosed, "Statement closed"); } /** * Executes the prepared statement update on the server. * * @param options RPC-layer hints for this call. * @return the count of updated records */ public long executeUpdate(final CallOption... options) { checkOpen(); final FlightDescriptor descriptor = FlightDescriptor.command( Any.pack( CommandPreparedStatementUpdate.newBuilder() .setPreparedStatementHandle( preparedStatementResult.getPreparedStatementHandle()) .build()) .toByteArray()); setParameters(parameterBindingRoot == null ? VectorSchemaRoot.of() : parameterBindingRoot); SyncPutListener putListener = putParameters(descriptor, options); try { final PutResult read = putListener.read(); try (final ArrowBuf metadata = read.getApplicationMetadata()) { final DoPutUpdateResult doPutUpdateResult = DoPutUpdateResult.parseFrom(metadata.nioBuffer()); return doPutUpdateResult.getRecordCount(); } } catch (final InterruptedException e) { throw CallStatus.CANCELLED.withCause(e).toRuntimeException(); } catch (final ExecutionException e) { throw CallStatus.CANCELLED.withCause(e.getCause()).toRuntimeException(); } catch (final InvalidProtocolBufferException e) { throw CallStatus.INTERNAL.withCause(e).toRuntimeException(); } } /** * Closes the client. * * @param options RPC-layer hints for this call. */ public void close(final CallOption... options) { if (isClosed) { return; } isClosed = true; final Action action = new Action( FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT.getType(), Any.pack( ActionClosePreparedStatementRequest.newBuilder() .setPreparedStatementHandle( preparedStatementResult.getPreparedStatementHandle()) .build()) .toByteArray()); final Iterator closePreparedStatementResults = client.doAction(action, options); closePreparedStatementResults.forEachRemaining(result -> {}); clearParameters(); } @Override public void close() { close(new CallOption[0]); } /** * Returns if the prepared statement is already closed. * * @return true if the prepared statement is already closed. */ public boolean isClosed() { return isClosed; } } /** A handle for an active savepoint. */ public static class Savepoint { private final byte[] transactionId; public Savepoint(byte[] transactionId) { this.transactionId = transactionId; } public byte[] getSavepointId() { return transactionId; } } /** A handle for an active transaction. */ public static class Transaction { private final byte[] transactionId; public Transaction(byte[] transactionId) { this.transactionId = transactionId; } public byte[] getTransactionId() { return transactionId; } } /** A wrapper around a Substrait plan and a Substrait version. */ public static final class SubstraitPlan { private final byte[] plan; private final String version; public SubstraitPlan(byte[] plan, String version) { this.plan = Preconditions.checkNotNull(plan); this.version = Preconditions.checkNotNull(version); } public byte[] getPlan() { return plan; } public String getVersion() { return version; } @Override public boolean equals(Object o) { if (this == o) { return true; } if (o == null || getClass() != o.getClass()) { return false; } SubstraitPlan that = (SubstraitPlan) o; if (!Arrays.equals(getPlan(), that.getPlan())) { return false; } return getVersion().equals(that.getVersion()); } @Override public int hashCode() { int result = Arrays.hashCode(getPlan()); result = 31 * result + getVersion().hashCode(); return result; } @Override public String toString() { return "SubstraitPlan{" + "plan=" + Arrays.toString(plan) + ", version='" + version + '\'' + '}'; } } } FlightSqlColumnMetadata.java000066400000000000000000000223421515451615200356020ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql; import java.util.Collections; import java.util.HashMap; import java.util.Map; /** * Metadata for a column in a Flight SQL query. * *

    This can be used with FlightSqlClient to access column's metadata contained in schemas * returned by GetTables and query execution as follows: * * 

     *   FlightSqlColumnMetadata metadata = new FlightSqlColumnMetadata(field.getMetadata());
 *   Integer precision = metadata.getPrecision();
 *
    * * FlightSqlProducer can use this to set metadata on a column in a schema as follows: * * 
     *   FlightSqlColumnMetadata metadata = new FlightSqlColumnMetadata.Builder()
 *         .precision(10)
 *         .scale(5)
 *         .build();
 *   Field field = new Field("column", new FieldType(..., metadata.getMetadataMap()), null);
 *
    */ public class FlightSqlColumnMetadata { private static final String CATALOG_NAME = "ARROW:FLIGHT:SQL:CATALOG_NAME"; private static final String SCHEMA_NAME = "ARROW:FLIGHT:SQL:SCHEMA_NAME"; private static final String TABLE_NAME = "ARROW:FLIGHT:SQL:TABLE_NAME"; private static final String TYPE_NAME = "ARROW:FLIGHT:SQL:TYPE_NAME"; private static final String PRECISION = "ARROW:FLIGHT:SQL:PRECISION"; private static final String SCALE = "ARROW:FLIGHT:SQL:SCALE"; private static final String IS_AUTO_INCREMENT = "ARROW:FLIGHT:SQL:IS_AUTO_INCREMENT"; 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"; private final Map metadataMap; /** Creates a new instance of FlightSqlColumnMetadata. */ public FlightSqlColumnMetadata(Map metadataMap) { this.metadataMap = new HashMap<>(metadataMap); } /** * Returns the metadata map. * * @return The metadata map. */ public Map getMetadataMap() { return Collections.unmodifiableMap(metadataMap); } /** * Returns the catalog name. * * @return The catalog name. */ public String getCatalogName() { return metadataMap.get(CATALOG_NAME); } /** * Returns the schema name. * * @return The schema name. */ public String getSchemaName() { return metadataMap.get(SCHEMA_NAME); } /** * Returns the table name. * * @return The table name. */ public String getTableName() { return metadataMap.get(TABLE_NAME); } /** * Returns the type name. * * @return The type name. */ public String getTypeName() { return metadataMap.get(TYPE_NAME); } /** * Returns the precision / column size. * * @return The precision / column size. */ public Integer getPrecision() { String value = metadataMap.get(PRECISION); if (value == null) { return null; } return Integer.valueOf(value); } /** * Returns the scale / decimal digits. * * @return The scale / decimal digits. */ public Integer getScale() { String value = metadataMap.get(SCALE); if (value == null) { return null; } return Integer.valueOf(value); } /** * Returns if the column is auto incremented. * * @return True if the column is auto incremented, false otherwise. */ public Boolean isAutoIncrement() { String value = metadataMap.get(IS_AUTO_INCREMENT); if (value == null) { return null; } return stringToBoolean(value); } /** * Returns if the column is case-sensitive. * * @return True if the column is case-sensitive, false otherwise. */ public Boolean isCaseSensitive() { String value = metadataMap.get(IS_CASE_SENSITIVE); if (value == null) { return null; } return stringToBoolean(value); } /** * Returns if the column is read only. * * @return True if the column is read only, false otherwise. */ public Boolean isReadOnly() { String value = metadataMap.get(IS_READ_ONLY); if (value == null) { return null; } return stringToBoolean(value); } /** * Returns if the column is searchable. * * @return True if the column is searchable, false otherwise. */ public Boolean isSearchable() { String value = metadataMap.get(IS_SEARCHABLE); if (value == null) { return null; } 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; /** Creates a new instance of FlightSqlColumnMetadata.Builder. */ public Builder() { this.metadataMap = new HashMap<>(); } /** * Sets the catalog name. * * @param catalogName the catalog name. * @return This builder. */ public Builder catalogName(String catalogName) { metadataMap.put(CATALOG_NAME, catalogName); return this; } /** * Sets the schema name. * * @param schemaName The schema name. * @return This builder. */ public Builder schemaName(String schemaName) { metadataMap.put(SCHEMA_NAME, schemaName); return this; } /** * Sets the table name. * * @param tableName The table name. * @return This builder. */ public Builder tableName(String tableName) { metadataMap.put(TABLE_NAME, tableName); return this; } /** * Sets the type name. * * @param typeName The type name. * @return This builder. */ public Builder typeName(String typeName) { metadataMap.put(TYPE_NAME, typeName); return this; } /** * Sets the precision / column size. * * @param precision The precision / column size. * @return This builder. */ public Builder precision(int precision) { metadataMap.put(PRECISION, Integer.toString(precision)); return this; } /** * Sets the scale / decimal digits. * * @param scale The scale / decimal digits. * @return This builder. */ public Builder scale(int scale) { metadataMap.put(SCALE, Integer.toString(scale)); return this; } /** * Sets if the column is auto incremented. * * @param isAutoIncrement True if the column is auto incremented. * @return This builder. */ public Builder isAutoIncrement(boolean isAutoIncrement) { metadataMap.put(IS_AUTO_INCREMENT, booleanToString(isAutoIncrement)); return this; } /** * Sets if the column is case-sensitive. * * @param isCaseSensitive If the column is case-sensitive. * @return This builder. */ public Builder isCaseSensitive(boolean isCaseSensitive) { metadataMap.put(IS_CASE_SENSITIVE, booleanToString(isCaseSensitive)); return this; } /** * Sets if the column is read only. * * @param isReadOnly If the column is read only. * @return This builder. */ public Builder isReadOnly(boolean isReadOnly) { metadataMap.put(IS_READ_ONLY, booleanToString(isReadOnly)); return this; } /** * Sets if the column is searchable. * * @param isSearchable If the column is searchable. * @return This builder. */ public Builder isSearchable(boolean isSearchable) { metadataMap.put(IS_SEARCHABLE, booleanToString(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. * * @return A new instance of FlightSqlColumnMetadata. */ public FlightSqlColumnMetadata build() { return new FlightSqlColumnMetadata(metadataMap); } } private static String booleanToString(boolean boolValue) { return boolValue ? BOOLEAN_TRUE_STR : BOOLEAN_FALSE_STR; } private static boolean stringToBoolean(String value) { return value.equals(BOOLEAN_TRUE_STR); } } FlightSqlProducer.java000066400000000000000000001453231515451615200344740ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql; import static java.util.Arrays.asList; import static java.util.Collections.singletonList; import static java.util.stream.IntStream.range; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionBeginSavepointRequest; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionBeginSavepointResult; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionBeginTransactionRequest; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionBeginTransactionResult; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionCancelQueryRequest; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementResult; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedSubstraitPlanRequest; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionEndSavepointRequest; import static org.apache.arrow.flight.sql.impl.FlightSql.ActionEndTransactionRequest; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCrossReference; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetDbSchemas; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetExportedKeys; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetImportedKeys; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetXdbcTypeInfo; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementSubstraitPlan; import static org.apache.arrow.vector.complex.MapVector.DATA_VECTOR_NAME; import static org.apache.arrow.vector.complex.MapVector.KEY_NAME; import static org.apache.arrow.vector.complex.MapVector.VALUE_NAME; import static org.apache.arrow.vector.types.Types.MinorType.BIGINT; import static org.apache.arrow.vector.types.Types.MinorType.BIT; import static org.apache.arrow.vector.types.Types.MinorType.INT; import static org.apache.arrow.vector.types.Types.MinorType.LIST; import static org.apache.arrow.vector.types.Types.MinorType.STRUCT; import static org.apache.arrow.vector.types.Types.MinorType.UINT4; import static org.apache.arrow.vector.types.Types.MinorType.VARCHAR; import com.google.common.collect.ImmutableList; import com.google.protobuf.Any; import com.google.protobuf.InvalidProtocolBufferException; import java.io.IOException; import java.net.URISyntaxException; import java.nio.ByteBuffer; import java.util.List; import org.apache.arrow.flight.Action; import org.apache.arrow.flight.ActionType; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.CancelFlightInfoRequest; import org.apache.arrow.flight.CancelStatus; import org.apache.arrow.flight.CloseSessionRequest; import org.apache.arrow.flight.CloseSessionResult; import org.apache.arrow.flight.FlightConstants; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.GetSessionOptionsRequest; import org.apache.arrow.flight.GetSessionOptionsResult; import org.apache.arrow.flight.PutResult; import org.apache.arrow.flight.RenewFlightEndpointRequest; import org.apache.arrow.flight.Result; import org.apache.arrow.flight.SchemaResult; import org.apache.arrow.flight.SetSessionOptionsRequest; import org.apache.arrow.flight.SetSessionOptionsResult; import org.apache.arrow.flight.Ticket; import org.apache.arrow.flight.sql.impl.FlightSql.ActionClosePreparedStatementRequest; import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementRequest; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCatalogs; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetPrimaryKeys; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetSqlInfo; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTableTypes; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTables; import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementQuery; import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementUpdate; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementQuery; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementUpdate; import org.apache.arrow.flight.sql.impl.FlightSql.DoPutUpdateResult; import org.apache.arrow.flight.sql.impl.FlightSql.TicketStatementQuery; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.UnionMode; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.ArrowType.Union; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; /** API to Implement an Arrow Flight SQL producer. */ public interface FlightSqlProducer extends FlightProducer, AutoCloseable { /** * Depending on the provided command, method either: 1. Return information about a SQL query, or * 2. Return information about a prepared statement. In this case, parameters binding is allowed. * * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return information about the given SQL query, or the given prepared statement. */ @Override default FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) { final Any command = FlightSqlUtils.parseOrThrow(descriptor.getCommand()); if (command.is(CommandStatementQuery.class)) { return getFlightInfoStatement( FlightSqlUtils.unpackOrThrow(command, CommandStatementQuery.class), context, descriptor); } else if (command.is(CommandStatementSubstraitPlan.class)) { return getFlightInfoSubstraitPlan( FlightSqlUtils.unpackOrThrow(command, CommandStatementSubstraitPlan.class), context, descriptor); } else if (command.is(CommandPreparedStatementQuery.class)) { return getFlightInfoPreparedStatement( FlightSqlUtils.unpackOrThrow(command, CommandPreparedStatementQuery.class), context, descriptor); } else if (command.is(CommandGetCatalogs.class)) { return getFlightInfoCatalogs( FlightSqlUtils.unpackOrThrow(command, CommandGetCatalogs.class), context, descriptor); } else if (command.is(CommandGetDbSchemas.class)) { return getFlightInfoSchemas( FlightSqlUtils.unpackOrThrow(command, CommandGetDbSchemas.class), context, descriptor); } else if (command.is(CommandGetTables.class)) { return getFlightInfoTables( FlightSqlUtils.unpackOrThrow(command, CommandGetTables.class), context, descriptor); } else if (command.is(CommandGetTableTypes.class)) { return getFlightInfoTableTypes( FlightSqlUtils.unpackOrThrow(command, CommandGetTableTypes.class), context, descriptor); } else if (command.is(CommandGetSqlInfo.class)) { return getFlightInfoSqlInfo( FlightSqlUtils.unpackOrThrow(command, CommandGetSqlInfo.class), context, descriptor); } else if (command.is(CommandGetPrimaryKeys.class)) { return getFlightInfoPrimaryKeys( FlightSqlUtils.unpackOrThrow(command, CommandGetPrimaryKeys.class), context, descriptor); } else if (command.is(CommandGetExportedKeys.class)) { return getFlightInfoExportedKeys( FlightSqlUtils.unpackOrThrow(command, CommandGetExportedKeys.class), context, descriptor); } else if (command.is(CommandGetImportedKeys.class)) { return getFlightInfoImportedKeys( FlightSqlUtils.unpackOrThrow(command, CommandGetImportedKeys.class), context, descriptor); } else if (command.is(CommandGetCrossReference.class)) { return getFlightInfoCrossReference( FlightSqlUtils.unpackOrThrow(command, CommandGetCrossReference.class), context, descriptor); } else if (command.is(CommandGetXdbcTypeInfo.class)) { return getFlightInfoTypeInfo( FlightSqlUtils.unpackOrThrow(command, CommandGetXdbcTypeInfo.class), context, descriptor); } throw CallStatus.INVALID_ARGUMENT .withDescription("Unrecognized request: " + command.getTypeUrl()) .toRuntimeException(); } /** * Returns the schema of the result produced by the SQL query. * * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return the result set schema. */ @Override default SchemaResult getSchema(CallContext context, FlightDescriptor descriptor) { final Any command = FlightSqlUtils.parseOrThrow(descriptor.getCommand()); if (command.is(CommandStatementQuery.class)) { return getSchemaStatement( FlightSqlUtils.unpackOrThrow(command, CommandStatementQuery.class), context, descriptor); } else if (command.is(CommandPreparedStatementQuery.class)) { return getSchemaPreparedStatement( FlightSqlUtils.unpackOrThrow(command, CommandPreparedStatementQuery.class), context, descriptor); } else if (command.is(CommandStatementSubstraitPlan.class)) { return getSchemaSubstraitPlan( FlightSqlUtils.unpackOrThrow(command, CommandStatementSubstraitPlan.class), context, descriptor); } else if (command.is(CommandGetCatalogs.class)) { return new SchemaResult(Schemas.GET_CATALOGS_SCHEMA); } else if (command.is(CommandGetCrossReference.class)) { return new SchemaResult(Schemas.GET_CROSS_REFERENCE_SCHEMA); } else if (command.is(CommandGetDbSchemas.class)) { return new SchemaResult(Schemas.GET_SCHEMAS_SCHEMA); } else if (command.is(CommandGetExportedKeys.class)) { return new SchemaResult(Schemas.GET_EXPORTED_KEYS_SCHEMA); } else if (command.is(CommandGetImportedKeys.class)) { return new SchemaResult(Schemas.GET_IMPORTED_KEYS_SCHEMA); } else if (command.is(CommandGetPrimaryKeys.class)) { return new SchemaResult(Schemas.GET_PRIMARY_KEYS_SCHEMA); } else if (command.is(CommandGetTables.class)) { if (FlightSqlUtils.unpackOrThrow(command, CommandGetTables.class).getIncludeSchema()) { return new SchemaResult(Schemas.GET_TABLES_SCHEMA); } return new SchemaResult(Schemas.GET_TABLES_SCHEMA_NO_SCHEMA); } else if (command.is(CommandGetTableTypes.class)) { return new SchemaResult(Schemas.GET_TABLE_TYPES_SCHEMA); } else if (command.is(CommandGetSqlInfo.class)) { return new SchemaResult(Schemas.GET_SQL_INFO_SCHEMA); } else if (command.is(CommandGetXdbcTypeInfo.class)) { return new SchemaResult(Schemas.GET_TYPE_INFO_SCHEMA); } throw CallStatus.INVALID_ARGUMENT .withDescription("Unrecognized request: " + command.getTypeUrl()) .toRuntimeException(); } /** * Depending on the provided command, method either: 1. Return data for a stream produced by * executing the provided SQL query, or 2. Return data for a prepared statement. In this case, * parameters binding is allowed. * * @param context Per-call context. * @param ticket The application-defined ticket identifying this stream. * @param listener An interface for sending data back to the client. */ @Override default void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { final Any command; try { command = Any.parseFrom(ticket.getBytes()); } catch (InvalidProtocolBufferException e) { listener.error(e); return; } if (command.is(TicketStatementQuery.class)) { getStreamStatement( FlightSqlUtils.unpackOrThrow(command, TicketStatementQuery.class), context, listener); } else if (command.is(CommandPreparedStatementQuery.class)) { getStreamPreparedStatement( FlightSqlUtils.unpackOrThrow(command, CommandPreparedStatementQuery.class), context, listener); } else if (command.is(CommandGetCatalogs.class)) { getStreamCatalogs(context, listener); } else if (command.is(CommandGetDbSchemas.class)) { getStreamSchemas( FlightSqlUtils.unpackOrThrow(command, CommandGetDbSchemas.class), context, listener); } else if (command.is(CommandGetTables.class)) { getStreamTables( FlightSqlUtils.unpackOrThrow(command, CommandGetTables.class), context, listener); } else if (command.is(CommandGetTableTypes.class)) { getStreamTableTypes(context, listener); } else if (command.is(CommandGetSqlInfo.class)) { getStreamSqlInfo( FlightSqlUtils.unpackOrThrow(command, CommandGetSqlInfo.class), context, listener); } else if (command.is(CommandGetPrimaryKeys.class)) { getStreamPrimaryKeys( FlightSqlUtils.unpackOrThrow(command, CommandGetPrimaryKeys.class), context, listener); } else if (command.is(CommandGetExportedKeys.class)) { getStreamExportedKeys( FlightSqlUtils.unpackOrThrow(command, CommandGetExportedKeys.class), context, listener); } else if (command.is(CommandGetImportedKeys.class)) { getStreamImportedKeys( FlightSqlUtils.unpackOrThrow(command, CommandGetImportedKeys.class), context, listener); } else if (command.is(CommandGetCrossReference.class)) { getStreamCrossReference( FlightSqlUtils.unpackOrThrow(command, CommandGetCrossReference.class), context, listener); } else if (command.is(CommandGetXdbcTypeInfo.class)) { getStreamTypeInfo( FlightSqlUtils.unpackOrThrow(command, CommandGetXdbcTypeInfo.class), context, listener); } else { throw CallStatus.INVALID_ARGUMENT .withDescription("The defined request is invalid.") .toRuntimeException(); } } /** * Depending on the provided command, method either: 1. Execute provided SQL query as an update * statement, or 2. Execute provided update SQL query prepared statement. In this case, parameters * binding is allowed, or 3. Binds parameters to the provided prepared statement, or 4. Bulk * ingests data provided through the flightStream. * * @param context Per-call context. * @param flightStream The data stream being uploaded. * @param ackStream The data stream listener for update result acknowledgement. * @return a Runnable to process the stream. */ @Override default Runnable acceptPut( CallContext context, FlightStream flightStream, StreamListener ackStream) { final Any command = FlightSqlUtils.parseOrThrow(flightStream.getDescriptor().getCommand()); if (command.is(CommandStatementUpdate.class)) { return acceptPutStatement( FlightSqlUtils.unpackOrThrow(command, CommandStatementUpdate.class), context, flightStream, ackStream); } else if (command.is(CommandStatementIngest.class)) { return acceptPutStatementBulkIngest( FlightSqlUtils.unpackOrThrow(command, CommandStatementIngest.class), context, flightStream, ackStream); } else if (command.is(CommandStatementSubstraitPlan.class)) { return acceptPutSubstraitPlan( FlightSqlUtils.unpackOrThrow(command, CommandStatementSubstraitPlan.class), context, flightStream, ackStream); } else if (command.is(CommandPreparedStatementUpdate.class)) { return acceptPutPreparedStatementUpdate( FlightSqlUtils.unpackOrThrow(command, CommandPreparedStatementUpdate.class), context, flightStream, ackStream); } else if (command.is(CommandPreparedStatementQuery.class)) { return acceptPutPreparedStatementQuery( FlightSqlUtils.unpackOrThrow(command, CommandPreparedStatementQuery.class), context, flightStream, ackStream); } throw CallStatus.INVALID_ARGUMENT .withDescription("The defined request is invalid.") .toRuntimeException(); } /** * Lists all available Flight SQL actions. * * @param context Per-call context. * @param listener An interface for sending data back to the client. */ @Override default void listActions(CallContext context, StreamListener listener) { FlightSqlUtils.FLIGHT_SQL_ACTIONS.forEach(listener::onNext); listener.onCompleted(); } /** * Performs the requested Flight SQL action. * * @param context Per-call context. * @param action Client-supplied parameters. * @param listener A stream of responses. */ @Override default void doAction(CallContext context, Action action, StreamListener listener) { final String actionType = action.getType(); if (actionType.equals(FlightSqlUtils.FLIGHT_SQL_BEGIN_SAVEPOINT.getType())) { final ActionBeginSavepointRequest request = FlightSqlUtils.unpackAndParseOrThrow(action.getBody(), ActionBeginSavepointRequest.class); beginSavepoint(request, context, new ProtoListener<>(listener)); } else if (actionType.equals(FlightSqlUtils.FLIGHT_SQL_BEGIN_TRANSACTION.getType())) { final ActionBeginTransactionRequest request = FlightSqlUtils.unpackAndParseOrThrow( action.getBody(), ActionBeginTransactionRequest.class); beginTransaction(request, context, new ProtoListener<>(listener)); } else if (actionType.equals(FlightSqlUtils.FLIGHT_SQL_CANCEL_QUERY.getType())) { //noinspection deprecation final ActionCancelQueryRequest request = FlightSqlUtils.unpackAndParseOrThrow(action.getBody(), ActionCancelQueryRequest.class); final FlightInfo info; try { info = FlightInfo.deserialize(request.getInfo().asReadOnlyByteBuffer()); } catch (IOException | URISyntaxException e) { listener.onError( CallStatus.INTERNAL .withDescription("Could not unpack FlightInfo: " + e) .withCause(e) .toRuntimeException()); return; } cancelQuery(info, context, new CancelListener(listener)); } else if (actionType.equals(FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_STATEMENT.getType())) { final ActionCreatePreparedStatementRequest request = FlightSqlUtils.unpackAndParseOrThrow( action.getBody(), ActionCreatePreparedStatementRequest.class); createPreparedStatement(request, context, listener); } else if (actionType.equals( FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_SUBSTRAIT_PLAN.getType())) { final ActionCreatePreparedSubstraitPlanRequest request = FlightSqlUtils.unpackAndParseOrThrow( action.getBody(), ActionCreatePreparedSubstraitPlanRequest.class); createPreparedSubstraitPlan(request, context, new ProtoListener<>(listener)); } else if (actionType.equals(FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT.getType())) { final ActionClosePreparedStatementRequest request = FlightSqlUtils.unpackAndParseOrThrow( action.getBody(), ActionClosePreparedStatementRequest.class); closePreparedStatement(request, context, new NoResultListener(listener)); } else if (actionType.equals(FlightSqlUtils.FLIGHT_SQL_END_SAVEPOINT.getType())) { ActionEndSavepointRequest request = FlightSqlUtils.unpackAndParseOrThrow(action.getBody(), ActionEndSavepointRequest.class); endSavepoint(request, context, new NoResultListener(listener)); } else if (actionType.equals(FlightSqlUtils.FLIGHT_SQL_END_TRANSACTION.getType())) { ActionEndTransactionRequest request = FlightSqlUtils.unpackAndParseOrThrow(action.getBody(), ActionEndTransactionRequest.class); endTransaction(request, context, new NoResultListener(listener)); } else if (actionType.equals(FlightConstants.CANCEL_FLIGHT_INFO.getType())) { final CancelFlightInfoRequest request; try { request = CancelFlightInfoRequest.deserialize(ByteBuffer.wrap(action.getBody())); } catch (IOException | URISyntaxException e) { listener.onError( CallStatus.INTERNAL .withDescription("Could not unpack FlightInfo: " + e) .withCause(e) .toRuntimeException()); return; } cancelFlightInfo(request, context, new CancelStatusListener(listener)); } else if (actionType.equals(FlightConstants.RENEW_FLIGHT_ENDPOINT.getType())) { final RenewFlightEndpointRequest request; try { request = RenewFlightEndpointRequest.deserialize(ByteBuffer.wrap(action.getBody())); } catch (IOException | URISyntaxException e) { listener.onError( CallStatus.INTERNAL .withDescription("Could not unpack FlightInfo: " + e) .withCause(e) .toRuntimeException()); return; } renewFlightEndpoint(request, context, new FlightEndpointListener(listener)); } else if (actionType.equals(FlightConstants.SET_SESSION_OPTIONS.getType())) { final SetSessionOptionsRequest request; try { request = SetSessionOptionsRequest.deserialize(ByteBuffer.wrap(action.getBody())); } catch (IOException e) { listener.onError( CallStatus.INTERNAL .withDescription("Could not unpack SetSessionOptionsRequest: " + e) .withCause(e) .toRuntimeException()); return; } setSessionOptions(request, context, new SetSessionOptionsResultListener(listener)); } else if (actionType.equals(FlightConstants.GET_SESSION_OPTIONS.getType())) { final GetSessionOptionsRequest request; try { request = GetSessionOptionsRequest.deserialize(ByteBuffer.wrap(action.getBody())); } catch (IOException e) { listener.onError( CallStatus.INTERNAL .withDescription("Could not unpack GetSessionOptionsRequest: " + e) .withCause(e) .toRuntimeException()); return; } getSessionOptions(request, context, new GetSessionOptionsResultListener(listener)); } else if (actionType.equals(FlightConstants.CLOSE_SESSION.getType())) { final CloseSessionRequest request; try { request = CloseSessionRequest.deserialize(ByteBuffer.wrap(action.getBody())); } catch (IOException e) { listener.onError( CallStatus.INTERNAL .withDescription("Could not unpack CloseSessionRequest: " + e) .withCause(e) .toRuntimeException()); return; } closeSession(request, context, new CloseSessionResultListener(listener)); } else { throw CallStatus.INVALID_ARGUMENT .withDescription("Unrecognized request: " + action.getType()) .toRuntimeException(); } } /** * Create a savepoint within a transaction. * * @param request The savepoint request. * @param context Per-call context. * @param listener The newly created savepoint ID. */ default void beginSavepoint( ActionBeginSavepointRequest request, CallContext context, StreamListener listener) { listener.onError(CallStatus.UNIMPLEMENTED.toRuntimeException()); } /** * Begin a transaction. * * @param request The transaction request. * @param context Per-call context. * @param listener The newly created transaction ID. */ default void beginTransaction( ActionBeginTransactionRequest request, CallContext context, StreamListener listener) { listener.onError(CallStatus.UNIMPLEMENTED.toRuntimeException()); } /** * Explicitly cancel a query. * * @param request The CancelFlightInfoRequest for the query to cancel. * @param context Per-call context. * @param listener An interface for sending data back to the client. */ default void cancelFlightInfo( CancelFlightInfoRequest request, CallContext context, StreamListener listener) { listener.onError(CallStatus.UNIMPLEMENTED.toRuntimeException()); } /** * Explicitly cancel a query. * * @param info The FlightInfo of the query to cancel. * @param context Per-call context. * @param listener Whether cancellation succeeded. * @deprecated Prefer {@link #cancelFlightInfo(CancelFlightInfoRequest, CallContext, * StreamListener)}. */ @Deprecated default void cancelQuery( FlightInfo info, CallContext context, StreamListener listener) { CancelFlightInfoRequest request = new CancelFlightInfoRequest(info); cancelFlightInfo( request, context, new StreamListener() { @Override public void onNext(CancelStatus val) { switch (val) { case UNSPECIFIED: listener.onNext(CancelResult.UNSPECIFIED); break; case CANCELLED: listener.onNext(CancelResult.CANCELLED); break; case CANCELLING: listener.onNext(CancelResult.CANCELLING); break; case NOT_CANCELLABLE: listener.onNext(CancelResult.NOT_CANCELLABLE); break; default: // XXX: CheckStyle requires a default clause which arguably makes the code worse. throw new AssertionError("Unknown enum variant " + val); } } @Override public void onError(Throwable t) { listener.onError(t); } @Override public void onCompleted() { listener.onCompleted(); } }); } /** * Set server session options(s). * * @param request The session options to set. For *DBC driver compatibility, servers should * support converting values from strings. * @param context Per-call context. * @param listener An interface for sending data back to the client. */ default void setSessionOptions( SetSessionOptionsRequest request, CallContext context, StreamListener listener) { listener.onError(CallStatus.UNIMPLEMENTED.toRuntimeException()); } /** * Get server session option(s). * * @param request The (empty) GetSessionOptionsRequest. * @param context Per-call context. * @param listener An interface for sending data back to the client. */ default void getSessionOptions( GetSessionOptionsRequest request, CallContext context, StreamListener listener) { listener.onError(CallStatus.UNIMPLEMENTED.toRuntimeException()); } /** * Close/invalidate the session. * * @param request The (empty) CloseSessionRequest. * @param context Per-call context. * @param listener An interface for sending data back to the client. */ default void closeSession( CloseSessionRequest request, CallContext context, StreamListener listener) { listener.onError(CallStatus.UNIMPLEMENTED.toRuntimeException()); } /** * Creates a prepared statement on the server and returns a handle and metadata for in a {@link * ActionCreatePreparedStatementResult} object in a {@link Result} object. * * @param request The sql command to generate the prepared statement. * @param context Per-call context. * @param listener A stream of responses. */ void createPreparedStatement( ActionCreatePreparedStatementRequest request, CallContext context, StreamListener listener); /** * Pre-compile a Substrait plan. * * @param request The plan. * @param context Per-call context. * @param listener The resulting prepared statement. */ default void createPreparedSubstraitPlan( ActionCreatePreparedSubstraitPlanRequest request, CallContext context, StreamListener listener) { listener.onError(CallStatus.UNIMPLEMENTED.toRuntimeException()); } /** * Closes a prepared statement on the server. No result is expected. * * @param request The sql command to generate the prepared statement. * @param context Per-call context. * @param listener A stream of responses. */ void closePreparedStatement( ActionClosePreparedStatementRequest request, CallContext context, StreamListener listener); /** * Release or roll back to a savepoint. * * @param request The savepoint, and whether to release/rollback. * @param context Per-call context. * @param listener Call {@link StreamListener#onCompleted()} or {@link * StreamListener#onError(Throwable)} when done; do not send a result. */ default void endSavepoint( ActionEndSavepointRequest request, CallContext context, StreamListener listener) { listener.onError(CallStatus.UNIMPLEMENTED.toRuntimeException()); } /** * Commit or roll back to a transaction. * * @param request The transaction, and whether to release/rollback. * @param context Per-call context. * @param listener Call {@link StreamListener#onCompleted()} or {@link * StreamListener#onError(Throwable)} when done; do not send a result. */ default void endTransaction( ActionEndTransactionRequest request, CallContext context, StreamListener listener) { listener.onError(CallStatus.UNIMPLEMENTED.toRuntimeException()); } /** * Evaluate a SQL query. * * @param command The SQL query. * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return Metadata about the stream. */ FlightInfo getFlightInfoStatement( CommandStatementQuery command, CallContext context, FlightDescriptor descriptor); /** * Evaluate a Substrait plan. * * @param command The Substrait plan. * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return Metadata about the stream. */ default FlightInfo getFlightInfoSubstraitPlan( CommandStatementSubstraitPlan command, CallContext context, FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED.toRuntimeException(); } /** * Gets information about a particular prepared statement data stream. * * @param command The prepared statement to generate the data stream. * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return Metadata about the stream. */ FlightInfo getFlightInfoPreparedStatement( CommandPreparedStatementQuery command, CallContext context, FlightDescriptor descriptor); /** * Get the result schema for a SQL query. * * @param command The SQL query. * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return the schema of the result set. */ SchemaResult getSchemaStatement( CommandStatementQuery command, CallContext context, FlightDescriptor descriptor); /** * Get the schema of the result set of a prepared statement. * * @param command The prepared statement handle. * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return the schema of the result set. */ default SchemaResult getSchemaPreparedStatement( CommandPreparedStatementQuery command, CallContext context, FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED .withDescription("GetSchema with CommandPreparedStatementQuery is not implemented") .toRuntimeException(); } /** * Get the result schema for a Substrait plan. * * @param command The Substrait plan. * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return Schema for the stream. */ default SchemaResult getSchemaSubstraitPlan( CommandStatementSubstraitPlan command, CallContext context, FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED.toRuntimeException(); } /** * Returns data for a SQL query based data stream. * * @param ticket Ticket message containing the statement handle. * @param context Per-call context. * @param listener An interface for sending data back to the client. */ void getStreamStatement( TicketStatementQuery ticket, CallContext context, ServerStreamListener listener); /** * Returns data for a particular prepared statement query instance. * * @param command The prepared statement to generate the data stream. * @param context Per-call context. * @param listener An interface for sending data back to the client. */ void getStreamPreparedStatement( CommandPreparedStatementQuery command, CallContext context, ServerStreamListener listener); /** * Accepts uploaded data for a particular SQL query based data stream. * *

    `PutResult`s must be in the form of a {@link DoPutUpdateResult}. * * @param command The sql command to generate the data stream. * @param context Per-call context. * @param flightStream The data stream being uploaded. * @param ackStream The result data stream. * @return A runnable to process the stream. */ Runnable acceptPutStatement( CommandStatementUpdate command, CallContext context, FlightStream flightStream, StreamListener ackStream); /** * Accepts uploaded data for a particular bulk ingest data stream. * *

    `PutResult`s must be in the form of a {@link DoPutUpdateResult}. * * @param command The bulk ingestion request. * @param context Per-call context. * @param flightStream The data stream being uploaded. * @param ackStream The result data stream. * @return A runnable to process the stream. */ default Runnable acceptPutStatementBulkIngest( CommandStatementIngest command, CallContext context, FlightStream flightStream, StreamListener ackStream) { return () -> { ackStream.onError(CallStatus.UNIMPLEMENTED.toRuntimeException()); }; } /** * Handle a Substrait plan with uploaded data. * * @param command The Substrait plan to evaluate. * @param context Per-call context. * @param flightStream The data stream being uploaded. * @param ackStream The result data stream. * @return A runnable to process the stream. */ default Runnable acceptPutSubstraitPlan( CommandStatementSubstraitPlan command, CallContext context, FlightStream flightStream, StreamListener ackStream) { return () -> { ackStream.onError(CallStatus.UNIMPLEMENTED.toRuntimeException()); }; } /** * Accepts uploaded data for a particular prepared statement data stream. * *

    `PutResult`s must be in the form of a {@link DoPutUpdateResult}. * * @param command The prepared statement to generate the data stream. * @param context Per-call context. * @param flightStream The data stream being uploaded. * @param ackStream The result data stream. * @return A runnable to process the stream. */ Runnable acceptPutPreparedStatementUpdate( CommandPreparedStatementUpdate command, CallContext context, FlightStream flightStream, StreamListener ackStream); /** * Accepts uploaded parameter values for a particular prepared statement query. * * @param command The prepared statement the parameter values will bind to. * @param context Per-call context. * @param flightStream The data stream being uploaded. * @param ackStream The result data stream. * @return A runnable to process the stream. */ Runnable acceptPutPreparedStatementQuery( CommandPreparedStatementQuery command, CallContext context, FlightStream flightStream, StreamListener ackStream); /** * Returns the SQL Info of the server by returning a {@link CommandGetSqlInfo} in a {@link * Result}. * * @param request request filter parameters. * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return Metadata about the stream. */ FlightInfo getFlightInfoSqlInfo( CommandGetSqlInfo request, CallContext context, FlightDescriptor descriptor); /** * Returns data for SQL info based data stream. * * @param command The command to generate the data stream. * @param context Per-call context. * @param listener An interface for sending data back to the client. */ void getStreamSqlInfo( CommandGetSqlInfo command, CallContext context, ServerStreamListener listener); /** * Returns a description of all the data types supported by source. * * @param request request filter parameters. * @param descriptor The descriptor identifying the data stream. * @return Metadata about the stream. */ FlightInfo getFlightInfoTypeInfo( CommandGetXdbcTypeInfo request, CallContext context, FlightDescriptor descriptor); /** * Returns data for type info based data stream. * * @param context Per-call context. * @param listener An interface for sending data back to the client. */ void getStreamTypeInfo( CommandGetXdbcTypeInfo request, CallContext context, ServerStreamListener listener); /** * Returns the available catalogs by returning a stream of {@link CommandGetCatalogs} objects in * {@link Result} objects. * * @param request request filter parameters. * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return Metadata about the stream. */ FlightInfo getFlightInfoCatalogs( CommandGetCatalogs request, CallContext context, FlightDescriptor descriptor); /** * Returns data for catalogs based data stream. * * @param context Per-call context. * @param listener An interface for sending data back to the client. */ void getStreamCatalogs(CallContext context, ServerStreamListener listener); /** * Returns the available schemas by returning a stream of {@link CommandGetDbSchemas} objects in * {@link Result} objects. * * @param request request filter parameters. * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return Metadata about the stream. */ FlightInfo getFlightInfoSchemas( CommandGetDbSchemas request, CallContext context, FlightDescriptor descriptor); /** * Returns data for schemas based data stream. * * @param command The command to generate the data stream. * @param context Per-call context. * @param listener An interface for sending data back to the client. */ void getStreamSchemas( CommandGetDbSchemas command, CallContext context, ServerStreamListener listener); /** * Returns the available tables by returning a stream of {@link CommandGetTables} objects in * {@link Result} objects. * * @param request request filter parameters. * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return Metadata about the stream. */ FlightInfo getFlightInfoTables( CommandGetTables request, CallContext context, FlightDescriptor descriptor); /** * Returns data for tables based data stream. * * @param command The command to generate the data stream. * @param context Per-call context. * @param listener An interface for sending data back to the client. */ void getStreamTables( CommandGetTables command, CallContext context, ServerStreamListener listener); /** * Returns the available table types by returning a stream of {@link CommandGetTableTypes} objects * in {@link Result} objects. * * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return Metadata about the stream. */ FlightInfo getFlightInfoTableTypes( CommandGetTableTypes request, CallContext context, FlightDescriptor descriptor); /** * Returns data for table types based data stream. * * @param context Per-call context. * @param listener An interface for sending data back to the client. */ void getStreamTableTypes(CallContext context, ServerStreamListener listener); /** * Returns the available primary keys by returning a stream of {@link CommandGetPrimaryKeys} * objects in {@link Result} objects. * * @param request request filter parameters. * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return Metadata about the stream. */ FlightInfo getFlightInfoPrimaryKeys( CommandGetPrimaryKeys request, CallContext context, FlightDescriptor descriptor); /** * Returns data for primary keys based data stream. * * @param command The command to generate the data stream. * @param context Per-call context. * @param listener An interface for sending data back to the client. */ void getStreamPrimaryKeys( CommandGetPrimaryKeys command, CallContext context, ServerStreamListener listener); /** * Retrieves a description of the foreign key columns that reference the given table's primary key * columns {@link CommandGetExportedKeys} objects in {@link Result} objects. * * @param request request filter parameters. * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return Metadata about the stream. */ FlightInfo getFlightInfoExportedKeys( CommandGetExportedKeys request, CallContext context, FlightDescriptor descriptor); /** * Retrieves a description of the primary key columns that are referenced by given table's foreign * key columns {@link CommandGetImportedKeys} objects in {@link Result} objects. * * @param request request filter parameters. * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return Metadata about the stream. */ FlightInfo getFlightInfoImportedKeys( CommandGetImportedKeys request, CallContext context, FlightDescriptor descriptor); /** * Retrieve a description of the foreign key columns that reference the given table's primary key * columns {@link CommandGetCrossReference} objects in {@link Result} objects. * * @param request request filter parameters. * @param context Per-call context. * @param descriptor The descriptor identifying the data stream. * @return Metadata about the stream. */ FlightInfo getFlightInfoCrossReference( CommandGetCrossReference request, CallContext context, FlightDescriptor descriptor); /** * Returns data for foreign keys based data stream. * * @param command The command to generate the data stream. * @param context Per-call context. * @param listener An interface for sending data back to the client. */ void getStreamExportedKeys( CommandGetExportedKeys command, CallContext context, ServerStreamListener listener); /** * Returns data for foreign keys based data stream. * * @param command The command to generate the data stream. * @param context Per-call context. * @param listener An interface for sending data back to the client. */ void getStreamImportedKeys( CommandGetImportedKeys command, CallContext context, ServerStreamListener listener); /** * Returns data for cross reference based data stream. * * @param command The command to generate the data stream. * @param context Per-call context. * @param listener An interface for sending data back to the client. */ void getStreamCrossReference( CommandGetCrossReference command, CallContext context, ServerStreamListener listener); /** * Renew the duration of the given endpoint. * * @param request The endpoint to renew. * @param context Per-call context. * @param listener An interface for sending data back to the client. */ default void renewFlightEndpoint( RenewFlightEndpointRequest request, CallContext context, StreamListener listener) { listener.onError(CallStatus.UNIMPLEMENTED.toRuntimeException()); } /** Default schema templates for the {@link FlightSqlProducer}. */ final class Schemas { public static final Schema GET_TABLES_SCHEMA = new Schema( asList( Field.nullable("catalog_name", VARCHAR.getType()), Field.nullable("db_schema_name", VARCHAR.getType()), Field.notNullable("table_name", VARCHAR.getType()), Field.notNullable("table_type", VARCHAR.getType()), Field.notNullable("table_schema", MinorType.VARBINARY.getType()))); public static final Schema GET_TABLES_SCHEMA_NO_SCHEMA = new Schema( asList( Field.nullable("catalog_name", VARCHAR.getType()), Field.nullable("db_schema_name", VARCHAR.getType()), Field.notNullable("table_name", VARCHAR.getType()), Field.notNullable("table_type", VARCHAR.getType()))); public static final Schema GET_CATALOGS_SCHEMA = new Schema(singletonList(Field.notNullable("catalog_name", VARCHAR.getType()))); public static final Schema GET_TABLE_TYPES_SCHEMA = new Schema(singletonList(Field.notNullable("table_type", VARCHAR.getType()))); public static final Schema GET_SCHEMAS_SCHEMA = new Schema( asList( Field.nullable("catalog_name", VARCHAR.getType()), Field.notNullable("db_schema_name", VARCHAR.getType()))); private static final Schema GET_IMPORTED_EXPORTED_AND_CROSS_REFERENCE_KEYS_SCHEMA = new Schema( asList( Field.nullable("pk_catalog_name", VARCHAR.getType()), Field.nullable("pk_db_schema_name", VARCHAR.getType()), Field.notNullable("pk_table_name", VARCHAR.getType()), Field.notNullable("pk_column_name", VARCHAR.getType()), Field.nullable("fk_catalog_name", VARCHAR.getType()), Field.nullable("fk_db_schema_name", VARCHAR.getType()), Field.notNullable("fk_table_name", VARCHAR.getType()), Field.notNullable("fk_column_name", VARCHAR.getType()), Field.notNullable("key_sequence", INT.getType()), Field.nullable("fk_key_name", VARCHAR.getType()), Field.nullable("pk_key_name", VARCHAR.getType()), Field.notNullable("update_rule", MinorType.UINT1.getType()), Field.notNullable("delete_rule", MinorType.UINT1.getType()))); public static final Schema GET_IMPORTED_KEYS_SCHEMA = GET_IMPORTED_EXPORTED_AND_CROSS_REFERENCE_KEYS_SCHEMA; public static final Schema GET_EXPORTED_KEYS_SCHEMA = GET_IMPORTED_EXPORTED_AND_CROSS_REFERENCE_KEYS_SCHEMA; public static final Schema GET_CROSS_REFERENCE_SCHEMA = GET_IMPORTED_EXPORTED_AND_CROSS_REFERENCE_KEYS_SCHEMA; private static final List GET_SQL_INFO_DENSE_UNION_SCHEMA_FIELDS = asList( Field.notNullable("string_value", VARCHAR.getType()), Field.notNullable("bool_value", BIT.getType()), Field.notNullable("bigint_value", BIGINT.getType()), Field.notNullable("int32_bitmask", INT.getType()), new Field( "string_list", FieldType.notNullable(LIST.getType()), singletonList(Field.nullable("item", VARCHAR.getType()))), new Field( "int32_to_int32_list_map", FieldType.notNullable(new ArrowType.Map(false)), singletonList( new Field( DATA_VECTOR_NAME, new FieldType(false, STRUCT.getType(), null), ImmutableList.of( Field.notNullable(KEY_NAME, INT.getType()), new Field( VALUE_NAME, FieldType.nullable(LIST.getType()), singletonList(Field.nullable("item", INT.getType())))))))); public static final Schema GET_SQL_INFO_SCHEMA = new Schema( asList( Field.notNullable("info_name", UINT4.getType()), new Field( "value", FieldType.notNullable( new Union( UnionMode.Dense, range(0, GET_SQL_INFO_DENSE_UNION_SCHEMA_FIELDS.size()).toArray())), GET_SQL_INFO_DENSE_UNION_SCHEMA_FIELDS))); public static final Schema GET_TYPE_INFO_SCHEMA = new Schema( asList( Field.notNullable("type_name", VARCHAR.getType()), Field.notNullable("data_type", INT.getType()), Field.nullable("column_size", INT.getType()), Field.nullable("literal_prefix", VARCHAR.getType()), Field.nullable("literal_suffix", VARCHAR.getType()), new Field( "create_params", FieldType.nullable(LIST.getType()), singletonList(Field.notNullable("item", VARCHAR.getType()))), Field.notNullable("nullable", INT.getType()), Field.notNullable("case_sensitive", BIT.getType()), Field.notNullable("searchable", INT.getType()), Field.nullable("unsigned_attribute", BIT.getType()), Field.notNullable("fixed_prec_scale", BIT.getType()), Field.nullable("auto_increment", BIT.getType()), Field.nullable("local_type_name", VARCHAR.getType()), Field.nullable("minimum_scale", INT.getType()), Field.nullable("maximum_scale", INT.getType()), Field.notNullable("sql_data_type", INT.getType()), Field.nullable("datetime_subcode", INT.getType()), Field.nullable("num_prec_radix", INT.getType()), Field.nullable("interval_precision", INT.getType()))); public static final Schema GET_PRIMARY_KEYS_SCHEMA = new Schema( asList( Field.nullable("catalog_name", VARCHAR.getType()), Field.nullable("db_schema_name", VARCHAR.getType()), Field.notNullable("table_name", VARCHAR.getType()), Field.notNullable("column_name", VARCHAR.getType()), Field.notNullable("key_sequence", INT.getType()), Field.nullable("key_name", VARCHAR.getType()))); private Schemas() { // Prevent instantiation. } } } FlightSqlUtils.java000066400000000000000000000130461515451615200340050ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql; import com.google.common.collect.ImmutableList; import com.google.protobuf.Any; import com.google.protobuf.InvalidProtocolBufferException; import com.google.protobuf.Message; import java.util.List; import org.apache.arrow.flight.ActionType; import org.apache.arrow.flight.CallStatus; /** Utilities to work with Flight SQL semantics. */ public final class FlightSqlUtils { public static final ActionType FLIGHT_SQL_BEGIN_SAVEPOINT = new ActionType( "BeginSavepoint", "Create a new savepoint.\n" + "Request Message: ActionBeginSavepointRequest\n" + "Response Message: ActionBeginSavepointResult"); public static final ActionType FLIGHT_SQL_BEGIN_TRANSACTION = new ActionType( "BeginTransaction", "Start a new transaction.\n" + "Request Message: ActionBeginTransactionRequest\n" + "Response Message: ActionBeginTransactionResult"); public static final ActionType FLIGHT_SQL_CREATE_PREPARED_STATEMENT = new ActionType( "CreatePreparedStatement", "Creates a reusable prepared statement resource on the server. \n" + "Request Message: ActionCreatePreparedStatementRequest\n" + "Response Message: ActionCreatePreparedStatementResult"); public static final ActionType FLIGHT_SQL_CLOSE_PREPARED_STATEMENT = new ActionType( "ClosePreparedStatement", "Closes a reusable prepared statement resource on the server. \n" + "Request Message: ActionClosePreparedStatementRequest\n" + "Response Message: N/A"); public static final ActionType FLIGHT_SQL_CREATE_PREPARED_SUBSTRAIT_PLAN = new ActionType( "CreatePreparedSubstraitPlan", "Creates a reusable prepared statement resource on the server.\n" + "Request Message: ActionCreatePreparedSubstraitPlanRequest\n" + "Response Message: ActionCreatePreparedStatementResult"); public static final ActionType FLIGHT_SQL_CANCEL_QUERY = new ActionType( "CancelQuery", "Explicitly cancel a running query.\n" + "Request Message: ActionCancelQueryRequest\n" + "Response Message: ActionCancelQueryResult"); public static final ActionType FLIGHT_SQL_END_SAVEPOINT = new ActionType( "EndSavepoint", "End a savepoint.\n" + "Request Message: ActionEndSavepointRequest\n" + "Response Message: N/A"); public static final ActionType FLIGHT_SQL_END_TRANSACTION = new ActionType( "EndTransaction", "End a transaction.\n" + "Request Message: ActionEndTransactionRequest\n" + "Response Message: N/A"); public static final List FLIGHT_SQL_ACTIONS = ImmutableList.of( FLIGHT_SQL_BEGIN_SAVEPOINT, FLIGHT_SQL_BEGIN_TRANSACTION, FLIGHT_SQL_CREATE_PREPARED_STATEMENT, FLIGHT_SQL_CLOSE_PREPARED_STATEMENT, FLIGHT_SQL_CREATE_PREPARED_SUBSTRAIT_PLAN, FLIGHT_SQL_CANCEL_QUERY, FLIGHT_SQL_END_SAVEPOINT, FLIGHT_SQL_END_TRANSACTION); /** * Helper to parse {@link com.google.protobuf.Any} objects to the specific protobuf object. * * @param source the raw bytes source value. * @return the materialized protobuf object. */ public static Any parseOrThrow(byte[] source) { try { return Any.parseFrom(source); } catch (final InvalidProtocolBufferException e) { throw CallStatus.INVALID_ARGUMENT .withDescription("Received invalid message from remote.") .withCause(e) .toRuntimeException(); } } /** * Helper to unpack {@link com.google.protobuf.Any} objects to the specific protobuf object. * * @param source the parsed Source value. * @param as the class to unpack as. * @param the class to unpack as. * @return the materialized protobuf object. */ public static T unpackOrThrow(Any source, Class as) { try { return source.unpack(as); } catch (final InvalidProtocolBufferException e) { throw CallStatus.INVALID_ARGUMENT .withDescription("Provided message cannot be unpacked as " + as.getName() + ": " + e) .withCause(e) .toRuntimeException(); } } /** * Helper to parse and unpack {@link com.google.protobuf.Any} objects to the specific protobuf * object. * * @param source the raw bytes source value. * @param as the class to unpack as. * @param the class to unpack as. * @return the materialized protobuf object. */ public static T unpackAndParseOrThrow(byte[] source, Class as) { return unpackOrThrow(parseOrThrow(source), as); } } GetSessionOptionsResultListener.java000066400000000000000000000031201515451615200374230ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.GetSessionOptionsResult; import org.apache.arrow.flight.Result; /** Typed StreamListener for getSessionOptions. */ public class GetSessionOptionsResultListener implements FlightProducer.StreamListener { private final FlightProducer.StreamListener listener; GetSessionOptionsResultListener(FlightProducer.StreamListener listener) { this.listener = listener; } @Override public void onNext(GetSessionOptionsResult val) { listener.onNext(new Result(val.serialize().array())); } @Override public void onError(Throwable t) { listener.onError(t); } @Override public void onCompleted() { listener.onCompleted(); } } NoOpFlightSqlProducer.java000066400000000000000000000234721515451615200352700ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.Criteria; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.PutResult; import org.apache.arrow.flight.Result; import org.apache.arrow.flight.SchemaResult; import org.apache.arrow.flight.sql.impl.FlightSql; /** A {@link FlightSqlProducer} that throws on all FlightSql-specific operations. */ public class NoOpFlightSqlProducer implements FlightSqlProducer { @Override public void createPreparedStatement( FlightSql.ActionCreatePreparedStatementRequest request, CallContext context, StreamListener listener) { listener.onError( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } @Override public void closePreparedStatement( FlightSql.ActionClosePreparedStatementRequest request, CallContext context, StreamListener listener) { listener.onError( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } @Override public FlightInfo getFlightInfoStatement( FlightSql.CommandStatementQuery command, CallContext context, FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public FlightInfo getFlightInfoPreparedStatement( FlightSql.CommandPreparedStatementQuery command, CallContext context, FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public SchemaResult getSchemaStatement( FlightSql.CommandStatementQuery command, CallContext context, FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public void getStreamStatement( FlightSql.TicketStatementQuery ticket, CallContext context, ServerStreamListener listener) { listener.error( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } @Override public void getStreamPreparedStatement( FlightSql.CommandPreparedStatementQuery command, CallContext context, ServerStreamListener listener) { listener.error( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } @Override public Runnable acceptPutStatement( FlightSql.CommandStatementUpdate command, CallContext context, FlightStream flightStream, StreamListener ackStream) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public Runnable acceptPutStatementBulkIngest( FlightSql.CommandStatementIngest command, CallContext context, FlightStream flightStream, StreamListener ackStream) { return () -> { ackStream.onError( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); }; } @Override public Runnable acceptPutPreparedStatementUpdate( FlightSql.CommandPreparedStatementUpdate command, CallContext context, FlightStream flightStream, StreamListener ackStream) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public Runnable acceptPutPreparedStatementQuery( FlightSql.CommandPreparedStatementQuery command, CallContext context, FlightStream flightStream, StreamListener ackStream) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public FlightInfo getFlightInfoSqlInfo( FlightSql.CommandGetSqlInfo request, CallContext context, FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public void getStreamSqlInfo( FlightSql.CommandGetSqlInfo command, CallContext context, ServerStreamListener listener) { listener.error( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } @Override public FlightInfo getFlightInfoTypeInfo( FlightSql.CommandGetXdbcTypeInfo request, CallContext context, FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public void getStreamTypeInfo( FlightSql.CommandGetXdbcTypeInfo request, CallContext context, ServerStreamListener listener) { listener.error( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } @Override public FlightInfo getFlightInfoCatalogs( FlightSql.CommandGetCatalogs request, CallContext context, FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public void getStreamCatalogs(CallContext context, ServerStreamListener listener) { listener.error( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } @Override public FlightInfo getFlightInfoSchemas( FlightSql.CommandGetDbSchemas request, CallContext context, FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public void getStreamSchemas( FlightSql.CommandGetDbSchemas command, CallContext context, ServerStreamListener listener) { listener.error( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } @Override public FlightInfo getFlightInfoTables( FlightSql.CommandGetTables request, CallContext context, FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public void getStreamTables( FlightSql.CommandGetTables command, CallContext context, ServerStreamListener listener) { listener.error( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } @Override public FlightInfo getFlightInfoTableTypes( FlightSql.CommandGetTableTypes request, CallContext context, FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public void getStreamTableTypes(CallContext context, ServerStreamListener listener) { listener.error( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } @Override public FlightInfo getFlightInfoPrimaryKeys( FlightSql.CommandGetPrimaryKeys request, CallContext context, FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public void getStreamPrimaryKeys( FlightSql.CommandGetPrimaryKeys command, CallContext context, ServerStreamListener listener) { listener.error( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } @Override public FlightInfo getFlightInfoExportedKeys( FlightSql.CommandGetExportedKeys request, CallContext context, FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public FlightInfo getFlightInfoImportedKeys( FlightSql.CommandGetImportedKeys request, CallContext context, FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public FlightInfo getFlightInfoCrossReference( FlightSql.CommandGetCrossReference request, CallContext context, FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public void getStreamExportedKeys( FlightSql.CommandGetExportedKeys command, CallContext context, ServerStreamListener listener) { listener.error( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } @Override public void getStreamImportedKeys( FlightSql.CommandGetImportedKeys command, CallContext context, ServerStreamListener listener) { listener.error( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } @Override public void getStreamCrossReference( FlightSql.CommandGetCrossReference command, CallContext context, ServerStreamListener listener) { listener.error( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } @Override public void close() throws Exception {} @Override public void listFlights( CallContext context, Criteria criteria, StreamListener listener) { throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } } NoResultListener.java000066400000000000000000000027651515451615200343560ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.Result; /** A StreamListener for actions that do not return results. */ class NoResultListener implements FlightProducer.StreamListener { private final FlightProducer.StreamListener listener; NoResultListener(FlightProducer.StreamListener listener) { this.listener = listener; } @Override public void onNext(Result val) { throw new UnsupportedOperationException("Do not call onNext on this listener."); } @Override public void onError(Throwable t) { listener.onError(t); } @Override public void onCompleted() { listener.onCompleted(); } } ProtoListener.java000066400000000000000000000031161515451615200336750ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql; import com.google.protobuf.Any; import com.google.protobuf.Message; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.Result; /** * A StreamListener that accepts a particular type. * * @param The message type to accept. */ class ProtoListener implements FlightProducer.StreamListener { private final FlightProducer.StreamListener listener; ProtoListener(FlightProducer.StreamListener listener) { this.listener = listener; } @Override public void onNext(T val) { listener.onNext(new Result(Any.pack(val).toByteArray())); } @Override public void onError(Throwable t) { listener.onError(t); } @Override public void onCompleted() { listener.onCompleted(); } } SetSessionOptionsResultListener.java000066400000000000000000000031201515451615200374370ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.Result; import org.apache.arrow.flight.SetSessionOptionsResult; /** Typed StreamListener for setSessionOptions. */ public class SetSessionOptionsResultListener implements FlightProducer.StreamListener { private final FlightProducer.StreamListener listener; SetSessionOptionsResultListener(FlightProducer.StreamListener listener) { this.listener = listener; } @Override public void onNext(SetSessionOptionsResult val) { listener.onNext(new Result(val.serialize().array())); } @Override public void onError(Throwable t) { listener.onError(t); } @Override public void onCompleted() { listener.onCompleted(); } } SqlInfoBuilder.java000066400000000000000000001262041515451615200337520ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql; import static java.util.stream.IntStream.range; import static org.apache.arrow.flight.FlightProducer.ServerStreamListener; import static org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedTransaction; import static org.apache.arrow.flight.sql.util.SqlInfoOptionsUtils.createBitmaskFromEnums; import com.google.protobuf.ProtocolMessageEnum; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.function.Consumer; import java.util.function.ObjIntConsumer; import org.apache.arrow.flight.sql.impl.FlightSql.SqlInfo; import org.apache.arrow.flight.sql.impl.FlightSql.SqlNullOrdering; import org.apache.arrow.flight.sql.impl.FlightSql.SqlOuterJoinsSupportLevel; import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedCaseSensitivity; import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedElementActions; import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedGroupBy; import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedPositionedCommands; import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedResultSetType; import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedSubqueries; import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedUnions; import org.apache.arrow.flight.sql.impl.FlightSql.SqlTransactionIsolationLevel; import org.apache.arrow.flight.sql.impl.FlightSql.SupportedAnsi92SqlGrammarLevel; import org.apache.arrow.flight.sql.impl.FlightSql.SupportedSqlGrammar; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.UInt4Vector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.impl.UnionListWriter; import org.apache.arrow.vector.complex.impl.UnionMapWriter; import org.apache.arrow.vector.complex.writer.BaseWriter; import org.apache.arrow.vector.holders.NullableBigIntHolder; import org.apache.arrow.vector.holders.NullableBitHolder; import org.apache.arrow.vector.holders.NullableIntHolder; import org.apache.arrow.vector.holders.NullableVarCharHolder; /** * Auxiliary class meant to facilitate the implementation of {@link * FlightSqlProducer#getStreamSqlInfo}. * *

    Usage requires the user to add the required SqlInfo values using the {@code with*} methods * like {@link SqlInfoBuilder#withFlightSqlServerName(String)}, and request it back through the * {@link SqlInfoBuilder#send(List, ServerStreamListener)} method. */ @SuppressWarnings({"unused"}) public class SqlInfoBuilder { private final Map> providers = new HashMap<>(); /** * Gets a {@link NullableVarCharHolder} from the provided {@code string} using the provided {@code * buf}. * * @param string the {@link StandardCharsets#UTF_8}-encoded text input to store onto the holder. * @param buf the {@link ArrowBuf} from which to create the new holder. * @return a new {@link NullableVarCharHolder} with the provided input data {@code string}. */ public static NullableVarCharHolder getHolderForUtf8(final String string, final ArrowBuf buf) { final byte[] bytes = string.getBytes(StandardCharsets.UTF_8); buf.setBytes(0, bytes); final NullableVarCharHolder holder = new NullableVarCharHolder(); holder.buffer = buf; holder.end = bytes.length; holder.isSet = 1; return holder; } /** * Sets a value for {@link SqlInfo#FLIGHT_SQL_SERVER_NAME} in the builder. * * @param value the value for {@link SqlInfo#FLIGHT_SQL_SERVER_NAME} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withFlightSqlServerName(final String value) { return withStringProvider(SqlInfo.FLIGHT_SQL_SERVER_NAME_VALUE, value); } /** * Sets a value for {@link SqlInfo#FLIGHT_SQL_SERVER_VERSION} in the builder. * * @param value the value for {@link SqlInfo#FLIGHT_SQL_SERVER_VERSION} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withFlightSqlServerVersion(final String value) { return withStringProvider(SqlInfo.FLIGHT_SQL_SERVER_VERSION_VALUE, value); } /** * Sets a value for {@link SqlInfo#FLIGHT_SQL_SERVER_ARROW_VERSION} in the builder. * * @param value the value for {@link SqlInfo#FLIGHT_SQL_SERVER_ARROW_VERSION} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withFlightSqlServerArrowVersion(final String value) { return withStringProvider(SqlInfo.FLIGHT_SQL_SERVER_ARROW_VERSION_VALUE, value); } /** Set a value for SQL support. */ public SqlInfoBuilder withFlightSqlServerSql(boolean value) { return withBooleanProvider(SqlInfo.FLIGHT_SQL_SERVER_SQL_VALUE, value); } /** Set a value for Substrait support. */ public SqlInfoBuilder withFlightSqlServerSubstrait(boolean value) { return withBooleanProvider(SqlInfo.FLIGHT_SQL_SERVER_SUBSTRAIT_VALUE, value); } /** Set a value for Substrait minimum version support. */ public SqlInfoBuilder withFlightSqlServerSubstraitMinVersion(String value) { return withStringProvider(SqlInfo.FLIGHT_SQL_SERVER_SUBSTRAIT_MIN_VERSION_VALUE, value); } /** Set a value for Substrait maximum version support. */ public SqlInfoBuilder withFlightSqlServerSubstraitMaxVersion(String value) { return withStringProvider(SqlInfo.FLIGHT_SQL_SERVER_SUBSTRAIT_MAX_VERSION_VALUE, value); } /** Set a value for transaction support. */ public SqlInfoBuilder withFlightSqlServerTransaction(SqlSupportedTransaction value) { return withIntProvider(SqlInfo.FLIGHT_SQL_SERVER_TRANSACTION_VALUE, value.getNumber()); } /** Set a value for query cancellation support. */ public SqlInfoBuilder withFlightSqlServerCancel(boolean value) { return withBooleanProvider(SqlInfo.FLIGHT_SQL_SERVER_CANCEL_VALUE, value); } /** Set a value for bulk ingestion support. */ public SqlInfoBuilder withFlightSqlServerBulkIngestion(boolean value) { return withBooleanProvider(SqlInfo.FLIGHT_SQL_SERVER_BULK_INGESTION_VALUE, value); } /** Set a value for transaction support for bulk ingestion. */ public SqlInfoBuilder withFlightSqlServerBulkIngestionTransaction(boolean value) { return withBooleanProvider( SqlInfo.FLIGHT_SQL_SERVER_INGEST_TRANSACTIONS_SUPPORTED_VALUE, value); } /** Set a value for statement timeouts. */ public SqlInfoBuilder withFlightSqlServerStatementTimeout(int value) { return withIntProvider(SqlInfo.FLIGHT_SQL_SERVER_STATEMENT_TIMEOUT_VALUE, value); } /** Set a value for transaction timeouts. */ public SqlInfoBuilder withFlightSqlServerTransactionTimeout(int value) { return withIntProvider(SqlInfo.FLIGHT_SQL_SERVER_TRANSACTION_TIMEOUT_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_IDENTIFIER_QUOTE_CHAR} in the builder. * * @param value the value for {@link SqlInfo#SQL_IDENTIFIER_QUOTE_CHAR} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlIdentifierQuoteChar(final String value) { return withStringProvider(SqlInfo.SQL_IDENTIFIER_QUOTE_CHAR_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_SEARCH_STRING_ESCAPE} in the builder. * * @param value the value for {@link SqlInfo#SQL_SEARCH_STRING_ESCAPE} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSearchStringEscape(final String value) { return withStringProvider(SqlInfo.SQL_SEARCH_STRING_ESCAPE_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_EXTRA_NAME_CHARACTERS} in the builder. * * @param value the value for {@link SqlInfo#SQL_EXTRA_NAME_CHARACTERS} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlExtraNameCharacters(final String value) { return withStringProvider(SqlInfo.SQL_EXTRA_NAME_CHARACTERS_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_SCHEMA_TERM} in the builder. * * @param value the value for {@link SqlInfo#SQL_SCHEMA_TERM} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSchemaTerm(final String value) { return withStringProvider(SqlInfo.SQL_SCHEMA_TERM_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_CATALOG_TERM} in the builder. * * @param value the value for {@link SqlInfo#SQL_CATALOG_TERM} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlCatalogTerm(final String value) { return withStringProvider(SqlInfo.SQL_CATALOG_TERM_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_PROCEDURE_TERM} in the builder. * * @param value the value for {@link SqlInfo#SQL_PROCEDURE_TERM} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlProcedureTerm(final String value) { return withStringProvider(SqlInfo.SQL_PROCEDURE_TERM_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_DDL_CATALOG} in the builder. * * @param value the value for {@link SqlInfo#SQL_DDL_CATALOG} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlDdlCatalog(final boolean value) { return withBooleanProvider(SqlInfo.SQL_DDL_CATALOG_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_DDL_SCHEMA} in the builder. * * @param value the value for {@link SqlInfo#SQL_DDL_SCHEMA} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlDdlSchema(final boolean value) { return withBooleanProvider(SqlInfo.SQL_DDL_SCHEMA_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_DDL_TABLE} in the builder. * * @param value the value for {@link SqlInfo#SQL_DDL_TABLE} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlDdlTable(final boolean value) { return withBooleanProvider(SqlInfo.SQL_DDL_TABLE_VALUE, value); } /** * Sets a value for {@link SqlInfo#FLIGHT_SQL_SERVER_READ_ONLY} in the builder. * * @param value the value for {@link SqlInfo#FLIGHT_SQL_SERVER_READ_ONLY} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withFlightSqlServerReadOnly(final boolean value) { return withBooleanProvider(SqlInfo.FLIGHT_SQL_SERVER_READ_ONLY_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_SUPPORTS_COLUMN_ALIASING} in the builder. * * @param value the value for {@link SqlInfo#SQL_SUPPORTS_COLUMN_ALIASING} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSupportsColumnAliasing(final boolean value) { return withBooleanProvider(SqlInfo.SQL_SUPPORTS_COLUMN_ALIASING_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_NULL_PLUS_NULL_IS_NULL} in the builder. * * @param value the value for {@link SqlInfo#SQL_NULL_PLUS_NULL_IS_NULL} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlNullPlusNullIsNull(final boolean value) { return withBooleanProvider(SqlInfo.SQL_NULL_PLUS_NULL_IS_NULL_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_SUPPORTS_TABLE_CORRELATION_NAMES} in the builder. * * @param value the value for {@link SqlInfo#SQL_SUPPORTS_TABLE_CORRELATION_NAMES} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSupportsTableCorrelationNames(final boolean value) { return withBooleanProvider(SqlInfo.SQL_SUPPORTS_TABLE_CORRELATION_NAMES_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_SUPPORTS_DIFFERENT_TABLE_CORRELATION_NAMES} in the builder. * * @param value the value for {@link SqlInfo#SQL_SUPPORTS_DIFFERENT_TABLE_CORRELATION_NAMES} to be * set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSupportsDifferentTableCorrelationNames(final boolean value) { return withBooleanProvider(SqlInfo.SQL_SUPPORTS_DIFFERENT_TABLE_CORRELATION_NAMES_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_SUPPORTS_EXPRESSIONS_IN_ORDER_BY} in the builder. * * @param value the value for {@link SqlInfo#SQL_SUPPORTS_EXPRESSIONS_IN_ORDER_BY} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSupportsExpressionsInOrderBy(final boolean value) { return withBooleanProvider(SqlInfo.SQL_SUPPORTS_EXPRESSIONS_IN_ORDER_BY_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_SUPPORTS_ORDER_BY_UNRELATED} in the builder. * * @param value the value for {@link SqlInfo#SQL_SUPPORTS_ORDER_BY_UNRELATED} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSupportsOrderByUnrelated(final boolean value) { return withBooleanProvider(SqlInfo.SQL_SUPPORTS_ORDER_BY_UNRELATED_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_SUPPORTS_LIKE_ESCAPE_CLAUSE} in the builder. * * @param value the value for {@link SqlInfo#SQL_SUPPORTS_LIKE_ESCAPE_CLAUSE} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSupportsLikeEscapeClause(final boolean value) { return withBooleanProvider(SqlInfo.SQL_SUPPORTS_LIKE_ESCAPE_CLAUSE_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_SUPPORTS_NON_NULLABLE_COLUMNS} in the builder. * * @param value the value for {@link SqlInfo#SQL_SUPPORTS_NON_NULLABLE_COLUMNS} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSupportsNonNullableColumns(final boolean value) { return withBooleanProvider(SqlInfo.SQL_SUPPORTS_NON_NULLABLE_COLUMNS_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_SUPPORTS_INTEGRITY_ENHANCEMENT_FACILITY} in the builder. * * @param value the value for {@link SqlInfo#SQL_SUPPORTS_INTEGRITY_ENHANCEMENT_FACILITY} to be * set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSupportsIntegrityEnhancementFacility(final boolean value) { return withBooleanProvider(SqlInfo.SQL_SUPPORTS_INTEGRITY_ENHANCEMENT_FACILITY_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_CATALOG_AT_START} in the builder. * * @param value the value for {@link SqlInfo#SQL_CATALOG_AT_START} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlCatalogAtStart(final boolean value) { return withBooleanProvider(SqlInfo.SQL_CATALOG_AT_START_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_SELECT_FOR_UPDATE_SUPPORTED} in the builder. * * @param value the value for {@link SqlInfo#SQL_SELECT_FOR_UPDATE_SUPPORTED} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSelectForUpdateSupported(final boolean value) { return withBooleanProvider(SqlInfo.SQL_SELECT_FOR_UPDATE_SUPPORTED_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_STORED_PROCEDURES_SUPPORTED} in the builder. * * @param value the value for {@link SqlInfo#SQL_STORED_PROCEDURES_SUPPORTED} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlStoredProceduresSupported(final boolean value) { return withBooleanProvider(SqlInfo.SQL_STORED_PROCEDURES_SUPPORTED_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_CORRELATED_SUBQUERIES_SUPPORTED} in the builder. * * @param value the value for {@link SqlInfo#SQL_CORRELATED_SUBQUERIES_SUPPORTED} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlCorrelatedSubqueriesSupported(final boolean value) { return withBooleanProvider(SqlInfo.SQL_CORRELATED_SUBQUERIES_SUPPORTED_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_ROW_SIZE_INCLUDES_BLOBS} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_ROW_SIZE_INCLUDES_BLOBS} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxRowSizeIncludesBlobs(final boolean value) { return withBooleanProvider(SqlInfo.SQL_MAX_ROW_SIZE_INCLUDES_BLOBS_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_TRANSACTIONS_SUPPORTED} in the builder. * * @param value the value for {@link SqlInfo#SQL_TRANSACTIONS_SUPPORTED} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlTransactionsSupported(final boolean value) { return withBooleanProvider(SqlInfo.SQL_TRANSACTIONS_SUPPORTED_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_DATA_DEFINITION_CAUSES_TRANSACTION_COMMIT} in the builder. * * @param value the value for {@link SqlInfo#SQL_DATA_DEFINITION_CAUSES_TRANSACTION_COMMIT} to be * set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlDataDefinitionCausesTransactionCommit(final boolean value) { return withBooleanProvider(SqlInfo.SQL_DATA_DEFINITION_CAUSES_TRANSACTION_COMMIT_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_DATA_DEFINITIONS_IN_TRANSACTIONS_IGNORED} in the builder. * * @param value the value for {@link SqlInfo#SQL_DATA_DEFINITIONS_IN_TRANSACTIONS_IGNORED} to be * set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlDataDefinitionsInTransactionsIgnored(final boolean value) { return withBooleanProvider(SqlInfo.SQL_DATA_DEFINITIONS_IN_TRANSACTIONS_IGNORED_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_BATCH_UPDATES_SUPPORTED} in the builder. * * @param value the value for {@link SqlInfo#SQL_BATCH_UPDATES_SUPPORTED} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlBatchUpdatesSupported(final boolean value) { return withBooleanProvider(SqlInfo.SQL_BATCH_UPDATES_SUPPORTED_VALUE, value); } /** * Sets a value for { @link SqlInfo#SQL_SAVEPOINTS_SUPPORTED} in the builder. * * @param value the value for {@link SqlInfo#SQL_SAVEPOINTS_SUPPORTED} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSavepointsSupported(final boolean value) { return withBooleanProvider(SqlInfo.SQL_SAVEPOINTS_SUPPORTED_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_NAMED_PARAMETERS_SUPPORTED} in the builder. * * @param value the value for {@link SqlInfo#SQL_NAMED_PARAMETERS_SUPPORTED} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlNamedParametersSupported(final boolean value) { return withBooleanProvider(SqlInfo.SQL_NAMED_PARAMETERS_SUPPORTED_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_LOCATORS_UPDATE_COPY} in the builder. * * @param value the value for {@link SqlInfo#SQL_LOCATORS_UPDATE_COPY} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlLocatorsUpdateCopy(final boolean value) { return withBooleanProvider(SqlInfo.SQL_LOCATORS_UPDATE_COPY_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_STORED_FUNCTIONS_USING_CALL_SYNTAX_SUPPORTED} in the * builder. * * @param value the value for {@link SqlInfo#SQL_STORED_FUNCTIONS_USING_CALL_SYNTAX_SUPPORTED} to * be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlStoredFunctionsUsingCallSyntaxSupported(final boolean value) { return withBooleanProvider( SqlInfo.SQL_STORED_FUNCTIONS_USING_CALL_SYNTAX_SUPPORTED_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_IDENTIFIER_CASE} in the builder. * * @param value the value for {@link SqlInfo#SQL_IDENTIFIER_CASE} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlIdentifierCase(final SqlSupportedCaseSensitivity value) { return withBitIntProvider(SqlInfo.SQL_IDENTIFIER_CASE_VALUE, value.getNumber()); } /** * Sets a value for {@link SqlInfo#SQL_QUOTED_IDENTIFIER_CASE} in the builder. * * @param value the value for {@link SqlInfo#SQL_QUOTED_IDENTIFIER_CASE} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlQuotedIdentifierCase(final SqlSupportedCaseSensitivity value) { return withBitIntProvider(SqlInfo.SQL_QUOTED_IDENTIFIER_CASE_VALUE, value.getNumber()); } /** * Sets a value for {@link SqlInfo#SQL_ALL_TABLES_ARE_SELECTABLE} in the builder. * * @param value the value for {@link SqlInfo#SQL_ALL_TABLES_ARE_SELECTABLE} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlAllTablesAreSelectable(final boolean value) { return withBooleanProvider(SqlInfo.SQL_ALL_TABLES_ARE_SELECTABLE_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_NULL_ORDERING} in the builder. * * @param value the value for {@link SqlInfo#SQL_NULL_ORDERING} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlNullOrdering(final SqlNullOrdering value) { return withBitIntProvider(SqlInfo.SQL_NULL_ORDERING_VALUE, value.getNumber()); } /** * Sets a value SqlInf @link SqlInfo#SQL_MAX_BINARY_LITERAL_LENGTH} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_BINARY_LITERAL_LENGTH} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxBinaryLiteralLength(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_BINARY_LITERAL_LENGTH_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_CHAR_LITERAL_LENGTH} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_CHAR_LITERAL_LENGTH} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxCharLiteralLength(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_CHAR_LITERAL_LENGTH_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_COLUMN_NAME_LENGTH} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_COLUMN_NAME_LENGTH} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxColumnNameLength(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_COLUMN_NAME_LENGTH_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_COLUMNS_IN_GROUP_BY} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_COLUMNS_IN_GROUP_BY} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxColumnsInGroupBy(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_COLUMNS_IN_GROUP_BY_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_COLUMNS_IN_INDEX} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_COLUMNS_IN_INDEX} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxColumnsInIndex(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_COLUMNS_IN_INDEX_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_COLUMNS_IN_ORDER_BY} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_COLUMNS_IN_ORDER_BY} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxColumnsInOrderBy(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_COLUMNS_IN_ORDER_BY_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_COLUMNS_IN_SELECT} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_COLUMNS_IN_SELECT} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxColumnsInSelect(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_COLUMNS_IN_SELECT_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_COLUMNS_IN_TABLE} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_COLUMNS_IN_TABLE} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxColumnsInTable(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_COLUMNS_IN_TABLE_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_CONNECTIONS} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_CONNECTIONS} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxConnections(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_CONNECTIONS_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_CURSOR_NAME_LENGTH} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_CURSOR_NAME_LENGTH} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxCursorNameLength(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_CURSOR_NAME_LENGTH_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_INDEX_LENGTH} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_INDEX_LENGTH} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxIndexLength(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_INDEX_LENGTH_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_DB_SCHEMA_NAME_LENGTH} in the builder. * * @param value the value for {@link SqlInfo#SQL_DB_SCHEMA_NAME_LENGTH} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlDbSchemaNameLength(final long value) { return withBitIntProvider(SqlInfo.SQL_DB_SCHEMA_NAME_LENGTH_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_PROCEDURE_NAME_LENGTH} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_PROCEDURE_NAME_LENGTH} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxProcedureNameLength(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_PROCEDURE_NAME_LENGTH_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_CATALOG_NAME_LENGTH} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_CATALOG_NAME_LENGTH} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxCatalogNameLength(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_CATALOG_NAME_LENGTH_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_ROW_SIZE} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_ROW_SIZE} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxRowSize(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_ROW_SIZE_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_STATEMENT_LENGTH} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_STATEMENT_LENGTH} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxStatementLength(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_STATEMENT_LENGTH_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_STATEMENTS} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_STATEMENTS} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxStatements(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_STATEMENTS_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_TABLE_NAME_LENGTH} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_TABLE_NAME_LENGTH} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxTableNameLength(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_TABLE_NAME_LENGTH_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_TABLES_IN_SELECT} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_TABLES_IN_SELECT} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxTablesInSelect(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_TABLES_IN_SELECT_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_MAX_USERNAME_LENGTH} in the builder. * * @param value the value for {@link SqlInfo#SQL_MAX_USERNAME_LENGTH} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlMaxUsernameLength(final long value) { return withBitIntProvider(SqlInfo.SQL_MAX_USERNAME_LENGTH_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_DEFAULT_TRANSACTION_ISOLATION} in the builder. * * @param value the value for {@link SqlInfo#SQL_DEFAULT_TRANSACTION_ISOLATION} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlDefaultTransactionIsolation(final long value) { return withBitIntProvider(SqlInfo.SQL_DEFAULT_TRANSACTION_ISOLATION_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_SUPPORTED_GROUP_BY} in the builder. * * @param values the value for {@link SqlInfo#SQL_SUPPORTED_GROUP_BY} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSupportedGroupBy(final SqlSupportedGroupBy... values) { return withEnumProvider(SqlInfo.SQL_SUPPORTED_GROUP_BY_VALUE, values); } /** * Sets a value for {@link SqlInfo#SQL_SUPPORTED_GRAMMAR} in the builder. * * @param values the value for {@link SqlInfo#SQL_SUPPORTED_GRAMMAR} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSupportedGrammar(final SupportedSqlGrammar... values) { return withEnumProvider(SqlInfo.SQL_SUPPORTED_GRAMMAR_VALUE, values); } /** * Sets a value for {@link SqlInfo#SQL_ANSI92_SUPPORTED_LEVEL} in the builder. * * @param values the value for {@link SqlInfo#SQL_ANSI92_SUPPORTED_LEVEL} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlAnsi92SupportedLevel( final SupportedAnsi92SqlGrammarLevel... values) { return withEnumProvider(SqlInfo.SQL_ANSI92_SUPPORTED_LEVEL_VALUE, values); } /** * Sets a value for {@link SqlInfo#SQL_SCHEMAS_SUPPORTED_ACTIONS} in the builder. * * @param values the value for {@link SqlInfo#SQL_SCHEMAS_SUPPORTED_ACTIONS} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSchemasSupportedActions(final SqlSupportedElementActions... values) { return withEnumProvider(SqlInfo.SQL_SCHEMAS_SUPPORTED_ACTIONS_VALUE, values); } /** * Sets a value for {@link SqlInfo#SQL_CATALOGS_SUPPORTED_ACTIONS} in the builder. * * @param values the value for {@link SqlInfo#SQL_CATALOGS_SUPPORTED_ACTIONS} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlCatalogsSupportedActions( final SqlSupportedElementActions... values) { return withEnumProvider(SqlInfo.SQL_CATALOGS_SUPPORTED_ACTIONS_VALUE, values); } /** * Sets a value for {@link SqlInfo#SQL_SUPPORTED_POSITIONED_COMMANDS} in the builder. * * @param values the value for {@link SqlInfo#SQL_SUPPORTED_POSITIONED_COMMANDS} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSupportedPositionedCommands( final SqlSupportedPositionedCommands... values) { return withEnumProvider(SqlInfo.SQL_SUPPORTED_POSITIONED_COMMANDS_VALUE, values); } /** * Sets a value for {@link SqlInfo#SQL_SUPPORTED_SUBQUERIES} in the builder. * * @param values the value for {@link SqlInfo#SQL_SUPPORTED_SUBQUERIES} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSubQueriesSupported(final SqlSupportedSubqueries... values) { return withEnumProvider(SqlInfo.SQL_SUPPORTED_SUBQUERIES_VALUE, values); } /** * Sets a value for {@link SqlInfo#SQL_SUPPORTED_UNIONS} in the builder. * * @param values the values for {@link SqlInfo#SQL_SUPPORTED_UNIONS} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSupportedUnions(final SqlSupportedUnions... values) { return withEnumProvider(SqlInfo.SQL_SUPPORTED_UNIONS_VALUE, values); } /** * Sets a value for {@link SqlInfo#SQL_OUTER_JOINS_SUPPORT_LEVEL} in the builder. * * @param value the value for {@link SqlInfo#SQL_OUTER_JOINS_SUPPORT_LEVEL} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlOuterJoinSupportLevel(final SqlOuterJoinsSupportLevel... value) { return withEnumProvider(SqlInfo.SQL_OUTER_JOINS_SUPPORT_LEVEL_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_SUPPORTED_TRANSACTIONS_ISOLATION_LEVELS} in the builder. * * @param values the values for {@link SqlInfo#SQL_SUPPORTED_TRANSACTIONS_ISOLATION_LEVELS} to be * set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSupportedTransactionsIsolationLevels( final SqlTransactionIsolationLevel... values) { return withEnumProvider(SqlInfo.SQL_SUPPORTED_TRANSACTIONS_ISOLATION_LEVELS_VALUE, values); } /** * Sets a value for {@link SqlInfo#SQL_SUPPORTED_RESULT_SET_TYPES} in the builder. * * @param values the values for {@link SqlInfo#SQL_SUPPORTED_RESULT_SET_TYPES} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSupportedResultSetTypes(final SqlSupportedResultSetType... values) { return withEnumProvider(SqlInfo.SQL_SUPPORTED_RESULT_SET_TYPES_VALUE, values); } /** * Sets a value for {@link SqlInfo#SQL_KEYWORDS} in the builder. * * @param value the values for {@link SqlInfo#SQL_KEYWORDS} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlKeywords(final String[] value) { return withStringArrayProvider(SqlInfo.SQL_KEYWORDS_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_NUMERIC_FUNCTIONS} in the builder. * * @param value the values for {@link SqlInfo#SQL_NUMERIC_FUNCTIONS} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlNumericFunctions(final String[] value) { return withStringArrayProvider(SqlInfo.SQL_NUMERIC_FUNCTIONS_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_STRING_FUNCTIONS} in the builder. * * @param value the values for {@link SqlInfo#SQL_STRING_FUNCTIONS} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlStringFunctions(final String[] value) { return withStringArrayProvider(SqlInfo.SQL_STRING_FUNCTIONS_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_SYSTEM_FUNCTIONS} in the builder. * * @param value the values for {@link SqlInfo#SQL_SYSTEM_FUNCTIONS} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSystemFunctions(final String[] value) { return withStringArrayProvider(SqlInfo.SQL_SYSTEM_FUNCTIONS_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_DATETIME_FUNCTIONS} in the builder. * * @param value the values for {@link SqlInfo#SQL_DATETIME_FUNCTIONS} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlDatetimeFunctions(final String[] value) { return withStringArrayProvider(SqlInfo.SQL_DATETIME_FUNCTIONS_VALUE, value); } /** * Sets a value for {@link SqlInfo#SQL_SUPPORTS_CONVERT} in the builder. * * @param value the values for {@link SqlInfo#SQL_SUPPORTS_CONVERT} to be set. * @return the SqlInfoBuilder itself. */ public SqlInfoBuilder withSqlSupportsConvert(final Map> value) { return withIntToIntListMapProvider(SqlInfo.SQL_SUPPORTS_CONVERT_VALUE, value); } private void addProvider(final int sqlInfo, final ObjIntConsumer provider) { providers.put(sqlInfo, provider); } private SqlInfoBuilder withEnumProvider(final int sqlInfo, final ProtocolMessageEnum[] values) { return withIntProvider(sqlInfo, (int) createBitmaskFromEnums(values)); } private SqlInfoBuilder withIntProvider(final int sqlInfo, final int value) { addProvider(sqlInfo, (root, index) -> setDataForIntField(root, index, sqlInfo, value)); return this; } private SqlInfoBuilder withBitIntProvider(final int sqlInfo, final long value) { addProvider(sqlInfo, (root, index) -> setDataForBigIntField(root, index, sqlInfo, value)); return this; } private SqlInfoBuilder withBooleanProvider(final int sqlInfo, final boolean value) { addProvider(sqlInfo, (root, index) -> setDataForBooleanField(root, index, sqlInfo, value)); return this; } private SqlInfoBuilder withStringProvider(final int sqlInfo, final String value) { addProvider(sqlInfo, (root, index) -> setDataForUtf8Field(root, index, sqlInfo, value)); return this; } private SqlInfoBuilder withStringArrayProvider(final int sqlInfo, final String[] value) { addProvider(sqlInfo, (root, index) -> setDataVarCharListField(root, index, sqlInfo, value)); return this; } private SqlInfoBuilder withIntToIntListMapProvider( final int sqlInfo, final Map> value) { addProvider(sqlInfo, (root, index) -> setIntToIntListMapField(root, index, sqlInfo, value)); return this; } /** * Send the requested information to given ServerStreamListener. * * @param infos List of SqlInfo to be sent. * @param listener ServerStreamListener to send data to. */ public void send(List infos, final ServerStreamListener listener) { if (infos == null || infos.isEmpty()) { infos = new ArrayList<>(providers.keySet()); } try (final BufferAllocator allocator = new RootAllocator(); final VectorSchemaRoot root = VectorSchemaRoot.create(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA, allocator)) { final int rows = infos.size(); for (int i = 0; i < rows; i++) { providers.get(infos.get(i)).accept(root, i); } root.setRowCount(rows); listener.start(root); listener.putNext(); } catch (final Throwable throwable) { listener.error(throwable); } finally { listener.completed(); } } private void setInfoName(final VectorSchemaRoot root, final int index, final int info) { final UInt4Vector infoName = (UInt4Vector) root.getVector("info_name"); infoName.setSafe(index, info); } private void setValues( final VectorSchemaRoot root, final int index, final byte typeId, final Consumer dataSetter) { final DenseUnionVector values = (DenseUnionVector) root.getVector("value"); values.setTypeId(index, typeId); dataSetter.accept(values); } /** * Executes the given action on an ad-hoc, newly created instance of {@link ArrowBuf}. * * @param executor the action to take. */ private void onCreateArrowBuf(final Consumer executor) { try (final BufferAllocator allocator = new RootAllocator(); final ArrowBuf buf = allocator.buffer(1024)) { executor.accept(buf); } } private void setDataForUtf8Field( final VectorSchemaRoot root, final int index, final int sqlInfo, final String value) { setInfoName(root, index, sqlInfo); onCreateArrowBuf( buf -> { final Consumer producer = values -> values.setSafe(index, getHolderForUtf8(value, buf)); setValues(root, index, (byte) 0, producer); }); } private void setDataForIntField( final VectorSchemaRoot root, final int index, final int sqlInfo, final int value) { setInfoName(root, index, sqlInfo); final NullableIntHolder dataHolder = new NullableIntHolder(); dataHolder.isSet = 1; dataHolder.value = value; setValues(root, index, (byte) 3, values -> values.setSafe(index, dataHolder)); } private void setDataForBigIntField( final VectorSchemaRoot root, final int index, final int sqlInfo, final long value) { setInfoName(root, index, sqlInfo); final NullableBigIntHolder dataHolder = new NullableBigIntHolder(); dataHolder.isSet = 1; dataHolder.value = value; setValues(root, index, (byte) 2, values -> values.setSafe(index, dataHolder)); } private void setDataForBooleanField( final VectorSchemaRoot root, final int index, final int sqlInfo, final boolean value) { setInfoName(root, index, sqlInfo); final NullableBitHolder dataHolder = new NullableBitHolder(); dataHolder.isSet = 1; dataHolder.value = value ? 1 : 0; setValues(root, index, (byte) 1, values -> values.setSafe(index, dataHolder)); } private void setDataVarCharListField( final VectorSchemaRoot root, final int index, final int sqlInfo, final String[] values) { final DenseUnionVector denseUnion = (DenseUnionVector) root.getVector("value"); final ListVector listVector = denseUnion.getList((byte) 4); final int listIndex = listVector.getValueCount(); final int denseUnionValueCount = index + 1; final int listVectorValueCount = listIndex + 1; denseUnion.setValueCount(denseUnionValueCount); listVector.setValueCount(listVectorValueCount); final UnionListWriter writer = listVector.getWriter(); writer.setPosition(listIndex); writer.startList(); final int length = values.length; range(0, length) .forEach( i -> onCreateArrowBuf( buf -> { final byte[] bytes = values[i].getBytes(StandardCharsets.UTF_8); buf.setBytes(0, bytes); writer.writeVarChar(0, bytes.length, buf); })); writer.endList(); writer.setValueCount(listVectorValueCount); denseUnion.setTypeId(index, (byte) 4); denseUnion.getOffsetBuffer().setInt(index * 4L, listIndex); setInfoName(root, index, sqlInfo); } private void setIntToIntListMapField( final VectorSchemaRoot root, final int index, final int sqlInfo, final Map> values) { final DenseUnionVector denseUnion = (DenseUnionVector) root.getVector("value"); final MapVector mapVector = denseUnion.getMap((byte) 5); final int mapIndex = mapVector.getValueCount(); denseUnion.setValueCount(index + 1); mapVector.setValueCount(mapIndex + 1); final UnionMapWriter mapWriter = mapVector.getWriter(); mapWriter.setPosition(mapIndex); mapWriter.startMap(); values.forEach( (key, value) -> { mapWriter.startEntry(); mapWriter.key().integer().writeInt(key); final BaseWriter.ListWriter listWriter = mapWriter.value().list(); listWriter.startList(); for (final int v : value) { listWriter.integer().writeInt(v); } listWriter.endList(); mapWriter.endEntry(); }); mapWriter.endMap(); mapWriter.setValueCount(mapIndex + 1); denseUnion.setTypeId(index, (byte) 5); denseUnion.getOffsetBuffer().setInt(index * 4L, mapIndex); setInfoName(root, index, sqlInfo); } } apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/example/000077500000000000000000000000001515451615200317325ustar00rootroot00000000000000FlightSqlClientDemoApp.java000066400000000000000000000213701515451615200370230ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/example/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql.example; import java.util.ArrayList; import java.util.List; import org.apache.arrow.flight.CallOption; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.sql.FlightSqlClient; import org.apache.arrow.flight.sql.util.TableRef; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.commons.cli.CommandLine; import org.apache.commons.cli.CommandLineParser; import org.apache.commons.cli.DefaultParser; import org.apache.commons.cli.HelpFormatter; import org.apache.commons.cli.Options; import org.apache.commons.cli.ParseException; /** Flight SQL Client Demo CLI Application. */ public class FlightSqlClientDemoApp implements AutoCloseable { public final List callOptions = new ArrayList<>(); public final BufferAllocator allocator; public FlightSqlClient flightSqlClient; public FlightSqlClientDemoApp(final BufferAllocator bufferAllocator) { allocator = bufferAllocator; } public static void main(final String[] args) throws Exception { final Options options = new Options(); options.addRequiredOption("host", "host", true, "Host to connect to"); options.addRequiredOption("port", "port", true, "Port to connect to"); options.addRequiredOption("command", "command", true, "Method to run"); options.addOption("query", "query", true, "Query"); options.addOption("catalog", "catalog", true, "Catalog"); options.addOption("schema", "schema", true, "Schema"); options.addOption("table", "table", true, "Table"); CommandLineParser parser = new DefaultParser(); HelpFormatter formatter = new HelpFormatter(); CommandLine cmd; try { cmd = parser.parse(options, args); try (final FlightSqlClientDemoApp thisApp = new FlightSqlClientDemoApp(new RootAllocator(Integer.MAX_VALUE))) { thisApp.executeApp(cmd); } } catch (final ParseException e) { System.out.println(e.getMessage()); formatter.printHelp("FlightSqlClientDemoApp -host localhost -port 32010 ...", options); throw e; } } /** * Gets the current {@link CallOption} as an array; usually used as an argument in {@link * FlightSqlClient} methods. * * @return current {@link CallOption} array. */ public CallOption[] getCallOptions() { return callOptions.toArray(new CallOption[0]); } /** * Calls {@link FlightSqlClientDemoApp#createFlightSqlClient(String, int)} in order to create a * {@link FlightSqlClient} to be used in future calls, and then calls {@link * FlightSqlClientDemoApp#executeCommand(CommandLine)} to execute the command parsed at execution. * * @param cmd parsed {@link CommandLine}; often the result of {@link DefaultParser#parse(Options, * String[])}. */ public void executeApp(final CommandLine cmd) throws Exception { final String host = cmd.getOptionValue("host").trim(); final int port = Integer.parseInt(cmd.getOptionValue("port").trim()); createFlightSqlClient(host, port); executeCommand(cmd); } /** * Parses the "{@code command}" CLI argument and redirects to the appropriate method. * * @param cmd parsed {@link CommandLine}; often the result of {@link DefaultParser#parse(Options, * String[])}. */ public void executeCommand(CommandLine cmd) throws Exception { switch (cmd.getOptionValue("command").trim()) { case "Execute": exampleExecute(cmd.getOptionValue("query")); break; case "ExecuteUpdate": exampleExecuteUpdate(cmd.getOptionValue("query")); break; case "GetCatalogs": exampleGetCatalogs(); break; case "GetSchemas": exampleGetSchemas(cmd.getOptionValue("catalog"), cmd.getOptionValue("schema")); break; case "GetTableTypes": exampleGetTableTypes(); break; case "GetTables": exampleGetTables( cmd.getOptionValue("catalog"), cmd.getOptionValue("schema"), cmd.getOptionValue("table")); break; case "GetExportedKeys": exampleGetExportedKeys( cmd.getOptionValue("catalog"), cmd.getOptionValue("schema"), cmd.getOptionValue("table")); break; case "GetImportedKeys": exampleGetImportedKeys( cmd.getOptionValue("catalog"), cmd.getOptionValue("schema"), cmd.getOptionValue("table")); break; case "GetPrimaryKeys": exampleGetPrimaryKeys( cmd.getOptionValue("catalog"), cmd.getOptionValue("schema"), cmd.getOptionValue("table")); break; default: System.out.println( "Command used is not valid! Please use one of: \n" + "[\"ExecuteUpdate\",\n" + "\"Execute\",\n" + "\"GetCatalogs\",\n" + "\"GetSchemas\",\n" + "\"GetTableTypes\",\n" + "\"GetTables\",\n" + "\"GetExportedKeys\",\n" + "\"GetImportedKeys\",\n" + "\"GetPrimaryKeys\"]"); } } /** * Creates a {@link FlightSqlClient} to be used with the example methods. * * @param host client's hostname. * @param port client's port. */ public void createFlightSqlClient(final String host, final int port) { final Location clientLocation = Location.forGrpcInsecure(host, port); flightSqlClient = new FlightSqlClient(FlightClient.builder(allocator, clientLocation).build()); } private void exampleExecute(final String query) throws Exception { printFlightInfoResults(flightSqlClient.execute(query, getCallOptions())); } private void exampleExecuteUpdate(final String query) { System.out.println( "Updated: " + flightSqlClient.executeUpdate(query, getCallOptions()) + "rows."); } private void exampleGetCatalogs() throws Exception { printFlightInfoResults(flightSqlClient.getCatalogs(getCallOptions())); } private void exampleGetSchemas(final String catalog, final String schema) throws Exception { printFlightInfoResults(flightSqlClient.getSchemas(catalog, schema, getCallOptions())); } private void exampleGetTableTypes() throws Exception { printFlightInfoResults(flightSqlClient.getTableTypes(getCallOptions())); } private void exampleGetTables(final String catalog, final String schema, final String table) throws Exception { // For now, this won't filter by table types. printFlightInfoResults( flightSqlClient.getTables(catalog, schema, table, null, false, getCallOptions())); } private void exampleGetExportedKeys(final String catalog, final String schema, final String table) throws Exception { printFlightInfoResults( flightSqlClient.getExportedKeys(TableRef.of(catalog, schema, table), getCallOptions())); } private void exampleGetImportedKeys(final String catalog, final String schema, final String table) throws Exception { printFlightInfoResults( flightSqlClient.getImportedKeys(TableRef.of(catalog, schema, table), getCallOptions())); } private void exampleGetPrimaryKeys(final String catalog, final String schema, final String table) throws Exception { printFlightInfoResults( flightSqlClient.getPrimaryKeys(TableRef.of(catalog, schema, table), getCallOptions())); } private void printFlightInfoResults(final FlightInfo flightInfo) throws Exception { final FlightStream stream = flightSqlClient.getStream(flightInfo.getEndpoints().get(0).getTicket(), getCallOptions()); while (stream.next()) { try (final VectorSchemaRoot root = stream.getRoot()) { System.out.println(root.contentToTSVString()); } } stream.close(); } @Override public void close() throws Exception { flightSqlClient.close(); allocator.close(); } } apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/util/000077500000000000000000000000001515451615200312545ustar00rootroot00000000000000SqlInfoOptionsUtils.java000066400000000000000000000053401515451615200360120ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql.util; import com.google.protobuf.ProtocolMessageEnum; import java.util.Arrays; import java.util.Collection; import org.apache.arrow.flight.sql.FlightSqlClient; import org.apache.arrow.flight.sql.impl.FlightSql.SqlInfo; /** Utility class for {@link SqlInfo} and {@link FlightSqlClient#getSqlInfo} option parsing. */ public final class SqlInfoOptionsUtils { private SqlInfoOptionsUtils() { // Prevent instantiation. } /** * Returns whether the provided {@code bitmask} points to the provided {@link ProtocolMessageEnum} * by comparing {@link ProtocolMessageEnum#getNumber} with the respective bit index of the {@code * bitmask}. * * @param enumInstance the protobuf message enum to use. * @param bitmask the bitmask response from {@link FlightSqlClient#getSqlInfo}. * @return whether the provided {@code bitmask} points to the specified {@code enumInstance}. */ public static boolean doesBitmaskTranslateToEnum( final ProtocolMessageEnum enumInstance, final long bitmask) { return ((bitmask >> enumInstance.getNumber()) & 1) == 1; } /** * Creates a bitmask that translates to the specified {@code enums}. * * @param enums the {@link ProtocolMessageEnum} instances to represent as bitmask. * @return the bitmask. */ public static long createBitmaskFromEnums(final ProtocolMessageEnum... enums) { return createBitmaskFromEnums(Arrays.asList(enums)); } /** * Creates a bitmask that translates to the specified {@code enums}. * * @param enums the {@link ProtocolMessageEnum} instances to represent as bitmask. * @return the bitmask. */ public static long createBitmaskFromEnums(final Collection enums) { return enums.stream() .mapToInt(ProtocolMessageEnum::getNumber) .map(bitIndexToSet -> 1 << bitIndexToSet) .reduce((firstBitmask, secondBitmask) -> firstBitmask | secondBitmask) .orElse(0); } } TableRef.java000066400000000000000000000043001515451615200335210ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql.util; /** A helper class to reference a table to be passed to the flight sql client. */ public class TableRef { private final String catalog; private final String dbSchema; private final String table; /** * The complete constructor for the TableRef class. * * @param catalog the catalog from a table. * @param dbSchema the database schema from a table. * @param table the table name from a table. */ public TableRef(String catalog, String dbSchema, String table) { this.catalog = catalog; this.dbSchema = dbSchema; this.table = table; } /** * A static initializer of the TableRef with all the arguments. * * @param catalog the catalog from a table. * @param dbSchema the database schema from a table. * @param table the table name from a table. * @return A TableRef object. */ public static TableRef of(String catalog, String dbSchema, String table) { return new TableRef(catalog, dbSchema, table); } /** * Retrieve the catalog from the object. * * @return the catalog. */ public String getCatalog() { return catalog; } /** * Retrieves the db schema from the object. * * @return the dbSchema */ public String getDbSchema() { return dbSchema; } /** * Retrieves the table from the object. * * @return the table. */ public String getTable() { return table; } } apache-arrow-java-19.0.0/flight/flight-sql/src/test/000077500000000000000000000000001515451615200221735ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/000077500000000000000000000000001515451615200231145ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/000077500000000000000000000000001515451615200237035ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/000077500000000000000000000000001515451615200251245ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200262565ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/arrow/flight/000077500000000000000000000000001515451615200275335ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/000077500000000000000000000000001515451615200303325ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/000077500000000000000000000000001515451615200317655ustar00rootroot00000000000000DoPutPreparedStatementResultPOJO.java000066400000000000000000000023621515451615200410460ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql.example; import java.io.Serializable; public class DoPutPreparedStatementResultPOJO implements Serializable { private String query; private byte[] parameters; public DoPutPreparedStatementResultPOJO(String query, byte[] parameters) { this.query = query; this.parameters = parameters.clone(); } public String getQuery() { return query; } public byte[] getParameters() { return parameters; } } FlightSqlExample.java000066400000000000000000002062441515451615200357720ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql.example; import static com.google.common.base.Strings.emptyToNull; import static com.google.protobuf.Any.pack; import static com.google.protobuf.ByteString.copyFrom; import static java.lang.String.format; import static java.util.Collections.singletonList; import static java.util.Objects.isNull; import static java.util.UUID.randomUUID; import static java.util.stream.IntStream.range; import static org.apache.arrow.adapter.jdbc.JdbcToArrow.sqlToArrowVectorIterator; import static org.apache.arrow.adapter.jdbc.JdbcToArrowUtils.jdbcToArrowSchema; import static org.apache.arrow.flight.sql.impl.FlightSql.*; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCrossReference; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetDbSchemas; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetExportedKeys; import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetImportedKeys; import static org.apache.arrow.flight.sql.impl.FlightSql.DoPutUpdateResult; import static org.apache.arrow.flight.sql.impl.FlightSql.TicketStatementQuery; import static org.apache.arrow.util.Preconditions.checkState; import static org.slf4j.LoggerFactory.getLogger; import com.google.common.cache.Cache; import com.google.common.cache.CacheBuilder; import com.google.common.cache.RemovalListener; import com.google.common.cache.RemovalNotification; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.protobuf.ByteString; import com.google.protobuf.Message; import com.google.protobuf.ProtocolStringList; import java.io.ByteArrayOutputStream; import java.io.File; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.channels.Channels; import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.NoSuchFileException; import java.nio.file.Path; import java.nio.file.Paths; import java.nio.file.StandardOpenOption; import java.sql.Connection; import java.sql.DatabaseMetaData; import java.sql.DriverManager; import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.SQLSyntaxErrorException; import java.sql.Statement; import java.util.ArrayList; import java.util.Arrays; import java.util.Calendar; import java.util.Comparator; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Map.Entry; import java.util.Objects; import java.util.Properties; import java.util.Set; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.function.BiConsumer; import java.util.function.Consumer; import java.util.function.Predicate; import java.util.stream.Collectors; import java.util.stream.Stream; import org.apache.arrow.adapter.jdbc.ArrowVectorIterator; import org.apache.arrow.adapter.jdbc.JdbcFieldInfo; import org.apache.arrow.adapter.jdbc.JdbcParameterBinder; import org.apache.arrow.adapter.jdbc.JdbcToArrowUtils; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.Criteria; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.PutResult; import org.apache.arrow.flight.Result; import org.apache.arrow.flight.SchemaResult; import org.apache.arrow.flight.Ticket; import org.apache.arrow.flight.sql.FlightSqlColumnMetadata; import org.apache.arrow.flight.sql.FlightSqlProducer; import org.apache.arrow.flight.sql.SqlInfoBuilder; import org.apache.arrow.flight.sql.impl.FlightSql.ActionClosePreparedStatementRequest; import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementRequest; import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementResult; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCatalogs; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetPrimaryKeys; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetSqlInfo; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTableTypes; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTables; import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementQuery; import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementUpdate; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions.TableExistsOption; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions.TableNotExistOption; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementQuery; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementUpdate; import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedCaseSensitivity; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.UInt1Vector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.impl.UnionListWriter; import org.apache.arrow.vector.ipc.WriteChannel; import org.apache.arrow.vector.ipc.message.MessageSerializer; 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.types.pojo.Schema; import org.apache.arrow.vector.util.Text; import org.apache.commons.dbcp2.ConnectionFactory; import org.apache.commons.dbcp2.DriverManagerConnectionFactory; import org.apache.commons.dbcp2.PoolableConnection; import org.apache.commons.dbcp2.PoolableConnectionFactory; import org.apache.commons.dbcp2.PoolingDataSource; import org.apache.commons.pool2.ObjectPool; import org.apache.commons.pool2.impl.GenericObjectPool; import org.apache.commons.text.StringEscapeUtils; import org.slf4j.Logger; /** * Example {@link FlightSqlProducer} implementation showing an Apache Derby backed Flight SQL server * that generally supports all current features of Flight SQL. */ public class FlightSqlExample implements FlightSqlProducer, AutoCloseable { private static final Logger LOGGER = getLogger(FlightSqlExample.class); protected static final Calendar DEFAULT_CALENDAR = JdbcToArrowUtils.getUtcCalendar(); public static final String DB_NAME = "derbyDB"; private final String databaseUri; // ARROW-15315: Use ExecutorService to simulate an async scenario private final ExecutorService executorService = Executors.newFixedThreadPool(10); private final Location location; protected final PoolingDataSource dataSource; protected final BufferAllocator rootAllocator = new RootAllocator(); private final Cache> preparedStatementLoadingCache; private final Cache> statementLoadingCache; private final SqlInfoBuilder sqlInfoBuilder; public static void main(String[] args) throws Exception { Location location = Location.forGrpcInsecure("localhost", 55555); final FlightSqlExample example = new FlightSqlExample(location, DB_NAME); Location listenLocation = Location.forGrpcInsecure("0.0.0.0", 55555); try (final BufferAllocator allocator = new RootAllocator(); final FlightServer server = FlightServer.builder(allocator, listenLocation, example).build()) { server.start(); server.awaitTermination(); } } public FlightSqlExample(final Location location, final String dbName) { // TODO Constructor should not be doing work. checkState(removeDerbyDatabaseIfExists(dbName), "Failed to clear Derby database!"); checkState(populateDerbyDatabase(dbName), "Failed to populate Derby database!"); databaseUri = "jdbc:derby:target/" + dbName; final ConnectionFactory connectionFactory = new DriverManagerConnectionFactory(databaseUri, new Properties()); final PoolableConnectionFactory poolableConnectionFactory = new PoolableConnectionFactory(connectionFactory, null); final ObjectPool connectionPool = new GenericObjectPool<>(poolableConnectionFactory); poolableConnectionFactory.setPool(connectionPool); // PoolingDataSource takes ownership of `connectionPool` dataSource = new PoolingDataSource<>(connectionPool); preparedStatementLoadingCache = CacheBuilder.newBuilder() .maximumSize(100) .expireAfterWrite(10, TimeUnit.MINUTES) .removalListener(new StatementRemovalListener()) .build(); statementLoadingCache = CacheBuilder.newBuilder() .maximumSize(100) .expireAfterWrite(10, TimeUnit.MINUTES) .removalListener(new StatementRemovalListener<>()) .build(); this.location = location; sqlInfoBuilder = new SqlInfoBuilder(); try (final Connection connection = dataSource.getConnection()) { final DatabaseMetaData metaData = connection.getMetaData(); sqlInfoBuilder .withFlightSqlServerName(metaData.getDatabaseProductName()) .withFlightSqlServerVersion(metaData.getDatabaseProductVersion()) .withFlightSqlServerArrowVersion(metaData.getDriverVersion()) .withFlightSqlServerReadOnly(metaData.isReadOnly()) .withFlightSqlServerSql(true) .withFlightSqlServerSubstrait(false) .withFlightSqlServerTransaction(SqlSupportedTransaction.SQL_SUPPORTED_TRANSACTION_NONE) .withSqlIdentifierQuoteChar(metaData.getIdentifierQuoteString()) .withSqlDdlCatalog(metaData.supportsCatalogsInDataManipulation()) .withSqlDdlSchema(metaData.supportsSchemasInDataManipulation()) .withSqlDdlTable(metaData.allTablesAreSelectable()) .withSqlIdentifierCase( metaData.storesMixedCaseIdentifiers() ? SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_CASE_INSENSITIVE : metaData.storesUpperCaseIdentifiers() ? SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_UPPERCASE : metaData.storesLowerCaseIdentifiers() ? SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_LOWERCASE : SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_UNKNOWN) .withSqlQuotedIdentifierCase( metaData.storesMixedCaseQuotedIdentifiers() ? SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_CASE_INSENSITIVE : metaData.storesUpperCaseQuotedIdentifiers() ? SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_UPPERCASE : metaData.storesLowerCaseQuotedIdentifiers() ? SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_LOWERCASE : SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_UNKNOWN) .withSqlAllTablesAreSelectable(true) .withSqlNullOrdering(SqlNullOrdering.SQL_NULLS_SORTED_AT_END) .withSqlMaxColumnsInTable(42) .withFlightSqlServerBulkIngestion(true) .withFlightSqlServerBulkIngestionTransaction(false); } catch (SQLException e) { throw new RuntimeException(e); } } public static boolean removeDerbyDatabaseIfExists(final String dbName) { final Path path = Paths.get("target" + File.separator + dbName); try (final Stream walk = Files.walk(path)) { /* * Iterate over all paths to delete, mapping each path to the outcome of its own * deletion as a boolean representing whether each individual operation was * successful; then reduce all booleans into a single answer. * If for whatever reason the resulting `Stream` is empty, throw an `IOException`; * this not expected. */ boolean unused = walk.sorted(Comparator.reverseOrder()) .map(Path::toFile) .map(File::delete) .reduce(Boolean::logicalAnd) .orElseThrow(IOException::new); } catch (NoSuchFileException e) { /* * The only acceptable scenario for an `IOException` to be thrown here is if * an attempt to delete an non-existing file takes place -- which should be * alright, since they would be deleted anyway. */ LOGGER.error(format("No existing Derby database to delete.: <%s>", e.getMessage()), e); return true; } catch (Exception e) { LOGGER.error(format("Failed attempt to clear DerbyDB.: <%s>", e.getMessage()), e); return false; } return true; } private static boolean populateDerbyDatabase(final String dbName) { try (final Connection connection = DriverManager.getConnection("jdbc:derby:target/" + dbName + ";create=true"); Statement statement = connection.createStatement()) { dropTable(statement, "intTable"); dropTable(statement, "foreignTable"); statement.execute( "CREATE TABLE foreignTable (" + "id INT not null primary key GENERATED ALWAYS AS IDENTITY (START WITH 1, INCREMENT BY 1), " + "foreignName varchar(100), " + "value int)"); statement.execute( "CREATE TABLE intTable (" + "id INT not null primary key GENERATED ALWAYS AS IDENTITY (START WITH 1, INCREMENT BY 1), " + "keyName varchar(100), " + "value int, " + "foreignId int references foreignTable(id))"); statement.execute("INSERT INTO foreignTable (foreignName, value) VALUES ('keyOne', 1)"); statement.execute("INSERT INTO foreignTable (foreignName, value) VALUES ('keyTwo', 0)"); statement.execute("INSERT INTO foreignTable (foreignName, value) VALUES ('keyThree', -1)"); statement.execute("INSERT INTO intTable (keyName, value, foreignId) VALUES ('one', 1, 1)"); statement.execute("INSERT INTO intTable (keyName, value, foreignId) VALUES ('zero', 0, 1)"); statement.execute( "INSERT INTO intTable (keyName, value, foreignId) VALUES ('negative one', -1, 1)"); } catch (final SQLException e) { LOGGER.error(format("Failed attempt to populate DerbyDB: <%s>", e.getMessage()), e); return false; } return true; } private static void dropTable(final Statement statement, final String tableName) throws SQLException { try { statement.execute("DROP TABLE " + tableName); } catch (SQLException e) { // sql error code for "object does not exist"; which is fine, we're trying to delete the table // see https://db.apache.org/derby/docs/10.17/ref/rrefexcept71493.html if (!"42Y55".equals(e.getSQLState())) { throw e; } } } private static ArrowType getArrowTypeFromJdbcType( final int jdbcDataType, final int precision, final int scale) { try { return JdbcToArrowUtils.getArrowTypeFromJdbcType( new JdbcFieldInfo(jdbcDataType, precision, scale), DEFAULT_CALENDAR); } catch (UnsupportedOperationException ignored) { return ArrowType.Utf8.INSTANCE; } } private static void saveToVector(final Byte data, final UInt1Vector vector, final int index) { vectorConsumer( data, vector, fieldVector -> fieldVector.setNull(index), (theData, fieldVector) -> fieldVector.setSafe(index, theData)); } private static void saveToVector(final Byte data, final BitVector vector, final int index) { vectorConsumer( data, vector, fieldVector -> fieldVector.setNull(index), (theData, fieldVector) -> fieldVector.setSafe(index, theData)); } private static void saveToVector(final String data, final VarCharVector vector, final int index) { preconditionCheckSaveToVector(vector, index); vectorConsumer( data, vector, fieldVector -> fieldVector.setNull(index), (theData, fieldVector) -> fieldVector.setSafe(index, new Text(theData))); } private static void saveToVector(final Integer data, final IntVector vector, final int index) { preconditionCheckSaveToVector(vector, index); vectorConsumer( data, vector, fieldVector -> fieldVector.setNull(index), (theData, fieldVector) -> fieldVector.setSafe(index, theData)); } private static void saveToVector( final byte[] data, final VarBinaryVector vector, final int index) { preconditionCheckSaveToVector(vector, index); vectorConsumer( data, vector, fieldVector -> fieldVector.setNull(index), (theData, fieldVector) -> fieldVector.setSafe(index, theData)); } private static void preconditionCheckSaveToVector(final FieldVector vector, final int index) { Objects.requireNonNull(vector, "vector cannot be null."); checkState(index >= 0, "Index must be a positive number!"); } private static void vectorConsumer( final T data, final V vector, final Consumer consumerIfNullable, final BiConsumer defaultConsumer) { if (isNull(data)) { consumerIfNullable.accept(vector); return; } defaultConsumer.accept(data, vector); } private static VectorSchemaRoot getSchemasRoot( final ResultSet data, final BufferAllocator allocator) throws SQLException { final VarCharVector catalogs = new VarCharVector("catalog_name", allocator); final VarCharVector schemas = new VarCharVector( "db_schema_name", FieldType.notNullable(MinorType.VARCHAR.getType()), allocator); final List vectors = ImmutableList.of(catalogs, schemas); vectors.forEach(FieldVector::allocateNew); final Map vectorToColumnName = ImmutableMap.of( catalogs, "TABLE_CATALOG", schemas, "TABLE_SCHEM"); saveToVectors(vectorToColumnName, data); final int rows = vectors.stream() .map(FieldVector::getValueCount) .findAny() .orElseThrow(IllegalStateException::new); vectors.forEach(vector -> vector.setValueCount(rows)); return new VectorSchemaRoot(vectors); } private static int saveToVectors( final Map vectorToColumnName, final ResultSet data, boolean emptyToNull) throws SQLException { Predicate alwaysTrue = (resultSet) -> true; return saveToVectors(vectorToColumnName, data, emptyToNull, alwaysTrue); } @SuppressWarnings("StringSplitter") private static int saveToVectors( final Map vectorToColumnName, final ResultSet data, boolean emptyToNull, Predicate resultSetPredicate) throws SQLException { Objects.requireNonNull(vectorToColumnName, "vectorToColumnName cannot be null."); Objects.requireNonNull(data, "data cannot be null."); final Set> entrySet = vectorToColumnName.entrySet(); int rows = 0; while (data.next()) { if (!resultSetPredicate.test(data)) { continue; } for (final Entry vectorToColumn : entrySet) { final T vector = vectorToColumn.getKey(); final String columnName = vectorToColumn.getValue(); if (vector instanceof VarCharVector) { String thisData = data.getString(columnName); saveToVector( emptyToNull ? emptyToNull(thisData) : thisData, (VarCharVector) vector, rows); } else if (vector instanceof IntVector) { final int intValue = data.getInt(columnName); saveToVector(data.wasNull() ? null : intValue, (IntVector) vector, rows); } else if (vector instanceof UInt1Vector) { final byte byteValue = data.getByte(columnName); saveToVector(data.wasNull() ? null : byteValue, (UInt1Vector) vector, rows); } else if (vector instanceof BitVector) { final byte byteValue = data.getByte(columnName); saveToVector(data.wasNull() ? null : byteValue, (BitVector) vector, rows); } else if (vector instanceof ListVector) { String createParamsValues = data.getString(columnName); UnionListWriter writer = ((ListVector) vector).getWriter(); BufferAllocator allocator = vector.getAllocator(); final ArrowBuf buf = allocator.buffer(1024); writer.setPosition(rows); writer.startList(); if (createParamsValues != null) { String[] split = createParamsValues.split(","); range(0, split.length) .forEach( i -> { byte[] bytes = split[i].getBytes(StandardCharsets.UTF_8); Preconditions.checkState( bytes.length < 1024, "The amount of bytes is greater than what the ArrowBuf supports"); buf.setBytes(0, bytes); writer.varChar().writeVarChar(0, bytes.length, buf); }); } buf.close(); writer.endList(); } else { throw CallStatus.INVALID_ARGUMENT .withDescription("Provided vector not supported") .toRuntimeException(); } } rows++; } for (final Entry vectorToColumn : entrySet) { vectorToColumn.getKey().setValueCount(rows); } return rows; } private static void saveToVectors( final Map vectorToColumnName, final ResultSet data) throws SQLException { saveToVectors(vectorToColumnName, data, false); } private static VectorSchemaRoot getTableTypesRoot( final ResultSet data, final BufferAllocator allocator) throws SQLException { return getRoot(data, allocator, "table_type", "TABLE_TYPE"); } private static VectorSchemaRoot getCatalogsRoot( final ResultSet data, final BufferAllocator allocator) throws SQLException { return getRoot(data, allocator, "catalog_name", "TABLE_CATALOG"); } private static VectorSchemaRoot getRoot( final ResultSet data, final BufferAllocator allocator, final String fieldVectorName, final String columnName) throws SQLException { final VarCharVector dataVector = new VarCharVector( fieldVectorName, FieldType.notNullable(MinorType.VARCHAR.getType()), allocator); saveToVectors(ImmutableMap.of(dataVector, columnName), data); final int rows = dataVector.getValueCount(); dataVector.setValueCount(rows); return new VectorSchemaRoot(singletonList(dataVector)); } private static VectorSchemaRoot getTypeInfoRoot( CommandGetXdbcTypeInfo request, ResultSet typeInfo, final BufferAllocator allocator) throws SQLException { Preconditions.checkNotNull(allocator, "BufferAllocator cannot be null."); VectorSchemaRoot root = VectorSchemaRoot.create(Schemas.GET_TYPE_INFO_SCHEMA, allocator); Map mapper = new HashMap<>(); mapper.put(root.getVector("type_name"), "TYPE_NAME"); mapper.put(root.getVector("data_type"), "DATA_TYPE"); mapper.put(root.getVector("column_size"), "PRECISION"); mapper.put(root.getVector("literal_prefix"), "LITERAL_PREFIX"); mapper.put(root.getVector("literal_suffix"), "LITERAL_SUFFIX"); mapper.put(root.getVector("create_params"), "CREATE_PARAMS"); mapper.put(root.getVector("nullable"), "NULLABLE"); mapper.put(root.getVector("case_sensitive"), "CASE_SENSITIVE"); mapper.put(root.getVector("searchable"), "SEARCHABLE"); mapper.put(root.getVector("unsigned_attribute"), "UNSIGNED_ATTRIBUTE"); mapper.put(root.getVector("fixed_prec_scale"), "FIXED_PREC_SCALE"); mapper.put(root.getVector("auto_increment"), "AUTO_INCREMENT"); mapper.put(root.getVector("local_type_name"), "LOCAL_TYPE_NAME"); mapper.put(root.getVector("minimum_scale"), "MINIMUM_SCALE"); mapper.put(root.getVector("maximum_scale"), "MAXIMUM_SCALE"); mapper.put(root.getVector("sql_data_type"), "SQL_DATA_TYPE"); mapper.put(root.getVector("datetime_subcode"), "SQL_DATETIME_SUB"); mapper.put(root.getVector("num_prec_radix"), "NUM_PREC_RADIX"); Predicate predicate; if (request.hasDataType()) { predicate = (resultSet) -> { try { return resultSet.getInt("DATA_TYPE") == request.getDataType(); } catch (SQLException e) { throw new RuntimeException(e); } }; } else { predicate = resultSet -> true; } int rows = saveToVectors(mapper, typeInfo, true, predicate); root.setRowCount(rows); return root; } private static VectorSchemaRoot getTablesRoot( final DatabaseMetaData databaseMetaData, final BufferAllocator allocator, final boolean includeSchema, final String catalog, final String schemaFilterPattern, final String tableFilterPattern, final String... tableTypes) throws SQLException, IOException { /* * TODO Fix DerbyDB inconsistency if possible. * During the early development of this prototype, an inconsistency has been found in the database * used for this demonstration; as DerbyDB does not operate with the concept of catalogs, fetching * the catalog name for a given table from `DatabaseMetadata#getColumns` and `DatabaseMetadata#getSchemas` * returns null, as expected. However, the inconsistency lies in the fact that accessing the same * information -- that is, the catalog name for a given table -- from `DatabaseMetadata#getSchemas` * returns an empty String.The temporary workaround for this was making sure we convert the empty Strings * to null using `com.google.common.base.Strings#emptyToNull`. */ Objects.requireNonNull(allocator, "BufferAllocator cannot be null."); final VarCharVector catalogNameVector = new VarCharVector("catalog_name", allocator); final VarCharVector schemaNameVector = new VarCharVector("db_schema_name", allocator); final VarCharVector tableNameVector = new VarCharVector( "table_name", FieldType.notNullable(MinorType.VARCHAR.getType()), allocator); final VarCharVector tableTypeVector = new VarCharVector( "table_type", FieldType.notNullable(MinorType.VARCHAR.getType()), allocator); final List vectors = new ArrayList<>(4); vectors.add(catalogNameVector); vectors.add(schemaNameVector); vectors.add(tableNameVector); vectors.add(tableTypeVector); vectors.forEach(FieldVector::allocateNew); final Map vectorToColumnName = ImmutableMap.of( catalogNameVector, "TABLE_CAT", schemaNameVector, "TABLE_SCHEM", tableNameVector, "TABLE_NAME", tableTypeVector, "TABLE_TYPE"); try (final ResultSet data = Objects.requireNonNull( databaseMetaData, format("%s cannot be null.", databaseMetaData.getClass().getName())) .getTables(catalog, schemaFilterPattern, tableFilterPattern, tableTypes)) { saveToVectors(vectorToColumnName, data, true); final int rows = vectors.stream() .map(FieldVector::getValueCount) .findAny() .orElseThrow(IllegalStateException::new); vectors.forEach(vector -> vector.setValueCount(rows)); if (includeSchema) { final VarBinaryVector tableSchemaVector = new VarBinaryVector( "table_schema", FieldType.notNullable(MinorType.VARBINARY.getType()), allocator); tableSchemaVector.allocateNew(rows); try (final ResultSet columnsData = databaseMetaData.getColumns(catalog, schemaFilterPattern, tableFilterPattern, null)) { final Map> tableToFields = new HashMap<>(); while (columnsData.next()) { final String catalogName = columnsData.getString("TABLE_CAT"); final String schemaName = columnsData.getString("TABLE_SCHEM"); final String tableName = columnsData.getString("TABLE_NAME"); final String typeName = columnsData.getString("TYPE_NAME"); final String fieldName = columnsData.getString("COLUMN_NAME"); final int dataType = columnsData.getInt("DATA_TYPE"); final boolean isNullable = columnsData.getInt("NULLABLE") != DatabaseMetaData.columnNoNulls; final int precision = columnsData.getInt("COLUMN_SIZE"); final int scale = columnsData.getInt("DECIMAL_DIGITS"); boolean isAutoIncrement = Objects.equals(columnsData.getString("IS_AUTOINCREMENT"), "YES"); final List fields = tableToFields.computeIfAbsent(tableName, tableName_ -> new ArrayList<>()); final FlightSqlColumnMetadata columnMetadata = new FlightSqlColumnMetadata.Builder() .catalogName(catalogName) .schemaName(schemaName) .tableName(tableName) .typeName(typeName) .precision(precision) .scale(scale) .isAutoIncrement(isAutoIncrement) .build(); final Field field = new Field( fieldName, new FieldType( isNullable, getArrowTypeFromJdbcType(dataType, precision, scale), null, columnMetadata.getMetadataMap()), null); fields.add(field); } for (int index = 0; index < rows; index++) { final String tableName = tableNameVector.getObject(index).toString(); final Schema schema = new Schema(tableToFields.get(tableName)); saveToVector( copyFrom(serializeMetadata(schema)).toByteArray(), tableSchemaVector, index); } } tableSchemaVector.setValueCount(rows); vectors.add(tableSchemaVector); } } return new VectorSchemaRoot(vectors); } private static ByteBuffer serializeMetadata(final Schema schema) { final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); try { MessageSerializer.serialize(new WriteChannel(Channels.newChannel(outputStream)), schema); return ByteBuffer.wrap(outputStream.toByteArray()); } catch (final IOException e) { throw new RuntimeException("Failed to serialize schema", e); } } private static String getRootAsCSVNoHeader(final VectorSchemaRoot root) { StringBuilder sb = new StringBuilder(); Schema schema = root.getSchema(); int rowCount = root.getRowCount(); List fieldVectors = root.getFieldVectors(); List row = new ArrayList<>(schema.getFields().size()); for (int i = 0; i < rowCount; i++) { if (i > 0) { sb.append("\n"); } row.clear(); for (FieldVector v : fieldVectors) { row.add(v.getObject(i)); } printRowAsCSV(sb, row); } return sb.toString(); } private static void printRowAsCSV(StringBuilder sb, List values) { sb.append( values.stream() .map(v -> isNull(v) ? "" : v.toString()) .map(StringEscapeUtils::escapeCsv) .collect(Collectors.joining(","))); } @Override public void getStreamPreparedStatement( final CommandPreparedStatementQuery command, final CallContext context, final ServerStreamListener listener) { final ByteString handle = command.getPreparedStatementHandle(); StatementContext statementContext = preparedStatementLoadingCache.getIfPresent(handle); Objects.requireNonNull(statementContext); final PreparedStatement statement = statementContext.getStatement(); try (final ResultSet resultSet = statement.executeQuery()) { final Schema schema = jdbcToArrowSchema(resultSet.getMetaData(), DEFAULT_CALENDAR); try (final VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.create(schema, rootAllocator)) { final VectorLoader loader = new VectorLoader(vectorSchemaRoot); listener.start(vectorSchemaRoot); final ArrowVectorIterator iterator = sqlToArrowVectorIterator(resultSet, rootAllocator); while (iterator.hasNext()) { final VectorSchemaRoot batch = iterator.next(); if (batch.getRowCount() == 0) { break; } final VectorUnloader unloader = new VectorUnloader(batch); loader.load(unloader.getRecordBatch()); listener.putNext(); vectorSchemaRoot.clear(); } listener.putNext(); } } catch (final SQLException | IOException e) { LOGGER.error(format("Failed to getStreamPreparedStatement: <%s>.", e.getMessage()), e); listener.error( CallStatus.INTERNAL .withDescription("Failed to prepare statement: " + e) .toRuntimeException()); } finally { listener.completed(); } } @Override public void closePreparedStatement( final ActionClosePreparedStatementRequest request, final CallContext context, final StreamListener listener) { // Running on another thread Future unused = executorService.submit( () -> { try { preparedStatementLoadingCache.invalidate(request.getPreparedStatementHandle()); } catch (final Exception e) { listener.onError(e); return; } listener.onCompleted(); }); } @Override public FlightInfo getFlightInfoStatement( final CommandStatementQuery request, final CallContext context, final FlightDescriptor descriptor) { ByteString handle = copyFrom(randomUUID().toString().getBytes(StandardCharsets.UTF_8)); try { // Ownership of the connection will be passed to the context. Do NOT close! final Connection connection = dataSource.getConnection(); final Statement statement = connection.createStatement(ResultSet.TYPE_SCROLL_INSENSITIVE, ResultSet.CONCUR_READ_ONLY); final String query = request.getQuery(); final StatementContext statementContext = new StatementContext<>(statement, query); statementLoadingCache.put(handle, statementContext); final ResultSet resultSet = statement.executeQuery(query); TicketStatementQuery ticket = TicketStatementQuery.newBuilder().setStatementHandle(handle).build(); return getFlightInfoForSchema( ticket, descriptor, jdbcToArrowSchema(resultSet.getMetaData(), DEFAULT_CALENDAR)); } catch (final SQLException e) { LOGGER.error( format("There was a problem executing the prepared statement: <%s>.", e.getMessage()), e); throw CallStatus.INTERNAL.withCause(e).toRuntimeException(); } } @Override public FlightInfo getFlightInfoPreparedStatement( final CommandPreparedStatementQuery command, final CallContext context, final FlightDescriptor descriptor) { final ByteString preparedStatementHandle = command.getPreparedStatementHandle(); StatementContext statementContext = preparedStatementLoadingCache.getIfPresent(preparedStatementHandle); try { assert statementContext != null; PreparedStatement statement = statementContext.getStatement(); ResultSetMetaData metaData = statement.getMetaData(); return getFlightInfoForSchema( command, descriptor, jdbcToArrowSchema(metaData, DEFAULT_CALENDAR)); } catch (final SQLException e) { LOGGER.error( format("There was a problem executing the prepared statement: <%s>.", e.getMessage()), e); throw CallStatus.INTERNAL.withCause(e).toRuntimeException(); } } @Override public SchemaResult getSchemaStatement( final CommandStatementQuery command, final CallContext context, final FlightDescriptor descriptor) { throw CallStatus.UNIMPLEMENTED.toRuntimeException(); } @Override public void close() throws Exception { try { preparedStatementLoadingCache.cleanUp(); } catch (Throwable t) { LOGGER.error(format("Failed to close resources: <%s>", t.getMessage()), t); } AutoCloseables.close(dataSource, rootAllocator); } @Override public void listFlights( CallContext context, Criteria criteria, StreamListener listener) { // TODO - build example implementation throw CallStatus.UNIMPLEMENTED.toRuntimeException(); } @Override public void createPreparedStatement( final ActionCreatePreparedStatementRequest request, final CallContext context, final StreamListener listener) { // Running on another thread Future unused = executorService.submit( () -> { try { final ByteString preparedStatementHandle = copyFrom(request.getQuery().getBytes(StandardCharsets.UTF_8)); // Ownership of the connection will be passed to the context. Do NOT close! final Connection connection = dataSource.getConnection(); final PreparedStatement preparedStatement = connection.prepareStatement( request.getQuery(), ResultSet.TYPE_SCROLL_INSENSITIVE, ResultSet.CONCUR_READ_ONLY); final StatementContext preparedStatementContext = new StatementContext<>(preparedStatement, request.getQuery()); preparedStatementLoadingCache.put( preparedStatementHandle, preparedStatementContext); final Schema parameterSchema = jdbcToArrowSchema(preparedStatement.getParameterMetaData(), DEFAULT_CALENDAR); final ResultSetMetaData metaData = preparedStatement.getMetaData(); final ByteString bytes = isNull(metaData) ? ByteString.EMPTY : ByteString.copyFrom( serializeMetadata(jdbcToArrowSchema(metaData, DEFAULT_CALENDAR))); final ActionCreatePreparedStatementResult result = ActionCreatePreparedStatementResult.newBuilder() .setDatasetSchema(bytes) .setParameterSchema(copyFrom(serializeMetadata(parameterSchema))) .setPreparedStatementHandle(preparedStatementHandle) .build(); listener.onNext(new Result(pack(result).toByteArray())); } catch (final SQLException e) { listener.onError( CallStatus.INTERNAL .withDescription("Failed to create prepared statement: " + e) .toRuntimeException()); return; } catch (final Throwable t) { listener.onError( CallStatus.INTERNAL .withDescription("Unknown error: " + t) .toRuntimeException()); return; } listener.onCompleted(); }); } @Override public void doExchange(CallContext context, FlightStream reader, ServerStreamListener writer) { // TODO - build example implementation throw CallStatus.UNIMPLEMENTED.toRuntimeException(); } @Override public Runnable acceptPutStatement( CommandStatementUpdate command, CallContext context, FlightStream flightStream, StreamListener ackStream) { final String query = command.getQuery(); return () -> { try (final Connection connection = dataSource.getConnection(); final Statement statement = connection.createStatement()) { final int result = statement.executeUpdate(query); final DoPutUpdateResult build = DoPutUpdateResult.newBuilder().setRecordCount(result).build(); try (final ArrowBuf buffer = rootAllocator.buffer(build.getSerializedSize())) { buffer.writeBytes(build.toByteArray()); ackStream.onNext(PutResult.metadata(buffer)); ackStream.onCompleted(); } } catch (SQLSyntaxErrorException e) { ackStream.onError( CallStatus.INVALID_ARGUMENT .withDescription("Failed to execute statement (invalid syntax): " + e) .toRuntimeException()); } catch (SQLException e) { ackStream.onError( CallStatus.INTERNAL .withDescription("Failed to execute statement: " + e) .toRuntimeException()); } }; } @Override public Runnable acceptPutStatementBulkIngest( CommandStatementIngest command, CallContext context, FlightStream flightStream, StreamListener ackStream) { final String schema = command.hasSchema() ? command.getSchema() : null; final String table = command.getTable(); final boolean temporary = command.getTemporary(); final boolean transactionId = command.hasTransactionId(); final TableDefinitionOptions tableDefinitionOptions = command.hasTableDefinitionOptions() ? command.getTableDefinitionOptions() : null; return () -> { TableExistsOption ifExists = TableExistsOption.TABLE_EXISTS_OPTION_APPEND; if (temporary) { ackStream.onError( CallStatus.UNIMPLEMENTED .withDescription("Bulk ingestion using temporary tables is not supported") .toRuntimeException()); } else if (transactionId) { ackStream.onError( CallStatus.UNIMPLEMENTED .withDescription( "Bulk ingestion automatically happens in a transaction. Specifying explicit transaction is not supported.") .toRuntimeException()); } else if (isNull(tableDefinitionOptions)) { ackStream.onError( CallStatus.INVALID_ARGUMENT .withDescription("TableDefinitionOptions not provided.") .toRuntimeException()); } else { TableNotExistOption ifNotExist = tableDefinitionOptions.getIfNotExist(); ifExists = tableDefinitionOptions.getIfExists(); if (!TableNotExistOption.TABLE_NOT_EXIST_OPTION_FAIL.equals(ifNotExist)) { ackStream.onError( CallStatus.UNIMPLEMENTED .withDescription( "Only supported option is TABLE_NOT_EXIST_OPTION_FAIL for TableNotExistsOption.") .toRuntimeException()); } else if (TableExistsOption.TABLE_EXISTS_OPTION_UNSPECIFIED.equals(ifExists)) { ackStream.onError( CallStatus.INVALID_ARGUMENT .withDescription("TableExistsOption must be specified") .toRuntimeException()); } else if (TableExistsOption.TABLE_EXISTS_OPTION_FAIL.equals(ifExists)) { ackStream.onError( CallStatus.UNIMPLEMENTED .withDescription("TABLE_EXISTS_OPTION_FAIL is not supported.") .toRuntimeException()); } } Path tempFile = null; try { tempFile = Files.createTempFile(null, null); VectorSchemaRoot root = null; int counter = 0; while (flightStream.next()) { if (counter > 0) { Files.writeString(tempFile, "\n", StandardCharsets.UTF_8, StandardOpenOption.APPEND); } counter += 1; root = flightStream.getRoot(); Files.writeString( tempFile, getRootAsCSVNoHeader(root), StandardCharsets.UTF_8, StandardOpenOption.APPEND); } if (counter > 0) { Files.writeString(tempFile, "\n", StandardCharsets.UTF_8, StandardOpenOption.APPEND); } if (!isNull(root)) { String header = root.getSchema().getFields().stream() .map(Field::getName) .collect(Collectors.joining(",")); try (final Connection connection = dataSource.getConnection(); final PreparedStatement preparedStatement = connection.prepareStatement( "CALL SYSCS_UTIL.SYSCS_IMPORT_DATA (?,?,?,null,?,?,?,?,?)")) { preparedStatement.setString(1, schema); preparedStatement.setString(2, table); preparedStatement.setString(3, header); preparedStatement.setString(4, tempFile.toString()); preparedStatement.setString(5, ","); preparedStatement.setString(6, "\""); preparedStatement.setString(7, "UTF-8"); preparedStatement.setInt( 8, TableExistsOption.TABLE_EXISTS_OPTION_REPLACE.equals(ifExists) ? 1 : 0); preparedStatement.execute(); final DoPutUpdateResult build = DoPutUpdateResult.newBuilder().setRecordCount(-1).build(); try (final ArrowBuf buffer = rootAllocator.buffer(build.getSerializedSize())) { buffer.writeBytes(build.toByteArray()); ackStream.onNext(PutResult.metadata(buffer)); ackStream.onCompleted(); } } catch (SQLException e) { ackStream.onError( CallStatus.INTERNAL .withDescription("Failed to execute bulk ingest: " + e) .toRuntimeException()); } } } catch (IOException e) { ackStream.onError( CallStatus.INTERNAL .withDescription("Failed to create temp file for bulk loading: " + e) .toRuntimeException()); } finally { if (!isNull(tempFile)) { try { Files.delete(tempFile); } catch (IOException e) { // } } } }; } @Override public Runnable acceptPutPreparedStatementUpdate( CommandPreparedStatementUpdate command, CallContext context, FlightStream flightStream, StreamListener ackStream) { final StatementContext statement = preparedStatementLoadingCache.getIfPresent(command.getPreparedStatementHandle()); return () -> { if (statement == null) { ackStream.onError( CallStatus.NOT_FOUND .withDescription("Prepared statement does not exist") .toRuntimeException()); return; } try { final PreparedStatement preparedStatement = statement.getStatement(); while (flightStream.next()) { final VectorSchemaRoot root = flightStream.getRoot(); final int rowCount = root.getRowCount(); final int recordCount; if (rowCount == 0) { preparedStatement.execute(); recordCount = preparedStatement.getUpdateCount(); } else { final JdbcParameterBinder binder = JdbcParameterBinder.builder(preparedStatement, root).bindAll().build(); while (binder.next()) { preparedStatement.addBatch(); } final int[] recordCounts = preparedStatement.executeBatch(); recordCount = Arrays.stream(recordCounts).sum(); } final DoPutUpdateResult build = DoPutUpdateResult.newBuilder().setRecordCount(recordCount).build(); try (final ArrowBuf buffer = rootAllocator.buffer(build.getSerializedSize())) { buffer.writeBytes(build.toByteArray()); ackStream.onNext(PutResult.metadata(buffer)); } } } catch (SQLException e) { ackStream.onError( CallStatus.INTERNAL .withDescription("Failed to execute update: " + e) .toRuntimeException()); return; } ackStream.onCompleted(); }; } @Override public Runnable acceptPutPreparedStatementQuery( CommandPreparedStatementQuery command, CallContext context, FlightStream flightStream, StreamListener ackStream) { final StatementContext statementContext = preparedStatementLoadingCache.getIfPresent(command.getPreparedStatementHandle()); return () -> { assert statementContext != null; PreparedStatement preparedStatement = statementContext.getStatement(); try { while (flightStream.next()) { final VectorSchemaRoot root = flightStream.getRoot(); final JdbcParameterBinder binder = JdbcParameterBinder.builder(preparedStatement, root).bindAll().build(); while (binder.next()) { // Do not execute() - will be done in a getStream call } } } catch (SQLException e) { ackStream.onError( CallStatus.INTERNAL .withDescription("Failed to bind parameters: " + e.getMessage()) .withCause(e) .toRuntimeException()); return; } ackStream.onCompleted(); }; } @Override public FlightInfo getFlightInfoSqlInfo( final CommandGetSqlInfo request, final CallContext context, final FlightDescriptor descriptor) { return getFlightInfoForSchema(request, descriptor, Schemas.GET_SQL_INFO_SCHEMA); } @Override public void getStreamSqlInfo( final CommandGetSqlInfo command, final CallContext context, final ServerStreamListener listener) { this.sqlInfoBuilder.send(command.getInfoList(), listener); } @Override public FlightInfo getFlightInfoTypeInfo( CommandGetXdbcTypeInfo request, CallContext context, FlightDescriptor descriptor) { return getFlightInfoForSchema(request, descriptor, Schemas.GET_TYPE_INFO_SCHEMA); } @Override public void getStreamTypeInfo( CommandGetXdbcTypeInfo request, CallContext context, ServerStreamListener listener) { try (final Connection connection = dataSource.getConnection(); final ResultSet typeInfo = connection.getMetaData().getTypeInfo(); final VectorSchemaRoot vectorSchemaRoot = getTypeInfoRoot(request, typeInfo, rootAllocator)) { listener.start(vectorSchemaRoot); listener.putNext(); } catch (SQLException e) { LOGGER.error(format("Failed to getStreamCatalogs: <%s>.", e.getMessage()), e); listener.error(e); } finally { listener.completed(); } } @Override public FlightInfo getFlightInfoCatalogs( final CommandGetCatalogs request, final CallContext context, final FlightDescriptor descriptor) { return getFlightInfoForSchema(request, descriptor, Schemas.GET_CATALOGS_SCHEMA); } @Override public void getStreamCatalogs(final CallContext context, final ServerStreamListener listener) { try (final Connection connection = dataSource.getConnection(); final ResultSet catalogs = connection.getMetaData().getCatalogs(); final VectorSchemaRoot vectorSchemaRoot = getCatalogsRoot(catalogs, rootAllocator)) { listener.start(vectorSchemaRoot); listener.putNext(); } catch (SQLException e) { LOGGER.error(format("Failed to getStreamCatalogs: <%s>.", e.getMessage()), e); listener.error(e); } finally { listener.completed(); } } @Override public FlightInfo getFlightInfoSchemas( final CommandGetDbSchemas request, final CallContext context, final FlightDescriptor descriptor) { return getFlightInfoForSchema(request, descriptor, Schemas.GET_SCHEMAS_SCHEMA); } @Override public void getStreamSchemas( final CommandGetDbSchemas command, final CallContext context, final ServerStreamListener listener) { final String catalog = command.hasCatalog() ? command.getCatalog() : null; final String schemaFilterPattern = command.hasDbSchemaFilterPattern() ? command.getDbSchemaFilterPattern() : null; try (final Connection connection = dataSource.getConnection(); final ResultSet schemas = connection.getMetaData().getSchemas(catalog, schemaFilterPattern); final VectorSchemaRoot vectorSchemaRoot = getSchemasRoot(schemas, rootAllocator)) { listener.start(vectorSchemaRoot); listener.putNext(); } catch (SQLException e) { LOGGER.error(format("Failed to getStreamSchemas: <%s>.", e.getMessage()), e); listener.error(e); } finally { listener.completed(); } } @Override public FlightInfo getFlightInfoTables( final CommandGetTables request, final CallContext context, final FlightDescriptor descriptor) { Schema schemaToUse = Schemas.GET_TABLES_SCHEMA; if (!request.getIncludeSchema()) { schemaToUse = Schemas.GET_TABLES_SCHEMA_NO_SCHEMA; } return getFlightInfoForSchema(request, descriptor, schemaToUse); } @Override public void getStreamTables( final CommandGetTables command, final CallContext context, final ServerStreamListener listener) { final String catalog = command.hasCatalog() ? command.getCatalog() : null; final String schemaFilterPattern = command.hasDbSchemaFilterPattern() ? command.getDbSchemaFilterPattern() : null; final String tableFilterPattern = command.hasTableNameFilterPattern() ? command.getTableNameFilterPattern() : null; final ProtocolStringList protocolStringList = command.getTableTypesList(); final int protocolSize = protocolStringList.size(); final String[] tableTypes = protocolSize == 0 ? null : protocolStringList.toArray(new String[protocolSize]); try (final Connection connection = DriverManager.getConnection(databaseUri); final VectorSchemaRoot vectorSchemaRoot = getTablesRoot( connection.getMetaData(), rootAllocator, command.getIncludeSchema(), catalog, schemaFilterPattern, tableFilterPattern, tableTypes)) { listener.start(vectorSchemaRoot); listener.putNext(); } catch (SQLException | IOException e) { LOGGER.error(format("Failed to getStreamTables: <%s>.", e.getMessage()), e); listener.error(e); } finally { listener.completed(); } } @Override public FlightInfo getFlightInfoTableTypes( final CommandGetTableTypes request, final CallContext context, final FlightDescriptor descriptor) { return getFlightInfoForSchema(request, descriptor, Schemas.GET_TABLE_TYPES_SCHEMA); } @Override public void getStreamTableTypes(final CallContext context, final ServerStreamListener listener) { try (final Connection connection = dataSource.getConnection(); final ResultSet tableTypes = connection.getMetaData().getTableTypes(); final VectorSchemaRoot vectorSchemaRoot = getTableTypesRoot(tableTypes, rootAllocator)) { listener.start(vectorSchemaRoot); listener.putNext(); } catch (SQLException e) { LOGGER.error(format("Failed to getStreamTableTypes: <%s>.", e.getMessage()), e); listener.error(e); } finally { listener.completed(); } } @Override public FlightInfo getFlightInfoPrimaryKeys( final CommandGetPrimaryKeys request, final CallContext context, final FlightDescriptor descriptor) { return getFlightInfoForSchema(request, descriptor, Schemas.GET_PRIMARY_KEYS_SCHEMA); } @Override public void getStreamPrimaryKeys( final CommandGetPrimaryKeys command, final CallContext context, final ServerStreamListener listener) { final String catalog = command.hasCatalog() ? command.getCatalog() : null; final String schema = command.hasDbSchema() ? command.getDbSchema() : null; final String table = command.getTable(); try (Connection connection = DriverManager.getConnection(databaseUri)) { final ResultSet primaryKeys = connection.getMetaData().getPrimaryKeys(catalog, schema, table); final VarCharVector catalogNameVector = new VarCharVector("catalog_name", rootAllocator); final VarCharVector schemaNameVector = new VarCharVector("db_schema_name", rootAllocator); final VarCharVector tableNameVector = new VarCharVector("table_name", rootAllocator); final VarCharVector columnNameVector = new VarCharVector("column_name", rootAllocator); final IntVector keySequenceVector = new IntVector("key_sequence", rootAllocator); final VarCharVector keyNameVector = new VarCharVector("key_name", rootAllocator); final List vectors = new ArrayList<>( ImmutableList.of( catalogNameVector, schemaNameVector, tableNameVector, columnNameVector, keySequenceVector, keyNameVector)); vectors.forEach(FieldVector::allocateNew); int rows = 0; for (; primaryKeys.next(); rows++) { saveToVector(primaryKeys.getString("TABLE_CAT"), catalogNameVector, rows); saveToVector(primaryKeys.getString("TABLE_SCHEM"), schemaNameVector, rows); saveToVector(primaryKeys.getString("TABLE_NAME"), tableNameVector, rows); saveToVector(primaryKeys.getString("COLUMN_NAME"), columnNameVector, rows); final int key_seq = primaryKeys.getInt("KEY_SEQ"); saveToVector(primaryKeys.wasNull() ? null : key_seq, keySequenceVector, rows); saveToVector(primaryKeys.getString("PK_NAME"), keyNameVector, rows); } try (final VectorSchemaRoot vectorSchemaRoot = new VectorSchemaRoot(vectors)) { vectorSchemaRoot.setRowCount(rows); listener.start(vectorSchemaRoot); listener.putNext(); } } catch (SQLException e) { listener.error(e); } finally { listener.completed(); } } @Override public FlightInfo getFlightInfoExportedKeys( final CommandGetExportedKeys request, final CallContext context, final FlightDescriptor descriptor) { return getFlightInfoForSchema(request, descriptor, Schemas.GET_EXPORTED_KEYS_SCHEMA); } @Override public void getStreamExportedKeys( final CommandGetExportedKeys command, final CallContext context, final ServerStreamListener listener) { String catalog = command.hasCatalog() ? command.getCatalog() : null; String schema = command.hasDbSchema() ? command.getDbSchema() : null; String table = command.getTable(); try (Connection connection = DriverManager.getConnection(databaseUri); ResultSet keys = connection.getMetaData().getExportedKeys(catalog, schema, table); VectorSchemaRoot vectorSchemaRoot = createVectors(keys)) { listener.start(vectorSchemaRoot); listener.putNext(); } catch (SQLException e) { listener.error(e); } finally { listener.completed(); } } @Override public FlightInfo getFlightInfoImportedKeys( final CommandGetImportedKeys request, final CallContext context, final FlightDescriptor descriptor) { return getFlightInfoForSchema(request, descriptor, Schemas.GET_IMPORTED_KEYS_SCHEMA); } @Override public void getStreamImportedKeys( final CommandGetImportedKeys command, final CallContext context, final ServerStreamListener listener) { String catalog = command.hasCatalog() ? command.getCatalog() : null; String schema = command.hasDbSchema() ? command.getDbSchema() : null; String table = command.getTable(); try (Connection connection = DriverManager.getConnection(databaseUri); ResultSet keys = connection.getMetaData().getImportedKeys(catalog, schema, table); VectorSchemaRoot vectorSchemaRoot = createVectors(keys)) { listener.start(vectorSchemaRoot); listener.putNext(); } catch (final SQLException e) { listener.error(e); } finally { listener.completed(); } } @Override public FlightInfo getFlightInfoCrossReference( CommandGetCrossReference request, CallContext context, FlightDescriptor descriptor) { return getFlightInfoForSchema(request, descriptor, Schemas.GET_CROSS_REFERENCE_SCHEMA); } @Override public void getStreamCrossReference( CommandGetCrossReference command, CallContext context, ServerStreamListener listener) { final String pkCatalog = command.hasPkCatalog() ? command.getPkCatalog() : null; final String pkSchema = command.hasPkDbSchema() ? command.getPkDbSchema() : null; final String fkCatalog = command.hasFkCatalog() ? command.getFkCatalog() : null; final String fkSchema = command.hasFkDbSchema() ? command.getFkDbSchema() : null; final String pkTable = command.getPkTable(); final String fkTable = command.getFkTable(); try (Connection connection = DriverManager.getConnection(databaseUri); ResultSet keys = connection .getMetaData() .getCrossReference(pkCatalog, pkSchema, pkTable, fkCatalog, fkSchema, fkTable); VectorSchemaRoot vectorSchemaRoot = createVectors(keys)) { listener.start(vectorSchemaRoot); listener.putNext(); } catch (final SQLException e) { listener.error(e); } finally { listener.completed(); } } private VectorSchemaRoot createVectors(ResultSet keys) throws SQLException { final VarCharVector pkCatalogNameVector = new VarCharVector("pk_catalog_name", rootAllocator); final VarCharVector pkSchemaNameVector = new VarCharVector("pk_db_schema_name", rootAllocator); final VarCharVector pkTableNameVector = new VarCharVector("pk_table_name", rootAllocator); final VarCharVector pkColumnNameVector = new VarCharVector("pk_column_name", rootAllocator); final VarCharVector fkCatalogNameVector = new VarCharVector("fk_catalog_name", rootAllocator); final VarCharVector fkSchemaNameVector = new VarCharVector("fk_db_schema_name", rootAllocator); final VarCharVector fkTableNameVector = new VarCharVector("fk_table_name", rootAllocator); final VarCharVector fkColumnNameVector = new VarCharVector("fk_column_name", rootAllocator); final IntVector keySequenceVector = new IntVector("key_sequence", rootAllocator); final VarCharVector fkKeyNameVector = new VarCharVector("fk_key_name", rootAllocator); final VarCharVector pkKeyNameVector = new VarCharVector("pk_key_name", rootAllocator); final UInt1Vector updateRuleVector = new UInt1Vector("update_rule", rootAllocator); final UInt1Vector deleteRuleVector = new UInt1Vector("delete_rule", rootAllocator); Map vectorToColumnName = new HashMap<>(); vectorToColumnName.put(pkCatalogNameVector, "PKTABLE_CAT"); vectorToColumnName.put(pkSchemaNameVector, "PKTABLE_SCHEM"); vectorToColumnName.put(pkTableNameVector, "PKTABLE_NAME"); vectorToColumnName.put(pkColumnNameVector, "PKCOLUMN_NAME"); vectorToColumnName.put(fkCatalogNameVector, "FKTABLE_CAT"); vectorToColumnName.put(fkSchemaNameVector, "FKTABLE_SCHEM"); vectorToColumnName.put(fkTableNameVector, "FKTABLE_NAME"); vectorToColumnName.put(fkColumnNameVector, "FKCOLUMN_NAME"); vectorToColumnName.put(keySequenceVector, "KEY_SEQ"); vectorToColumnName.put(updateRuleVector, "UPDATE_RULE"); vectorToColumnName.put(deleteRuleVector, "DELETE_RULE"); vectorToColumnName.put(fkKeyNameVector, "FK_NAME"); vectorToColumnName.put(pkKeyNameVector, "PK_NAME"); final VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.of( pkCatalogNameVector, pkSchemaNameVector, pkTableNameVector, pkColumnNameVector, fkCatalogNameVector, fkSchemaNameVector, fkTableNameVector, fkColumnNameVector, keySequenceVector, fkKeyNameVector, pkKeyNameVector, updateRuleVector, deleteRuleVector); vectorSchemaRoot.allocateNew(); final int rowCount = saveToVectors(vectorToColumnName, keys, true); vectorSchemaRoot.setRowCount(rowCount); return vectorSchemaRoot; } @Override public void getStreamStatement( final TicketStatementQuery ticketStatementQuery, final CallContext context, final ServerStreamListener listener) { final ByteString handle = ticketStatementQuery.getStatementHandle(); final StatementContext statementContext = Objects.requireNonNull(statementLoadingCache.getIfPresent(handle)); try (final ResultSet resultSet = statementContext.getStatement().getResultSet()) { final Schema schema = jdbcToArrowSchema(resultSet.getMetaData(), DEFAULT_CALENDAR); try (VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.create(schema, rootAllocator)) { final VectorLoader loader = new VectorLoader(vectorSchemaRoot); listener.start(vectorSchemaRoot); final ArrowVectorIterator iterator = sqlToArrowVectorIterator(resultSet, rootAllocator); while (iterator.hasNext()) { final VectorUnloader unloader = new VectorUnloader(iterator.next()); loader.load(unloader.getRecordBatch()); listener.putNext(); vectorSchemaRoot.clear(); } listener.putNext(); } } catch (SQLException | IOException e) { LOGGER.error(format("Failed to getStreamPreparedStatement: <%s>.", e.getMessage()), e); listener.error(e); } finally { listener.completed(); statementLoadingCache.invalidate(handle); } } protected FlightInfo getFlightInfoForSchema( final T request, final FlightDescriptor descriptor, final Schema schema) { final Ticket ticket = new Ticket(pack(request).toByteArray()); // TODO Support multiple endpoints. final List endpoints = singletonList(new FlightEndpoint(ticket, location)); return new FlightInfo(schema, descriptor, endpoints, -1, -1); } private static class StatementRemovalListener implements RemovalListener> { @Override public void onRemoval(final RemovalNotification> notification) { try { AutoCloseables.close(notification.getValue()); } catch (final Exception e) { // swallow } } } } FlightSqlStatelessExample.java000066400000000000000000000251071515451615200376570ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql.example; import static java.lang.String.format; import static org.apache.arrow.adapter.jdbc.JdbcToArrow.sqlToArrowVectorIterator; import static org.apache.arrow.adapter.jdbc.JdbcToArrowUtils.jdbcToArrowSchema; import static org.apache.arrow.flight.sql.impl.FlightSql.*; import static org.slf4j.LoggerFactory.getLogger; import com.google.protobuf.ByteString; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.ObjectInputStream; import java.io.ObjectOutputStream; import java.io.StreamCorruptedException; import java.nio.ByteBuffer; import java.nio.channels.Channels; import java.sql.Connection; import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import org.apache.arrow.adapter.jdbc.ArrowVectorIterator; import org.apache.arrow.adapter.jdbc.JdbcParameterBinder; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.PutResult; import org.apache.arrow.flight.sql.FlightSqlProducer; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; import org.apache.arrow.vector.ipc.ArrowFileReader; import org.apache.arrow.vector.ipc.ArrowFileWriter; import org.apache.arrow.vector.ipc.SeekableReadChannel; import org.apache.arrow.vector.ipc.message.ArrowBlock; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel; import org.slf4j.Logger; /** * Example {@link FlightSqlProducer} implementation showing an Apache Derby backed Flight SQL server * that generally supports all current features of Flight SQL. */ public class FlightSqlStatelessExample extends FlightSqlExample { private static final Logger LOGGER = getLogger(FlightSqlStatelessExample.class); public static final String DB_NAME = "derbyStatelessDB"; public FlightSqlStatelessExample(final Location location, final String dbName) { super(location, dbName); } @Override public Runnable acceptPutPreparedStatementQuery( CommandPreparedStatementQuery command, CallContext context, FlightStream flightStream, StreamListener ackStream) { return () -> { final String query = new String(command.getPreparedStatementHandle().toStringUtf8()); try (Connection connection = dataSource.getConnection(); PreparedStatement preparedStatement = createPreparedStatement(connection, query)) { while (flightStream.next()) { final VectorSchemaRoot root = flightStream.getRoot(); final JdbcParameterBinder binder = JdbcParameterBinder.builder(preparedStatement, root).bindAll().build(); while (binder.next()) { // Do not execute() - will be done in a getStream call } final ByteArrayOutputStream parametersStream = new ByteArrayOutputStream(); try (ArrowFileWriter writer = new ArrowFileWriter(root, null, Channels.newChannel(parametersStream))) { writer.start(); writer.writeBatch(); } if (parametersStream.size() > 0) { final DoPutPreparedStatementResultPOJO doPutPreparedStatementResultPOJO = new DoPutPreparedStatementResultPOJO(query, parametersStream.toByteArray()); final byte[] doPutPreparedStatementResultPOJOArr = serializePOJO(doPutPreparedStatementResultPOJO); final DoPutPreparedStatementResult doPutPreparedStatementResult = DoPutPreparedStatementResult.newBuilder() .setPreparedStatementHandle( ByteString.copyFrom(ByteBuffer.wrap(doPutPreparedStatementResultPOJOArr))) .build(); try (final ArrowBuf buffer = rootAllocator.buffer(doPutPreparedStatementResult.getSerializedSize())) { buffer.writeBytes(doPutPreparedStatementResult.toByteArray()); ackStream.onNext(PutResult.metadata(buffer)); } } } } catch (SQLException | IOException e) { ackStream.onError( CallStatus.INTERNAL .withDescription("Failed to bind parameters: " + e.getMessage()) .withCause(e) .toRuntimeException()); return; } ackStream.onCompleted(); }; } @Override public void getStreamPreparedStatement( final CommandPreparedStatementQuery command, final CallContext context, final ServerStreamListener listener) { final byte[] handle = command.getPreparedStatementHandle().toByteArray(); try { // Case where there are parameters try { final DoPutPreparedStatementResultPOJO doPutPreparedStatementResultPOJO = deserializePOJO(handle); final String query = doPutPreparedStatementResultPOJO.getQuery(); try (Connection connection = dataSource.getConnection(); PreparedStatement statement = createPreparedStatement(connection, query); ArrowFileReader reader = new ArrowFileReader( new SeekableReadChannel( new ByteArrayReadableSeekableByteChannel( doPutPreparedStatementResultPOJO.getParameters())), rootAllocator)) { for (ArrowBlock arrowBlock : reader.getRecordBlocks()) { reader.loadRecordBatch(arrowBlock); VectorSchemaRoot vectorSchemaRootRecover = reader.getVectorSchemaRoot(); JdbcParameterBinder binder = JdbcParameterBinder.builder(statement, vectorSchemaRootRecover).bindAll().build(); while (binder.next()) { executeQuery(statement, listener); } } } } catch (StreamCorruptedException e) { // Case where there are no parameters final String query = new String(command.getPreparedStatementHandle().toStringUtf8()); try (Connection connection = dataSource.getConnection(); PreparedStatement preparedStatement = createPreparedStatement(connection, query)) { executeQuery(preparedStatement, listener); } } } catch (final SQLException | IOException | ClassNotFoundException e) { LOGGER.error(format("Failed to getStreamPreparedStatement: <%s>.", e.getMessage()), e); listener.error( CallStatus.INTERNAL .withDescription("Failed to prepare statement: " + e) .toRuntimeException()); } finally { listener.completed(); } } private void executeQuery(PreparedStatement statement, final ServerStreamListener listener) throws IOException, SQLException { try (final ResultSet resultSet = statement.executeQuery()) { final Schema schema = jdbcToArrowSchema(resultSet.getMetaData(), DEFAULT_CALENDAR); try (final VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.create(schema, rootAllocator)) { final VectorLoader loader = new VectorLoader(vectorSchemaRoot); listener.start(vectorSchemaRoot); final ArrowVectorIterator iterator = sqlToArrowVectorIterator(resultSet, rootAllocator); while (iterator.hasNext()) { final VectorSchemaRoot batch = iterator.next(); if (batch.getRowCount() == 0) { break; } final VectorUnloader unloader = new VectorUnloader(batch); loader.load(unloader.getRecordBatch()); listener.putNext(); vectorSchemaRoot.clear(); } listener.putNext(); } } } @Override public FlightInfo getFlightInfoPreparedStatement( final CommandPreparedStatementQuery command, final CallContext context, final FlightDescriptor descriptor) { final byte[] handle = command.getPreparedStatementHandle().toByteArray(); try { String query; try { query = deserializePOJO(handle).getQuery(); } catch (StreamCorruptedException e) { query = new String(command.getPreparedStatementHandle().toStringUtf8()); } try (Connection connection = dataSource.getConnection(); PreparedStatement statement = createPreparedStatement(connection, query)) { ResultSetMetaData metaData = statement.getMetaData(); return getFlightInfoForSchema( command, descriptor, jdbcToArrowSchema(metaData, DEFAULT_CALENDAR)); } } catch (final SQLException | IOException | ClassNotFoundException e) { LOGGER.error( format("There was a problem executing the prepared statement: <%s>.", e.getMessage()), e); throw CallStatus.INTERNAL.withCause(e).toRuntimeException(); } } private DoPutPreparedStatementResultPOJO deserializePOJO(byte[] handle) throws IOException, ClassNotFoundException { try (ByteArrayInputStream bis = new ByteArrayInputStream(handle); ObjectInputStream ois = new ObjectInputStream(bis)) { return (DoPutPreparedStatementResultPOJO) ois.readObject(); } } private byte[] serializePOJO(DoPutPreparedStatementResultPOJO doPutPreparedStatementResultPOJO) throws IOException { try (ByteArrayOutputStream bos = new ByteArrayOutputStream(); ObjectOutputStream oos = new ObjectOutputStream(bos)) { oos.writeObject(doPutPreparedStatementResultPOJO); return bos.toByteArray(); } } private PreparedStatement createPreparedStatement(Connection connection, String query) throws SQLException { return connection.prepareStatement( query, ResultSet.TYPE_SCROLL_INSENSITIVE, ResultSet.CONCUR_READ_ONLY); } } StatementContext.java000066400000000000000000000045711515451615200360710ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql.example; import java.sql.Connection; import java.sql.Statement; import java.util.Objects; import org.apache.arrow.flight.sql.FlightSqlProducer; import org.apache.arrow.util.AutoCloseables; /** * Context for {@link T} to be persisted in memory in between {@link FlightSqlProducer} calls. * * @param the {@link Statement} to be persisted. */ public final class StatementContext implements AutoCloseable { private final T statement; private final String query; public StatementContext(final T statement, final String query) { this.statement = Objects.requireNonNull(statement, "statement cannot be null."); this.query = query; } /** * Gets the statement wrapped by this {@link StatementContext}. * * @return the inner statement. */ public T getStatement() { return statement; } /** * Gets the optional SQL query wrapped by this {@link StatementContext}. * * @return the SQL query if present; empty otherwise. */ public String getQuery() { return query; } @Override public void close() throws Exception { Connection connection = statement.getConnection(); AutoCloseables.close(statement, connection); } @Override public boolean equals(final Object other) { if (this == other) { return true; } if (!(other instanceof StatementContext)) { return false; } final StatementContext that = (StatementContext) other; return statement.equals(that.statement); } @Override public int hashCode() { return Objects.hash(statement); } } apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/000077500000000000000000000000001515451615200313115ustar00rootroot00000000000000TestFlightSql.java000066400000000000000000001735331515451615200346460ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql.test; import static java.util.Arrays.asList; import static java.util.Collections.emptyList; 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.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; import com.google.common.collect.ImmutableList; import java.io.IOException; import java.io.PipedInputStream; import java.io.PipedOutputStream; import java.nio.charset.StandardCharsets; import java.sql.SQLException; import java.util.ArrayList; import java.util.Arrays; 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; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightStatusCode; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.RenewFlightEndpointRequest; import org.apache.arrow.flight.sql.FlightSqlClient; import org.apache.arrow.flight.sql.FlightSqlClient.PreparedStatement; import org.apache.arrow.flight.sql.FlightSqlColumnMetadata; import org.apache.arrow.flight.sql.FlightSqlProducer; import org.apache.arrow.flight.sql.example.FlightSqlExample; import org.apache.arrow.flight.sql.impl.FlightSql; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions.TableExistsOption; import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions.TableNotExistOption; import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedCaseSensitivity; import org.apache.arrow.flight.sql.util.TableRef; 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.ArrowStreamReader; import org.apache.arrow.vector.ipc.ArrowStreamWriter; 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.types.pojo.Schema; import org.apache.arrow.vector.util.Text; import org.apache.arrow.vector.util.VectorBatchAppender; import org.assertj.core.api.Condition; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.function.Executable; /** Test direct usage of Flight SQL workflows. */ public class TestFlightSql { protected static final Schema SCHEMA_INT_TABLE = new Schema( asList( new Field("ID", new FieldType(false, MinorType.INT.getType(), null), null), Field.nullable("KEYNAME", MinorType.VARCHAR.getType()), Field.nullable("VALUE", MinorType.INT.getType()), Field.nullable("FOREIGNID", MinorType.INT.getType()))); private static final List> EXPECTED_RESULTS_FOR_STAR_SELECT_QUERY = ImmutableList.of( asList("1", "one", "1", "1"), asList("2", "zero", "0", "1"), asList("3", "negative one", "-1", "1")); protected static final List> EXPECTED_RESULTS_FOR_PARAMETER_BINDING = ImmutableList.of(asList("1", "one", "1", "1")); private static final Map GET_SQL_INFO_EXPECTED_RESULTS_MAP = new LinkedHashMap<>(); protected static final String LOCALHOST = "localhost"; protected static BufferAllocator allocator; protected static FlightServer server; protected static FlightSqlClient sqlClient; private static void populateNext10RowsInIngestRootBatch( int startRowNumber, IntVector valueVector, VarCharVector keyNameVector, IntVector foreignIdVector, VarCharVector keyNamesToBeDeletedVector, VectorSchemaRoot ingestRoot) { final int NumRowsInBatch = 10; valueVector.reset(); keyNameVector.reset(); foreignIdVector.reset(); final IntStream range = IntStream.range(1, NumRowsInBatch); range.forEach( i -> { valueVector.setSafe(i - 1, (i + startRowNumber - 1) * NumRowsInBatch); keyNameVector.setSafe(i - 1, new Text("value" + (i + startRowNumber - 1))); foreignIdVector.setSafe(i - 1, 1); }); // put some comma and double-quote containing string as well valueVector.setSafe(NumRowsInBatch - 1, (NumRowsInBatch + startRowNumber - 1) * NumRowsInBatch); keyNameVector.setSafe( NumRowsInBatch - 1, new Text( String.format( "value%d, is \"%d\"", (NumRowsInBatch + startRowNumber - 1), (NumRowsInBatch + startRowNumber - 1) * NumRowsInBatch))); foreignIdVector.setSafe(NumRowsInBatch - 1, 1); ingestRoot.setRowCount(NumRowsInBatch); VectorBatchAppender.batchAppend(keyNamesToBeDeletedVector, keyNameVector); } @BeforeAll public static void setUp() throws Exception { setUpClientServer(); setUpExpectedResultsMap(); } private static void setUpClientServer() throws Exception { allocator = new RootAllocator(Integer.MAX_VALUE); final Location serverLocation = Location.forGrpcInsecure(LOCALHOST, 0); server = FlightServer.builder( allocator, serverLocation, new FlightSqlExample(serverLocation, FlightSqlExample.DB_NAME)) .build() .start(); final Location clientLocation = Location.forGrpcInsecure(LOCALHOST, server.getPort()); sqlClient = new FlightSqlClient(FlightClient.builder(allocator, clientLocation).build()); } protected static void setUpExpectedResultsMap() { GET_SQL_INFO_EXPECTED_RESULTS_MAP.put( Integer.toString(FlightSql.SqlInfo.FLIGHT_SQL_SERVER_NAME_VALUE), "Apache Derby"); GET_SQL_INFO_EXPECTED_RESULTS_MAP.put( Integer.toString(FlightSql.SqlInfo.FLIGHT_SQL_SERVER_VERSION_VALUE), "10.15.2.0 - (1873585)"); GET_SQL_INFO_EXPECTED_RESULTS_MAP.put( Integer.toString(FlightSql.SqlInfo.FLIGHT_SQL_SERVER_ARROW_VERSION_VALUE), "10.15.2.0 - (1873585)"); GET_SQL_INFO_EXPECTED_RESULTS_MAP.put( Integer.toString(FlightSql.SqlInfo.FLIGHT_SQL_SERVER_READ_ONLY_VALUE), "false"); GET_SQL_INFO_EXPECTED_RESULTS_MAP.put( Integer.toString(FlightSql.SqlInfo.SQL_ALL_TABLES_ARE_SELECTABLE_VALUE), "true"); GET_SQL_INFO_EXPECTED_RESULTS_MAP.put( Integer.toString(FlightSql.SqlInfo.SQL_NULL_ORDERING_VALUE), Integer.toString(FlightSql.SqlNullOrdering.SQL_NULLS_SORTED_AT_END_VALUE)); GET_SQL_INFO_EXPECTED_RESULTS_MAP.put( Integer.toString(FlightSql.SqlInfo.SQL_DDL_CATALOG_VALUE), "false"); GET_SQL_INFO_EXPECTED_RESULTS_MAP.put( Integer.toString(FlightSql.SqlInfo.SQL_DDL_SCHEMA_VALUE), "true"); GET_SQL_INFO_EXPECTED_RESULTS_MAP.put( Integer.toString(FlightSql.SqlInfo.SQL_DDL_TABLE_VALUE), "true"); GET_SQL_INFO_EXPECTED_RESULTS_MAP.put( Integer.toString(FlightSql.SqlInfo.SQL_IDENTIFIER_CASE_VALUE), Integer.toString(SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_UPPERCASE_VALUE)); GET_SQL_INFO_EXPECTED_RESULTS_MAP.put( Integer.toString(FlightSql.SqlInfo.SQL_IDENTIFIER_QUOTE_CHAR_VALUE), "\""); GET_SQL_INFO_EXPECTED_RESULTS_MAP.put( Integer.toString(FlightSql.SqlInfo.SQL_QUOTED_IDENTIFIER_CASE_VALUE), Integer.toString(SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_CASE_INSENSITIVE_VALUE)); GET_SQL_INFO_EXPECTED_RESULTS_MAP.put( Integer.toString(FlightSql.SqlInfo.SQL_MAX_COLUMNS_IN_TABLE_VALUE), "42"); } @AfterAll public static void tearDown() throws Exception { close(sqlClient, server, allocator); FlightSqlExample.removeDerbyDatabaseIfExists(FlightSqlExample.DB_NAME); } private static List> getNonConformingResultsForGetSqlInfo( final List> results) { return getNonConformingResultsForGetSqlInfo( results, FlightSql.SqlInfo.FLIGHT_SQL_SERVER_NAME, FlightSql.SqlInfo.FLIGHT_SQL_SERVER_VERSION, FlightSql.SqlInfo.FLIGHT_SQL_SERVER_ARROW_VERSION, FlightSql.SqlInfo.FLIGHT_SQL_SERVER_READ_ONLY, FlightSql.SqlInfo.SQL_ALL_TABLES_ARE_SELECTABLE, FlightSql.SqlInfo.SQL_NULL_ORDERING, FlightSql.SqlInfo.SQL_DDL_CATALOG, FlightSql.SqlInfo.SQL_DDL_SCHEMA, FlightSql.SqlInfo.SQL_DDL_TABLE, FlightSql.SqlInfo.SQL_IDENTIFIER_CASE, FlightSql.SqlInfo.SQL_IDENTIFIER_QUOTE_CHAR, FlightSql.SqlInfo.SQL_QUOTED_IDENTIFIER_CASE, FlightSql.SqlInfo.SQL_MAX_COLUMNS_IN_TABLE); } private static List> getNonConformingResultsForGetSqlInfo( final List> results, final FlightSql.SqlInfo... args) { final List> nonConformingResults = new ArrayList<>(); if (results.size() == args.length) { for (int index = 0; index < results.size(); index++) { final List result = results.get(index); final String providedName = result.get(0); final String expectedName = Integer.toString(args[index].getNumber()); System.err.println(expectedName); if (!(GET_SQL_INFO_EXPECTED_RESULTS_MAP.get(providedName).equals(result.get(1)) && providedName.equals(expectedName))) { nonConformingResults.add(result); break; } } } return nonConformingResults; } @Test public void testGetTablesSchema() { final FlightInfo info = sqlClient.getTables(null, null, null, null, true); 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); assertThat(info.getSchemaOptional()) .isEqualTo(Optional.of(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA)); } @Test public void testGetTablesResultNoSchema() throws Exception { try (final FlightStream stream = sqlClient.getStream( sqlClient.getTables(null, null, null, null, false).getEndpoints().get(0).getTicket())) { assertAll( () -> { 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 /* TODO No catalog yet */, "SYS", "SYSALIASES", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSCHECKS", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSCOLPERMS", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSCOLUMNS", "SYSTEM TABLE"), asList( null /* TODO No catalog yet */, "SYS", "SYSCONGLOMERATES", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSCONSTRAINTS", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSDEPENDS", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSFILES", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSFOREIGNKEYS", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSKEYS", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSPERMS", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSROLES", "SYSTEM TABLE"), asList( null /* TODO No catalog yet */, "SYS", "SYSROUTINEPERMS", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSSCHEMAS", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSSEQUENCES", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSSTATEMENTS", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSSTATISTICS", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSTABLEPERMS", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSTABLES", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSTRIGGERS", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSUSERS", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "SYS", "SYSVIEWS", "SYSTEM TABLE"), 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")); assertThat(results).isEqualTo(expectedResults); }); } } @Test public void testGetTablesResultFilteredNoSchema() throws Exception { try (final FlightStream stream = sqlClient.getStream( sqlClient .getTables(null, null, null, singletonList("TABLE"), false) .getEndpoints() .get(0) .getTicket())) { assertAll( () -> 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 /* TODO No catalog yet */, "APP", "FOREIGNTABLE", "TABLE"), asList(null /* TODO No catalog yet */, "APP", "INTTABLE", "TABLE")); assertThat(results).isEqualTo(expectedResults); }); } } @Test public void testGetTablesResultFilteredWithSchema() throws Exception { try (final FlightStream stream = sqlClient.getStream( sqlClient .getTables(null, null, null, singletonList("TABLE"), true) .getEndpoints() .get(0) .getTicket())) { assertAll( () -> assertThat(stream.getSchema()).isEqualTo(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA), () -> { assertThat(stream.getSchema()).isEqualTo(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA); final List> results = getResults(stream); final List> expectedResults = ImmutableList.of( // catalog_name | schema_name | table_name | table_type | table_schema asList( null /* TODO No catalog yet */, "APP", "FOREIGNTABLE", "TABLE", new Schema( asList( new Field( "ID", new FieldType( false, MinorType.INT.getType(), null, new FlightSqlColumnMetadata.Builder() .catalogName("") .typeName("INTEGER") .schemaName("APP") .tableName("FOREIGNTABLE") .precision(10) .scale(0) .isAutoIncrement(true) .build() .getMetadataMap()), null), new Field( "FOREIGNNAME", new FieldType( true, MinorType.VARCHAR.getType(), null, new FlightSqlColumnMetadata.Builder() .catalogName("") .typeName("VARCHAR") .schemaName("APP") .tableName("FOREIGNTABLE") .precision(100) .scale(0) .isAutoIncrement(false) .build() .getMetadataMap()), null), new Field( "VALUE", new FieldType( true, MinorType.INT.getType(), null, new FlightSqlColumnMetadata.Builder() .catalogName("") .typeName("INTEGER") .schemaName("APP") .tableName("FOREIGNTABLE") .precision(10) .scale(0) .isAutoIncrement(false) .build() .getMetadataMap()), null))) .toJson()), asList( null /* TODO No catalog yet */, "APP", "INTTABLE", "TABLE", new Schema( asList( new Field( "ID", new FieldType( false, MinorType.INT.getType(), null, new FlightSqlColumnMetadata.Builder() .catalogName("") .typeName("INTEGER") .schemaName("APP") .tableName("INTTABLE") .precision(10) .scale(0) .isAutoIncrement(true) .build() .getMetadataMap()), null), new Field( "KEYNAME", new FieldType( true, MinorType.VARCHAR.getType(), null, new FlightSqlColumnMetadata.Builder() .catalogName("") .typeName("VARCHAR") .schemaName("APP") .tableName("INTTABLE") .precision(100) .scale(0) .isAutoIncrement(false) .build() .getMetadataMap()), null), new Field( "VALUE", new FieldType( true, MinorType.INT.getType(), null, new FlightSqlColumnMetadata.Builder() .catalogName("") .typeName("INTEGER") .schemaName("APP") .tableName("INTTABLE") .precision(10) .scale(0) .isAutoIncrement(false) .build() .getMetadataMap()), null), new Field( "FOREIGNID", new FieldType( true, MinorType.INT.getType(), null, new FlightSqlColumnMetadata.Builder() .catalogName("") .typeName("INTEGER") .schemaName("APP") .tableName("INTTABLE") .precision(10) .scale(0) .isAutoIncrement(false) .build() .getMetadataMap()), null))) .toJson())); assertThat(results).isEqualTo(expectedResults); }); } } @Test public void testSimplePreparedStatementSchema() throws Exception { try (final PreparedStatement preparedStatement = sqlClient.prepare("SELECT * FROM intTable")) { assertAll( () -> { final Schema actualSchema = preparedStatement.getResultSetSchema(); assertThat(actualSchema).isEqualTo(SCHEMA_INT_TABLE); }, () -> { final FlightInfo info = preparedStatement.execute(); assertThat(info.getSchemaOptional()).isEqualTo(Optional.of(SCHEMA_INT_TABLE)); }); } } @Test public void testSimplePreparedStatementResults() throws Exception { try (final PreparedStatement preparedStatement = sqlClient.prepare("SELECT * FROM intTable"); final FlightStream stream = sqlClient.getStream(preparedStatement.execute().getEndpoints().get(0).getTicket())) { assertAll( () -> assertThat(stream.getSchema()).isEqualTo(SCHEMA_INT_TABLE), () -> assertThat(getResults(stream)).isEqualTo(EXPECTED_RESULTS_FOR_STAR_SELECT_QUERY)); } } @Test public void testSimplePreparedStatementResultsWithParameterBinding() throws Exception { try (PreparedStatement prepare = sqlClient.prepare("SELECT * FROM intTable WHERE id = ?")) { final Schema parameterSchema = prepare.getParameterSchema(); try (final VectorSchemaRoot insertRoot = VectorSchemaRoot.create(parameterSchema, allocator)) { insertRoot.allocateNew(); final IntVector valueVector = (IntVector) insertRoot.getVector(0); valueVector.setSafe(0, 1); insertRoot.setRowCount(1); prepare.setParameters(insertRoot); FlightInfo flightInfo = prepare.execute(); FlightStream stream = sqlClient.getStream(flightInfo.getEndpoints().get(0).getTicket()); assertAll( () -> assertThat(stream.getSchema()).isEqualTo(SCHEMA_INT_TABLE), () -> assertThat(getResults(stream)).isEqualTo(EXPECTED_RESULTS_FOR_PARAMETER_BINDING)); } } } @Test public void testSimplePreparedStatementUpdateResults() throws SQLException { try (PreparedStatement prepare = sqlClient.prepare("INSERT INTO INTTABLE (keyName, value ) VALUES (?, ?)"); PreparedStatement deletePrepare = sqlClient.prepare("DELETE FROM INTTABLE WHERE keyName = ?")) { final Schema parameterSchema = prepare.getParameterSchema(); try (final VectorSchemaRoot insertRoot = VectorSchemaRoot.create(parameterSchema, allocator)) { final VarCharVector varCharVector = (VarCharVector) insertRoot.getVector(0); final IntVector valueVector = (IntVector) insertRoot.getVector(1); final int counter = 10; insertRoot.allocateNew(); final IntStream range = IntStream.range(0, counter); range.forEach( i -> { valueVector.setSafe(i, i * counter); varCharVector.setSafe(i, new Text("value" + i)); }); insertRoot.setRowCount(counter); prepare.setParameters(insertRoot); final long updatedRows = prepare.executeUpdate(); final long deletedRows; try (final VectorSchemaRoot deleteRoot = VectorSchemaRoot.of(varCharVector)) { deletePrepare.setParameters(deleteRoot); deletedRows = deletePrepare.executeUpdate(); } assertAll( () -> assertThat(updatedRows).isEqualTo(10L), () -> assertThat(deletedRows).isEqualTo(10L)); } } } @Test public void testBulkIngest() throws IOException { // For bulk ingest DerbyDB requires uppercase column names var keyName = new Field("KEYNAME", FieldType.nullable(new ArrowType.Utf8()), null); var value = new Field("VALUE", FieldType.nullable(new ArrowType.Int(32, true)), null); var foreignId = new Field("FOREIGNID", FieldType.nullable(new ArrowType.Int(32, true)), null); Schema dataSchema = new Schema(List.of(keyName, value, foreignId)); try (final VectorSchemaRoot ingestRoot = VectorSchemaRoot.create(dataSchema, allocator); final VarCharVector keyNamesToBeDeletedVector = new VarCharVector(keyName, allocator)) { final VarCharVector keyNameVector = (VarCharVector) ingestRoot.getVector(0); final IntVector valueVector = (IntVector) ingestRoot.getVector(1); final IntVector foreignIdVector = (IntVector) ingestRoot.getVector(2); ingestRoot.allocateNew(); keyNamesToBeDeletedVector.allocateNew(); try (PipedInputStream inPipe = new PipedInputStream(1024); PipedOutputStream outPipe = new PipedOutputStream(inPipe); ArrowStreamReader reader = new ArrowStreamReader(inPipe, allocator)) { new Thread( () -> { try (ArrowStreamWriter writer = new ArrowStreamWriter(ingestRoot, null, outPipe)) { writer.start(); populateNext10RowsInIngestRootBatch( 1, valueVector, keyNameVector, foreignIdVector, keyNamesToBeDeletedVector, ingestRoot); writer.writeBatch(); populateNext10RowsInIngestRootBatch( 11, valueVector, keyNameVector, foreignIdVector, keyNamesToBeDeletedVector, ingestRoot); writer.writeBatch(); } catch (Exception e) { throw new RuntimeException(e); } }) .start(); // Ingest from a stream final long updatedRows = sqlClient.executeIngest( reader, new FlightSqlClient.ExecuteIngestOptions( "INTTABLE", TableDefinitionOptions.newBuilder() .setIfExists(TableExistsOption.TABLE_EXISTS_OPTION_APPEND) .setIfNotExist(TableNotExistOption.TABLE_NOT_EXIST_OPTION_FAIL) .build(), null, null, null)); assertThat(updatedRows).isEqualTo(-1L); // Ingest directly using VectorSchemaRoot populateNext10RowsInIngestRootBatch( 21, valueVector, keyNameVector, foreignIdVector, keyNamesToBeDeletedVector, ingestRoot); sqlClient.executeIngest( ingestRoot, new FlightSqlClient.ExecuteIngestOptions( "INTTABLE", TableDefinitionOptions.newBuilder() .setIfExists(TableExistsOption.TABLE_EXISTS_OPTION_APPEND) .setIfNotExist(TableNotExistOption.TABLE_NOT_EXIST_OPTION_FAIL) .build(), null, null, null)); try (PreparedStatement deletePrepare = sqlClient.prepare("DELETE FROM INTTABLE WHERE keyName = ?")) { final long deletedRows; try (final VectorSchemaRoot deleteRoot = VectorSchemaRoot.of(keyNamesToBeDeletedVector)) { deletePrepare.setParameters(deleteRoot); deletedRows = deletePrepare.executeUpdate(); } assertThat(deletedRows).isEqualTo(30L); } } } } @Test public void testBulkIngestTransaction() { assertThrows( RuntimeException.class, () -> { sqlClient.executeIngest( VectorSchemaRoot.create(new Schema(List.of()), allocator), new FlightSqlClient.ExecuteIngestOptions( "INTTABLE", TableDefinitionOptions.newBuilder() .setIfExists(TableExistsOption.TABLE_EXISTS_OPTION_APPEND) .setIfNotExist(TableNotExistOption.TABLE_NOT_EXIST_OPTION_FAIL) .build(), null, null, null), new FlightSqlClient.Transaction("123".getBytes(StandardCharsets.UTF_8))); }); } @Test public void testSimplePreparedStatementUpdateResultsWithoutParameters() throws SQLException { try (PreparedStatement prepare = sqlClient.prepare("INSERT INTO INTTABLE (keyName, value ) VALUES ('test', 1000)"); PreparedStatement deletePrepare = sqlClient.prepare("DELETE FROM INTTABLE WHERE keyName = 'test'")) { final long updatedRows = prepare.executeUpdate(); final long deletedRows = deletePrepare.executeUpdate(); assertAll( () -> assertThat(updatedRows).isEqualTo(1L), () -> assertThat(deletedRows).isEqualTo(1L)); } } @Test public void testSimplePreparedStatementClosesProperly() { final PreparedStatement preparedStatement = sqlClient.prepare("SELECT * FROM intTable"); assertAll( () -> { assertThat(preparedStatement.isClosed()).isEqualTo(false); }, () -> { preparedStatement.close(); assertThat(preparedStatement.isClosed()).isEqualTo(true); }); } @Test public void testGetCatalogsSchema() { final FlightInfo info = sqlClient.getCatalogs(); assertThat(info.getSchemaOptional()) .isEqualTo(Optional.of(FlightSqlProducer.Schemas.GET_CATALOGS_SCHEMA)); } @Test public void testGetCatalogsResults() throws Exception { try (final FlightStream stream = sqlClient.getStream(sqlClient.getCatalogs().getEndpoints().get(0).getTicket())) { assertAll( () -> assertThat(stream.getSchema()) .isEqualTo(FlightSqlProducer.Schemas.GET_CATALOGS_SCHEMA), () -> { List> catalogs = getResults(stream); assertThat(catalogs).isEqualTo(emptyList()); }); } } @Test public void testGetTableTypesSchema() { final FlightInfo info = sqlClient.getTableTypes(); assertThat(info.getSchemaOptional()) .isEqualTo(Optional.of(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA)); } @Test public void testGetTableTypesResult() throws Exception { try (final FlightStream stream = sqlClient.getStream(sqlClient.getTableTypes().getEndpoints().get(0).getTicket())) { assertAll( () -> { assertThat(stream.getSchema()) .isEqualTo(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA); }, () -> { final List> tableTypes = getResults(stream); final List> expectedTableTypes = ImmutableList.of( // table_type singletonList("SYNONYM"), singletonList("SYSTEM TABLE"), singletonList("TABLE"), singletonList("VIEW")); assertThat(tableTypes).isEqualTo(expectedTableTypes); }); } } @Test public void testGetSchemasSchema() { final FlightInfo info = sqlClient.getSchemas(null, null); assertThat(info.getSchemaOptional()) .isEqualTo(Optional.of(FlightSqlProducer.Schemas.GET_SCHEMAS_SCHEMA)); } @Test public void testGetSchemasResult() throws Exception { try (final FlightStream stream = sqlClient.getStream(sqlClient.getSchemas(null, null).getEndpoints().get(0).getTicket())) { assertAll( () -> { assertThat(stream.getSchema()).isEqualTo(FlightSqlProducer.Schemas.GET_SCHEMAS_SCHEMA); }, () -> { final List> schemas = getResults(stream); final List> expectedSchemas = ImmutableList.of( // catalog_name | schema_name asList(null /* TODO Add catalog. */, "APP"), asList(null /* TODO Add catalog. */, "NULLID"), asList(null /* TODO Add catalog. */, "SQLJ"), asList(null /* TODO Add catalog. */, "SYS"), asList(null /* TODO Add catalog. */, "SYSCAT"), asList(null /* TODO Add catalog. */, "SYSCS_DIAG"), asList(null /* TODO Add catalog. */, "SYSCS_UTIL"), asList(null /* TODO Add catalog. */, "SYSFUN"), asList(null /* TODO Add catalog. */, "SYSIBM"), asList(null /* TODO Add catalog. */, "SYSPROC"), asList(null /* TODO Add catalog. */, "SYSSTAT")); assertThat(schemas).isEqualTo(expectedSchemas); }); } } @Test public void testGetPrimaryKey() { final FlightInfo flightInfo = sqlClient.getPrimaryKeys(TableRef.of(null, null, "INTTABLE")); final FlightStream stream = sqlClient.getStream(flightInfo.getEndpoints().get(0).getTicket()); final List> results = getResults(stream); assertAll( () -> assertThat(results.size()).isEqualTo(1), () -> { final List result = results.get(0); assertAll( () -> 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(); assertThat(info.getSchemaOptional()) .isEqualTo(Optional.of(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA)); } @Test public void testGetSqlInfoResults() throws Exception { final FlightInfo info = sqlClient.getSqlInfo(); try (final FlightStream stream = sqlClient.getStream(info.getEndpoints().get(0).getTicket())) { assertAll( () -> assertThat(stream.getSchema()) .isEqualTo(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA), () -> assertThat(getNonConformingResultsForGetSqlInfo(getResults(stream))) .isEqualTo(emptyList())); } } @Test public void testGetSqlInfoResultsWithSingleArg() throws Exception { final FlightSql.SqlInfo arg = FlightSql.SqlInfo.FLIGHT_SQL_SERVER_NAME; final FlightInfo info = sqlClient.getSqlInfo(arg); try (final FlightStream stream = sqlClient.getStream(info.getEndpoints().get(0).getTicket())) { assertAll( () -> assertThat(stream.getSchema()) .isEqualTo(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA), () -> assertThat(getNonConformingResultsForGetSqlInfo(getResults(stream), arg)) .isEqualTo(emptyList())); } } @Test public void testGetSqlInfoResultsWithManyArgs() throws Exception { final FlightSql.SqlInfo[] args = { FlightSql.SqlInfo.FLIGHT_SQL_SERVER_NAME, FlightSql.SqlInfo.FLIGHT_SQL_SERVER_VERSION, FlightSql.SqlInfo.FLIGHT_SQL_SERVER_ARROW_VERSION, FlightSql.SqlInfo.FLIGHT_SQL_SERVER_READ_ONLY, FlightSql.SqlInfo.SQL_ALL_TABLES_ARE_SELECTABLE, FlightSql.SqlInfo.SQL_NULL_ORDERING, FlightSql.SqlInfo.SQL_DDL_CATALOG, FlightSql.SqlInfo.SQL_DDL_SCHEMA, FlightSql.SqlInfo.SQL_DDL_TABLE, FlightSql.SqlInfo.SQL_IDENTIFIER_CASE, FlightSql.SqlInfo.SQL_IDENTIFIER_QUOTE_CHAR, FlightSql.SqlInfo.SQL_QUOTED_IDENTIFIER_CASE, FlightSql.SqlInfo.SQL_MAX_COLUMNS_IN_TABLE }; final FlightInfo info = sqlClient.getSqlInfo(args); try (final FlightStream stream = sqlClient.getStream(info.getEndpoints().get(0).getTicket())) { assertAll( () -> assertThat(stream.getSchema()) .isEqualTo(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA), () -> assertThat(getNonConformingResultsForGetSqlInfo(getResults(stream), args)) .isEqualTo(emptyList())); } } @Test public void testGetCommandExportedKeys() throws Exception { try (final FlightStream stream = sqlClient.getStream( sqlClient .getExportedKeys(TableRef.of(null, null, "FOREIGNTABLE")) .getEndpoints() .get(0) .getTicket())) { final List> results = getResults(stream); final List> matchers = asList( 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 Condition expected = matchers.get(i); assertions.add(() -> assertThat(actual).satisfies(expected)); } assertAll(assertions); } } @Test public void testGetCommandImportedKeys() throws Exception { try (final FlightStream stream = sqlClient.getStream( sqlClient .getImportedKeys(TableRef.of(null, null, "INTTABLE")) .getEndpoints() .get(0) .getTicket())) { final List> results = getResults(stream); final List> matchers = asList( 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 Condition expected = matchers.get(i); assertions.add(() -> assertThat(actual).satisfies(expected)); } assertAll(assertions); } } @Test public void testGetTypeInfo() throws Exception { FlightInfo flightInfo = sqlClient.getXdbcTypeInfo(); try (FlightStream stream = sqlClient.getStream(flightInfo.getEndpoints().get(0).getTicket())) { final List> results = getResults(stream); final List> matchers = ImmutableList.of( asList( "BIGINT", "-5", "19", null, null, emptyList().toString(), "1", "false", "2", "false", "false", "true", "BIGINT", "0", "0", null, null, "10", null), asList( "LONG VARCHAR FOR BIT DATA", "-4", "32700", "X'", "'", emptyList().toString(), "1", "false", "0", "true", "false", "false", "LONG VARCHAR FOR BIT DATA", null, null, null, null, null, null), asList( "VARCHAR () FOR BIT DATA", "-3", "32672", "X'", "'", singletonList("length").toString(), "1", "false", "2", "true", "false", "false", "VARCHAR () FOR BIT DATA", null, null, null, null, null, null), asList( "CHAR () FOR BIT DATA", "-2", "254", "X'", "'", singletonList("length").toString(), "1", "false", "2", "true", "false", "false", "CHAR () FOR BIT DATA", null, null, null, null, null, null), asList( "LONG VARCHAR", "-1", "32700", "'", "'", emptyList().toString(), "1", "true", "1", "true", "false", "false", "LONG VARCHAR", null, null, null, null, null, null), asList( "CHAR", "1", "254", "'", "'", singletonList("length").toString(), "1", "true", "3", "true", "false", "false", "CHAR", null, null, null, null, null, null), asList( "NUMERIC", "2", "31", null, null, Arrays.asList("precision", "scale").toString(), "1", "false", "2", "false", "true", "false", "NUMERIC", "0", "31", null, null, "10", null), asList( "DECIMAL", "3", "31", null, null, Arrays.asList("precision", "scale").toString(), "1", "false", "2", "false", "true", "false", "DECIMAL", "0", "31", null, null, "10", null), asList( "INTEGER", "4", "10", null, null, emptyList().toString(), "1", "false", "2", "false", "false", "true", "INTEGER", "0", "0", null, null, "10", null), asList( "SMALLINT", "5", "5", null, null, emptyList().toString(), "1", "false", "2", "false", "false", "true", "SMALLINT", "0", "0", null, null, "10", null), asList( "FLOAT", "6", "52", null, null, singletonList("precision").toString(), "1", "false", "2", "false", "false", "false", "FLOAT", null, null, null, null, "2", null), asList( "REAL", "7", "23", null, null, emptyList().toString(), "1", "false", "2", "false", "false", "false", "REAL", null, null, null, null, "2", null), asList( "DOUBLE", "8", "52", null, null, emptyList().toString(), "1", "false", "2", "false", "false", "false", "DOUBLE", null, null, null, null, "2", null), asList( "VARCHAR", "12", "32672", "'", "'", singletonList("length").toString(), "1", "true", "3", "true", "false", "false", "VARCHAR", null, null, null, null, null, null), asList( "BOOLEAN", "16", "1", null, null, emptyList().toString(), "1", "false", "2", "true", "false", "false", "BOOLEAN", null, null, null, null, null, null), asList( "DATE", "91", "10", "DATE'", "'", emptyList().toString(), "1", "false", "2", "true", "false", "false", "DATE", "0", "0", null, null, "10", null), asList( "TIME", "92", "8", "TIME'", "'", emptyList().toString(), "1", "false", "2", "true", "false", "false", "TIME", "0", "0", null, null, "10", null), asList( "TIMESTAMP", "93", "29", "TIMESTAMP'", "'", emptyList().toString(), "1", "false", "2", "true", "false", "false", "TIMESTAMP", "0", "9", null, null, "10", null), asList( "OBJECT", "2000", null, null, null, emptyList().toString(), "1", "false", "2", "true", "false", "false", "OBJECT", null, null, null, null, null, null), asList( "BLOB", "2004", "2147483647", null, null, singletonList("length").toString(), "1", "false", "0", null, "false", null, "BLOB", null, null, null, null, null, null), asList( "CLOB", "2005", "2147483647", "'", "'", singletonList("length").toString(), "1", "true", "1", null, "false", null, "CLOB", null, null, null, null, null, null), asList( "XML", "2009", null, null, null, emptyList().toString(), "1", "true", "0", "false", "false", "false", "XML", null, null, null, null, null, null)); assertThat(results).isEqualTo(matchers); } } @Test public void testGetTypeInfoWithFiltering() throws Exception { FlightInfo flightInfo = sqlClient.getXdbcTypeInfo(-5); try (FlightStream stream = sqlClient.getStream(flightInfo.getEndpoints().get(0).getTicket())) { final List> results = getResults(stream); final List> matchers = ImmutableList.of( asList( "BIGINT", "-5", "19", null, null, emptyList().toString(), "1", "false", "2", "false", "false", "true", "BIGINT", "0", "0", null, null, "10", null)); assertThat(results).isEqualTo(matchers); } } @Test public void testGetCommandCrossReference() throws Exception { final FlightInfo flightInfo = sqlClient.getCrossReference( TableRef.of(null, null, "FOREIGNTABLE"), TableRef.of(null, null, "INTTABLE")); try (final FlightStream stream = sqlClient.getStream(flightInfo.getEndpoints().get(0).getTicket())) { final List> results = getResults(stream); final List> matchers = asList( 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 Condition expected = matchers.get(i); assertions.add(() -> assertThat(actual).satisfies(expected)); } assertAll(assertions); } } @Test public void testCreateStatementSchema() throws Exception { final FlightInfo info = sqlClient.execute("SELECT * FROM intTable"); 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())) { while (stream.next()) { // Do nothing } } } @Test public void testCreateStatementResults() throws Exception { try (final FlightStream stream = sqlClient.getStream( sqlClient.execute("SELECT * FROM intTable").getEndpoints().get(0).getTicket())) { assertAll( () -> { assertThat(stream.getSchema()).isEqualTo(SCHEMA_INT_TABLE); }, () -> { assertThat(getResults(stream)).isEqualTo(EXPECTED_RESULTS_FOR_STAR_SELECT_QUERY); }); } } @Test public void testExecuteUpdate() { assertAll( () -> { long insertedCount = sqlClient.executeUpdate( "INSERT INTO INTTABLE (keyName, value) VALUES " + "('KEYNAME1', 1001), ('KEYNAME2', 1002), ('KEYNAME3', 1003)"); assertThat(insertedCount).isEqualTo(3L); }, () -> { long updatedCount = sqlClient.executeUpdate( "UPDATE INTTABLE SET keyName = 'KEYNAME1' " + "WHERE keyName = 'KEYNAME2' OR keyName = 'KEYNAME3'"); assertThat(updatedCount).isEqualTo(2L); }, () -> { long deletedCount = sqlClient.executeUpdate("DELETE FROM INTTABLE WHERE keyName = 'KEYNAME1'"); assertThat(deletedCount).isEqualTo(3L); }); } @Test public void testQueryWithNoResultsShouldNotHang() throws Exception { try (final PreparedStatement preparedStatement = sqlClient.prepare("SELECT * FROM intTable WHERE 1 = 0"); final FlightStream stream = sqlClient.getStream(preparedStatement.execute().getEndpoints().get(0).getTicket())) { assertAll( () -> assertThat(stream.getSchema()).isEqualTo(SCHEMA_INT_TABLE), () -> { final List> result = getResults(stream); assertThat(result).isEqualTo(emptyList()); }); } } @Test public void testCancelFlightInfo() { FlightInfo info = sqlClient.getSqlInfo(); CancelFlightInfoRequest request = new CancelFlightInfoRequest(info); FlightRuntimeException fre = assertThrows(FlightRuntimeException.class, () -> sqlClient.cancelFlightInfo(request)); assertEquals(FlightStatusCode.UNIMPLEMENTED, fre.status().code()); } @Test public void testCancelQuery() { FlightInfo info = sqlClient.getSqlInfo(); FlightRuntimeException fre = assertThrows(FlightRuntimeException.class, () -> sqlClient.cancelQuery(info)); assertEquals(FlightStatusCode.UNIMPLEMENTED, fre.status().code()); } @Test public void testRenewEndpoint() { FlightInfo info = sqlClient.getSqlInfo(); FlightRuntimeException fre = assertThrows( FlightRuntimeException.class, () -> sqlClient.renewFlightEndpoint( new RenewFlightEndpointRequest(info.getEndpoints().get(0)))); assertEquals(FlightStatusCode.UNIMPLEMENTED, fre.status().code()); } } TestFlightSqlStateless.java000066400000000000000000000077011515451615200365270ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql.test; import static org.apache.arrow.flight.sql.util.FlightStreamUtils.getResults; import static org.apache.arrow.util.AutoCloseables.close; import static org.assertj.core.api.Assertions.assertThat; import static org.junit.jupiter.api.Assertions.assertAll; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.sql.FlightSqlClient; import org.apache.arrow.flight.sql.FlightSqlClient.PreparedStatement; import org.apache.arrow.flight.sql.example.FlightSqlStatelessExample; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; /** Test direct usage of Flight SQL workflows. */ public class TestFlightSqlStateless extends TestFlightSql { @BeforeAll public static void setUp() throws Exception { setUpClientServer(); setUpExpectedResultsMap(); } @AfterAll public static void tearDown() throws Exception { close(sqlClient, server, allocator); FlightSqlStatelessExample.removeDerbyDatabaseIfExists(FlightSqlStatelessExample.DB_NAME); } private static void setUpClientServer() throws Exception { allocator = new RootAllocator(Integer.MAX_VALUE); final Location serverLocation = Location.forGrpcInsecure(LOCALHOST, 0); server = FlightServer.builder( allocator, serverLocation, new FlightSqlStatelessExample(serverLocation, FlightSqlStatelessExample.DB_NAME)) .build() .start(); final Location clientLocation = Location.forGrpcInsecure(LOCALHOST, server.getPort()); sqlClient = new FlightSqlClient(FlightClient.builder(allocator, clientLocation).build()); } @Override @Test public void testSimplePreparedStatementResultsWithParameterBinding() throws Exception { try (PreparedStatement prepare = sqlClient.prepare("SELECT * FROM intTable WHERE id = ?")) { final Schema parameterSchema = prepare.getParameterSchema(); try (final VectorSchemaRoot insertRoot = VectorSchemaRoot.create(parameterSchema, allocator)) { insertRoot.allocateNew(); final IntVector valueVector = (IntVector) insertRoot.getVector(0); valueVector.setSafe(0, 1); insertRoot.setRowCount(1); prepare.setParameters(insertRoot); final FlightInfo flightInfo = prepare.execute(); for (FlightEndpoint endpoint : flightInfo.getEndpoints()) { try (FlightStream stream = sqlClient.getStream(endpoint.getTicket())) { assertAll( () -> assertThat(stream.getSchema()).isEqualTo(SCHEMA_INT_TABLE), () -> assertThat(getResults(stream)) .isEqualTo(EXPECTED_RESULTS_FOR_PARAMETER_BINDING)); } } } } } } TestFlightSqlStreams.java000066400000000000000000000314431515451615200361760ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql.test; import static java.util.Arrays.asList; import static java.util.Collections.emptyList; 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.apache.arrow.vector.types.Types.MinorType.INT; import static org.assertj.core.api.Assertions.assertThat; import static org.junit.jupiter.api.Assertions.assertAll; import com.google.common.collect.ImmutableList; import com.google.protobuf.Any; import com.google.protobuf.Message; import java.util.Collections; import java.util.List; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.Ticket; import org.apache.arrow.flight.sql.BasicFlightSqlProducer; import org.apache.arrow.flight.sql.FlightSqlClient; import org.apache.arrow.flight.sql.FlightSqlProducer; import org.apache.arrow.flight.sql.impl.FlightSql; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.Types; 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.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; public class TestFlightSqlStreams { /** * A limited {@link FlightSqlProducer} for testing GetTables, GetTableTypes, GetSqlInfo, and * limited SQL commands. */ private static class FlightSqlTestProducer extends BasicFlightSqlProducer { // Note that for simplicity the getStream* implementations are blocking, but a proper // FlightSqlProducer should // have non-blocking implementations of getStream*. private static final String FIXED_QUERY = "SELECT 1 AS c1 FROM test_table"; private static final Schema FIXED_SCHEMA = new Schema(asList(Field.nullable("c1", Types.MinorType.INT.getType()))); private BufferAllocator allocator; FlightSqlTestProducer(BufferAllocator allocator) { this.allocator = allocator; } @Override protected List determineEndpoints( T request, FlightDescriptor flightDescriptor, Schema schema) { if (request instanceof FlightSql.CommandGetTables || request instanceof FlightSql.CommandGetTableTypes || request instanceof FlightSql.CommandGetXdbcTypeInfo || request instanceof FlightSql.CommandGetSqlInfo) { return Collections.singletonList( new FlightEndpoint(new Ticket(Any.pack(request).toByteArray()))); } else if (request instanceof FlightSql.CommandStatementQuery && ((FlightSql.CommandStatementQuery) request).getQuery().equals(FIXED_QUERY)) { // Tickets from CommandStatementQuery requests should be built using TicketStatementQuery // then packed() into // a ticket. The content of the statement handle is specific to the FlightSqlProducer. It // does not need to // be the query. It can be a query ID for example. FlightSql.TicketStatementQuery ticketStatementQuery = FlightSql.TicketStatementQuery.newBuilder() .setStatementHandle(((FlightSql.CommandStatementQuery) request).getQueryBytes()) .build(); return Collections.singletonList( new FlightEndpoint(new Ticket(Any.pack(ticketStatementQuery).toByteArray()))); } throw CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException(); } @Override public FlightInfo getFlightInfoStatement( FlightSql.CommandStatementQuery command, CallContext context, FlightDescriptor descriptor) { return generateFlightInfo(command, descriptor, FIXED_SCHEMA); } @Override public void getStreamStatement( FlightSql.TicketStatementQuery ticket, CallContext context, ServerStreamListener listener) { final String query = ticket.getStatementHandle().toStringUtf8(); if (!query.equals(FIXED_QUERY)) { listener.error( CallStatus.UNIMPLEMENTED.withDescription("Not implemented.").toRuntimeException()); } try (VectorSchemaRoot root = VectorSchemaRoot.create(FIXED_SCHEMA, allocator)) { root.setRowCount(1); ((IntVector) root.getVector("c1")).setSafe(0, 1); listener.start(root); listener.putNext(); listener.completed(); } } @Override public void getStreamSqlInfo( FlightSql.CommandGetSqlInfo command, CallContext context, ServerStreamListener listener) { try (VectorSchemaRoot root = VectorSchemaRoot.create(Schemas.GET_SQL_INFO_SCHEMA, allocator)) { root.setRowCount(0); listener.start(root); listener.putNext(); listener.completed(); } } @Override public void getStreamTypeInfo( FlightSql.CommandGetXdbcTypeInfo request, CallContext context, ServerStreamListener listener) { try (VectorSchemaRoot root = VectorSchemaRoot.create(Schemas.GET_TYPE_INFO_SCHEMA, allocator)) { root.setRowCount(1); ((VarCharVector) root.getVector("type_name")).setSafe(0, new Text("Integer")); ((IntVector) root.getVector("data_type")).setSafe(0, INT.ordinal()); ((IntVector) root.getVector("column_size")).setSafe(0, 400); root.getVector("literal_prefix").setNull(0); root.getVector("literal_suffix").setNull(0); root.getVector("create_params").setNull(0); ((IntVector) root.getVector("nullable")) .setSafe(0, FlightSql.Nullable.NULLABILITY_NULLABLE.getNumber()); ((BitVector) root.getVector("case_sensitive")).setSafe(0, 1); ((IntVector) root.getVector("nullable")) .setSafe(0, FlightSql.Searchable.SEARCHABLE_FULL.getNumber()); ((BitVector) root.getVector("unsigned_attribute")).setSafe(0, 1); root.getVector("fixed_prec_scale").setNull(0); ((BitVector) root.getVector("auto_increment")).setSafe(0, 1); ((VarCharVector) root.getVector("local_type_name")).setSafe(0, new Text("Integer")); root.getVector("minimum_scale").setNull(0); root.getVector("maximum_scale").setNull(0); ((IntVector) root.getVector("sql_data_type")).setSafe(0, INT.ordinal()); root.getVector("datetime_subcode").setNull(0); ((IntVector) root.getVector("num_prec_radix")).setSafe(0, 10); root.getVector("interval_precision").setNull(0); listener.start(root); listener.putNext(); listener.completed(); } } @Override public void getStreamTables( FlightSql.CommandGetTables command, CallContext context, ServerStreamListener listener) { try (VectorSchemaRoot root = VectorSchemaRoot.create(Schemas.GET_TABLES_SCHEMA_NO_SCHEMA, allocator)) { root.setRowCount(1); root.getVector("catalog_name").setNull(0); root.getVector("db_schema_name").setNull(0); ((VarCharVector) root.getVector("table_name")).setSafe(0, new Text("test_table")); ((VarCharVector) root.getVector("table_type")).setSafe(0, new Text("TABLE")); listener.start(root); listener.putNext(); listener.completed(); } } @Override public void getStreamTableTypes(CallContext context, ServerStreamListener listener) { try (VectorSchemaRoot root = VectorSchemaRoot.create(Schemas.GET_TABLE_TYPES_SCHEMA, allocator)) { root.setRowCount(1); ((VarCharVector) root.getVector("table_type")).setSafe(0, new Text("TABLE")); listener.start(root); listener.putNext(); listener.completed(); } } } private static BufferAllocator allocator; private static FlightServer server; private static FlightSqlClient sqlClient; @BeforeAll public static void setUp() throws Exception { allocator = new RootAllocator(Integer.MAX_VALUE); final Location serverLocation = Location.forGrpcInsecure("localhost", 0); server = FlightServer.builder(allocator, serverLocation, new FlightSqlTestProducer(allocator)) .build() .start(); final Location clientLocation = Location.forGrpcInsecure("localhost", server.getPort()); sqlClient = new FlightSqlClient(FlightClient.builder(allocator, clientLocation).build()); } @AfterAll public static void tearDown() throws Exception { close(sqlClient, server); // Manually close all child allocators. allocator.getChildAllocators().forEach(BufferAllocator::close); close(allocator); } @Test public void testGetTablesResultNoSchema() throws Exception { try (final FlightStream stream = sqlClient.getStream( sqlClient.getTables(null, null, null, null, false).getEndpoints().get(0).getTicket())) { assertAll( () -> 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")); assertThat(results).isEqualTo(expectedResults); }); } } @Test public void testGetTableTypesResult() throws Exception { try (final FlightStream stream = sqlClient.getStream(sqlClient.getTableTypes().getEndpoints().get(0).getTicket())) { assertAll( () -> assertThat(stream.getSchema()) .isEqualTo(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA), () -> { final List> tableTypes = getResults(stream); final List> expectedTableTypes = ImmutableList.of( // table_type singletonList("TABLE")); assertThat(tableTypes).isEqualTo(expectedTableTypes); }); } } @Test public void testGetSqlInfoResults() throws Exception { final FlightInfo info = sqlClient.getSqlInfo(); try (final FlightStream stream = sqlClient.getStream(info.getEndpoints().get(0).getTicket())) { assertAll( () -> assertThat(stream.getSchema()) .isEqualTo(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA), () -> assertThat(getResults(stream)).isEqualTo(emptyList())); } } @Test public void testGetTypeInfo() throws Exception { FlightInfo flightInfo = sqlClient.getXdbcTypeInfo(); try (FlightStream stream = sqlClient.getStream(flightInfo.getEndpoints().get(0).getTicket())) { final List> results = getResults(stream); final List> matchers = ImmutableList.of( asList( "Integer", "4", "400", null, null, "3", "true", null, "true", null, "true", "Integer", null, null, "4", null, "10", null)); assertThat(results).isEqualTo(matchers); } } @Test public void testExecuteQuery() throws Exception { try (final FlightStream stream = sqlClient.getStream( sqlClient .execute(FlightSqlTestProducer.FIXED_QUERY) .getEndpoints() .get(0) .getTicket())) { assertAll( () -> assertThat(stream.getSchema()).isEqualTo(FlightSqlTestProducer.FIXED_SCHEMA), () -> assertThat(getResults(stream)).isEqualTo(singletonList(singletonList("1")))); } } } apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/util/000077500000000000000000000000001515451615200313075ustar00rootroot00000000000000AdhocTestOption.java000066400000000000000000000030411515451615200351400ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql.util; import com.google.protobuf.Descriptors.EnumDescriptor; import com.google.protobuf.Descriptors.EnumValueDescriptor; import com.google.protobuf.ProtocolMessageEnum; enum AdhocTestOption implements ProtocolMessageEnum { OPTION_A, OPTION_B, OPTION_C; @Override public int getNumber() { return ordinal(); } @Override public EnumValueDescriptor getValueDescriptor() { throw getUnsupportedException(); } @Override public EnumDescriptor getDescriptorForType() { throw getUnsupportedException(); } private UnsupportedOperationException getUnsupportedException() { return new UnsupportedOperationException( "Unimplemented method is irrelevant for the scope of this test."); } } FlightStreamUtils.java000066400000000000000000000150031515451615200355040ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql.util; import static java.util.Objects.isNull; import java.io.ByteArrayInputStream; import java.io.IOException; import java.nio.channels.Channels; import java.util.ArrayList; import java.util.List; import java.util.Objects; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.UInt1Vector; import org.apache.arrow.vector.UInt4Vector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.ipc.ReadChannel; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.util.Text; public class FlightStreamUtils { public static List> getResults(FlightStream stream) { final List> results = new ArrayList<>(); while (stream.next()) { try (final VectorSchemaRoot root = stream.getRoot()) { final long rowCount = root.getRowCount(); for (int i = 0; i < rowCount; ++i) { results.add(new ArrayList<>()); } root.getSchema() .getFields() .forEach( field -> { try (final FieldVector fieldVector = root.getVector(field.getName())) { if (fieldVector instanceof VarCharVector) { final VarCharVector varcharVector = (VarCharVector) fieldVector; for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) { final Text data = varcharVector.getObject(rowIndex); results.get(rowIndex).add(isNull(data) ? null : data.toString()); } } else if (fieldVector instanceof IntVector) { for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) { Object data = fieldVector.getObject(rowIndex); results.get(rowIndex).add(isNull(data) ? null : Objects.toString(data)); } } else if (fieldVector instanceof VarBinaryVector) { final VarBinaryVector varbinaryVector = (VarBinaryVector) fieldVector; for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) { final byte[] data = varbinaryVector.getObject(rowIndex); final String output; try { output = isNull(data) ? null : MessageSerializer.deserializeSchema( new ReadChannel( Channels.newChannel(new ByteArrayInputStream(data)))) .toJson(); } catch (final IOException e) { throw new RuntimeException("Failed to deserialize schema", e); } results.get(rowIndex).add(output); } } else if (fieldVector instanceof DenseUnionVector) { final DenseUnionVector denseUnionVector = (DenseUnionVector) fieldVector; for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) { final Object data = denseUnionVector.getObject(rowIndex); results.get(rowIndex).add(isNull(data) ? null : Objects.toString(data)); } } else if (fieldVector instanceof ListVector) { for (int i = 0; i < fieldVector.getValueCount(); i++) { if (!fieldVector.isNull(i)) { List elements = (List) ((ListVector) fieldVector).getObject(i); List values = new ArrayList<>(); for (Text element : elements) { values.add(element.toString()); } results.get(i).add(values.toString()); } } } else if (fieldVector instanceof UInt4Vector) { final UInt4Vector uInt4Vector = (UInt4Vector) fieldVector; for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) { final Object data = uInt4Vector.getObject(rowIndex); results.get(rowIndex).add(isNull(data) ? null : Objects.toString(data)); } } else if (fieldVector instanceof UInt1Vector) { final UInt1Vector uInt1Vector = (UInt1Vector) fieldVector; for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) { final Object data = uInt1Vector.getObject(rowIndex); results.get(rowIndex).add(isNull(data) ? null : Objects.toString(data)); } } else if (fieldVector instanceof BitVector) { for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) { Object data = fieldVector.getObject(rowIndex); results.get(rowIndex).add(isNull(data) ? null : Objects.toString(data)); } } else { throw new UnsupportedOperationException("Not yet implemented"); } } }); } } return results; } } SqlInfoOptionsUtilsBitmaskCreationTest.java000066400000000000000000000043101515451615200417010ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql.util; import static java.util.Arrays.asList; import static org.apache.arrow.flight.sql.util.AdhocTestOption.OPTION_A; import static org.apache.arrow.flight.sql.util.AdhocTestOption.OPTION_B; import static org.apache.arrow.flight.sql.util.AdhocTestOption.OPTION_C; import static org.apache.arrow.flight.sql.util.SqlInfoOptionsUtils.createBitmaskFromEnums; import static org.junit.jupiter.api.Assertions.assertEquals; import java.util.List; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.MethodSource; public final class SqlInfoOptionsUtilsBitmaskCreationTest { public static List provideParameters() { return asList( new Object[][] { {new AdhocTestOption[0], 0L}, {new AdhocTestOption[] {OPTION_A}, 1L}, {new AdhocTestOption[] {OPTION_B}, 0b10L}, {new AdhocTestOption[] {OPTION_A, OPTION_B}, 0b11L}, {new AdhocTestOption[] {OPTION_C}, 0b100L}, {new AdhocTestOption[] {OPTION_A, OPTION_C}, 0b101L}, {new AdhocTestOption[] {OPTION_B, OPTION_C}, 0b110L}, {AdhocTestOption.values(), 0b111L}, }); } @ParameterizedTest @MethodSource("provideParameters") public void testShouldBuildBitmaskFromEnums( AdhocTestOption[] adhocTestOptions, long expectedBitmask) { assertEquals(createBitmaskFromEnums(adhocTestOptions), expectedBitmask); } } SqlInfoOptionsUtilsBitmaskParsingTest.java000066400000000000000000000050231515451615200415420ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.sql.util; import static java.util.Arrays.asList; import static java.util.Arrays.stream; import static java.util.stream.Collectors.toCollection; import static org.apache.arrow.flight.sql.util.AdhocTestOption.OPTION_A; import static org.apache.arrow.flight.sql.util.AdhocTestOption.OPTION_B; import static org.apache.arrow.flight.sql.util.AdhocTestOption.OPTION_C; import static org.apache.arrow.flight.sql.util.SqlInfoOptionsUtils.doesBitmaskTranslateToEnum; import static org.junit.jupiter.api.Assertions.assertEquals; import java.util.EnumSet; import java.util.List; import java.util.Set; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.MethodSource; public final class SqlInfoOptionsUtilsBitmaskParsingTest { public static List provideParameters() { return asList( new Object[][] { {0L, EnumSet.noneOf(AdhocTestOption.class)}, {1L, EnumSet.of(OPTION_A)}, {0b10L, EnumSet.of(OPTION_B)}, {0b11L, EnumSet.of(OPTION_A, OPTION_B)}, {0b100L, EnumSet.of(OPTION_C)}, {0b101L, EnumSet.of(OPTION_A, OPTION_C)}, {0b110L, EnumSet.of(OPTION_B, OPTION_C)}, {0b111L, EnumSet.allOf(AdhocTestOption.class)}, }); } @ParameterizedTest @MethodSource("provideParameters") public void testShouldFilterOutEnumsBasedOnBitmask( long bitmask, Set expectedOptions) { final Set actualOptions = stream(AdhocTestOption.values()) .filter(enumInstance -> doesBitmaskTranslateToEnum(enumInstance, bitmask)) .collect(toCollection(() -> EnumSet.noneOf(AdhocTestOption.class))); assertEquals(actualOptions, expectedOptions); } } apache-arrow-java-19.0.0/flight/flight-sql/src/test/resources/000077500000000000000000000000001515451615200242055ustar00rootroot00000000000000apache-arrow-java-19.0.0/flight/flight-sql/src/test/resources/logback.xml000066400000000000000000000025221515451615200263320ustar00rootroot00000000000000 %d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n apache-arrow-java-19.0.0/flight/pom.xml000066400000000000000000000027561515451615200177020ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-java-root 19.0.0 arrow-flight pom Arrow Flight flight-core flight-sql flight-sql-jdbc-core flight-sql-jdbc-driver flight-integration-tests apache-arrow-java-19.0.0/format/000077500000000000000000000000001515451615200163665ustar00rootroot00000000000000apache-arrow-java-19.0.0/format/pom.xml000066400000000000000000000051361515451615200177100ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-java-root 19.0.0 arrow-format jar Arrow Format Generated Java files from the IPC Flatbuffer definitions. com.google.flatbuffers flatbuffers-java org.apache.maven.plugins maven-checkstyle-plugin true com.diffplug.spotless spotless-maven-plugin ${maven.multiModuleProjectDirectory}/dev/license/asf-java.license package org.apache.maven.plugins maven-compiler-plugin -Werror apache-arrow-java-19.0.0/format/src/000077500000000000000000000000001515451615200171555ustar00rootroot00000000000000apache-arrow-java-19.0.0/format/src/main/000077500000000000000000000000001515451615200201015ustar00rootroot00000000000000apache-arrow-java-19.0.0/format/src/main/java/000077500000000000000000000000001515451615200210225ustar00rootroot00000000000000apache-arrow-java-19.0.0/format/src/main/java/module-info.java000066400000000000000000000020741515451615200241060ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ // TODO(https://github.com/apache/arrow/issues/44037): Google hasn't reviewed Flatbuffers fix @SuppressWarnings({ "requires-automatic", "requires-transitive-automatic" }) module org.apache.arrow.format { exports org.apache.arrow.flatbuf; requires transitive flatbuffers.java; } apache-arrow-java-19.0.0/format/src/main/java/org/000077500000000000000000000000001515451615200216115ustar00rootroot00000000000000apache-arrow-java-19.0.0/format/src/main/java/org/apache/000077500000000000000000000000001515451615200230325ustar00rootroot00000000000000apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200241645ustar00rootroot00000000000000apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/000077500000000000000000000000001515451615200256075ustar00rootroot00000000000000apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Binary.java000066400000000000000000000051311515451615200276760ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * Opaque binary data */ @SuppressWarnings("unused") public final class Binary extends Table { 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); } public Binary __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public static void startBinary(FlatBufferBuilder builder) { builder.startTable(0); } public static int endBinary(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Binary get(int j) { return get(new Binary(), j); } public Binary get(Binary obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/BinaryView.java000066400000000000000000000061061515451615200305340ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * Logically the same as Binary, but the internal representation uses a view * struct that contains the string length and either the string's entire data * inline (for small strings) or an inlined prefix, an index of another buffer, * and an offset pointing to a slice in that buffer (for non-small strings). * * Since it uses a variable number of data buffers, each Field with this type * must have a corresponding entry in `variadicBufferCounts`. */ @SuppressWarnings("unused") public final class BinaryView extends Table { 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); } public BinaryView __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public static void startBinaryView(FlatBufferBuilder builder) { builder.startTable(0); } public static int endBinaryView(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public BinaryView get(int j) { return get(new BinaryView(), j); } public BinaryView get(BinaryView obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Block.java000066400000000000000000000053121515451615200275050ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; @SuppressWarnings("unused") public final class Block extends Struct { public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } public Block __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } /** * Index to the start of the RecordBlock (note this is past the Message header) */ public long offset() { return bb.getLong(bb_pos + 0); } /** * Length of the metadata */ public int metaDataLength() { return bb.getInt(bb_pos + 8); } /** * Length of the data (this is aligned so there can be a gap between this and * the metadata). */ public long bodyLength() { return bb.getLong(bb_pos + 16); } public static int createBlock(FlatBufferBuilder builder, long offset, int metaDataLength, long bodyLength) { builder.prep(8, 24); builder.putLong(bodyLength); builder.pad(4); builder.putInt(metaDataLength); builder.putLong(offset); return builder.offset(); } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Block get(int j) { return get(new Block(), j); } public Block get(Block obj, int j) { return obj.__assign(__element(j), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/BodyCompression.java000066400000000000000000000073061515451615200315770ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * Optional compression for the memory buffers constituting IPC message * bodies. Intended for use with RecordBatch but could be used for other * message types */ @SuppressWarnings("unused") public final class BodyCompression extends Table { 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); } public BodyCompression __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } /** * Compressor library. * For LZ4_FRAME, each compressed buffer must consist of a single frame. */ public byte codec() { int o = __offset(4); return o != 0 ? bb.get(o + bb_pos) : 0; } /** * Indicates the way the record batch body was compressed */ public byte method() { int o = __offset(6); return o != 0 ? bb.get(o + bb_pos) : 0; } public static int createBodyCompression(FlatBufferBuilder builder, byte codec, byte method) { builder.startTable(2); BodyCompression.addMethod(builder, method); BodyCompression.addCodec(builder, codec); return BodyCompression.endBodyCompression(builder); } public static void startBodyCompression(FlatBufferBuilder builder) { builder.startTable(2); } public static void addCodec(FlatBufferBuilder builder, byte codec) { builder.addByte(0, codec, 0); } public static void addMethod(FlatBufferBuilder builder, byte method) { builder.addByte(1, method, 0); } public static int endBodyCompression(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public BodyCompression get(int j) { return get(new BodyCompression(), j); } public BodyCompression get(BodyCompression obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/BodyCompressionMethod.java000066400000000000000000000034151515451615200327350ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; /** * Provided for forward compatibility in case we need to support different * strategies for compressing the IPC message body (like whole-body * compression rather than buffer-level) in the future */ @SuppressWarnings("unused") public final class BodyCompressionMethod { private BodyCompressionMethod() { } /** * Each constituent buffer is first compressed with the indicated * compressor, and then written with the uncompressed length in the first 8 * bytes as a 64-bit little-endian signed integer followed by the compressed * buffer bytes (and then padding as required by the protocol). The * uncompressed length may be set to -1 to indicate that the data that * follows is not compressed, which can be useful for cases where * compression does not yield appreciable savings. */ public static final byte BUFFER = 0; public static final String[] names = { "BUFFER", }; public static String name(int e) { return names[e]; } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Bool.java000066400000000000000000000050351515451615200273500ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; @SuppressWarnings("unused") public final class Bool extends Table { 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); } public Bool __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public static void startBool(FlatBufferBuilder builder) { builder.startTable(0); } public static int endBool(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Bool get(int j) { return get(new Bool(), j); } public Bool get(Bool obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Buffer.java000066400000000000000000000056261515451615200276740ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * ---------------------------------------------------------------------- * A Buffer represents a single contiguous memory segment */ @SuppressWarnings("unused") public final class Buffer extends Struct { public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } public Buffer __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } /** * The relative offset into the shared memory page where the bytes for this * buffer starts */ public long offset() { return bb.getLong(bb_pos + 0); } /** * The absolute length (in bytes) of the memory buffer. The memory is found * from offset (inclusive) to offset + length (non-inclusive). When building * messages using the encapsulated IPC message, padding bytes may be written * after a buffer, but such padding bytes do not need to be accounted for in * the size here. */ public long length() { return bb.getLong(bb_pos + 8); } public static int createBuffer(FlatBufferBuilder builder, long offset, long length) { builder.prep(8, 16); builder.putLong(length); builder.putLong(offset); return builder.offset(); } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Buffer get(int j) { return get(new Buffer(), j); } public Buffer get(Buffer obj, int j) { return obj.__assign(__element(j), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/CompressionType.java000066400000000000000000000021611515451615200316150ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; @SuppressWarnings("unused") public final class CompressionType { private CompressionType() { } public static final byte LZ4_FRAME = 0; public static final byte ZSTD = 1; public static final String[] names = { "LZ4_FRAME", "ZSTD", }; public static String name(int e) { return names[e]; } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Date.java000066400000000000000000000063561515451615200273410ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * Date is either a 32-bit or 64-bit signed integer type representing an * elapsed time since UNIX epoch (1970-01-01), stored in either of two units: * * * Milliseconds (64 bits) indicating UNIX time elapsed since the epoch (no * leap seconds), where the values are evenly divisible by 86400000 * * Days (32 bits) since the UNIX epoch */ @SuppressWarnings("unused") public final class Date extends Table { 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); } public Date __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public short unit() { int o = __offset(4); return o != 0 ? bb.getShort(o + bb_pos) : 1; } public static int createDate(FlatBufferBuilder builder, short unit) { builder.startTable(1); Date.addUnit(builder, unit); return Date.endDate(builder); } public static void startDate(FlatBufferBuilder builder) { builder.startTable(1); } public static void addUnit(FlatBufferBuilder builder, short unit) { builder.addShort(0, unit, 1); } public static int endDate(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Date get(int j) { return get(new Date(), j); } public Date get(Date obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/DateUnit.java000066400000000000000000000021471515451615200301730ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; @SuppressWarnings("unused") public final class DateUnit { private DateUnit() { } public static final short DAY = 0; public static final short MILLISECOND = 1; public static final String[] names = { "DAY", "MILLISECOND", }; public static String name(int e) { return names[e]; } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Decimal.java000066400000000000000000000077101515451615200300150ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * Exact decimal value represented as an integer value in two's * complement. Currently only 128-bit (16-byte) and 256-bit (32-byte) integers * are used. The representation uses the endianness indicated * in the Schema. */ @SuppressWarnings("unused") public final class Decimal extends Table { 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); } public Decimal __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } /** * Total number of decimal digits */ public int precision() { int o = __offset(4); return o != 0 ? bb.getInt(o + bb_pos) : 0; } /** * Number of digits after the decimal point "." */ public int scale() { int o = __offset(6); return o != 0 ? bb.getInt(o + bb_pos) : 0; } /** * Number of bits per value. The only accepted widths are 128 and 256. * We use bitWidth for consistency with Int::bitWidth. */ public int bitWidth() { int o = __offset(8); return o != 0 ? bb.getInt(o + bb_pos) : 128; } public static int createDecimal(FlatBufferBuilder builder, int precision, int scale, int bitWidth) { builder.startTable(3); Decimal.addBitWidth(builder, bitWidth); Decimal.addScale(builder, scale); Decimal.addPrecision(builder, precision); return Decimal.endDecimal(builder); } public static void startDecimal(FlatBufferBuilder builder) { builder.startTable(3); } public static void addPrecision(FlatBufferBuilder builder, int precision) { builder.addInt(0, precision, 0); } public static void addScale(FlatBufferBuilder builder, int scale) { builder.addInt(1, scale, 0); } public static void addBitWidth(FlatBufferBuilder builder, int bitWidth) { builder.addInt(2, bitWidth, 128); } public static int endDecimal(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Decimal get(int j) { return get(new Decimal(), j); } public Decimal get(Decimal obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/DictionaryBatch.java000066400000000000000000000105441515451615200315250ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * For sending dictionary encoding information. Any Field can be * dictionary-encoded, but in this case none of its children may be * dictionary-encoded. * There is one vector / column per dictionary, but that vector / column * may be spread across multiple dictionary batches by using the isDelta * flag */ @SuppressWarnings("unused") public final class DictionaryBatch extends Table { 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); } public DictionaryBatch __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public long id() { int o = __offset(4); return o != 0 ? bb.getLong(o + bb_pos) : 0L; } public org.apache.arrow.flatbuf.RecordBatch data() { return data(new org.apache.arrow.flatbuf.RecordBatch()); } public org.apache.arrow.flatbuf.RecordBatch data(org.apache.arrow.flatbuf.RecordBatch obj) { int o = __offset(6); return o != 0 ? obj.__assign(__indirect(o + bb_pos), bb) : null; } /** * If isDelta is true the values in the dictionary are to be appended to a * dictionary with the indicated id. If isDelta is false this dictionary * should replace the existing dictionary. */ public boolean isDelta() { int o = __offset(8); return o != 0 ? 0!=bb.get(o + bb_pos) : false; } public static int createDictionaryBatch(FlatBufferBuilder builder, long id, int dataOffset, boolean isDelta) { builder.startTable(3); DictionaryBatch.addId(builder, id); DictionaryBatch.addData(builder, dataOffset); DictionaryBatch.addIsDelta(builder, isDelta); return DictionaryBatch.endDictionaryBatch(builder); } public static void startDictionaryBatch(FlatBufferBuilder builder) { builder.startTable(3); } public static void addId(FlatBufferBuilder builder, long id) { builder.addLong(0, id, 0L); } public static void addData(FlatBufferBuilder builder, int dataOffset) { builder.addOffset(1, dataOffset, 0); } public static void addIsDelta(FlatBufferBuilder builder, boolean isDelta) { builder.addBoolean(2, isDelta, false); } public static int endDictionaryBatch(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public DictionaryBatch get(int j) { return get(new DictionaryBatch(), j); } public DictionaryBatch get(DictionaryBatch obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/DictionaryEncoding.java000066400000000000000000000121071515451615200322270ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; @SuppressWarnings("unused") public final class DictionaryEncoding extends Table { 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); } public DictionaryEncoding __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } /** * The known dictionary id in the application where this data is used. In * the file or streaming formats, the dictionary ids are found in the * DictionaryBatch messages */ public long id() { int o = __offset(4); return o != 0 ? bb.getLong(o + bb_pos) : 0L; } /** * The dictionary indices are constrained to be non-negative integers. If * this field is null, the indices must be signed int32. To maximize * cross-language compatibility and performance, implementations are * recommended to prefer signed integer types over unsigned integer types * and to avoid uint64 indices unless they are required by an application. */ public org.apache.arrow.flatbuf.Int indexType() { return indexType(new org.apache.arrow.flatbuf.Int()); } public org.apache.arrow.flatbuf.Int indexType(org.apache.arrow.flatbuf.Int obj) { int o = __offset(6); return o != 0 ? obj.__assign(__indirect(o + bb_pos), bb) : null; } /** * By default, dictionaries are not ordered, or the order does not have * semantic meaning. In some statistical, applications, dictionary-encoding * is used to represent ordered categorical data, and we provide a way to * preserve that metadata here */ public boolean isOrdered() { int o = __offset(8); return o != 0 ? 0!=bb.get(o + bb_pos) : false; } public short dictionaryKind() { int o = __offset(10); return o != 0 ? bb.getShort(o + bb_pos) : 0; } public static int createDictionaryEncoding(FlatBufferBuilder builder, long id, int indexTypeOffset, boolean isOrdered, short dictionaryKind) { builder.startTable(4); DictionaryEncoding.addId(builder, id); DictionaryEncoding.addIndexType(builder, indexTypeOffset); DictionaryEncoding.addDictionaryKind(builder, dictionaryKind); DictionaryEncoding.addIsOrdered(builder, isOrdered); return DictionaryEncoding.endDictionaryEncoding(builder); } public static void startDictionaryEncoding(FlatBufferBuilder builder) { builder.startTable(4); } public static void addId(FlatBufferBuilder builder, long id) { builder.addLong(0, id, 0L); } public static void addIndexType(FlatBufferBuilder builder, int indexTypeOffset) { builder.addOffset(1, indexTypeOffset, 0); } public static void addIsOrdered(FlatBufferBuilder builder, boolean isOrdered) { builder.addBoolean(2, isOrdered, false); } public static void addDictionaryKind(FlatBufferBuilder builder, short dictionaryKind) { builder.addShort(3, dictionaryKind, 0); } public static int endDictionaryEncoding(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public DictionaryEncoding get(int j) { return get(new DictionaryEncoding(), j); } public DictionaryEncoding get(DictionaryEncoding obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/DictionaryKind.java000066400000000000000000000025361515451615200313730ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; /** * ---------------------------------------------------------------------- * Dictionary encoding metadata * Maintained for forwards compatibility, in the future * Dictionaries might be explicit maps between integers and values * allowing for non-contiguous index values */ @SuppressWarnings("unused") public final class DictionaryKind { private DictionaryKind() { } public static final short DenseArray = 0; public static final String[] names = { "DenseArray", }; public static String name(int e) { return names[e]; } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Duration.java000066400000000000000000000057341515451615200302500ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; @SuppressWarnings("unused") public final class Duration extends Table { 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); } public Duration __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public short unit() { int o = __offset(4); return o != 0 ? bb.getShort(o + bb_pos) : 1; } public static int createDuration(FlatBufferBuilder builder, short unit) { builder.startTable(1); Duration.addUnit(builder, unit); return Duration.endDuration(builder); } public static void startDuration(FlatBufferBuilder builder) { builder.startTable(1); } public static void addUnit(FlatBufferBuilder builder, short unit) { builder.addShort(0, unit, 1); } public static int endDuration(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Duration get(int j) { return get(new Duration(), j); } public Duration get(Duration obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Endianness.java000066400000000000000000000023441515451615200305440ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; /** * ---------------------------------------------------------------------- * Endianness of the platform producing the data */ @SuppressWarnings("unused") public final class Endianness { private Endianness() { } public static final short Little = 0; public static final short Big = 1; public static final String[] names = { "Little", "Big", }; public static String name(int e) { return names[e]; } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Feature.java000066400000000000000000000042561515451615200300540ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; /** * Represents Arrow Features that might not have full support * within implementations. This is intended to be used in * two scenarios: * 1. A mechanism for readers of Arrow Streams * and files to understand that the stream or file makes * use of a feature that isn't supported or unknown to * the implementation (and therefore can meet the Arrow * forward compatibility guarantees). * 2. A means of negotiating between a client and server * what features a stream is allowed to use. The enums * values here are intented to represent higher level * features, additional details maybe negotiated * with key-value pairs specific to the protocol. * * Enums added to this list should be assigned power-of-two values * to facilitate exchanging and comparing bitmaps for supported * features. */ @SuppressWarnings("unused") public final class Feature { private Feature() { } /** * Needed to make flatbuffers happy. */ public static final long UNUSED = 0L; /** * The stream makes use of multiple full dictionaries with the * same ID and assumes clients implement dictionary replacement * correctly. */ public static final long DICTIONARY_REPLACEMENT = 1L; /** * The stream makes use of compressed bodies as described * in Message.fbs. */ public static final long COMPRESSED_BODY = 2L; } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Field.java000066400000000000000000000173101515451615200274770ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * ---------------------------------------------------------------------- * A field represents a named column in a record / row batch or child of a * nested type. */ @SuppressWarnings("unused") public final class Field extends Table { 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); } public Field __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } /** * Name is not required, in i.e. a List */ public String name() { int o = __offset(4); return o != 0 ? __string(o + bb_pos) : null; } public ByteBuffer nameAsByteBuffer() { return __vector_as_bytebuffer(4, 1); } public ByteBuffer nameInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 4, 1); } /** * Whether or not this field can contain nulls. Should be true in general. */ public boolean nullable() { int o = __offset(6); return o != 0 ? 0!=bb.get(o + bb_pos) : false; } public byte typeType() { int o = __offset(8); return o != 0 ? bb.get(o + bb_pos) : 0; } /** * This is the type of the decoded value if the field is dictionary encoded. */ public Table type(Table obj) { int o = __offset(10); return o != 0 ? __union(obj, o + bb_pos) : null; } /** * Present only if the field is dictionary encoded. */ public org.apache.arrow.flatbuf.DictionaryEncoding dictionary() { return dictionary(new org.apache.arrow.flatbuf.DictionaryEncoding()); } public org.apache.arrow.flatbuf.DictionaryEncoding dictionary(org.apache.arrow.flatbuf.DictionaryEncoding obj) { int o = __offset(12); return o != 0 ? obj.__assign(__indirect(o + bb_pos), bb) : null; } /** * children apply only to nested data types like Struct, List and Union. For * primitive types children will have length 0. */ public org.apache.arrow.flatbuf.Field children(int j) { return children(new org.apache.arrow.flatbuf.Field(), j); } public org.apache.arrow.flatbuf.Field children(org.apache.arrow.flatbuf.Field obj, int j) { int o = __offset(14); return o != 0 ? obj.__assign(__indirect(__vector(o) + j * 4), bb) : null; } public int childrenLength() { int o = __offset(14); return o != 0 ? __vector_len(o) : 0; } public org.apache.arrow.flatbuf.Field.Vector childrenVector() { return childrenVector(new org.apache.arrow.flatbuf.Field.Vector()); } public org.apache.arrow.flatbuf.Field.Vector childrenVector(org.apache.arrow.flatbuf.Field.Vector obj) { int o = __offset(14); return o != 0 ? obj.__assign(__vector(o), 4, bb) : null; } /** * User-defined metadata */ public org.apache.arrow.flatbuf.KeyValue customMetadata(int j) { return customMetadata(new org.apache.arrow.flatbuf.KeyValue(), j); } public org.apache.arrow.flatbuf.KeyValue customMetadata(org.apache.arrow.flatbuf.KeyValue obj, int j) { int o = __offset(16); return o != 0 ? obj.__assign(__indirect(__vector(o) + j * 4), bb) : null; } public int customMetadataLength() { int o = __offset(16); return o != 0 ? __vector_len(o) : 0; } public org.apache.arrow.flatbuf.KeyValue.Vector customMetadataVector() { return customMetadataVector(new org.apache.arrow.flatbuf.KeyValue.Vector()); } public org.apache.arrow.flatbuf.KeyValue.Vector customMetadataVector(org.apache.arrow.flatbuf.KeyValue.Vector obj) { int o = __offset(16); return o != 0 ? obj.__assign(__vector(o), 4, bb) : null; } public static int createField(FlatBufferBuilder builder, int nameOffset, boolean nullable, byte typeType, int typeOffset, int dictionaryOffset, int childrenOffset, int customMetadataOffset) { builder.startTable(7); Field.addCustomMetadata(builder, customMetadataOffset); Field.addChildren(builder, childrenOffset); Field.addDictionary(builder, dictionaryOffset); Field.addType(builder, typeOffset); Field.addName(builder, nameOffset); Field.addTypeType(builder, typeType); Field.addNullable(builder, nullable); return Field.endField(builder); } public static void startField(FlatBufferBuilder builder) { builder.startTable(7); } public static void addName(FlatBufferBuilder builder, int nameOffset) { builder.addOffset(0, nameOffset, 0); } public static void addNullable(FlatBufferBuilder builder, boolean nullable) { builder.addBoolean(1, nullable, false); } public static void addTypeType(FlatBufferBuilder builder, byte typeType) { builder.addByte(2, typeType, 0); } public static void addType(FlatBufferBuilder builder, int typeOffset) { builder.addOffset(3, typeOffset, 0); } public static void addDictionary(FlatBufferBuilder builder, int dictionaryOffset) { builder.addOffset(4, dictionaryOffset, 0); } public static void addChildren(FlatBufferBuilder builder, int childrenOffset) { builder.addOffset(5, childrenOffset, 0); } public static int createChildrenVector(FlatBufferBuilder builder, int[] data) { builder.startVector(4, data.length, 4); for (int i = data.length - 1; i >= 0; i--) builder.addOffset(data[i]); return builder.endVector(); } public static void startChildrenVector(FlatBufferBuilder builder, int numElems) { builder.startVector(4, numElems, 4); } public static void addCustomMetadata(FlatBufferBuilder builder, int customMetadataOffset) { builder.addOffset(6, customMetadataOffset, 0); } public static int createCustomMetadataVector(FlatBufferBuilder builder, int[] data) { builder.startVector(4, data.length, 4); for (int i = data.length - 1; i >= 0; i--) builder.addOffset(data[i]); return builder.endVector(); } public static void startCustomMetadataVector(FlatBufferBuilder builder, int numElems) { builder.startVector(4, numElems, 4); } public static int endField(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Field get(int j) { return get(new Field(), j); } public Field get(Field obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/FieldNode.java000066400000000000000000000062241515451615200303070ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * ---------------------------------------------------------------------- * Data structures for describing a table row batch (a collection of * equal-length Arrow arrays) * Metadata about a field at some level of a nested type tree (but not * its children). * * For example, a List with values `[[1, 2, 3], null, [4], [5, 6], null]` * would have {length: 5, null_count: 2} for its List node, and {length: 6, * null_count: 0} for its Int16 node, as separate FieldNode structs */ @SuppressWarnings("unused") public final class FieldNode extends Struct { public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } public FieldNode __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } /** * The number of value slots in the Arrow array at this level of a nested * tree */ public long length() { return bb.getLong(bb_pos + 0); } /** * The number of observed nulls. Fields with null_count == 0 may choose not * to write their physical validity bitmap out as a materialized buffer, * instead setting the length of the bitmap buffer to 0. */ public long nullCount() { return bb.getLong(bb_pos + 8); } public static int createFieldNode(FlatBufferBuilder builder, long length, long nullCount) { builder.prep(8, 16); builder.putLong(nullCount); builder.putLong(length); return builder.offset(); } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public FieldNode get(int j) { return get(new FieldNode(), j); } public FieldNode get(FieldNode obj, int j) { return obj.__assign(__element(j), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/FixedSizeBinary.java000066400000000000000000000062451515451615200315200ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; @SuppressWarnings("unused") public final class FixedSizeBinary extends Table { 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); } public FixedSizeBinary __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } /** * Number of bytes per value */ public int byteWidth() { int o = __offset(4); return o != 0 ? bb.getInt(o + bb_pos) : 0; } public static int createFixedSizeBinary(FlatBufferBuilder builder, int byteWidth) { builder.startTable(1); FixedSizeBinary.addByteWidth(builder, byteWidth); return FixedSizeBinary.endFixedSizeBinary(builder); } public static void startFixedSizeBinary(FlatBufferBuilder builder) { builder.startTable(1); } public static void addByteWidth(FlatBufferBuilder builder, int byteWidth) { builder.addInt(0, byteWidth, 0); } public static int endFixedSizeBinary(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public FixedSizeBinary get(int j) { return get(new FixedSizeBinary(), j); } public FixedSizeBinary get(FixedSizeBinary obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/FixedSizeList.java000066400000000000000000000061751515451615200312110ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; @SuppressWarnings("unused") public final class FixedSizeList extends Table { 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); } public FixedSizeList __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } /** * Number of list items per value */ public int listSize() { int o = __offset(4); return o != 0 ? bb.getInt(o + bb_pos) : 0; } public static int createFixedSizeList(FlatBufferBuilder builder, int listSize) { builder.startTable(1); FixedSizeList.addListSize(builder, listSize); return FixedSizeList.endFixedSizeList(builder); } public static void startFixedSizeList(FlatBufferBuilder builder) { builder.startTable(1); } public static void addListSize(FlatBufferBuilder builder, int listSize) { builder.addInt(0, listSize, 0); } public static int endFixedSizeList(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public FixedSizeList get(int j) { return get(new FixedSizeList(), j); } public FixedSizeList get(FixedSizeList obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/FloatingPoint.java000066400000000000000000000061361515451615200312350ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; @SuppressWarnings("unused") public final class FloatingPoint extends Table { 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); } public FloatingPoint __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public short precision() { int o = __offset(4); return o != 0 ? bb.getShort(o + bb_pos) : 0; } public static int createFloatingPoint(FlatBufferBuilder builder, short precision) { builder.startTable(1); FloatingPoint.addPrecision(builder, precision); return FloatingPoint.endFloatingPoint(builder); } public static void startFloatingPoint(FlatBufferBuilder builder) { builder.startTable(1); } public static void addPrecision(FlatBufferBuilder builder, short precision) { builder.addShort(0, precision, 0); } public static int endFloatingPoint(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public FloatingPoint get(int j) { return get(new FloatingPoint(), j); } public FloatingPoint get(FloatingPoint obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Footer.java000066400000000000000000000164531515451615200277210ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * ---------------------------------------------------------------------- * Arrow File metadata * */ @SuppressWarnings("unused") public final class Footer extends Table { 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); } public Footer __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public short version() { int o = __offset(4); return o != 0 ? bb.getShort(o + bb_pos) : 0; } public org.apache.arrow.flatbuf.Schema schema() { return schema(new org.apache.arrow.flatbuf.Schema()); } public org.apache.arrow.flatbuf.Schema schema(org.apache.arrow.flatbuf.Schema obj) { int o = __offset(6); return o != 0 ? obj.__assign(__indirect(o + bb_pos), bb) : null; } public org.apache.arrow.flatbuf.Block dictionaries(int j) { return dictionaries(new org.apache.arrow.flatbuf.Block(), j); } public org.apache.arrow.flatbuf.Block dictionaries(org.apache.arrow.flatbuf.Block obj, int j) { int o = __offset(8); return o != 0 ? obj.__assign(__vector(o) + j * 24, bb) : null; } public int dictionariesLength() { int o = __offset(8); return o != 0 ? __vector_len(o) : 0; } public org.apache.arrow.flatbuf.Block.Vector dictionariesVector() { return dictionariesVector(new org.apache.arrow.flatbuf.Block.Vector()); } public org.apache.arrow.flatbuf.Block.Vector dictionariesVector(org.apache.arrow.flatbuf.Block.Vector obj) { int o = __offset(8); return o != 0 ? obj.__assign(__vector(o), 24, bb) : null; } public org.apache.arrow.flatbuf.Block recordBatches(int j) { return recordBatches(new org.apache.arrow.flatbuf.Block(), j); } public org.apache.arrow.flatbuf.Block recordBatches(org.apache.arrow.flatbuf.Block obj, int j) { int o = __offset(10); return o != 0 ? obj.__assign(__vector(o) + j * 24, bb) : null; } public int recordBatchesLength() { int o = __offset(10); return o != 0 ? __vector_len(o) : 0; } public org.apache.arrow.flatbuf.Block.Vector recordBatchesVector() { return recordBatchesVector(new org.apache.arrow.flatbuf.Block.Vector()); } public org.apache.arrow.flatbuf.Block.Vector recordBatchesVector(org.apache.arrow.flatbuf.Block.Vector obj) { int o = __offset(10); return o != 0 ? obj.__assign(__vector(o), 24, bb) : null; } /** * User-defined metadata */ public org.apache.arrow.flatbuf.KeyValue customMetadata(int j) { return customMetadata(new org.apache.arrow.flatbuf.KeyValue(), j); } public org.apache.arrow.flatbuf.KeyValue customMetadata(org.apache.arrow.flatbuf.KeyValue obj, int j) { int o = __offset(12); return o != 0 ? obj.__assign(__indirect(__vector(o) + j * 4), bb) : null; } public int customMetadataLength() { int o = __offset(12); return o != 0 ? __vector_len(o) : 0; } public org.apache.arrow.flatbuf.KeyValue.Vector customMetadataVector() { return customMetadataVector(new org.apache.arrow.flatbuf.KeyValue.Vector()); } public org.apache.arrow.flatbuf.KeyValue.Vector customMetadataVector(org.apache.arrow.flatbuf.KeyValue.Vector obj) { int o = __offset(12); return o != 0 ? obj.__assign(__vector(o), 4, bb) : null; } public static int createFooter(FlatBufferBuilder builder, short version, int schemaOffset, int dictionariesOffset, int recordBatchesOffset, int customMetadataOffset) { builder.startTable(5); Footer.addCustomMetadata(builder, customMetadataOffset); Footer.addRecordBatches(builder, recordBatchesOffset); Footer.addDictionaries(builder, dictionariesOffset); Footer.addSchema(builder, schemaOffset); Footer.addVersion(builder, version); return Footer.endFooter(builder); } public static void startFooter(FlatBufferBuilder builder) { builder.startTable(5); } public static void addVersion(FlatBufferBuilder builder, short version) { builder.addShort(0, version, 0); } public static void addSchema(FlatBufferBuilder builder, int schemaOffset) { builder.addOffset(1, schemaOffset, 0); } public static void addDictionaries(FlatBufferBuilder builder, int dictionariesOffset) { builder.addOffset(2, dictionariesOffset, 0); } public static void startDictionariesVector(FlatBufferBuilder builder, int numElems) { builder.startVector(24, numElems, 8); } public static void addRecordBatches(FlatBufferBuilder builder, int recordBatchesOffset) { builder.addOffset(3, recordBatchesOffset, 0); } public static void startRecordBatchesVector(FlatBufferBuilder builder, int numElems) { builder.startVector(24, numElems, 8); } public static void addCustomMetadata(FlatBufferBuilder builder, int customMetadataOffset) { builder.addOffset(4, customMetadataOffset, 0); } public static int createCustomMetadataVector(FlatBufferBuilder builder, int[] data) { builder.startVector(4, data.length, 4); for (int i = data.length - 1; i >= 0; i--) builder.addOffset(data[i]); return builder.endVector(); } public static void startCustomMetadataVector(FlatBufferBuilder builder, int numElems) { builder.startVector(4, numElems, 4); } public static int endFooter(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static void finishFooterBuffer(FlatBufferBuilder builder, int offset) { builder.finish(offset); } public static void finishSizePrefixedFooterBuffer(FlatBufferBuilder builder, int offset) { builder.finishSizePrefixed(offset); } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Footer get(int j) { return get(new Footer(), j); } public Footer get(Footer obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Int.java000066400000000000000000000062551515451615200272140ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; @SuppressWarnings("unused") public final class Int extends Table { 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); } public Int __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public int bitWidth() { int o = __offset(4); return o != 0 ? bb.getInt(o + bb_pos) : 0; } public boolean isSigned() { int o = __offset(6); return o != 0 ? 0!=bb.get(o + bb_pos) : false; } public static int createInt(FlatBufferBuilder builder, int bitWidth, boolean isSigned) { builder.startTable(2); Int.addBitWidth(builder, bitWidth); Int.addIsSigned(builder, isSigned); return Int.endInt(builder); } public static void startInt(FlatBufferBuilder builder) { builder.startTable(2); } public static void addBitWidth(FlatBufferBuilder builder, int bitWidth) { builder.addInt(0, bitWidth, 0); } public static void addIsSigned(FlatBufferBuilder builder, boolean isSigned) { builder.addBoolean(1, isSigned, false); } public static int endInt(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Int get(int j) { return get(new Int(), j); } public Int get(Int obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Interval.java000066400000000000000000000057341515451615200302470ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; @SuppressWarnings("unused") public final class Interval extends Table { 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); } public Interval __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public short unit() { int o = __offset(4); return o != 0 ? bb.getShort(o + bb_pos) : 0; } public static int createInterval(FlatBufferBuilder builder, short unit) { builder.startTable(1); Interval.addUnit(builder, unit); return Interval.endInterval(builder); } public static void startInterval(FlatBufferBuilder builder) { builder.startTable(1); } public static void addUnit(FlatBufferBuilder builder, short unit) { builder.addShort(0, unit, 0); } public static int endInterval(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Interval get(int j) { return get(new Interval(), j); } public Interval get(Interval obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/IntervalUnit.java000066400000000000000000000022711515451615200311000ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; @SuppressWarnings("unused") public final class IntervalUnit { private IntervalUnit() { } public static final short YEAR_MONTH = 0; public static final short DAY_TIME = 1; public static final short MONTH_DAY_NANO = 2; public static final String[] names = { "YEAR_MONTH", "DAY_TIME", "MONTH_DAY_NANO", }; public static String name(int e) { return names[e]; } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/KeyValue.java000066400000000000000000000074561515451615200302130ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * ---------------------------------------------------------------------- * user defined key value pairs to add custom metadata to arrow * key namespacing is the responsibility of the user */ @SuppressWarnings("unused") public final class KeyValue extends Table { 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); } public KeyValue __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public String key() { int o = __offset(4); return o != 0 ? __string(o + bb_pos) : null; } public ByteBuffer keyAsByteBuffer() { return __vector_as_bytebuffer(4, 1); } public ByteBuffer keyInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 4, 1); } public String value() { int o = __offset(6); return o != 0 ? __string(o + bb_pos) : null; } public ByteBuffer valueAsByteBuffer() { return __vector_as_bytebuffer(6, 1); } public ByteBuffer valueInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 6, 1); } public static int createKeyValue(FlatBufferBuilder builder, int keyOffset, int valueOffset) { builder.startTable(2); KeyValue.addValue(builder, valueOffset); KeyValue.addKey(builder, keyOffset); return KeyValue.endKeyValue(builder); } public static void startKeyValue(FlatBufferBuilder builder) { builder.startTable(2); } public static void addKey(FlatBufferBuilder builder, int keyOffset) { builder.addOffset(0, keyOffset, 0); } public static void addValue(FlatBufferBuilder builder, int valueOffset) { builder.addOffset(1, valueOffset, 0); } public static int endKeyValue(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public KeyValue get(int j) { return get(new KeyValue(), j); } public KeyValue get(KeyValue obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/LargeBinary.java000066400000000000000000000053601515451615200306550ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * Same as Binary, but with 64-bit offsets, allowing to represent * extremely large data values. */ @SuppressWarnings("unused") public final class LargeBinary extends Table { 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); } public LargeBinary __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public static void startLargeBinary(FlatBufferBuilder builder) { builder.startTable(0); } public static int endLargeBinary(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public LargeBinary get(int j) { return get(new LargeBinary(), j); } public LargeBinary get(LargeBinary obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/LargeList.java000066400000000000000000000053201515451615200303400ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * Same as List, but with 64-bit offsets, allowing to represent * extremely large data values. */ @SuppressWarnings("unused") public final class LargeList extends Table { 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); } public LargeList __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public static void startLargeList(FlatBufferBuilder builder) { builder.startTable(0); } public static int endLargeList(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public LargeList get(int j) { return get(new LargeList(), j); } public LargeList get(LargeList obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/LargeListView.java000066400000000000000000000054321515451615200311770ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * Same as ListView, but with 64-bit offsets and sizes, allowing to represent * extremely large data values. */ @SuppressWarnings("unused") public final class LargeListView extends Table { 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); } public LargeListView __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public static void startLargeListView(FlatBufferBuilder builder) { builder.startTable(0); } public static int endLargeListView(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public LargeListView get(int j) { return get(new LargeListView(), j); } public LargeListView get(LargeListView obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/LargeUtf8.java000066400000000000000000000053201515451615200302530ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * Same as Utf8, but with 64-bit offsets, allowing to represent * extremely large data values. */ @SuppressWarnings("unused") public final class LargeUtf8 extends Table { 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); } public LargeUtf8 __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public static void startLargeUtf8(FlatBufferBuilder builder) { builder.startTable(0); } public static int endLargeUtf8(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public LargeUtf8 get(int j) { return get(new LargeUtf8(), j); } public LargeUtf8 get(LargeUtf8 obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/List.java000066400000000000000000000050351515451615200273700ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; @SuppressWarnings("unused") public final class List extends Table { 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); } public List __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public static void startList(FlatBufferBuilder builder) { builder.startTable(0); } public static int endList(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public List get(int j) { return get(new List(), j); } public List get(List obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/ListView.java000066400000000000000000000054121515451615200302220ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * Represents the same logical types that List can, but contains offsets and * sizes allowing for writes in any order and sharing of child values among * list values. */ @SuppressWarnings("unused") public final class ListView extends Table { 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); } public ListView __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public static void startListView(FlatBufferBuilder builder) { builder.startTable(0); } public static int endListView(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public ListView get(int j) { return get(new ListView(), j); } public ListView get(ListView obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Map.java000066400000000000000000000100101515451615200271570ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * A Map is a logical nested type that is represented as * * List> * * In this layout, the keys and values are each respectively contiguous. We do * not constrain the key and value types, so the application is responsible * for ensuring that the keys are hashable and unique. Whether the keys are sorted * may be set in the metadata for this field. * * In a field with Map type, the field has a child Struct field, which then * has two children: key type and the second the value type. The names of the * child fields may be respectively "entries", "key", and "value", but this is * not enforced. * * Map * ```text * - child[0] entries: Struct * - child[0] key: K * - child[1] value: V * ``` * Neither the "entries" field nor the "key" field may be nullable. * * The metadata is structured so that Arrow systems without special handling * for Map can make Map an alias for List. The "layout" attribute for the Map * field must have the same contents as a List. */ @SuppressWarnings("unused") public final class Map extends Table { 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); } public Map __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } /** * Set to true if the keys within each value are sorted */ public boolean keysSorted() { int o = __offset(4); return o != 0 ? 0!=bb.get(o + bb_pos) : false; } public static int createMap(FlatBufferBuilder builder, boolean keysSorted) { builder.startTable(1); Map.addKeysSorted(builder, keysSorted); return Map.endMap(builder); } public static void startMap(FlatBufferBuilder builder) { builder.startTable(1); } public static void addKeysSorted(FlatBufferBuilder builder, boolean keysSorted) { builder.addBoolean(0, keysSorted, false); } public static int endMap(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Map get(int j) { return get(new Map(), j); } public Map get(Map obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Message.java000066400000000000000000000126331515451615200300430ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; @SuppressWarnings("unused") public final class Message extends Table { 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); } public Message __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public short version() { int o = __offset(4); return o != 0 ? bb.getShort(o + bb_pos) : 0; } public byte headerType() { int o = __offset(6); return o != 0 ? bb.get(o + bb_pos) : 0; } public Table header(Table obj) { int o = __offset(8); return o != 0 ? __union(obj, o + bb_pos) : null; } public long bodyLength() { int o = __offset(10); return o != 0 ? bb.getLong(o + bb_pos) : 0L; } public org.apache.arrow.flatbuf.KeyValue customMetadata(int j) { return customMetadata(new org.apache.arrow.flatbuf.KeyValue(), j); } public org.apache.arrow.flatbuf.KeyValue customMetadata(org.apache.arrow.flatbuf.KeyValue obj, int j) { int o = __offset(12); return o != 0 ? obj.__assign(__indirect(__vector(o) + j * 4), bb) : null; } public int customMetadataLength() { int o = __offset(12); return o != 0 ? __vector_len(o) : 0; } public org.apache.arrow.flatbuf.KeyValue.Vector customMetadataVector() { return customMetadataVector(new org.apache.arrow.flatbuf.KeyValue.Vector()); } public org.apache.arrow.flatbuf.KeyValue.Vector customMetadataVector(org.apache.arrow.flatbuf.KeyValue.Vector obj) { int o = __offset(12); return o != 0 ? obj.__assign(__vector(o), 4, bb) : null; } public static int createMessage(FlatBufferBuilder builder, short version, byte headerType, int headerOffset, long bodyLength, int customMetadataOffset) { builder.startTable(5); Message.addBodyLength(builder, bodyLength); Message.addCustomMetadata(builder, customMetadataOffset); Message.addHeader(builder, headerOffset); Message.addVersion(builder, version); Message.addHeaderType(builder, headerType); return Message.endMessage(builder); } public static void startMessage(FlatBufferBuilder builder) { builder.startTable(5); } public static void addVersion(FlatBufferBuilder builder, short version) { builder.addShort(0, version, 0); } public static void addHeaderType(FlatBufferBuilder builder, byte headerType) { builder.addByte(1, headerType, 0); } public static void addHeader(FlatBufferBuilder builder, int headerOffset) { builder.addOffset(2, headerOffset, 0); } public static void addBodyLength(FlatBufferBuilder builder, long bodyLength) { builder.addLong(3, bodyLength, 0L); } public static void addCustomMetadata(FlatBufferBuilder builder, int customMetadataOffset) { builder.addOffset(4, customMetadataOffset, 0); } public static int createCustomMetadataVector(FlatBufferBuilder builder, int[] data) { builder.startVector(4, data.length, 4); for (int i = data.length - 1; i >= 0; i--) builder.addOffset(data[i]); return builder.endVector(); } public static void startCustomMetadataVector(FlatBufferBuilder builder, int numElems) { builder.startVector(4, numElems, 4); } public static int endMessage(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static void finishMessageBuffer(FlatBufferBuilder builder, int offset) { builder.finish(offset); } public static void finishSizePrefixedMessageBuffer(FlatBufferBuilder builder, int offset) { builder.finishSizePrefixed(offset); } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Message get(int j) { return get(new Message(), j); } public Message get(Message obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/MessageHeader.java000066400000000000000000000034341515451615200311530ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; /** * ---------------------------------------------------------------------- * The root Message type * This union enables us to easily send different message types without * redundant storage, and in the future we can easily add new message types. * * Arrow implementations do not need to implement all of the message types, * which may include experimental metadata types. For maximum compatibility, * it is best to send data using RecordBatch */ @SuppressWarnings("unused") public final class MessageHeader { private MessageHeader() { } public static final byte NONE = 0; public static final byte Schema = 1; public static final byte DictionaryBatch = 2; public static final byte RecordBatch = 3; public static final byte Tensor = 4; public static final byte SparseTensor = 5; public static final String[] names = { "NONE", "Schema", "DictionaryBatch", "RecordBatch", "Tensor", "SparseTensor", }; public static String name(int e) { return names[e]; } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/MetadataVersion.java000066400000000000000000000035601515451615200315440ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; @SuppressWarnings("unused") public final class MetadataVersion { private MetadataVersion() { } /** * 0.1.0 (October 2016). */ public static final short V1 = 0; /** * 0.2.0 (February 2017). Non-backwards compatible with V1. */ public static final short V2 = 1; /** * 0.3.0 -> 0.7.1 (May - December 2017). Non-backwards compatible with V2. */ public static final short V3 = 2; /** * >= 0.8.0 (December 2017). Non-backwards compatible with V3. */ public static final short V4 = 3; /** * >= 1.0.0 (July 2020). Backwards compatible with V4 (V5 readers can read V4 * metadata and IPC messages). Implementations are recommended to provide a * V4 compatibility mode with V5 format changes disabled. * * Incompatible changes between V4 and V5: * - Union buffer layout has changed. In V5, Unions don't have a validity * bitmap buffer. */ public static final short V5 = 4; public static final String[] names = { "V1", "V2", "V3", "V4", "V5", }; public static String name(int e) { return names[e]; } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Null.java000066400000000000000000000051431515451615200273670ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * These are stored in the flatbuffer in the Type union below */ @SuppressWarnings("unused") public final class Null extends Table { 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); } public Null __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public static void startNull(FlatBufferBuilder builder) { builder.startTable(0); } public static int endNull(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Null get(int j) { return get(new Null(), j); } public Null get(Null obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Precision.java000066400000000000000000000022231515451615200304040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; @SuppressWarnings("unused") public final class Precision { private Precision() { } public static final short HALF = 0; public static final short SINGLE = 1; public static final short DOUBLE = 2; public static final String[] names = { "HALF", "SINGLE", "DOUBLE", }; public static String name(int e) { return names[e]; } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/RecordBatch.java000066400000000000000000000207021515451615200306330ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * A data header describing the shared memory layout of a "record" or "row" * batch. Some systems call this a "row batch" internally and others a "record * batch". */ @SuppressWarnings("unused") public final class RecordBatch extends Table { 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); } public RecordBatch __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } /** * number of records / rows. The arrays in the batch should all have this * length */ public long length() { int o = __offset(4); return o != 0 ? bb.getLong(o + bb_pos) : 0L; } /** * Nodes correspond to the pre-ordered flattened logical schema */ public org.apache.arrow.flatbuf.FieldNode nodes(int j) { return nodes(new org.apache.arrow.flatbuf.FieldNode(), j); } public org.apache.arrow.flatbuf.FieldNode nodes(org.apache.arrow.flatbuf.FieldNode obj, int j) { int o = __offset(6); return o != 0 ? obj.__assign(__vector(o) + j * 16, bb) : null; } public int nodesLength() { int o = __offset(6); return o != 0 ? __vector_len(o) : 0; } public org.apache.arrow.flatbuf.FieldNode.Vector nodesVector() { return nodesVector(new org.apache.arrow.flatbuf.FieldNode.Vector()); } public org.apache.arrow.flatbuf.FieldNode.Vector nodesVector(org.apache.arrow.flatbuf.FieldNode.Vector obj) { int o = __offset(6); return o != 0 ? obj.__assign(__vector(o), 16, bb) : null; } /** * Buffers correspond to the pre-ordered flattened buffer tree * * The number of buffers appended to this list depends on the schema. For * example, most primitive arrays will have 2 buffers, 1 for the validity * bitmap and 1 for the values. For struct arrays, there will only be a * single buffer for the validity (nulls) bitmap */ public org.apache.arrow.flatbuf.Buffer buffers(int j) { return buffers(new org.apache.arrow.flatbuf.Buffer(), j); } public org.apache.arrow.flatbuf.Buffer buffers(org.apache.arrow.flatbuf.Buffer obj, int j) { int o = __offset(8); return o != 0 ? obj.__assign(__vector(o) + j * 16, bb) : null; } public int buffersLength() { int o = __offset(8); return o != 0 ? __vector_len(o) : 0; } public org.apache.arrow.flatbuf.Buffer.Vector buffersVector() { return buffersVector(new org.apache.arrow.flatbuf.Buffer.Vector()); } public org.apache.arrow.flatbuf.Buffer.Vector buffersVector(org.apache.arrow.flatbuf.Buffer.Vector obj) { int o = __offset(8); return o != 0 ? obj.__assign(__vector(o), 16, bb) : null; } /** * Optional compression of the message body */ public org.apache.arrow.flatbuf.BodyCompression compression() { return compression(new org.apache.arrow.flatbuf.BodyCompression()); } public org.apache.arrow.flatbuf.BodyCompression compression(org.apache.arrow.flatbuf.BodyCompression obj) { int o = __offset(10); return o != 0 ? obj.__assign(__indirect(o + bb_pos), bb) : null; } /** * Some types such as Utf8View are represented using a variable number of buffers. * For each such Field in the pre-ordered flattened logical schema, there will be * an entry in variadicBufferCounts to indicate the number of number of variadic * buffers which belong to that Field in the current RecordBatch. * * For example, the schema * col1: Struct * col2: Utf8View * contains two Fields with variadic buffers so variadicBufferCounts will have * two entries, the first counting the variadic buffers of `col1.beta` and the * second counting `col2`'s. * * This field may be omitted if and only if the schema contains no Fields with * a variable number of buffers, such as BinaryView and Utf8View. */ public long variadicBufferCounts(int j) { int o = __offset(12); return o != 0 ? bb.getLong(__vector(o) + j * 8) : 0; } public int variadicBufferCountsLength() { int o = __offset(12); return o != 0 ? __vector_len(o) : 0; } public LongVector variadicBufferCountsVector() { return variadicBufferCountsVector(new LongVector()); } public LongVector variadicBufferCountsVector(LongVector obj) { int o = __offset(12); return o != 0 ? obj.__assign(__vector(o), bb) : null; } public ByteBuffer variadicBufferCountsAsByteBuffer() { return __vector_as_bytebuffer(12, 8); } public ByteBuffer variadicBufferCountsInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 12, 8); } public static int createRecordBatch(FlatBufferBuilder builder, long length, int nodesOffset, int buffersOffset, int compressionOffset, int variadicBufferCountsOffset) { builder.startTable(5); RecordBatch.addLength(builder, length); RecordBatch.addVariadicBufferCounts(builder, variadicBufferCountsOffset); RecordBatch.addCompression(builder, compressionOffset); RecordBatch.addBuffers(builder, buffersOffset); RecordBatch.addNodes(builder, nodesOffset); return RecordBatch.endRecordBatch(builder); } public static void startRecordBatch(FlatBufferBuilder builder) { builder.startTable(5); } public static void addLength(FlatBufferBuilder builder, long length) { builder.addLong(0, length, 0L); } public static void addNodes(FlatBufferBuilder builder, int nodesOffset) { builder.addOffset(1, nodesOffset, 0); } public static void startNodesVector(FlatBufferBuilder builder, int numElems) { builder.startVector(16, numElems, 8); } public static void addBuffers(FlatBufferBuilder builder, int buffersOffset) { builder.addOffset(2, buffersOffset, 0); } public static void startBuffersVector(FlatBufferBuilder builder, int numElems) { builder.startVector(16, numElems, 8); } public static void addCompression(FlatBufferBuilder builder, int compressionOffset) { builder.addOffset(3, compressionOffset, 0); } public static void addVariadicBufferCounts(FlatBufferBuilder builder, int variadicBufferCountsOffset) { builder.addOffset(4, variadicBufferCountsOffset, 0); } public static int createVariadicBufferCountsVector(FlatBufferBuilder builder, long[] data) { builder.startVector(8, data.length, 8); for (int i = data.length - 1; i >= 0; i--) builder.addLong(data[i]); return builder.endVector(); } public static void startVariadicBufferCountsVector(FlatBufferBuilder builder, int numElems) { builder.startVector(8, numElems, 8); } public static int endRecordBatch(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public RecordBatch get(int j) { return get(new RecordBatch(), j); } public RecordBatch get(RecordBatch obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/RunEndEncoded.java000066400000000000000000000057331515451615200311370ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * Contains two child arrays, run_ends and values. * The run_ends child array must be a 16/32/64-bit integer array * which encodes the indices at which the run with the value in * each corresponding index in the values child array ends. * Like list/struct types, the value array can be of any type. */ @SuppressWarnings("unused") public final class RunEndEncoded extends Table { 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); } public RunEndEncoded __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public static void startRunEndEncoded(FlatBufferBuilder builder) { builder.startTable(0); } public static int endRunEndEncoded(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public RunEndEncoded get(int j) { return get(new RunEndEncoded(), j); } public RunEndEncoded get(RunEndEncoded obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Schema.java000066400000000000000000000163541515451615200276630ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * ---------------------------------------------------------------------- * A Schema describes the columns in a row batch */ @SuppressWarnings("unused") public final class Schema extends Table { 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); } public Schema __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } /** * endianness of the buffer * it is Little Endian by default * if endianness doesn't match the underlying system then the vectors need to be converted */ public short endianness() { int o = __offset(4); return o != 0 ? bb.getShort(o + bb_pos) : 0; } public org.apache.arrow.flatbuf.Field fields(int j) { return fields(new org.apache.arrow.flatbuf.Field(), j); } public org.apache.arrow.flatbuf.Field fields(org.apache.arrow.flatbuf.Field obj, int j) { int o = __offset(6); return o != 0 ? obj.__assign(__indirect(__vector(o) + j * 4), bb) : null; } public int fieldsLength() { int o = __offset(6); return o != 0 ? __vector_len(o) : 0; } public org.apache.arrow.flatbuf.Field.Vector fieldsVector() { return fieldsVector(new org.apache.arrow.flatbuf.Field.Vector()); } public org.apache.arrow.flatbuf.Field.Vector fieldsVector(org.apache.arrow.flatbuf.Field.Vector obj) { int o = __offset(6); return o != 0 ? obj.__assign(__vector(o), 4, bb) : null; } public org.apache.arrow.flatbuf.KeyValue customMetadata(int j) { return customMetadata(new org.apache.arrow.flatbuf.KeyValue(), j); } public org.apache.arrow.flatbuf.KeyValue customMetadata(org.apache.arrow.flatbuf.KeyValue obj, int j) { int o = __offset(8); return o != 0 ? obj.__assign(__indirect(__vector(o) + j * 4), bb) : null; } public int customMetadataLength() { int o = __offset(8); return o != 0 ? __vector_len(o) : 0; } public org.apache.arrow.flatbuf.KeyValue.Vector customMetadataVector() { return customMetadataVector(new org.apache.arrow.flatbuf.KeyValue.Vector()); } public org.apache.arrow.flatbuf.KeyValue.Vector customMetadataVector(org.apache.arrow.flatbuf.KeyValue.Vector obj) { int o = __offset(8); return o != 0 ? obj.__assign(__vector(o), 4, bb) : null; } /** * Features used in the stream/file. */ public long features(int j) { int o = __offset(10); return o != 0 ? bb.getLong(__vector(o) + j * 8) : 0; } public int featuresLength() { int o = __offset(10); return o != 0 ? __vector_len(o) : 0; } public LongVector featuresVector() { return featuresVector(new LongVector()); } public LongVector featuresVector(LongVector obj) { int o = __offset(10); return o != 0 ? obj.__assign(__vector(o), bb) : null; } public ByteBuffer featuresAsByteBuffer() { return __vector_as_bytebuffer(10, 8); } public ByteBuffer featuresInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 10, 8); } public static int createSchema(FlatBufferBuilder builder, short endianness, int fieldsOffset, int customMetadataOffset, int featuresOffset) { builder.startTable(4); Schema.addFeatures(builder, featuresOffset); Schema.addCustomMetadata(builder, customMetadataOffset); Schema.addFields(builder, fieldsOffset); Schema.addEndianness(builder, endianness); return Schema.endSchema(builder); } public static void startSchema(FlatBufferBuilder builder) { builder.startTable(4); } public static void addEndianness(FlatBufferBuilder builder, short endianness) { builder.addShort(0, endianness, 0); } public static void addFields(FlatBufferBuilder builder, int fieldsOffset) { builder.addOffset(1, fieldsOffset, 0); } public static int createFieldsVector(FlatBufferBuilder builder, int[] data) { builder.startVector(4, data.length, 4); for (int i = data.length - 1; i >= 0; i--) builder.addOffset(data[i]); return builder.endVector(); } public static void startFieldsVector(FlatBufferBuilder builder, int numElems) { builder.startVector(4, numElems, 4); } public static void addCustomMetadata(FlatBufferBuilder builder, int customMetadataOffset) { builder.addOffset(2, customMetadataOffset, 0); } public static int createCustomMetadataVector(FlatBufferBuilder builder, int[] data) { builder.startVector(4, data.length, 4); for (int i = data.length - 1; i >= 0; i--) builder.addOffset(data[i]); return builder.endVector(); } public static void startCustomMetadataVector(FlatBufferBuilder builder, int numElems) { builder.startVector(4, numElems, 4); } public static void addFeatures(FlatBufferBuilder builder, int featuresOffset) { builder.addOffset(3, featuresOffset, 0); } public static int createFeaturesVector(FlatBufferBuilder builder, long[] data) { builder.startVector(8, data.length, 8); for (int i = data.length - 1; i >= 0; i--) builder.addLong(data[i]); return builder.endVector(); } public static void startFeaturesVector(FlatBufferBuilder builder, int numElems) { builder.startVector(8, numElems, 8); } public static int endSchema(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static void finishSchemaBuffer(FlatBufferBuilder builder, int offset) { builder.finish(offset); } public static void finishSizePrefixedSchemaBuffer(FlatBufferBuilder builder, int offset) { builder.finishSizePrefixed(offset); } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Schema get(int j) { return get(new Schema(), j); } public Schema get(Schema obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } SparseMatrixCompressedAxis.java000066400000000000000000000022011515451615200336620ustar00rootroot00000000000000apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; @SuppressWarnings("unused") public final class SparseMatrixCompressedAxis { private SparseMatrixCompressedAxis() { } public static final short Row = 0; public static final short Column = 1; public static final String[] names = { "Row", "Column", }; public static String name(int e) { return names[e]; } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/SparseMatrixIndexCSX.java000066400000000000000000000142241515451615200324450ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * Compressed Sparse format, that is matrix-specific. */ @SuppressWarnings("unused") public final class SparseMatrixIndexCSX extends Table { 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); } public SparseMatrixIndexCSX __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } /** * Which axis, row or column, is compressed */ public short compressedAxis() { int o = __offset(4); return o != 0 ? bb.getShort(o + bb_pos) : 0; } /** * The type of values in indptrBuffer */ public org.apache.arrow.flatbuf.Int indptrType() { return indptrType(new org.apache.arrow.flatbuf.Int()); } public org.apache.arrow.flatbuf.Int indptrType(org.apache.arrow.flatbuf.Int obj) { int o = __offset(6); return o != 0 ? obj.__assign(__indirect(o + bb_pos), bb) : null; } /** * indptrBuffer stores the location and size of indptr array that * represents the range of the rows. * The i-th row spans from `indptr[i]` to `indptr[i+1]` in the data. * The length of this array is 1 + (the number of rows), and the type * of index value is long. * * For example, let X be the following 6x4 matrix: * ```text * X := [[0, 1, 2, 0], * [0, 0, 3, 0], * [0, 4, 0, 5], * [0, 0, 0, 0], * [6, 0, 7, 8], * [0, 9, 0, 0]]. * ``` * The array of non-zero values in X is: * ```text * values(X) = [1, 2, 3, 4, 5, 6, 7, 8, 9]. * ``` * And the indptr of X is: * ```text * indptr(X) = [0, 2, 3, 5, 5, 8, 10]. * ``` */ public org.apache.arrow.flatbuf.Buffer indptrBuffer() { return indptrBuffer(new org.apache.arrow.flatbuf.Buffer()); } public org.apache.arrow.flatbuf.Buffer indptrBuffer(org.apache.arrow.flatbuf.Buffer obj) { int o = __offset(8); return o != 0 ? obj.__assign(o + bb_pos, bb) : null; } /** * The type of values in indicesBuffer */ public org.apache.arrow.flatbuf.Int indicesType() { return indicesType(new org.apache.arrow.flatbuf.Int()); } public org.apache.arrow.flatbuf.Int indicesType(org.apache.arrow.flatbuf.Int obj) { int o = __offset(10); return o != 0 ? obj.__assign(__indirect(o + bb_pos), bb) : null; } /** * indicesBuffer stores the location and size of the array that * contains the column indices of the corresponding non-zero values. * The type of index value is long. * * For example, the indices of the above X is: * ```text * indices(X) = [1, 2, 2, 1, 3, 0, 2, 3, 1]. * ``` * Note that the indices are sorted in lexicographical order for each row. */ public org.apache.arrow.flatbuf.Buffer indicesBuffer() { return indicesBuffer(new org.apache.arrow.flatbuf.Buffer()); } public org.apache.arrow.flatbuf.Buffer indicesBuffer(org.apache.arrow.flatbuf.Buffer obj) { int o = __offset(12); return o != 0 ? obj.__assign(o + bb_pos, bb) : null; } public static void startSparseMatrixIndexCSX(FlatBufferBuilder builder) { builder.startTable(5); } public static void addCompressedAxis(FlatBufferBuilder builder, short compressedAxis) { builder.addShort(0, compressedAxis, 0); } public static void addIndptrType(FlatBufferBuilder builder, int indptrTypeOffset) { builder.addOffset(1, indptrTypeOffset, 0); } public static void addIndptrBuffer(FlatBufferBuilder builder, int indptrBufferOffset) { builder.addStruct(2, indptrBufferOffset, 0); } public static void addIndicesType(FlatBufferBuilder builder, int indicesTypeOffset) { builder.addOffset(3, indicesTypeOffset, 0); } public static void addIndicesBuffer(FlatBufferBuilder builder, int indicesBufferOffset) { builder.addStruct(4, indicesBufferOffset, 0); } public static int endSparseMatrixIndexCSX(FlatBufferBuilder builder) { int o = builder.endTable(); builder.required(o, 6); // indptrType builder.required(o, 8); // indptrBuffer builder.required(o, 10); // indicesType builder.required(o, 12); // indicesBuffer return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public SparseMatrixIndexCSX get(int j) { return get(new SparseMatrixIndexCSX(), j); } public SparseMatrixIndexCSX get(SparseMatrixIndexCSX obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/SparseTensor.java000066400000000000000000000140641515451615200311070ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; @SuppressWarnings("unused") public final class SparseTensor extends Table { 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); } public SparseTensor __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public byte typeType() { int o = __offset(4); return o != 0 ? bb.get(o + bb_pos) : 0; } /** * The type of data contained in a value cell. * Currently only fixed-width value types are supported, * no strings or nested types. */ public Table type(Table obj) { int o = __offset(6); return o != 0 ? __union(obj, o + bb_pos) : null; } /** * The dimensions of the tensor, optionally named. */ public org.apache.arrow.flatbuf.TensorDim shape(int j) { return shape(new org.apache.arrow.flatbuf.TensorDim(), j); } public org.apache.arrow.flatbuf.TensorDim shape(org.apache.arrow.flatbuf.TensorDim obj, int j) { int o = __offset(8); return o != 0 ? obj.__assign(__indirect(__vector(o) + j * 4), bb) : null; } public int shapeLength() { int o = __offset(8); return o != 0 ? __vector_len(o) : 0; } public org.apache.arrow.flatbuf.TensorDim.Vector shapeVector() { return shapeVector(new org.apache.arrow.flatbuf.TensorDim.Vector()); } public org.apache.arrow.flatbuf.TensorDim.Vector shapeVector(org.apache.arrow.flatbuf.TensorDim.Vector obj) { int o = __offset(8); return o != 0 ? obj.__assign(__vector(o), 4, bb) : null; } /** * The number of non-zero values in a sparse tensor. */ public long nonZeroLength() { int o = __offset(10); return o != 0 ? bb.getLong(o + bb_pos) : 0L; } public byte sparseIndexType() { int o = __offset(12); return o != 0 ? bb.get(o + bb_pos) : 0; } /** * Sparse tensor index */ public Table sparseIndex(Table obj) { int o = __offset(14); return o != 0 ? __union(obj, o + bb_pos) : null; } /** * The location and size of the tensor's data */ public org.apache.arrow.flatbuf.Buffer data() { return data(new org.apache.arrow.flatbuf.Buffer()); } public org.apache.arrow.flatbuf.Buffer data(org.apache.arrow.flatbuf.Buffer obj) { int o = __offset(16); return o != 0 ? obj.__assign(o + bb_pos, bb) : null; } public static void startSparseTensor(FlatBufferBuilder builder) { builder.startTable(7); } public static void addTypeType(FlatBufferBuilder builder, byte typeType) { builder.addByte(0, typeType, 0); } public static void addType(FlatBufferBuilder builder, int typeOffset) { builder.addOffset(1, typeOffset, 0); } public static void addShape(FlatBufferBuilder builder, int shapeOffset) { builder.addOffset(2, shapeOffset, 0); } public static int createShapeVector(FlatBufferBuilder builder, int[] data) { builder.startVector(4, data.length, 4); for (int i = data.length - 1; i >= 0; i--) builder.addOffset(data[i]); return builder.endVector(); } public static void startShapeVector(FlatBufferBuilder builder, int numElems) { builder.startVector(4, numElems, 4); } public static void addNonZeroLength(FlatBufferBuilder builder, long nonZeroLength) { builder.addLong(3, nonZeroLength, 0L); } public static void addSparseIndexType(FlatBufferBuilder builder, byte sparseIndexType) { builder.addByte(4, sparseIndexType, 0); } public static void addSparseIndex(FlatBufferBuilder builder, int sparseIndexOffset) { builder.addOffset(5, sparseIndexOffset, 0); } public static void addData(FlatBufferBuilder builder, int dataOffset) { builder.addStruct(6, dataOffset, 0); } public static int endSparseTensor(FlatBufferBuilder builder) { int o = builder.endTable(); builder.required(o, 6); // type builder.required(o, 8); // shape builder.required(o, 14); // sparseIndex builder.required(o, 16); // data return o; } public static void finishSparseTensorBuffer(FlatBufferBuilder builder, int offset) { builder.finish(offset); } public static void finishSizePrefixedSparseTensorBuffer(FlatBufferBuilder builder, int offset) { builder.finishSizePrefixed(offset); } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public SparseTensor get(int j) { return get(new SparseTensor(), j); } public SparseTensor get(SparseTensor obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/SparseTensorIndex.java000066400000000000000000000024451515451615200320770ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; @SuppressWarnings("unused") public final class SparseTensorIndex { private SparseTensorIndex() { } public static final byte NONE = 0; public static final byte SparseTensorIndexCOO = 1; public static final byte SparseMatrixIndexCSX = 2; public static final byte SparseTensorIndexCSF = 3; public static final String[] names = { "NONE", "SparseTensorIndexCOO", "SparseMatrixIndexCSX", "SparseTensorIndexCSF", }; public static String name(int e) { return names[e]; } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/SparseTensorIndexCOO.java000066400000000000000000000154421515451615200324410ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * ---------------------------------------------------------------------- * EXPERIMENTAL: Data structures for sparse tensors * Coordinate (COO) format of sparse tensor index. * * COO's index list are represented as a NxM matrix, * where N is the number of non-zero values, * and M is the number of dimensions of a sparse tensor. * * indicesBuffer stores the location and size of the data of this indices * matrix. The value type and the stride of the indices matrix is * specified in indicesType and indicesStrides fields. * * For example, let X be a 2x3x4x5 tensor, and it has the following * 6 non-zero values: * ```text * X[0, 1, 2, 0] := 1 * X[1, 1, 2, 3] := 2 * X[0, 2, 1, 0] := 3 * X[0, 1, 3, 0] := 4 * X[0, 1, 2, 1] := 5 * X[1, 2, 0, 4] := 6 * ``` * In COO format, the index matrix of X is the following 4x6 matrix: * ```text * [[0, 0, 0, 0, 1, 1], * [1, 1, 1, 2, 1, 2], * [2, 2, 3, 1, 2, 0], * [0, 1, 0, 0, 3, 4]] * ``` * When isCanonical is true, the indices is sorted in lexicographical order * (row-major order), and it does not have duplicated entries. Otherwise, * the indices may not be sorted, or may have duplicated entries. */ @SuppressWarnings("unused") public final class SparseTensorIndexCOO extends Table { 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); } public SparseTensorIndexCOO __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } /** * The type of values in indicesBuffer */ public org.apache.arrow.flatbuf.Int indicesType() { return indicesType(new org.apache.arrow.flatbuf.Int()); } public org.apache.arrow.flatbuf.Int indicesType(org.apache.arrow.flatbuf.Int obj) { int o = __offset(4); return o != 0 ? obj.__assign(__indirect(o + bb_pos), bb) : null; } /** * Non-negative byte offsets to advance one value cell along each dimension * If omitted, default to row-major order (C-like). */ public long indicesStrides(int j) { int o = __offset(6); return o != 0 ? bb.getLong(__vector(o) + j * 8) : 0; } public int indicesStridesLength() { int o = __offset(6); return o != 0 ? __vector_len(o) : 0; } public LongVector indicesStridesVector() { return indicesStridesVector(new LongVector()); } public LongVector indicesStridesVector(LongVector obj) { int o = __offset(6); return o != 0 ? obj.__assign(__vector(o), bb) : null; } public ByteBuffer indicesStridesAsByteBuffer() { return __vector_as_bytebuffer(6, 8); } public ByteBuffer indicesStridesInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 6, 8); } /** * The location and size of the indices matrix's data */ public org.apache.arrow.flatbuf.Buffer indicesBuffer() { return indicesBuffer(new org.apache.arrow.flatbuf.Buffer()); } public org.apache.arrow.flatbuf.Buffer indicesBuffer(org.apache.arrow.flatbuf.Buffer obj) { int o = __offset(8); return o != 0 ? obj.__assign(o + bb_pos, bb) : null; } /** * This flag is true if and only if the indices matrix is sorted in * row-major order, and does not have duplicated entries. * This sort order is the same as of Tensorflow's SparseTensor, * but it is inverse order of SciPy's canonical coo_matrix * (SciPy employs column-major order for its coo_matrix). */ public boolean isCanonical() { int o = __offset(10); return o != 0 ? 0!=bb.get(o + bb_pos) : false; } public static void startSparseTensorIndexCOO(FlatBufferBuilder builder) { builder.startTable(4); } public static void addIndicesType(FlatBufferBuilder builder, int indicesTypeOffset) { builder.addOffset(0, indicesTypeOffset, 0); } public static void addIndicesStrides(FlatBufferBuilder builder, int indicesStridesOffset) { builder.addOffset(1, indicesStridesOffset, 0); } public static int createIndicesStridesVector(FlatBufferBuilder builder, long[] data) { builder.startVector(8, data.length, 8); for (int i = data.length - 1; i >= 0; i--) builder.addLong(data[i]); return builder.endVector(); } public static void startIndicesStridesVector(FlatBufferBuilder builder, int numElems) { builder.startVector(8, numElems, 8); } public static void addIndicesBuffer(FlatBufferBuilder builder, int indicesBufferOffset) { builder.addStruct(2, indicesBufferOffset, 0); } public static void addIsCanonical(FlatBufferBuilder builder, boolean isCanonical) { builder.addBoolean(3, isCanonical, false); } public static int endSparseTensorIndexCOO(FlatBufferBuilder builder) { int o = builder.endTable(); builder.required(o, 4); // indicesType builder.required(o, 8); // indicesBuffer return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public SparseTensorIndexCOO get(int j) { return get(new SparseTensorIndexCOO(), j); } public SparseTensorIndexCOO get(SparseTensorIndexCOO obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/SparseTensorIndexCSF.java000066400000000000000000000236141515451615200324340ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * Compressed Sparse Fiber (CSF) sparse tensor index. */ @SuppressWarnings("unused") public final class SparseTensorIndexCSF extends Table { 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); } public SparseTensorIndexCSF __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } /** * CSF is a generalization of compressed sparse row (CSR) index. * See [smith2017knl](http://shaden.io/pub-files/smith2017knl.pdf) * * CSF index recursively compresses each dimension of a tensor into a set * of prefix trees. Each path from a root to leaf forms one tensor * non-zero index. CSF is implemented with two arrays of buffers and one * arrays of integers. * * For example, let X be a 2x3x4x5 tensor and let it have the following * 8 non-zero values: * ```text * X[0, 0, 0, 1] := 1 * X[0, 0, 0, 2] := 2 * X[0, 1, 0, 0] := 3 * X[0, 1, 0, 2] := 4 * X[0, 1, 1, 0] := 5 * X[1, 1, 1, 0] := 6 * X[1, 1, 1, 1] := 7 * X[1, 1, 1, 2] := 8 * ``` * As a prefix tree this would be represented as: * ```text * 0 1 * / \ | * 0 1 1 * / / \ | * 0 0 1 1 * /| /| | /| | * 1 2 0 2 0 0 1 2 * ``` * The type of values in indptrBuffers */ public org.apache.arrow.flatbuf.Int indptrType() { return indptrType(new org.apache.arrow.flatbuf.Int()); } public org.apache.arrow.flatbuf.Int indptrType(org.apache.arrow.flatbuf.Int obj) { int o = __offset(4); return o != 0 ? obj.__assign(__indirect(o + bb_pos), bb) : null; } /** * indptrBuffers stores the sparsity structure. * Each two consecutive dimensions in a tensor correspond to a buffer in * indptrBuffers. A pair of consecutive values at `indptrBuffers[dim][i]` * and `indptrBuffers[dim][i + 1]` signify a range of nodes in * `indicesBuffers[dim + 1]` who are children of `indicesBuffers[dim][i]` node. * * For example, the indptrBuffers for the above X is: * ```text * indptrBuffer(X) = [ * [0, 2, 3], * [0, 1, 3, 4], * [0, 2, 4, 5, 8] * ]. * ``` */ public org.apache.arrow.flatbuf.Buffer indptrBuffers(int j) { return indptrBuffers(new org.apache.arrow.flatbuf.Buffer(), j); } public org.apache.arrow.flatbuf.Buffer indptrBuffers(org.apache.arrow.flatbuf.Buffer obj, int j) { int o = __offset(6); return o != 0 ? obj.__assign(__vector(o) + j * 16, bb) : null; } public int indptrBuffersLength() { int o = __offset(6); return o != 0 ? __vector_len(o) : 0; } public org.apache.arrow.flatbuf.Buffer.Vector indptrBuffersVector() { return indptrBuffersVector(new org.apache.arrow.flatbuf.Buffer.Vector()); } public org.apache.arrow.flatbuf.Buffer.Vector indptrBuffersVector(org.apache.arrow.flatbuf.Buffer.Vector obj) { int o = __offset(6); return o != 0 ? obj.__assign(__vector(o), 16, bb) : null; } /** * The type of values in indicesBuffers */ public org.apache.arrow.flatbuf.Int indicesType() { return indicesType(new org.apache.arrow.flatbuf.Int()); } public org.apache.arrow.flatbuf.Int indicesType(org.apache.arrow.flatbuf.Int obj) { int o = __offset(8); return o != 0 ? obj.__assign(__indirect(o + bb_pos), bb) : null; } /** * indicesBuffers stores values of nodes. * Each tensor dimension corresponds to a buffer in indicesBuffers. * For example, the indicesBuffers for the above X is: * ```text * indicesBuffer(X) = [ * [0, 1], * [0, 1, 1], * [0, 0, 1, 1], * [1, 2, 0, 2, 0, 0, 1, 2] * ]. * ``` */ public org.apache.arrow.flatbuf.Buffer indicesBuffers(int j) { return indicesBuffers(new org.apache.arrow.flatbuf.Buffer(), j); } public org.apache.arrow.flatbuf.Buffer indicesBuffers(org.apache.arrow.flatbuf.Buffer obj, int j) { int o = __offset(10); return o != 0 ? obj.__assign(__vector(o) + j * 16, bb) : null; } public int indicesBuffersLength() { int o = __offset(10); return o != 0 ? __vector_len(o) : 0; } public org.apache.arrow.flatbuf.Buffer.Vector indicesBuffersVector() { return indicesBuffersVector(new org.apache.arrow.flatbuf.Buffer.Vector()); } public org.apache.arrow.flatbuf.Buffer.Vector indicesBuffersVector(org.apache.arrow.flatbuf.Buffer.Vector obj) { int o = __offset(10); return o != 0 ? obj.__assign(__vector(o), 16, bb) : null; } /** * axisOrder stores the sequence in which dimensions were traversed to * produce the prefix tree. * For example, the axisOrder for the above X is: * ```text * axisOrder(X) = [0, 1, 2, 3]. * ``` */ public int axisOrder(int j) { int o = __offset(12); return o != 0 ? bb.getInt(__vector(o) + j * 4) : 0; } public int axisOrderLength() { int o = __offset(12); return o != 0 ? __vector_len(o) : 0; } public IntVector axisOrderVector() { return axisOrderVector(new IntVector()); } public IntVector axisOrderVector(IntVector obj) { int o = __offset(12); return o != 0 ? obj.__assign(__vector(o), bb) : null; } public ByteBuffer axisOrderAsByteBuffer() { return __vector_as_bytebuffer(12, 4); } public ByteBuffer axisOrderInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 12, 4); } public static int createSparseTensorIndexCSF(FlatBufferBuilder builder, int indptrTypeOffset, int indptrBuffersOffset, int indicesTypeOffset, int indicesBuffersOffset, int axisOrderOffset) { builder.startTable(5); SparseTensorIndexCSF.addAxisOrder(builder, axisOrderOffset); SparseTensorIndexCSF.addIndicesBuffers(builder, indicesBuffersOffset); SparseTensorIndexCSF.addIndicesType(builder, indicesTypeOffset); SparseTensorIndexCSF.addIndptrBuffers(builder, indptrBuffersOffset); SparseTensorIndexCSF.addIndptrType(builder, indptrTypeOffset); return SparseTensorIndexCSF.endSparseTensorIndexCSF(builder); } public static void startSparseTensorIndexCSF(FlatBufferBuilder builder) { builder.startTable(5); } public static void addIndptrType(FlatBufferBuilder builder, int indptrTypeOffset) { builder.addOffset(0, indptrTypeOffset, 0); } public static void addIndptrBuffers(FlatBufferBuilder builder, int indptrBuffersOffset) { builder.addOffset(1, indptrBuffersOffset, 0); } public static void startIndptrBuffersVector(FlatBufferBuilder builder, int numElems) { builder.startVector(16, numElems, 8); } public static void addIndicesType(FlatBufferBuilder builder, int indicesTypeOffset) { builder.addOffset(2, indicesTypeOffset, 0); } public static void addIndicesBuffers(FlatBufferBuilder builder, int indicesBuffersOffset) { builder.addOffset(3, indicesBuffersOffset, 0); } public static void startIndicesBuffersVector(FlatBufferBuilder builder, int numElems) { builder.startVector(16, numElems, 8); } public static void addAxisOrder(FlatBufferBuilder builder, int axisOrderOffset) { builder.addOffset(4, axisOrderOffset, 0); } public static int createAxisOrderVector(FlatBufferBuilder builder, int[] data) { builder.startVector(4, data.length, 4); for (int i = data.length - 1; i >= 0; i--) builder.addInt(data[i]); return builder.endVector(); } public static void startAxisOrderVector(FlatBufferBuilder builder, int numElems) { builder.startVector(4, numElems, 4); } public static int endSparseTensorIndexCSF(FlatBufferBuilder builder) { int o = builder.endTable(); builder.required(o, 4); // indptrType builder.required(o, 6); // indptrBuffers builder.required(o, 8); // indicesType builder.required(o, 10); // indicesBuffers builder.required(o, 12); // axisOrder return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public SparseTensorIndexCSF get(int j) { return get(new SparseTensorIndexCSF(), j); } public SparseTensorIndexCSF get(SparseTensorIndexCSF obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Struct_.java000066400000000000000000000054131515451615200301000ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * A Struct_ in the flatbuffer metadata is the same as an Arrow Struct * (according to the physical memory layout). We used Struct_ here as * Struct is a reserved word in Flatbuffers */ @SuppressWarnings("unused") public final class Struct_ extends Table { 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); } public Struct_ __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public static void startStruct_(FlatBufferBuilder builder) { builder.startTable(0); } public static int endStruct_(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Struct_ get(int j) { return get(new Struct_(), j); } public Struct_ get(Struct_ obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Tensor.java000066400000000000000000000144451515451615200277340ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; @SuppressWarnings("unused") public final class Tensor extends Table { 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); } public Tensor __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public byte typeType() { int o = __offset(4); return o != 0 ? bb.get(o + bb_pos) : 0; } /** * The type of data contained in a value cell. Currently only fixed-width * value types are supported, no strings or nested types */ public Table type(Table obj) { int o = __offset(6); return o != 0 ? __union(obj, o + bb_pos) : null; } /** * The dimensions of the tensor, optionally named */ public org.apache.arrow.flatbuf.TensorDim shape(int j) { return shape(new org.apache.arrow.flatbuf.TensorDim(), j); } public org.apache.arrow.flatbuf.TensorDim shape(org.apache.arrow.flatbuf.TensorDim obj, int j) { int o = __offset(8); return o != 0 ? obj.__assign(__indirect(__vector(o) + j * 4), bb) : null; } public int shapeLength() { int o = __offset(8); return o != 0 ? __vector_len(o) : 0; } public org.apache.arrow.flatbuf.TensorDim.Vector shapeVector() { return shapeVector(new org.apache.arrow.flatbuf.TensorDim.Vector()); } public org.apache.arrow.flatbuf.TensorDim.Vector shapeVector(org.apache.arrow.flatbuf.TensorDim.Vector obj) { int o = __offset(8); return o != 0 ? obj.__assign(__vector(o), 4, bb) : null; } /** * Non-negative byte offsets to advance one value cell along each dimension * If omitted, default to row-major order (C-like). */ public long strides(int j) { int o = __offset(10); return o != 0 ? bb.getLong(__vector(o) + j * 8) : 0; } public int stridesLength() { int o = __offset(10); return o != 0 ? __vector_len(o) : 0; } public LongVector stridesVector() { return stridesVector(new LongVector()); } public LongVector stridesVector(LongVector obj) { int o = __offset(10); return o != 0 ? obj.__assign(__vector(o), bb) : null; } public ByteBuffer stridesAsByteBuffer() { return __vector_as_bytebuffer(10, 8); } public ByteBuffer stridesInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 10, 8); } /** * The location and size of the tensor's data */ public org.apache.arrow.flatbuf.Buffer data() { return data(new org.apache.arrow.flatbuf.Buffer()); } public org.apache.arrow.flatbuf.Buffer data(org.apache.arrow.flatbuf.Buffer obj) { int o = __offset(12); return o != 0 ? obj.__assign(o + bb_pos, bb) : null; } public static void startTensor(FlatBufferBuilder builder) { builder.startTable(5); } public static void addTypeType(FlatBufferBuilder builder, byte typeType) { builder.addByte(0, typeType, 0); } public static void addType(FlatBufferBuilder builder, int typeOffset) { builder.addOffset(1, typeOffset, 0); } public static void addShape(FlatBufferBuilder builder, int shapeOffset) { builder.addOffset(2, shapeOffset, 0); } public static int createShapeVector(FlatBufferBuilder builder, int[] data) { builder.startVector(4, data.length, 4); for (int i = data.length - 1; i >= 0; i--) builder.addOffset(data[i]); return builder.endVector(); } public static void startShapeVector(FlatBufferBuilder builder, int numElems) { builder.startVector(4, numElems, 4); } public static void addStrides(FlatBufferBuilder builder, int stridesOffset) { builder.addOffset(3, stridesOffset, 0); } public static int createStridesVector(FlatBufferBuilder builder, long[] data) { builder.startVector(8, data.length, 8); for (int i = data.length - 1; i >= 0; i--) builder.addLong(data[i]); return builder.endVector(); } public static void startStridesVector(FlatBufferBuilder builder, int numElems) { builder.startVector(8, numElems, 8); } public static void addData(FlatBufferBuilder builder, int dataOffset) { builder.addStruct(4, dataOffset, 0); } public static int endTensor(FlatBufferBuilder builder) { int o = builder.endTable(); builder.required(o, 6); // type builder.required(o, 8); // shape builder.required(o, 12); // data return o; } public static void finishTensorBuffer(FlatBufferBuilder builder, int offset) { builder.finish(offset); } public static void finishSizePrefixedTensorBuffer(FlatBufferBuilder builder, int offset) { builder.finishSizePrefixed(offset); } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Tensor get(int j) { return get(new Tensor(), j); } public Tensor get(Tensor obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/TensorDim.java000066400000000000000000000072501515451615200303620ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * ---------------------------------------------------------------------- * Data structures for dense tensors * Shape data for a single axis in a tensor */ @SuppressWarnings("unused") public final class TensorDim extends Table { 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); } public TensorDim __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } /** * Length of dimension */ public long size() { int o = __offset(4); return o != 0 ? bb.getLong(o + bb_pos) : 0L; } /** * Name of the dimension, optional */ public String name() { int o = __offset(6); return o != 0 ? __string(o + bb_pos) : null; } public ByteBuffer nameAsByteBuffer() { return __vector_as_bytebuffer(6, 1); } public ByteBuffer nameInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 6, 1); } public static int createTensorDim(FlatBufferBuilder builder, long size, int nameOffset) { builder.startTable(2); TensorDim.addSize(builder, size); TensorDim.addName(builder, nameOffset); return TensorDim.endTensorDim(builder); } public static void startTensorDim(FlatBufferBuilder builder) { builder.startTable(2); } public static void addSize(FlatBufferBuilder builder, long size) { builder.addLong(0, size, 0L); } public static void addName(FlatBufferBuilder builder, int nameOffset) { builder.addOffset(1, nameOffset, 0); } public static int endTensorDim(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public TensorDim get(int j) { return get(new TensorDim(), j); } public TensorDim get(TensorDim obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Time.java000066400000000000000000000076311515451615200273570ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * Time is either a 32-bit or 64-bit signed integer type representing an * elapsed time since midnight, stored in either of four units: seconds, * milliseconds, microseconds or nanoseconds. * * The integer `bitWidth` depends on the `unit` and must be one of the following: * * SECOND and MILLISECOND: 32 bits * * MICROSECOND and NANOSECOND: 64 bits * * The allowed values are between 0 (inclusive) and 86400 (=24*60*60) seconds * (exclusive), adjusted for the time unit (for example, up to 86400000 * exclusive for the MILLISECOND unit). * This definition doesn't allow for leap seconds. Time values from * measurements with leap seconds will need to be corrected when ingesting * into Arrow (for example by replacing the value 86400 with 86399). */ @SuppressWarnings("unused") public final class Time extends Table { 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); } public Time __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public short unit() { int o = __offset(4); return o != 0 ? bb.getShort(o + bb_pos) : 1; } public int bitWidth() { int o = __offset(6); return o != 0 ? bb.getInt(o + bb_pos) : 32; } public static int createTime(FlatBufferBuilder builder, short unit, int bitWidth) { builder.startTable(2); Time.addBitWidth(builder, bitWidth); Time.addUnit(builder, unit); return Time.endTime(builder); } public static void startTime(FlatBufferBuilder builder) { builder.startTable(2); } public static void addUnit(FlatBufferBuilder builder, short unit) { builder.addShort(0, unit, 1); } public static void addBitWidth(FlatBufferBuilder builder, int bitWidth) { builder.addInt(1, bitWidth, 32); } public static int endTime(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Time get(int j) { return get(new Time(), j); } public Time get(Time obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/TimeUnit.java000066400000000000000000000023431515451615200302120ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; @SuppressWarnings("unused") public final class TimeUnit { private TimeUnit() { } public static final short SECOND = 0; public static final short MILLISECOND = 1; public static final short MICROSECOND = 2; public static final short NANOSECOND = 3; public static final String[] names = { "SECOND", "MILLISECOND", "MICROSECOND", "NANOSECOND", }; public static String name(int e) { return names[e]; } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Timestamp.java000066400000000000000000000226141515451615200304220ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * Timestamp is a 64-bit signed integer representing an elapsed time since a * fixed epoch, stored in either of four units: seconds, milliseconds, * microseconds or nanoseconds, and is optionally annotated with a timezone. * * Timestamp values do not include any leap seconds (in other words, all * days are considered 86400 seconds long). * * Timestamps with a non-empty timezone * ------------------------------------ * * If a Timestamp column has a non-empty timezone value, its epoch is * 1970-01-01 00:00:00 (January 1st 1970, midnight) in the *UTC* timezone * (the Unix epoch), regardless of the Timestamp's own timezone. * * Therefore, timestamp values with a non-empty timezone correspond to * physical points in time together with some additional information about * how the data was obtained and/or how to display it (the timezone). * * For example, the timestamp value 0 with the timezone string "Europe/Paris" * corresponds to "January 1st 1970, 00h00" in the UTC timezone, but the * application may prefer to display it as "January 1st 1970, 01h00" in * the Europe/Paris timezone (which is the same physical point in time). * * One consequence is that timestamp values with a non-empty timezone * can be compared and ordered directly, since they all share the same * well-known point of reference (the Unix epoch). * * Timestamps with an unset / empty timezone * ----------------------------------------- * * If a Timestamp column has no timezone value, its epoch is * 1970-01-01 00:00:00 (January 1st 1970, midnight) in an *unknown* timezone. * * Therefore, timestamp values without a timezone cannot be meaningfully * interpreted as physical points in time, but only as calendar / clock * indications ("wall clock time") in an unspecified timezone. * * For example, the timestamp value 0 with an empty timezone string * corresponds to "January 1st 1970, 00h00" in an unknown timezone: there * is not enough information to interpret it as a well-defined physical * point in time. * * One consequence is that timestamp values without a timezone cannot * be reliably compared or ordered, since they may have different points of * reference. In particular, it is *not* possible to interpret an unset * or empty timezone as the same as "UTC". * * Conversion between timezones * ---------------------------- * * If a Timestamp column has a non-empty timezone, changing the timezone * to a different non-empty value is a metadata-only operation: * the timestamp values need not change as their point of reference remains * the same (the Unix epoch). * * However, if a Timestamp column has no timezone value, changing it to a * non-empty value requires to think about the desired semantics. * One possibility is to assume that the original timestamp values are * relative to the epoch of the timezone being set; timestamp values should * then adjusted to the Unix epoch (for example, changing the timezone from * empty to "Europe/Paris" would require converting the timestamp values * from "Europe/Paris" to "UTC", which seems counter-intuitive but is * nevertheless correct). * * Guidelines for encoding data from external libraries * ---------------------------------------------------- * * Date & time libraries often have multiple different data types for temporal * data. In order to ease interoperability between different implementations the * Arrow project has some recommendations for encoding these types into a Timestamp * column. * * An "instant" represents a physical point in time that has no relevant timezone * (for example, astronomical data). To encode an instant, use a Timestamp with * the timezone string set to "UTC", and make sure the Timestamp values * are relative to the UTC epoch (January 1st 1970, midnight). * * A "zoned date-time" represents a physical point in time annotated with an * informative timezone (for example, the timezone in which the data was * recorded). To encode a zoned date-time, use a Timestamp with the timezone * string set to the name of the timezone, and make sure the Timestamp values * are relative to the UTC epoch (January 1st 1970, midnight). * * (There is some ambiguity between an instant and a zoned date-time with the * UTC timezone. Both of these are stored the same in Arrow. Typically, * this distinction does not matter. If it does, then an application should * use custom metadata or an extension type to distinguish between the two cases.) * * An "offset date-time" represents a physical point in time combined with an * explicit offset from UTC. To encode an offset date-time, use a Timestamp * with the timezone string set to the numeric timezone offset string * (e.g. "+03:00"), and make sure the Timestamp values are relative to * the UTC epoch (January 1st 1970, midnight). * * A "naive date-time" (also called "local date-time" in some libraries) * represents a wall clock time combined with a calendar date, but with * no indication of how to map this information to a physical point in time. * Naive date-times must be handled with care because of this missing * information, and also because daylight saving time (DST) may make * some values ambiguous or nonexistent. A naive date-time may be * stored as a struct with Date and Time fields. However, it may also be * encoded into a Timestamp column with an empty timezone. The timestamp * values should be computed "as if" the timezone of the date-time values * was UTC; for example, the naive date-time "January 1st 1970, 00h00" would * be encoded as timestamp value 0. */ @SuppressWarnings("unused") public final class Timestamp extends Table { 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); } public Timestamp __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public short unit() { int o = __offset(4); return o != 0 ? bb.getShort(o + bb_pos) : 0; } /** * The timezone is an optional string indicating the name of a timezone, * one of: * * * As used in the Olson timezone database (the "tz database" or * "tzdata"), such as "America/New_York". * * An absolute timezone offset of the form "+XX:XX" or "-XX:XX", * such as "+07:30". * * Whether a timezone string is present indicates different semantics about * the data (see above). */ public String timezone() { int o = __offset(6); return o != 0 ? __string(o + bb_pos) : null; } public ByteBuffer timezoneAsByteBuffer() { return __vector_as_bytebuffer(6, 1); } public ByteBuffer timezoneInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 6, 1); } public static int createTimestamp(FlatBufferBuilder builder, short unit, int timezoneOffset) { builder.startTable(2); Timestamp.addTimezone(builder, timezoneOffset); Timestamp.addUnit(builder, unit); return Timestamp.endTimestamp(builder); } public static void startTimestamp(FlatBufferBuilder builder) { builder.startTable(2); } public static void addUnit(FlatBufferBuilder builder, short unit) { builder.addShort(0, unit, 0); } public static void addTimezone(FlatBufferBuilder builder, int timezoneOffset) { builder.addOffset(1, timezoneOffset, 0); } public static int endTimestamp(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Timestamp get(int j) { return get(new Timestamp(), j); } public Timestamp get(Timestamp obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Type.java000066400000000000000000000051631515451615200274000ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; /** * ---------------------------------------------------------------------- * Top-level Type value, enabling extensible type-specific metadata. We can * add new logical types to Type without breaking backwards compatibility */ @SuppressWarnings("unused") public final class Type { private Type() { } public static final byte NONE = 0; public static final byte Null = 1; public static final byte Int = 2; public static final byte FloatingPoint = 3; public static final byte Binary = 4; public static final byte Utf8 = 5; public static final byte Bool = 6; public static final byte Decimal = 7; public static final byte Date = 8; public static final byte Time = 9; public static final byte Timestamp = 10; public static final byte Interval = 11; public static final byte List = 12; public static final byte Struct_ = 13; public static final byte Union = 14; public static final byte FixedSizeBinary = 15; public static final byte FixedSizeList = 16; public static final byte Map = 17; public static final byte Duration = 18; public static final byte LargeBinary = 19; public static final byte LargeUtf8 = 20; public static final byte LargeList = 21; public static final byte RunEndEncoded = 22; public static final byte BinaryView = 23; public static final byte Utf8View = 24; public static final byte ListView = 25; public static final byte LargeListView = 26; public static final String[] names = { "NONE", "Null", "Int", "FloatingPoint", "Binary", "Utf8", "Bool", "Decimal", "Date", "Time", "Timestamp", "Interval", "List", "Struct_", "Union", "FixedSizeBinary", "FixedSizeList", "Map", "Duration", "LargeBinary", "LargeUtf8", "LargeList", "RunEndEncoded", "BinaryView", "Utf8View", "ListView", "LargeListView", }; public static String name(int e) { return names[e]; } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Union.java000066400000000000000000000104461515451615200275470ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * A union is a complex type with children in Field * By default ids in the type vector refer to the offsets in the children * optionally typeIds provides an indirection between the child offset and the type id * for each child `typeIds[offset]` is the id used in the type vector */ @SuppressWarnings("unused") public final class Union extends Table { 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); } public Union __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public short mode() { int o = __offset(4); return o != 0 ? bb.getShort(o + bb_pos) : 0; } public int typeIds(int j) { int o = __offset(6); return o != 0 ? bb.getInt(__vector(o) + j * 4) : 0; } public int typeIdsLength() { int o = __offset(6); return o != 0 ? __vector_len(o) : 0; } public IntVector typeIdsVector() { return typeIdsVector(new IntVector()); } public IntVector typeIdsVector(IntVector obj) { int o = __offset(6); return o != 0 ? obj.__assign(__vector(o), bb) : null; } public ByteBuffer typeIdsAsByteBuffer() { return __vector_as_bytebuffer(6, 4); } public ByteBuffer typeIdsInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 6, 4); } public static int createUnion(FlatBufferBuilder builder, short mode, int typeIdsOffset) { builder.startTable(2); Union.addTypeIds(builder, typeIdsOffset); Union.addMode(builder, mode); return Union.endUnion(builder); } public static void startUnion(FlatBufferBuilder builder) { builder.startTable(2); } public static void addMode(FlatBufferBuilder builder, short mode) { builder.addShort(0, mode, 0); } public static void addTypeIds(FlatBufferBuilder builder, int typeIdsOffset) { builder.addOffset(1, typeIdsOffset, 0); } public static int createTypeIdsVector(FlatBufferBuilder builder, int[] data) { builder.startVector(4, data.length, 4); for (int i = data.length - 1; i >= 0; i--) builder.addInt(data[i]); return builder.endVector(); } public static void startTypeIdsVector(FlatBufferBuilder builder, int numElems) { builder.startVector(4, numElems, 4); } public static int endUnion(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Union get(int j) { return get(new Union(), j); } public Union get(Union obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/UnionMode.java000066400000000000000000000021431515451615200303470ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; @SuppressWarnings("unused") public final class UnionMode { private UnionMode() { } public static final short Sparse = 0; public static final short Dense = 1; public static final String[] names = { "Sparse", "Dense", }; public static String name(int e) { return names[e]; } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Utf8.java000066400000000000000000000051041515451615200273000ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * Unicode with UTF-8 encoding */ @SuppressWarnings("unused") public final class Utf8 extends Table { 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); } public Utf8 __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public static void startUtf8(FlatBufferBuilder builder) { builder.startTable(0); } public static int endUtf8(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Utf8 get(int j) { return get(new Utf8(), j); } public Utf8 get(Utf8 obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/format/src/main/java/org/apache/arrow/flatbuf/Utf8View.java000066400000000000000000000060461515451615200301410ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.flatbuf; import com.google.flatbuffers.BaseVector; import com.google.flatbuffers.BooleanVector; import com.google.flatbuffers.ByteVector; import com.google.flatbuffers.Constants; import com.google.flatbuffers.DoubleVector; import com.google.flatbuffers.FlatBufferBuilder; import com.google.flatbuffers.FloatVector; import com.google.flatbuffers.IntVector; import com.google.flatbuffers.LongVector; import com.google.flatbuffers.ShortVector; import com.google.flatbuffers.StringVector; import com.google.flatbuffers.Struct; import com.google.flatbuffers.Table; import com.google.flatbuffers.UnionVector; import java.nio.ByteBuffer; import java.nio.ByteOrder; /** * Logically the same as Utf8, but the internal representation uses a view * struct that contains the string length and either the string's entire data * inline (for small strings) or an inlined prefix, an index of another buffer, * and an offset pointing to a slice in that buffer (for non-small strings). * * Since it uses a variable number of data buffers, each Field with this type * must have a corresponding entry in `variadicBufferCounts`. */ @SuppressWarnings("unused") public final class Utf8View extends Table { 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); } public Utf8View __assign(int _i, ByteBuffer _bb) { __init(_i, _bb); return this; } public static void startUtf8View(FlatBufferBuilder builder) { builder.startTable(0); } public static int endUtf8View(FlatBufferBuilder builder) { int o = builder.endTable(); return o; } public static final class Vector extends BaseVector { public Vector __assign(int _vector, int _element_size, ByteBuffer _bb) { __reset(_vector, _element_size, _bb); return this; } public Utf8View get(int j) { return get(new Utf8View(), j); } public Utf8View get(Utf8View obj, int j) { return obj.__assign(__indirect(__element(j), bb), bb); } } } apache-arrow-java-19.0.0/gandiva/000077500000000000000000000000001515451615200165075ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/CMakeLists.txt000066400000000000000000000103111515451615200212430ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. find_package(Gandiva REQUIRED) include_directories(${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_SOURCE_DIR} ${JNI_INCLUDE_DIRS} ${JNI_HEADERS_DIR}) add_jar(arrow_java_jni_gandiva_jar src/main/java/org/apache/arrow/gandiva/evaluator/ConfigurationBuilder.java src/main/java/org/apache/arrow/gandiva/evaluator/JniWrapper.java src/main/java/org/apache/arrow/gandiva/evaluator/ExpressionRegistryJniHelper.java src/main/java/org/apache/arrow/gandiva/exceptions/GandivaException.java GENERATE_NATIVE_HEADERS arrow_java_jni_gandiva_headers) set(GANDIVA_PROTO_OUTPUT_DIR ${CMAKE_CURRENT_BINARY_DIR}) set(GANDIVA_PROTO_OUTPUT_FILES "${GANDIVA_PROTO_OUTPUT_DIR}/gandiva/types.pb.cc" "${GANDIVA_PROTO_OUTPUT_DIR}/gandiva/types.pb.h") set_source_files_properties(${GANDIVA_PROTO_OUTPUT_FILES} PROPERTIES GENERATED TRUE) set(GANDIVA_PROTO_DIR ${CMAKE_CURRENT_SOURCE_DIR}/proto) get_filename_component(GANDIVA_PROTO_FILE_ABSOLUTE ${GANDIVA_PROTO_DIR}/gandiva/types.proto ABSOLUTE) find_package(Protobuf REQUIRED) add_custom_command(OUTPUT ${GANDIVA_PROTO_OUTPUT_FILES} COMMAND protobuf::protoc --proto_path ${GANDIVA_PROTO_DIR} --cpp_out ${GANDIVA_PROTO_OUTPUT_DIR} ${GANDIVA_PROTO_FILE_ABSOLUTE} DEPENDS ${GANDIVA_PROTO_FILE_ABSOLUTE} COMMENT "Running Protobuf compiler on gandiva/types.proto" VERBATIM) add_custom_target(garrow_java_jni_gandiva_proto ALL DEPENDS ${GANDIVA_PROTO_OUTPUT_FILES}) add_library(arrow_java_jni_gandiva SHARED src/main/cpp/config_builder.cc src/main/cpp/config_holder.cc src/main/cpp/expression_registry_helper.cc src/main/cpp/jni_common.cc ${GANDIVA_PROTO_OUTPUT_FILES}) set_property(TARGET arrow_java_jni_gandiva PROPERTY OUTPUT_NAME "gandiva_jni") target_link_libraries(arrow_java_jni_gandiva arrow_java_jni_gandiva_headers jni protobuf::libprotobuf Gandiva::gandiva_static) # Localize thirdparty symbols using a linker version script. This hides them # from the client application. The OS X linker does not support the # version-script option. if(CMAKE_VERSION VERSION_LESS 3.18) if(APPLE OR WIN32) set(CXX_LINKER_SUPPORTS_VERSION_SCRIPT FALSE) else() set(CXX_LINKER_SUPPORTS_VERSION_SCRIPT TRUE) endif() else() include(CheckLinkerFlag) check_linker_flag(CXX "-Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/src/main/cpp/symbols.map" CXX_LINKER_SUPPORTS_VERSION_SCRIPT) endif() # filter out everything that is not needed for the jni bridge # statically linked stdc++ has conflicts with stdc++ loaded by other libraries. if(CXX_LINKER_SUPPORTS_VERSION_SCRIPT) set_target_properties(arrow_java_jni_gandiva PROPERTIES LINK_FLAGS "-Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/src/main/cpp/symbols.map" ) endif() set(ARROW_JAVA_JNI_GANDIVA_LIBDIR "${CMAKE_INSTALL_PREFIX}/lib/gandiva_jni/${ARROW_JAVA_JNI_ARCH_DIR}") set(ARROW_JAVA_JNI_GANDIVA_BINDIR "${CMAKE_INSTALL_PREFIX}/bin/gandiva_jni/${ARROW_JAVA_JNI_ARCH_DIR}") install(TARGETS arrow_java_jni_gandiva LIBRARY DESTINATION ${ARROW_JAVA_JNI_GANDIVA_LIBDIR} RUNTIME DESTINATION ${ARROW_JAVA_JNI_GANDIVA_BINDIR}) apache-arrow-java-19.0.0/gandiva/README.md000066400000000000000000000017671515451615200200010ustar00rootroot00000000000000/** * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ # Gandiva Java ## Setup Build Environment install: - java 7 or later - maven 3.3 or later ## Building and running tests ``` cd java mvn install -Dgandiva.cpp.build.dir= ``` apache-arrow-java-19.0.0/gandiva/pom.xml000066400000000000000000000107771515451615200200400ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-java-root 19.0.0 org.apache.arrow.gandiva arrow-gandiva jar Arrow Gandiva Java wrappers around the native Gandiva SQL expression compiler. true ../../../cpp/release-build org.apache.arrow arrow-memory-core org.immutables value-annotations org.apache.arrow arrow-memory-netty test org.apache.arrow arrow-vector ${arrow.vector.classifier} com.google.protobuf protobuf-java com.google.guava guava org.slf4j slf4j-api ${arrow.cpp.build.dir} **/*gandiva_jni.* org.xolstice.maven.plugins protobuf-maven-plugin src compile proto release org.apache.maven.plugins maven-source-plugin attach-sources jar-no-fork org.apache.maven.plugins maven-javadoc-plugin attach-javadocs jar org.apache.maven.plugins maven-gpg-plugin sign-artifacts sign verify apache-arrow-java-19.0.0/gandiva/proto/000077500000000000000000000000001515451615200176525ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/proto/gandiva/000077500000000000000000000000001515451615200212635ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/proto/gandiva/types.proto000066400000000000000000000125071515451615200235210ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. syntax = "proto3"; package gandiva.types; option java_package = "org.apache.arrow.gandiva.ipc"; option java_outer_classname = "GandivaTypes"; option optimize_for = SPEED; enum GandivaType { NONE = 0; // arrow::Type::NA BOOL = 1; // arrow::Type::BOOL UINT8 = 2; // arrow::Type::UINT8 INT8 = 3; // arrow::Type::INT8 UINT16 = 4; // represents arrow::Type fields in src/arrow/type.h INT16 = 5; UINT32 = 6; INT32 = 7; UINT64 = 8; INT64 = 9; HALF_FLOAT = 10; FLOAT = 11; DOUBLE = 12; UTF8 = 13; BINARY = 14; FIXED_SIZE_BINARY = 15; DATE32 = 16; DATE64 = 17; TIMESTAMP = 18; TIME32 = 19; TIME64 = 20; INTERVAL = 21; DECIMAL = 22; LIST = 23; STRUCT = 24; UNION = 25; DICTIONARY = 26; MAP = 27; } enum DateUnit { DAY = 0; MILLI = 1; } enum TimeUnit { SEC = 0; MILLISEC = 1; MICROSEC = 2; NANOSEC = 3; } enum IntervalType { YEAR_MONTH = 0; DAY_TIME = 1; } enum SelectionVectorType { SV_NONE = 0; SV_INT16 = 1; SV_INT32 = 2; } message ExtGandivaType { optional GandivaType type = 1; optional uint32 width = 2; // used by FIXED_SIZE_BINARY optional int32 precision = 3; // used by DECIMAL optional int32 scale = 4; // used by DECIMAL optional DateUnit dateUnit = 5; // used by DATE32/DATE64 optional TimeUnit timeUnit = 6; // used by TIME32/TIME64 optional string timeZone = 7; // used by TIMESTAMP optional IntervalType intervalType = 8; // used by INTERVAL } message Field { // name of the field optional string name = 1; optional ExtGandivaType type = 2; optional bool nullable = 3; // for complex data types like structs, unions repeated Field children = 4; } message FieldNode { optional Field field = 1; } message FunctionNode { optional string functionName = 1; repeated TreeNode inArgs = 2; optional ExtGandivaType returnType = 3; } message IfNode { optional TreeNode cond = 1; optional TreeNode thenNode = 2; optional TreeNode elseNode = 3; optional ExtGandivaType returnType = 4; } message AndNode { repeated TreeNode args = 1; } message OrNode { repeated TreeNode args = 1; } message NullNode { optional ExtGandivaType type = 1; } message IntNode { optional int32 value = 1; } message FloatNode { optional float value = 1; } message DoubleNode { optional double value = 1; } message BooleanNode { optional bool value = 1; } message LongNode { optional int64 value = 1; } message StringNode { optional bytes value = 1; } message BinaryNode { optional bytes value = 1; } message DecimalNode { optional string value = 1; optional int32 precision = 2; optional int32 scale = 3; } message TreeNode { optional FieldNode fieldNode = 1; optional FunctionNode fnNode = 2; // control expressions optional IfNode ifNode = 6; optional AndNode andNode = 7; optional OrNode orNode = 8; // literals optional NullNode nullNode = 11; optional IntNode intNode = 12; optional FloatNode floatNode = 13; optional LongNode longNode = 14; optional BooleanNode booleanNode = 15; optional DoubleNode doubleNode = 16; optional StringNode stringNode = 17; optional BinaryNode binaryNode = 18; optional DecimalNode decimalNode = 19; // in expr optional InNode inNode = 21; } message ExpressionRoot { optional TreeNode root = 1; optional Field resultType = 2; } message ExpressionList { repeated ExpressionRoot exprs = 2; } message Condition { optional TreeNode root = 1; } message Schema { repeated Field columns = 1; } message GandivaDataTypes { repeated ExtGandivaType dataType = 1; } message GandivaFunctions { repeated FunctionSignature function = 1; } message FunctionSignature { optional string name = 1; optional ExtGandivaType returnType = 2; repeated ExtGandivaType paramTypes = 3; } message InNode { optional TreeNode node = 1; optional IntConstants intValues = 2; optional LongConstants longValues = 3; optional StringConstants stringValues = 4; optional BinaryConstants binaryValues = 5; optional DecimalConstants decimalValues = 6; optional FloatConstants floatValues = 7; optional DoubleConstants doubleValues = 8; } message IntConstants { repeated IntNode intValues = 1; } message LongConstants { repeated LongNode longValues = 1; } message DecimalConstants { repeated DecimalNode decimalValues = 1; } message FloatConstants { repeated FloatNode floatValues = 1; } message DoubleConstants { repeated DoubleNode doubleValues = 1; } message StringConstants { repeated StringNode stringValues = 1; } message BinaryConstants { repeated BinaryNode binaryValues = 1; } apache-arrow-java-19.0.0/gandiva/src/000077500000000000000000000000001515451615200172765ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/000077500000000000000000000000001515451615200202225ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/cpp/000077500000000000000000000000001515451615200210045ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/cpp/config_builder.cc000066400000000000000000000036541515451615200242760ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #include #include #include "config_holder.h" #include "env_helper.h" #include "org_apache_arrow_gandiva_evaluator_ConfigurationBuilder.h" using gandiva::ConfigHolder; using gandiva::Configuration; using gandiva::ConfigurationBuilder; /* * Class: org_apache_arrow_gandiva_evaluator_ConfigBuilder * Method: buildConfigInstance * Signature: (ZZ)J */ JNIEXPORT jlong JNICALL Java_org_apache_arrow_gandiva_evaluator_ConfigurationBuilder_buildConfigInstance( JNIEnv* env, jobject configuration, jboolean optimize, jboolean target_host_cpu) { ConfigurationBuilder configuration_builder; std::shared_ptr config = configuration_builder.build(); config->set_optimize(optimize); config->target_host_cpu(target_host_cpu); return ConfigHolder::MapInsert(config); } /* * Class: org_apache_arrow_gandiva_evaluator_ConfigBuilder * Method: releaseConfigInstance * Signature: (J)V */ JNIEXPORT void JNICALL Java_org_apache_arrow_gandiva_evaluator_ConfigurationBuilder_releaseConfigInstance( JNIEnv* env, jobject configuration, jlong config_id) { ConfigHolder::MapErase(config_id); } apache-arrow-java-19.0.0/gandiva/src/main/cpp/config_holder.cc000066400000000000000000000021301515451615200241110ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #include #include "config_holder.h" namespace gandiva { int64_t ConfigHolder::config_id_ = 1; // map of configuration objects created so far std::unordered_map> ConfigHolder::configuration_map_; std::mutex ConfigHolder::g_mtx_; } // namespace gandiva apache-arrow-java-19.0.0/gandiva/src/main/cpp/config_holder.h000066400000000000000000000036111515451615200237600ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #pragma once #include #include #include #include #include namespace gandiva { class ConfigHolder { public: static int64_t MapInsert(std::shared_ptr config) { g_mtx_.lock(); int64_t result = config_id_++; configuration_map_.insert( std::pair>(result, config)); g_mtx_.unlock(); return result; } static void MapErase(int64_t config_id_) { g_mtx_.lock(); configuration_map_.erase(config_id_); g_mtx_.unlock(); } static std::shared_ptr MapLookup(int64_t config_id_) { std::shared_ptr result = nullptr; try { result = configuration_map_.at(config_id_); } catch (const std::out_of_range&) { } return result; } private: // map of configuration objects created so far static std::unordered_map> configuration_map_; static std::mutex g_mtx_; // atomic counter for projector module ids static int64_t config_id_; }; } // namespace gandiva apache-arrow-java-19.0.0/gandiva/src/main/cpp/env_helper.h000066400000000000000000000016021515451615200233030ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #pragma once #include // class references extern jclass configuration_builder_class_; apache-arrow-java-19.0.0/gandiva/src/main/cpp/expression_registry_helper.cc000066400000000000000000000173071515451615200270110ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #include #include #include #include #include #include using gandiva::DataTypePtr; using gandiva::ExpressionRegistry; gandiva::types::TimeUnit MapTimeUnit(arrow::TimeUnit::type& unit) { switch (unit) { case arrow::TimeUnit::MILLI: return gandiva::types::TimeUnit::MILLISEC; case arrow::TimeUnit::SECOND: return gandiva::types::TimeUnit::SEC; case arrow::TimeUnit::MICRO: return gandiva::types::TimeUnit::MICROSEC; case arrow::TimeUnit::NANO: return gandiva::types::TimeUnit::NANOSEC; } // satisfy gcc. should be unreachable. return gandiva::types::TimeUnit::SEC; } void ArrowToProtobuf(DataTypePtr type, gandiva::types::ExtGandivaType* gandiva_data_type) { switch (type->id()) { case arrow::Type::BOOL: gandiva_data_type->set_type(gandiva::types::GandivaType::BOOL); break; case arrow::Type::UINT8: gandiva_data_type->set_type(gandiva::types::GandivaType::UINT8); break; case arrow::Type::INT8: gandiva_data_type->set_type(gandiva::types::GandivaType::INT8); break; case arrow::Type::UINT16: gandiva_data_type->set_type(gandiva::types::GandivaType::UINT16); break; case arrow::Type::INT16: gandiva_data_type->set_type(gandiva::types::GandivaType::INT16); break; case arrow::Type::UINT32: gandiva_data_type->set_type(gandiva::types::GandivaType::UINT32); break; case arrow::Type::INT32: gandiva_data_type->set_type(gandiva::types::GandivaType::INT32); break; case arrow::Type::UINT64: gandiva_data_type->set_type(gandiva::types::GandivaType::UINT64); break; case arrow::Type::INT64: gandiva_data_type->set_type(gandiva::types::GandivaType::INT64); break; case arrow::Type::HALF_FLOAT: gandiva_data_type->set_type(gandiva::types::GandivaType::HALF_FLOAT); break; case arrow::Type::FLOAT: gandiva_data_type->set_type(gandiva::types::GandivaType::FLOAT); break; case arrow::Type::DOUBLE: gandiva_data_type->set_type(gandiva::types::GandivaType::DOUBLE); break; case arrow::Type::STRING: gandiva_data_type->set_type(gandiva::types::GandivaType::UTF8); break; case arrow::Type::BINARY: gandiva_data_type->set_type(gandiva::types::GandivaType::BINARY); break; case arrow::Type::DATE32: gandiva_data_type->set_type(gandiva::types::GandivaType::DATE32); break; case arrow::Type::DATE64: gandiva_data_type->set_type(gandiva::types::GandivaType::DATE64); break; case arrow::Type::TIMESTAMP: { gandiva_data_type->set_type(gandiva::types::GandivaType::TIMESTAMP); std::shared_ptr cast_time_stamp_type = std::dynamic_pointer_cast(type); arrow::TimeUnit::type unit = cast_time_stamp_type->unit(); gandiva::types::TimeUnit time_unit = MapTimeUnit(unit); gandiva_data_type->set_timeunit(time_unit); break; } case arrow::Type::TIME32: { gandiva_data_type->set_type(gandiva::types::GandivaType::TIME32); std::shared_ptr cast_time_32_type = std::dynamic_pointer_cast(type); arrow::TimeUnit::type unit = cast_time_32_type->unit(); gandiva::types::TimeUnit time_unit = MapTimeUnit(unit); gandiva_data_type->set_timeunit(time_unit); break; } case arrow::Type::TIME64: { gandiva_data_type->set_type(gandiva::types::GandivaType::TIME32); std::shared_ptr cast_time_64_type = std::dynamic_pointer_cast(type); arrow::TimeUnit::type unit = cast_time_64_type->unit(); gandiva::types::TimeUnit time_unit = MapTimeUnit(unit); gandiva_data_type->set_timeunit(time_unit); break; } case arrow::Type::NA: gandiva_data_type->set_type(gandiva::types::GandivaType::NONE); break; case arrow::Type::DECIMAL: { gandiva_data_type->set_type(gandiva::types::GandivaType::DECIMAL); gandiva_data_type->set_precision(0); gandiva_data_type->set_scale(0); break; } case arrow::Type::INTERVAL_MONTHS: gandiva_data_type->set_type(gandiva::types::GandivaType::INTERVAL); gandiva_data_type->set_intervaltype(gandiva::types::IntervalType::YEAR_MONTH); break; case arrow::Type::INTERVAL_DAY_TIME: gandiva_data_type->set_type(gandiva::types::GandivaType::INTERVAL); gandiva_data_type->set_intervaltype(gandiva::types::IntervalType::DAY_TIME); break; default: // un-supported types. test ensures that // when one of these are added build breaks. ARROW_DCHECK(false); } } JNIEXPORT jbyteArray JNICALL Java_org_apache_arrow_gandiva_evaluator_ExpressionRegistryJniHelper_getGandivaSupportedDataTypes( // NOLINT JNIEnv* env, jobject types_helper) { gandiva::types::GandivaDataTypes gandiva_data_types; auto supported_types = ExpressionRegistry::supported_types(); for (auto const& type : supported_types) { gandiva::types::ExtGandivaType* gandiva_data_type = gandiva_data_types.add_datatype(); ArrowToProtobuf(type, gandiva_data_type); } auto size = static_cast(gandiva_data_types.ByteSizeLong()); std::unique_ptr buffer{new jbyte[size]}; gandiva_data_types.SerializeToArray(reinterpret_cast(buffer.get()), size); jbyteArray ret = env->NewByteArray(size); env->SetByteArrayRegion(ret, 0, size, buffer.get()); return ret; } /* * Class: org_apache_arrow_gandiva_types_ExpressionRegistryJniHelper * Method: getGandivaSupportedFunctions * Signature: ()[B */ JNIEXPORT jbyteArray JNICALL Java_org_apache_arrow_gandiva_evaluator_ExpressionRegistryJniHelper_getGandivaSupportedFunctions( // NOLINT JNIEnv* env, jobject types_helper) { ExpressionRegistry expr_registry; gandiva::types::GandivaFunctions gandiva_functions; for (auto function = expr_registry.function_signature_begin(); function != expr_registry.function_signature_end(); function++) { gandiva::types::FunctionSignature* function_signature = gandiva_functions.add_function(); function_signature->set_name((*function).base_name()); gandiva::types::ExtGandivaType* return_type = function_signature->mutable_returntype(); ArrowToProtobuf((*function).ret_type(), return_type); for (auto& param_type : (*function).param_types()) { gandiva::types::ExtGandivaType* proto_param_type = function_signature->add_paramtypes(); ArrowToProtobuf(param_type, proto_param_type); } } auto size = static_cast(gandiva_functions.ByteSizeLong()); std::unique_ptr buffer{new jbyte[size]}; gandiva_functions.SerializeToArray(reinterpret_cast(buffer.get()), size); jbyteArray ret = env->NewByteArray(size); env->SetByteArrayRegion(ret, 0, size, buffer.get()); return ret; } apache-arrow-java-19.0.0/gandiva/src/main/cpp/id_to_module_map.h000066400000000000000000000033251515451615200244600ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #pragma once #include #include #include #include namespace gandiva { template class IdToModuleMap { public: IdToModuleMap() : module_id_(kInitModuleId) {} jlong Insert(HOLDER holder) { mtx_.lock(); jlong result = module_id_++; map_.insert(std::pair(result, holder)); mtx_.unlock(); return result; } void Erase(jlong module_id) { mtx_.lock(); map_.erase(module_id); mtx_.unlock(); } HOLDER Lookup(jlong module_id) { HOLDER result = nullptr; mtx_.lock(); try { result = map_.at(module_id); } catch (const std::out_of_range&) { } mtx_.unlock(); return result; } private: static const int kInitModuleId = 4; int64_t module_id_; std::mutex mtx_; // map from module ids returned to Java and module pointers std::unordered_map map_; }; } // namespace gandiva apache-arrow-java-19.0.0/gandiva/src/main/cpp/jni_common.cc000066400000000000000000001067211515451615200234520ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include "config_holder.h" #include "env_helper.h" #include "id_to_module_map.h" #include "module_holder.h" using gandiva::ConditionPtr; using gandiva::DataTypePtr; using gandiva::ExpressionPtr; using gandiva::ExpressionVector; using gandiva::FieldPtr; using gandiva::FieldVector; using gandiva::Filter; using gandiva::NodePtr; using gandiva::NodeVector; using gandiva::Projector; using gandiva::SchemaPtr; using gandiva::Status; using gandiva::TreeExprBuilder; using gandiva::ArrayDataVector; using gandiva::ConfigHolder; using gandiva::Configuration; using gandiva::ConfigurationBuilder; using gandiva::FilterHolder; using gandiva::ProjectorHolder; // forward declarations NodePtr ProtoTypeToNode(const gandiva::types::TreeNode& node); static jint JNI_VERSION = JNI_VERSION_10; // extern refs - initialized for other modules. jclass configuration_builder_class_; // refs for self. static jclass gandiva_exception_; static jclass vector_expander_class_; static jclass vector_expander_ret_class_; static jmethodID vector_expander_method_; static jfieldID vector_expander_ret_address_; static jfieldID vector_expander_ret_capacity_; // module maps gandiva::IdToModuleMap> projector_modules_; gandiva::IdToModuleMap> filter_modules_; jint JNI_OnLoad(JavaVM* vm, void* reserved) { JNIEnv* env; if (vm->GetEnv(reinterpret_cast(&env), JNI_VERSION) != JNI_OK) { return JNI_ERR; } jclass local_configuration_builder_class_ = env->FindClass("org/apache/arrow/gandiva/evaluator/ConfigurationBuilder"); configuration_builder_class_ = (jclass)env->NewGlobalRef(local_configuration_builder_class_); env->DeleteLocalRef(local_configuration_builder_class_); jclass localExceptionClass = env->FindClass("org/apache/arrow/gandiva/exceptions/GandivaException"); gandiva_exception_ = (jclass)env->NewGlobalRef(localExceptionClass); env->ExceptionDescribe(); env->DeleteLocalRef(localExceptionClass); jclass local_expander_class = env->FindClass("org/apache/arrow/gandiva/evaluator/VectorExpander"); vector_expander_class_ = (jclass)env->NewGlobalRef(local_expander_class); env->DeleteLocalRef(local_expander_class); vector_expander_method_ = env->GetMethodID( vector_expander_class_, "expandOutputVectorAtIndex", "(IJ)Lorg/apache/arrow/gandiva/evaluator/VectorExpander$ExpandResult;"); jclass local_expander_ret_class = env->FindClass("org/apache/arrow/gandiva/evaluator/VectorExpander$ExpandResult"); vector_expander_ret_class_ = (jclass)env->NewGlobalRef(local_expander_ret_class); env->DeleteLocalRef(local_expander_ret_class); vector_expander_ret_address_ = env->GetFieldID(vector_expander_ret_class_, "address", "J"); vector_expander_ret_capacity_ = env->GetFieldID(vector_expander_ret_class_, "capacity", "J"); return JNI_VERSION; } void JNI_OnUnload(JavaVM* vm, void* reserved) { JNIEnv* env; vm->GetEnv(reinterpret_cast(&env), JNI_VERSION); env->DeleteGlobalRef(configuration_builder_class_); env->DeleteGlobalRef(gandiva_exception_); env->DeleteGlobalRef(vector_expander_class_); env->DeleteGlobalRef(vector_expander_ret_class_); } DataTypePtr ProtoTypeToTime32(const gandiva::types::ExtGandivaType& ext_type) { switch (ext_type.timeunit()) { case gandiva::types::SEC: return arrow::time32(arrow::TimeUnit::SECOND); case gandiva::types::MILLISEC: return arrow::time32(arrow::TimeUnit::MILLI); default: std::cerr << "Unknown time unit: " << ext_type.timeunit() << " for time32\n"; return nullptr; } } DataTypePtr ProtoTypeToTime64(const gandiva::types::ExtGandivaType& ext_type) { switch (ext_type.timeunit()) { case gandiva::types::MICROSEC: return arrow::time64(arrow::TimeUnit::MICRO); case gandiva::types::NANOSEC: return arrow::time64(arrow::TimeUnit::NANO); default: std::cerr << "Unknown time unit: " << ext_type.timeunit() << " for time64\n"; return nullptr; } } DataTypePtr ProtoTypeToTimestamp(const gandiva::types::ExtGandivaType& ext_type) { switch (ext_type.timeunit()) { case gandiva::types::SEC: return arrow::timestamp(arrow::TimeUnit::SECOND); case gandiva::types::MILLISEC: return arrow::timestamp(arrow::TimeUnit::MILLI); case gandiva::types::MICROSEC: return arrow::timestamp(arrow::TimeUnit::MICRO); case gandiva::types::NANOSEC: return arrow::timestamp(arrow::TimeUnit::NANO); default: std::cerr << "Unknown time unit: " << ext_type.timeunit() << " for timestamp\n"; return nullptr; } } DataTypePtr ProtoTypeToInterval(const gandiva::types::ExtGandivaType& ext_type) { switch (ext_type.intervaltype()) { case gandiva::types::YEAR_MONTH: return arrow::month_interval(); case gandiva::types::DAY_TIME: return arrow::day_time_interval(); default: std::cerr << "Unknown interval type: " << ext_type.intervaltype() << "\n"; return nullptr; } } DataTypePtr ProtoTypeToDataType(const gandiva::types::ExtGandivaType& ext_type) { switch (ext_type.type()) { case gandiva::types::NONE: return arrow::null(); case gandiva::types::BOOL: return arrow::boolean(); case gandiva::types::UINT8: return arrow::uint8(); case gandiva::types::INT8: return arrow::int8(); case gandiva::types::UINT16: return arrow::uint16(); case gandiva::types::INT16: return arrow::int16(); case gandiva::types::UINT32: return arrow::uint32(); case gandiva::types::INT32: return arrow::int32(); case gandiva::types::UINT64: return arrow::uint64(); case gandiva::types::INT64: return arrow::int64(); case gandiva::types::HALF_FLOAT: return arrow::float16(); case gandiva::types::FLOAT: return arrow::float32(); case gandiva::types::DOUBLE: return arrow::float64(); case gandiva::types::UTF8: return arrow::utf8(); case gandiva::types::BINARY: return arrow::binary(); case gandiva::types::DATE32: return arrow::date32(); case gandiva::types::DATE64: return arrow::date64(); case gandiva::types::DECIMAL: // TODO: error handling return arrow::decimal(ext_type.precision(), ext_type.scale()); case gandiva::types::TIME32: return ProtoTypeToTime32(ext_type); case gandiva::types::TIME64: return ProtoTypeToTime64(ext_type); case gandiva::types::TIMESTAMP: return ProtoTypeToTimestamp(ext_type); case gandiva::types::INTERVAL: return ProtoTypeToInterval(ext_type); case gandiva::types::FIXED_SIZE_BINARY: case gandiva::types::LIST: case gandiva::types::STRUCT: case gandiva::types::UNION: case gandiva::types::DICTIONARY: case gandiva::types::MAP: std::cerr << "Unhandled data type: " << ext_type.type() << "\n"; return nullptr; default: std::cerr << "Unknown data type: " << ext_type.type() << "\n"; return nullptr; } } FieldPtr ProtoTypeToField(const gandiva::types::Field& f) { const std::string& name = f.name(); DataTypePtr type = ProtoTypeToDataType(f.type()); bool nullable = true; if (f.has_nullable()) { nullable = f.nullable(); } return field(name, type, nullable); } NodePtr ProtoTypeToFieldNode(const gandiva::types::FieldNode& node) { FieldPtr field_ptr = ProtoTypeToField(node.field()); if (field_ptr == nullptr) { std::cerr << "Unable to create field node from protobuf\n"; return nullptr; } return TreeExprBuilder::MakeField(field_ptr); } NodePtr ProtoTypeToFnNode(const gandiva::types::FunctionNode& node) { const std::string& name = node.functionname(); NodeVector children; for (int i = 0; i < node.inargs_size(); i++) { const gandiva::types::TreeNode& arg = node.inargs(i); NodePtr n = ProtoTypeToNode(arg); if (n == nullptr) { std::cerr << "Unable to create argument for function: " << name << "\n"; return nullptr; } children.push_back(n); } DataTypePtr return_type = ProtoTypeToDataType(node.returntype()); if (return_type == nullptr) { std::cerr << "Unknown return type for function: " << name << "\n"; return nullptr; } return TreeExprBuilder::MakeFunction(name, children, return_type); } NodePtr ProtoTypeToIfNode(const gandiva::types::IfNode& node) { NodePtr cond = ProtoTypeToNode(node.cond()); if (cond == nullptr) { std::cerr << "Unable to create cond node for if node\n"; return nullptr; } NodePtr then_node = ProtoTypeToNode(node.thennode()); if (then_node == nullptr) { std::cerr << "Unable to create then node for if node\n"; return nullptr; } NodePtr else_node = ProtoTypeToNode(node.elsenode()); if (else_node == nullptr) { std::cerr << "Unable to create else node for if node\n"; return nullptr; } DataTypePtr return_type = ProtoTypeToDataType(node.returntype()); if (return_type == nullptr) { std::cerr << "Unknown return type for if node\n"; return nullptr; } return TreeExprBuilder::MakeIf(cond, then_node, else_node, return_type); } NodePtr ProtoTypeToAndNode(const gandiva::types::AndNode& node) { NodeVector children; for (int i = 0; i < node.args_size(); i++) { const gandiva::types::TreeNode& arg = node.args(i); NodePtr n = ProtoTypeToNode(arg); if (n == nullptr) { std::cerr << "Unable to create argument for boolean and\n"; return nullptr; } children.push_back(n); } return TreeExprBuilder::MakeAnd(children); } NodePtr ProtoTypeToOrNode(const gandiva::types::OrNode& node) { NodeVector children; for (int i = 0; i < node.args_size(); i++) { const gandiva::types::TreeNode& arg = node.args(i); NodePtr n = ProtoTypeToNode(arg); if (n == nullptr) { std::cerr << "Unable to create argument for boolean or\n"; return nullptr; } children.push_back(n); } return TreeExprBuilder::MakeOr(children); } NodePtr ProtoTypeToInNode(const gandiva::types::InNode& node) { NodePtr field = ProtoTypeToNode(node.node()); if (node.has_intvalues()) { std::unordered_set int_values; for (int i = 0; i < node.intvalues().intvalues_size(); i++) { int_values.insert(node.intvalues().intvalues(i).value()); } return TreeExprBuilder::MakeInExpressionInt32(field, int_values); } if (node.has_longvalues()) { std::unordered_set long_values; for (int i = 0; i < node.longvalues().longvalues_size(); i++) { long_values.insert(node.longvalues().longvalues(i).value()); } return TreeExprBuilder::MakeInExpressionInt64(field, long_values); } if (node.has_decimalvalues()) { std::unordered_set decimal_values; for (int i = 0; i < node.decimalvalues().decimalvalues_size(); i++) { decimal_values.insert( gandiva::DecimalScalar128(node.decimalvalues().decimalvalues(i).value(), node.decimalvalues().decimalvalues(i).precision(), node.decimalvalues().decimalvalues(i).scale())); } return TreeExprBuilder::MakeInExpressionDecimal(field, decimal_values); } if (node.has_floatvalues()) { std::unordered_set float_values; for (int i = 0; i < node.floatvalues().floatvalues_size(); i++) { float_values.insert(node.floatvalues().floatvalues(i).value()); } return TreeExprBuilder::MakeInExpressionFloat(field, float_values); } if (node.has_doublevalues()) { std::unordered_set double_values; for (int i = 0; i < node.doublevalues().doublevalues_size(); i++) { double_values.insert(node.doublevalues().doublevalues(i).value()); } return TreeExprBuilder::MakeInExpressionDouble(field, double_values); } if (node.has_stringvalues()) { std::unordered_set stringvalues; for (int i = 0; i < node.stringvalues().stringvalues_size(); i++) { stringvalues.insert(node.stringvalues().stringvalues(i).value()); } return TreeExprBuilder::MakeInExpressionString(field, stringvalues); } if (node.has_binaryvalues()) { std::unordered_set stringvalues; for (int i = 0; i < node.binaryvalues().binaryvalues_size(); i++) { stringvalues.insert(node.binaryvalues().binaryvalues(i).value()); } return TreeExprBuilder::MakeInExpressionBinary(field, stringvalues); } // not supported yet. std::cerr << "Unknown constant type for in expression.\n"; return nullptr; } NodePtr ProtoTypeToNullNode(const gandiva::types::NullNode& node) { DataTypePtr data_type = ProtoTypeToDataType(node.type()); if (data_type == nullptr) { std::cerr << "Unknown type " << data_type->ToString() << " for null node\n"; return nullptr; } return TreeExprBuilder::MakeNull(data_type); } NodePtr ProtoTypeToNode(const gandiva::types::TreeNode& node) { if (node.has_fieldnode()) { return ProtoTypeToFieldNode(node.fieldnode()); } if (node.has_fnnode()) { return ProtoTypeToFnNode(node.fnnode()); } if (node.has_ifnode()) { return ProtoTypeToIfNode(node.ifnode()); } if (node.has_andnode()) { return ProtoTypeToAndNode(node.andnode()); } if (node.has_ornode()) { return ProtoTypeToOrNode(node.ornode()); } if (node.has_innode()) { return ProtoTypeToInNode(node.innode()); } if (node.has_nullnode()) { return ProtoTypeToNullNode(node.nullnode()); } if (node.has_intnode()) { return TreeExprBuilder::MakeLiteral(node.intnode().value()); } if (node.has_floatnode()) { return TreeExprBuilder::MakeLiteral(node.floatnode().value()); } if (node.has_longnode()) { return TreeExprBuilder::MakeLiteral(node.longnode().value()); } if (node.has_booleannode()) { return TreeExprBuilder::MakeLiteral(node.booleannode().value()); } if (node.has_doublenode()) { return TreeExprBuilder::MakeLiteral(node.doublenode().value()); } if (node.has_stringnode()) { return TreeExprBuilder::MakeStringLiteral(node.stringnode().value()); } if (node.has_binarynode()) { return TreeExprBuilder::MakeBinaryLiteral(node.binarynode().value()); } if (node.has_decimalnode()) { std::string value = node.decimalnode().value(); gandiva::DecimalScalar128 literal(value, node.decimalnode().precision(), node.decimalnode().scale()); return TreeExprBuilder::MakeDecimalLiteral(literal); } std::cerr << "Unknown node type in protobuf\n"; return nullptr; } ExpressionPtr ProtoTypeToExpression(const gandiva::types::ExpressionRoot& root) { NodePtr root_node = ProtoTypeToNode(root.root()); if (root_node == nullptr) { std::cerr << "Unable to create expression node from expression protobuf\n"; return nullptr; } FieldPtr field = ProtoTypeToField(root.resulttype()); if (field == nullptr) { std::cerr << "Unable to extra return field from expression protobuf\n"; return nullptr; } return TreeExprBuilder::MakeExpression(root_node, field); } ConditionPtr ProtoTypeToCondition(const gandiva::types::Condition& condition) { NodePtr root_node = ProtoTypeToNode(condition.root()); if (root_node == nullptr) { return nullptr; } return TreeExprBuilder::MakeCondition(root_node); } SchemaPtr ProtoTypeToSchema(const gandiva::types::Schema& schema) { std::vector fields; for (int i = 0; i < schema.columns_size(); i++) { FieldPtr field = ProtoTypeToField(schema.columns(i)); if (field == nullptr) { std::cerr << "Unable to extract arrow field from schema\n"; return nullptr; } fields.push_back(field); } return arrow::schema(fields); } // Common for both projector and filters. bool ParseProtobuf(uint8_t* buf, int bufLen, google::protobuf::Message* msg) { google::protobuf::io::CodedInputStream cis(buf, bufLen); cis.SetRecursionLimit(2000); return msg->ParseFromCodedStream(&cis); } Status make_record_batch_with_buf_addrs(SchemaPtr schema, int num_rows, jlong* in_buf_addrs, jlong* in_buf_sizes, int in_bufs_len, std::shared_ptr* batch) { std::vector> columns; auto num_fields = schema->num_fields(); int buf_idx = 0; int sz_idx = 0; for (int i = 0; i < num_fields; i++) { auto field = schema->field(i); std::vector> buffers; if (buf_idx >= in_bufs_len) { return Status::Invalid("insufficient number of in_buf_addrs"); } jlong validity_addr = in_buf_addrs[buf_idx++]; jlong validity_size = in_buf_sizes[sz_idx++]; auto validity = std::shared_ptr( new arrow::Buffer(reinterpret_cast(validity_addr), validity_size)); buffers.push_back(validity); if (buf_idx >= in_bufs_len) { return Status::Invalid("insufficient number of in_buf_addrs"); } jlong value_addr = in_buf_addrs[buf_idx++]; jlong value_size = in_buf_sizes[sz_idx++]; auto data = std::shared_ptr( new arrow::Buffer(reinterpret_cast(value_addr), value_size)); buffers.push_back(data); if (arrow::is_binary_like(field->type()->id())) { if (buf_idx >= in_bufs_len) { return Status::Invalid("insufficient number of in_buf_addrs"); } // add offsets buffer for variable-len fields. jlong offsets_addr = in_buf_addrs[buf_idx++]; jlong offsets_size = in_buf_sizes[sz_idx++]; auto offsets = std::shared_ptr( new arrow::Buffer(reinterpret_cast(offsets_addr), offsets_size)); buffers.push_back(offsets); } auto array_data = arrow::ArrayData::Make(field->type(), num_rows, std::move(buffers)); columns.push_back(array_data); } *batch = arrow::RecordBatch::Make(schema, num_rows, columns); return Status::OK(); } // projector related functions. void releaseProjectorInput(jbyteArray schema_arr, jbyte* schema_bytes, jbyteArray exprs_arr, jbyte* exprs_bytes, JNIEnv* env) { env->ReleaseByteArrayElements(schema_arr, schema_bytes, JNI_ABORT); env->ReleaseByteArrayElements(exprs_arr, exprs_bytes, JNI_ABORT); } JNIEXPORT jlong JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_buildProjector( JNIEnv* env, jobject obj, jbyteArray schema_arr, jbyteArray exprs_arr, jint selection_vector_type, jlong configuration_id) { jlong module_id = 0LL; std::shared_ptr projector; std::shared_ptr holder; gandiva::types::Schema schema; jsize schema_len = env->GetArrayLength(schema_arr); jbyte* schema_bytes = env->GetByteArrayElements(schema_arr, 0); gandiva::types::ExpressionList exprs; jsize exprs_len = env->GetArrayLength(exprs_arr); jbyte* exprs_bytes = env->GetByteArrayElements(exprs_arr, 0); ExpressionVector expr_vector; SchemaPtr schema_ptr; FieldVector ret_types; gandiva::Status status; auto mode = gandiva::SelectionVector::MODE_NONE; std::shared_ptr config = ConfigHolder::MapLookup(configuration_id); std::stringstream ss; if (config == nullptr) { ss << "configuration is mandatory."; releaseProjectorInput(schema_arr, schema_bytes, exprs_arr, exprs_bytes, env); goto err_out; } if (!ParseProtobuf(reinterpret_cast(schema_bytes), schema_len, &schema)) { ss << "Unable to parse schema protobuf\n"; releaseProjectorInput(schema_arr, schema_bytes, exprs_arr, exprs_bytes, env); goto err_out; } if (!ParseProtobuf(reinterpret_cast(exprs_bytes), exprs_len, &exprs)) { releaseProjectorInput(schema_arr, schema_bytes, exprs_arr, exprs_bytes, env); ss << "Unable to parse expressions protobuf\n"; goto err_out; } // convert gandiva::types::Schema to arrow::Schema schema_ptr = ProtoTypeToSchema(schema); if (schema_ptr == nullptr) { ss << "Unable to construct arrow schema object from schema protobuf\n"; releaseProjectorInput(schema_arr, schema_bytes, exprs_arr, exprs_bytes, env); goto err_out; } // create Expression out of the list of exprs for (int i = 0; i < exprs.exprs_size(); i++) { ExpressionPtr root = ProtoTypeToExpression(exprs.exprs(i)); if (root == nullptr) { ss << "Unable to construct expression object from expression protobuf\n"; releaseProjectorInput(schema_arr, schema_bytes, exprs_arr, exprs_bytes, env); goto err_out; } expr_vector.push_back(root); ret_types.push_back(root->result()); } switch (selection_vector_type) { case gandiva::types::SV_NONE: mode = gandiva::SelectionVector::MODE_NONE; break; case gandiva::types::SV_INT16: mode = gandiva::SelectionVector::MODE_UINT16; break; case gandiva::types::SV_INT32: mode = gandiva::SelectionVector::MODE_UINT32; break; } // good to invoke the evaluator now status = Projector::Make(schema_ptr, expr_vector, mode, config, &projector); if (!status.ok()) { ss << "Failed to make LLVM module due to " << status.message() << "\n"; releaseProjectorInput(schema_arr, schema_bytes, exprs_arr, exprs_bytes, env); goto err_out; } // store the result in a map holder = std::shared_ptr( new ProjectorHolder(schema_ptr, ret_types, std::move(projector))); module_id = projector_modules_.Insert(holder); releaseProjectorInput(schema_arr, schema_bytes, exprs_arr, exprs_bytes, env); return module_id; err_out: env->ThrowNew(gandiva_exception_, ss.str().c_str()); return module_id; } /// /// \brief Resizable buffer which resizes by doing a callback into java. /// class JavaResizableBuffer : public arrow::ResizableBuffer { public: JavaResizableBuffer(JNIEnv* env, jobject jexpander, int32_t vector_idx, uint8_t* buffer, int32_t len) : ResizableBuffer(buffer, len), env_(env), jexpander_(jexpander), vector_idx_(vector_idx) { size_ = 0; } Status Resize(const int64_t new_size, bool shrink_to_fit) override; Status Reserve(const int64_t new_capacity) override; private: JNIEnv* env_; jobject jexpander_; int32_t vector_idx_; }; Status JavaResizableBuffer::Reserve(const int64_t new_capacity) { // callback into java to expand the buffer jobject ret = env_->CallObjectMethod(jexpander_, vector_expander_method_, vector_idx_, new_capacity); if (env_->ExceptionCheck()) { env_->ExceptionDescribe(); env_->ExceptionClear(); return Status::OutOfMemory("buffer expand failed in java"); } jlong ret_address = env_->GetLongField(ret, vector_expander_ret_address_); jlong ret_capacity = env_->GetLongField(ret, vector_expander_ret_capacity_); data_ = reinterpret_cast(ret_address); capacity_ = ret_capacity; return Status::OK(); } Status JavaResizableBuffer::Resize(const int64_t new_size, bool shrink_to_fit) { if (shrink_to_fit == true) { return Status::NotImplemented("shrink not implemented"); } if (ARROW_PREDICT_TRUE(new_size <= capacity())) { // no need to expand. size_ = new_size; return Status::OK(); } RETURN_NOT_OK(Reserve(new_size)); ARROW_DCHECK_GE(capacity_, new_size); size_ = new_size; return Status::OK(); } #define CHECK_OUT_BUFFER_IDX_AND_BREAK(idx, len) \ if (idx >= len) { \ status = gandiva::Status::Invalid("insufficient number of out_buf_addrs"); \ break; \ } JNIEXPORT void JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_evaluateProjector( JNIEnv* env, jobject object, jobject jexpander, jlong module_id, jint num_rows, jlongArray buf_addrs, jlongArray buf_sizes, jint sel_vec_type, jint sel_vec_rows, jlong sel_vec_addr, jlong sel_vec_size, jlongArray out_buf_addrs, jlongArray out_buf_sizes) { Status status; std::shared_ptr holder = projector_modules_.Lookup(module_id); if (holder == nullptr) { std::stringstream ss; ss << "Unknown module id " << module_id; env->ThrowNew(gandiva_exception_, ss.str().c_str()); return; } int in_bufs_len = env->GetArrayLength(buf_addrs); if (in_bufs_len != env->GetArrayLength(buf_sizes)) { env->ThrowNew(gandiva_exception_, "mismatch in arraylen of buf_addrs and buf_sizes"); return; } int out_bufs_len = env->GetArrayLength(out_buf_addrs); if (out_bufs_len != env->GetArrayLength(out_buf_sizes)) { env->ThrowNew(gandiva_exception_, "mismatch in arraylen of out_buf_addrs and out_buf_sizes"); return; } jlong* in_buf_addrs = env->GetLongArrayElements(buf_addrs, 0); jlong* in_buf_sizes = env->GetLongArrayElements(buf_sizes, 0); jlong* out_bufs = env->GetLongArrayElements(out_buf_addrs, 0); jlong* out_sizes = env->GetLongArrayElements(out_buf_sizes, 0); do { std::shared_ptr in_batch; status = make_record_batch_with_buf_addrs(holder->schema(), num_rows, in_buf_addrs, in_buf_sizes, in_bufs_len, &in_batch); if (!status.ok()) { break; } std::shared_ptr selection_vector; auto selection_buffer = std::make_shared( reinterpret_cast(sel_vec_addr), sel_vec_size); int output_row_count = 0; switch (sel_vec_type) { case gandiva::types::SV_NONE: { output_row_count = num_rows; break; } case gandiva::types::SV_INT16: { status = gandiva::SelectionVector::MakeImmutableInt16( sel_vec_rows, selection_buffer, &selection_vector); output_row_count = sel_vec_rows; break; } case gandiva::types::SV_INT32: { status = gandiva::SelectionVector::MakeImmutableInt32( sel_vec_rows, selection_buffer, &selection_vector); output_row_count = sel_vec_rows; break; } } if (!status.ok()) { break; } auto ret_types = holder->rettypes(); ArrayDataVector output; int buf_idx = 0; int sz_idx = 0; int output_vector_idx = 0; for (FieldPtr field : ret_types) { std::vector> buffers; CHECK_OUT_BUFFER_IDX_AND_BREAK(buf_idx, out_bufs_len); uint8_t* validity_buf = reinterpret_cast(out_bufs[buf_idx++]); jlong bitmap_sz = out_sizes[sz_idx++]; buffers.push_back(std::make_shared(validity_buf, bitmap_sz)); if (arrow::is_binary_like(field->type()->id())) { CHECK_OUT_BUFFER_IDX_AND_BREAK(buf_idx, out_bufs_len); uint8_t* offsets_buf = reinterpret_cast(out_bufs[buf_idx++]); jlong offsets_sz = out_sizes[sz_idx++]; buffers.push_back( std::make_shared(offsets_buf, offsets_sz)); } CHECK_OUT_BUFFER_IDX_AND_BREAK(buf_idx, out_bufs_len); uint8_t* value_buf = reinterpret_cast(out_bufs[buf_idx++]); jlong data_sz = out_sizes[sz_idx++]; if (arrow::is_binary_like(field->type()->id())) { if (jexpander == nullptr) { status = Status::Invalid( "expression has variable len output columns, but the expander object is " "null"); break; } buffers.push_back(std::make_shared( env, jexpander, output_vector_idx, value_buf, data_sz)); } else { buffers.push_back(std::make_shared(value_buf, data_sz)); } auto array_data = arrow::ArrayData::Make(field->type(), output_row_count, buffers); output.push_back(array_data); ++output_vector_idx; } if (!status.ok()) { break; } status = holder->projector()->Evaluate(*in_batch, selection_vector.get(), output); } while (0); env->ReleaseLongArrayElements(buf_addrs, in_buf_addrs, JNI_ABORT); env->ReleaseLongArrayElements(buf_sizes, in_buf_sizes, JNI_ABORT); env->ReleaseLongArrayElements(out_buf_addrs, out_bufs, JNI_ABORT); env->ReleaseLongArrayElements(out_buf_sizes, out_sizes, JNI_ABORT); if (!status.ok()) { std::stringstream ss; ss << "Evaluate returned " << status.message() << "\n"; env->ThrowNew(gandiva_exception_, status.message().c_str()); return; } } JNIEXPORT void JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_closeProjector( JNIEnv* env, jobject cls, jlong module_id) { projector_modules_.Erase(module_id); } // filter related functions. void releaseFilterInput(jbyteArray schema_arr, jbyte* schema_bytes, jbyteArray condition_arr, jbyte* condition_bytes, JNIEnv* env) { env->ReleaseByteArrayElements(schema_arr, schema_bytes, JNI_ABORT); env->ReleaseByteArrayElements(condition_arr, condition_bytes, JNI_ABORT); } JNIEXPORT jlong JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_buildFilter( JNIEnv* env, jobject obj, jbyteArray schema_arr, jbyteArray condition_arr, jlong configuration_id) { jlong module_id = 0LL; std::shared_ptr filter; std::shared_ptr holder; gandiva::types::Schema schema; jsize schema_len = env->GetArrayLength(schema_arr); jbyte* schema_bytes = env->GetByteArrayElements(schema_arr, 0); gandiva::types::Condition condition; jsize condition_len = env->GetArrayLength(condition_arr); jbyte* condition_bytes = env->GetByteArrayElements(condition_arr, 0); ConditionPtr condition_ptr; SchemaPtr schema_ptr; gandiva::Status status; std::shared_ptr config = ConfigHolder::MapLookup(configuration_id); std::stringstream ss; if (config == nullptr) { ss << "configuration is mandatory."; releaseFilterInput(schema_arr, schema_bytes, condition_arr, condition_bytes, env); goto err_out; } if (!ParseProtobuf(reinterpret_cast(schema_bytes), schema_len, &schema)) { ss << "Unable to parse schema protobuf\n"; releaseFilterInput(schema_arr, schema_bytes, condition_arr, condition_bytes, env); goto err_out; } if (!ParseProtobuf(reinterpret_cast(condition_bytes), condition_len, &condition)) { ss << "Unable to parse condition protobuf\n"; releaseFilterInput(schema_arr, schema_bytes, condition_arr, condition_bytes, env); goto err_out; } // convert gandiva::types::Schema to arrow::Schema schema_ptr = ProtoTypeToSchema(schema); if (schema_ptr == nullptr) { ss << "Unable to construct arrow schema object from schema protobuf\n"; releaseFilterInput(schema_arr, schema_bytes, condition_arr, condition_bytes, env); goto err_out; } condition_ptr = ProtoTypeToCondition(condition); if (condition_ptr == nullptr) { ss << "Unable to construct condition object from condition protobuf\n"; releaseFilterInput(schema_arr, schema_bytes, condition_arr, condition_bytes, env); goto err_out; } // good to invoke the filter builder now status = Filter::Make(schema_ptr, condition_ptr, config, &filter); if (!status.ok()) { ss << "Failed to make LLVM module due to " << status.message() << "\n"; releaseFilterInput(schema_arr, schema_bytes, condition_arr, condition_bytes, env); goto err_out; } // store the result in a map holder = std::shared_ptr(new FilterHolder(schema_ptr, std::move(filter))); module_id = filter_modules_.Insert(holder); releaseFilterInput(schema_arr, schema_bytes, condition_arr, condition_bytes, env); return module_id; err_out: env->ThrowNew(gandiva_exception_, ss.str().c_str()); return module_id; } JNIEXPORT jint JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_evaluateFilter( JNIEnv* env, jobject cls, jlong module_id, jint num_rows, jlongArray buf_addrs, jlongArray buf_sizes, jint jselection_vector_type, jlong out_buf_addr, jlong out_buf_size) { gandiva::Status status; std::shared_ptr holder = filter_modules_.Lookup(module_id); if (holder == nullptr) { env->ThrowNew(gandiva_exception_, "Unknown module id\n"); return -1; } int in_bufs_len = env->GetArrayLength(buf_addrs); if (in_bufs_len != env->GetArrayLength(buf_sizes)) { env->ThrowNew(gandiva_exception_, "mismatch in arraylen of buf_addrs and buf_sizes"); return -1; } jlong* in_buf_addrs = env->GetLongArrayElements(buf_addrs, 0); jlong* in_buf_sizes = env->GetLongArrayElements(buf_sizes, 0); std::shared_ptr selection_vector; do { std::shared_ptr in_batch; status = make_record_batch_with_buf_addrs(holder->schema(), num_rows, in_buf_addrs, in_buf_sizes, in_bufs_len, &in_batch); if (!status.ok()) { break; } auto selection_vector_type = static_cast(jselection_vector_type); auto out_buffer = std::make_shared( reinterpret_cast(out_buf_addr), out_buf_size); switch (selection_vector_type) { case gandiva::types::SV_INT16: status = gandiva::SelectionVector::MakeInt16(num_rows, out_buffer, &selection_vector); break; case gandiva::types::SV_INT32: status = gandiva::SelectionVector::MakeInt32(num_rows, out_buffer, &selection_vector); break; default: status = gandiva::Status::Invalid("unknown selection vector type"); } if (!status.ok()) { break; } status = holder->filter()->Evaluate(*in_batch, selection_vector); } while (0); env->ReleaseLongArrayElements(buf_addrs, in_buf_addrs, JNI_ABORT); env->ReleaseLongArrayElements(buf_sizes, in_buf_sizes, JNI_ABORT); if (!status.ok()) { std::stringstream ss; ss << "Evaluate returned " << status.message() << "\n"; env->ThrowNew(gandiva_exception_, status.message().c_str()); return -1; } else { int64_t num_slots = selection_vector->GetNumSlots(); // Check integer overflow if (num_slots > INT_MAX) { std::stringstream ss; ss << "The selection vector has " << num_slots << " slots, which is larger than the " << INT_MAX << " limit.\n"; const std::string message = ss.str(); env->ThrowNew(gandiva_exception_, message.c_str()); return -1; } return static_cast(num_slots); } } JNIEXPORT void JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_closeFilter( JNIEnv* env, jobject cls, jlong module_id) { filter_modules_.Erase(module_id); } apache-arrow-java-19.0.0/gandiva/src/main/cpp/module_holder.h000066400000000000000000000033461515451615200240050ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #pragma once #include #include #include namespace gandiva { class Projector; class Filter; class ProjectorHolder { public: ProjectorHolder(SchemaPtr schema, FieldVector ret_types, std::shared_ptr projector) : schema_(schema), ret_types_(ret_types), projector_(std::move(projector)) {} SchemaPtr schema() { return schema_; } FieldVector rettypes() { return ret_types_; } std::shared_ptr projector() { return projector_; } private: SchemaPtr schema_; FieldVector ret_types_; std::shared_ptr projector_; }; class FilterHolder { public: FilterHolder(SchemaPtr schema, std::shared_ptr filter) : schema_(schema), filter_(std::move(filter)) {} SchemaPtr schema() { return schema_; } std::shared_ptr filter() { return filter_; } private: SchemaPtr schema_; std::shared_ptr filter_; }; } // namespace gandiva apache-arrow-java-19.0.0/gandiva/src/main/cpp/symbols.map000066400000000000000000000015251515451615200231760ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. { global: extern "C++" { gandiva*; }; Java*; JNI*; local: *; }; apache-arrow-java-19.0.0/gandiva/src/main/java/000077500000000000000000000000001515451615200211435ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/module-info.java000066400000000000000000000023041515451615200242230ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ open module org.apache.arrow.gandiva { exports org.apache.arrow.gandiva.expression; exports org.apache.arrow.gandiva.exceptions; exports org.apache.arrow.gandiva.evaluator; exports org.apache.arrow.gandiva.ipc; requires com.google.common; requires com.google.protobuf; requires org.apache.arrow.format; requires org.apache.arrow.memory.core; requires org.apache.arrow.vector; requires org.slf4j; } apache-arrow-java-19.0.0/gandiva/src/main/java/org/000077500000000000000000000000001515451615200217325ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/000077500000000000000000000000001515451615200231535ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200243055ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/000077500000000000000000000000001515451615200257165ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/000077500000000000000000000000001515451615200277205ustar00rootroot00000000000000ConfigurationBuilder.java000066400000000000000000000042021515451615200346200ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import java.util.Objects; /** Used to construct gandiva configuration objects. */ public class ConfigurationBuilder { public long buildConfigInstance(ConfigOptions configOptions) { return buildConfigInstance(configOptions.optimize, configOptions.targetCPU); } private native long buildConfigInstance(boolean optimize, boolean detectHostCPU); public native void releaseConfigInstance(long configId); /** ConfigOptions contains the configuration parameters to provide to gandiva. */ public static class ConfigOptions { private boolean optimize = true; private boolean targetCPU = true; public static ConfigOptions getDefault() { return new ConfigOptions(); } public ConfigOptions() {} public ConfigOptions withOptimize(boolean optimize) { this.optimize = optimize; return this; } public ConfigOptions withTargetCPU(boolean targetCPU) { this.targetCPU = targetCPU; return this; } @Override public int hashCode() { return Objects.hash(optimize, targetCPU); } @Override public boolean equals(Object obj) { if (!(obj instanceof ConfigOptions)) { return false; } return this.optimize == ((ConfigOptions) obj).optimize && this.targetCPU == ((ConfigOptions) obj).targetCPU; } } } DecimalTypeUtil.java000066400000000000000000000063761515451615200335560ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import org.apache.arrow.vector.types.pojo.ArrowType.Decimal; /** Utility methods for working with {@link Decimal} values. */ public class DecimalTypeUtil { private DecimalTypeUtil() {} /** Enum for supported mathematical operations. */ public enum OperationType { ADD, SUBTRACT, MULTIPLY, DIVIDE, MOD } private static final int MIN_ADJUSTED_SCALE = 6; /// The maximum precision representable by a 16-byte decimal private static final int MAX_PRECISION = 38; /** Determines the scale and precision of applying the given operation to the operands. */ public static Decimal getResultTypeForOperation( OperationType operation, Decimal operand1, Decimal operand2) { int s1 = operand1.getScale(); int s2 = operand2.getScale(); int p1 = operand1.getPrecision(); int p2 = operand2.getPrecision(); int resultScale = 0; int resultPrecision = 0; switch (operation) { case ADD: case SUBTRACT: resultScale = Math.max(operand1.getScale(), operand2.getScale()); resultPrecision = resultScale + Math.max( operand1.getPrecision() - operand1.getScale(), operand2.getPrecision() - operand2.getScale()) + 1; break; case MULTIPLY: resultScale = s1 + s2; resultPrecision = p1 + p2 + 1; break; case DIVIDE: resultScale = Math.max(MIN_ADJUSTED_SCALE, operand1.getScale() + operand2.getPrecision() + 1); resultPrecision = operand1.getPrecision() - operand1.getScale() + operand2.getScale() + resultScale; break; case MOD: resultScale = Math.max(operand1.getScale(), operand2.getScale()); resultPrecision = Math.min( operand1.getPrecision() - operand1.getScale(), operand2.getPrecision() - operand2.getScale()) + resultScale; break; default: throw new RuntimeException("Needs support"); } return adjustScaleIfNeeded(resultPrecision, resultScale); } private static Decimal adjustScaleIfNeeded(int precision, int scale) { if (precision > MAX_PRECISION) { int minScale = Math.min(scale, MIN_ADJUSTED_SCALE); int delta = precision - MAX_PRECISION; precision = MAX_PRECISION; scale = Math.max(scale - delta, minScale); } return new Decimal(precision, scale, 128); } } ExpressionRegistry.java000066400000000000000000000210421515451615200343730ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import com.google.common.collect.Lists; import com.google.common.collect.Sets; import com.google.protobuf.InvalidProtocolBufferException; import java.util.List; import java.util.Set; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; import org.apache.arrow.gandiva.ipc.GandivaTypes.ExtGandivaType; import org.apache.arrow.gandiva.ipc.GandivaTypes.GandivaDataTypes; import org.apache.arrow.gandiva.ipc.GandivaTypes.GandivaFunctions; import org.apache.arrow.gandiva.ipc.GandivaTypes.GandivaType; import org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.IntervalUnit; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.pojo.ArrowType; /** * Used to get the functions and data types supported by Gandiva. All types are in Arrow namespace. */ public class ExpressionRegistry { private static final int BIT_WIDTH8 = 8; private static final int BIT_WIDTH_16 = 16; private static final int BIT_WIDTH_32 = 32; private static final int BIT_WIDTH_64 = 64; private static final boolean IS_SIGNED_FALSE = false; private static final boolean IS_SIGNED_TRUE = true; private final Set supportedTypes; private final Set functionSignatures; private static volatile ExpressionRegistry INSTANCE; private ExpressionRegistry( Set supportedTypes, Set functionSignatures) { this.supportedTypes = supportedTypes; this.functionSignatures = functionSignatures; } /** * Returns a singleton instance of the class. * * @return singleton instance * @throws GandivaException if error in Gandiva Library integration. */ public static ExpressionRegistry getInstance() throws GandivaException { if (INSTANCE == null) { synchronized (ExpressionRegistry.class) { if (INSTANCE == null) { // ensure library is setup. JniLoader.getInstance(); Set typesFromGandiva = getSupportedTypesFromGandiva(); Set functionsFromGandiva = getSupportedFunctionsFromGandiva(); INSTANCE = new ExpressionRegistry(typesFromGandiva, functionsFromGandiva); } } } return INSTANCE; } public Set getSupportedFunctions() { return functionSignatures; } public Set getSupportedTypes() { return supportedTypes; } private static Set getSupportedTypesFromGandiva() throws GandivaException { Set supportedTypes = Sets.newHashSet(); try { byte[] gandivaSupportedDataTypes = new ExpressionRegistryJniHelper().getGandivaSupportedDataTypes(); GandivaDataTypes gandivaDataTypes = GandivaDataTypes.parseFrom(gandivaSupportedDataTypes); for (ExtGandivaType type : gandivaDataTypes.getDataTypeList()) { supportedTypes.add(getArrowType(type)); } } catch (InvalidProtocolBufferException invalidProtException) { throw new GandivaException("Could not get supported types.", invalidProtException); } return supportedTypes; } private static Set getSupportedFunctionsFromGandiva() throws GandivaException { Set supportedTypes = Sets.newHashSet(); try { byte[] gandivaSupportedFunctions = new ExpressionRegistryJniHelper().getGandivaSupportedFunctions(); GandivaFunctions gandivaFunctions = GandivaFunctions.parseFrom(gandivaSupportedFunctions); for (GandivaTypes.FunctionSignature protoFunctionSignature : gandivaFunctions.getFunctionList()) { String functionName = protoFunctionSignature.getName(); ArrowType returnType = getArrowType(protoFunctionSignature.getReturnType()); List paramTypes = Lists.newArrayList(); for (ExtGandivaType type : protoFunctionSignature.getParamTypesList()) { paramTypes.add(getArrowType(type)); } FunctionSignature functionSignature = new FunctionSignature(functionName, returnType, paramTypes); supportedTypes.add(functionSignature); } } catch (InvalidProtocolBufferException invalidProtException) { throw new GandivaException("Could not get supported functions.", invalidProtException); } return supportedTypes; } private static ArrowType getArrowType(ExtGandivaType type) { switch (type.getType().getNumber()) { case GandivaType.BOOL_VALUE: return ArrowType.Bool.INSTANCE; case GandivaType.UINT8_VALUE: return new ArrowType.Int(BIT_WIDTH8, IS_SIGNED_FALSE); case GandivaType.INT8_VALUE: return new ArrowType.Int(BIT_WIDTH8, IS_SIGNED_TRUE); case GandivaType.UINT16_VALUE: return new ArrowType.Int(BIT_WIDTH_16, IS_SIGNED_FALSE); case GandivaType.INT16_VALUE: return new ArrowType.Int(BIT_WIDTH_16, IS_SIGNED_TRUE); case GandivaType.UINT32_VALUE: return new ArrowType.Int(BIT_WIDTH_32, IS_SIGNED_FALSE); case GandivaType.INT32_VALUE: return new ArrowType.Int(BIT_WIDTH_32, IS_SIGNED_TRUE); case GandivaType.UINT64_VALUE: return new ArrowType.Int(BIT_WIDTH_64, IS_SIGNED_FALSE); case GandivaType.INT64_VALUE: return new ArrowType.Int(BIT_WIDTH_64, IS_SIGNED_TRUE); case GandivaType.HALF_FLOAT_VALUE: return new ArrowType.FloatingPoint(FloatingPointPrecision.HALF); case GandivaType.FLOAT_VALUE: return new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE); case GandivaType.DOUBLE_VALUE: return new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE); case GandivaType.UTF8_VALUE: return new ArrowType.Utf8(); case GandivaType.BINARY_VALUE: return new ArrowType.Binary(); case GandivaType.DATE32_VALUE: return new ArrowType.Date(DateUnit.DAY); case GandivaType.DATE64_VALUE: return new ArrowType.Date(DateUnit.MILLISECOND); case GandivaType.TIMESTAMP_VALUE: return new ArrowType.Timestamp(mapArrowTimeUnit(type.getTimeUnit()), null); case GandivaType.TIME32_VALUE: return new ArrowType.Time(mapArrowTimeUnit(type.getTimeUnit()), BIT_WIDTH_32); case GandivaType.TIME64_VALUE: return new ArrowType.Time(mapArrowTimeUnit(type.getTimeUnit()), BIT_WIDTH_64); case GandivaType.NONE_VALUE: return new ArrowType.Null(); case GandivaType.DECIMAL_VALUE: return new ArrowType.Decimal(0, 0, 128); case GandivaType.INTERVAL_VALUE: return new ArrowType.Interval(mapArrowIntervalUnit(type.getIntervalType())); case GandivaType.FIXED_SIZE_BINARY_VALUE: case GandivaType.MAP_VALUE: case GandivaType.DICTIONARY_VALUE: case GandivaType.LIST_VALUE: case GandivaType.STRUCT_VALUE: case GandivaType.UNION_VALUE: default: assert false; } return null; } private static TimeUnit mapArrowTimeUnit(GandivaTypes.TimeUnit timeUnit) { switch (timeUnit.getNumber()) { case GandivaTypes.TimeUnit.MICROSEC_VALUE: return TimeUnit.MICROSECOND; case GandivaTypes.TimeUnit.MILLISEC_VALUE: return TimeUnit.MILLISECOND; case GandivaTypes.TimeUnit.NANOSEC_VALUE: return TimeUnit.NANOSECOND; case GandivaTypes.TimeUnit.SEC_VALUE: return TimeUnit.SECOND; default: return null; } } private static IntervalUnit mapArrowIntervalUnit(GandivaTypes.IntervalType intervalType) { switch (intervalType.getNumber()) { case GandivaTypes.IntervalType.YEAR_MONTH_VALUE: return IntervalUnit.YEAR_MONTH; case GandivaTypes.IntervalType.DAY_TIME_VALUE: return IntervalUnit.DAY_TIME; default: return null; } } } ExpressionRegistryJniHelper.java000066400000000000000000000020411515451615200361720ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; /** JNI Adapter used to get supported types and functions from Gandiva. */ class ExpressionRegistryJniHelper { native byte[] getGandivaSupportedDataTypes(); native byte[] getGandivaSupportedFunctions(); } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Filter.java000066400000000000000000000203451515451615200320140ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import java.util.ArrayList; import java.util.List; import org.apache.arrow.gandiva.exceptions.EvaluatorClosedException; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.expression.ArrowTypeHelper; import org.apache.arrow.gandiva.expression.Condition; import org.apache.arrow.gandiva.ipc.GandivaTypes; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.ipc.message.ArrowBuffer; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.types.pojo.Schema; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * This class provides a mechanism to filter a RecordBatch by evaluating a condition expression. * Follow these steps to use this class: 1) Use the static method make() to create an instance of * this class that evaluates a condition. 2) Invoke the method evaluate() to evaluate the filter * against a RecordBatch 3) Invoke close() to release resources */ public class Filter { private static final Logger logger = LoggerFactory.getLogger(Filter.class); private final JniWrapper wrapper; private final long moduleId; private final Schema schema; private boolean closed; private Filter(JniWrapper wrapper, long moduleId, Schema schema) { this.wrapper = wrapper; this.moduleId = moduleId; this.schema = schema; this.closed = false; } /** * Invoke this function to generate LLVM code to evaluate the condition expression. Invoke * Filter::Evaluate() against a RecordBatch to evaluate the filter on this record batch * * @param schema Table schema. The field names in the schema should match the fields used to * create the TreeNodes * @param condition condition to be evaluated against data * @return A native filter object that can be used to invoke on a RecordBatch */ public static Filter make(Schema schema, Condition condition) throws GandivaException { return make(schema, condition, JniLoader.getDefaultConfiguration()); } /** * Invoke this function to generate LLVM code to evaluate the condition expression. Invoke * Filter::Evaluate() against a RecordBatch to evaluate the filter on this record batch * * @param schema Table schema. The field names in the schema should match the fields used to * create the TreeNodes * @param condition condition to be evaluated against data * @param configOptions ConfigOptions parameter * @return A native filter object that can be used to invoke on a RecordBatch */ public static Filter make( Schema schema, Condition condition, ConfigurationBuilder.ConfigOptions configOptions) throws GandivaException { return make(schema, condition, JniLoader.getConfiguration(configOptions)); } /** * Invoke this function to generate LLVM code to evaluate the condition expression. Invoke * Filter::Evaluate() against a RecordBatch to evaluate the filter on this record batch * * @param schema Table schema. The field names in the schema should match the fields used to * create the TreeNodes * @param condition condition to be evaluated against data * @param optimize Flag to choose if the generated llvm code is to be optimized * @return A native filter object that can be used to invoke on a RecordBatch */ @Deprecated public static Filter make(Schema schema, Condition condition, boolean optimize) throws GandivaException { return make( schema, condition, JniLoader.getConfiguration( (new ConfigurationBuilder.ConfigOptions()).withOptimize(optimize))); } /** * Invoke this function to generate LLVM code to evaluate the condition expression. Invoke * Filter::Evaluate() against a RecordBatch to evaluate the filter on this record batch * * @param schema Table schema. The field names in the schema should match the fields used to * create the TreeNodes * @param condition condition to be evaluated against data * @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 Filter make(Schema schema, Condition condition, long configurationId) throws GandivaException { // Invoke the JNI layer to create the LLVM module representing the filter. GandivaTypes.Condition conditionBuf = condition.toProtobuf(); GandivaTypes.Schema schemaBuf = ArrowTypeHelper.arrowSchemaToProtobuf(schema); JniWrapper wrapper = JniLoader.getInstance().getWrapper(); long moduleId = wrapper.buildFilter(schemaBuf.toByteArray(), conditionBuf.toByteArray(), configurationId); logger.debug("Created module for the filter with id {}", moduleId); return new Filter(wrapper, moduleId, schema); } /** * Invoke this function to evaluate a filter against a recordBatch. * * @param recordBatch Record batch including the data * @param selectionVector Result of applying the filter on the data */ public void evaluate(ArrowRecordBatch recordBatch, SelectionVector selectionVector) throws GandivaException { evaluate( recordBatch.getLength(), recordBatch.getBuffers(), recordBatch.getBuffersLayout(), selectionVector); } /** * Invoke this function to evaluate filter against a set of arrow buffers. (this is an optimised * version that skips taking references). * * @param numRows number of rows. * @param buffers List of input arrow buffers * @param selectionVector Result of applying the filter on the data */ public void evaluate(int numRows, List buffers, SelectionVector selectionVector) throws GandivaException { List buffersLayout = new ArrayList<>(); long offset = 0; for (ArrowBuf arrowBuf : buffers) { long size = arrowBuf.readableBytes(); buffersLayout.add(new ArrowBuffer(offset, size)); offset += size; } evaluate(numRows, buffers, buffersLayout, selectionVector); } private void evaluate( int numRows, List buffers, List buffersLayout, SelectionVector selectionVector) throws GandivaException { if (this.closed) { throw new EvaluatorClosedException(); } if (selectionVector.getMaxRecords() < numRows) { logger.error( "selectionVector has capacity for " + selectionVector.getMaxRecords() + " rows, minimum required " + numRows); throw new GandivaException("SelectionVector too small"); } long[] bufAddrs = new long[buffers.size()]; long[] bufSizes = new long[buffers.size()]; int idx = 0; for (ArrowBuf buf : buffers) { bufAddrs[idx++] = buf.memoryAddress(); } idx = 0; for (ArrowBuffer bufLayout : buffersLayout) { bufSizes[idx++] = bufLayout.getSize(); } int numRecords = wrapper.evaluateFilter( this.moduleId, numRows, bufAddrs, bufSizes, selectionVector.getType().getNumber(), selectionVector.getBuffer().memoryAddress(), selectionVector.getBuffer().capacity()); if (numRecords >= 0) { selectionVector.setRecordCount(numRecords); } } /** Closes the LLVM module representing this filter. */ public void close() throws GandivaException { if (this.closed) { return; } wrapper.closeFilter(this.moduleId); this.closed = true; } } FunctionSignature.java000066400000000000000000000051031515451615200341520ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import com.google.common.base.MoreObjects; import com.google.common.base.Objects; import java.util.List; import org.apache.arrow.vector.types.pojo.ArrowType; /** POJO to define a function signature. */ public class FunctionSignature { private final String name; private final ArrowType returnType; private final List paramTypes; public ArrowType getReturnType() { return returnType; } public List getParamTypes() { return paramTypes; } public String getName() { return name; } /** * Ctor. * * @param name - name of the function. * @param returnType - data type of return * @param paramTypes - data type of input args. */ public FunctionSignature(String name, ArrowType returnType, List paramTypes) { this.name = name; this.returnType = returnType; this.paramTypes = paramTypes; } /** * Override equals. * * @param signature - signature to compare * @return true if equal and false if not. */ public boolean equals(Object signature) { if (signature == null) { return false; } if (getClass() != signature.getClass()) { return false; } final FunctionSignature other = (FunctionSignature) signature; return this.name.equalsIgnoreCase(other.name) && Objects.equal(this.returnType, other.returnType) && Objects.equal(this.paramTypes, other.paramTypes); } @Override public int hashCode() { return Objects.hashCode(this.name.toLowerCase(), this.returnType, this.paramTypes); } @Override public String toString() { return MoreObjects.toStringHelper(this) .add("name ", name) .add("return type ", returnType) .add("param types ", paramTypes) .toString(); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/JniLoader.java000066400000000000000000000145761515451615200324470ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import static java.util.UUID.randomUUID; import java.io.File; import java.io.IOException; import java.io.InputStream; import java.nio.file.Files; import java.nio.file.StandardCopyOption; import java.util.Locale; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import org.apache.arrow.gandiva.exceptions.GandivaException; /** This class handles loading of the jni library, and acts as a bridge for the native functions. */ class JniLoader { private static final String LIBRARY_NAME = "gandiva_jni"; private static volatile JniLoader INSTANCE; private static volatile long defaultConfiguration = 0L; private static final ConcurrentMap configurationMap = new ConcurrentHashMap<>(); private final JniWrapper wrapper; private JniLoader() { this.wrapper = new JniWrapper(); } static JniLoader getInstance() throws GandivaException { if (INSTANCE == null) { synchronized (JniLoader.class) { if (INSTANCE == null) { INSTANCE = setupInstance(); } } } return INSTANCE; } private static JniLoader setupInstance() throws GandivaException { try { String tempDir = System.getProperty("java.io.tmpdir"); loadGandivaLibraryFromJar(tempDir); return new JniLoader(); } catch (IOException ioException) { throw new GandivaException("unable to create native instance", ioException); } } private static void loadGandivaLibraryFromJar(final String tmpDir) throws IOException, GandivaException { final String libraryToLoad = LIBRARY_NAME + "/" + getNormalizedArch() + "/" + System.mapLibraryName(LIBRARY_NAME); final File libraryFile = moveFileFromJarToTemp(tmpDir, libraryToLoad, LIBRARY_NAME); System.load(libraryFile.getAbsolutePath()); } private static String getNormalizedArch() { String arch = System.getProperty("os.arch").toLowerCase(Locale.US); switch (arch) { case "amd64": arch = "x86_64"; break; case "aarch64": arch = "aarch_64"; break; default: break; } return arch; } private static File moveFileFromJarToTemp( final String tmpDir, String libraryToLoad, String libraryName) throws IOException, GandivaException { final File temp = setupFile(tmpDir, libraryName); try (final InputStream is = JniLoader.class.getClassLoader().getResourceAsStream(libraryToLoad)) { if (is == null) { throw new GandivaException(libraryToLoad + " was not found inside JAR."); } else { Files.copy(is, temp.toPath(), StandardCopyOption.REPLACE_EXISTING); } } return temp; } private static File setupFile(String tmpDir, String libraryToLoad) throws IOException, GandivaException { // accommodate multiple processes running with gandiva jar. // length should be ok since uuid is only 36 characters. final String randomizeFileName = libraryToLoad + randomUUID(); final File temp = new File(tmpDir, randomizeFileName); if (temp.exists() && !temp.delete()) { throw new GandivaException( "File: " + temp.getAbsolutePath() + " already exists and cannot be removed."); } if (!temp.createNewFile()) { throw new GandivaException("File: " + temp.getAbsolutePath() + " could not be created."); } temp.deleteOnExit(); return temp; } /** Returns the jni wrapper. */ JniWrapper getWrapper() throws GandivaException { return wrapper; } static long getConfiguration(ConfigurationBuilder.ConfigOptions configOptions) throws GandivaException { if (!configurationMap.containsKey(configOptions)) { synchronized (ConfigurationBuilder.class) { if (!configurationMap.containsKey(configOptions)) { JniLoader.getInstance(); // setup long configInstance = new ConfigurationBuilder().buildConfigInstance(configOptions); configurationMap.put(configOptions, configInstance); if (ConfigurationBuilder.ConfigOptions.getDefault().equals(configOptions)) { defaultConfiguration = configInstance; } return configInstance; } } } return configurationMap.get(configOptions); } /** * Get the default configuration to invoke gandiva. * * @return default configuration * @throws GandivaException if unable to get native builder instance. */ static long getDefaultConfiguration() throws GandivaException { if (defaultConfiguration == 0L) { synchronized (ConfigurationBuilder.class) { if (defaultConfiguration == 0L) { JniLoader.getInstance(); // setup ConfigurationBuilder.ConfigOptions defaultConfigOptions = ConfigurationBuilder.ConfigOptions.getDefault(); defaultConfiguration = new ConfigurationBuilder().buildConfigInstance(defaultConfigOptions); configurationMap.put(defaultConfigOptions, defaultConfiguration); } } } return defaultConfiguration; } /** Remove the configuration. */ static void removeConfiguration(ConfigurationBuilder.ConfigOptions configOptions) { if (configurationMap.containsKey(configOptions)) { synchronized (ConfigurationBuilder.class) { if (configurationMap.containsKey(configOptions)) { (new ConfigurationBuilder()) .releaseConfigInstance(configurationMap.remove(configOptions)); if (configOptions.equals(ConfigurationBuilder.ConfigOptions.getDefault())) { defaultConfiguration = 0; } } } } } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/JniWrapper.java000066400000000000000000000122071515451615200326460ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import org.apache.arrow.gandiva.exceptions.GandivaException; /** * This class is implemented in JNI. This provides the Java interface to invoke functions in JNI. * This file is used to generated the .h files required for jni. Avoid all external dependencies in * this file. */ public class JniWrapper { /** * Generates the projector module to evaluate the expressions with custom configuration. * * @param schemaBuf The schema serialized as a protobuf. See Types.proto to see the protobuf * specification * @param exprListBuf The serialized protobuf of the expression vector. Each expression is created * using TreeBuilder::MakeExpression. * @param selectionVectorType type of selection vector * @param configId Configuration to gandiva. * @return A moduleId that is passed to the evaluateProjector() and closeProjector() methods */ native long buildProjector( byte[] schemaBuf, byte[] exprListBuf, int selectionVectorType, long configId) throws GandivaException; /** * Evaluate the expressions represented by the moduleId on a record batch and store the output in * ValueVectors. Throws an exception in case of errors * * @param expander VectorExpander object. Used for callbacks from cpp. * @param moduleId moduleId representing expressions. Created using a call to buildNativeCode * @param numRows Number of rows in the record batch * @param bufAddrs An array of memory addresses. Each memory address points to a validity vector * or a data vector (will add support for offset vectors later). * @param bufSizes An array of buffer sizes. For each memory address in bufAddrs, the size of the * buffer is present in bufSizes * @param outAddrs An array of output buffers, including the validity and data addresses. * @param outSizes The allocated size of the output buffers. On successful evaluation, the result * is stored in the output buffers */ native void evaluateProjector( Object expander, long moduleId, int numRows, long[] bufAddrs, long[] bufSizes, int selectionVectorType, int selectionVectorSize, long selectionVectorBufferAddr, long selectionVectorBufferSize, long[] outAddrs, long[] outSizes) throws GandivaException; /** * Closes the projector referenced by moduleId. * * @param moduleId moduleId that needs to be closed */ native void closeProjector(long moduleId); /** * Generates the filter module to evaluate the condition expression with custom configuration. * * @param schemaBuf The schema serialized as a protobuf. See Types.proto to see the protobuf * specification * @param conditionBuf The serialized protobuf of the condition expression. Each expression is * created using TreeBuilder::MakeCondition * @param configId Configuration to gandiva. * @return A moduleId that is passed to the evaluateFilter() and closeFilter() methods */ native long buildFilter(byte[] schemaBuf, byte[] conditionBuf, long configId) throws GandivaException; /** * Evaluate the filter represented by the moduleId on a record batch and store the output in * buffer 'outAddr'. Throws an exception in case of errors * * @param moduleId moduleId representing expressions. Created using a call to buildNativeCode * @param numRows Number of rows in the record batch * @param bufAddrs An array of memory addresses. Each memory address points to a validity vector * or a data vector (will add support for offset vectors later). * @param bufSizes An array of buffer sizes. For each memory address in bufAddrs, the size of the * buffer is present in bufSizes * @param selectionVectorType type of selection vector * @param outAddr output buffer, whose type is represented by selectionVectorType * @param outSize The allocated size of the output buffer. On successful evaluation, the result is * stored in the output buffer */ native int evaluateFilter( long moduleId, int numRows, long[] bufAddrs, long[] bufSizes, int selectionVectorType, long outAddr, long outSize) throws GandivaException; /** * Closes the filter referenced by moduleId. * * @param moduleId moduleId that needs to be closed */ native void closeFilter(long moduleId); } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Projector.java000066400000000000000000000371271515451615200325440ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import java.util.ArrayList; import java.util.List; import org.apache.arrow.gandiva.exceptions.EvaluatorClosedException; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.exceptions.UnsupportedTypeException; import org.apache.arrow.gandiva.expression.ArrowTypeHelper; import org.apache.arrow.gandiva.expression.ExpressionTree; import org.apache.arrow.gandiva.ipc.GandivaTypes; import org.apache.arrow.gandiva.ipc.GandivaTypes.SelectionVectorType; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.FixedWidthVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VariableWidthVector; import org.apache.arrow.vector.ipc.message.ArrowBuffer; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.types.pojo.Schema; /** * This class provides a mechanism to evaluate a set of expressions against a RecordBatch. Follow * these steps to use this class: 1) Use the static method make() to create an instance of this * class that evaluates a set of expressions 2) Invoke the method evaluate() to evaluate these * expressions against a RecordBatch 3) Invoke close() to release resources */ public class Projector { private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(Projector.class); private JniWrapper wrapper; private final long moduleId; private final Schema schema; private final int numExprs; private boolean closed; private Projector(JniWrapper wrapper, long moduleId, Schema schema, int numExprs) { this.wrapper = wrapper; this.moduleId = moduleId; this.schema = schema; this.numExprs = numExprs; this.closed = false; } /** * Invoke this function to generate LLVM code to evaluate the list of project expressions. Invoke * Projector::Evaluate() against a RecordBatch to evaluate the record batch against these * projections. * * @param schema Table schema. The field names in the schema should match the fields used to * create the TreeNodes * @param exprs List of expressions to be evaluated against data * @return A native evaluator object that can be used to invoke these projections on a RecordBatch */ public static Projector make(Schema schema, List exprs) throws GandivaException { return make(schema, exprs, SelectionVectorType.SV_NONE, JniLoader.getDefaultConfiguration()); } /** * Invoke this function to generate LLVM code to evaluate the list of project expressions. Invoke * Projector::Evaluate() against a RecordBatch to evaluate the record batch against these * projections. * * @param schema Table schema. The field names in the schema should match the fields used to * create the TreeNodes * @param exprs List of expressions to be evaluated against data * @param configOptions ConfigOptions parameter * @return A native evaluator object that can be used to invoke these projections on a RecordBatch */ public static Projector make( Schema schema, List exprs, ConfigurationBuilder.ConfigOptions configOptions) throws GandivaException { return make( schema, exprs, SelectionVectorType.SV_NONE, JniLoader.getConfiguration(configOptions)); } /** * Invoke this function to generate LLVM code to evaluate the list of project expressions. Invoke * Projector::Evaluate() against a RecordBatch to evaluate the record batch against these * projections. * * @param schema Table schema. The field names in the schema should match the fields used to * create the TreeNodes * @param exprs List of expressions to be evaluated against data * @param optimize Flag to choose if the generated llvm code is to be optimized * @return A native evaluator object that can be used to invoke these projections on a RecordBatch */ @Deprecated public static Projector make(Schema schema, List exprs, boolean optimize) throws GandivaException { return make( schema, exprs, SelectionVectorType.SV_NONE, JniLoader.getConfiguration( (new ConfigurationBuilder.ConfigOptions()).withOptimize(optimize))); } /** * Invoke this function to generate LLVM code to evaluate the list of project expressions. Invoke * Projector::Evaluate() against a RecordBatch to evaluate the record batch against these * projections. * * @param schema Table schema. The field names in the schema should match the fields used to * create the TreeNodes * @param exprs List of expressions to be evaluated against data * @param selectionVectorType type of selection vector * @return A native evaluator object that can be used to invoke these projections on a RecordBatch */ public static Projector make( Schema schema, List exprs, SelectionVectorType selectionVectorType) throws GandivaException { return make(schema, exprs, selectionVectorType, JniLoader.getDefaultConfiguration()); } /** * Invoke this function to generate LLVM code to evaluate the list of project expressions. Invoke * Projector::Evaluate() against a RecordBatch to evaluate the record batch against these * projections. * * @param schema Table schema. The field names in the schema should match the fields used to * create the TreeNodes * @param exprs List of expressions to be evaluated against data * @param selectionVectorType type of selection vector * @param configOptions ConfigOptions parameter * @return A native evaluator object that can be used to invoke these projections on a RecordBatch */ public static Projector make( Schema schema, List exprs, SelectionVectorType selectionVectorType, ConfigurationBuilder.ConfigOptions configOptions) throws GandivaException { return make(schema, exprs, selectionVectorType, JniLoader.getConfiguration(configOptions)); } /** * Invoke this function to generate LLVM code to evaluate the list of project expressions. Invoke * Projector::Evaluate() against a RecordBatch to evaluate the record batch against these * projections. * * @param schema Table schema. The field names in the schema should match the fields used to * create the TreeNodes * @param exprs List of expressions to be evaluated against data * @param selectionVectorType type of selection vector * @param optimize Flag to choose if the generated llvm code is to be optimized * @return A native evaluator object that can be used to invoke these projections on a RecordBatch */ @Deprecated public static Projector make( Schema schema, List exprs, SelectionVectorType selectionVectorType, boolean optimize) throws GandivaException { return make( schema, exprs, selectionVectorType, JniLoader.getConfiguration( (new ConfigurationBuilder.ConfigOptions()).withOptimize(optimize))); } /** * Invoke this function to generate LLVM code to evaluate the list of project expressions. Invoke * Projector::Evaluate() against a RecordBatch to evaluate the record batch against these * projections. * * @param schema Table schema. The field names in the schema should match the fields used to * create the TreeNodes * @param exprs List of expressions to be evaluated against data * @param selectionVectorType type of selection vector * @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 synchronized Projector make( Schema schema, List exprs, SelectionVectorType selectionVectorType, long configurationId) throws GandivaException { // serialize the schema and the list of expressions as a protobuf GandivaTypes.ExpressionList.Builder builder = GandivaTypes.ExpressionList.newBuilder(); for (ExpressionTree expr : exprs) { builder.addExprs(expr.toProtobuf()); } // Invoke the JNI layer to create the LLVM module representing the expressions GandivaTypes.Schema schemaBuf = ArrowTypeHelper.arrowSchemaToProtobuf(schema); JniWrapper wrapper = JniLoader.getInstance().getWrapper(); long moduleId = wrapper.buildProjector( schemaBuf.toByteArray(), builder.build().toByteArray(), selectionVectorType.getNumber(), configurationId); logger.debug("Created module for the projector with id {}", moduleId); return new Projector(wrapper, moduleId, schema, exprs.size()); } /** * Invoke this function to evaluate a set of expressions against a recordBatch. * * @param recordBatch Record batch including the data * @param outColumns Result of applying the project on the data */ public void evaluate(ArrowRecordBatch recordBatch, List outColumns) throws GandivaException { evaluate( recordBatch.getLength(), recordBatch.getBuffers(), recordBatch.getBuffersLayout(), SelectionVectorType.SV_NONE.getNumber(), recordBatch.getLength(), 0, 0, outColumns); } /** * Invoke this function to evaluate a set of expressions against a set of arrow buffers. (this is * an optimised version that skips taking references). * * @param numRows number of rows. * @param buffers List of input arrow buffers * @param outColumns Result of applying the project on the data */ public void evaluate(int numRows, List buffers, List outColumns) throws GandivaException { List buffersLayout = new ArrayList<>(); long offset = 0; for (ArrowBuf arrowBuf : buffers) { long size = arrowBuf.readableBytes(); buffersLayout.add(new ArrowBuffer(offset, size)); offset += size; } evaluate( numRows, buffers, buffersLayout, SelectionVectorType.SV_NONE.getNumber(), numRows, 0, 0, outColumns); } /** * Invoke this function to evaluate a set of expressions against a {@link ArrowRecordBatch}. * * @param recordBatch The data to evaluate against. * @param selectionVector Selection vector which stores the selected rows. * @param outColumns Result of applying the project on the data */ public void evaluate( ArrowRecordBatch recordBatch, SelectionVector selectionVector, List outColumns) throws GandivaException { evaluate( recordBatch.getLength(), recordBatch.getBuffers(), recordBatch.getBuffersLayout(), selectionVector.getType().getNumber(), selectionVector.getRecordCount(), selectionVector.getBuffer().memoryAddress(), selectionVector.getBuffer().capacity(), outColumns); } /** * Invoke this function to evaluate a set of expressions against a set of arrow buffers on the * selected positions. (this is an optimised version that skips taking references). * * @param numRows number of rows. * @param buffers List of input arrow buffers * @param selectionVector Selection vector which stores the selected rows. * @param outColumns Result of applying the project on the data */ public void evaluate( int numRows, List buffers, SelectionVector selectionVector, List outColumns) throws GandivaException { List buffersLayout = new ArrayList<>(); long offset = 0; for (ArrowBuf arrowBuf : buffers) { long size = arrowBuf.readableBytes(); buffersLayout.add(new ArrowBuffer(offset, size)); offset += size; } evaluate( numRows, buffers, buffersLayout, selectionVector.getType().getNumber(), selectionVector.getRecordCount(), selectionVector.getBuffer().memoryAddress(), selectionVector.getBuffer().capacity(), outColumns); } private synchronized void evaluate( int numRows, List buffers, List buffersLayout, int selectionVectorType, int selectionVectorRecordCount, long selectionVectorAddr, long selectionVectorSize, List outColumns) throws GandivaException { if (this.closed) { throw new EvaluatorClosedException(); } if (numExprs != outColumns.size()) { logger.info("Expected " + numExprs + " columns, got " + outColumns.size()); throw new GandivaException("Incorrect number of columns for the output vector"); } long[] bufAddrs = new long[buffers.size()]; long[] bufSizes = new long[buffers.size()]; int idx = 0; for (ArrowBuf buf : buffers) { bufAddrs[idx++] = buf.memoryAddress(); } idx = 0; for (ArrowBuffer bufLayout : buffersLayout) { bufSizes[idx++] = bufLayout.getSize(); } boolean hasVariableWidthColumns = false; BaseVariableWidthVector[] resizableVectors = new BaseVariableWidthVector[outColumns.size()]; long[] outAddrs = new long[3 * outColumns.size()]; long[] outSizes = new long[3 * outColumns.size()]; idx = 0; int outColumnIdx = 0; for (ValueVector valueVector : outColumns) { boolean isFixedWith = valueVector instanceof FixedWidthVector; boolean isVarWidth = valueVector instanceof VariableWidthVector; if (!isFixedWith && !isVarWidth) { throw new UnsupportedTypeException( "Unsupported value vector type " + valueVector.getField().getFieldType()); } outAddrs[idx] = valueVector.getValidityBuffer().memoryAddress(); outSizes[idx++] = valueVector.getValidityBuffer().capacity(); if (isVarWidth) { outAddrs[idx] = valueVector.getOffsetBuffer().memoryAddress(); outSizes[idx++] = valueVector.getOffsetBuffer().capacity(); hasVariableWidthColumns = true; // save vector to allow for resizing. resizableVectors[outColumnIdx] = (BaseVariableWidthVector) valueVector; } outAddrs[idx] = valueVector.getDataBuffer().memoryAddress(); outSizes[idx++] = valueVector.getDataBuffer().capacity(); valueVector.setValueCount(selectionVectorRecordCount); outColumnIdx++; } wrapper.evaluateProjector( hasVariableWidthColumns ? new VectorExpander(resizableVectors) : null, this.moduleId, numRows, bufAddrs, bufSizes, selectionVectorType, selectionVectorRecordCount, selectionVectorAddr, selectionVectorSize, outAddrs, outSizes); } /** Closes the LLVM module representing this evaluator. */ public void close() throws GandivaException { if (this.closed) { return; } wrapper.closeProjector(this.moduleId); this.closed = true; } } SelectionVector.java000066400000000000000000000051701515451615200336170ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import static org.apache.arrow.memory.util.LargeMemoryUtil.capAtMaxInt; import org.apache.arrow.gandiva.ipc.GandivaTypes.SelectionVectorType; import org.apache.arrow.memory.ArrowBuf; /** * A selection vector contains the indexes of "selected" records in a row batch. It is backed by an * arrow buffer. Client manages the lifecycle of the arrow buffer - to release the reference. */ public abstract class SelectionVector { private int recordCount; private ArrowBuf buffer; public SelectionVector(ArrowBuf buffer) { this.buffer = buffer; } public final ArrowBuf getBuffer() { return this.buffer; } /* * The maximum number of records that the selection vector can hold. */ public final int getMaxRecords() { return capAtMaxInt(buffer.capacity() / getRecordSize()); } /* * The number of records held by the selection vector. */ public final int getRecordCount() { return this.recordCount; } /* * Set the number of records in the selection vector. */ final void setRecordCount(int recordCount) { if (recordCount * getRecordSize() > buffer.capacity()) { throw new IllegalArgumentException( "recordCount " + recordCount + " of size " + getRecordSize() + " exceeds buffer capacity " + buffer.capacity()); } this.recordCount = recordCount; } /* * Get the value at specified index. */ public abstract int getIndex(int index); /* * Get the record size of the selection vector itself. */ abstract int getRecordSize(); abstract SelectionVectorType getType(); final void checkReadBounds(int index) { if (index >= this.recordCount) { throw new IllegalArgumentException("index " + index + " is >= recordCount " + recordCount); } } } SelectionVectorInt16.java000066400000000000000000000027101515451615200344360ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import org.apache.arrow.gandiva.ipc.GandivaTypes.SelectionVectorType; import org.apache.arrow.memory.ArrowBuf; /** Selection vector with records of arrow type INT16. */ public class SelectionVectorInt16 extends SelectionVector { public SelectionVectorInt16(ArrowBuf buffer) { super(buffer); } @Override public int getRecordSize() { return 2; } @Override public SelectionVectorType getType() { return SelectionVectorType.SV_INT16; } @Override public int getIndex(int index) { checkReadBounds(index); char value = getBuffer().getChar(index * getRecordSize()); return (int) value; } } SelectionVectorInt32.java000066400000000000000000000026511515451615200344400ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import org.apache.arrow.gandiva.ipc.GandivaTypes.SelectionVectorType; import org.apache.arrow.memory.ArrowBuf; /** Selection vector with records of arrow type INT32. */ public class SelectionVectorInt32 extends SelectionVector { public SelectionVectorInt32(ArrowBuf buffer) { super(buffer); } @Override public int getRecordSize() { return 4; } @Override public SelectionVectorType getType() { return SelectionVectorType.SV_INT32; } @Override public int getIndex(int index) { checkReadBounds(index); return getBuffer().getInt(index * getRecordSize()); } } VectorExpander.java000066400000000000000000000043531515451615200334420ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import org.apache.arrow.vector.BaseVariableWidthVector; /** * This class provides the functionality to expand output vectors using a callback mechanism from * gandiva. */ public class VectorExpander { private final BaseVariableWidthVector[] vectors; public VectorExpander(BaseVariableWidthVector[] vectors) { this.vectors = vectors; } /** Result of vector expansion. */ public static class ExpandResult { public long address; public long capacity; public ExpandResult(long address, long capacity) { this.address = address; this.capacity = capacity; } } /** * Expand vector at specified index. This is used as a back call from jni, and is only relevant * for variable width vectors. * * @param index index of buffer in the list passed to jni. * @param toCapacity the size to which the buffer should be expanded to. * @return address and size of the buffer after expansion. */ public ExpandResult expandOutputVectorAtIndex(int index, long toCapacity) { if (index >= vectors.length || vectors[index] == null) { throw new IllegalArgumentException("invalid index " + index); } BaseVariableWidthVector vector = vectors[index]; while (vector.getDataBuffer().capacity() < toCapacity) { vector.reallocDataBuffer(); } return new ExpandResult( vector.getDataBuffer().memoryAddress(), vector.getDataBuffer().capacity()); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/exceptions/000077500000000000000000000000001515451615200300775ustar00rootroot00000000000000EvaluatorClosedException.java000066400000000000000000000021031515451615200356320ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/exceptions/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.exceptions; /** Indicates an attempted call to methods on a closed evaluator. */ public class EvaluatorClosedException extends GandivaException { public EvaluatorClosedException() { super("Cannot invoke methods on evaluator after closing it"); } } GandivaException.java000066400000000000000000000022401515451615200341110ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/exceptions/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.exceptions; /** Base class for all specialized exceptions this package uses. */ public class GandivaException extends Exception { public GandivaException(String msg) { super(msg); } public GandivaException(String msg, Exception cause) { super(msg, cause); } @Override public String toString() { return getMessage(); } } UnsupportedTypeException.java000066400000000000000000000020421515451615200357320ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/exceptions/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.exceptions; /** Represents an exception thrown while dealing with unsupported types. */ public class UnsupportedTypeException extends GandivaException { public UnsupportedTypeException(String msg) { super(msg); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/000077500000000000000000000000001515451615200301155ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/AndNode.java000066400000000000000000000030571515451615200322750ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import java.util.List; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; /** Node representing a logical And expression. */ class AndNode implements TreeNode { private final List children; AndNode(List children) { this.children = children; } @Override public GandivaTypes.TreeNode toProtobuf() throws GandivaException { GandivaTypes.AndNode.Builder andNode = GandivaTypes.AndNode.newBuilder(); for (TreeNode arg : children) { andNode.addArgs(arg.toProtobuf()); } GandivaTypes.TreeNode.Builder builder = GandivaTypes.TreeNode.newBuilder(); builder.setAndNode(andNode.build()); return builder.build(); } } ArrowTypeHelper.java000066400000000000000000000267001515451615200340020ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import org.apache.arrow.flatbuf.DateUnit; import org.apache.arrow.flatbuf.IntervalUnit; import org.apache.arrow.flatbuf.TimeUnit; import org.apache.arrow.flatbuf.Type; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.exceptions.UnsupportedTypeException; import org.apache.arrow.gandiva.ipc.GandivaTypes; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; /** Utility methods to convert between Arrow and Gandiva types. */ public class ArrowTypeHelper { private ArrowTypeHelper() {} static final int WIDTH_8 = 8; static final int WIDTH_16 = 16; static final int WIDTH_32 = 32; static final int WIDTH_64 = 64; private static void initArrowTypeInt( ArrowType.Int intType, GandivaTypes.ExtGandivaType.Builder builder) throws GandivaException { int width = intType.getBitWidth(); if (intType.getIsSigned()) { switch (width) { case WIDTH_8: { builder.setType(GandivaTypes.GandivaType.INT8); return; } case WIDTH_16: { builder.setType(GandivaTypes.GandivaType.INT16); return; } case WIDTH_32: { builder.setType(GandivaTypes.GandivaType.INT32); return; } case WIDTH_64: { builder.setType(GandivaTypes.GandivaType.INT64); return; } default: { throw new UnsupportedTypeException("Unsupported width for integer type"); } } } // unsigned int switch (width) { case WIDTH_8: { builder.setType(GandivaTypes.GandivaType.UINT8); return; } case WIDTH_16: { builder.setType(GandivaTypes.GandivaType.UINT16); return; } case WIDTH_32: { builder.setType(GandivaTypes.GandivaType.UINT32); return; } case WIDTH_64: { builder.setType(GandivaTypes.GandivaType.UINT64); return; } default: { throw new UnsupportedTypeException("Unsupported width for integer type"); } } } private static void initArrowTypeFloat( ArrowType.FloatingPoint floatType, GandivaTypes.ExtGandivaType.Builder builder) throws GandivaException { switch (floatType.getPrecision()) { case HALF: { builder.setType(GandivaTypes.GandivaType.HALF_FLOAT); break; } case SINGLE: { builder.setType(GandivaTypes.GandivaType.FLOAT); break; } case DOUBLE: { builder.setType(GandivaTypes.GandivaType.DOUBLE); break; } default: { throw new UnsupportedTypeException("Floating point type with unknown precision"); } } } private static void initArrowTypeDecimal( ArrowType.Decimal decimalType, GandivaTypes.ExtGandivaType.Builder builder) { Preconditions.checkArgument( decimalType.getPrecision() > 0 && decimalType.getPrecision() <= 38, "Gandiva only supports decimals of upto 38 " + "precision. Input precision : " + decimalType.getPrecision()); builder.setPrecision(decimalType.getPrecision()); builder.setScale(decimalType.getScale()); builder.setType(GandivaTypes.GandivaType.DECIMAL); } private static void initArrowTypeDate( ArrowType.Date dateType, GandivaTypes.ExtGandivaType.Builder builder) { short dateUnit = dateType.getUnit().getFlatbufID(); switch (dateUnit) { case DateUnit.DAY: { builder.setType(GandivaTypes.GandivaType.DATE32); break; } case DateUnit.MILLISECOND: { builder.setType(GandivaTypes.GandivaType.DATE64); break; } default: { // not supported break; } } } private static void initArrowTypeTime( ArrowType.Time timeType, GandivaTypes.ExtGandivaType.Builder builder) { short timeUnit = timeType.getUnit().getFlatbufID(); switch (timeUnit) { case TimeUnit.SECOND: { builder.setType(GandivaTypes.GandivaType.TIME32); builder.setTimeUnit(GandivaTypes.TimeUnit.SEC); break; } case TimeUnit.MILLISECOND: { builder.setType(GandivaTypes.GandivaType.TIME32); builder.setTimeUnit(GandivaTypes.TimeUnit.MILLISEC); break; } case TimeUnit.MICROSECOND: { builder.setType(GandivaTypes.GandivaType.TIME64); builder.setTimeUnit(GandivaTypes.TimeUnit.MICROSEC); break; } case TimeUnit.NANOSECOND: { builder.setType(GandivaTypes.GandivaType.TIME64); builder.setTimeUnit(GandivaTypes.TimeUnit.NANOSEC); break; } default: { // not supported } } } private static void initArrowTypeTimestamp( ArrowType.Timestamp timestampType, GandivaTypes.ExtGandivaType.Builder builder) { short timeUnit = timestampType.getUnit().getFlatbufID(); switch (timeUnit) { case TimeUnit.SECOND: { builder.setType(GandivaTypes.GandivaType.TIMESTAMP); builder.setTimeUnit(GandivaTypes.TimeUnit.SEC); break; } case TimeUnit.MILLISECOND: { builder.setType(GandivaTypes.GandivaType.TIMESTAMP); builder.setTimeUnit(GandivaTypes.TimeUnit.MILLISEC); break; } case TimeUnit.MICROSECOND: { builder.setType(GandivaTypes.GandivaType.TIMESTAMP); builder.setTimeUnit(GandivaTypes.TimeUnit.MICROSEC); break; } case TimeUnit.NANOSECOND: { builder.setType(GandivaTypes.GandivaType.TIMESTAMP); builder.setTimeUnit(GandivaTypes.TimeUnit.NANOSEC); break; } default: { // not supported } } } private static void initArrowTypeInterval( ArrowType.Interval interval, GandivaTypes.ExtGandivaType.Builder builder) { short intervalUnit = interval.getUnit().getFlatbufID(); switch (intervalUnit) { case IntervalUnit.YEAR_MONTH: { builder.setType(GandivaTypes.GandivaType.INTERVAL); builder.setIntervalType(GandivaTypes.IntervalType.YEAR_MONTH); break; } case IntervalUnit.DAY_TIME: { builder.setType(GandivaTypes.GandivaType.INTERVAL); builder.setIntervalType(GandivaTypes.IntervalType.DAY_TIME); break; } default: { // not supported } } } /** * Converts an arrow type into a protobuf. * * @param arrowType Arrow type to be converted * @return Protobuf representing the arrow type */ public static GandivaTypes.ExtGandivaType arrowTypeToProtobuf(ArrowType arrowType) throws GandivaException { GandivaTypes.ExtGandivaType.Builder builder = GandivaTypes.ExtGandivaType.newBuilder(); byte typeId = arrowType.getTypeID().getFlatbufID(); switch (typeId) { case Type.NONE: { // 0 builder.setType(GandivaTypes.GandivaType.NONE); break; } case Type.Null: { // 1 // TODO: Need to handle this later break; } case Type.Int: { // 2 ArrowTypeHelper.initArrowTypeInt((ArrowType.Int) arrowType, builder); break; } case Type.FloatingPoint: { // 3 ArrowTypeHelper.initArrowTypeFloat((ArrowType.FloatingPoint) arrowType, builder); break; } case Type.Binary: { // 4 builder.setType(GandivaTypes.GandivaType.BINARY); break; } case Type.Utf8: { // 5 builder.setType(GandivaTypes.GandivaType.UTF8); break; } case Type.Bool: { // 6 builder.setType(GandivaTypes.GandivaType.BOOL); break; } case Type.Decimal: { // 7 ArrowTypeHelper.initArrowTypeDecimal((ArrowType.Decimal) arrowType, builder); break; } case Type.Date: { // 8 ArrowTypeHelper.initArrowTypeDate((ArrowType.Date) arrowType, builder); break; } case Type.Time: { // 9 ArrowTypeHelper.initArrowTypeTime((ArrowType.Time) arrowType, builder); break; } case Type.Timestamp: { // 10 ArrowTypeHelper.initArrowTypeTimestamp((ArrowType.Timestamp) arrowType, builder); break; } case Type.Interval: { // 11 ArrowTypeHelper.initArrowTypeInterval((ArrowType.Interval) arrowType, builder); break; } case Type.List: { // 12 break; } case Type.Struct_: { // 13 break; } case Type.Union: { // 14 break; } case Type.FixedSizeBinary: { // 15 break; } case Type.FixedSizeList: { // 16 break; } case Type.Map: { // 17 break; } default: { break; } } if (!builder.hasType()) { // type has not been set // throw an exception throw new UnsupportedTypeException("Unsupported type " + arrowType.toString()); } return builder.build(); } /** * Converts an arrow field object to a protobuf. * * @param field Arrow field to be converted * @return Protobuf representing the arrow field */ public static GandivaTypes.Field arrowFieldToProtobuf(Field field) throws GandivaException { GandivaTypes.Field.Builder builder = GandivaTypes.Field.newBuilder(); builder.setName(field.getName()); builder.setType(ArrowTypeHelper.arrowTypeToProtobuf(field.getType())); builder.setNullable(field.isNullable()); for (Field child : field.getChildren()) { builder.addChildren(ArrowTypeHelper.arrowFieldToProtobuf(child)); } return builder.build(); } /** * Converts a schema object to a protobuf. * * @param schema Schema object to be converted * @return Protobuf representing a schema object */ public static GandivaTypes.Schema arrowSchemaToProtobuf(Schema schema) throws GandivaException { GandivaTypes.Schema.Builder builder = GandivaTypes.Schema.newBuilder(); for (Field field : schema.getFields()) { builder.addColumns(ArrowTypeHelper.arrowFieldToProtobuf(field)); } return builder.build(); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/BinaryNode.java000066400000000000000000000027511515451615200330170ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import com.google.protobuf.ByteString; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; /** Used to represent expression tree nodes representing binary constants. */ class BinaryNode implements TreeNode { private final byte[] value; public BinaryNode(byte[] value) { this.value = value; } @Override public GandivaTypes.TreeNode toProtobuf() throws GandivaException { GandivaTypes.BinaryNode binaryNode = GandivaTypes.BinaryNode.newBuilder().setValue(ByteString.copyFrom(value)).build(); return GandivaTypes.TreeNode.newBuilder().setBinaryNode(binaryNode).build(); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/BooleanNode.java000066400000000000000000000031071515451615200331460ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; /** * Used to represent expression tree nodes representing boolean constants. Used while creating * expressions like if (!x). */ class BooleanNode implements TreeNode { private final Boolean value; BooleanNode(Boolean value) { this.value = value; } @Override public GandivaTypes.TreeNode toProtobuf() throws GandivaException { GandivaTypes.BooleanNode.Builder boolBuilder = GandivaTypes.BooleanNode.newBuilder(); boolBuilder.setValue(value.booleanValue()); GandivaTypes.TreeNode.Builder builder = GandivaTypes.TreeNode.newBuilder(); builder.setBooleanNode(boolBuilder.build()); return builder.build(); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/Condition.java000066400000000000000000000027231515451615200327120ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; /** Opaque class representing a filter condition. */ public class Condition { private final TreeNode root; Condition(TreeNode root) { this.root = root; } /** * Converts an condition expression into a protobuf. * * @return A protobuf representing the condition expression tree */ public GandivaTypes.Condition toProtobuf() throws GandivaException { GandivaTypes.Condition.Builder builder = GandivaTypes.Condition.newBuilder(); builder.setRoot(root.toProtobuf()); return builder.build(); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/DecimalNode.java000066400000000000000000000033771515451615200331360ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; /** * Used to represent expression tree nodes representing decimal constants. Used in the expression (x * + 5.0) */ class DecimalNode implements TreeNode { private final String value; private final int precision; private final int scale; DecimalNode(String value, int precision, int scale) { this.value = value; this.precision = precision; this.scale = scale; } @Override public GandivaTypes.TreeNode toProtobuf() throws GandivaException { GandivaTypes.DecimalNode.Builder decimalNode = GandivaTypes.DecimalNode.newBuilder(); decimalNode.setValue(value); decimalNode.setPrecision(precision); decimalNode.setScale(scale); GandivaTypes.TreeNode.Builder builder = GandivaTypes.TreeNode.newBuilder(); builder.setDecimalNode(decimalNode.build()); return builder.build(); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/DoubleNode.java000066400000000000000000000030671515451615200330060ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; /** * Used to represent expression tree nodes representing double constants. Used in the expression (x * + 5.0) */ class DoubleNode implements TreeNode { private final Double value; DoubleNode(Double value) { this.value = value; } @Override public GandivaTypes.TreeNode toProtobuf() throws GandivaException { GandivaTypes.DoubleNode.Builder doubleBuilder = GandivaTypes.DoubleNode.newBuilder(); doubleBuilder.setValue(value.doubleValue()); GandivaTypes.TreeNode.Builder builder = GandivaTypes.TreeNode.newBuilder(); builder.setDoubleNode(doubleBuilder.build()); return builder.build(); } } ExpressionTree.java000066400000000000000000000032611515451615200336620ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; import org.apache.arrow.vector.types.pojo.Field; /** Opaque class representing an expression. */ public class ExpressionTree { private final TreeNode root; private final Field resultField; ExpressionTree(TreeNode root, Field resultField) { this.root = root; this.resultField = resultField; } /** * Converts an expression tree into a protobuf. * * @return A protobuf representing the expression tree */ public GandivaTypes.ExpressionRoot toProtobuf() throws GandivaException { GandivaTypes.ExpressionRoot.Builder builder = GandivaTypes.ExpressionRoot.newBuilder(); builder.setRoot(root.toProtobuf()); builder.setResultType(ArrowTypeHelper.arrowFieldToProtobuf(resultField)); return builder.build(); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/FieldNode.java000066400000000000000000000030501515451615200326070ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; import org.apache.arrow.vector.types.pojo.Field; /** Opaque class that represents a tree node. */ class FieldNode implements TreeNode { private final Field field; FieldNode(Field field) { this.field = field; } @Override public GandivaTypes.TreeNode toProtobuf() throws GandivaException { GandivaTypes.FieldNode.Builder fieldNode = GandivaTypes.FieldNode.newBuilder(); fieldNode.setField(ArrowTypeHelper.arrowFieldToProtobuf(field)); GandivaTypes.TreeNode.Builder builder = GandivaTypes.TreeNode.newBuilder(); builder.setFieldNode(fieldNode.build()); return builder.build(); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/FloatNode.java000066400000000000000000000030621515451615200326340ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; /** * Used to represent expression tree nodes representing float constants. Used in the expression (x + * 5.0) */ class FloatNode implements TreeNode { private final Float value; public FloatNode(Float value) { this.value = value; } @Override public GandivaTypes.TreeNode toProtobuf() throws GandivaException { GandivaTypes.FloatNode.Builder floatBuilder = GandivaTypes.FloatNode.newBuilder(); floatBuilder.setValue(value.floatValue()); GandivaTypes.TreeNode.Builder builder = GandivaTypes.TreeNode.newBuilder(); builder.setFloatNode(floatBuilder.build()); return builder.build(); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/FunctionNode.java000066400000000000000000000036241515451615200333600ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import java.util.List; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; import org.apache.arrow.vector.types.pojo.ArrowType; /** Node representing an arbitrary function in an expression. */ class FunctionNode implements TreeNode { private final String function; private final List children; private final ArrowType retType; FunctionNode(String function, List children, ArrowType retType) { this.function = function; this.children = children; this.retType = retType; } @Override public GandivaTypes.TreeNode toProtobuf() throws GandivaException { GandivaTypes.FunctionNode.Builder fnNode = GandivaTypes.FunctionNode.newBuilder(); fnNode.setFunctionName(function); fnNode.setReturnType(ArrowTypeHelper.arrowTypeToProtobuf(retType)); for (TreeNode arg : children) { fnNode.addInArgs(arg.toProtobuf()); } GandivaTypes.TreeNode.Builder builder = GandivaTypes.TreeNode.newBuilder(); builder.setFnNode(fnNode.build()); return builder.build(); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/IfNode.java000066400000000000000000000037471515451615200321370ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; import org.apache.arrow.vector.types.pojo.ArrowType; /** Node representing a if-then-else block expression. */ class IfNode implements TreeNode { private final TreeNode condition; private final TreeNode thenNode; private final TreeNode elseNode; private final ArrowType retType; IfNode(TreeNode condition, TreeNode thenNode, TreeNode elseNode, ArrowType retType) { this.condition = condition; this.thenNode = thenNode; this.elseNode = elseNode; this.retType = retType; } @Override public GandivaTypes.TreeNode toProtobuf() throws GandivaException { GandivaTypes.IfNode.Builder ifNodeBuilder = GandivaTypes.IfNode.newBuilder(); ifNodeBuilder.setCond(condition.toProtobuf()); ifNodeBuilder.setThenNode(thenNode.toProtobuf()); ifNodeBuilder.setElseNode(elseNode.toProtobuf()); ifNodeBuilder.setReturnType(ArrowTypeHelper.arrowTypeToProtobuf(retType)); GandivaTypes.TreeNode.Builder builder = GandivaTypes.TreeNode.newBuilder(); builder.setIfNode(ifNodeBuilder.build()); return builder.build(); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/InNode.java000066400000000000000000000172501515451615200321410ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import com.google.protobuf.ByteString; import java.math.BigDecimal; import java.nio.charset.Charset; import java.util.Set; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; /** In Node representation in java. */ public class InNode implements TreeNode { private static final Charset charset = Charset.forName("UTF-8"); private final Set intValues; private final Set longValues; private final Set floatValues; private final Set doubleValues; private final Set decimalValues; private final Set stringValues; private final Set binaryValues; private final TreeNode input; private final Integer precision; private final Integer scale; private InNode( Set values, Set longValues, Set stringValues, Set binaryValues, Set decimalValues, Integer precision, Integer scale, Set floatValues, Set doubleValues, TreeNode node) { this.intValues = values; this.longValues = longValues; this.decimalValues = decimalValues; this.precision = precision; this.scale = scale; this.stringValues = stringValues; this.binaryValues = binaryValues; this.floatValues = floatValues; this.doubleValues = doubleValues; this.input = node; } /** * Makes an IN node for int values. * * @param node Node with the 'IN' clause. * @param intValues Int values to build the IN node. * @return InNode referring to tree node. */ public static InNode makeIntInExpr(TreeNode node, Set intValues) { return new InNode(intValues, null, null, null, null, null, null, null, null, node); } /** * Makes an IN node for long values. * * @param node Node with the 'IN' clause. * @param longValues Long values to build the IN node. * @return InNode referring to tree node. */ public static InNode makeLongInExpr(TreeNode node, Set longValues) { return new InNode(null, longValues, null, null, null, null, null, null, null, node); } /** * Makes an IN node for float values. * * @param node Node with the 'IN' clause. * @param floatValues Float values to build the IN node. * @return InNode referring to tree node. */ public static InNode makeFloatInExpr(TreeNode node, Set floatValues) { return new InNode(null, null, null, null, null, null, null, floatValues, null, node); } /** * Makes an IN node for double values. * * @param node Node with the 'IN' clause. * @param doubleValues Double values to build the IN node. * @return InNode referring to tree node. */ public static InNode makeDoubleInExpr(TreeNode node, Set doubleValues) { return new InNode(null, null, null, null, null, null, null, null, doubleValues, node); } public static InNode makeDecimalInExpr( TreeNode node, Set decimalValues, Integer precision, Integer scale) { return new InNode(null, null, null, null, decimalValues, precision, scale, null, null, node); } public static InNode makeStringInExpr(TreeNode node, Set stringValues) { return new InNode(null, null, stringValues, null, null, null, null, null, null, node); } public static InNode makeBinaryInExpr(TreeNode node, Set binaryValues) { return new InNode(null, null, null, binaryValues, null, null, null, null, null, node); } @Override public GandivaTypes.TreeNode toProtobuf() throws GandivaException { GandivaTypes.InNode.Builder inNode = GandivaTypes.InNode.newBuilder(); inNode.setNode(input.toProtobuf()); if (intValues != null) { GandivaTypes.IntConstants.Builder intConstants = GandivaTypes.IntConstants.newBuilder(); intValues.stream() .forEach( val -> intConstants.addIntValues( GandivaTypes.IntNode.newBuilder().setValue(val).build())); inNode.setIntValues(intConstants.build()); } else if (longValues != null) { GandivaTypes.LongConstants.Builder longConstants = GandivaTypes.LongConstants.newBuilder(); longValues.stream() .forEach( val -> longConstants.addLongValues( GandivaTypes.LongNode.newBuilder().setValue(val).build())); inNode.setLongValues(longConstants.build()); } else if (floatValues != null) { GandivaTypes.FloatConstants.Builder floatConstants = GandivaTypes.FloatConstants.newBuilder(); floatValues.stream() .forEach( val -> floatConstants.addFloatValues( GandivaTypes.FloatNode.newBuilder().setValue(val).build())); inNode.setFloatValues(floatConstants.build()); } else if (doubleValues != null) { GandivaTypes.DoubleConstants.Builder doubleConstants = GandivaTypes.DoubleConstants.newBuilder(); doubleValues.stream() .forEach( val -> doubleConstants.addDoubleValues( GandivaTypes.DoubleNode.newBuilder().setValue(val).build())); inNode.setDoubleValues(doubleConstants.build()); } else if (decimalValues != null) { GandivaTypes.DecimalConstants.Builder decimalConstants = GandivaTypes.DecimalConstants.newBuilder(); decimalValues.stream() .forEach( val -> decimalConstants.addDecimalValues( GandivaTypes.DecimalNode.newBuilder() .setValue(val.toPlainString()) .setPrecision(precision) .setScale(scale) .build())); inNode.setDecimalValues(decimalConstants.build()); } else if (stringValues != null) { GandivaTypes.StringConstants.Builder stringConstants = GandivaTypes.StringConstants.newBuilder(); stringValues.stream() .forEach( val -> stringConstants.addStringValues( GandivaTypes.StringNode.newBuilder() .setValue(ByteString.copyFrom(val.getBytes(charset))) .build())); inNode.setStringValues(stringConstants.build()); } else if (binaryValues != null) { GandivaTypes.BinaryConstants.Builder binaryConstants = GandivaTypes.BinaryConstants.newBuilder(); binaryValues.stream() .forEach( val -> binaryConstants.addBinaryValues( GandivaTypes.BinaryNode.newBuilder() .setValue(ByteString.copyFrom(val)) .build())); inNode.setBinaryValues(binaryConstants.build()); } GandivaTypes.TreeNode.Builder builder = GandivaTypes.TreeNode.newBuilder(); builder.setInNode(inNode.build()); return builder.build(); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/IntNode.java000066400000000000000000000030311515451615200323150ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; /** * Used to represent expression tree nodes representing int constants. Used in the expression (x + * 5) */ class IntNode implements TreeNode { private final Integer value; IntNode(Integer value) { this.value = value; } @Override public GandivaTypes.TreeNode toProtobuf() throws GandivaException { GandivaTypes.IntNode.Builder intBuilder = GandivaTypes.IntNode.newBuilder(); intBuilder.setValue(value.intValue()); GandivaTypes.TreeNode.Builder builder = GandivaTypes.TreeNode.newBuilder(); builder.setIntNode(intBuilder.build()); return builder.build(); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/LongNode.java000066400000000000000000000030361515451615200324670ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; /** * Used to represent expression tree nodes representing long constants. Used in the expression (x + * 5L) */ class LongNode implements TreeNode { private final Long value; LongNode(Long value) { this.value = value; } @Override public GandivaTypes.TreeNode toProtobuf() throws GandivaException { GandivaTypes.LongNode.Builder longBuilder = GandivaTypes.LongNode.newBuilder(); longBuilder.setValue(value.longValue()); GandivaTypes.TreeNode.Builder builder = GandivaTypes.TreeNode.newBuilder(); builder.setLongNode(longBuilder.build()); return builder.build(); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/NullNode.java000066400000000000000000000030421515451615200324770ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; import org.apache.arrow.vector.types.pojo.ArrowType; /** An expression indicating a null value. */ class NullNode implements TreeNode { private final ArrowType type; NullNode(ArrowType type) { this.type = type; } @Override public GandivaTypes.TreeNode toProtobuf() throws GandivaException { GandivaTypes.NullNode.Builder nullNode = GandivaTypes.NullNode.newBuilder(); nullNode.setType(ArrowTypeHelper.arrowTypeToProtobuf(type)); GandivaTypes.TreeNode.Builder builder = GandivaTypes.TreeNode.newBuilder(); builder.setNullNode(nullNode.build()); return builder.build(); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/OrNode.java000066400000000000000000000030311515451615200321430ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import java.util.List; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; /** Represents a logical OR Node. */ class OrNode implements TreeNode { private final List children; OrNode(List children) { this.children = children; } @Override public GandivaTypes.TreeNode toProtobuf() throws GandivaException { GandivaTypes.OrNode.Builder orNode = GandivaTypes.OrNode.newBuilder(); for (TreeNode arg : children) { orNode.addArgs(arg.toProtobuf()); } GandivaTypes.TreeNode.Builder builder = GandivaTypes.TreeNode.newBuilder(); builder.setOrNode(orNode.build()); return builder.build(); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/StringNode.java000066400000000000000000000031671515451615200330430ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import com.google.protobuf.ByteString; import java.nio.charset.Charset; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; /** Used to represent expression tree nodes representing utf8 constants. */ class StringNode implements TreeNode { private static final Charset charset = Charset.forName("UTF-8"); private final String value; public StringNode(String value) { this.value = value; } @Override public GandivaTypes.TreeNode toProtobuf() throws GandivaException { GandivaTypes.StringNode stringNode = GandivaTypes.StringNode.newBuilder() .setValue(ByteString.copyFrom(value.getBytes(charset))) .build(); return GandivaTypes.TreeNode.newBuilder().setStringNode(stringNode).build(); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/TreeBuilder.java000066400000000000000000000163541515451615200331770ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import java.math.BigDecimal; import java.util.ArrayList; import java.util.List; import java.util.Set; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; /** Contains helper functions for constructing expression trees. */ public class TreeBuilder { private TreeBuilder() {} /** Helper functions to create literal constants. */ public static TreeNode makeLiteral(Boolean booleanConstant) { return new BooleanNode(booleanConstant); } public static TreeNode makeLiteral(Float floatConstant) { return new FloatNode(floatConstant); } public static TreeNode makeLiteral(Double doubleConstant) { return new DoubleNode(doubleConstant); } public static TreeNode makeLiteral(Integer integerConstant) { return new IntNode(integerConstant); } public static TreeNode makeLiteral(Long longConstant) { return new LongNode(longConstant); } public static TreeNode makeStringLiteral(String stringConstant) { return new StringNode(stringConstant); } public static TreeNode makeBinaryLiteral(byte[] binaryConstant) { return new BinaryNode(binaryConstant); } public static TreeNode makeDecimalLiteral(String decimalConstant, int precision, int scale) { return new DecimalNode(decimalConstant, precision, scale); } /** create a null literal. */ public static TreeNode makeNull(ArrowType type) { return new NullNode(type); } /** * Invoke this function to create a node representing a field, e.g. a column name. * * @param field represents the input argument - includes the name and type of the field * @return Node representing a field */ public static TreeNode makeField(Field field) { return new FieldNode(field); } /** * Invoke this function to create a node representing a function. * * @param function Name of the function, e.g. add * @param children The arguments to the function * @param retType The type of the return value of the operator * @return Node representing a function */ public static TreeNode makeFunction(String function, List children, ArrowType retType) { return new FunctionNode(function, children, retType); } /** * Invoke this function to create a node representing an if-clause. * * @param condition Node representing the condition * @param thenNode Node representing the if-block * @param elseNode Node representing the else-block * @param retType Return type of the node * @return Node representing an if-clause */ public static TreeNode makeIf( TreeNode condition, TreeNode thenNode, TreeNode elseNode, ArrowType retType) { return new IfNode(condition, thenNode, elseNode, retType); } /** * Invoke this function to create a node representing an and-clause. * * @param nodes Nodes in the 'and' clause. * @return Node representing an and-clause */ public static TreeNode makeAnd(List nodes) { return new AndNode(nodes); } /** * Invoke this function to create a node representing an or-clause. * * @param nodes Nodes in the 'or' clause. * @return Node representing an or-clause */ public static TreeNode makeOr(List nodes) { return new OrNode(nodes); } /** * Invoke this function to create an expression tree. * * @param root is returned by a call to MakeField, MakeFunction, or MakeIf * @param resultField represents the return value of the expression * @return ExpressionTree referring to the root of an expression tree */ public static ExpressionTree makeExpression(TreeNode root, Field resultField) { return new ExpressionTree(root, resultField); } /** * Short cut to create an expression tree involving a single function, e.g. a+b+c. * * @param function Name of the function, e.g. add() * @param inFields In arguments to the function * @param resultField represents the return value of the expression * @return ExpressionTree referring to the root of an expression tree */ public static ExpressionTree makeExpression( String function, List inFields, Field resultField) { List children = new ArrayList(inFields.size()); for (Field field : inFields) { children.add(makeField(field)); } TreeNode root = makeFunction(function, children, resultField.getType()); return makeExpression(root, resultField); } /** * Invoke this function to create a condition. * * @param root is returned by a call to MakeField, MakeFunction, MakeIf, .. * @return condition referring to the root of an expression tree */ public static Condition makeCondition(TreeNode root) { return new Condition(root); } /** * Short cut to create an expression tree involving a single function, e.g. a+b+c. * * @param function Name of the function, e.g. add() * @param inFields In arguments to the function * @return condition referring to the root of an expression tree */ public static Condition makeCondition(String function, List inFields) { List children = new ArrayList<>(inFields.size()); for (Field field : inFields) { children.add(makeField(field)); } TreeNode root = makeFunction(function, children, new ArrowType.Bool()); return makeCondition(root); } public static TreeNode makeInExpressionInt32(TreeNode resultNode, Set intValues) { return InNode.makeIntInExpr(resultNode, intValues); } public static TreeNode makeInExpressionBigInt(TreeNode resultNode, Set longValues) { return InNode.makeLongInExpr(resultNode, longValues); } public static TreeNode makeInExpressionDecimal( TreeNode resultNode, Set decimalValues, Integer precision, Integer scale) { return InNode.makeDecimalInExpr(resultNode, decimalValues, precision, scale); } public static TreeNode makeInExpressionFloat(TreeNode resultNode, Set floatValues) { return InNode.makeFloatInExpr(resultNode, floatValues); } public static TreeNode makeInExpressionDouble(TreeNode resultNode, Set doubleValues) { return InNode.makeDoubleInExpr(resultNode, doubleValues); } public static TreeNode makeInExpressionString(TreeNode resultNode, Set stringValues) { return InNode.makeStringInExpr(resultNode, stringValues); } public static TreeNode makeInExpressionBinary(TreeNode resultNode, Set binaryValues) { return InNode.makeBinaryInExpr(resultNode, binaryValues); } } apache-arrow-java-19.0.0/gandiva/src/main/java/org/apache/arrow/gandiva/expression/TreeNode.java000066400000000000000000000023661515451615200324740ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; /** Defines an internal node in the expression tree. */ public interface TreeNode { /** * Converts a TreeNode into a protobuf. * * @return A treenode protobuf * @throws GandivaException in case the TreeNode cannot be processed */ GandivaTypes.TreeNode toProtobuf() throws GandivaException; } apache-arrow-java-19.0.0/gandiva/src/test/000077500000000000000000000000001515451615200202555ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/test/java/000077500000000000000000000000001515451615200211765ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/test/java/org/000077500000000000000000000000001515451615200217655ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/test/java/org/apache/000077500000000000000000000000001515451615200232065ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200243405ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/test/java/org/apache/arrow/gandiva/000077500000000000000000000000001515451615200257515ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/000077500000000000000000000000001515451615200277535ustar00rootroot00000000000000BaseEvaluatorTest.java000066400000000000000000000300371515451615200341370ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import java.math.BigDecimal; import java.time.Instant; import java.util.ArrayList; import java.util.HashSet; import java.util.List; import java.util.Random; import java.util.Set; import java.util.concurrent.TimeUnit; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.expression.Condition; import org.apache.arrow.gandiva.expression.ExpressionTree; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; class BaseEvaluatorTest { interface BaseEvaluator { void evaluate(ArrowRecordBatch recordBatch, BufferAllocator allocator) throws GandivaException; long getElapsedMillis(); } class ProjectEvaluator implements BaseEvaluator { private Projector projector; private DataAndVectorGenerator generator; private int numExprs; private int maxRowsInBatch; private long elapsedTime = 0; private List outputVectors = new ArrayList<>(); public ProjectEvaluator( Projector projector, DataAndVectorGenerator generator, int numExprs, int maxRowsInBatch) { this.projector = projector; this.generator = generator; this.numExprs = numExprs; this.maxRowsInBatch = maxRowsInBatch; } @Override public void evaluate(ArrowRecordBatch recordBatch, BufferAllocator allocator) throws GandivaException { // set up output vectors // for each expression, generate the output vector for (int i = 0; i < numExprs; i++) { ValueVector valueVector = generator.generateOutputVector(maxRowsInBatch); outputVectors.add(valueVector); } try { long start = System.nanoTime(); projector.evaluate(recordBatch, outputVectors); long finish = System.nanoTime(); elapsedTime += (finish - start); } finally { for (ValueVector valueVector : outputVectors) { valueVector.close(); } } outputVectors.clear(); } @Override public long getElapsedMillis() { return TimeUnit.NANOSECONDS.toMillis(elapsedTime); } } class FilterEvaluator implements BaseEvaluator { private Filter filter; private long elapsedTime = 0; public FilterEvaluator(Filter filter) { this.filter = filter; } @Override public void evaluate(ArrowRecordBatch recordBatch, BufferAllocator allocator) throws GandivaException { ArrowBuf selectionBuffer = allocator.buffer(recordBatch.getLength() * 2); SelectionVectorInt16 selectionVector = new SelectionVectorInt16(selectionBuffer); try { long start = System.nanoTime(); filter.evaluate(recordBatch, selectionVector); long finish = System.nanoTime(); elapsedTime += (finish - start); } finally { selectionBuffer.close(); } } @Override public long getElapsedMillis() { return TimeUnit.NANOSECONDS.toMillis(elapsedTime); } } interface DataAndVectorGenerator { void writeData(ArrowBuf buffer); ValueVector generateOutputVector(int numRowsInBatch); } class Int32DataAndVectorGenerator implements DataAndVectorGenerator { protected final BufferAllocator allocator; protected final Random rand; Int32DataAndVectorGenerator(BufferAllocator allocator) { this.allocator = allocator; this.rand = new Random(); } @Override public void writeData(ArrowBuf buffer) { buffer.writeInt(rand.nextInt()); } @Override public ValueVector generateOutputVector(int numRowsInBatch) { IntVector intVector = new IntVector(BaseEvaluatorTest.EMPTY_SCHEMA_PATH, allocator); intVector.allocateNew(numRowsInBatch); return intVector; } } class BoundedInt32DataAndVectorGenerator extends Int32DataAndVectorGenerator { private final int upperBound; BoundedInt32DataAndVectorGenerator(BufferAllocator allocator, int upperBound) { super(allocator); this.upperBound = upperBound; } @Override public void writeData(ArrowBuf buffer) { buffer.writeInt(rand.nextInt(upperBound)); } } protected static final int THOUSAND = 1000; protected static final int MILLION = THOUSAND * THOUSAND; protected static final String EMPTY_SCHEMA_PATH = ""; protected BufferAllocator allocator; protected ArrowType boolType; protected ArrowType int8; protected ArrowType int32; protected ArrowType int64; protected ArrowType float64; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); boolType = new ArrowType.Bool(); int8 = new ArrowType.Int(8, true); int32 = new ArrowType.Int(32, true); int64 = new ArrowType.Int(64, true); float64 = new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE); } @AfterEach public void tearDown() { allocator.close(); } ArrowBuf buf(int length) { ArrowBuf buffer = allocator.buffer(length); return buffer; } ArrowBuf buf(byte[] bytes) { ArrowBuf buffer = allocator.buffer(bytes.length); buffer.writeBytes(bytes); return buffer; } ArrowBuf arrowBufWithAllValid(int size) { int bufLen = (size + 7) / 8; ArrowBuf buffer = allocator.buffer(bufLen); for (int i = 0; i < bufLen; i++) { buffer.writeByte(255); } return buffer; } ArrowBuf intBuf(int[] ints) { ArrowBuf buffer = allocator.buffer(ints.length * 4); for (int i = 0; i < ints.length; i++) { buffer.writeInt(ints[i]); } return buffer; } DecimalVector decimalVector(String[] values, int precision, int scale) { DecimalVector vector = new DecimalVector("decimal" + Math.random(), allocator, precision, scale); vector.allocateNew(); for (int i = 0; i < values.length; i++) { BigDecimal decimal = new BigDecimal(values[i]).setScale(scale); vector.setSafe(i, decimal); } vector.setValueCount(values.length); return vector; } Set decimalSet(String[] values, Integer scale) { Set decimalSet = new HashSet<>(); for (int i = 0; i < values.length; i++) { decimalSet.add(new BigDecimal(values[i]).setScale(scale)); } return decimalSet; } VarCharVector varcharVector(String[] values) { VarCharVector vector = new VarCharVector("VarCharVector" + Math.random(), allocator); vector.allocateNew(); for (int i = 0; i < values.length; i++) { vector.setSafe(i, values[i].getBytes(), 0, values[i].length()); } vector.setValueCount(values.length); return vector; } ArrowBuf longBuf(long[] longs) { ArrowBuf buffer = allocator.buffer(longs.length * 8); for (int i = 0; i < longs.length; i++) { buffer.writeLong(longs[i]); } return buffer; } ArrowBuf doubleBuf(double[] data) { ArrowBuf buffer = allocator.buffer(data.length * 8); for (int i = 0; i < data.length; i++) { buffer.writeDouble(data[i]); } return buffer; } ArrowBuf stringToMillis(String[] dates) { ArrowBuf buffer = allocator.buffer(dates.length * 8); for (int i = 0; i < dates.length; i++) { Instant instant = Instant.parse(dates[i]); buffer.writeLong(instant.toEpochMilli()); } return buffer; } ArrowBuf stringToDayInterval(String[] values) { ArrowBuf buffer = allocator.buffer(values.length * 8); for (int i = 0; i < values.length; i++) { buffer.writeInt(Integer.parseInt(values[i].split(" ")[0])); // days buffer.writeInt(Integer.parseInt(values[i].split(" ")[1])); // millis } return buffer; } void releaseRecordBatch(ArrowRecordBatch recordBatch) { // There are 2 references to the buffers // One in the recordBatch - release that by calling close() // One in the allocator - release that explicitly List buffers = recordBatch.getBuffers(); recordBatch.close(); for (ArrowBuf buf : buffers) { buf.getReferenceManager().release(); } } void releaseValueVectors(List valueVectors) { for (ValueVector valueVector : valueVectors) { valueVector.close(); } } void generateData(DataAndVectorGenerator generator, int numRecords, ArrowBuf buffer) { for (int i = 0; i < numRecords; i++) { generator.writeData(buffer); } } private void generateDataAndEvaluate( DataAndVectorGenerator generator, BaseEvaluator evaluator, int numFields, int numRows, int maxRowsInBatch, int inputFieldSize) throws GandivaException, Exception { int numRemaining = numRows; List inputData = new ArrayList(); List fieldNodes = new ArrayList(); // set the bitmap while (numRemaining > 0) { int numRowsInBatch = maxRowsInBatch; if (numRowsInBatch > numRemaining) { numRowsInBatch = numRemaining; } // generate data for (int i = 0; i < numFields; i++) { ArrowBuf buf = allocator.buffer(numRowsInBatch * inputFieldSize); ArrowBuf validity = arrowBufWithAllValid(maxRowsInBatch); generateData(generator, numRowsInBatch, buf); fieldNodes.add(new ArrowFieldNode(numRowsInBatch, 0)); inputData.add(validity); inputData.add(buf); } // create record batch ArrowRecordBatch recordBatch = new ArrowRecordBatch(numRowsInBatch, fieldNodes, inputData); evaluator.evaluate(recordBatch, allocator); // fix numRemaining numRemaining -= numRowsInBatch; // release refs releaseRecordBatch(recordBatch); inputData.clear(); fieldNodes.clear(); } } long timedProject( DataAndVectorGenerator generator, Schema schema, List exprs, int numRows, int maxRowsInBatch, int inputFieldSize) throws GandivaException, Exception { Projector projector = Projector.make(schema, exprs); try { ProjectEvaluator evaluator = new ProjectEvaluator(projector, generator, exprs.size(), maxRowsInBatch); generateDataAndEvaluate( generator, evaluator, schema.getFields().size(), numRows, maxRowsInBatch, inputFieldSize); return evaluator.getElapsedMillis(); } finally { projector.close(); } } long timedFilter( DataAndVectorGenerator generator, Schema schema, Condition condition, int numRows, int maxRowsInBatch, int inputFieldSize) throws GandivaException, Exception { Filter filter = Filter.make(schema, condition); try { FilterEvaluator evaluator = new FilterEvaluator(filter); generateDataAndEvaluate( generator, evaluator, schema.getFields().size(), numRows, maxRowsInBatch, inputFieldSize); return evaluator.getElapsedMillis(); } finally { filter.close(); } } } DecimalTypeUtilTest.java000066400000000000000000000067721515451615200344510ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import static org.junit.jupiter.api.Assertions.assertTrue; import org.apache.arrow.vector.types.pojo.ArrowType; import org.junit.jupiter.api.Test; public class DecimalTypeUtilTest { @Test public void testOutputTypesForAdd() { ArrowType.Decimal operand1 = getDecimal(30, 10); ArrowType.Decimal operand2 = getDecimal(30, 10); ArrowType.Decimal resultType = DecimalTypeUtil.getResultTypeForOperation( DecimalTypeUtil.OperationType.ADD, operand1, operand2); assertTrue(getDecimal(31, 10).equals(resultType)); operand1 = getDecimal(30, 6); operand2 = getDecimal(30, 5); resultType = DecimalTypeUtil.getResultTypeForOperation( DecimalTypeUtil.OperationType.ADD, operand1, operand2); assertTrue(getDecimal(32, 6).equals(resultType)); operand1 = getDecimal(30, 10); operand2 = getDecimal(38, 10); resultType = DecimalTypeUtil.getResultTypeForOperation( DecimalTypeUtil.OperationType.ADD, operand1, operand2); assertTrue(getDecimal(38, 9).equals(resultType)); operand1 = getDecimal(38, 10); operand2 = getDecimal(38, 38); resultType = DecimalTypeUtil.getResultTypeForOperation( DecimalTypeUtil.OperationType.ADD, operand1, operand2); assertTrue(getDecimal(38, 9).equals(resultType)); operand1 = getDecimal(38, 10); operand2 = getDecimal(38, 2); resultType = DecimalTypeUtil.getResultTypeForOperation( DecimalTypeUtil.OperationType.ADD, operand1, operand2); assertTrue(getDecimal(38, 6).equals(resultType)); } @Test public void testOutputTypesForMultiply() { ArrowType.Decimal operand1 = getDecimal(30, 10); ArrowType.Decimal operand2 = getDecimal(30, 10); ArrowType.Decimal resultType = DecimalTypeUtil.getResultTypeForOperation( DecimalTypeUtil.OperationType.MULTIPLY, operand1, operand2); assertTrue(getDecimal(38, 6).equals(resultType)); operand1 = getDecimal(38, 10); operand2 = getDecimal(9, 2); resultType = DecimalTypeUtil.getResultTypeForOperation( DecimalTypeUtil.OperationType.MULTIPLY, operand1, operand2); assertTrue(getDecimal(38, 6).equals(resultType)); } @Test public void testOutputTypesForMod() { ArrowType.Decimal operand1 = getDecimal(30, 10); ArrowType.Decimal operand2 = getDecimal(28, 7); ArrowType.Decimal resultType = DecimalTypeUtil.getResultTypeForOperation( DecimalTypeUtil.OperationType.MOD, operand1, operand2); assertTrue(getDecimal(30, 10).equals(resultType)); } private ArrowType.Decimal getDecimal(int precision, int scale) { return new ArrowType.Decimal(precision, scale, 128); } } ExpressionRegistryTest.java000066400000000000000000000050761515451615200352770ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import static org.junit.jupiter.api.Assertions.assertTrue; import com.google.common.collect.Lists; import java.util.Set; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.vector.types.pojo.ArrowType; import org.junit.jupiter.api.Test; public class ExpressionRegistryTest { @Test public void testTypes() throws GandivaException { Set types = ExpressionRegistry.getInstance().getSupportedTypes(); ArrowType.Int uint8 = new ArrowType.Int(8, false); assertTrue(types.contains(uint8)); } @Test public void testFunctions() throws GandivaException { ArrowType.Int uint8 = new ArrowType.Int(8, false); FunctionSignature signature = new FunctionSignature("add", uint8, Lists.newArrayList(uint8, uint8)); Set functions = ExpressionRegistry.getInstance().getSupportedFunctions(); assertTrue(functions.contains(signature)); } @Test public void testFunctionAliases() throws GandivaException { ArrowType.Int int64 = new ArrowType.Int(64, true); FunctionSignature signature = new FunctionSignature("modulo", int64, Lists.newArrayList(int64, int64)); Set functions = ExpressionRegistry.getInstance().getSupportedFunctions(); assertTrue(functions.contains(signature)); } @Test public void testCaseInsensitiveFunctionName() throws GandivaException { ArrowType.Utf8 utf8 = new ArrowType.Utf8(); ArrowType.Int int64 = new ArrowType.Int(64, true); FunctionSignature signature = new FunctionSignature("castvarchar", utf8, Lists.newArrayList(utf8, int64)); Set functions = ExpressionRegistry.getInstance().getSupportedFunctions(); assertTrue(functions.contains(signature)); } } FilterProjectTest.java000066400000000000000000000102471515451615200341570ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import com.google.common.collect.Lists; import java.util.ArrayList; import java.util.List; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.expression.Condition; import org.apache.arrow.gandiva.expression.ExpressionTree; import org.apache.arrow.gandiva.expression.TreeBuilder; import org.apache.arrow.gandiva.ipc.GandivaTypes.SelectionVectorType; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; 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.Test; public class FilterProjectTest extends BaseEvaluatorTest { @Test public void testSimpleSV16() throws GandivaException, Exception { Field a = Field.nullable("a", int32); Field b = Field.nullable("b", int32); Field c = Field.nullable("c", int32); List args = Lists.newArrayList(a, b); Condition condition = TreeBuilder.makeCondition("less_than", args); Schema schema = new Schema(args); Filter filter = Filter.make(schema, condition); ExpressionTree expression = TreeBuilder.makeExpression("add", Lists.newArrayList(a, b), c); Projector projector = Projector.make(schema, Lists.newArrayList(expression), SelectionVectorType.SV_INT16); int numRows = 16; byte[] validity = new byte[] {(byte) 255, 0}; // second half is "undefined" int[] aValues = new int[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}; int[] bValues = new int[] {2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 14, 15}; int[] expected = {3, 7, 11, 15}; verifyTestCaseFor16(filter, projector, numRows, validity, aValues, bValues, expected); } private void verifyTestCaseFor16( Filter filter, Projector projector, int numRows, byte[] validity, int[] aValues, int[] bValues, int[] expected) throws GandivaException { ArrowBuf validitya = buf(validity); ArrowBuf valuesa = intBuf(aValues); ArrowBuf validityb = buf(validity); ArrowBuf valuesb = intBuf(bValues); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 0), new ArrowFieldNode(numRows, 0)), Lists.newArrayList(validitya, valuesa, validityb, valuesb)); ArrowBuf selectionBuffer = buf(numRows * 2); SelectionVectorInt16 selectionVector = new SelectionVectorInt16(selectionBuffer); filter.evaluate(batch, selectionVector); IntVector intVector = new IntVector(EMPTY_SCHEMA_PATH, allocator); intVector.allocateNew(selectionVector.getRecordCount()); List output = new ArrayList(); output.add(intVector); projector.evaluate(batch, selectionVector, output); for (int i = 0; i < selectionVector.getRecordCount(); i++) { assertFalse(intVector.isNull(i)); assertEquals(expected[i], intVector.get(i)); } // free buffers releaseRecordBatch(batch); releaseValueVectors(output); selectionBuffer.close(); filter.close(); projector.close(); } } apache-arrow-java-19.0.0/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/FilterTest.java000066400000000000000000000263121515451615200327070ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import com.google.common.collect.Lists; import com.google.common.collect.Sets; import java.nio.charset.Charset; import java.util.Arrays; import java.util.List; import java.util.stream.IntStream; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.expression.Condition; import org.apache.arrow.gandiva.expression.TreeBuilder; import org.apache.arrow.gandiva.expression.TreeNode; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; 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; public class FilterTest extends BaseEvaluatorTest { private int[] selectionVectorToArray(SelectionVector vector) { int[] actual = new int[vector.getRecordCount()]; for (int i = 0; i < vector.getRecordCount(); ++i) { actual[i] = vector.getIndex(i); } return actual; } private Charset utf8Charset = Charset.forName("UTF-8"); private Charset utf16Charset = Charset.forName("UTF-16"); List varBufs(String[] strings, Charset charset) { ArrowBuf offsetsBuffer = allocator.buffer((strings.length + 1) * 4); ArrowBuf dataBuffer = allocator.buffer(strings.length * 8); int startOffset = 0; for (int i = 0; i < strings.length; i++) { offsetsBuffer.writeInt(startOffset); final byte[] bytes = strings[i].getBytes(charset); dataBuffer = dataBuffer.reallocIfNeeded(dataBuffer.writerIndex() + bytes.length); dataBuffer.setBytes(startOffset, bytes, 0, bytes.length); startOffset += bytes.length; } offsetsBuffer.writeInt(startOffset); // offset for the last element return Arrays.asList(offsetsBuffer, dataBuffer); } List stringBufs(String[] strings) { return varBufs(strings, utf8Charset); } @Test public void testSimpleInString() throws GandivaException, Exception { Field c1 = Field.nullable("c1", new ArrowType.Utf8()); TreeNode l1 = TreeBuilder.makeLiteral(1L); TreeNode l2 = TreeBuilder.makeLiteral(3L); List argsSchema = Lists.newArrayList(c1); List args = Lists.newArrayList(TreeBuilder.makeField(c1), l1, l2); TreeNode substr = TreeBuilder.makeFunction("substr", args, new ArrowType.Utf8()); TreeNode inExpr = TreeBuilder.makeInExpressionString(substr, Sets.newHashSet("one", "two", "thr", "fou")); Condition condition = TreeBuilder.makeCondition(inExpr); Schema schema = new Schema(argsSchema); Filter filter = Filter.make(schema, condition); int numRows = 16; byte[] validity = new byte[] {(byte) 255, 0}; // second half is "undefined" String[] c1Values = new String[] { "one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten", "eleven", "twelve", "thirteen", "fourteen", "fifteen", "sixteen" }; int[] expected = {0, 1, 2, 3}; ArrowBuf c1Validity = buf(validity); ArrowBuf c2Validity = buf(validity); List dataBufsX = stringBufs(c1Values); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode), Lists.newArrayList(c1Validity, dataBufsX.get(0), dataBufsX.get(1), c2Validity)); ArrowBuf selectionBuffer = buf(numRows * 2); SelectionVectorInt16 selectionVector = new SelectionVectorInt16(selectionBuffer); filter.evaluate(batch, selectionVector); int[] actual = selectionVectorToArray(selectionVector); releaseRecordBatch(batch); selectionBuffer.close(); filter.close(); assertArrayEquals(expected, actual); } @Test public void testSimpleInInt() throws GandivaException, Exception { Field c1 = Field.nullable("c1", int32); List argsSchema = Lists.newArrayList(c1); TreeNode inExpr = TreeBuilder.makeInExpressionInt32(TreeBuilder.makeField(c1), Sets.newHashSet(1, 2, 3, 4)); Condition condition = TreeBuilder.makeCondition(inExpr); Schema schema = new Schema(argsSchema); Filter filter = Filter.make(schema, condition); int numRows = 16; byte[] validity = new byte[] {(byte) 255, 0}; // second half is "undefined" int[] aValues = new int[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}; int[] expected = {0, 1, 2, 3}; ArrowBuf validitya = buf(validity); ArrowBuf validityb = buf(validity); ArrowBuf valuesa = intBuf(aValues); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode), Lists.newArrayList(validitya, valuesa, validityb)); ArrowBuf selectionBuffer = buf(numRows * 2); SelectionVectorInt16 selectionVector = new SelectionVectorInt16(selectionBuffer); filter.evaluate(batch, selectionVector); // free buffers int[] actual = selectionVectorToArray(selectionVector); releaseRecordBatch(batch); selectionBuffer.close(); filter.close(); assertArrayEquals(expected, actual); } @Test public void testSimpleSV16() throws GandivaException, Exception { Field a = Field.nullable("a", int32); Field b = Field.nullable("b", int32); List args = Lists.newArrayList(a, b); Condition condition = TreeBuilder.makeCondition("less_than", args); Schema schema = new Schema(args); Filter filter = Filter.make(schema, condition); int numRows = 16; byte[] validity = new byte[] {(byte) 255, 0}; // second half is "undefined" int[] aValues = new int[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}; int[] bValues = new int[] {2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 14, 15}; int[] expected = {0, 2, 4, 6}; verifyTestCase(filter, numRows, validity, aValues, bValues, expected); } @Test public void testSimpleSV16_AllMatched() throws GandivaException, Exception { Field a = Field.nullable("a", int32); Field b = Field.nullable("b", int32); List args = Lists.newArrayList(a, b); Condition condition = TreeBuilder.makeCondition("less_than", args); Schema schema = new Schema(args); Filter filter = Filter.make(schema, condition); int numRows = 32; byte[] validity = new byte[numRows / 8]; IntStream.range(0, numRows / 8).forEach(i -> validity[i] = (byte) 255); int[] aValues = new int[numRows]; IntStream.range(0, numRows).forEach(i -> aValues[i] = i); int[] bValues = new int[numRows]; IntStream.range(0, numRows).forEach(i -> bValues[i] = i + 1); int[] expected = new int[numRows]; IntStream.range(0, numRows).forEach(i -> expected[i] = i); verifyTestCase(filter, numRows, validity, aValues, bValues, expected); } @Test public void testSimpleSV16_GreaterThan64Recs() throws GandivaException, Exception { Field a = Field.nullable("a", int32); Field b = Field.nullable("b", int32); List args = Lists.newArrayList(a, b); Condition condition = TreeBuilder.makeCondition("greater_than", args); Schema schema = new Schema(args); Filter filter = Filter.make(schema, condition); int numRows = 1000; byte[] validity = new byte[numRows / 8]; IntStream.range(0, numRows / 8).forEach(i -> validity[i] = (byte) 255); int[] aValues = new int[numRows]; IntStream.range(0, numRows).forEach(i -> aValues[i] = i); int[] bValues = new int[numRows]; IntStream.range(0, numRows).forEach(i -> bValues[i] = i + 1); aValues[0] = 5; bValues[0] = 0; int[] expected = {0}; verifyTestCase(filter, numRows, validity, aValues, bValues, expected); } @Test public void testSimpleSV32() throws GandivaException, Exception { Field a = Field.nullable("a", int32); Field b = Field.nullable("b", int32); List args = Lists.newArrayList(a, b); Condition condition = TreeBuilder.makeCondition("less_than", args); Schema schema = new Schema(args); Filter filter = Filter.make(schema, condition); int numRows = 16; byte[] validity = new byte[] {(byte) 255, 0}; // second half is "undefined" int[] aValues = new int[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}; int[] bValues = new int[] {2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 14, 15}; int[] expected = {0, 2, 4, 6}; verifyTestCase(filter, numRows, validity, aValues, bValues, expected); } @Test public void testSimpleFilterWithNoOptimisation() throws GandivaException, Exception { Field a = Field.nullable("a", int32); Field b = Field.nullable("b", int32); List args = Lists.newArrayList(a, b); Condition condition = TreeBuilder.makeCondition("less_than", args); Schema schema = new Schema(args); Filter filter = Filter.make(schema, condition, false); int numRows = 16; byte[] validity = new byte[] {(byte) 255, 0}; // second half is "undefined" int[] aValues = new int[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}; int[] bValues = new int[] {2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 14, 15}; int[] expected = {0, 2, 4, 6}; verifyTestCase(filter, numRows, validity, aValues, bValues, expected); } private void verifyTestCase( Filter filter, int numRows, byte[] validity, int[] aValues, int[] bValues, int[] expected) throws GandivaException { ArrowBuf validitya = buf(validity); ArrowBuf valuesa = intBuf(aValues); ArrowBuf validityb = buf(validity); ArrowBuf valuesb = intBuf(bValues); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 0), new ArrowFieldNode(numRows, 0)), Lists.newArrayList(validitya, valuesa, validityb, valuesb)); ArrowBuf selectionBuffer = buf(numRows * 2); SelectionVectorInt16 selectionVector = new SelectionVectorInt16(selectionBuffer); filter.evaluate(batch, selectionVector); // free buffers int[] actual = selectionVectorToArray(selectionVector); releaseRecordBatch(batch); selectionBuffer.close(); filter.close(); assertArrayEquals(expected, actual); } } MicroBenchmarkTest.java000066400000000000000000000125161515451615200342700ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import static org.junit.jupiter.api.Assertions.assertTrue; import com.google.common.collect.Lists; import java.util.List; import org.apache.arrow.gandiva.expression.Condition; import org.apache.arrow.gandiva.expression.ExpressionTree; import org.apache.arrow.gandiva.expression.TreeBuilder; import org.apache.arrow.gandiva.expression.TreeNode; 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; @Disabled public class MicroBenchmarkTest extends BaseEvaluatorTest { private double toleranceRatio = 4.0; @Test public void testAdd3() throws Exception { Field x = Field.nullable("x", int32); Field n2x = Field.nullable("n2x", int32); Field n3x = Field.nullable("n3x", int32); // x + n2x + n3x TreeNode add1 = TreeBuilder.makeFunction( "add", Lists.newArrayList(TreeBuilder.makeField(x), TreeBuilder.makeField(n2x)), int32); TreeNode add = TreeBuilder.makeFunction( "add", Lists.newArrayList(add1, TreeBuilder.makeField(n3x)), int32); ExpressionTree expr = TreeBuilder.makeExpression(add, x); List cols = Lists.newArrayList(x, n2x, n3x); Schema schema = new Schema(cols); long timeTaken = timedProject( new Int32DataAndVectorGenerator(allocator), schema, Lists.newArrayList(expr), 1 * MILLION, 16 * THOUSAND, 4); System.out.println("Time taken for projecting 1m records of add3 is " + timeTaken + "ms"); assertTrue(timeTaken <= 13 * toleranceRatio); } @Test public void testIf() throws Exception { /* * when x < 10 then 0 * when x < 20 then 1 * when x < 30 then 2 * when x < 40 then 3 * when x < 50 then 4 * when x < 60 then 5 * when x < 70 then 6 * when x < 80 then 7 * when x < 90 then 8 * when x < 100 then 9 * when x < 110 then 10 * when x < 120 then 11 * when x < 130 then 12 * when x < 140 then 13 * when x < 150 then 14 * when x < 160 then 15 * when x < 170 then 16 * when x < 180 then 17 * when x < 190 then 18 * when x < 200 then 19 * else 20 */ Field x = Field.nullable("x", int32); TreeNode xNode = TreeBuilder.makeField(x); // if (x < 100) then 9 else 10 int returnValue = 20; TreeNode topNode = TreeBuilder.makeLiteral(returnValue); int compareWith = 200; while (compareWith >= 10) { // cond (x < compareWith) TreeNode condNode = TreeBuilder.makeFunction( "less_than", Lists.newArrayList(xNode, TreeBuilder.makeLiteral(compareWith)), boolType); topNode = TreeBuilder.makeIf( condNode, // cond (x < compareWith) TreeBuilder.makeLiteral(returnValue), // then returnValue topNode, // else topNode int32); compareWith -= 10; returnValue--; } ExpressionTree expr = TreeBuilder.makeExpression(topNode, x); Schema schema = new Schema(Lists.newArrayList(x)); long timeTaken = timedProject( new BoundedInt32DataAndVectorGenerator(allocator, 250), schema, Lists.newArrayList(expr), 1 * MILLION, 16 * THOUSAND, 4); System.out.println("Time taken for projecting 10m records of nestedIf is " + timeTaken + "ms"); assertTrue(timeTaken <= 15 * toleranceRatio); } @Test public void testFilterAdd2() throws Exception { Field x = Field.nullable("x", int32); Field n2x = Field.nullable("n2x", int32); Field n3x = Field.nullable("n3x", int32); // x + n2x < n3x TreeNode add = TreeBuilder.makeFunction( "add", Lists.newArrayList(TreeBuilder.makeField(x), TreeBuilder.makeField(n2x)), int32); TreeNode lessThan = TreeBuilder.makeFunction( "less_than", Lists.newArrayList(add, TreeBuilder.makeField(n3x)), boolType); Condition condition = TreeBuilder.makeCondition(lessThan); List cols = Lists.newArrayList(x, n2x, n3x); Schema schema = new Schema(cols); long timeTaken = timedFilter( new Int32DataAndVectorGenerator(allocator), schema, condition, 1 * MILLION, 16 * THOUSAND, 4); System.out.println("Time taken for filtering 10m records of a+b args = Lists.newArrayList(a, b); ArrowType.Decimal outputType = DecimalTypeUtil.getResultTypeForOperation( DecimalTypeUtil.OperationType.ADD, decimal, decimal); Field retType = Field.nullable("c", outputType); ExpressionTree root = TreeBuilder.makeExpression("add", args, retType); List exprs = Lists.newArrayList(root); Schema schema = new Schema(args); Projector eval = Projector.make(schema, exprs); int numRows = 4; byte[] validity = new byte[] {(byte) 255}; String[] aValues = new String[] {"1.12345678", "2.12345678", "3.12345678", "4.12345678"}; String[] bValues = new String[] {"2.12345678", "3.12345678", "4.12345678", "5.12345678"}; DecimalVector valuesa = decimalVector(aValues, precision, scale); DecimalVector valuesb = decimalVector(bValues, precision, scale); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 0), new ArrowFieldNode(numRows, 0)), Lists.newArrayList( valuesa.getValidityBuffer(), valuesa.getDataBuffer(), valuesb.getValidityBuffer(), valuesb.getDataBuffer())); DecimalVector outVector = new DecimalVector( "decimal_output", allocator, outputType.getPrecision(), outputType.getScale()); outVector.allocateNew(numRows); List output = new ArrayList(); output.add(outVector); eval.evaluate(batch, output); // should have scaled down. BigDecimal[] expOutput = new BigDecimal[] { BigDecimal.valueOf(3.2469136), BigDecimal.valueOf(5.2469136), BigDecimal.valueOf(7.2469136), BigDecimal.valueOf(9.2469136) }; for (int i = 0; i < 4; i++) { assertFalse(outVector.isNull(i)); assertTrue( expOutput[i].compareTo(outVector.getObject(i)) == 0, "index : " + i + " failed compare"); } // free buffers releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void test_add_literal() throws GandivaException { int precision = 2; int scale = 0; ArrowType.Decimal decimal = new ArrowType.Decimal(precision, scale, 128); ArrowType.Decimal literalType = new ArrowType.Decimal(2, 1, 128); Field a = Field.nullable("a", decimal); ArrowType.Decimal outputType = DecimalTypeUtil.getResultTypeForOperation( DecimalTypeUtil.OperationType.ADD, decimal, literalType); Field retType = Field.nullable("c", outputType); TreeNode field = TreeBuilder.makeField(a); TreeNode literal = TreeBuilder.makeDecimalLiteral("6", 2, 1); List args = Lists.newArrayList(field, literal); TreeNode root = TreeBuilder.makeFunction("add", args, outputType); ExpressionTree tree = TreeBuilder.makeExpression(root, retType); List exprs = Lists.newArrayList(tree); Schema schema = new Schema(Lists.newArrayList(a)); Projector eval = Projector.make(schema, exprs); int numRows = 4; String[] aValues = new String[] {"1", "2", "3", "4"}; DecimalVector valuesa = decimalVector(aValues, precision, scale); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 0)), Lists.newArrayList(valuesa.getValidityBuffer(), valuesa.getDataBuffer())); DecimalVector outVector = new DecimalVector( "decimal_output", allocator, outputType.getPrecision(), outputType.getScale()); outVector.allocateNew(numRows); List output = new ArrayList(); output.add(outVector); eval.evaluate(batch, output); BigDecimal[] expOutput = new BigDecimal[] { BigDecimal.valueOf(1.6), BigDecimal.valueOf(2.6), BigDecimal.valueOf(3.6), BigDecimal.valueOf(4.6) }; for (int i = 0; i < 4; i++) { assertFalse(outVector.isNull(i)); assertTrue(expOutput[i].compareTo(outVector.getObject(i)) == 0); } // free buffers releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void test_multiply() throws GandivaException { int precision = 38; int scale = 8; ArrowType.Decimal decimal = new ArrowType.Decimal(precision, scale, 128); Field a = Field.nullable("a", decimal); Field b = Field.nullable("b", decimal); List args = Lists.newArrayList(a, b); ArrowType.Decimal outputType = DecimalTypeUtil.getResultTypeForOperation( DecimalTypeUtil.OperationType.MULTIPLY, decimal, decimal); Field retType = Field.nullable("c", outputType); ExpressionTree root = TreeBuilder.makeExpression("multiply", args, retType); List exprs = Lists.newArrayList(root); Schema schema = new Schema(args); Projector eval = Projector.make(schema, exprs); int numRows = 4; byte[] validity = new byte[] {(byte) 255}; String[] aValues = new String[] {"1.12345678", "2.12345678", "3.12345678", "999999999999.99999999"}; String[] bValues = new String[] {"2.12345678", "3.12345678", "4.12345678", "999999999999.99999999"}; DecimalVector valuesa = decimalVector(aValues, precision, scale); DecimalVector valuesb = decimalVector(bValues, precision, scale); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 0), new ArrowFieldNode(numRows, 0)), Lists.newArrayList( valuesa.getValidityBuffer(), valuesa.getDataBuffer(), valuesb.getValidityBuffer(), valuesb.getDataBuffer())); DecimalVector outVector = new DecimalVector( "decimal_output", allocator, outputType.getPrecision(), outputType.getScale()); outVector.allocateNew(numRows); List output = new ArrayList(); output.add(outVector); eval.evaluate(batch, output); // should have scaled down. BigDecimal[] expOutput = new BigDecimal[] { BigDecimal.valueOf(2.385612), BigDecimal.valueOf(6.632525), BigDecimal.valueOf(12.879439), new BigDecimal("999999999999999999980000.000000") }; for (int i = 0; i < 4; i++) { assertFalse(outVector.isNull(i)); assertTrue( expOutput[i].compareTo(outVector.getObject(i)) == 0, "index : " + i + " failed compare"); } // free buffers releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testCompare() throws GandivaException { Decimal aType = new Decimal(38, 3, 128); Decimal bType = new Decimal(38, 2, 128); Field a = Field.nullable("a", aType); Field b = Field.nullable("b", bType); List args = Lists.newArrayList(a, b); List exprs = new ArrayList<>( Arrays.asList( TreeBuilder.makeExpression("equal", args, Field.nullable("eq", boolType)), TreeBuilder.makeExpression("not_equal", args, Field.nullable("ne", boolType)), TreeBuilder.makeExpression("less_than", args, Field.nullable("lt", boolType)), TreeBuilder.makeExpression( "less_than_or_equal_to", args, Field.nullable("le", boolType)), TreeBuilder.makeExpression("greater_than", args, Field.nullable("gt", boolType)), TreeBuilder.makeExpression( "greater_than_or_equal_to", args, Field.nullable("ge", boolType)))); Schema schema = new Schema(args); Projector eval = Projector.make(schema, exprs); List output = null; ArrowRecordBatch batch = null; try { int numRows = 4; String[] aValues = new String[] {"7.620", "2.380", "3.860", "-18.160"}; String[] bValues = new String[] {"7.62", "3.50", "1.90", "-1.45"}; DecimalVector valuesa = decimalVector(aValues, aType.getPrecision(), aType.getScale()); DecimalVector valuesb = decimalVector(bValues, bType.getPrecision(), bType.getScale()); batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 0), new ArrowFieldNode(numRows, 0)), Lists.newArrayList( valuesa.getValidityBuffer(), valuesa.getDataBuffer(), valuesb.getValidityBuffer(), valuesb.getDataBuffer())); // expected results. boolean[][] expected = { {true, false, false, false}, // eq {false, true, true, true}, // ne {false, true, false, true}, // lt {true, true, false, true}, // le {false, false, true, false}, // gt {true, false, true, false}, // ge }; // Allocate output vectors. output = new ArrayList<>( Arrays.asList( new BitVector("eq", allocator), new BitVector("ne", allocator), new BitVector("lt", allocator), new BitVector("le", allocator), new BitVector("gt", allocator), new BitVector("ge", allocator))); for (ValueVector v : output) { v.allocateNew(); } // evaluate expressions. eval.evaluate(batch, output); // compare the outputs. for (int idx = 0; idx < output.size(); ++idx) { boolean[] expectedArray = expected[idx]; BitVector resultVector = (BitVector) output.get(idx); for (int i = 0; i < numRows; i++) { assertFalse(resultVector.isNull(i)); assertEquals( expectedArray[i], resultVector.getObject(i).booleanValue(), "mismatch in result for expr at idx " + idx + " for row " + i); } } } finally { // free buffers if (batch != null) { releaseRecordBatch(batch); } if (output != null) { releaseValueVectors(output); } eval.close(); } } @Test public void testRound() throws GandivaException { Decimal aType = new Decimal(38, 2, 128); Decimal aWithScaleZero = new Decimal(38, 0, 128); Decimal aWithScaleOne = new Decimal(38, 1, 128); Field a = Field.nullable("a", aType); List args = Lists.newArrayList(a); List exprs = new ArrayList<>( Arrays.asList( TreeBuilder.makeExpression("abs", args, Field.nullable("abs", aType)), TreeBuilder.makeExpression("ceil", args, Field.nullable("ceil", aWithScaleZero)), TreeBuilder.makeExpression("floor", args, Field.nullable("floor", aWithScaleZero)), TreeBuilder.makeExpression("round", args, Field.nullable("round", aWithScaleZero)), TreeBuilder.makeExpression( "truncate", args, Field.nullable("truncate", aWithScaleZero)), TreeBuilder.makeExpression( TreeBuilder.makeFunction( "round", Lists.newArrayList(TreeBuilder.makeField(a), TreeBuilder.makeLiteral(1)), aWithScaleOne), Field.nullable("round_scale_1", aWithScaleOne)), TreeBuilder.makeExpression( TreeBuilder.makeFunction( "truncate", Lists.newArrayList(TreeBuilder.makeField(a), TreeBuilder.makeLiteral(1)), aWithScaleOne), Field.nullable("truncate_scale_1", aWithScaleOne)))); Schema schema = new Schema(args); Projector eval = Projector.make(schema, exprs); List output = null; ArrowRecordBatch batch = null; try { int numRows = 4; String[] aValues = new String[] {"1.23", "1.58", "-1.23", "-1.58"}; DecimalVector valuesa = decimalVector(aValues, aType.getPrecision(), aType.getScale()); batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 0)), Lists.newArrayList(valuesa.getValidityBuffer(), valuesa.getDataBuffer())); // expected results. BigDecimal[][] expected = { { BigDecimal.valueOf(1.23), BigDecimal.valueOf(1.58), BigDecimal.valueOf(1.23), BigDecimal.valueOf(1.58) }, // abs { BigDecimal.valueOf(2), BigDecimal.valueOf(2), BigDecimal.valueOf(-1), BigDecimal.valueOf(-1) }, // ceil { BigDecimal.valueOf(1), BigDecimal.valueOf(1), BigDecimal.valueOf(-2), BigDecimal.valueOf(-2) }, // floor { BigDecimal.valueOf(1), BigDecimal.valueOf(2), BigDecimal.valueOf(-1), BigDecimal.valueOf(-2) }, // round { BigDecimal.valueOf(1), BigDecimal.valueOf(1), BigDecimal.valueOf(-1), BigDecimal.valueOf(-1) }, // truncate { BigDecimal.valueOf(1.2), BigDecimal.valueOf(1.6), BigDecimal.valueOf(-1.2), BigDecimal.valueOf(-1.6) }, // round-to-scale-1 { BigDecimal.valueOf(1.2), BigDecimal.valueOf(1.5), BigDecimal.valueOf(-1.2), BigDecimal.valueOf(-1.5) }, // truncate-to-scale-1 }; // Allocate output vectors. output = new ArrayList<>( Arrays.asList( new DecimalVector("abs", allocator, aType.getPrecision(), aType.getScale()), new DecimalVector("ceil", allocator, aType.getPrecision(), 0), new DecimalVector("floor", allocator, aType.getPrecision(), 0), new DecimalVector("round", allocator, aType.getPrecision(), 0), new DecimalVector("truncate", allocator, aType.getPrecision(), 0), new DecimalVector("round_to_scale_1", allocator, aType.getPrecision(), 1), new DecimalVector("truncate_to_scale_1", allocator, aType.getPrecision(), 1))); for (ValueVector v : output) { v.allocateNew(); } // evaluate expressions. eval.evaluate(batch, output); // compare the outputs. for (int idx = 0; idx < output.size(); ++idx) { BigDecimal[] expectedArray = expected[idx]; DecimalVector resultVector = (DecimalVector) output.get(idx); for (int i = 0; i < numRows; i++) { assertFalse(resultVector.isNull(i)); assertTrue( expectedArray[i].compareTo(resultVector.getObject(i)) == 0, "mismatch in result for " + "field " + resultVector.getField().getName() + " for row " + i + " expected " + expectedArray[i] + ", got " + resultVector.getObject(i)); } } } finally { // free buffers if (batch != null) { releaseRecordBatch(batch); } if (output != null) { releaseValueVectors(output); } eval.close(); } } @Test public void testCastToDecimal() throws GandivaException { Decimal decimalType = new Decimal(38, 2, 128); Decimal decimalWithScaleOne = new Decimal(38, 1, 128); Field dec = Field.nullable("dec", decimalType); Field int64f = Field.nullable("int64", int64); Field doublef = Field.nullable("float64", float64); List exprs = new ArrayList<>( Arrays.asList( TreeBuilder.makeExpression( "castDECIMAL", Lists.newArrayList(int64f), Field.nullable("int64_to_dec", decimalType)), TreeBuilder.makeExpression( "castDECIMAL", Lists.newArrayList(doublef), Field.nullable("float64_to_dec", decimalType)), TreeBuilder.makeExpression( "castDECIMAL", Lists.newArrayList(dec), Field.nullable("dec_to_dec", decimalWithScaleOne)))); Schema schema = new Schema(Lists.newArrayList(int64f, doublef, dec)); Projector eval = Projector.make(schema, exprs); List output = null; ArrowRecordBatch batch = null; try { int numRows = 4; String[] aValues = new String[] {"1.23", "1.58", "-1.23", "-1.58"}; DecimalVector valuesa = decimalVector(aValues, decimalType.getPrecision(), decimalType.getScale()); batch = new ArrowRecordBatch( numRows, Lists.newArrayList( new ArrowFieldNode(numRows, 0), new ArrowFieldNode(numRows, 0), new ArrowFieldNode(numRows, 0)), Lists.newArrayList( arrowBufWithAllValid(4), longBuf(new long[] {123, 158, -123, -158}), arrowBufWithAllValid(4), doubleBuf(new double[] {1.23, 1.58, -1.23, -1.58}), valuesa.getValidityBuffer(), valuesa.getDataBuffer())); // Allocate output vectors. output = new ArrayList<>( Arrays.asList( new DecimalVector( "int64_to_dec", allocator, decimalType.getPrecision(), decimalType.getScale()), new DecimalVector( "float64_to_dec", allocator, decimalType.getPrecision(), decimalType.getScale()), new DecimalVector( "dec_to_dec", allocator, decimalWithScaleOne.getPrecision(), decimalWithScaleOne.getScale()))); for (ValueVector v : output) { v.allocateNew(); } // evaluate expressions. eval.evaluate(batch, output); // compare the outputs. BigDecimal[][] expected = { { BigDecimal.valueOf(123), BigDecimal.valueOf(158), BigDecimal.valueOf(-123), BigDecimal.valueOf(-158) }, { BigDecimal.valueOf(1.23), BigDecimal.valueOf(1.58), BigDecimal.valueOf(-1.23), BigDecimal.valueOf(-1.58) }, { BigDecimal.valueOf(1.2), BigDecimal.valueOf(1.6), BigDecimal.valueOf(-1.2), BigDecimal.valueOf(-1.6) } }; for (int idx = 0; idx < output.size(); ++idx) { BigDecimal[] expectedArray = expected[idx]; DecimalVector resultVector = (DecimalVector) output.get(idx); for (int i = 0; i < numRows; i++) { assertFalse(resultVector.isNull(i)); assertTrue( expectedArray[i].compareTo(resultVector.getObject(i)) == 0, "mismatch in result for " + "field " + resultVector.getField().getName() + " for row " + i + " expected " + expectedArray[i] + ", got " + resultVector.getObject(i)); } } } finally { // free buffers if (batch != null) { releaseRecordBatch(batch); } if (output != null) { releaseValueVectors(output); } eval.close(); } } @Test public void testCastToLong() throws GandivaException { Decimal decimalType = new Decimal(38, 2, 128); Field dec = Field.nullable("dec", decimalType); Schema schema = new Schema(Lists.newArrayList(dec)); Projector eval = Projector.make( schema, Lists.newArrayList( TreeBuilder.makeExpression( "castBIGINT", Lists.newArrayList(dec), Field.nullable("dec_to_int64", int64)))); List output = null; ArrowRecordBatch batch = null; try { int numRows = 5; String[] aValues = new String[] {"1.23", "1.50", "98765.78", "-1.23", "-1.58"}; DecimalVector valuesa = decimalVector(aValues, decimalType.getPrecision(), decimalType.getScale()); batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 0)), Lists.newArrayList(valuesa.getValidityBuffer(), valuesa.getDataBuffer())); // Allocate output vectors. BigIntVector resultVector = new BigIntVector("dec_to_int64", allocator); resultVector.allocateNew(); output = new ArrayList<>(Arrays.asList(resultVector)); // evaluate expressions. eval.evaluate(batch, output); // compare the outputs. long[] expected = {1, 2, 98766, -1, -2}; for (int i = 0; i < numRows; i++) { assertFalse(resultVector.isNull(i)); assertEquals(expected[i], resultVector.get(i)); } } finally { // free buffers if (batch != null) { releaseRecordBatch(batch); } if (output != null) { releaseValueVectors(output); } eval.close(); } } @Test public void testCastToDouble() throws GandivaException { Decimal decimalType = new Decimal(38, 2, 128); Field dec = Field.nullable("dec", decimalType); Schema schema = new Schema(Lists.newArrayList(dec)); Projector eval = Projector.make( schema, Lists.newArrayList( TreeBuilder.makeExpression( "castFLOAT8", Lists.newArrayList(dec), Field.nullable("dec_to_float64", float64)))); List output = null; ArrowRecordBatch batch = null; try { int numRows = 4; String[] aValues = new String[] {"1.23", "1.58", "-1.23", "-1.58"}; DecimalVector valuesa = decimalVector(aValues, decimalType.getPrecision(), decimalType.getScale()); batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 0)), Lists.newArrayList(valuesa.getValidityBuffer(), valuesa.getDataBuffer())); // Allocate output vectors. Float8Vector resultVector = new Float8Vector("dec_to_float64", allocator); resultVector.allocateNew(); output = new ArrayList<>(Arrays.asList(resultVector)); // evaluate expressions. eval.evaluate(batch, output); // compare the outputs. double[] expected = {1.23, 1.58, -1.23, -1.58}; for (int i = 0; i < numRows; i++) { assertFalse(resultVector.isNull(i)); assertEquals(expected[i], resultVector.get(i), 0); } } finally { // free buffers if (batch != null) { releaseRecordBatch(batch); } if (output != null) { releaseValueVectors(output); } eval.close(); } } @Test public void testCastToString() throws GandivaException { Decimal decimalType = new Decimal(38, 2, 128); Field dec = Field.nullable("dec", decimalType); Field str = Field.nullable("str", new ArrowType.Utf8()); TreeNode field = TreeBuilder.makeField(dec); TreeNode literal = TreeBuilder.makeLiteral(5L); List args = Lists.newArrayList(field, literal); TreeNode cast = TreeBuilder.makeFunction("castVARCHAR", args, new ArrowType.Utf8()); TreeNode root = TreeBuilder.makeFunction( "equal", Lists.newArrayList(cast, TreeBuilder.makeField(str)), new ArrowType.Bool()); ExpressionTree tree = TreeBuilder.makeExpression(root, Field.nullable("are_equal", new ArrowType.Bool())); Schema schema = new Schema(Lists.newArrayList(dec, str)); Projector eval = Projector.make(schema, Lists.newArrayList(tree)); List output = null; ArrowRecordBatch batch = null; try { int numRows = 4; String[] aValues = new String[] {"10.51", "100.23", "-1000.23", "-0000.10"}; String[] expected = {"10.51", "100.2", "-1000", "-0.10"}; DecimalVector valuesa = decimalVector(aValues, decimalType.getPrecision(), decimalType.getScale()); VarCharVector result = varcharVector(expected); batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 0)), Lists.newArrayList( valuesa.getValidityBuffer(), valuesa.getDataBuffer(), result.getValidityBuffer(), result.getOffsetBuffer(), result.getDataBuffer())); BitVector resultVector = new BitVector("res", allocator); resultVector.allocateNew(); output = new ArrayList<>(Arrays.asList(resultVector)); // evaluate expressions. eval.evaluate(batch, output); // compare the outputs. for (int i = 0; i < numRows; i++) { assertTrue(resultVector.getObject(i).booleanValue()); } } finally { // free buffers if (batch != null) { releaseRecordBatch(batch); } if (output != null) { releaseValueVectors(output); } eval.close(); } } @Test public void testCastStringToDecimal() throws GandivaException { Decimal decimalType = new Decimal(4, 2, 128); Field dec = Field.nullable("dec", decimalType); Field str = Field.nullable("str", new ArrowType.Utf8()); TreeNode field = TreeBuilder.makeField(str); List args = Lists.newArrayList(field); TreeNode cast = TreeBuilder.makeFunction("castDECIMAL", args, decimalType); ExpressionTree tree = TreeBuilder.makeExpression(cast, Field.nullable("dec_str", decimalType)); Schema schema = new Schema(Lists.newArrayList(str)); Projector eval = Projector.make(schema, Lists.newArrayList(tree)); List output = null; ArrowRecordBatch batch = null; try { int numRows = 4; String[] aValues = new String[] {"10.5134", "-0.1", "10.516", "-1000"}; VarCharVector valuesa = varcharVector(aValues); batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 0)), Lists.newArrayList( valuesa.getValidityBuffer(), valuesa.getOffsetBuffer(), valuesa.getDataBuffer())); DecimalVector resultVector = new DecimalVector("res", allocator, decimalType.getPrecision(), decimalType.getScale()); resultVector.allocateNew(); output = new ArrayList<>(Arrays.asList(resultVector)); BigDecimal[] expected = { BigDecimal.valueOf(10.51), BigDecimal.valueOf(-0.10), BigDecimal.valueOf(10.52), BigDecimal.valueOf(0.00) }; // evaluate expressions. eval.evaluate(batch, output); // compare the outputs. for (int i = 0; i < numRows; i++) { assertTrue( expected[i].compareTo(resultVector.getObject(i)) == 0, "mismatch in result for " + "field " + resultVector.getField().getName() + " for row " + i + " expected " + expected[i] + ", got " + resultVector.getObject(i)); } } finally { // free buffers if (batch != null) { releaseRecordBatch(batch); } if (output != null) { releaseValueVectors(output); } eval.close(); } } @Test public void testInvalidDecimal() throws GandivaException { Decimal decimalType = new Decimal(0, 0, 128); Field int64f = Field.nullable("int64", int64); Schema schema = new Schema(Lists.newArrayList(int64f)); IllegalArgumentException exception = assertThrows( IllegalArgumentException.class, () -> { Projector eval = Projector.make( schema, Lists.newArrayList( TreeBuilder.makeExpression( "castDECIMAL", Lists.newArrayList(int64f), Field.nullable("invalid_dec", decimalType)))); }); assertEquals( "Gandiva only supports decimals of upto 38 precision. Input precision : 0", exception.getMessage()); } @Test public void testInvalidDecimalGt38() throws GandivaException { Decimal decimalType = new Decimal(42, 0, 128); Field int64f = Field.nullable("int64", int64); Schema schema = new Schema(Lists.newArrayList(int64f)); IllegalArgumentException exception = assertThrows( IllegalArgumentException.class, () -> { Projector eval = Projector.make( schema, Lists.newArrayList( TreeBuilder.makeExpression( "castDECIMAL", Lists.newArrayList(int64f), Field.nullable("invalid_dec", decimalType)))); }); assertEquals( "Gandiva only supports decimals of upto 38 precision. Input precision : 42", exception.getMessage()); } } apache-arrow-java-19.0.0/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/ProjectorTest.java000066400000000000000000002651541515451615200334420ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; 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 com.google.common.collect.Lists; import com.google.common.collect.Sets; import java.math.BigDecimal; import java.nio.charset.Charset; import java.time.Instant; import java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.Set; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.IntStream; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.expression.ExpressionTree; import org.apache.arrow.gandiva.expression.TreeBuilder; import org.apache.arrow.gandiva.expression.TreeNode; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.IntervalDayVector; import org.apache.arrow.vector.IntervalYearVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.holders.NullableIntervalDayHolder; import org.apache.arrow.vector.holders.NullableIntervalYearHolder; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.IntervalUnit; 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.Disabled; import org.junit.jupiter.api.Test; public class ProjectorTest extends BaseEvaluatorTest { private Charset utf8Charset = Charset.forName("UTF-8"); private Charset utf16Charset = Charset.forName("UTF-16"); List varBufs(String[] strings, Charset charset) { ArrowBuf offsetsBuffer = allocator.buffer((strings.length + 1) * 4); long dataBufferSize = 0L; for (String string : strings) { dataBufferSize += string.getBytes(charset).length; } ArrowBuf dataBuffer = allocator.buffer(dataBufferSize); int startOffset = 0; for (int i = 0; i < strings.length; i++) { offsetsBuffer.writeInt(startOffset); final byte[] bytes = strings[i].getBytes(charset); dataBuffer = dataBuffer.reallocIfNeeded(dataBuffer.writerIndex() + bytes.length); dataBuffer.setBytes(startOffset, bytes, 0, bytes.length); startOffset += bytes.length; } offsetsBuffer.writeInt(startOffset); // offset for the last element return Arrays.asList(offsetsBuffer, dataBuffer); } List stringBufs(String[] strings) { return varBufs(strings, utf8Charset); } List binaryBufs(String[] strings) { return varBufs(strings, utf16Charset); } private void testMakeProjectorParallel(ConfigurationBuilder.ConfigOptions configOptions) throws InterruptedException { List schemas = Lists.newArrayList(); Field a = Field.nullable("a", int64); Field b = Field.nullable("b", int64); IntStream.range(0, 1000) .forEach( i -> { Field c = Field.nullable("" + i, int64); List cols = Lists.newArrayList(a, b, c); schemas.add(new Schema(cols)); }); TreeNode aNode = TreeBuilder.makeField(a); TreeNode bNode = TreeBuilder.makeField(b); List args = Lists.newArrayList(aNode, bNode); TreeNode cond = TreeBuilder.makeFunction("greater_than", args, boolType); TreeNode ifNode = TreeBuilder.makeIf(cond, aNode, bNode, int64); ExpressionTree expr = TreeBuilder.makeExpression(ifNode, Field.nullable("c", int64)); List exprs = Lists.newArrayList(expr); // build projectors in parallel choosing schema at random // this should hit the same cache entry thus exposing // any threading issues. ExecutorService executors = Executors.newFixedThreadPool(16); IntStream.range(0, 1000) .forEach( i -> { executors.submit( () -> { try { Projector evaluator = configOptions == null ? Projector.make(schemas.get((int) (Math.random() * 100)), exprs) : Projector.make( schemas.get((int) (Math.random() * 100)), exprs, configOptions); evaluator.close(); } catch (GandivaException e) { e.printStackTrace(); } }); }); executors.shutdown(); executors.awaitTermination(100, java.util.concurrent.TimeUnit.SECONDS); } @Test public void testMakeProjectorParallel() throws Exception { testMakeProjectorParallel(null); testMakeProjectorParallel(new ConfigurationBuilder.ConfigOptions().withTargetCPU(false)); testMakeProjectorParallel( new ConfigurationBuilder.ConfigOptions().withTargetCPU(false).withOptimize(false)); } // Will be fixed by https://issues.apache.org/jira/browse/ARROW-4371 @Disabled @Test public void testMakeProjector() throws GandivaException { Field a = Field.nullable("a", int64); Field b = Field.nullable("b", int64); TreeNode aNode = TreeBuilder.makeField(a); TreeNode bNode = TreeBuilder.makeField(b); List args = Lists.newArrayList(aNode, bNode); List cols = Lists.newArrayList(a, b); Schema schema = new Schema(cols); TreeNode cond = TreeBuilder.makeFunction("greater_than", args, boolType); TreeNode ifNode = TreeBuilder.makeIf(cond, aNode, bNode, int64); ExpressionTree expr = TreeBuilder.makeExpression(ifNode, Field.nullable("c", int64)); List exprs = Lists.newArrayList(expr); long startTime = System.currentTimeMillis(); Projector evaluator1 = Projector.make(schema, exprs); System.out.println( "Projector build: iteration 1 took " + (System.currentTimeMillis() - startTime) + " ms"); startTime = System.currentTimeMillis(); Projector evaluator2 = Projector.make(schema, exprs); System.out.println( "Projector build: iteration 2 took " + (System.currentTimeMillis() - startTime) + " ms"); startTime = System.currentTimeMillis(); Projector evaluator3 = Projector.make(schema, exprs); long timeToMakeProjector = (System.currentTimeMillis() - startTime); // should be getting the projector from the cache; // giving 5ms for varying system load. assertTrue(timeToMakeProjector < 5L); evaluator1.close(); evaluator2.close(); evaluator3.close(); } @Test public void testMakeProjectorValidationError() throws InterruptedException { Field a = Field.nullable("a", int64); TreeNode aNode = TreeBuilder.makeField(a); List args = Lists.newArrayList(aNode); List cols = Lists.newArrayList(a); Schema schema = new Schema(cols); TreeNode cond = TreeBuilder.makeFunction("non_existent_fn", args, boolType); ExpressionTree expr = TreeBuilder.makeExpression(cond, Field.nullable("c", int64)); List exprs = Lists.newArrayList(expr); boolean exceptionThrown = false; try { Projector evaluator1 = Projector.make(schema, exprs); } catch (GandivaException e) { exceptionThrown = true; } assertTrue(exceptionThrown); // allow GC to collect any temp resources. Thread.sleep(1000); // try again to ensure no temporary resources. exceptionThrown = false; try { Projector evaluator1 = Projector.make(schema, exprs); } catch (GandivaException e) { exceptionThrown = true; } assertTrue(exceptionThrown); } @Test public void testEvaluate() throws GandivaException, Exception { Field a = Field.nullable("a", int32); Field b = Field.nullable("b", int32); List args = Lists.newArrayList(a, b); Field retType = Field.nullable("c", int32); ExpressionTree root = TreeBuilder.makeExpression("add", args, retType); List exprs = Lists.newArrayList(root); Schema schema = new Schema(args); Projector eval = Projector.make(schema, exprs); int numRows = 16; byte[] validity = new byte[] {(byte) 255, 0}; // second half is "undefined" int[] aValues = new int[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}; int[] bValues = new int[] {16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1}; ArrowBuf validitya = buf(validity); ArrowBuf valuesa = intBuf(aValues); ArrowBuf validityb = buf(validity); ArrowBuf valuesb = intBuf(bValues); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 8), new ArrowFieldNode(numRows, 8)), Lists.newArrayList(validitya, valuesa, validityb, valuesb)); IntVector intVector = new IntVector(EMPTY_SCHEMA_PATH, allocator); intVector.allocateNew(numRows); List output = new ArrayList(); output.add(intVector); eval.evaluate(batch, output); for (int i = 0; i < 8; i++) { assertFalse(intVector.isNull(i)); assertEquals(17, intVector.get(i)); } for (int i = 8; i < 16; i++) { assertTrue(intVector.isNull(i)); } // free buffers releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testEvaluateDivZero() throws GandivaException, Exception { Field a = Field.nullable("a", int32); Field b = Field.nullable("b", int32); List args = Lists.newArrayList(a, b); Field retType = Field.nullable("c", int32); ExpressionTree root = TreeBuilder.makeExpression("divide", args, retType); List exprs = Lists.newArrayList(root); Schema schema = new Schema(args); Projector eval = Projector.make(schema, exprs); int numRows = 2; byte[] validity = new byte[] {(byte) 255}; // second half is "undefined" int[] aValues = new int[] {2, 2}; int[] bValues = new int[] {1, 0}; ArrowBuf validitya = buf(validity); ArrowBuf valuesa = intBuf(aValues); ArrowBuf validityb = buf(validity); ArrowBuf valuesb = intBuf(bValues); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 0), new ArrowFieldNode(numRows, 0)), Lists.newArrayList(validitya, valuesa, validityb, valuesb)); IntVector intVector = new IntVector(EMPTY_SCHEMA_PATH, allocator); intVector.allocateNew(numRows); List output = new ArrayList(); output.add(intVector); boolean exceptionThrown = false; try { eval.evaluate(batch, output); } catch (GandivaException e) { assertTrue(e.getMessage().contains("divide by zero")); exceptionThrown = true; } assertTrue(exceptionThrown); // free buffers releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testDivZeroParallel() throws GandivaException, InterruptedException { Field a = Field.nullable("a", int32); Field b = Field.nullable("b", int32); Field c = Field.nullable("c", int32); List cols = Lists.newArrayList(a, b); Schema s = new Schema(cols); List args = Lists.newArrayList(a, b); ExpressionTree expr = TreeBuilder.makeExpression("divide", args, c); List exprs = Lists.newArrayList(expr); ExecutorService executors = Executors.newFixedThreadPool(16); AtomicInteger errorCount = new AtomicInteger(0); AtomicInteger errorCountExp = new AtomicInteger(0); // pre-build the projector so that same projector is used for all executions. Projector test = Projector.make(s, exprs); IntStream.range(0, 1000) .forEach( i -> { executors.submit( () -> { try { Projector evaluator = Projector.make(s, exprs); int numRows = 2; byte[] validity = new byte[] {(byte) 255}; int[] aValues = new int[] {2, 2}; int[] bValues; if (i % 2 == 0) { errorCountExp.incrementAndGet(); bValues = new int[] {1, 0}; } else { bValues = new int[] {1, 1}; } ArrowBuf validitya = buf(validity); ArrowBuf valuesa = intBuf(aValues); ArrowBuf validityb = buf(validity); ArrowBuf valuesb = intBuf(bValues); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList( new ArrowFieldNode(numRows, 0), new ArrowFieldNode(numRows, 0)), Lists.newArrayList(validitya, valuesa, validityb, valuesb)); IntVector intVector = new IntVector(EMPTY_SCHEMA_PATH, allocator); intVector.allocateNew(numRows); List output = new ArrayList(); output.add(intVector); try { evaluator.evaluate(batch, output); } catch (GandivaException e) { errorCount.incrementAndGet(); } // free buffers releaseRecordBatch(batch); releaseValueVectors(output); evaluator.close(); } catch (GandivaException ignore) { } }); }); executors.shutdown(); executors.awaitTermination(100, java.util.concurrent.TimeUnit.SECONDS); test.close(); assertEquals(errorCountExp.intValue(), errorCount.intValue()); } @Test public void testAdd3() throws GandivaException, Exception { Field x = Field.nullable("x", int32); Field n2x = Field.nullable("n2x", int32); Field n3x = Field.nullable("n3x", int32); List args = new ArrayList(); // x + n2x + n3x TreeNode add1 = TreeBuilder.makeFunction( "add", Lists.newArrayList(TreeBuilder.makeField(x), TreeBuilder.makeField(n2x)), int32); TreeNode add = TreeBuilder.makeFunction( "add", Lists.newArrayList(add1, TreeBuilder.makeField(n3x)), int32); ExpressionTree expr = TreeBuilder.makeExpression(add, x); List cols = Lists.newArrayList(x, n2x, n3x); Schema schema = new Schema(cols); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 16; byte[] validity = new byte[] {(byte) 255, 0}; // second half is "undefined" int[] xValues = new int[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}; int[] n2xValues = new int[] {16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1}; int[] n3xValues = new int[] {1, 2, 3, 4, 4, 3, 2, 1, 5, 6, 7, 8, 8, 7, 6, 5}; int[] expected = new int[] {18, 19, 20, 21, 21, 20, 19, 18, 18, 19, 20, 21, 21, 20, 19, 18}; ArrowBuf xValidity = buf(validity); ArrowBuf xData = intBuf(xValues); ArrowBuf n2xValidity = buf(validity); ArrowBuf n2xData = intBuf(n2xValues); ArrowBuf n3xValidity = buf(validity); ArrowBuf n3xData = intBuf(n3xValues); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 8); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode, fieldNode, fieldNode), Lists.newArrayList(xValidity, xData, n2xValidity, n2xData, n3xValidity, n3xData)); IntVector intVector = new IntVector(EMPTY_SCHEMA_PATH, allocator); intVector.allocateNew(numRows); List output = new ArrayList(); output.add(intVector); eval.evaluate(batch, output); for (int i = 0; i < 8; i++) { assertFalse(intVector.isNull(i)); assertEquals(expected[i], intVector.get(i)); } for (int i = 8; i < 16; i++) { assertTrue(intVector.isNull(i)); } releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testStringFields() throws GandivaException { /* * when x < "hello" then octet_length(x) + a * else octet_length(x) + b */ Field x = Field.nullable("x", new ArrowType.Utf8()); Field a = Field.nullable("a", new ArrowType.Int(32, true)); Field b = Field.nullable("b", new ArrowType.Int(32, true)); ArrowType retType = new ArrowType.Int(32, true); TreeNode cond = TreeBuilder.makeFunction( "less_than", Lists.newArrayList(TreeBuilder.makeField(x), TreeBuilder.makeStringLiteral("hello")), boolType); TreeNode octetLenFuncNode = TreeBuilder.makeFunction( "octet_length", Lists.newArrayList(TreeBuilder.makeField(x)), retType); TreeNode octetLenPlusANode = TreeBuilder.makeFunction( "add", Lists.newArrayList(TreeBuilder.makeField(a), octetLenFuncNode), retType); TreeNode octetLenPlusBNode = TreeBuilder.makeFunction( "add", Lists.newArrayList(TreeBuilder.makeField(b), octetLenFuncNode), retType); TreeNode ifHello = TreeBuilder.makeIf(cond, octetLenPlusANode, octetLenPlusBNode, retType); ExpressionTree expr = TreeBuilder.makeExpression(ifHello, Field.nullable("res", retType)); Schema schema = new Schema(Lists.newArrayList(a, x, b)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 5; byte[] validity = new byte[] {(byte) 255, 0}; // "A função" means "The function" in portugese String[] valuesX = new String[] {"hell", "abc", "hellox", "ijk", "A função"}; int[] valuesA = new int[] {10, 20, 30, 40, 50}; int[] valuesB = new int[] {110, 120, 130, 140, 150}; int[] expected = new int[] {14, 23, 136, 143, 60}; ArrowBuf validityX = buf(validity); List dataBufsX = stringBufs(valuesX); ArrowBuf validityA = buf(validity); ArrowBuf dataA = intBuf(valuesA); ArrowBuf validityB = buf(validity); ArrowBuf dataB = intBuf(valuesB); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 0), new ArrowFieldNode(numRows, 0)), Lists.newArrayList( validityA, dataA, validityX, dataBufsX.get(0), dataBufsX.get(1), validityB, dataB)); IntVector intVector = new IntVector(EMPTY_SCHEMA_PATH, allocator); intVector.allocateNew(numRows); List output = new ArrayList(); output.add(intVector); eval.evaluate(batch, output); for (int i = 0; i < numRows; i++) { assertFalse(intVector.isNull(i)); assertEquals(expected[i], intVector.get(i)); } releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testStringOutput() throws GandivaException { /* * if (x >= 0) "hi" else "bye" */ Field x = Field.nullable("x", new ArrowType.Int(32, true)); ArrowType retType = new ArrowType.Utf8(); TreeNode ifHiBye = TreeBuilder.makeIf( TreeBuilder.makeFunction( "greater_than_or_equal_to", Lists.newArrayList(TreeBuilder.makeField(x), TreeBuilder.makeLiteral(0)), boolType), TreeBuilder.makeStringLiteral("hi"), TreeBuilder.makeStringLiteral("bye"), retType); ExpressionTree expr = TreeBuilder.makeExpression(ifHiBye, Field.nullable("res", retType)); Schema schema = new Schema(Lists.newArrayList(x)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); // fill up input record batch int numRows = 4; byte[] validity = new byte[] {(byte) 255, 0}; int[] xValues = new int[] {10, -10, 20, -20}; String[] expected = new String[] {"hi", "bye", "hi", "bye"}; ArrowBuf validityX = buf(validity); ArrowBuf dataX = intBuf(xValues); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 0)), Lists.newArrayList(validityX, dataX)); // allocate data for output vector. VarCharVector outVector = new VarCharVector(EMPTY_SCHEMA_PATH, allocator); outVector.allocateNew(64, numRows); // evaluate expression List output = new ArrayList<>(); output.add(outVector); eval.evaluate(batch, output); // match expected output. for (int i = 0; i < numRows; i++) { assertFalse(outVector.isNull(i)); assertEquals(expected[i], new String(outVector.get(i))); } // test with insufficient data buffer. try { outVector.allocateNew(4, numRows); eval.evaluate(batch, output); } finally { releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } } @Test public void testRegex() throws GandivaException { /* * like "%map%" */ Field x = Field.nullable("x", new ArrowType.Utf8()); TreeNode cond = TreeBuilder.makeFunction( "like", Lists.newArrayList(TreeBuilder.makeField(x), TreeBuilder.makeStringLiteral("%map%")), boolType); ExpressionTree expr = TreeBuilder.makeExpression(cond, Field.nullable("res", boolType)); Schema schema = new Schema(Lists.newArrayList(x)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 5; byte[] validity = new byte[] {(byte) 255, 0}; String[] valuesX = new String[] {"mapD", "maps", "google maps", "map", "MapR"}; boolean[] expected = new boolean[] {true, true, true, true, false}; ArrowBuf validityX = buf(validity); List dataBufsX = stringBufs(valuesX); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 0)), Lists.newArrayList(validityX, dataBufsX.get(0), dataBufsX.get(1))); BitVector bitVector = new BitVector(EMPTY_SCHEMA_PATH, allocator); bitVector.allocateNew(numRows); List output = new ArrayList(); output.add(bitVector); eval.evaluate(batch, output); for (int i = 0; i < numRows; i++) { assertFalse(bitVector.isNull(i)); assertEquals(expected[i], bitVector.getObject(i).booleanValue()); } releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testRegexpReplace() throws GandivaException { Field x = Field.nullable("x", new ArrowType.Utf8()); Field replaceString = Field.nullable("replaceString", new ArrowType.Utf8()); Field retType = Field.nullable("c", new ArrowType.Utf8()); TreeNode cond = TreeBuilder.makeFunction( "regexp_replace", Lists.newArrayList( TreeBuilder.makeField(x), TreeBuilder.makeStringLiteral("ana"), TreeBuilder.makeField(replaceString)), new ArrowType.Utf8()); ExpressionTree expr = TreeBuilder.makeExpression(cond, retType); Schema schema = new Schema(Lists.newArrayList(x, replaceString)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 5; byte[] validity = new byte[] {(byte) 15, 0}; String[] valuesX = new String[] {"banana", "bananaana", "bananana", "anaana", "anaana"}; String[] valuesReplace = new String[] {"ue", "", "", "c", ""}; String[] expected = new String[] {"buena", "bna", "bn", "cc", null}; ArrowBuf validityX = buf(validity); ArrowBuf validityReplace = buf(validity); List dataBufsX = stringBufs(valuesX); List dataBufsReplace = stringBufs(valuesReplace); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode, fieldNode), Lists.newArrayList( validityX, dataBufsX.get(0), dataBufsX.get(1), validityReplace, dataBufsReplace.get(0), dataBufsReplace.get(1))); // allocate data for output vector. VarCharVector outVector = new VarCharVector(EMPTY_SCHEMA_PATH, allocator); outVector.allocateNew(numRows * 15, numRows); // evaluate expression List output = new ArrayList<>(); output.add(outVector); eval.evaluate(batch, output); eval.close(); // match expected output. for (int i = 0; i < numRows - 1; i++) { assertFalse(outVector.isNull(i), "Expect none value equals null"); assertEquals(expected[i], new String(outVector.get(i))); } assertTrue(outVector.isNull(numRows - 1), "Last value must be null"); releaseRecordBatch(batch); releaseValueVectors(output); } @Test public void testCastIntervalDay() throws GandivaException { Field x = Field.nullable("x", new ArrowType.Utf8()); Field retType = Field.nullable("c", new ArrowType.Interval(IntervalUnit.DAY_TIME)); TreeNode cond = TreeBuilder.makeFunction( "castintervalday", Lists.newArrayList(TreeBuilder.makeField(x)), new ArrowType.Interval(IntervalUnit.DAY_TIME)); ExpressionTree expr = TreeBuilder.makeExpression(cond, retType); Schema schema = new Schema(Lists.newArrayList(x)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 4; byte[] validity = new byte[] {(byte) 7, 0}; String[] valuesX = new String[] {"1742461111", "P1Y1M1DT1H1M1S", "PT48H1M1S", "test"}; int[][] expected = new int[][] { // day and millis {20, 14461111}, {1, 3661000}, {2, 61000}, null }; ArrowBuf validityX = buf(validity); List dataBufsX = stringBufs(valuesX); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode, fieldNode), Lists.newArrayList(validityX, dataBufsX.get(0), dataBufsX.get(1))); // allocate data for output vector. IntervalDayVector outVector = new IntervalDayVector(EMPTY_SCHEMA_PATH, allocator); outVector.allocateNew(); // evaluate expression List output = new ArrayList<>(); output.add(outVector); eval.evaluate(batch, output); eval.close(); // match expected output. NullableIntervalDayHolder holder = new NullableIntervalDayHolder(); for (int i = 0; i < numRows - 1; i++) { assertFalse(outVector.isNull(i), "Expect none value equals null"); outVector.get(i, holder); assertEquals(expected[i][0], holder.days); assertEquals(expected[i][1], holder.milliseconds); } assertTrue(outVector.isNull(numRows - 1), "Last value must be null"); releaseRecordBatch(batch); releaseValueVectors(output); } @Test public void testCastIntervalYear() throws GandivaException { Field x = Field.nullable("x", new ArrowType.Utf8()); Field retType = Field.nullable("c", new ArrowType.Interval(IntervalUnit.YEAR_MONTH)); TreeNode cond = TreeBuilder.makeFunction( "castintervalyear", Lists.newArrayList(TreeBuilder.makeField(x)), new ArrowType.Interval(IntervalUnit.YEAR_MONTH)); ExpressionTree expr = TreeBuilder.makeExpression(cond, retType); Schema schema = new Schema(Lists.newArrayList(x)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 4; byte[] validity = new byte[] {(byte) 7, 0}; String[] valuesX = new String[] {"65851111", "P1Y1M1DT1H1M1S", "P1Y", "test"}; int[][] expected = new int[][] { // year and month {0, 65851111}, {1, 1}, {1, 0}, null }; ArrowBuf validityX = buf(validity); List dataBufsX = stringBufs(valuesX); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode, fieldNode), Lists.newArrayList(validityX, dataBufsX.get(0), dataBufsX.get(1))); // allocate data for output vector. IntervalYearVector outVector = new IntervalYearVector(EMPTY_SCHEMA_PATH, allocator); outVector.allocateNew(); // evaluate expression List output = new ArrayList<>(); output.add(outVector); eval.evaluate(batch, output); eval.close(); // match expected output. NullableIntervalYearHolder holder = new NullableIntervalYearHolder(); for (int i = 0; i < numRows - 1; i++) { assertFalse(outVector.isNull(i), "Expect none value equals null"); outVector.get(i, holder); int numberMonths = expected[i][0] * 12 + // number of years expected[i][1]; // number of months assertEquals(numberMonths, holder.value); } assertTrue(outVector.isNull(numRows - 1), "Last value must be null"); releaseRecordBatch(batch); releaseValueVectors(output); } @Test public void testRand() throws GandivaException { TreeNode randWithSeed = TreeBuilder.makeFunction("rand", Lists.newArrayList(TreeBuilder.makeLiteral(12)), float64); TreeNode rand = TreeBuilder.makeFunction("rand", Lists.newArrayList(), float64); ExpressionTree exprWithSeed = TreeBuilder.makeExpression(randWithSeed, Field.nullable("res", float64)); ExpressionTree expr = TreeBuilder.makeExpression(rand, Field.nullable("res2", float64)); Field x = Field.nullable("x", new ArrowType.Utf8()); Schema schema = new Schema(Lists.newArrayList(x)); Projector evalWithSeed = Projector.make(schema, Lists.newArrayList(exprWithSeed)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 5; byte[] validity = new byte[] {(byte) 255, 0}; String[] valuesX = new String[] {"mapD", "maps", "google maps", "map", "MapR"}; double[] expected = new double[] { 0.1597116001879662D, 0.7347813877263527D, 0.6069965050584282D, 0.7240285696335824D, 0.09975540272957834D }; ArrowBuf validityX = buf(validity); List dataBufsX = stringBufs(valuesX); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 0)), Lists.newArrayList(validityX, dataBufsX.get(0), dataBufsX.get(1))); Float8Vector float8Vector = new Float8Vector(EMPTY_SCHEMA_PATH, allocator); float8Vector.allocateNew(numRows); List output = new ArrayList(); output.add(float8Vector); evalWithSeed.evaluate(batch, output); for (int i = 0; i < numRows; i++) { assertFalse(float8Vector.isNull(i)); assertEquals(expected[i], float8Vector.getObject(i), 0.000000001); } eval.evaluate(batch, output); // without seed assertNotEquals(float8Vector.getObject(0), float8Vector.getObject(1), 0.000000001); releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); evalWithSeed.close(); } @Test public void testBinaryFields() throws GandivaException { Field a = Field.nullable("a", new ArrowType.Binary()); Field b = Field.nullable("b", new ArrowType.Binary()); List args = Lists.newArrayList(a, b); ArrowType retType = new ArrowType.Bool(); ExpressionTree expr = TreeBuilder.makeExpression("equal", args, Field.nullable("res", retType)); Schema schema = new Schema(Lists.newArrayList(args)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 5; byte[] validity = new byte[] {(byte) 255, 0}; String[] valuesA = new String[] {"a", "aa", "aaa", "aaaa", "A função"}; String[] valuesB = new String[] {"a", "bb", "aaa", "bbbbb", "A função"}; boolean[] expected = new boolean[] {true, false, true, false, true}; ArrowBuf validitya = buf(validity); ArrowBuf validityb = buf(validity); List inBufsA = binaryBufs(valuesA); List inBufsB = binaryBufs(valuesB); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 8), new ArrowFieldNode(numRows, 8)), Lists.newArrayList( validitya, inBufsA.get(0), inBufsA.get(1), validityb, inBufsB.get(0), inBufsB.get(1))); BitVector bitVector = new BitVector(EMPTY_SCHEMA_PATH, allocator); bitVector.allocateNew(numRows); List output = new ArrayList(); output.add(bitVector); eval.evaluate(batch, output); for (int i = 0; i < numRows; i++) { assertFalse(bitVector.isNull(i)); assertEquals(expected[i], bitVector.getObject(i).booleanValue()); } releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } private TreeNode makeLongLessThanCond(TreeNode arg, long value) { return TreeBuilder.makeFunction( "less_than", Lists.newArrayList(arg, TreeBuilder.makeLiteral(value)), boolType); } private TreeNode makeLongGreaterThanCond(TreeNode arg, long value) { return TreeBuilder.makeFunction( "greater_than", Lists.newArrayList(arg, TreeBuilder.makeLiteral(value)), boolType); } private TreeNode ifLongLessThanElse( TreeNode arg, long value, long thenValue, TreeNode elseNode, ArrowType type) { return TreeBuilder.makeIf( makeLongLessThanCond(arg, value), TreeBuilder.makeLiteral(thenValue), elseNode, type); } @Test public void testIf() throws GandivaException, Exception { /* * when x < 10 then 0 * when x < 20 then 1 * when x < 30 then 2 * when x < 40 then 3 * when x < 50 then 4 * when x < 60 then 5 * when x < 70 then 6 * when x < 80 then 7 * when x < 90 then 8 * when x < 100 then 9 * else 10 */ Field x = Field.nullable("x", int64); TreeNode xNode = TreeBuilder.makeField(x); // if (x < 100) then 9 else 10 TreeNode ifLess100 = ifLongLessThanElse(xNode, 100L, 9L, TreeBuilder.makeLiteral(10L), int64); // if (x < 90) then 8 else ifLess100 TreeNode ifLess90 = ifLongLessThanElse(xNode, 90L, 8L, ifLess100, int64); // if (x < 80) then 7 else ifLess90 TreeNode ifLess80 = ifLongLessThanElse(xNode, 80L, 7L, ifLess90, int64); // if (x < 70) then 6 else ifLess80 TreeNode ifLess70 = ifLongLessThanElse(xNode, 70L, 6L, ifLess80, int64); // if (x < 60) then 5 else ifLess70 TreeNode ifLess60 = ifLongLessThanElse(xNode, 60L, 5L, ifLess70, int64); // if (x < 50) then 4 else ifLess60 TreeNode ifLess50 = ifLongLessThanElse(xNode, 50L, 4L, ifLess60, int64); // if (x < 40) then 3 else ifLess50 TreeNode ifLess40 = ifLongLessThanElse(xNode, 40L, 3L, ifLess50, int64); // if (x < 30) then 2 else ifLess40 TreeNode ifLess30 = ifLongLessThanElse(xNode, 30L, 2L, ifLess40, int64); // if (x < 20) then 1 else ifLess30 TreeNode ifLess20 = ifLongLessThanElse(xNode, 20L, 1L, ifLess30, int64); // if (x < 10) then 0 else ifLess20 TreeNode ifLess10 = ifLongLessThanElse(xNode, 10L, 0L, ifLess20, int64); ExpressionTree expr = TreeBuilder.makeExpression(ifLess10, x); Schema schema = new Schema(Lists.newArrayList(x)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 16; byte[] validity = new byte[] {(byte) 255, (byte) 255}; long[] xValues = new long[] {9, 15, 21, 32, 43, 54, 65, 76, 87, 98, 109, 200, -10, 60, 77, 80}; long[] expected = new long[] {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 10, 0, 6, 7, 8}; ArrowBuf bufValidity = buf(validity); ArrowBuf xData = longBuf(xValues); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode), Lists.newArrayList(bufValidity, xData)); BigIntVector bigIntVector = new BigIntVector(EMPTY_SCHEMA_PATH, allocator); bigIntVector.allocateNew(numRows); List output = new ArrayList(); output.add(bigIntVector); eval.evaluate(batch, output); for (int i = 0; i < numRows; i++) { assertFalse(bigIntVector.isNull(i)); assertEquals(expected[i], bigIntVector.get(i)); } releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testAnd() throws GandivaException, Exception { /* * x > 10 AND x < 20 */ ArrowType int64 = new ArrowType.Int(64, true); Field x = Field.nullable("x", int64); TreeNode xNode = TreeBuilder.makeField(x); TreeNode gt10 = makeLongGreaterThanCond(xNode, 10); TreeNode lt20 = makeLongLessThanCond(xNode, 20); TreeNode and = TreeBuilder.makeAnd(Lists.newArrayList(gt10, lt20)); Field res = Field.nullable("res", boolType); ExpressionTree expr = TreeBuilder.makeExpression(and, res); Schema schema = new Schema(Lists.newArrayList(x)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 4; byte[] validity = new byte[] {(byte) 255}; long[] xValues = new long[] {9, 15, 17, 25}; boolean[] expected = new boolean[] {false, true, true, false}; ArrowBuf bufValidity = buf(validity); ArrowBuf xData = longBuf(xValues); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode), Lists.newArrayList(bufValidity, xData)); BitVector bitVector = new BitVector(EMPTY_SCHEMA_PATH, allocator); bitVector.allocateNew(numRows); List output = new ArrayList(); output.add(bitVector); eval.evaluate(batch, output); for (int i = 0; i < numRows; i++) { assertFalse(bitVector.isNull(i)); assertEquals(expected[i], bitVector.getObject(i).booleanValue()); } releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testOr() throws GandivaException, Exception { /* * x > 10 OR x < 5 */ ArrowType int64 = new ArrowType.Int(64, true); Field x = Field.nullable("x", int64); TreeNode xNode = TreeBuilder.makeField(x); TreeNode gt10 = makeLongGreaterThanCond(xNode, 10); TreeNode lt5 = makeLongLessThanCond(xNode, 5); TreeNode or = TreeBuilder.makeOr(Lists.newArrayList(gt10, lt5)); Field res = Field.nullable("res", boolType); ExpressionTree expr = TreeBuilder.makeExpression(or, res); Schema schema = new Schema(Lists.newArrayList(x)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 4; byte[] validity = new byte[] {(byte) 255}; long[] xValues = new long[] {4, 9, 15, 17}; boolean[] expected = new boolean[] {true, false, true, true}; ArrowBuf bufValidity = buf(validity); ArrowBuf xData = longBuf(xValues); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode), Lists.newArrayList(bufValidity, xData)); BitVector bitVector = new BitVector(EMPTY_SCHEMA_PATH, allocator); bitVector.allocateNew(numRows); List output = new ArrayList(); output.add(bitVector); eval.evaluate(batch, output); for (int i = 0; i < numRows; i++) { assertFalse(bitVector.isNull(i)); assertEquals(expected[i], bitVector.getObject(i).booleanValue()); } releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testNull() throws GandivaException, Exception { /* * when x < 10 then 1 * else null */ ArrowType int64 = new ArrowType.Int(64, true); Field x = Field.nullable("x", int64); TreeNode xNode = TreeBuilder.makeField(x); // if (x < 10) then 1 else null TreeNode ifLess10 = ifLongLessThanElse(xNode, 10L, 1L, TreeBuilder.makeNull(int64), int64); ExpressionTree expr = TreeBuilder.makeExpression(ifLess10, x); Schema schema = new Schema(Lists.newArrayList(x)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 2; byte[] validity = new byte[] {(byte) 255}; long[] xValues = new long[] {5, 32}; long[] expected = new long[] {1, 0}; ArrowBuf bufValidity = buf(validity); ArrowBuf xData = longBuf(xValues); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode), Lists.newArrayList(bufValidity, xData)); BigIntVector bigIntVector = new BigIntVector(EMPTY_SCHEMA_PATH, allocator); bigIntVector.allocateNew(numRows); List output = new ArrayList(); output.add(bigIntVector); eval.evaluate(batch, output); // first element should be 1 assertFalse(bigIntVector.isNull(0)); assertEquals(expected[0], bigIntVector.get(0)); // second element should be null assertTrue(bigIntVector.isNull(1)); releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testTimeNull() throws GandivaException, Exception { ArrowType time64 = new ArrowType.Time(TimeUnit.MICROSECOND, 64); Field x = Field.nullable("x", time64); TreeNode xNode = TreeBuilder.makeNull(time64); ExpressionTree expr = TreeBuilder.makeExpression(xNode, x); Schema schema = new Schema(Lists.newArrayList(x)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 2; byte[] validity = new byte[] {(byte) 255}; int[] xValues = new int[] {5, 32}; ArrowBuf bufValidity = buf(validity); ArrowBuf xData = intBuf(xValues); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode), Lists.newArrayList(bufValidity, xData)); BigIntVector bigIntVector = new BigIntVector(EMPTY_SCHEMA_PATH, allocator); bigIntVector.allocateNew(numRows); List output = new ArrayList(); output.add(bigIntVector); eval.evaluate(batch, output); assertTrue(bigIntVector.isNull(0)); assertTrue(bigIntVector.isNull(1)); releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testTimeEquals() throws GandivaException, Exception { /* * when isnotnull(x) then x * else y */ Field x = Field.nullable("x", new ArrowType.Time(TimeUnit.MILLISECOND, 32)); TreeNode xNode = TreeBuilder.makeField(x); Field y = Field.nullable("y", new ArrowType.Time(TimeUnit.MILLISECOND, 32)); TreeNode yNode = TreeBuilder.makeField(y); // if isnotnull(x) then x else y TreeNode condition = TreeBuilder.makeFunction("isnotnull", Lists.newArrayList(xNode), boolType); TreeNode ifCoalesce = TreeBuilder.makeIf(condition, xNode, yNode, new ArrowType.Time(TimeUnit.MILLISECOND, 32)); ExpressionTree expr = TreeBuilder.makeExpression(ifCoalesce, x); Schema schema = new Schema(Lists.newArrayList(x, y)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 2; byte[] validity = new byte[] {(byte) 1}; byte[] yValidity = new byte[] {(byte) 3}; int[] xValues = new int[] {5, 1}; int[] yValues = new int[] {10, 2}; int[] expected = new int[] {5, 2}; ArrowBuf bufValidity = buf(validity); ArrowBuf xData = intBuf(xValues); ArrowBuf yBufValidity = buf(yValidity); ArrowBuf yData = intBuf(yValues); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode), Lists.newArrayList(bufValidity, xData, yBufValidity, yData)); IntVector intVector = new IntVector(EMPTY_SCHEMA_PATH, allocator); intVector.allocateNew(numRows); List output = new ArrayList(); output.add(intVector); eval.evaluate(batch, output); // output should be 5 and 2 assertFalse(intVector.isNull(0)); assertEquals(expected[0], intVector.get(0)); assertEquals(expected[1], intVector.get(1)); releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testIsNull() throws GandivaException, Exception { Field x = Field.nullable("x", float64); TreeNode xNode = TreeBuilder.makeField(x); TreeNode isNull = TreeBuilder.makeFunction("isnull", Lists.newArrayList(xNode), boolType); ExpressionTree expr = TreeBuilder.makeExpression(isNull, Field.nullable("result", boolType)); Schema schema = new Schema(Lists.newArrayList(x)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 16; byte[] validity = new byte[] {(byte) 255, 0}; double[] xValues = new double[] { 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0 }; ArrowBuf bufValidity = buf(validity); ArrowBuf xData = doubleBuf(xValues); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode), Lists.newArrayList(bufValidity, xData)); BitVector bitVector = new BitVector(EMPTY_SCHEMA_PATH, allocator); bitVector.allocateNew(numRows); List output = new ArrayList(); output.add(bitVector); eval.evaluate(batch, output); for (int i = 0; i < 8; i++) { assertFalse(bitVector.getObject(i).booleanValue()); } for (int i = 8; i < numRows; i++) { assertTrue(bitVector.getObject(i).booleanValue()); } releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testEquals() throws GandivaException, Exception { Field c1 = Field.nullable("c1", int32); Field c2 = Field.nullable("c2", int32); TreeNode c1Node = TreeBuilder.makeField(c1); TreeNode c2Node = TreeBuilder.makeField(c2); TreeNode equals = TreeBuilder.makeFunction("equal", Lists.newArrayList(c1Node, c2Node), boolType); ExpressionTree expr = TreeBuilder.makeExpression(equals, Field.nullable("result", boolType)); Schema schema = new Schema(Lists.newArrayList(c1, c2)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 16; byte[] validity = new byte[] {(byte) 255, 0}; int[] c1Values = new int[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}; int[] c2Values = new int[] {1, 2, 3, 4, 8, 7, 6, 5, 16, 15, 14, 13, 12, 11, 10, 9}; ArrowBuf c1Validity = buf(validity); ArrowBuf c1Data = intBuf(c1Values); ArrowBuf c2Validity = buf(validity); ArrowBuf c2Data = intBuf(c2Values); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode, fieldNode), Lists.newArrayList(c1Validity, c1Data, c2Validity, c2Data)); BitVector bitVector = new BitVector(EMPTY_SCHEMA_PATH, allocator); bitVector.allocateNew(numRows); List output = new ArrayList(); output.add(bitVector); eval.evaluate(batch, output); for (int i = 0; i < 4; i++) { assertTrue(bitVector.getObject(i).booleanValue()); } for (int i = 4; i < 8; i++) { assertFalse(bitVector.getObject(i).booleanValue()); } for (int i = 8; i < 16; i++) { assertTrue(bitVector.isNull(i)); } releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testInExpr() throws GandivaException, Exception { Field c1 = Field.nullable("c1", int32); TreeNode inExpr = TreeBuilder.makeInExpressionInt32( TreeBuilder.makeField(c1), Sets.newHashSet(1, 2, 3, 4, 5, 15, 16)); ExpressionTree expr = TreeBuilder.makeExpression(inExpr, Field.nullable("result", boolType)); Schema schema = new Schema(Lists.newArrayList(c1)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 16; byte[] validity = new byte[] {(byte) 255, 0}; int[] c1Values = new int[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}; ArrowBuf c1Validity = buf(validity); ArrowBuf c1Data = intBuf(c1Values); ArrowBuf c2Validity = buf(validity); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode, fieldNode), Lists.newArrayList(c1Validity, c1Data, c2Validity)); BitVector bitVector = new BitVector(EMPTY_SCHEMA_PATH, allocator); bitVector.allocateNew(numRows); List output = new ArrayList(); output.add(bitVector); eval.evaluate(batch, output); for (int i = 0; i < 5; i++) { assertTrue(bitVector.getObject(i).booleanValue()); } for (int i = 5; i < 16; i++) { assertFalse(bitVector.getObject(i).booleanValue()); } releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testInExprDecimal() throws GandivaException, Exception { Integer precision = 26; Integer scale = 5; ArrowType.Decimal decimal = new ArrowType.Decimal(precision, scale, 128); Field c1 = Field.nullable("c1", decimal); String[] values = new String[] {"1", "2", "3", "4"}; Set decimalSet = decimalSet(values, scale); decimalSet.add(new BigDecimal(-0.0)); decimalSet.add(new BigDecimal(Long.MAX_VALUE)); decimalSet.add(new BigDecimal(Long.MIN_VALUE)); TreeNode inExpr = TreeBuilder.makeInExpressionDecimal( TreeBuilder.makeField(c1), decimalSet, precision, scale); ExpressionTree expr = TreeBuilder.makeExpression(inExpr, Field.nullable("result", boolType)); Schema schema = new Schema(Lists.newArrayList(c1)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 16; byte[] validity = new byte[] {(byte) 255, 0}; String[] c1Values = new String[] { "1", "2", "3", "4", "-0.0", "6", "7", "8", "9", "10", "11", "12", "13", "14", String.valueOf(Long.MAX_VALUE), String.valueOf(Long.MIN_VALUE) }; DecimalVector c1Data = decimalVector(c1Values, precision, scale); ArrowBuf c1Validity = buf(validity); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode, fieldNode), Lists.newArrayList(c1Validity, c1Data.getDataBuffer(), c1Data.getValidityBuffer())); BitVector bitVector = new BitVector(EMPTY_SCHEMA_PATH, allocator); bitVector.allocateNew(numRows); List output = new ArrayList(); output.add(bitVector); eval.evaluate(batch, output); for (int i = 0; i < 5; i++) { assertTrue(bitVector.getObject(i).booleanValue()); } for (int i = 5; i < 16; i++) { assertFalse(bitVector.getObject(i).booleanValue()); } releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testInExprDouble() throws GandivaException, Exception { Field c1 = Field.nullable("c1", float64); TreeNode inExpr = TreeBuilder.makeInExpressionDouble( TreeBuilder.makeField(c1), Sets.newHashSet( 1.0, -0.0, 3.0, 4.0, Double.NaN, Double.POSITIVE_INFINITY, Double.NEGATIVE_INFINITY)); ExpressionTree expr = TreeBuilder.makeExpression(inExpr, Field.nullable("result", boolType)); Schema schema = new Schema(Lists.newArrayList(c1)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); // Create a row-batch with some sample data to look for int numRows = 16; // Only the first 8 values will be valid. byte[] validity = new byte[] {(byte) 255, 0}; double[] c1Values = new double[] { 1, -0.0, Double.NEGATIVE_INFINITY, Double.POSITIVE_INFINITY, Double.NaN, 6, 7, 8, 9, 10, 11, 12, 13, 14, 4, 3 }; ArrowBuf c1Validity = buf(validity); ArrowBuf c1Data = doubleBuf(c1Values); ArrowBuf c2Validity = buf(validity); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode, fieldNode), Lists.newArrayList(c1Validity, c1Data, c2Validity)); BitVector bitVector = new BitVector(EMPTY_SCHEMA_PATH, allocator); bitVector.allocateNew(numRows); List output = new ArrayList(); output.add(bitVector); eval.evaluate(batch, output); // The first four values in the vector must match the expression, but not the other ones. for (int i = 0; i < 4; i++) { assertTrue(bitVector.getObject(i).booleanValue()); } for (int i = 4; i < 16; i++) { assertFalse(bitVector.getObject(i).booleanValue()); } releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testInExprStrings() throws GandivaException, Exception { Field c1 = Field.nullable("c1", new ArrowType.Utf8()); TreeNode l1 = TreeBuilder.makeLiteral(1L); TreeNode l2 = TreeBuilder.makeLiteral(3L); List args = Lists.newArrayList(TreeBuilder.makeField(c1), l1, l2); TreeNode substr = TreeBuilder.makeFunction("substr", args, new ArrowType.Utf8()); TreeNode inExpr = TreeBuilder.makeInExpressionString(substr, Sets.newHashSet("one", "two", "thr", "fou")); ExpressionTree expr = TreeBuilder.makeExpression(inExpr, Field.nullable("result", boolType)); Schema schema = new Schema(Lists.newArrayList(c1)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 16; byte[] validity = new byte[] {(byte) 255, 0}; String[] c1Values = new String[] { "one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten", "eleven", "twelve", "thirteen", "fourteen", "fifteen", "sixteen" }; ArrowBuf c1Validity = buf(validity); List dataBufsX = stringBufs(c1Values); ArrowBuf c2Validity = buf(validity); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode, fieldNode), Lists.newArrayList(c1Validity, dataBufsX.get(0), dataBufsX.get(1), c2Validity)); BitVector bitVector = new BitVector(EMPTY_SCHEMA_PATH, allocator); bitVector.allocateNew(numRows); List output = new ArrayList(); output.add(bitVector); eval.evaluate(batch, output); for (int i = 0; i < 4; i++) { assertTrue(bitVector.getObject(i).booleanValue()); } for (int i = 5; i < 16; i++) { assertFalse(bitVector.getObject(i).booleanValue()); } releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testSmallOutputVectors() throws GandivaException, Exception { Field a = Field.nullable("a", int32); Field b = Field.nullable("b", int32); List args = Lists.newArrayList(a, b); Field retType = Field.nullable("c", int32); ExpressionTree root = TreeBuilder.makeExpression("add", args, retType); List exprs = Lists.newArrayList(root); Schema schema = new Schema(args); Projector eval = Projector.make(schema, exprs); int numRows = 16; byte[] validity = new byte[] {(byte) 255, 0}; // second half is "undefined" int[] aValues = new int[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}; int[] bValues = new int[] {16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1}; ArrowBuf aValidity = buf(validity); ArrowBuf aData = intBuf(aValues); ArrowBuf bValidity = buf(validity); ArrowBuf b2Validity = buf(validity); ArrowBuf bData = intBuf(bValues); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 8), new ArrowFieldNode(numRows, 8)), Lists.newArrayList(aValidity, aData, bValidity, bData, b2Validity)); IntVector intVector = new IntVector(EMPTY_SCHEMA_PATH, allocator); List output = new ArrayList(); output.add(intVector); try { eval.evaluate(batch, output); } catch (Throwable t) { intVector.allocateNew(numRows); eval.evaluate(batch, output); } for (int i = 0; i < 8; i++) { assertFalse(intVector.isNull(i)); assertEquals(17, intVector.get(i)); } for (int i = 8; i < 16; i++) { assertTrue(intVector.isNull(i)); } releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testDateTime() throws GandivaException, Exception { ArrowType date64 = new ArrowType.Date(DateUnit.MILLISECOND); // ArrowType time32 = new ArrowType.Time(TimeUnit.MILLISECOND, 32); ArrowType timeStamp = new ArrowType.Timestamp(TimeUnit.MILLISECOND, "TZ"); Field dateField = Field.nullable("date", date64); // Field timeField = Field.nullable("time", time32); Field tsField = Field.nullable("timestamp", timeStamp); TreeNode dateNode = TreeBuilder.makeField(dateField); TreeNode tsNode = TreeBuilder.makeField(tsField); List dateArgs = Lists.newArrayList(dateNode); TreeNode dateToYear = TreeBuilder.makeFunction("extractYear", dateArgs, int64); TreeNode dateToMonth = TreeBuilder.makeFunction("extractMonth", dateArgs, int64); TreeNode dateToDay = TreeBuilder.makeFunction("extractDay", dateArgs, int64); TreeNode dateToHour = TreeBuilder.makeFunction("extractHour", dateArgs, int64); TreeNode dateToMin = TreeBuilder.makeFunction("extractMinute", dateArgs, int64); List tsArgs = Lists.newArrayList(tsNode); TreeNode tsToYear = TreeBuilder.makeFunction("extractYear", tsArgs, int64); TreeNode tsToMonth = TreeBuilder.makeFunction("extractMonth", tsArgs, int64); TreeNode tsToDay = TreeBuilder.makeFunction("extractDay", tsArgs, int64); TreeNode tsToHour = TreeBuilder.makeFunction("extractHour", tsArgs, int64); TreeNode tsToMin = TreeBuilder.makeFunction("extractMinute", tsArgs, int64); Field resultField = Field.nullable("result", int64); List exprs = Lists.newArrayList( TreeBuilder.makeExpression(dateToYear, resultField), TreeBuilder.makeExpression(dateToMonth, resultField), TreeBuilder.makeExpression(dateToDay, resultField), TreeBuilder.makeExpression(dateToHour, resultField), TreeBuilder.makeExpression(dateToMin, resultField), TreeBuilder.makeExpression(tsToYear, resultField), TreeBuilder.makeExpression(tsToMonth, resultField), TreeBuilder.makeExpression(tsToDay, resultField), TreeBuilder.makeExpression(tsToHour, resultField), TreeBuilder.makeExpression(tsToMin, resultField)); Schema schema = new Schema(Lists.newArrayList(dateField, tsField)); Projector eval = Projector.make(schema, exprs); int numRows = 8; byte[] validity = new byte[] {(byte) 255}; String[] values = new String[] { "2007-01-01T01:00:00.00Z", "2007-03-05T03:40:00.00Z", "2008-05-31T13:55:00.00Z", "2000-06-30T23:20:00.00Z", "2000-07-10T20:30:00.00Z", "2000-08-20T00:14:00.00Z", "2000-09-30T02:29:00.00Z", "2000-10-31T05:33:00.00Z" }; long[] expYearFromDate = new long[] {2007, 2007, 2008, 2000, 2000, 2000, 2000, 2000}; long[] expMonthFromDate = new long[] {1, 3, 5, 6, 7, 8, 9, 10}; long[] expDayFromDate = new long[] {1, 5, 31, 30, 10, 20, 30, 31}; long[] expHourFromDate = new long[] {1, 3, 13, 23, 20, 0, 2, 5}; long[] expMinFromDate = new long[] {0, 40, 55, 20, 30, 14, 29, 33}; long[][] expValues = new long[][] { expYearFromDate, expMonthFromDate, expDayFromDate, expHourFromDate, expMinFromDate }; ArrowBuf bufValidity = buf(validity); ArrowBuf millisData = stringToMillis(values); ArrowBuf buf2Validity = buf(validity); ArrowBuf millis2Data = stringToMillis(values); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode, fieldNode), Lists.newArrayList(bufValidity, millisData, buf2Validity, millis2Data)); List output = new ArrayList(); for (int i = 0; i < exprs.size(); i++) { BigIntVector bigIntVector = new BigIntVector(EMPTY_SCHEMA_PATH, allocator); bigIntVector.allocateNew(numRows); output.add(bigIntVector); } eval.evaluate(batch, output); eval.close(); for (int i = 0; i < output.size(); i++) { long[] expected = expValues[i % 5]; BigIntVector bigIntVector = (BigIntVector) output.get(i); for (int j = 0; j < numRows; j++) { assertFalse(bigIntVector.isNull(j)); assertEquals(expected[j], bigIntVector.get(j)); } } releaseRecordBatch(batch); releaseValueVectors(output); } @Test public void testDateTrunc() throws Exception { ArrowType date64 = new ArrowType.Date(DateUnit.MILLISECOND); Field dateField = Field.nullable("date", date64); TreeNode dateNode = TreeBuilder.makeField(dateField); List dateArgs = Lists.newArrayList(dateNode); TreeNode dateToYear = TreeBuilder.makeFunction("date_trunc_Year", dateArgs, date64); TreeNode dateToMonth = TreeBuilder.makeFunction("date_trunc_Month", dateArgs, date64); Field resultField = Field.nullable("result", date64); List exprs = Lists.newArrayList( TreeBuilder.makeExpression(dateToYear, resultField), TreeBuilder.makeExpression(dateToMonth, resultField)); Schema schema = new Schema(Lists.newArrayList(dateField)); Projector eval = Projector.make(schema, exprs); int numRows = 4; byte[] validity = new byte[] {(byte) 255}; String[] values = new String[] { "2007-01-01T01:00:00.00Z", "2007-03-05T03:40:00.00Z", "2008-05-31T13:55:00.00Z", "2000-06-30T23:20:00.00Z", }; String[] expYearFromDate = new String[] { "2007-01-01T00:00:00.00Z", "2007-01-01T00:00:00.00Z", "2008-01-01T00:00:00.00Z", "2000-01-01T00:00:00.00Z", }; String[] expMonthFromDate = new String[] { "2007-01-01T00:00:00.00Z", "2007-03-01T00:00:00.00Z", "2008-05-01T00:00:00.00Z", "2000-06-01T00:00:00.00Z", }; String[][] expValues = new String[][] {expYearFromDate, expMonthFromDate}; ArrowBuf bufValidity = buf(validity); ArrowBuf millisData = stringToMillis(values); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode), Lists.newArrayList(bufValidity, millisData)); List output = new ArrayList(); for (int i = 0; i < exprs.size(); i++) { BigIntVector bigIntVector = new BigIntVector(EMPTY_SCHEMA_PATH, allocator); bigIntVector.allocateNew(numRows); output.add(bigIntVector); } eval.evaluate(batch, output); eval.close(); for (int i = 0; i < output.size(); i++) { String[] expected = expValues[i]; BigIntVector bigIntVector = (BigIntVector) output.get(i); for (int j = 0; j < numRows; j++) { assertFalse(bigIntVector.isNull(j)); assertEquals(Instant.parse(expected[j]).toEpochMilli(), bigIntVector.get(j)); } } releaseRecordBatch(batch); releaseValueVectors(output); } @Test public void testUnknownFunction() { Field c1 = Field.nullable("c1", int8); Field c2 = Field.nullable("c2", int8); TreeNode c1Node = TreeBuilder.makeField(c1); TreeNode c2Node = TreeBuilder.makeField(c2); TreeNode unknown = TreeBuilder.makeFunction("xxx_yyy", Lists.newArrayList(c1Node, c2Node), int8); ExpressionTree expr = TreeBuilder.makeExpression(unknown, Field.nullable("result", int8)); Schema schema = new Schema(Lists.newArrayList(c1, c2)); boolean caughtException = false; try { Projector eval = Projector.make(schema, Lists.newArrayList(expr)); } catch (GandivaException ge) { caughtException = true; } assertTrue(caughtException); } @Test public void testCastTimestampToString() throws Exception { ArrowType timeStamp = new ArrowType.Timestamp(TimeUnit.MILLISECOND, "TZ"); Field tsField = Field.nullable("timestamp", timeStamp); Field lenField = Field.nullable("outLength", int64); TreeNode tsNode = TreeBuilder.makeField(tsField); TreeNode lenNode = TreeBuilder.makeField(lenField); TreeNode tsToString = TreeBuilder.makeFunction( "castVARCHAR", Lists.newArrayList(tsNode, lenNode), new ArrowType.Utf8()); Field resultField = Field.nullable("result", new ArrowType.Utf8()); List exprs = Lists.newArrayList(TreeBuilder.makeExpression(tsToString, resultField)); Schema schema = new Schema(Lists.newArrayList(tsField, lenField)); Projector eval = Projector.make(schema, exprs); int numRows = 5; byte[] validity = new byte[] {(byte) 255}; String[] values = new String[] { "0007-01-01T01:00:00Z", "2007-03-05T03:40:00Z", "2008-05-31T13:55:00Z", "2000-06-30T23:20:00Z", "2000-07-10T20:30:00Z", }; long[] lenValues = new long[] {23L, 24L, 22L, 0L, 4L}; String[] expValues = new String[] { "0007-01-01 01:00:00.000", "2007-03-05 03:40:00.000", "2008-05-31 13:55:00.00", "", "2000", }; ArrowBuf bufValidity = buf(validity); ArrowBuf millisData = stringToMillis(values); ArrowBuf lenValidity = buf(validity); ArrowBuf lenData = longBuf(lenValues); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode, fieldNode), Lists.newArrayList(bufValidity, millisData, lenValidity, lenData)); List output = new ArrayList<>(); for (int i = 0; i < exprs.size(); i++) { VarCharVector charVector = new VarCharVector(EMPTY_SCHEMA_PATH, allocator); charVector.allocateNew(numRows * 23, numRows); output.add(charVector); } eval.evaluate(batch, output); eval.close(); for (ValueVector valueVector : output) { VarCharVector charVector = (VarCharVector) valueVector; for (int j = 0; j < numRows; j++) { assertFalse(charVector.isNull(j)); assertEquals(expValues[j], new String(charVector.get(j))); } } releaseRecordBatch(batch); releaseValueVectors(output); } @Test public void testCastDayIntervalToBigInt() throws Exception { ArrowType dayIntervalType = new ArrowType.Interval(IntervalUnit.DAY_TIME); Field dayIntervalField = Field.nullable("dayInterval", dayIntervalType); TreeNode intervalNode = TreeBuilder.makeField(dayIntervalField); TreeNode intervalToBigint = TreeBuilder.makeFunction("castBIGINT", Lists.newArrayList(intervalNode), int64); Field resultField = Field.nullable("result", int64); List exprs = Lists.newArrayList(TreeBuilder.makeExpression(intervalToBigint, resultField)); Schema schema = new Schema(Lists.newArrayList(dayIntervalField)); Projector eval = Projector.make(schema, exprs); int numRows = 5; byte[] validity = new byte[] {(byte) 255}; String[] values = new String[] { "1 0", // "days millis" "2 0", "1 1", "10 5000", "11 86400001", }; Long[] expValues = new Long[] { 86400000L, 2 * 86400000L, 86400000L + 1L, 10 * 86400000L + 5000L, 11 * 86400000L + 86400001L }; ArrowBuf bufValidity = buf(validity); ArrowBuf intervalsData = stringToDayInterval(values); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode, fieldNode), Lists.newArrayList(bufValidity, intervalsData)); List output = new ArrayList<>(); for (int i = 0; i < exprs.size(); i++) { BigIntVector bigIntVector = new BigIntVector(EMPTY_SCHEMA_PATH, allocator); bigIntVector.allocateNew(numRows); output.add(bigIntVector); } eval.evaluate(batch, output); eval.close(); for (ValueVector valueVector : output) { BigIntVector bigintVector = (BigIntVector) valueVector; for (int j = 0; j < numRows; j++) { assertFalse(bigintVector.isNull(j)); assertEquals(expValues[j], Long.valueOf(bigintVector.get(j))); } } releaseRecordBatch(batch); releaseValueVectors(output); } @Test public void testCaseInsensitiveFunctions() throws Exception { ArrowType timeStamp = new ArrowType.Timestamp(TimeUnit.MILLISECOND, "TZ"); Field tsField = Field.nullable("timestamp", timeStamp); TreeNode tsNode = TreeBuilder.makeField(tsField); TreeNode extractday = TreeBuilder.makeFunction("extractday", Lists.newArrayList(tsNode), int64); ExpressionTree expr = TreeBuilder.makeExpression(extractday, Field.nullable("result", int64)); Schema schema = new Schema(Lists.newArrayList(tsField)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 5; byte[] validity = new byte[] {(byte) 255}; String[] values = new String[] { "0007-01-01T01:00:00Z", "2007-03-05T03:40:00Z", "2008-05-31T13:55:00Z", "2000-06-30T23:20:00Z", "2000-07-10T20:30:00Z", }; long[] expValues = new long[] {1, 5, 31, 30, 10}; ArrowBuf bufValidity = buf(validity); ArrowBuf millisData = stringToMillis(values); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode), Lists.newArrayList(bufValidity, millisData)); List output = new ArrayList<>(); BigIntVector bigIntVector = new BigIntVector(EMPTY_SCHEMA_PATH, allocator); bigIntVector.allocateNew(numRows); output.add(bigIntVector); eval.evaluate(batch, output); eval.close(); for (ValueVector valueVector : output) { BigIntVector vector = (BigIntVector) valueVector; for (int j = 0; j < numRows; j++) { assertFalse(vector.isNull(j)); assertEquals(expValues[j], vector.get(j)); } } releaseRecordBatch(batch); releaseValueVectors(output); } @Test public void testCastInt() throws Exception { Field inField = Field.nullable("input", new ArrowType.Utf8()); TreeNode inNode = TreeBuilder.makeField(inField); TreeNode castINTFn = TreeBuilder.makeFunction("castINT", Lists.newArrayList(inNode), int32); Field resultField = Field.nullable("result", int32); List exprs = Lists.newArrayList(TreeBuilder.makeExpression(castINTFn, resultField)); Schema schema = new Schema(Lists.newArrayList(inField)); Projector eval = Projector.make(schema, exprs); int numRows = 5; byte[] validity = new byte[] {(byte) 255}; String[] values = new String[] {"0", "123", "-123", "-1", "1"}; int[] expValues = new int[] {0, 123, -123, -1, 1}; ArrowBuf bufValidity = buf(validity); List bufData = stringBufs(values); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode), Lists.newArrayList(bufValidity, bufData.get(0), bufData.get(1))); List output = new ArrayList<>(); for (int i = 0; i < exprs.size(); i++) { IntVector intVector = new IntVector(EMPTY_SCHEMA_PATH, allocator); intVector.allocateNew(numRows); output.add(intVector); } eval.evaluate(batch, output); eval.close(); for (ValueVector valueVector : output) { IntVector intVector = (IntVector) valueVector; for (int j = 0; j < numRows; j++) { assertFalse(intVector.isNull(j)); assertTrue(expValues[j] == intVector.get(j)); } } releaseRecordBatch(batch); releaseValueVectors(output); } @Test public void testCastIntInvalidValue() throws Exception { Field inField = Field.nullable("input", new ArrowType.Utf8()); TreeNode inNode = TreeBuilder.makeField(inField); TreeNode castINTFn = TreeBuilder.makeFunction("castINT", Lists.newArrayList(inNode), int32); Field resultField = Field.nullable("result", int32); List exprs = Lists.newArrayList(TreeBuilder.makeExpression(castINTFn, resultField)); Schema schema = new Schema(Lists.newArrayList(inField)); Projector eval = Projector.make(schema, exprs); int numRows = 1; byte[] validity = new byte[] {(byte) 255}; String[] values = new String[] {"abc"}; ArrowBuf bufValidity = buf(validity); List bufData = stringBufs(values); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode), Lists.newArrayList(bufValidity, bufData.get(0), bufData.get(1))); List output = new ArrayList<>(); for (int i = 0; i < exprs.size(); i++) { IntVector intVector = new IntVector(EMPTY_SCHEMA_PATH, allocator); intVector.allocateNew(numRows); output.add(intVector); } assertThrows( GandivaException.class, () -> { try { eval.evaluate(batch, output); } finally { eval.close(); releaseRecordBatch(batch); releaseValueVectors(output); } }); } @Test public void testCastFloat() throws Exception { Field inField = Field.nullable("input", new ArrowType.Utf8()); TreeNode inNode = TreeBuilder.makeField(inField); TreeNode castFLOAT8Fn = TreeBuilder.makeFunction("castFLOAT8", Lists.newArrayList(inNode), float64); Field resultField = Field.nullable("result", float64); List exprs = Lists.newArrayList(TreeBuilder.makeExpression(castFLOAT8Fn, resultField)); Schema schema = new Schema(Lists.newArrayList(inField)); Projector eval = Projector.make(schema, exprs); int numRows = 5; byte[] validity = new byte[] {(byte) 255}; String[] values = new String[] {"2.3", "-11.11", "0", "111", "12345.67"}; double[] expValues = new double[] {2.3, -11.11, 0, 111, 12345.67}; ArrowBuf bufValidity = buf(validity); List bufData = stringBufs(values); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode), Lists.newArrayList(bufValidity, bufData.get(0), bufData.get(1))); List output = new ArrayList<>(); for (int i = 0; i < exprs.size(); i++) { Float8Vector float8Vector = new Float8Vector(EMPTY_SCHEMA_PATH, allocator); float8Vector.allocateNew(numRows); output.add(float8Vector); } eval.evaluate(batch, output); eval.close(); for (ValueVector valueVector : output) { Float8Vector float8Vector = (Float8Vector) valueVector; for (int j = 0; j < numRows; j++) { assertFalse(float8Vector.isNull(j)); assertTrue(expValues[j] == float8Vector.get(j)); } } releaseRecordBatch(batch); releaseValueVectors(output); } @Test public void testCastFloatVarbinary() throws Exception { Field inField = Field.nullable("input", new ArrowType.Binary()); TreeNode inNode = TreeBuilder.makeField(inField); TreeNode castFLOAT8Fn = TreeBuilder.makeFunction("castFLOAT8", Lists.newArrayList(inNode), float64); Field resultField = Field.nullable("result", float64); List exprs = Lists.newArrayList(TreeBuilder.makeExpression(castFLOAT8Fn, resultField)); Schema schema = new Schema(Lists.newArrayList(inField)); Projector eval = Projector.make(schema, exprs); int numRows = 5; byte[] validity = new byte[] {(byte) 255}; String[] values = new String[] {"2.3", "-11.11", "0", "111", "12345.67"}; double[] expValues = new double[] {2.3, -11.11, 0, 111, 12345.67}; ArrowBuf bufValidity = buf(validity); List bufData = stringBufs(values); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode), Lists.newArrayList(bufValidity, bufData.get(0), bufData.get(1))); List output = new ArrayList<>(); for (int i = 0; i < exprs.size(); i++) { Float8Vector float8Vector = new Float8Vector(EMPTY_SCHEMA_PATH, allocator); float8Vector.allocateNew(numRows); output.add(float8Vector); } eval.evaluate(batch, output); eval.close(); for (ValueVector valueVector : output) { Float8Vector float8Vector = (Float8Vector) valueVector; for (int j = 0; j < numRows; j++) { assertFalse(float8Vector.isNull(j)); assertTrue(expValues[j] == float8Vector.get(j)); } } releaseRecordBatch(batch); releaseValueVectors(output); } @Test public void testCastFloatInvalidValue() throws Exception { Field inField = Field.nullable("input", new ArrowType.Utf8()); TreeNode inNode = TreeBuilder.makeField(inField); TreeNode castFLOAT8Fn = TreeBuilder.makeFunction("castFLOAT8", Lists.newArrayList(inNode), float64); Field resultField = Field.nullable("result", float64); List exprs = Lists.newArrayList(TreeBuilder.makeExpression(castFLOAT8Fn, resultField)); Schema schema = new Schema(Lists.newArrayList(inField)); Projector eval = Projector.make(schema, exprs); int numRows = 5; byte[] validity = new byte[] {(byte) 255}; String[] values = new String[] {"2.3", "-11.11", "abc", "111", "12345.67"}; ArrowBuf bufValidity = buf(validity); List bufData = stringBufs(values); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode), Lists.newArrayList(bufValidity, bufData.get(0), bufData.get(1))); List output = new ArrayList<>(); for (int i = 0; i < exprs.size(); i++) { Float8Vector float8Vector = new Float8Vector(EMPTY_SCHEMA_PATH, allocator); float8Vector.allocateNew(numRows); output.add(float8Vector); } assertThrows( GandivaException.class, () -> { try { eval.evaluate(batch, output); } finally { eval.close(); releaseRecordBatch(batch); releaseValueVectors(output); } }); } @Test public void testEvaluateWithUnsetTargetHostCPU() throws Exception { Field a = Field.nullable("a", int32); Field b = Field.nullable("b", int32); List args = Lists.newArrayList(a, b); Field retType = Field.nullable("c", int32); ExpressionTree root = TreeBuilder.makeExpression("add", args, retType); List exprs = Lists.newArrayList(root); Schema schema = new Schema(args); Projector eval = Projector.make( schema, exprs, new ConfigurationBuilder.ConfigOptions().withTargetCPU(false)); int numRows = 16; byte[] validity = new byte[] {(byte) 255, 0}; // second half is "undefined" int[] aValues = new int[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}; int[] bValues = new int[] {16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1}; ArrowBuf validitya = buf(validity); ArrowBuf valuesa = intBuf(aValues); ArrowBuf validityb = buf(validity); ArrowBuf valuesb = intBuf(bValues); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 8), new ArrowFieldNode(numRows, 8)), Lists.newArrayList(validitya, valuesa, validityb, valuesb)); IntVector intVector = new IntVector(EMPTY_SCHEMA_PATH, allocator); intVector.allocateNew(numRows); List output = new ArrayList(); output.add(intVector); eval.evaluate(batch, output); for (int i = 0; i < 8; i++) { assertFalse(intVector.isNull(i)); assertEquals(17, intVector.get(i)); } for (int i = 8; i < 16; i++) { assertTrue(intVector.isNull(i)); } // free buffers releaseRecordBatch(batch); releaseValueVectors(output); eval.close(); } @Test public void testCastVarcharFromInteger() throws Exception { Field inField = Field.nullable("input", int32); Field lenField = Field.nullable("outLength", int64); TreeNode inNode = TreeBuilder.makeField(inField); TreeNode lenNode = TreeBuilder.makeField(lenField); TreeNode tsToString = TreeBuilder.makeFunction( "castVARCHAR", Lists.newArrayList(inNode, lenNode), new ArrowType.Utf8()); Field resultField = Field.nullable("result", new ArrowType.Utf8()); List exprs = Lists.newArrayList(TreeBuilder.makeExpression(tsToString, resultField)); Schema schema = new Schema(Lists.newArrayList(inField, lenField)); Projector eval = Projector.make(schema, exprs); int numRows = 5; byte[] validity = new byte[] {(byte) 255}; int[] values = new int[] { 2345, 2345, 2345, 2345, -2345, }; long[] lenValues = new long[] {0L, 4L, 2L, 6L, 5L}; String[] expValues = new String[] { "", Integer.toString(2345).substring(0, 4), Integer.toString(2345).substring(0, 2), Integer.toString(2345), Integer.toString(-2345) }; ArrowBuf bufValidity = buf(validity); ArrowBuf bufData = intBuf(values); ArrowBuf lenValidity = buf(validity); ArrowBuf lenData = longBuf(lenValues); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode, fieldNode), Lists.newArrayList(bufValidity, bufData, lenValidity, lenData)); List output = new ArrayList<>(); for (int i = 0; i < exprs.size(); i++) { VarCharVector charVector = new VarCharVector(EMPTY_SCHEMA_PATH, allocator); charVector.allocateNew(numRows * 5, numRows); output.add(charVector); } eval.evaluate(batch, output); eval.close(); for (ValueVector valueVector : output) { VarCharVector charVector = (VarCharVector) valueVector; for (int j = 0; j < numRows; j++) { assertFalse(charVector.isNull(j)); assertEquals(expValues[j], new String(charVector.get(j))); } } releaseRecordBatch(batch); releaseValueVectors(output); } @Test public void testCastVarcharFromFloat() throws Exception { Field inField = Field.nullable("input", float64); Field lenField = Field.nullable("outLength", int64); TreeNode inNode = TreeBuilder.makeField(inField); TreeNode lenNode = TreeBuilder.makeField(lenField); TreeNode tsToString = TreeBuilder.makeFunction( "castVARCHAR", Lists.newArrayList(inNode, lenNode), new ArrowType.Utf8()); Field resultField = Field.nullable("result", new ArrowType.Utf8()); List exprs = Lists.newArrayList(TreeBuilder.makeExpression(tsToString, resultField)); Schema schema = new Schema(Lists.newArrayList(inField, lenField)); Projector eval = Projector.make(schema, exprs); int numRows = 5; byte[] validity = new byte[] {(byte) 255}; double[] values = new double[] { 0.0, -0.0, 1.0, 0.001, 0.0009, 0.00099893, 999999.9999, 10000000.0, 23943410000000.343434, Double.POSITIVE_INFINITY, Double.NEGATIVE_INFINITY, Double.NaN, 23.45, 23.45, -23.45, }; long[] lenValues = new long[] {6L, 6L, 6L, 6L, 10L, 15L, 15L, 15L, 30L, 15L, 15L, 15L, 0L, 6L, 6L}; /* The Java real numbers are represented in two ways and Gandiva must * follow the same rules: * - If the number is greater or equals than 10^7 and less than 10^(-3) * it will be represented using scientific notation, e.g: * - 0.000012 -> 1.2E-5 * - 10000002.3 -> 1.00000023E7 * - If the numbers are between that interval above, they are showed as is. * * The test checks if the Gandiva function casts the number with the same notation of the * Java. * */ String[] expValues = new String[] { Double.toString(0.0), // must be cast to -> "0.0" Double.toString(-0.0), // must be cast to -> "-0.0" Double.toString(1.0), // must be cast to -> "1.0" Double.toString(0.001), // must be cast to -> "0.001" Double.toString(0.0009), // must be cast to -> "9E-4" Double.toString(0.00099893), // must be cast to -> "9E-4" Double.toString(999999.9999), // must be cast to -> "999999.9999" Double.toString(10000000.0), // must be cast to 1E7 Double.toString(23943410000000.343434), Double.toString(Double.POSITIVE_INFINITY), Double.toString(Double.NEGATIVE_INFINITY), Double.toString(Double.NaN), "", Double.toString(23.45), Double.toString(-23.45) }; ArrowBuf bufValidity = buf(validity); ArrowBuf bufData = doubleBuf(values); ArrowBuf lenValidity = buf(validity); ArrowBuf lenData = longBuf(lenValues); ArrowFieldNode fieldNode = new ArrowFieldNode(numRows, 0); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(fieldNode, fieldNode), Lists.newArrayList(bufValidity, bufData, lenValidity, lenData)); List output = new ArrayList<>(); for (int i = 0; i < exprs.size(); i++) { VarCharVector charVector = new VarCharVector(EMPTY_SCHEMA_PATH, allocator); charVector.allocateNew(numRows * 5, numRows); output.add(charVector); } eval.evaluate(batch, output); eval.close(); for (ValueVector valueVector : output) { VarCharVector charVector = (VarCharVector) valueVector; for (int j = 0; j < numRows; j++) { assertFalse(charVector.isNull(j)); assertEquals(expValues[j], new String(charVector.get(j))); } } releaseRecordBatch(batch); releaseValueVectors(output); } @Test public void testInitCap() throws Exception { Field x = Field.nullable("x", new ArrowType.Utf8()); Field retType = Field.nullable("c", new ArrowType.Utf8()); TreeNode cond = TreeBuilder.makeFunction( "initcap", Lists.newArrayList(TreeBuilder.makeField(x)), new ArrowType.Utf8()); ExpressionTree expr = TreeBuilder.makeExpression(cond, retType); Schema schema = new Schema(Lists.newArrayList(x)); Projector eval = Projector.make(schema, Lists.newArrayList(expr)); int numRows = 5; byte[] validity = new byte[] {(byte) 15, 0}; String[] valuesX = new String[] { " øhpqršvñ \n\n", "möbelträger1füße \nmöbelträge'rfüße", "ÂbĆDËFgh\néll", "citroën CaR", "kjk" }; String[] expected = new String[] { " Øhpqršvñ \n\n", "Möbelträger1füße \nMöbelträge'Rfüße", "Âbćdëfgh\nÉll", "Citroën Car", null }; ArrowBuf validityX = buf(validity); List dataBufsX = stringBufs(valuesX); ArrowRecordBatch batch = new ArrowRecordBatch( numRows, Lists.newArrayList(new ArrowFieldNode(numRows, 0)), Lists.newArrayList(validityX, dataBufsX.get(0), dataBufsX.get(1))); // allocate data for output vector. VarCharVector outVector = new VarCharVector(EMPTY_SCHEMA_PATH, allocator); outVector.allocateNew(numRows * 100, numRows); // evaluate expression List output = new ArrayList<>(); output.add(outVector); eval.evaluate(batch, output); eval.close(); // match expected output. for (int i = 0; i < numRows - 1; i++) { assertFalse(outVector.isNull(i), "Expect none value equals null"); assertEquals(expected[i], new String(outVector.get(i))); } assertTrue(outVector.isNull(numRows - 1), "Last value must be null"); releaseRecordBatch(batch); releaseValueVectors(output); } } apache-arrow-java-19.0.0/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/TestJniLoader.java000066400000000000000000000055511515451615200333330ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.evaluator; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import org.junit.jupiter.api.Test; public class TestJniLoader { @Test public void testDefaultConfiguration() throws Exception { long configId = JniLoader.getConfiguration(ConfigurationBuilder.ConfigOptions.getDefault()); assertEquals(configId, JniLoader.getDefaultConfiguration()); assertEquals( configId, JniLoader.getConfiguration(ConfigurationBuilder.ConfigOptions.getDefault())); long configId2 = JniLoader.getConfiguration(new ConfigurationBuilder.ConfigOptions().withOptimize(false)); long configId3 = JniLoader.getConfiguration(new ConfigurationBuilder.ConfigOptions().withTargetCPU(false)); long configId4 = JniLoader.getConfiguration( new ConfigurationBuilder.ConfigOptions().withOptimize(false).withTargetCPU(false)); assertTrue(configId != configId2 && configId2 != configId3 && configId3 != configId4); assertEquals( configId2, JniLoader.getConfiguration(new ConfigurationBuilder.ConfigOptions().withOptimize(false))); assertEquals( configId3, JniLoader.getConfiguration(new ConfigurationBuilder.ConfigOptions().withTargetCPU(false))); assertEquals( configId4, JniLoader.getConfiguration( new ConfigurationBuilder.ConfigOptions().withOptimize(false).withTargetCPU(false))); JniLoader.removeConfiguration(new ConfigurationBuilder.ConfigOptions().withOptimize(false)); // configids are monotonically updated. after a config is removed, new one is assigned with // higher id assertNotEquals( configId2, JniLoader.getConfiguration(new ConfigurationBuilder.ConfigOptions().withOptimize(false))); JniLoader.removeConfiguration(new ConfigurationBuilder.ConfigOptions()); assertNotEquals( configId, JniLoader.getConfiguration(ConfigurationBuilder.ConfigOptions.getDefault())); } } apache-arrow-java-19.0.0/gandiva/src/test/java/org/apache/arrow/gandiva/expression/000077500000000000000000000000001515451615200301505ustar00rootroot00000000000000ArrowTypeHelperTest.java000066400000000000000000000105461515451615200346760ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/test/java/org/apache/arrow/gandiva/expression/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import static org.junit.jupiter.api.Assertions.assertEquals; import java.util.ArrayList; import java.util.List; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; import org.apache.arrow.vector.types.FloatingPointPrecision; 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; public class ArrowTypeHelperTest { private void testInt(int width, boolean isSigned, int expected) throws GandivaException { ArrowType arrowType = new ArrowType.Int(width, isSigned); GandivaTypes.ExtGandivaType gandivaType = ArrowTypeHelper.arrowTypeToProtobuf(arrowType); assertEquals(expected, gandivaType.getType().getNumber()); } @Test public void testAllInts() throws GandivaException { testInt(8, false, GandivaTypes.GandivaType.UINT8_VALUE); testInt(8, true, GandivaTypes.GandivaType.INT8_VALUE); testInt(16, false, GandivaTypes.GandivaType.UINT16_VALUE); testInt(16, true, GandivaTypes.GandivaType.INT16_VALUE); testInt(32, false, GandivaTypes.GandivaType.UINT32_VALUE); testInt(32, true, GandivaTypes.GandivaType.INT32_VALUE); testInt(64, false, GandivaTypes.GandivaType.UINT64_VALUE); testInt(64, true, GandivaTypes.GandivaType.INT64_VALUE); } private void testFloat(FloatingPointPrecision precision, int expected) throws GandivaException { ArrowType arrowType = new ArrowType.FloatingPoint(precision); GandivaTypes.ExtGandivaType gandivaType = ArrowTypeHelper.arrowTypeToProtobuf(arrowType); assertEquals(expected, gandivaType.getType().getNumber()); } @Test public void testAllFloats() throws GandivaException { testFloat(FloatingPointPrecision.HALF, GandivaTypes.GandivaType.HALF_FLOAT_VALUE); testFloat(FloatingPointPrecision.SINGLE, GandivaTypes.GandivaType.FLOAT_VALUE); testFloat(FloatingPointPrecision.DOUBLE, GandivaTypes.GandivaType.DOUBLE_VALUE); } private void testBasic(ArrowType arrowType, int expected) throws GandivaException { GandivaTypes.ExtGandivaType gandivaType = ArrowTypeHelper.arrowTypeToProtobuf(arrowType); assertEquals(expected, gandivaType.getType().getNumber()); } @Test public void testSimpleTypes() throws GandivaException { testBasic(new ArrowType.Bool(), GandivaTypes.GandivaType.BOOL_VALUE); testBasic(new ArrowType.Binary(), GandivaTypes.GandivaType.BINARY_VALUE); testBasic(new ArrowType.Utf8(), GandivaTypes.GandivaType.UTF8_VALUE); } @Test public void testField() throws GandivaException { Field field = Field.nullable("col1", new ArrowType.Bool()); GandivaTypes.Field f = ArrowTypeHelper.arrowFieldToProtobuf(field); assertEquals(field.getName(), f.getName()); assertEquals(true, f.getNullable()); assertEquals(GandivaTypes.GandivaType.BOOL_VALUE, f.getType().getType().getNumber()); } @Test public void testSchema() throws GandivaException { Field a = Field.nullable("a", new ArrowType.Int(16, false)); Field b = Field.nullable("b", new ArrowType.Int(32, true)); Field c = Field.nullable("c", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)); List fields = new ArrayList(); fields.add(a); fields.add(b); fields.add(c); GandivaTypes.Schema schema = ArrowTypeHelper.arrowSchemaToProtobuf(new Schema(fields)); int idx = 0; for (GandivaTypes.Field f : schema.getColumnsList()) { assertEquals(fields.get(idx).getName(), f.getName()); idx++; } } } TreeBuilderTest.java000066400000000000000000000326111515451615200340050ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/test/java/org/apache/arrow/gandiva/expression/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.gandiva.expression; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.ArrayList; import java.util.Arrays; import java.util.List; import org.apache.arrow.gandiva.exceptions.GandivaException; import org.apache.arrow.gandiva.ipc.GandivaTypes; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.junit.jupiter.api.Test; public class TreeBuilderTest { @Test public void testMakeLiteral() throws GandivaException { TreeNode n = TreeBuilder.makeLiteral(Boolean.TRUE); GandivaTypes.TreeNode node = n.toProtobuf(); assertEquals(true, node.getBooleanNode().getValue()); n = TreeBuilder.makeLiteral(new Integer(10)); node = n.toProtobuf(); assertEquals(10, node.getIntNode().getValue()); n = TreeBuilder.makeLiteral(new Long(50)); node = n.toProtobuf(); assertEquals(50, node.getLongNode().getValue()); Float f = new Float(2.5); n = TreeBuilder.makeLiteral(f); node = n.toProtobuf(); assertEquals(f.floatValue(), node.getFloatNode().getValue(), 0.1); Double d = new Double(3.3); n = TreeBuilder.makeLiteral(d); node = n.toProtobuf(); assertEquals(d.doubleValue(), node.getDoubleNode().getValue(), 0.1); String s = new String("hello"); n = TreeBuilder.makeStringLiteral(s); node = n.toProtobuf(); assertArrayEquals(s.getBytes(), node.getStringNode().getValue().toByteArray()); byte[] b = new String("hello").getBytes(); n = TreeBuilder.makeBinaryLiteral(b); node = n.toProtobuf(); assertArrayEquals(b, node.getBinaryNode().getValue().toByteArray()); } @Test public void testMakeNull() throws GandivaException { TreeNode n = TreeBuilder.makeNull(new ArrowType.Bool()); GandivaTypes.TreeNode node = n.toProtobuf(); assertEquals( GandivaTypes.GandivaType.BOOL_VALUE, node.getNullNode().getType().getType().getNumber()); n = TreeBuilder.makeNull(new ArrowType.Int(32, true)); node = n.toProtobuf(); assertEquals( GandivaTypes.GandivaType.INT32_VALUE, node.getNullNode().getType().getType().getNumber()); n = TreeBuilder.makeNull(new ArrowType.Int(64, false)); node = n.toProtobuf(); assertEquals( GandivaTypes.GandivaType.UINT64_VALUE, node.getNullNode().getType().getType().getNumber()); n = TreeBuilder.makeNull(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)); node = n.toProtobuf(); assertEquals( GandivaTypes.GandivaType.FLOAT_VALUE, node.getNullNode().getType().getType().getNumber()); } @Test public void testMakeField() throws GandivaException { TreeNode n = TreeBuilder.makeField(Field.nullable("a", new ArrowType.Int(32, true))); GandivaTypes.TreeNode node = n.toProtobuf(); assertEquals("a", node.getFieldNode().getField().getName()); assertEquals( GandivaTypes.GandivaType.INT32_VALUE, node.getFieldNode().getField().getType().getType().getNumber()); } @Test public void testMakeFunction() throws GandivaException { TreeNode a = TreeBuilder.makeField(Field.nullable("a", new ArrowType.Int(64, false))); TreeNode b = TreeBuilder.makeField(Field.nullable("b", new ArrowType.Int(64, false))); List args = new ArrayList(2); args.add(a); args.add(b); TreeNode addNode = TreeBuilder.makeFunction("add", args, new ArrowType.Int(64, false)); GandivaTypes.TreeNode node = addNode.toProtobuf(); assertTrue(node.hasFnNode()); assertEquals("add", node.getFnNode().getFunctionName()); assertEquals("a", node.getFnNode().getInArgsList().get(0).getFieldNode().getField().getName()); assertEquals("b", node.getFnNode().getInArgsList().get(1).getFieldNode().getField().getName()); assertEquals( GandivaTypes.GandivaType.UINT64_VALUE, node.getFnNode().getReturnType().getType().getNumber()); } @Test public void testMakeIf() throws GandivaException { Field a = Field.nullable("a", new ArrowType.Int(64, false)); Field b = Field.nullable("b", new ArrowType.Int(64, false)); TreeNode aNode = TreeBuilder.makeField(a); TreeNode bNode = TreeBuilder.makeField(b); List args = new ArrayList(2); args.add(aNode); args.add(bNode); ArrowType retType = new ArrowType.Bool(); TreeNode cond = TreeBuilder.makeFunction("greater_than", args, retType); TreeNode ifNode = TreeBuilder.makeIf(cond, aNode, bNode, retType); GandivaTypes.TreeNode node = ifNode.toProtobuf(); assertTrue(node.hasIfNode()); assertEquals("greater_than", node.getIfNode().getCond().getFnNode().getFunctionName()); assertEquals(a.getName(), node.getIfNode().getThenNode().getFieldNode().getField().getName()); assertEquals(b.getName(), node.getIfNode().getElseNode().getFieldNode().getField().getName()); assertEquals( GandivaTypes.GandivaType.BOOL_VALUE, node.getIfNode().getReturnType().getType().getNumber()); } @Test public void testMakeAnd() throws GandivaException { TreeNode a = TreeBuilder.makeField(Field.nullable("a", new ArrowType.Bool())); TreeNode b = TreeBuilder.makeField(Field.nullable("b", new ArrowType.Bool())); List args = new ArrayList(2); args.add(a); args.add(b); TreeNode andNode = TreeBuilder.makeAnd(args); GandivaTypes.TreeNode node = andNode.toProtobuf(); assertTrue(node.hasAndNode()); assertEquals(2, node.getAndNode().getArgsList().size()); assertEquals("a", node.getAndNode().getArgsList().get(0).getFieldNode().getField().getName()); assertEquals("b", node.getAndNode().getArgsList().get(1).getFieldNode().getField().getName()); } @Test public void testMakeOr() throws GandivaException { TreeNode a = TreeBuilder.makeField(Field.nullable("a", new ArrowType.Bool())); TreeNode b = TreeBuilder.makeField(Field.nullable("b", new ArrowType.Bool())); List args = new ArrayList(2); args.add(a); args.add(b); TreeNode orNode = TreeBuilder.makeOr(args); GandivaTypes.TreeNode node = orNode.toProtobuf(); assertTrue(node.hasOrNode()); assertEquals(2, node.getOrNode().getArgsList().size()); assertEquals("a", node.getOrNode().getArgsList().get(0).getFieldNode().getField().getName()); assertEquals("b", node.getOrNode().getArgsList().get(1).getFieldNode().getField().getName()); } @Test public void testExpression() throws GandivaException { Field a = Field.nullable("a", new ArrowType.Int(64, false)); Field b = Field.nullable("b", new ArrowType.Int(64, false)); TreeNode aNode = TreeBuilder.makeField(a); TreeNode bNode = TreeBuilder.makeField(b); List args = new ArrayList(2); args.add(aNode); args.add(bNode); ArrowType retType = new ArrowType.Bool(); TreeNode cond = TreeBuilder.makeFunction("greater_than", args, retType); TreeNode ifNode = TreeBuilder.makeIf(cond, aNode, bNode, retType); ExpressionTree expr = TreeBuilder.makeExpression(ifNode, Field.nullable("c", retType)); GandivaTypes.ExpressionRoot root = expr.toProtobuf(); assertTrue(root.getRoot().hasIfNode()); assertEquals( "greater_than", root.getRoot().getIfNode().getCond().getFnNode().getFunctionName()); assertEquals("c", root.getResultType().getName()); assertEquals( GandivaTypes.GandivaType.BOOL_VALUE, root.getResultType().getType().getType().getNumber()); } @Test public void testExpression2() throws GandivaException { Field a = Field.nullable("a", new ArrowType.Int(64, false)); Field b = Field.nullable("b", new ArrowType.Int(64, false)); List args = new ArrayList(2); args.add(a); args.add(b); Field c = Field.nullable("c", new ArrowType.Int(64, false)); ExpressionTree expr = TreeBuilder.makeExpression("add", args, c); GandivaTypes.ExpressionRoot root = expr.toProtobuf(); GandivaTypes.TreeNode node = root.getRoot(); assertEquals("c", root.getResultType().getName()); assertTrue(node.hasFnNode()); assertEquals("add", node.getFnNode().getFunctionName()); assertEquals("a", node.getFnNode().getInArgsList().get(0).getFieldNode().getField().getName()); assertEquals("b", node.getFnNode().getInArgsList().get(1).getFieldNode().getField().getName()); assertEquals( GandivaTypes.GandivaType.UINT64_VALUE, node.getFnNode().getReturnType().getType().getNumber()); } @Test public void testExpressionWithAnd() throws GandivaException { TreeNode a = TreeBuilder.makeField(Field.nullable("a", new ArrowType.Bool())); TreeNode b = TreeBuilder.makeField(Field.nullable("b", new ArrowType.Bool())); List args = new ArrayList(2); args.add(a); args.add(b); TreeNode andNode = TreeBuilder.makeAnd(args); ExpressionTree expr = TreeBuilder.makeExpression(andNode, Field.nullable("c", new ArrowType.Bool())); GandivaTypes.ExpressionRoot root = expr.toProtobuf(); assertTrue(root.getRoot().hasAndNode()); assertEquals( "a", root.getRoot().getAndNode().getArgsList().get(0).getFieldNode().getField().getName()); assertEquals( "b", root.getRoot().getAndNode().getArgsList().get(1).getFieldNode().getField().getName()); assertEquals("c", root.getResultType().getName()); assertEquals( GandivaTypes.GandivaType.BOOL_VALUE, root.getResultType().getType().getType().getNumber()); } @Test public void testExpressionWithOr() throws GandivaException { TreeNode a = TreeBuilder.makeField(Field.nullable("a", new ArrowType.Bool())); TreeNode b = TreeBuilder.makeField(Field.nullable("b", new ArrowType.Bool())); List args = new ArrayList(2); args.add(a); args.add(b); TreeNode orNode = TreeBuilder.makeOr(args); ExpressionTree expr = TreeBuilder.makeExpression(orNode, Field.nullable("c", new ArrowType.Bool())); GandivaTypes.ExpressionRoot root = expr.toProtobuf(); assertTrue(root.getRoot().hasOrNode()); assertEquals( "a", root.getRoot().getOrNode().getArgsList().get(0).getFieldNode().getField().getName()); assertEquals( "b", root.getRoot().getOrNode().getArgsList().get(1).getFieldNode().getField().getName()); assertEquals("c", root.getResultType().getName()); assertEquals( GandivaTypes.GandivaType.BOOL_VALUE, root.getResultType().getType().getType().getNumber()); } @Test public void testCondition() throws GandivaException { Field a = Field.nullable("a", new ArrowType.Int(64, false)); Field b = Field.nullable("b", new ArrowType.Int(64, false)); TreeNode aNode = TreeBuilder.makeField(a); TreeNode bNode = TreeBuilder.makeField(b); List args = new ArrayList(2); args.add(aNode); args.add(bNode); TreeNode root = TreeBuilder.makeFunction("greater_than", args, new ArrowType.Bool()); Condition condition = TreeBuilder.makeCondition(root); GandivaTypes.Condition conditionProto = condition.toProtobuf(); assertTrue(conditionProto.getRoot().hasFnNode()); assertEquals("greater_than", conditionProto.getRoot().getFnNode().getFunctionName()); assertEquals( "a", conditionProto .getRoot() .getFnNode() .getInArgsList() .get(0) .getFieldNode() .getField() .getName()); assertEquals( "b", conditionProto .getRoot() .getFnNode() .getInArgsList() .get(1) .getFieldNode() .getField() .getName()); } @Test public void testCondition2() throws GandivaException { Field a = Field.nullable("a", new ArrowType.Int(64, false)); Field b = Field.nullable("b", new ArrowType.Int(64, false)); Condition condition = TreeBuilder.makeCondition("greater_than", Arrays.asList(a, b)); GandivaTypes.Condition conditionProto = condition.toProtobuf(); assertTrue(conditionProto.getRoot().hasFnNode()); assertEquals("greater_than", conditionProto.getRoot().getFnNode().getFunctionName()); assertEquals( "a", conditionProto .getRoot() .getFnNode() .getInArgsList() .get(0) .getFieldNode() .getField() .getName()); assertEquals( "b", conditionProto .getRoot() .getFnNode() .getInArgsList() .get(1) .getFieldNode() .getField() .getName()); } } apache-arrow-java-19.0.0/gandiva/src/test/resources/000077500000000000000000000000001515451615200222675ustar00rootroot00000000000000apache-arrow-java-19.0.0/gandiva/src/test/resources/logback.xml000066400000000000000000000025221515451615200244140ustar00rootroot00000000000000 %d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n apache-arrow-java-19.0.0/memory/000077500000000000000000000000001515451615200164065ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/000077500000000000000000000000001515451615200206445ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/pom.xml000066400000000000000000000106261515451615200221660ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-memory 19.0.0 arrow-memory-core Arrow Memory - Core Core off-heap memory management libraries for Arrow ValueVectors. --add-reads=org.apache.arrow.memory.core=ch.qos.logback.classic --add-opens=java.base/java.lang.reflect=org.apache.arrow.memory.core --add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED org.slf4j slf4j-api org.immutables value-annotations org.checkerframework checker-qual com.google.errorprone error_prone_annotations org.apache.maven.plugins maven-compiler-plugin -Xmaxerrs 10000 -Xmaxwarns 10000 -AskipDefs=.*Test -AatfDoNotCache org.checkerframework checker ${checker.framework.version} org.apache.maven.plugins maven-surefire-plugin --add-reads=org.apache.arrow.memory.core=ch.qos.logback.classic --add-opens=java.base/java.lang.reflect=org.apache.arrow.memory.core --add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED **/TestOpens.java opens-tests test test **/TestOpens.java apache-arrow-java-19.0.0/memory/memory-core/src/000077500000000000000000000000001515451615200214335ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/000077500000000000000000000000001515451615200223575ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/000077500000000000000000000000001515451615200233005ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/module-info.java000066400000000000000000000024051515451615200263620ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ module org.apache.arrow.memory.core { exports org.apache.arrow.memory; exports org.apache.arrow.memory.rounding; exports org.apache.arrow.memory.util; exports org.apache.arrow.memory.util.hash; exports org.apache.arrow.util; requires java.compiler; requires transitive jdk.unsupported; requires static org.checkerframework.checker.qual; requires static org.immutables.value.annotations; requires static com.google.errorprone.annotations; requires org.slf4j; } apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/000077500000000000000000000000001515451615200240675ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/000077500000000000000000000000001515451615200253105ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200264425ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/000077500000000000000000000000001515451615200277525ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java000066400000000000000000000271661515451615200327300ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.concurrent.atomic.AtomicLongFieldUpdater; import org.apache.arrow.util.Preconditions; import org.checkerframework.checker.nullness.qual.Nullable; /** * Provides a concurrent way to manage account for memory usage without locking. Used as basis for * Allocators. All operations are threadsafe (except for close). */ class Accountant implements AutoCloseable { /** The parent allocator. */ protected final @Nullable Accountant parent; private final String name; /** * The amount of memory reserved for this allocator. Releases below this amount of memory will not * be returned to the parent Accountant until this Accountant is closed. */ protected final long reservation; // 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 volatile long allocationLimit = 0; /** Currently allocated amount of memory. */ private volatile long locallyHeldMemory = 0; public Accountant( @Nullable Accountant parent, String name, long reservation, long maxAllocation) { Preconditions.checkNotNull(name, "name must not be null"); Preconditions.checkArgument( reservation >= 0, "The initial reservation size must be non-negative."); Preconditions.checkArgument( maxAllocation >= 0, "The maximum allocation limit must be non-negative."); Preconditions.checkArgument( reservation <= maxAllocation, "The initial reservation size must be <= the maximum allocation."); Preconditions.checkArgument( reservation == 0 || parent != null, "The root accountant can't reserve memory."); this.parent = parent; this.name = name; this.reservation = reservation; ALLOCATION_LIMIT_UPDATER.set(this, maxAllocation); if (reservation != 0) { Preconditions.checkArgument(parent != null, "parent must not be null"); // we will allocate a reservation from our parent. final AllocationOutcome outcome = parent.allocateBytes(reservation); if (!outcome.isOk()) { throw new OutOfMemoryException( String.format( "Failure trying to allocate initial reservation for Allocator. " + "Attempted to allocate %d bytes.", reservation), outcome.getDetails()); } } } /** * Attempt to allocate the requested amount of memory. Either completely succeeds or completely * fails. If it fails, no changes are made to accounting. * * @param size The amount of memory to reserve in bytes. * @return the status and details of allocation at each allocator in the chain. */ AllocationOutcome allocateBytes(long size) { AllocationOutcome.Status status = allocateBytesInternal(size); if (status.isOk()) { return AllocationOutcome.SUCCESS_INSTANCE; } else { // Try again, but with details this time. // Populating details only on failures avoids performance overhead in the common case (success // case). AllocationOutcomeDetails details = new AllocationOutcomeDetails(); status = allocateBytesInternal(size, details); return new AllocationOutcome(status, details); } } private AllocationOutcome.Status allocateBytesInternal( long size, @Nullable AllocationOutcomeDetails details) { final AllocationOutcome.Status status = allocate(size, true /*incomingUpdatePeek*/, false /*forceAllocation*/, details); if (!status.isOk()) { releaseBytes(size); } return status; } private AllocationOutcome.Status allocateBytesInternal(long size) { return allocateBytesInternal(size, null /*details*/); } private void updatePeak() { final long currentMemory = locallyHeldMemory; while (true) { final long previousPeak = peakAllocation; if (currentMemory > previousPeak) { if (!PEAK_ALLOCATION_UPDATER.compareAndSet(this, previousPeak, currentMemory)) { // peak allocation changed underneath us. try again. continue; } } // we either succeeded to set peak allocation or we weren't above the previous peak, exit. return; } } /** * Increase the accounting. Returns whether the allocation fit within limits. * * @param size to increase * @return Whether the allocation fit within limits. */ public boolean forceAllocate(long size) { final AllocationOutcome.Status outcome = allocate(size, true, true, null); return outcome.isOk(); } /** * Internal method for allocation. This takes a forced approach to allocation to ensure that we * manage reservation boundary issues consistently. Allocation is always done through the entire * tree. The two options that we influence are whether the allocation should be forced and whether * or not the peak memory allocation should be updated. If at some point during allocation * escalation we determine that the allocation is no longer possible, we will continue to do a * complete and consistent allocation but we will stop updating the peak allocation. We do this * because we know that we will be directly unwinding this allocation (and thus never actually * making the allocation). If force allocation is passed, then we continue to update the peak * limits since we now know that this allocation will occur despite our moving past one or more * limits. * * @param size The size of the allocation. * @param incomingUpdatePeak Whether we should update the local peak for this allocation. * @param forceAllocation Whether we should force the allocation. * @return The outcome of the allocation. */ private AllocationOutcome.Status allocate( final long size, final boolean incomingUpdatePeak, final boolean forceAllocation, @Nullable AllocationOutcomeDetails details) { 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 // No need to reset locallyHeldMemory on overflow; allocateBytesInternal will releaseBytes on // failure final boolean overflow = ((oldLocal ^ newLocal) & (size ^ newLocal)) < 0; final long beyondReservation = newLocal - reservation; final boolean beyondLimit = overflow || newLocal > allocationLimit; final boolean updatePeak = forceAllocation || (incomingUpdatePeak && !beyondLimit); if (details != null) { // Add details if required (used in exceptions and debugging). boolean allocationFailed = true; long allocatedLocal = 0; if (!beyondLimit) { allocatedLocal = size - Math.min(beyondReservation, size); allocationFailed = false; } details.pushEntry(this, newLocal - size, size, allocatedLocal, allocationFailed); } AllocationOutcome.Status parentOutcome = AllocationOutcome.Status.SUCCESS; if (beyondReservation > 0 && parent != null) { // we need to get memory from our parent. final long parentRequest = Math.min(beyondReservation, size); parentOutcome = parent.allocate(parentRequest, updatePeak, forceAllocation, details); } final AllocationOutcome.Status finalOutcome; if (beyondLimit) { finalOutcome = AllocationOutcome.Status.FAILED_LOCAL; } else { finalOutcome = parentOutcome.isOk() ? AllocationOutcome.Status.SUCCESS : AllocationOutcome.Status.FAILED_PARENT; } if (updatePeak) { updatePeak(); } return finalOutcome; } public void releaseBytes(long size) { // reduce local memory. all memory released above reservation should be released up the tree. final long newSize = LOCALLY_HELD_MEMORY_UPDATER.addAndGet(this, -size); Preconditions.checkArgument(newSize >= 0, "Accounted size went negative."); final long originalSize = newSize + size; if (originalSize > reservation && parent != null) { // we deallocated memory that we should release to our parent. final long possibleAmountToReleaseToParent = originalSize - reservation; final long actualToReleaseToParent = Math.min(size, possibleAmountToReleaseToParent); parent.releaseBytes(actualToReleaseToParent); } } public boolean isOverLimit() { return getAllocatedMemory() > getLimit() || (parent != null && parent.isOverLimit()); } /** Close this Accountant. This will release any reservation bytes back to a parent Accountant. */ @Override public void close() { // return memory reservation to parent allocator. if (parent != null) { parent.releaseBytes(reservation); } } /** * Return the name of the accountant. * * @return name of accountant */ public String getName() { return name; } /** * Return the current limit of this Accountant. * * @return Limit in bytes. */ public long getLimit() { return allocationLimit; } /** * Return the initial reservation. * * @return reservation in bytes. */ public long getInitReservation() { return reservation; } /** * Set the maximum amount of memory that can be allocated in the this Accountant before failing an * allocation. * * @param newLimit The limit in bytes. */ public void setLimit(long newLimit) { ALLOCATION_LIMIT_UPDATER.set(this, newLimit); } /** * Return the current amount of allocated memory that this Accountant is managing accounting for. * Note this does not include reservation memory that hasn't been allocated. * * @return Currently allocate memory in bytes. */ public long getAllocatedMemory() { return locallyHeldMemory; } /** * The peak memory allocated by this Accountant. * * @return The peak allocated memory in bytes. */ public long getPeakMemoryAllocation() { return peakAllocation; } public long getHeadroom() { 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); return Math.min(localHeadroom, parent.getHeadroom() + reservedHeadroom); } } AllocationListener.java000066400000000000000000000062071515451615200343360ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * An allocation listener being notified for allocation/deallocation * *

    It might be called from multiple threads if the allocator hierarchy shares a listener, in * which case, the provider should take care of making the implementation thread-safe. */ public interface AllocationListener { AllocationListener NOOP = new AllocationListener() {}; /** * Called each time a new buffer has been requested. * *

    An exception can be safely thrown by this method to terminate the allocation. * * @param size the buffer size being allocated */ default void onPreAllocation(long size) {} /** * Called each time a new buffer has been allocated. * *

    An exception cannot be thrown by this method. * * @param size the buffer size being allocated */ default void onAllocation(long size) {} /** * Informed each time a buffer is released from allocation. * *

    An exception cannot be thrown by this method. * * @param size The size of the buffer being released. */ default void onRelease(long size) {} /** * Called whenever an allocation failed, giving the caller a chance to create some space in the * allocator (either by freeing some resource, or by changing the limit), and, if successful, * allowing the allocator to retry the allocation. * * @param size the buffer size that was being allocated * @param outcome the outcome of the failed allocation. Carries information of what failed * @return true, if the allocation can be retried; false if the allocation should fail */ default boolean onFailedAllocation(long size, AllocationOutcome outcome) { return false; } /** * Called immediately after a child allocator was added to the parent allocator. * * @param parentAllocator The parent allocator to which a child was added * @param childAllocator The child allocator that was just added */ default void onChildAdded(BufferAllocator parentAllocator, BufferAllocator childAllocator) {} /** * Called immediately after a child allocator was removed from the parent allocator. * * @param parentAllocator The parent allocator from which a child was removed * @param childAllocator The child allocator that was just removed */ default void onChildRemoved(BufferAllocator parentAllocator, BufferAllocator childAllocator) {} } AllocationManager.java000066400000000000000000000234051515451615200341220ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.util.Preconditions; import org.checkerframework.checker.nullness.qual.Nullable; /** * An AllocationManager is the implementation of a physical memory allocation. * *

    Manages the relationship between the allocators and a particular memory allocation. Ensures * that one allocator owns the memory that multiple allocators may be referencing. Manages a * BufferLedger between each of its associated allocators. It does not track the reference count; * that is the role of {@link BufferLedger} (aka {@link ReferenceManager}). * *

    This is a public interface implemented by concrete allocator implementations (e.g. Netty or * Unsafe). * *

    Threading: AllocationManager manages thread-safety internally. Operations within the context * of a single BufferLedger are lockless in nature and can be leveraged by multiple threads. * Operations that cross the context of two ledgers will acquire a lock on the AllocationManager * instance. Important note, there is one AllocationManager per physical buffer allocation. As such, * there will be thousands of these in a typical query. The contention of acquiring a lock on * AllocationManager should be very low. */ public abstract class AllocationManager { // The RootAllocator we are associated with. An allocation can only ever be associated with a // single RootAllocator. private final BufferAllocator root; // An allocation can be tracked by multiple allocators. (This is because an allocator is more like // a ledger.) // All such allocators track reference counts individually, via BufferLedger instances. When an // individual // reference count reaches zero, the allocator will be dissociated from this allocation. If that // was via the // owningLedger, then no more allocators should be tracking this allocation, and the allocation // will be freed. // ARROW-1627: Trying to minimize memory overhead caused by previously used IdentityHashMap private final LowCostIdentityHashMap map = new LowCostIdentityHashMap<>(); // The primary BufferLedger (i.e. reference count) tracking this allocation. // This is mostly a semantic constraint on the API user: if the reference count reaches 0 in the // owningLedger, then // there are not supposed to be any references through other allocators. In practice, this doesn't // do anything // as the implementation just forces ownership to be transferred to one of the other extant // references. private volatile @Nullable BufferLedger owningLedger; @SuppressWarnings( "nullness:method.invocation") // call to associate(a, b) not allowed on the given receiver protected AllocationManager(BufferAllocator accountingAllocator) { Preconditions.checkNotNull(accountingAllocator); accountingAllocator.assertOpen(); this.root = accountingAllocator.getRoot(); // we do a no retain association since our creator will want to retrieve the newly created // ledger and will create a reference count at that point this.owningLedger = associate(accountingAllocator, false); } @Nullable BufferLedger getOwningLedger() { return owningLedger; } void setOwningLedger(final BufferLedger ledger) { this.owningLedger = ledger; } /** * Associate the existing underlying buffer with a new allocator. This will increase the reference * count on the corresponding buffer ledger by 1. * * @param allocator The target allocator to associate this buffer with. * @return The reference manager (new or existing) that associates the underlying buffer to this * new ledger. */ BufferLedger associate(final BufferAllocator allocator) { return associate(allocator, true); } private BufferLedger associate(final BufferAllocator allocator, final boolean retain) { allocator.assertOpen(); Preconditions.checkState( root == allocator.getRoot(), "A buffer can only be associated between two allocators that share the same root"); synchronized (this) { BufferLedger ledger = map.get(allocator); if (ledger != null) { // We were already being tracked by the given allocator, just return it if (retain) { // bump the ref count for the ledger ledger.increment(); } return ledger; } // We weren't previously being tracked by the given allocator; create a new ledger ledger = new BufferLedger(allocator, this); if (retain) { // the new reference manager will have a ref count of 1 ledger.increment(); } // store the mapping for BufferLedger oldLedger = map.put(ledger); Preconditions.checkState( oldLedger == null, "Detected inconsistent state: A reference manager already exists for this allocator"); if (allocator instanceof BaseAllocator) { // needed for debugging only: keep a pointer to reference manager inside allocator // to dump state, verify allocator state etc ((BaseAllocator) allocator).associateLedger(ledger); } return ledger; } } /** * The way that a particular ReferenceManager (BufferLedger) communicates back to the * AllocationManager that it no longer needs to hold a reference to a particular piece of memory. * Reference manager needs to hold a lock to invoke this method It is called when the shared * refcount of all the ArrowBufs managed by the calling ReferenceManager drops to 0. */ void release(final BufferLedger ledger) { final BufferAllocator allocator = ledger.getAllocator(); allocator.assertOpen(); // remove the mapping for the allocator // of calling BufferLedger Preconditions.checkState( map.containsKey(allocator), "Expecting a mapping for allocator and reference manager"); final BufferLedger oldLedger = map.remove(allocator); Preconditions.checkState( oldLedger != null, "Expecting a mapping for allocator and reference manager"); BufferAllocator oldAllocator = oldLedger.getAllocator(); if (oldAllocator instanceof BaseAllocator) { // needed for debug only: tell the allocator that AllocationManager is removing a // reference manager associated with this particular allocator ((BaseAllocator) oldAllocator).dissociateLedger(oldLedger); } if (oldLedger == owningLedger) { // the release call was made by the owning reference manager if (map.isEmpty()) { // the only mapping was for the owner // which now has been removed, it implies we can safely destroy the // underlying memory chunk as it is no longer being referenced oldAllocator.releaseBytes(getSize()); // free the memory chunk associated with the allocation manager release0(); oldAllocator.getListener().onRelease(getSize()); owningLedger = null; } else { // since the refcount dropped to 0 for the owning reference manager and allocation // manager will no longer keep a mapping for it, we need to change the owning // reference manager to whatever the next available // mapping exists. BufferLedger newOwningLedger = map.getNextValue(); // we'll forcefully transfer the ownership and not worry about whether we // exceeded the limit since this consumer can't do anything with this. oldLedger.transferBalance(newOwningLedger); } } else { // the release call was made by a non-owning reference manager, so after remove there have // to be 1 or more mappings Preconditions.checkState( map.size() > 0, "The final removal of reference manager should be connected to owning reference manager"); } } /** * Return the size of underlying chunk of memory managed by this Allocation Manager. * *

    The underlying memory chunk managed can be different from the original requested size. * * @return size of underlying memory chunk */ public abstract long getSize(); /** Return the absolute memory address pointing to the fist byte of underlying memory chunk. */ protected abstract long memoryAddress(); /** Release the underlying memory chunk. */ protected abstract void release0(); /** * A factory interface for creating {@link AllocationManager}. One may extend this interface to * use a user-defined AllocationManager implementation. */ public interface Factory { /** * Create an {@link AllocationManager}. * * @param accountingAllocator The allocator that are expected to be associated with newly * created AllocationManager. Currently it is always equivalent to "this" * @param size Size (in bytes) of memory managed by the AllocationManager * @return The created AllocationManager used by this allocator */ AllocationManager create(BufferAllocator accountingAllocator, long size); ArrowBuf empty(); } } AllocationOutcome.java000066400000000000000000000050701515451615200341610ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.Optional; import org.checkerframework.checker.nullness.qual.Nullable; /** Describes the type of outcome that occurred when trying to account for allocation of memory. */ public class AllocationOutcome { private final Status status; private final @Nullable AllocationOutcomeDetails details; static final AllocationOutcome SUCCESS_INSTANCE = new AllocationOutcome(Status.SUCCESS); AllocationOutcome(Status status, @Nullable AllocationOutcomeDetails details) { this.status = status; this.details = details; } AllocationOutcome(Status status) { this(status, null); } /** * Get the status of the allocation. * * @return status code. */ public Status getStatus() { return status; } /** * Get additional details of the allocation (like the status at each allocator in the hierarchy). * * @return details of allocation */ public Optional getDetails() { return Optional.ofNullable(details); } /** * Returns true if the allocation was a success. * * @return true if allocation was successful, false otherwise. */ public boolean isOk() { return status.isOk(); } /** Allocation status code. */ public enum Status { /** Allocation succeeded. */ SUCCESS(true), /** Allocation succeeded but only because the allocator was forced to move beyond a limit. */ FORCED_SUCCESS(true), /** Allocation failed because the local allocator's limits were exceeded. */ FAILED_LOCAL(false), /** Allocation failed because a parent allocator's limits were exceeded. */ FAILED_PARENT(false); private final boolean ok; Status(boolean ok) { this.ok = ok; } public boolean isOk() { return ok; } } } AllocationOutcomeDetails.java000066400000000000000000000077651515451615200355040ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.ArrayDeque; import java.util.Deque; import org.checkerframework.checker.nullness.qual.Nullable; /** Captures details of allocation for each accountant in the hierarchical chain. */ public class AllocationOutcomeDetails { Deque allocEntries; AllocationOutcomeDetails() { allocEntries = new ArrayDeque<>(); } void pushEntry( Accountant accountant, long totalUsedBeforeAllocation, long requestedSize, long allocatedSize, boolean allocationFailed) { Entry top = allocEntries.peekLast(); if (top != null && top.allocationFailed) { // if the allocation has already failed, stop saving the entries. return; } allocEntries.addLast( new Entry( accountant, totalUsedBeforeAllocation, requestedSize, allocatedSize, allocationFailed)); } /** * Get the allocator that caused the failure. * * @return the allocator that caused failure, null if there was no failure. */ public @Nullable BufferAllocator getFailedAllocator() { Entry top = allocEntries.peekLast(); if (top != null && top.allocationFailed && (top.accountant instanceof BufferAllocator)) { return (BufferAllocator) top.accountant; } else { return null; } } @Override public String toString() { StringBuilder sb = new StringBuilder(); sb.append("Allocation outcome details:\n"); allocEntries.forEach(sb::append); return sb.toString(); } /** Outcome of the allocation request at one accountant in the hierarchy. */ public static class Entry { private final Accountant accountant; // Remember allocator attributes at the time of the request. private final long limit; private final long used; // allocation outcome private final long requestedSize; private final long allocatedSize; private final boolean allocationFailed; Entry( Accountant accountant, long totalUsedBeforeAllocation, long requestedSize, long allocatedSize, boolean allocationFailed) { this.accountant = accountant; this.limit = accountant.getLimit(); this.used = totalUsedBeforeAllocation; this.requestedSize = requestedSize; this.allocatedSize = allocatedSize; this.allocationFailed = allocationFailed; } public Accountant getAccountant() { return accountant; } public long getLimit() { return limit; } public long getUsed() { return used; } public long getRequestedSize() { return requestedSize; } public long getAllocatedSize() { return allocatedSize; } public boolean isAllocationFailed() { return allocationFailed; } @Override public String toString() { return "allocator[" + accountant.getName() + "]" + " reservation: " + accountant.getInitReservation() + " limit: " + limit + " used: " + used + " requestedSize: " + requestedSize + " allocatedSize: " + allocatedSize + " localAllocationStatus: " + (allocationFailed ? "fail" : "success") + "\n"; } } } AllocationReservation.java000066400000000000000000000074641515451615200350600ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * Supports cumulative allocation reservation. Clients may increase the size of the reservation * repeatedly until they call for an allocation of the current total size. The reservation can only * be used once, and will throw an exception if it is used more than once. * *

    For the purposes of airtight memory accounting, the reservation must be close()d whether it is * used or not. This is not threadsafe. */ public interface AllocationReservation extends AutoCloseable { /** * Add to the current reservation. * *

    Adding may fail if the allocator is not allowed to consume any more space. * * @param nBytes the number of bytes to add * @return true if the addition is possible, false otherwise * @throws IllegalStateException if called after buffer() is used to allocate the reservation * @deprecated use {@link #add(long)} instead */ @Deprecated(forRemoval = true) boolean add(int nBytes); /** * Add to the current reservation. * *

    Adding may fail if the allocator is not allowed to consume any more space. * * @param nBytes the number of bytes to add * @return true if the addition is possible, false otherwise * @throws IllegalStateException if called after buffer() is used to allocate the reservation */ boolean add(long nBytes); /** * Requests a reservation of additional space. * *

    The implementation of the allocator's inner class provides this. * * @param nBytes the amount to reserve * @return true if the reservation can be satisfied, false otherwise * @deprecated use {@link #reserve(long)} instead */ @Deprecated(forRemoval = true) boolean reserve(int nBytes); /** * Requests a reservation of additional space. * *

    The implementation of the allocator's inner class provides this. * * @param nBytes the amount to reserve * @return true if the reservation can be satisfied, false otherwise */ boolean reserve(long nBytes); /** * Allocate a buffer whose size is the total of all the add()s made. * *

    The allocation request can still fail, even if the amount of space requested is available, * if the allocation cannot be made contiguously. * * @return the buffer, or null, if the request cannot be satisfied * @throws IllegalStateException if called more than once */ ArrowBuf allocateBuffer(); /** * Get the current size of the reservation (the sum of all the add()s). * * @return size of the current reservation */ int getSize(); /** * Get the current size of the reservation (the sum of all the add()s) as a long value. * * @return size of the current reservation */ long getSizeLong(); /** * Return whether or not the reservation has been used. * * @return whether or not the reservation has been used */ boolean isUsed(); /** * Return whether or not the reservation has been closed. * * @return whether or not the reservation has been closed */ boolean isClosed(); @Override void close(); } AllocatorClosedException.java000066400000000000000000000023401515451615200354660ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** Exception thrown when a closed BufferAllocator is used. Note this is an unchecked exception. */ @SuppressWarnings("serial") public class AllocatorClosedException extends RuntimeException { /** * Constructs a new allocator closed exception with a given message. * * @param message string associated with the cause */ public AllocatorClosedException(String message) { super(message); } } apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java000066400000000000000000001237761515451615200323640ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.nio.ReadOnlyBufferException; import org.apache.arrow.memory.BaseAllocator.Verbosity; import org.apache.arrow.memory.util.CommonUtil; import org.apache.arrow.memory.util.HistoricalLog; import org.apache.arrow.memory.util.MemoryUtil; import org.apache.arrow.util.Preconditions; import org.apache.arrow.util.VisibleForTesting; import org.checkerframework.checker.nullness.qual.Nullable; /** * ArrowBuf serves as a facade over underlying memory by providing several access APIs to read/write * data into a chunk of direct memory. All the accounting, ownership and reference management is * done by {@link ReferenceManager} and ArrowBuf can work with a custom user provided implementation * of ReferenceManager * *

    Two important instance variables of an ArrowBuf: (1) address - starting virtual address in the * underlying memory chunk that this ArrowBuf has access to (2) length - length (in bytes) in the * underlying memory chunk that this ArrowBuf has access to * *

    The management (allocation, deallocation, reference counting etc) for the memory chunk is not * done by ArrowBuf. Default implementation of ReferenceManager, allocation is in {@link * BaseAllocator}, {@link BufferLedger} and {@link AllocationManager} */ public final class ArrowBuf implements AutoCloseable { private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(ArrowBuf.class); private static final int SHORT_SIZE = Short.BYTES; private static final int INT_SIZE = Integer.BYTES; private static final int FLOAT_SIZE = Float.BYTES; private static final int DOUBLE_SIZE = Double.BYTES; private static final int LONG_SIZE = Long.BYTES; private static final int LOG_BYTES_PER_ROW = 10; private final ReferenceManager referenceManager; private final @Nullable BufferManager bufferManager; private final long addr; private long readerIndex; private long writerIndex; private final @Nullable HistoricalLog historicalLog = BaseAllocator.DEBUG ? new HistoricalLog( BaseAllocator.DEBUG_LOG_LENGTH, "ArrowBuf[%d]", System.identityHashCode(this)) : null; private volatile long capacity; /** * Constructs a new ArrowBuf. * * @param referenceManager The memory manager to track memory usage and reference count of this * buffer * @param capacity The capacity in bytes of this buffer */ public ArrowBuf( final ReferenceManager referenceManager, final @Nullable BufferManager bufferManager, final long capacity, final long memoryAddress) { this.referenceManager = referenceManager; this.bufferManager = bufferManager; this.addr = memoryAddress; this.capacity = capacity; this.readerIndex = 0; this.writerIndex = 0; if (historicalLog != null) { historicalLog.recordEvent("create()"); } } public int refCnt() { return referenceManager.getRefCount(); } /** * Allows a function to determine whether not reading a particular string of bytes is valid. * *

    Will throw an exception if the memory is not readable for some reason. Only doesn't * something in the case that AssertionUtil.BOUNDS_CHECKING_ENABLED is true. * * @param start The starting position of the bytes to be read. * @param end The exclusive endpoint of the bytes to be read. */ public void checkBytes(long start, long end) { if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { checkIndexD(start, end - start); } } /** For get/set operations, reference count should be >= 1. */ private void ensureAccessible() { if (this.refCnt() == 0) { throw new IllegalStateException("Ref count should be >= 1 for accessing the ArrowBuf"); } } /** * Get reference manager for this ArrowBuf. * * @return user provided implementation of {@link ReferenceManager} */ public ReferenceManager getReferenceManager() { return referenceManager; } public long capacity() { return capacity; } /** * Adjusts the capacity of this buffer. Size increases are NOT supported. * * @param newCapacity Must be in the range [0, length). */ public synchronized ArrowBuf capacity(long newCapacity) { if (newCapacity == capacity) { return this; } Preconditions.checkArgument(newCapacity >= 0); if (newCapacity < capacity) { capacity = newCapacity; return this; } throw new UnsupportedOperationException( "Buffers don't support resizing that increases the size."); } /** Returns the byte order of elements in this buffer. */ public ByteOrder order() { return ByteOrder.nativeOrder(); } /** Returns the number of bytes still available to read in this buffer. */ public long readableBytes() { Preconditions.checkState( writerIndex >= readerIndex, "Writer index cannot be less than reader index"); return writerIndex - readerIndex; } /** * Returns the number of bytes still available to write into this buffer before capacity is * reached. */ public long writableBytes() { return capacity() - writerIndex; } /** Returns a slice of only the readable bytes in the buffer. */ public ArrowBuf slice() { return slice(readerIndex, readableBytes()); } /** Returns a slice (view) starting at index with the given length. */ public ArrowBuf slice(long index, long length) { Preconditions.checkPositionIndex(index, this.capacity); Preconditions.checkPositionIndex(index + length, this.capacity); /* * Re the behavior of reference counting, see http://netty.io/wiki/reference-counted-objects * .html#wiki-h3-5, which * explains that derived buffers share their reference count with their parent */ final ArrowBuf newBuf = referenceManager.deriveBuffer(this, index, length); newBuf.writerIndex(length); return newBuf; } /** Make a nio byte buffer from this arrowbuf. */ public ByteBuffer nioBuffer() { return nioBuffer(readerIndex, checkedCastToInt(readableBytes())); } /** Make a nio byte buffer from this ArrowBuf. */ public ByteBuffer nioBuffer(long index, int length) { chk(index, length); return getDirectBuffer(index, length); } private ByteBuffer getDirectBuffer(long index, int length) { long address = addr(index); return MemoryUtil.directBuffer(address, length); } public long memoryAddress() { return this.addr; } @Override public String toString() { return String.format( "ArrowBuf[%d], address:%d, capacity:%d", getId(), memoryAddress(), capacity); } @Override public int hashCode() { return System.identityHashCode(this); } @Override public boolean equals(@Nullable Object obj) { // identity equals only. return this == obj; } /* * IMPORTANT NOTE * The data getters and setters work with a caller provided * index. This index is 0 based and since ArrowBuf has access * to a portion of underlying chunk of memory starting at * some address, we convert the given relative index into * absolute index as memory address + index. * * Example: * * Let's say we have an underlying chunk of memory of length 64 bytes * Now let's say we have an ArrowBuf that has access to the chunk * from offset 4 for length of 16 bytes. * * If the starting virtual address of chunk is MAR, then memory * address of this ArrowBuf is MAR + offset -- this is what is stored * in variable addr. See the BufferLedger and AllocationManager code * for the implementation of ReferenceManager that manages a * chunk of memory and creates ArrowBuf with access to a range of * bytes within the chunk (or the entire chunk) * * So now to get/set data, we will do => addr + index * This logic is put in method addr(index) and is frequently * used in get/set data methods to compute the absolute * byte address for get/set operation in the underlying chunk * * @param index the index at which we the user wants to read/write * @return the absolute address within the memory */ private long addr(long index) { return addr + index; } /*-------------------------------------------------* | Following are a set of fast path data set and | | get APIs to write/read data from ArrowBuf | | at a given index (0 based relative to this | | ArrowBuf and not relative to the underlying | | memory chunk). | | | *-------------------------------------------------*/ /** * Helper function to do bounds checking at a particular index for particular length of data. * * @param index index (0 based relative to this ArrowBuf) * @param length provided length of data for get/set */ private void chk(long index, long length) { if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { checkIndexD(index, length); } } private void checkIndexD(long index, long fieldLength) { // check reference count ensureAccessible(); // check bounds Preconditions.checkArgument(fieldLength >= 0, "expecting non-negative data length"); if (index < 0 || index > capacity() - fieldLength) { if (historicalLog != null) { historicalLog.logHistory(logger); } throw new IndexOutOfBoundsException( String.format( "index: %d, length: %d (expected: range(0, %d))", index, fieldLength, capacity())); } } /** * Get long value stored at a particular index in the underlying memory chunk this ArrowBuf has * access to. * * @param index index (0 based relative to this ArrowBuf) where the value will be read from * @return 8 byte long value */ public long getLong(long index) { chk(index, LONG_SIZE); return MemoryUtil.getLong(addr(index)); } /** * Set long value at a particular index in the underlying memory chunk this ArrowBuf has access * to. * * @param index index (0 based relative to this ArrowBuf) where the value will be written * @param value value to write */ public void setLong(long index, long value) { chk(index, LONG_SIZE); MemoryUtil.putLong(addr(index), value); } /** * Get float value stored at a particular index in the underlying memory chunk this ArrowBuf has * access to. * * @param index index (0 based relative to this ArrowBuf) where the value will be read from * @return 4 byte float value */ public float getFloat(long index) { return Float.intBitsToFloat(getInt(index)); } /** * Set float value at a particular index in the underlying memory chunk this ArrowBuf has access * to. * * @param index index (0 based relative to this ArrowBuf) where the value will be written * @param value value to write */ public void setFloat(long index, float value) { chk(index, FLOAT_SIZE); MemoryUtil.putInt(addr(index), Float.floatToRawIntBits(value)); } /** * Get double value stored at a particular index in the underlying memory chunk this ArrowBuf has * access to. * * @param index index (0 based relative to this ArrowBuf) where the value will be read from * @return 8 byte double value */ public double getDouble(long index) { return Double.longBitsToDouble(getLong(index)); } /** * Set double value at a particular index in the underlying memory chunk this ArrowBuf has access * to. * * @param index index (0 based relative to this ArrowBuf) where the value will be written * @param value value to write */ public void setDouble(long index, double value) { chk(index, DOUBLE_SIZE); MemoryUtil.putLong(addr(index), Double.doubleToRawLongBits(value)); } /** * Get char value stored at a particular index in the underlying memory chunk this ArrowBuf has * access to. * * @param index index (0 based relative to this ArrowBuf) where the value will be read from * @return 2 byte char value */ public char getChar(long index) { return (char) getShort(index); } /** * Set char value at a particular index in the underlying memory chunk this ArrowBuf has access * to. * * @param index index (0 based relative to this ArrowBuf) where the value will be written * @param value value to write */ public void setChar(long index, int value) { chk(index, SHORT_SIZE); MemoryUtil.putShort(addr(index), (short) value); } /** * Get int value stored at a particular index in the underlying memory chunk this ArrowBuf has * access to. * * @param index index (0 based relative to this ArrowBuf) where the value will be read from * @return 4 byte int value */ public int getInt(long index) { chk(index, INT_SIZE); return MemoryUtil.getInt(addr(index)); } /** * Set int value at a particular index in the underlying memory chunk this ArrowBuf has access to. * * @param index index (0 based relative to this ArrowBuf) where the value will be written * @param value value to write */ public void setInt(long index, int value) { chk(index, INT_SIZE); MemoryUtil.putInt(addr(index), value); } /** * Get short value stored at a particular index in the underlying memory chunk this ArrowBuf has * access to. * * @param index index (0 based relative to this ArrowBuf) where the value will be read from * @return 2 byte short value */ public short getShort(long index) { chk(index, SHORT_SIZE); return MemoryUtil.getShort(addr(index)); } /** * Set short value at a particular index in the underlying memory chunk this ArrowBuf has access * to. * * @param index index (0 based relative to this ArrowBuf) where the value will be written * @param value value to write */ public void setShort(long index, int value) { setShort(index, (short) value); } /** * Set short value at a particular index in the underlying memory chunk this ArrowBuf has access * to. * * @param index index (0 based relative to this ArrowBuf) where the value will be written * @param value value to write */ public void setShort(long index, short value) { chk(index, SHORT_SIZE); MemoryUtil.putShort(addr(index), value); } /** * Set byte value at a particular index in the underlying memory chunk this ArrowBuf has access * to. * * @param index index (0 based relative to this ArrowBuf) where the value will be written * @param value value to write */ public void setByte(long index, int value) { chk(index, 1); MemoryUtil.putByte(addr(index), (byte) value); } /** * Set byte value at a particular index in the underlying memory chunk this ArrowBuf has access * to. * * @param index index (0 based relative to this ArrowBuf) where the value will be written * @param value value to write */ public void setByte(long index, byte value) { chk(index, 1); MemoryUtil.putByte(addr(index), value); } /** * Get byte value stored at a particular index in the underlying memory chunk this ArrowBuf has * access to. * * @param index index (0 based relative to this ArrowBuf) where the value will be read from * @return byte value */ public byte getByte(long index) { chk(index, 1); return MemoryUtil.getByte(addr(index)); } /*--------------------------------------------------* | Following are another set of data set APIs | | that directly work with writerIndex | | | *--------------------------------------------------*/ /** * Helper function to do bound checking w.r.t writerIndex by checking if we can set "length" bytes * of data at the writerIndex in this ArrowBuf. * * @param length provided length of data for set */ private void ensureWritable(final int length) { if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { Preconditions.checkArgument(length >= 0, "expecting non-negative length"); // check reference count this.ensureAccessible(); // check bounds if (length > writableBytes()) { throw new IndexOutOfBoundsException( String.format( "writerIndex(%d) + length(%d) exceeds capacity(%d)", writerIndex, length, capacity())); } } } /** * Helper function to do bound checking w.r.t readerIndex by checking if we can read "length" * bytes of data at the readerIndex in this ArrowBuf. * * @param length provided length of data for get */ private void ensureReadable(final int length) { if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { Preconditions.checkArgument(length >= 0, "expecting non-negative length"); // check reference count this.ensureAccessible(); // check bounds if (length > readableBytes()) { throw new IndexOutOfBoundsException( String.format( "readerIndex(%d) + length(%d) exceeds writerIndex(%d)", readerIndex, length, writerIndex)); } } } /** * Read the byte at readerIndex. * * @return byte value */ public byte readByte() { ensureReadable(1); final byte b = getByte(readerIndex); ++readerIndex; return b; } /** * Read dst.length bytes at readerIndex into dst byte array. * * @param dst byte array where the data will be written */ public void readBytes(byte[] dst) { Preconditions.checkArgument(dst != null, "expecting valid dst bytearray"); ensureReadable(dst.length); getBytes(readerIndex, dst, 0, checkedCastToInt(dst.length)); } /** * Set the provided byte value at the writerIndex. * * @param value value to set */ public void writeByte(byte value) { ensureWritable(1); MemoryUtil.putByte(addr(writerIndex), value); ++writerIndex; } /** * Set the lower order byte for the provided value at the writerIndex. * * @param value value to be set */ public void writeByte(int value) { ensureWritable(1); MemoryUtil.putByte(addr(writerIndex), (byte) value); ++writerIndex; } /** * Write the bytes from given byte array into this ArrowBuf starting at writerIndex. * * @param src src byte array */ public void writeBytes(byte[] src) { Preconditions.checkArgument(src != null, "expecting valid src array"); writeBytes(src, 0, src.length); } /** * Write the bytes from given byte array starting at srcIndex into this ArrowBuf starting at * writerIndex. * * @param src src byte array * @param srcIndex index in the byte array where the copy will being from * @param length length of data to copy */ public void writeBytes(byte[] src, int srcIndex, int length) { ensureWritable(length); setBytes(writerIndex, src, srcIndex, length); writerIndex += length; } /** * Set the provided int value as short at the writerIndex. * * @param value value to set */ public void writeShort(int value) { ensureWritable(SHORT_SIZE); MemoryUtil.putShort(addr(writerIndex), (short) value); writerIndex += SHORT_SIZE; } /** * Set the provided int value at the writerIndex. * * @param value value to set */ public void writeInt(int value) { ensureWritable(INT_SIZE); MemoryUtil.putInt(addr(writerIndex), value); writerIndex += INT_SIZE; } /** * Set the provided long value at the writerIndex. * * @param value value to set */ public void writeLong(long value) { ensureWritable(LONG_SIZE); MemoryUtil.putLong(addr(writerIndex), value); writerIndex += LONG_SIZE; } /** * Set the provided float value at the writerIndex. * * @param value value to set */ public void writeFloat(float value) { ensureWritable(FLOAT_SIZE); MemoryUtil.putInt(addr(writerIndex), Float.floatToRawIntBits(value)); writerIndex += FLOAT_SIZE; } /** * Set the provided double value at the writerIndex. * * @param value value to set */ public void writeDouble(double value) { ensureWritable(DOUBLE_SIZE); MemoryUtil.putLong(addr(writerIndex), Double.doubleToRawLongBits(value)); writerIndex += DOUBLE_SIZE; } /*--------------------------------------------------* | Following are another set of data set/get APIs | | that read and write stream of bytes from/to byte | | arrays, ByteBuffer, ArrowBuf etc | | | *--------------------------------------------------*/ /** * Determine if the requested {@code index} and {@code length} will fit within {@code capacity}. * * @param index The starting index. * @param length The length which will be utilized (starting from {@code index}). * @param capacity The capacity that {@code index + length} is allowed to be within. * @return {@code true} if the requested {@code index} and {@code length} will fit within {@code * capacity}. {@code false} if this would result in an index out of bounds exception. */ private static boolean isOutOfBounds(long index, long length, long capacity) { return (index | length | (index + length) | (capacity - (index + length))) < 0; } private void checkIndex(long index, long fieldLength) { if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { // check reference count this.ensureAccessible(); // check bounds if (isOutOfBounds(index, fieldLength, this.capacity())) { throw new IndexOutOfBoundsException( String.format( "index: %d, length: %d (expected: range(0, %d))", index, fieldLength, this.capacity())); } } } /** * Copy data from this ArrowBuf at a given index in into destination byte array. * * @param index starting index (0 based relative to the portion of memory) this ArrowBuf has * access to * @param dst byte array to copy the data into */ public void getBytes(long index, byte[] dst) { getBytes(index, dst, 0, dst.length); } /** * Copy data from this ArrowBuf at a given index into destination byte array. * * @param index index (0 based relative to the portion of memory this ArrowBuf has access to) * @param dst byte array to copy the data into * @param dstIndex starting index in dst byte array to copy into * @param length length of data to copy from this ArrowBuf */ public void getBytes(long index, byte[] dst, int dstIndex, int length) { // bound check for this ArrowBuf where the data will be copied from checkIndex(index, length); // null check Preconditions.checkArgument(dst != null, "expecting a valid dst byte array"); // bound check for dst byte array where the data will be copied to if (isOutOfBounds(dstIndex, length, dst.length)) { // not enough space to copy "length" bytes into dst array from dstIndex onwards throw new IndexOutOfBoundsException( "Not enough space to copy data into destination" + dstIndex); } if (length != 0) { // copy "length" bytes from this ArrowBuf starting at addr(index) address // into dst byte array at dstIndex onwards MemoryUtil.copyFromMemory(addr(index), dst, dstIndex, length); } } /** * Copy data from a given byte array into this ArrowBuf starting at a given index. * * @param index starting index (0 based relative to the portion of memory) this ArrowBuf has * access to * @param src byte array to copy the data from */ public void setBytes(long index, byte[] src) { setBytes(index, src, 0, src.length); } /** * Copy data from a given byte array starting at the given source index into this ArrowBuf at a * given index. * * @param index index (0 based relative to the portion of memory this ArrowBuf has access to) * @param src src byte array to copy the data from * @param srcIndex index in the byte array where the copy will start from * @param length length of data to copy from byte array */ public void setBytes(long index, byte[] src, int srcIndex, long length) { // bound check for this ArrowBuf where the data will be copied into checkIndex(index, length); // null check Preconditions.checkArgument(src != null, "expecting a valid src byte array"); // bound check for src byte array where the data will be copied from if (isOutOfBounds(srcIndex, length, src.length)) { // not enough space to copy "length" bytes into dst array from dstIndex onwards throw new IndexOutOfBoundsException( "Not enough space to copy data from byte array" + srcIndex); } if (length > 0) { // copy "length" bytes from src byte array at the starting index (srcIndex) // into this ArrowBuf starting at address "addr(index)" MemoryUtil.copyToMemory(src, srcIndex, addr(index), length); } } /** * Copy data from this ArrowBuf at a given index into the destination ByteBuffer. * * @param index index (0 based relative to the portion of memory this ArrowBuf has access to) * @param dst dst ByteBuffer where the data will be copied into */ public void getBytes(long index, ByteBuffer dst) { // bound check for this ArrowBuf where the data will be copied from checkIndex(index, dst.remaining()); // dst.remaining() bytes of data will be copied into dst ByteBuffer if (dst.remaining() != 0) { // address in this ArrowBuf where the copy will begin from final long srcAddress = addr(index); if (dst.isDirect()) { if (dst.isReadOnly()) { throw new ReadOnlyBufferException(); } // copy dst.remaining() bytes of data from this ArrowBuf starting // at address srcAddress into the dst ByteBuffer starting at // address dstAddress final long dstAddress = MemoryUtil.getByteBufferAddress(dst) + dst.position(); MemoryUtil.copyMemory(srcAddress, dstAddress, dst.remaining()); // after copy, bump the next write position for the dst ByteBuffer dst.position(dst.position() + dst.remaining()); } else if (dst.hasArray()) { // copy dst.remaining() bytes of data from this ArrowBuf starting // at address srcAddress into the dst ByteBuffer starting at // index dstIndex final int dstIndex = dst.arrayOffset() + dst.position(); MemoryUtil.copyFromMemory(srcAddress, dst.array(), dstIndex, dst.remaining()); // after copy, bump the next write position for the dst ByteBuffer dst.position(dst.position() + dst.remaining()); } else { throw new UnsupportedOperationException( "Copy from this ArrowBuf to ByteBuffer is not supported"); } } } /** * Copy data into this ArrowBuf at a given index onwards from a source ByteBuffer. * * @param index index index (0 based relative to the portion of memory this ArrowBuf has access * to) * @param src src ByteBuffer where the data will be copied from */ public void setBytes(long index, ByteBuffer src) { // bound check for this ArrowBuf where the data will be copied into checkIndex(index, src.remaining()); // length of data to copy int length = src.remaining(); // address in this ArrowBuf where the data will be copied to long dstAddress = addr(index); if (length != 0) { if (src.isDirect()) { // copy src.remaining() bytes of data from src ByteBuffer starting at // address srcAddress into this ArrowBuf starting at address dstAddress final long srcAddress = MemoryUtil.getByteBufferAddress(src) + src.position(); MemoryUtil.copyMemory(srcAddress, dstAddress, length); // after copy, bump the next read position for the src ByteBuffer src.position(src.position() + length); } else if (src.hasArray()) { // copy src.remaining() bytes of data from src ByteBuffer starting at // index srcIndex into this ArrowBuf starting at address dstAddress final int srcIndex = src.arrayOffset() + src.position(); MemoryUtil.copyToMemory(src.array(), srcIndex, dstAddress, length); // after copy, bump the next read position for the src ByteBuffer src.position(src.position() + length); } else { final ByteOrder originalByteOrder = src.order(); src.order(order()); try { // copy word at a time while (length - 128 >= LONG_SIZE) { for (int x = 0; x < 16; x++) { MemoryUtil.putLong(dstAddress, src.getLong()); length -= LONG_SIZE; dstAddress += LONG_SIZE; } } while (length >= LONG_SIZE) { MemoryUtil.putLong(dstAddress, src.getLong()); length -= LONG_SIZE; dstAddress += LONG_SIZE; } // copy last byte while (length > 0) { MemoryUtil.putByte(dstAddress, src.get()); --length; ++dstAddress; } } finally { src.order(originalByteOrder); } } } } /** * Copy data into this ArrowBuf at a given index onwards from a source ByteBuffer starting at a * given srcIndex for a certain length. * * @param index index (0 based relative to the portion of memory this ArrowBuf has access to) * @param src src ByteBuffer where the data will be copied from * @param srcIndex starting index in the src ByteBuffer where the data copy will start from * @param length length of data to copy from src ByteBuffer */ public void setBytes(long index, ByteBuffer src, int srcIndex, int length) { // bound check for this ArrowBuf where the data will be copied into checkIndex(index, length); if (src.isDirect()) { // copy length bytes of data from src ByteBuffer starting at address // srcAddress into this ArrowBuf at address dstAddress final long srcAddress = MemoryUtil.getByteBufferAddress(src) + srcIndex; final long dstAddress = addr(index); MemoryUtil.copyMemory(srcAddress, dstAddress, length); } else { if (srcIndex == 0 && src.capacity() == length) { // copy the entire ByteBuffer from start to end of length setBytes(index, src); } else { ByteBuffer newBuf = src.duplicate(); newBuf.position(srcIndex); newBuf.limit(srcIndex + length); setBytes(index, newBuf); } } } /** * Copy a given length of data from this ArrowBuf starting at a given index into a dst ArrowBuf at * dstIndex. * * @param index index (0 based relative to the portion of memory this ArrowBuf has access to) * @param dst dst ArrowBuf where the data will be copied into * @param dstIndex index (0 based relative to the portion of memory dst ArrowBuf has access to) * @param length length of data to copy */ public void getBytes(long index, ArrowBuf dst, long dstIndex, int length) { // bound check for this ArrowBuf where the data will be copied from checkIndex(index, length); // bound check for this ArrowBuf where the data will be copied into Preconditions.checkArgument(dst != null, "expecting a valid ArrowBuf"); // bound check for dst ArrowBuf if (isOutOfBounds(dstIndex, length, dst.capacity())) { throw new IndexOutOfBoundsException( String.format( "index: %d, length: %d (expected: range(0, %d))", dstIndex, length, dst.capacity())); } if (length != 0) { // copy length bytes of data from this ArrowBuf starting at // address srcAddress into dst ArrowBuf starting at address // dstAddress final long srcAddress = addr(index); final long dstAddress = dst.memoryAddress() + (long) dstIndex; MemoryUtil.copyMemory(srcAddress, dstAddress, length); } } /** * Copy data from src ArrowBuf starting at index srcIndex into this ArrowBuf at given index. * * @param index index index (0 based relative to the portion of memory this ArrowBuf has access * to) * @param src src ArrowBuf where the data will be copied from * @param srcIndex starting index in the src ArrowBuf where the copy will begin from * @param length length of data to copy from src ArrowBuf */ public void setBytes(long index, ArrowBuf src, long srcIndex, long length) { // bound check for this ArrowBuf where the data will be copied into checkIndex(index, length); // null check Preconditions.checkArgument(src != null, "expecting a valid ArrowBuf"); // bound check for src ArrowBuf if (isOutOfBounds(srcIndex, length, src.capacity())) { throw new IndexOutOfBoundsException( String.format( "index: %d, length: %d (expected: range(0, %d))", srcIndex, length, src.capacity())); } if (length != 0) { // copy length bytes of data from src ArrowBuf starting at // address srcAddress into this ArrowBuf starting at address // dstAddress final long srcAddress = src.memoryAddress() + srcIndex; final long dstAddress = addr(index); MemoryUtil.copyMemory(srcAddress, dstAddress, length); } } /** * Copy readableBytes() number of bytes from src ArrowBuf starting from its readerIndex into this * ArrowBuf starting at the given index. * * @param index index index (0 based relative to the portion of memory this ArrowBuf has access * to) * @param src src ArrowBuf where the data will be copied from */ public void setBytes(long index, ArrowBuf src) { // null check Preconditions.checkArgument(src != null, "expecting valid ArrowBuf"); final long length = src.readableBytes(); // bound check for this ArrowBuf where the data will be copied into checkIndex(index, length); final long srcAddress = src.memoryAddress() + src.readerIndex; final long dstAddress = addr(index); MemoryUtil.copyMemory(srcAddress, dstAddress, length); src.readerIndex(src.readerIndex + length); } /** * Copy a certain length of bytes from given InputStream into this ArrowBuf at the provided index. * * @param index index index (0 based relative to the portion of memory this ArrowBuf has access * to) * @param in src stream to copy from * @param length length of data to copy * @return number of bytes copied from stream into ArrowBuf * @throws IOException on failing to read from stream */ public int setBytes(long index, InputStream in, int length) throws IOException { Preconditions.checkArgument(in != null, "expecting valid input stream"); checkIndex(index, length); int readBytes = 0; if (length > 0) { byte[] tmp = new byte[length]; // read the data from input stream into tmp byte array readBytes = in.read(tmp); if (readBytes > 0) { // copy readBytes length of data from the tmp byte array starting // at srcIndex 0 into this ArrowBuf starting at address addr(index) MemoryUtil.copyToMemory(tmp, 0, addr(index), readBytes); } } return readBytes; } /** * Copy a certain length of bytes from this ArrowBuf at a given index into the given OutputStream. * * @param index index index (0 based relative to the portion of memory this ArrowBuf has access * to) * @param out dst stream to copy data into * @param length length of data to copy * @throws IOException on failing to write to stream */ public void getBytes(long index, OutputStream out, int length) throws IOException { Preconditions.checkArgument(out != null, "expecting valid output stream"); checkIndex(index, length); if (length > 0) { // copy length bytes of data from this ArrowBuf starting at // address addr(index) into the tmp byte array starting at index 0 byte[] tmp = new byte[length]; MemoryUtil.copyFromMemory(addr(index), tmp, 0, length); // write the copied data to output stream out.write(tmp); } } @Override public void close() { referenceManager.release(); } /** * Returns the possible memory consumed by this ArrowBuf in the worse case scenario. (not shared, * connected to larger underlying buffer of allocated memory) * * @return Size in bytes. */ public long getPossibleMemoryConsumed() { return referenceManager.getSize(); } /** * Return that is Accounted for by this buffer (and its potentially shared siblings within the * context of the associated allocator). * * @return Size in bytes. */ public long getActualMemoryConsumed() { return referenceManager.getAccountedSize(); } /** * Return the buffer's byte contents in the form of a hex dump. * * @param start the starting byte index * @param length how many bytes to log * @return A hex dump in a String. */ public String toHexString(final long start, final int length) { final long roundedStart = (start / LOG_BYTES_PER_ROW) * LOG_BYTES_PER_ROW; final StringBuilder sb = new StringBuilder("buffer byte dump\n"); long index = roundedStart; for (long nLogged = 0; nLogged < length; nLogged += LOG_BYTES_PER_ROW) { sb.append(String.format(" [%05d-%05d]", index, index + LOG_BYTES_PER_ROW - 1)); for (int i = 0; i < LOG_BYTES_PER_ROW; ++i) { try { final byte b = getByte(index++); sb.append(String.format(" 0x%02x", b)); } catch (IndexOutOfBoundsException ioob) { sb.append(" "); } } sb.append('\n'); } return sb.toString(); } /** * 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 the identity hash code for this buffer */ public long getId() { return System.identityHashCode(this); } /** * Print information of this buffer into sb at the given indentation and verbosity * level. * *

    It will include history if BaseAllocator.DEBUG is true and the * verbosity.includeHistoricalLog are true. */ @VisibleForTesting public void print(StringBuilder sb, int indent, Verbosity verbosity) { CommonUtil.indent(sb, indent).append(toString()); if (historicalLog != null && verbosity.includeHistoricalLog) { sb.append("\n"); historicalLog.buildHistory(sb, indent + 1, verbosity.includeStackTraces); } } /** * Print detailed information of this buffer into sb. * *

    Most information will only be present if BaseAllocator.DEBUG is true. */ public void print(StringBuilder sb, int indent) { print(sb, indent, Verbosity.LOG_WITH_STACKTRACE); } /** * Get the index at which the next byte will be read from. * * @return reader index */ public long readerIndex() { return readerIndex; } /** * Get the index at which next byte will be written to. * * @return writer index */ public long writerIndex() { return writerIndex; } /** * Set the reader index for this ArrowBuf. * * @param readerIndex new reader index * @return this ArrowBuf */ public ArrowBuf readerIndex(long readerIndex) { this.readerIndex = readerIndex; return this; } /** * Set the writer index for this ArrowBuf. * * @param writerIndex new writer index * @return this ArrowBuf */ public ArrowBuf writerIndex(long writerIndex) { this.writerIndex = writerIndex; return this; } /** * Zero-out the bytes in this ArrowBuf starting at the given index for the given length. * * @param index index index (0 based relative to the portion of memory this ArrowBuf has access * to) * @param length length of bytes to zero-out * @return this ArrowBuf */ public ArrowBuf setZero(long index, long length) { if (length != 0) { this.checkIndex(index, length); MemoryUtil.setMemory(this.addr + index, length, (byte) 0); } return this; } /** * Sets all bits to one in the specified range. * * @param index index index (0 based relative to the portion of memory this ArrowBuf has access * to) * @param length length of bytes to set. * @return this ArrowBuf * @deprecated use {@link ArrowBuf#setOne(long, long)} instead. */ @Deprecated public ArrowBuf setOne(int index, int length) { if (length != 0) { this.checkIndex(index, length); MemoryUtil.setMemory(this.addr + index, length, (byte) 0xff); } return this; } /** * Sets all bits to one in the specified range. * * @param index index index (0 based relative to the portion of memory this ArrowBuf has access * to) * @param length length of bytes to set. * @return this ArrowBuf */ public ArrowBuf setOne(long index, long length) { if (length != 0) { this.checkIndex(index, length); MemoryUtil.setMemory(this.addr + index, length, (byte) 0xff); } return this; } /** * Returns this if size is less than {@link #capacity()}, otherwise delegates to * {@link BufferManager#replace(ArrowBuf, long)} to get a new buffer. */ public ArrowBuf reallocIfNeeded(final long size) { Preconditions.checkArgument(size >= 0, "reallocation size must be non-negative"); if (this.capacity() >= size) { return this; } if (bufferManager != null) { return bufferManager.replace(this, size); } else { throw new UnsupportedOperationException( "Realloc is only available in the context of operator's UDFs"); } } public ArrowBuf clear() { this.readerIndex = this.writerIndex = 0; return this; } } apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/BaseAllocator.java000066400000000000000000001036101515451615200333310ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 com.google.errorprone.annotations.FormatMethod; import com.google.errorprone.annotations.FormatString; import java.util.Collection; import java.util.Collections; import java.util.HashSet; import java.util.IdentityHashMap; import java.util.Map; import java.util.Set; import org.apache.arrow.memory.rounding.DefaultRoundingPolicy; import org.apache.arrow.memory.rounding.RoundingPolicy; import org.apache.arrow.memory.util.AssertionUtil; import org.apache.arrow.memory.util.CommonUtil; import org.apache.arrow.memory.util.HistoricalLog; import org.apache.arrow.memory.util.LargeMemoryUtil; import org.apache.arrow.util.Preconditions; import org.checkerframework.checker.initialization.qual.Initialized; import org.checkerframework.checker.nullness.qual.KeyFor; import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; import org.immutables.value.Value; /** * A base-class that implements all functionality of {@linkplain BufferAllocator}s. * *

    The class is abstract to enforce usage of {@linkplain RootAllocator}/{@linkplain * ChildAllocator} facades. */ abstract class BaseAllocator extends Accountant implements BufferAllocator { public static final String DEBUG_ALLOCATOR = "arrow.memory.debug.allocator"; public static final int DEBUG_LOG_LENGTH = 6; public static final boolean DEBUG; private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(BaseAllocator.class); // Initialize this before DEFAULT_CONFIG as DEFAULT_CONFIG will eventually initialize the // allocation manager, // which in turn allocates an ArrowBuf, which requires DEBUG to have been properly initialized static { // the system property takes precedence. String propValue = System.getProperty(DEBUG_ALLOCATOR); if (propValue != null) { DEBUG = Boolean.parseBoolean(propValue); } else { DEBUG = false; } logger.info( "Debug mode " + (DEBUG ? "enabled." : "disabled. Enable with the VM option -Darrow.memory.debug.allocator=true.")); } public static final Config DEFAULT_CONFIG = ImmutableConfig.builder().build(); // Package exposed for sharing between AllocatorManger and BaseAllocator objects private final String name; private final RootAllocator root; private final Object DEBUG_LOCK = new Object(); private final AllocationListener listener; private final @Nullable BaseAllocator parentAllocator; private final Map childAllocators; private final ArrowBuf empty; // members used purely for debugging private final @Nullable IdentityHashMap childLedgers; private final @Nullable IdentityHashMap reservations; private final @Nullable HistoricalLog historicalLog; private final RoundingPolicy roundingPolicy; private final AllocationManager.@NonNull Factory allocationManagerFactory; private volatile boolean isClosed = false; // the allocator has been closed /** * Initialize an allocator. * * @param parentAllocator parent allocator. null if defining a root allocator * @param name name of this allocator * @param config configuration including other options of this allocator * @see Config */ @SuppressWarnings({"nullness:method.invocation", "nullness:cast.unsafe"}) // {"call to hist(,...) not allowed on the given receiver.", "cast cannot be statically verified"} protected BaseAllocator( final @Nullable BaseAllocator parentAllocator, final String name, final Config config) throws OutOfMemoryException { super(parentAllocator, name, config.getInitReservation(), config.getMaxAllocation()); this.listener = config.getListener(); this.allocationManagerFactory = config.getAllocationManagerFactory(); if (parentAllocator != null) { this.root = parentAllocator.root; empty = parentAllocator.empty; } else if (this instanceof RootAllocator) { this.root = (@Initialized RootAllocator) this; empty = createEmpty(); } else { throw new IllegalStateException( "An parent allocator must either carry a root or be the " + "root."); } this.parentAllocator = parentAllocator; this.name = name; this.childAllocators = Collections.synchronizedMap(new IdentityHashMap<>()); if (DEBUG) { reservations = new IdentityHashMap<>(); childLedgers = new IdentityHashMap<>(); historicalLog = new HistoricalLog(DEBUG_LOG_LENGTH, "allocator[%s]", name); hist("created by \"%s\", owned = %d", name, this.getAllocatedMemory()); } else { reservations = null; historicalLog = null; childLedgers = null; } this.roundingPolicy = config.getRoundingPolicy(); } @Override public AllocationListener getListener() { return listener; } @Override public @Nullable BaseAllocator getParentAllocator() { return parentAllocator; } @Override public Collection getChildAllocators() { synchronized (childAllocators) { return new HashSet<>(childAllocators.keySet()); } } private static String createErrorMsg( final BufferAllocator allocator, final long rounded, final long requested) { if (rounded != requested) { return String.format( "Unable to allocate buffer of size %d (rounded from %d) due to memory limit. Current " + "allocation: %d", rounded, requested, allocator.getAllocatedMemory()); } else { return String.format( "Unable to allocate buffer of size %d due to memory limit. Current " + "allocation: %d", rounded, allocator.getAllocatedMemory()); } } public static boolean isDebug() { return DEBUG; } @Override public void assertOpen() { if (AssertionUtil.ASSERT_ENABLED) { if (isClosed) { throw new IllegalStateException( "Attempting operation on allocator when allocator is closed.\n" + toVerboseString()); } } } @Override public String getName() { return name; } @Override public ArrowBuf getEmpty() { return empty; } /** * For debug/verification purposes only. Allows an AllocationManager to tell the allocator that we * have a new ledger associated with this allocator. */ void associateLedger(BufferLedger ledger) { assertOpen(); if (DEBUG) { synchronized (DEBUG_LOCK) { if (childLedgers != null) { childLedgers.put(ledger, null); } } } } /** * For debug/verification purposes only. Allows an AllocationManager to tell the allocator that we * are removing a ledger associated with this allocator */ void dissociateLedger(BufferLedger ledger) { assertOpen(); if (DEBUG) { synchronized (DEBUG_LOCK) { Preconditions.checkState(childLedgers != null, "childLedgers must not be null"); if (!childLedgers.containsKey(ledger)) { throw new IllegalStateException("Trying to remove a child ledger that doesn't exist."); } childLedgers.remove(ledger); } } } /** * Track when a ChildAllocator of this BaseAllocator is closed. Used for debugging purposes. * * @param childAllocator The child allocator that has been closed. */ private void childClosed(final BaseAllocator childAllocator) { assertOpen(); if (DEBUG) { Preconditions.checkArgument(childAllocator != null, "child allocator can't be null"); synchronized (DEBUG_LOCK) { final Object object = childAllocators.remove(childAllocator); if (object == null) { if (childAllocator.historicalLog != null) { childAllocator.historicalLog.logHistory(logger); } throw new IllegalStateException( "Child allocator[" + childAllocator.name + "] not found in parent allocator[" + name + "]'s childAllocators"); } } } else { childAllocators.remove(childAllocator); } listener.onChildRemoved(this, childAllocator); } @Override public ArrowBuf wrapForeignAllocation(ForeignAllocation allocation) { assertOpen(); final long size = allocation.getSize(); listener.onPreAllocation(size); AllocationOutcome outcome = this.allocateBytes(size); if (!outcome.isOk()) { if (listener.onFailedAllocation(size, outcome)) { // Second try, in case the listener can do something about it outcome = this.allocateBytes(size); } if (!outcome.isOk()) { throw new OutOfMemoryException(createErrorMsg(this, size, size), outcome.getDetails()); } } try { final AllocationManager manager = new ForeignAllocationManager(this, allocation); final BufferLedger ledger = manager.associate(this); final ArrowBuf buf = new ArrowBuf(ledger, /* bufferManager= */ null, size, allocation.memoryAddress()); buf.writerIndex(size); listener.onAllocation(size); return buf; } catch (Throwable t) { try { releaseBytes(size); } catch (Throwable e) { t.addSuppressed(e); } try { allocation.release0(); } catch (Throwable e) { t.addSuppressed(e); } throw t; } } @Override public ArrowBuf buffer(final long initialRequestSize) { assertOpen(); return buffer(initialRequestSize, null); } @SuppressWarnings("nullness:dereference.of.nullable") // dereference of possibly-null reference // allocationManagerFactory private ArrowBuf createEmpty() { return allocationManagerFactory.empty(); } @Override public ArrowBuf buffer(final long initialRequestSize, @Nullable BufferManager manager) { assertOpen(); Preconditions.checkArgument(initialRequestSize >= 0, "the requested size must be non-negative"); if (initialRequestSize == 0) { return getEmpty(); } // round the request size according to the rounding policy final long actualRequestSize = roundingPolicy.getRoundedSize(initialRequestSize); listener.onPreAllocation(actualRequestSize); AllocationOutcome outcome = this.allocateBytes(actualRequestSize); if (!outcome.isOk()) { if (listener.onFailedAllocation(actualRequestSize, outcome)) { // Second try, in case the listener can do something about it outcome = this.allocateBytes(actualRequestSize); } if (!outcome.isOk()) { throw new OutOfMemoryException( createErrorMsg(this, actualRequestSize, initialRequestSize), outcome.getDetails()); } } boolean success = false; try { ArrowBuf buffer = bufferWithoutReservation(actualRequestSize, manager); success = true; listener.onAllocation(actualRequestSize); return buffer; } catch (OutOfMemoryError e) { throw e; } finally { if (!success) { releaseBytes(actualRequestSize); } } } /** * Used by usual allocation as well as for allocating a pre-reserved buffer. Skips the typical * accounting associated with creating a new buffer. */ private ArrowBuf bufferWithoutReservation(final long size, @Nullable BufferManager bufferManager) throws OutOfMemoryException { assertOpen(); final AllocationManager manager = newAllocationManager(size); final BufferLedger ledger = manager.associate(this); // +1 ref cnt (required) final ArrowBuf buffer = ledger.newArrowBuf(size, bufferManager); // make sure that our allocation is equal to what we expected. Preconditions.checkArgument( buffer.capacity() == size, "Allocated capacity %d was not equal to requested capacity %d.", buffer.capacity(), size); return buffer; } private AllocationManager newAllocationManager(long size) { return newAllocationManager(this, size); } private AllocationManager newAllocationManager(BaseAllocator accountingAllocator, long size) { return allocationManagerFactory.create(accountingAllocator, size); } @Override public BufferAllocator getRoot() { return root; } @Override public BufferAllocator newChildAllocator( final String name, final long initReservation, final long maxAllocation) { return newChildAllocator(name, this.listener, initReservation, maxAllocation); } @Override public BufferAllocator newChildAllocator( final String name, final AllocationListener listener, final long initReservation, final long maxAllocation) { assertOpen(); final ChildAllocator childAllocator = new ChildAllocator( this, name, configBuilder() .listener(listener) .initReservation(initReservation) .maxAllocation(maxAllocation) .roundingPolicy(roundingPolicy) .allocationManagerFactory(allocationManagerFactory) .build()); if (DEBUG) { synchronized (DEBUG_LOCK) { childAllocators.put(childAllocator, childAllocator); if (historicalLog != null) { historicalLog.recordEvent( "allocator[%s] created new child allocator[%s]", name, childAllocator.getName()); } } } else { childAllocators.put(childAllocator, childAllocator); } this.listener.onChildAdded(this, childAllocator); return childAllocator; } @Override public AllocationReservation newReservation() { assertOpen(); return new Reservation(); } @Override public synchronized void close() { /* * Some owners may close more than once because of complex cleanup and shutdown * procedures. */ if (isClosed) { return; } isClosed = true; StringBuilder outstandingChildAllocators = new StringBuilder(); if (DEBUG) { synchronized (DEBUG_LOCK) { verifyAllocator(); // are there outstanding child allocators? if (!childAllocators.isEmpty()) { for (final BaseAllocator childAllocator : childAllocators.keySet()) { if (childAllocator.isClosed) { logger.warn( String.format( "Closed child allocator[%s] on parent allocator[%s]'s child list.\n%s", childAllocator.name, name, toString())); } } throw new IllegalStateException( String.format( "Allocator[%s] closed with outstanding child allocators.\n%s", name, toString())); } // are there outstanding buffers? final int allocatedCount = childLedgers != null ? childLedgers.size() : 0; if (allocatedCount > 0) { throw new IllegalStateException( String.format( "Allocator[%s] closed with outstanding buffers allocated (%d).\n%s", name, allocatedCount, toString())); } if (reservations != null && reservations.size() != 0) { throw new IllegalStateException( String.format( "Allocator[%s] closed with outstanding reservations (%d).\n%s", name, reservations.size(), toString())); } } } else { if (!childAllocators.isEmpty()) { outstandingChildAllocators.append("Outstanding child allocators : \n"); synchronized (childAllocators) { for (final BaseAllocator childAllocator : childAllocators.keySet()) { outstandingChildAllocators.append(String.format(" %s", childAllocator.toString())); } } } } // Is there unaccounted-for outstanding allocation? final long allocated = getAllocatedMemory(); if (allocated > 0) { if (parent != null && reservation > allocated) { parent.releaseBytes(reservation - allocated); } String msg = String.format( "Memory was leaked by query. Memory leaked: (%d)\n%s%s", allocated, outstandingChildAllocators.toString(), toString()); logger.error(msg); throw new IllegalStateException(msg); } // we need to release our memory to our parent before we tell it we've closed. super.close(); // Inform our parent allocator that we've closed if (parentAllocator != null) { parentAllocator.childClosed(this); } if (DEBUG) { if (historicalLog != null) { historicalLog.recordEvent("closed"); } logger.debug(String.format("closed allocator[%s].", name)); } } @Override public String toString() { final Verbosity verbosity = logger.isTraceEnabled() ? Verbosity.LOG_WITH_STACKTRACE : Verbosity.BASIC; final StringBuilder sb = new StringBuilder(); print(sb, 0, verbosity); return sb.toString(); } /** * Provide a verbose string of the current allocator state. Includes the state of all child * allocators, along with historical logs of each object and including stacktraces. * * @return A Verbose string of current allocator state. */ @Override public String toVerboseString() { final StringBuilder sb = new StringBuilder(); print(sb, 0, Verbosity.LOG_WITH_STACKTRACE); return sb.toString(); } @FormatMethod private void hist(@FormatString String noteFormat, Object... args) { if (historicalLog != null) { historicalLog.recordEvent(noteFormat, args); } } /** * Verifies the accounting state of the allocator. Only works for DEBUG. * * @throws IllegalStateException when any problems are found */ void verifyAllocator() { final IdentityHashMap seen = new IdentityHashMap<>(); verifyAllocator(seen); } /** * Verifies the accounting state of the allocator (Only works for DEBUG) This overload is used for * recursive calls, allowing for checking that ArrowBufs are unique across all allocators that are * checked. * * @param buffersSeen a map of buffers that have already been seen when walking a tree of * allocators * @throws IllegalStateException when any problems are found */ private void verifyAllocator( final IdentityHashMap buffersSeen) { // The remaining tests can only be performed if we're in debug mode. if (!DEBUG) { return; } synchronized (DEBUG_LOCK) { final long allocated = getAllocatedMemory(); // verify my direct descendants final Set childSet = childAllocators.keySet(); for (final BaseAllocator childAllocator : childSet) { childAllocator.verifyAllocator(buffersSeen); } /* * Verify my relationships with my descendants. * * The sum of direct child allocators' owned memory must be <= my allocated memory; my * allocated memory also * includes ArrowBuf's directly allocated by me. */ long childTotal = 0; for (final BaseAllocator childAllocator : childSet) { childTotal += Math.max(childAllocator.getAllocatedMemory(), childAllocator.reservation); } if (childTotal > getAllocatedMemory()) { if (historicalLog != null) { historicalLog.logHistory(logger); } logger.debug("allocator[" + name + "] child event logs BEGIN"); for (final BaseAllocator childAllocator : childSet) { if (childAllocator.historicalLog != null) { childAllocator.historicalLog.logHistory(logger); } } logger.debug("allocator[" + name + "] child event logs END"); throw new IllegalStateException( "Child allocators own more memory (" + childTotal + ") than their parent (name = " + name + " ) has allocated (" + getAllocatedMemory() + ')'); } // Furthermore, the amount I've allocated should be that plus buffers I've allocated. long bufferTotal = 0; final Set<@KeyFor("this.childLedgers") BufferLedger> ledgerSet = childLedgers != null ? childLedgers.keySet() : null; if (ledgerSet != null) { for (final BufferLedger ledger : ledgerSet) { if (!ledger.isOwningLedger()) { continue; } final AllocationManager am = ledger.getAllocationManager(); /* * Even when shared, ArrowBufs are rewrapped, so we should never see the same instance * twice. */ final BaseAllocator otherOwner = buffersSeen.get(am); if (otherOwner != null) { throw new IllegalStateException( "This allocator's ArrowBuf already owned by another " + "allocator"); } buffersSeen.put(am, this); bufferTotal += am.getSize(); } } // Preallocated space has to be accounted for final Set<@KeyFor("this.reservations") Reservation> reservationSet = reservations != null ? reservations.keySet() : null; long reservedTotal = 0; if (reservationSet != null) { for (final Reservation reservation : reservationSet) { if (!reservation.isUsed()) { reservedTotal += reservation.getSize(); } } } if (bufferTotal + reservedTotal + childTotal != getAllocatedMemory()) { final StringBuilder sb = new StringBuilder(); sb.append("allocator["); sb.append(name); sb.append("]\nallocated: "); sb.append(Long.toString(allocated)); sb.append(" allocated - (bufferTotal + reservedTotal + childTotal): "); sb.append(Long.toString(allocated - (bufferTotal + reservedTotal + childTotal))); sb.append('\n'); if (bufferTotal != 0) { sb.append("buffer total: "); sb.append(Long.toString(bufferTotal)); sb.append('\n'); dumpBuffers(sb, ledgerSet); } if (childTotal != 0) { sb.append("child total: "); sb.append(Long.toString(childTotal)); sb.append('\n'); for (final BaseAllocator childAllocator : childSet) { sb.append("child allocator["); sb.append(childAllocator.name); sb.append("] owned "); sb.append(Long.toString(childAllocator.getAllocatedMemory())); sb.append('\n'); } } if (reservedTotal != 0) { sb.append(String.format("reserved total : %d bytes.", reservedTotal)); if (reservationSet != null) { for (final Reservation reservation : reservationSet) { if (reservation.historicalLog != null) { reservation.historicalLog.buildHistory(sb, 0, true); } sb.append('\n'); } } } logger.debug(sb.toString()); final long allocated2 = getAllocatedMemory(); if (allocated2 != allocated) { throw new IllegalStateException( String.format( "allocator[%s]: allocated t1 (%d) + allocated t2 (%d). Someone released memory while in verification.", name, allocated, allocated2)); } throw new IllegalStateException( String.format( "allocator[%s]: buffer space (%d) + prealloc space (%d) + child space (%d) != allocated (%d)", name, bufferTotal, reservedTotal, childTotal, allocated)); } } } void print(StringBuilder sb, int level, Verbosity verbosity) { CommonUtil.indent(sb, level) .append("Allocator(") .append(name) .append(") ") .append(reservation) .append('/') .append(getAllocatedMemory()) .append('/') .append(getPeakMemoryAllocation()) .append('/') .append(getLimit()) .append(" (res/actual/peak/limit)") .append('\n'); if (DEBUG) { CommonUtil.indent(sb, level + 1) .append(String.format("child allocators: %d\n", childAllocators.size())); for (BaseAllocator child : childAllocators.keySet()) { child.print(sb, level + 2, verbosity); } CommonUtil.indent(sb, level + 1) .append(String.format("ledgers: %d\n", childLedgers != null ? childLedgers.size() : 0)); if (childLedgers != null) { for (BufferLedger ledger : childLedgers.keySet()) { ledger.print(sb, level + 2, verbosity); } } final Set<@KeyFor("this.reservations") Reservation> reservations = this.reservations != null ? this.reservations.keySet() : null; CommonUtil.indent(sb, level + 1) .append( String.format("reservations: %d\n", reservations != null ? reservations.size() : 0)); if (reservations != null) { for (final Reservation reservation : reservations) { if (verbosity.includeHistoricalLog) { if (reservation.historicalLog != null) { reservation.historicalLog.buildHistory(sb, level + 3, true); } } } } } } private void dumpBuffers( final StringBuilder sb, final @Nullable Set<@KeyFor("this.childLedgers") BufferLedger> ledgerSet) { if (ledgerSet != null) { for (final BufferLedger ledger : ledgerSet) { if (!ledger.isOwningLedger()) { continue; } final AllocationManager am = ledger.getAllocationManager(); sb.append("UnsafeDirectLittleEndian[identityHashCode == "); sb.append(Integer.toString(System.identityHashCode(am))); sb.append("] size "); sb.append(Long.toString(am.getSize())); sb.append('\n'); } } } /** Enum for logging verbosity. */ public enum Verbosity { BASIC(false, false), // only include basic information LOG(true, false), // include basic LOG_WITH_STACKTRACE(true, true) // ; public final boolean includeHistoricalLog; public final boolean includeStackTraces; Verbosity(boolean includeHistoricalLog, boolean includeStackTraces) { this.includeHistoricalLog = includeHistoricalLog; this.includeStackTraces = includeStackTraces; } } /** * Returns a default {@link Config} instance. * * @see ImmutableConfig.Builder */ public static Config defaultConfig() { return DEFAULT_CONFIG; } /** Returns a builder class for configuring BaseAllocator's options. */ public static ImmutableConfig.Builder configBuilder() { return ImmutableConfig.builder(); } @Override public RoundingPolicy getRoundingPolicy() { return roundingPolicy; } /** Config class of {@link BaseAllocator}. */ @Value.Immutable abstract static class Config { /** Factory for creating {@link AllocationManager} instances. */ @Value.Default AllocationManager.Factory getAllocationManagerFactory() { return DefaultAllocationManagerOption.getDefaultAllocationManagerFactory(); } /** Listener callback. Must be non-null. */ @Value.Default AllocationListener getListener() { return AllocationListener.NOOP; } /** Initial reservation size (in bytes) for this allocator. */ @Value.Default long getInitReservation() { return 0; } /** * Max allocation size (in bytes) for this allocator, allocations past this limit fail. Can be * modified after construction. */ @Value.Default long getMaxAllocation() { return Long.MAX_VALUE; } /** The policy for rounding the buffer size. */ @Value.Default RoundingPolicy getRoundingPolicy() { return DefaultRoundingPolicy.DEFAULT_ROUNDING_POLICY; } } /** * Implementation of {@link AllocationReservation} that supports history tracking under * {@linkplain #DEBUG} is true. */ public class Reservation implements AllocationReservation { private final @Nullable HistoricalLog historicalLog; private long nBytes = 0; private boolean used = false; private boolean closed = false; /** * Creates a new reservation. * *

    If {@linkplain #DEBUG} is true this will capture a historical log of events relevant to * this Reservation. */ @SuppressWarnings("nullness:argument") // to handle null assignment on third party dependency: // System.identityHashCode public Reservation() { if (DEBUG) { historicalLog = new HistoricalLog( "Reservation[allocator[%s], %d]", name, System.identityHashCode(this)); historicalLog.recordEvent("created"); synchronized (DEBUG_LOCK) { if (reservations != null) { reservations.put(this, this); } } } else { historicalLog = null; } } @SuppressWarnings({"removal", "InlineMeSuggester"}) @Deprecated(forRemoval = true) @Override public boolean add(final int nBytes) { return add((long) nBytes); } @Override public boolean add(final long nBytes) { assertOpen(); Preconditions.checkArgument(nBytes >= 0, "nBytes(%d) < 0", nBytes); Preconditions.checkState( !closed, "Attempt to increase reservation after reservation has been closed"); Preconditions.checkState( !used, "Attempt to increase reservation after reservation has been used"); // we round up to next power of two since all reservations are done in powers of two. This // may overestimate the // preallocation since someone may perceive additions to be power of two. If this becomes a // problem, we can look // at // modifying this behavior so that we maintain what we reserve and what the user asked for // and make sure to only // round to power of two as necessary. final long nBytesTwo = CommonUtil.nextPowerOfTwo(nBytes); if (!reserve(nBytesTwo)) { return false; } this.nBytes += nBytesTwo; return true; } @Override public ArrowBuf allocateBuffer() { assertOpen(); Preconditions.checkState(!closed, "Attempt to allocate after closed"); Preconditions.checkState(!used, "Attempt to allocate more than once"); final ArrowBuf arrowBuf = allocate(nBytes); used = true; return arrowBuf; } @Override public int getSize() { return LargeMemoryUtil.checkedCastToInt(nBytes); } @Override public long getSizeLong() { return nBytes; } @Override public boolean isUsed() { return used; } @Override public boolean isClosed() { return closed; } @Override public void close() { assertOpen(); if (closed) { return; } if (DEBUG) { if (!isClosed()) { final Object object; synchronized (DEBUG_LOCK) { object = reservations != null ? reservations.remove(this) : null; } if (object == null) { final StringBuilder sb = new StringBuilder(); print(sb, 0, Verbosity.LOG_WITH_STACKTRACE); logger.debug(sb.toString()); throw new IllegalStateException( String.format( "Didn't find closing reservation[%d]", System.identityHashCode(this))); } if (historicalLog != null) { historicalLog.recordEvent("closed"); } } } if (!used) { releaseReservation(nBytes); } closed = true; } @SuppressWarnings({"removal", "InlineMeSuggester"}) @Deprecated(forRemoval = true) @Override public boolean reserve(int nBytes) { return reserve((long) nBytes); } @Override public boolean reserve(long nBytes) { assertOpen(); final AllocationOutcome outcome = BaseAllocator.this.allocateBytes(nBytes); if (historicalLog != null) { historicalLog.recordEvent("reserve(%d) => %s", nBytes, Boolean.toString(outcome.isOk())); } return outcome.isOk(); } /** * Allocate a buffer of the requested size. * *

    The implementation of the allocator's inner class provides this. * * @param nBytes the size of the buffer requested * @return the buffer, or null, if the request cannot be satisfied */ private ArrowBuf allocate(long nBytes) { assertOpen(); boolean success = false; /* * The reservation already added the requested bytes to the allocators owned and allocated * bytes via reserve(). * This ensures that they can't go away. But when we ask for the buffer here, that will add * to the allocated bytes * as well, so we need to return the same number back to avoid double-counting them. */ try { final ArrowBuf arrowBuf = BaseAllocator.this.bufferWithoutReservation(nBytes, null); listener.onAllocation(nBytes); if (historicalLog != null) { historicalLog.recordEvent( "allocate() => %s", String.format("ArrowBuf[%d]", arrowBuf.getId())); } success = true; return arrowBuf; } finally { if (!success) { releaseBytes(nBytes); } } } /** * Return the reservation back to the allocator without having used it. * * @param nBytes the size of the reservation */ private void releaseReservation(long nBytes) { assertOpen(); releaseBytes(nBytes); if (historicalLog != null) { historicalLog.recordEvent("releaseReservation(%d)", nBytes); } } } } BoundsChecking.java000066400000000000000000000051451515451615200334310ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * Configuration class to determine if bounds checking should be turned on or off. * *

    Bounds checking is on by default. You can disable it by setting either the system property or * the environmental variable to "true". The system property can be * "arrow.enable_unsafe_memory_access" or "drill.enable_unsafe_memory_access". The latter is * deprecated. The environmental variable is named "ARROW_ENABLE_UNSAFE_MEMORY_ACCESS". When both * the system property and the environmental variable are set, the system property takes precedence. */ public class BoundsChecking { public static final boolean BOUNDS_CHECKING_ENABLED; static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(BoundsChecking.class); static { String envProperty = System.getenv("ARROW_ENABLE_UNSAFE_MEMORY_ACCESS"); String oldProperty = System.getProperty("drill.enable_unsafe_memory_access"); if (oldProperty != null) { logger.warn( "\"drill.enable_unsafe_memory_access\" has been renamed to \"arrow.enable_unsafe_memory_access\""); logger.warn( "\"arrow.enable_unsafe_memory_access\" can be set to: " + " true (to not check) or false (to check, default)"); } String newProperty = System.getProperty("arrow.enable_unsafe_memory_access"); // The priority of determining the unsafe flag: // 1. The system properties take precedence over the environmental variable. // 2. The new system property takes precedence over the new system property. String unsafeFlagValue = newProperty; if (unsafeFlagValue == null) { unsafeFlagValue = oldProperty; } if (unsafeFlagValue == null) { unsafeFlagValue = envProperty; } BOUNDS_CHECKING_ENABLED = !"true".equals(unsafeFlagValue); } private BoundsChecking() {} } BufferAllocator.java000066400000000000000000000200521515451615200336070ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.Collection; import org.apache.arrow.memory.rounding.DefaultRoundingPolicy; import org.apache.arrow.memory.rounding.RoundingPolicy; import org.checkerframework.checker.nullness.qual.Nullable; /** Wrapper class to deal with byte buffer allocation. Ensures users only use designated methods. */ public interface BufferAllocator extends AutoCloseable { /** * 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 * @throws OutOfMemoryException if buffer cannot be allocated */ ArrowBuf buffer(long 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. * @return a new ArrowBuf, or null if the request can't be satisfied * @throws OutOfMemoryException if buffer cannot be allocated */ ArrowBuf buffer(long size, BufferManager manager); /** * Get the root allocator of this allocator. If this allocator is already a root, return this * directly. * * @return The root allocator */ BufferAllocator getRoot(); /** * Create a new child allocator. * * @param name the name of the allocator. * @param initReservation the initial space reservation (obtained from this allocator) * @param maxAllocation maximum amount of space the new allocator can allocate * @return the new allocator, or null if it can't be created */ BufferAllocator newChildAllocator(String name, long initReservation, long maxAllocation); /** * Create a new child allocator. * * @param name the name of the allocator. * @param listener allocation listener for the newly created child * @param initReservation the initial space reservation (obtained from this allocator) * @param maxAllocation maximum amount of space the new allocator can allocate * @return the new allocator, or null if it can't be created */ BufferAllocator newChildAllocator( String name, AllocationListener listener, long initReservation, long maxAllocation); /** * Close and release all buffers generated from this buffer pool. * *

    When assertions are on, complains if there are any outstanding buffers; to avoid that, * release all buffers before the allocator is closed. */ @Override void close(); /** * Returns the amount of memory currently allocated from this allocator. * * @return the amount of memory currently allocated */ long getAllocatedMemory(); /** * Return the current maximum limit this allocator imposes. * * @return Limit in number of bytes. */ long getLimit(); /** * Return the initial reservation. * * @return reservation in bytes. */ long getInitReservation(); /** * Set the maximum amount of memory this allocator is allowed to allocate. * * @param newLimit The new Limit to apply to allocations */ void setLimit(long newLimit); /** * Returns the peak amount of memory allocated from this allocator. * * @return the peak amount of memory allocated */ long getPeakMemoryAllocation(); /** * Returns the amount of memory that can probably be allocated at this moment without exceeding * this or any parents allocation maximum. * * @return Headroom in bytes */ long getHeadroom(); /** * Forcibly allocate bytes. Returns whether the allocation fit within limits. * * @param size to increase * @return Whether the allocation fit within limits. */ boolean forceAllocate(long size); /** * Release bytes from this allocator. * * @param size to release */ void releaseBytes(long size); /** * Returns the allocation listener used by this allocator. * * @return the {@link AllocationListener} instance. Or {@link AllocationListener#NOOP} by default * if no listener is configured when this allocator was created. */ AllocationListener getListener(); /** * Returns the parent allocator. * * @return parent allocator */ @Nullable BufferAllocator getParentAllocator(); /** * Returns the set of child allocators. * * @return set of child allocators */ Collection getChildAllocators(); /** * Create an allocation reservation. A reservation is a way of building up a request for a buffer * whose size is not known in advance. See * * @return the newly created reservation * @see AllocationReservation */ AllocationReservation newReservation(); /** * Get a reference to the empty buffer associated with this allocator. Empty buffers are special * because we don't worry about them leaking or managing reference counts on them since they don't * actually point to any memory. * * @return the empty buffer */ ArrowBuf getEmpty(); /** * Return the name of this allocator. This is a human readable name that can help debugging. * Typically provides coordinates about where this allocator was created * * @return the name of the allocator */ String getName(); /** * Return whether or not this allocator (or one if its parents) is over its limits. In the case * that an allocator is over its limit, all consumers of that allocator should aggressively try to * address the overlimit situation. * * @return whether or not this allocator (or one if its parents) is over its limits */ boolean isOverLimit(); /** * Return a verbose string describing this allocator. If in DEBUG mode, this will also include * relevant stacktraces and historical logs for underlying objects * * @return A very verbose description of the allocator hierarchy. */ String toVerboseString(); /** * Asserts (using java assertions) that the provided allocator is currently open. If assertions * are disabled, this is a no-op. */ void assertOpen(); /** Gets the rounding policy of the allocator. */ default RoundingPolicy getRoundingPolicy() { return DefaultRoundingPolicy.DEFAULT_ROUNDING_POLICY; } /** * EXPERIMENTAL: Wrap an allocation created outside this BufferAllocator. * *

    This is useful to integrate allocations from native code into the same memory management * framework as Java-allocated buffers, presenting users a consistent API. The created buffer will * be tracked by this allocator and can be transferred like Java-allocated buffers. * *

    The underlying allocation will be closed when all references to the buffer are released. If * this method throws, the underlying allocation will also be closed. * * @param allocation The underlying allocation. */ default ArrowBuf wrapForeignAllocation(ForeignAllocation allocation) { throw new UnsupportedOperationException(); } } apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java000066400000000000000000000526541515451615200331650ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.IdentityHashMap; 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; import org.checkerframework.checker.nullness.qual.Nullable; /** * The reference manager that binds an {@link AllocationManager} to {@link BufferAllocator} and a * set of {@link ArrowBuf}. The set of ArrowBufs managed by this reference manager share a common * fate (same reference count). */ public class BufferLedger implements ValueWithKeyIncluded, ReferenceManager { private final @Nullable IdentityHashMap buffers = BaseAllocator.DEBUG ? new IdentityHashMap<>() : null; // 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; private final @Nullable HistoricalLog historicalLog = BaseAllocator.DEBUG ? new HistoricalLog(BaseAllocator.DEBUG_LOG_LENGTH, "BufferLedger[%d]", 1) : null; private volatile long lDestructionTime = 0; BufferLedger(final BufferAllocator allocator, final AllocationManager allocationManager) { this.allocator = allocator; this.allocationManager = allocationManager; } boolean isOwningLedger() { return this == allocationManager.getOwningLedger(); } @Override public BufferAllocator getKey() { return allocator; } /** * Get the buffer allocator associated with this reference manager. * * @return buffer allocator */ @Override public BufferAllocator getAllocator() { return allocator; } /** * Get this ledger's reference count. * * @return reference count */ @Override public int getRefCount() { return bufRefCnt; } /** * Increment the ledger's reference count for the associated underlying memory chunk. All * ArrowBufs managed by this ledger will share the ref count. */ void increment() { BUF_REF_CNT_UPDATER.incrementAndGet(this); } /** * Decrement the ledger's reference count by 1 for the associated underlying memory chunk. If the * reference count drops to 0, it implies that no ArrowBufs managed by this reference manager need * access to the memory chunk. In that case, the ledger should inform the allocation manager about * releasing its ownership for the chunk. Whether or not the memory chunk will be released is * something that {@link AllocationManager} will decide since tracks the usage of memory chunk * across multiple reference managers and allocators. * * @return true if the new ref count has dropped to 0, false otherwise */ @Override public boolean release() { return release(1); } /** * Decrement the ledger's reference count for the associated underlying memory chunk. If the * reference count drops to 0, it implies that no ArrowBufs managed by this reference manager need * access to the memory chunk. In that case, the ledger should inform the allocation manager about * releasing its ownership for the chunk. Whether or not the memory chunk will be released is * something that {@link AllocationManager} will decide since tracks the usage of memory chunk * across multiple reference managers and allocators. * * @param decrement amount to decrease the reference count by * @return true if the new ref count has dropped to 0, false otherwise */ @Override public boolean release(int decrement) { Preconditions.checkState( decrement >= 1, "ref count decrement should be greater than or equal to 1"); // decrement the ref count final int refCnt = decrement(decrement); if (historicalLog != null) { historicalLog.recordEvent("release(%d). original value: %d", decrement, refCnt + decrement); } // the new ref count should be >= 0 Preconditions.checkState(refCnt >= 0, "RefCnt has gone negative"); return refCnt == 0; } /** * Decrement the ledger's reference count for the associated underlying memory chunk. If the * reference count drops to 0, it implies that no ArrowBufs managed by this reference manager need * access to the memory chunk. In that case, the ledger should inform the allocation manager about * releasing its ownership for the chunk. Whether or not the memory chunk will be released is * something that {@link AllocationManager} will decide since tracks the usage of memory chunk * across multiple reference managers and allocators. * * @param decrement amount to decrease the reference count by * @return the new reference count */ private int decrement(int decrement) { allocator.assertOpen(); final int outcome; synchronized (allocationManager) { outcome = BUF_REF_CNT_UPDATER.addAndGet(this, -decrement); if (outcome == 0) { lDestructionTime = System.nanoTime(); // refcount of this reference manager has dropped to 0 // inform the allocation manager that this reference manager // no longer holds references to underlying memory allocationManager.release(this); } } return outcome; } /** Increment the ledger's reference count for associated underlying memory chunk by 1. */ @Override public void retain() { retain(1); } /** * Increment the ledger's reference count for associated underlying memory chunk by the given * amount. * * @param increment amount to increase the reference count by */ @Override public void retain(int increment) { Preconditions.checkArgument(increment > 0, "retain(%s) argument is not positive", increment); if (historicalLog != null) { historicalLog.recordEvent("retain(%d)", increment); } final int originalReferenceCount = BUF_REF_CNT_UPDATER.getAndAdd(this, increment); Preconditions.checkArgument(originalReferenceCount > 0); } /** * Derive a new ArrowBuf from a given source ArrowBuf. The new derived ArrowBuf will share the * same reference count as rest of the ArrowBufs associated with this ledger. This operation is * typically used for slicing -- creating new ArrowBufs from a compound ArrowBuf starting at a * particular index in the underlying memory and having access to a particular length (in bytes) * of data in memory chunk. * *

    This method is also used as a helper for transferring ownership and retain to target * allocator. * * @param sourceBuffer source ArrowBuf * @param index index (relative to source ArrowBuf) new ArrowBuf should be derived from * @param length length (bytes) of data in underlying memory that derived buffer will have access * to in underlying memory * @return derived buffer */ @Override public ArrowBuf deriveBuffer(final ArrowBuf sourceBuffer, long index, long length) { /* * Usage type 1 for deriveBuffer(): * Used for slicing where index represents a relative index in the source ArrowBuf * as the slice start point. This is why we need to add the source buffer offset * to compute the start virtual address of derived buffer within the * underlying chunk. * * Usage type 2 for deriveBuffer(): * Used for retain(target allocator) and transferOwnership(target allocator) * where index is 0 since these operations simply create a new ArrowBuf associated * with another combination of allocator buffer ledger for the same underlying memory */ // the memory address stored inside ArrowBuf is its starting virtual // address in the underlying memory chunk from the point it has // access. so it is already accounting for the offset of the source buffer // we can simply add the index to get the starting address of new buffer. final long derivedBufferAddress = sourceBuffer.memoryAddress() + index; // create new ArrowBuf final ArrowBuf derivedBuf = new ArrowBuf( this, null, length, // length (in bytes) in the underlying memory chunk for this new ArrowBuf derivedBufferAddress // starting byte address in the underlying memory for this new // ArrowBuf ); // logging return loggingArrowBufHistoricalLog(derivedBuf); } /** * Used by an allocator to create a new ArrowBuf. This is provided as a helper method for the * allocator when it allocates a new memory chunk using a new instance of allocation manager and * creates a new reference manager too. * * @param length The length in bytes that this ArrowBuf will provide access to. * @param manager An optional BufferManager argument that can be used to manage expansion of this * ArrowBuf * @return A new ArrowBuf that shares references with all ArrowBufs associated with this * BufferLedger */ ArrowBuf newArrowBuf(final long length, final @Nullable BufferManager manager) { allocator.assertOpen(); // the start virtual address of the ArrowBuf will be same as address of memory chunk final long startAddress = allocationManager.memoryAddress(); // create ArrowBuf final ArrowBuf buf = new ArrowBuf(this, manager, length, startAddress); // logging return loggingArrowBufHistoricalLog(buf); } private ArrowBuf loggingArrowBufHistoricalLog(ArrowBuf buf) { if (historicalLog != null) { historicalLog.recordEvent( "ArrowBuf(BufferLedger, BufferAllocator[%s], " + "UnsafeDirectLittleEndian[identityHashCode == " + "%d](%s)) => ledger hc == %d", allocator.getName(), System.identityHashCode(buf), buf.toString(), System.identityHashCode(this)); Preconditions.checkState(buffers != null, "IdentityHashMap of buffers must not be null"); synchronized (buffers) { buffers.put(buf, null); } } return buf; } /** * Create a new ArrowBuf that is associated with an alternative allocator for the purposes of * memory ownership and accounting. This has no impact on the reference counting for the current * ArrowBuf except in the situation where the passed in Allocator is the same as the current * buffer. * *

    This operation has no impact on the reference count of this ArrowBuf. The newly created * ArrowBuf with either have a reference count of 1 (in the case that this is the first time this * memory is being associated with the target allocator or in other words allocation manager * currently doesn't hold a mapping for the target allocator) or the current value of the * reference count for the target allocator-reference manager combination + 1 in the case that the * provided allocator already had an association to this underlying memory. * * @param srcBuffer source ArrowBuf * @param target The target allocator to create an association with. * @return A new ArrowBuf which shares the same underlying memory as the provided ArrowBuf. */ @Override public ArrowBuf retain(final ArrowBuf srcBuffer, BufferAllocator target) { if (historicalLog != null) { historicalLog.recordEvent("retain(%s)", target.getName()); } // the call to associate will return the corresponding reference manager (buffer ledger) for // the target allocator. if the allocation manager didn't already have a mapping // for the target allocator, it will create one and return the new reference manager with a // reference count of 1. Thus the newly created buffer in this case will have a ref count of 1. // alternatively, if there was already a mapping for in // allocation manager, the ref count of the new buffer will be targetrefmanager.refcount() + 1 // and this will be true for all the existing buffers currently managed by targetrefmanager final BufferLedger targetRefManager = allocationManager.associate(target); // create a new ArrowBuf to associate with new allocator and target ref manager final long targetBufLength = srcBuffer.capacity(); ArrowBuf targetArrowBuf = targetRefManager.deriveBuffer(srcBuffer, 0, targetBufLength); targetArrowBuf.readerIndex(srcBuffer.readerIndex()); targetArrowBuf.writerIndex(srcBuffer.writerIndex()); return targetArrowBuf; } /** * Transfer any balance the current ledger has to the target ledger. In the case that the current * ledger holds no memory, no transfer is made to the new ledger. * * @param targetReferenceManager The ledger to transfer ownership account to. * @return Whether transfer fit within target ledgers limits. */ boolean transferBalance(final @Nullable ReferenceManager targetReferenceManager) { Preconditions.checkArgument( targetReferenceManager != null, "Expecting valid target reference manager"); boolean overlimit = false; if (targetReferenceManager != null) { final BufferAllocator targetAllocator = targetReferenceManager.getAllocator(); Preconditions.checkArgument( allocator.getRoot() == targetAllocator.getRoot(), "You can only transfer between two allocators that share the same root."); allocator.assertOpen(); targetReferenceManager.getAllocator().assertOpen(); // if we're transferring to ourself, just return. if (targetReferenceManager == this) { return true; } // since two balance transfers out from the allocation manager could cause incorrect // accounting, we need to ensure // that this won't happen by synchronizing on the allocation manager instance. synchronized (allocationManager) { if (allocationManager.getOwningLedger() != this) { // since the calling reference manager is not the owning // reference manager for the underlying memory, transfer is // a NO-OP return true; } if (BaseAllocator.DEBUG && this.historicalLog != null) { this.historicalLog.recordEvent( "transferBalance(%s)", targetReferenceManager.getAllocator().getName()); } overlimit = targetAllocator.forceAllocate(allocationManager.getSize()); allocator.releaseBytes(allocationManager.getSize()); // since the transfer can only happen from the owning reference manager, // we need to set the target ref manager as the new owning ref manager // for the chunk of memory in allocation manager allocationManager.setOwningLedger((BufferLedger) targetReferenceManager); } } return overlimit; } /** * Transfer the memory accounting ownership of this ArrowBuf to another allocator. This will * generate a new ArrowBuf that carries an association with the underlying memory of this * ArrowBuf. If this ArrowBuf is connected to the owning BufferLedger of this memory, that memory * ownership/accounting will be transferred to the target allocator. If this ArrowBuf does not * currently own the memory underlying it (and is only associated with it), this does not transfer * any ownership to the newly created ArrowBuf. * *

    This operation has no impact on the reference count of this ArrowBuf. The newly created * ArrowBuf with either have a reference count of 1 (in the case that this is the first time this * memory is being associated with the new allocator) or the current value of the reference count * for the other AllocationManager/BufferLedger combination + 1 in the case that the provided * allocator already had an association to this underlying memory. * *

    Transfers will always succeed, even if that puts the other allocator into an overlimit * situation. This is possible due to the fact that the original owning allocator may have * allocated this memory out of a local reservation whereas the target allocator may need to * allocate new memory from a parent or RootAllocator. This operation is done n a mostly-lockless * but consistent manner. As such, the overlimit==true situation could occur slightly prematurely * to an actual overlimit==true condition. This is simply conservative behavior which means we may * return overlimit slightly sooner than is necessary. * * @param target The allocator to transfer ownership to. * @return A new transfer result with the impact of the transfer (whether it was overlimit) as * well as the newly created ArrowBuf. */ @Override public TransferResult transferOwnership(final ArrowBuf srcBuffer, final BufferAllocator target) { // the call to associate will return the corresponding reference manager (buffer ledger) for // the target allocator. if the allocation manager didn't already have a mapping // for the target allocator, it will create one and return the new reference manager with a // reference count of 1. Thus the newly created buffer in this case will have a ref count of 1. // alternatively, if there was already a mapping for in // allocation manager, the ref count of the new buffer will be targetrefmanager.refcount() + 1 // and this will be true for all the existing buffers currently managed by targetrefmanager final BufferLedger targetRefManager = allocationManager.associate(target); // create a new ArrowBuf to associate with new allocator and target ref manager final long targetBufLength = srcBuffer.capacity(); final ArrowBuf targetArrowBuf = targetRefManager.deriveBuffer(srcBuffer, 0, targetBufLength); targetArrowBuf.readerIndex(srcBuffer.readerIndex()); targetArrowBuf.writerIndex(srcBuffer.writerIndex()); final boolean allocationFit = transferBalance(targetRefManager); return new TransferResult(allocationFit, targetArrowBuf); } /** The outcome of a Transfer. */ public static class TransferResult implements OwnershipTransferResult { // Whether this transfer fit within the target allocator's capacity. final boolean allocationFit; // The newly created buffer associated with the target allocator public final ArrowBuf buffer; private TransferResult(boolean allocationFit, ArrowBuf buffer) { this.allocationFit = allocationFit; this.buffer = buffer; } @Override public ArrowBuf getTransferredBuffer() { return buffer; } @Override public boolean getAllocationFit() { return allocationFit; } } /** * Total size (in bytes) of memory underlying this reference manager. * * @return Size (in bytes) of the memory chunk */ @Override public long getSize() { return allocationManager.getSize(); } /** * How much memory is accounted for by this ledger. This is either getSize() if this is the owning * ledger for the memory or zero in the case that this is not the owning ledger associated with * this memory. * * @return Amount of accounted(owned) memory associated with this ledger. */ @Override public long getAccountedSize() { synchronized (allocationManager) { if (allocationManager.getOwningLedger() == this) { return allocationManager.getSize(); } else { return 0; } } } /** * Print the current ledger state to the provided StringBuilder. * * @param sb The StringBuilder to populate. * @param indent The level of indentation to position the data. * @param verbosity The level of verbosity to print. */ void print(StringBuilder sb, int indent, BaseAllocator.Verbosity verbosity) { CommonUtil.indent(sb, indent) .append("ledger[") .append(System.identityHashCode(this)) .append("] allocator: ") .append(allocator.getName()) .append("), isOwning: ") .append(", size: ") .append(", references: ") .append(bufRefCnt) .append(", life: ") .append(lCreationTime) .append("..") .append(lDestructionTime) .append(", allocatorManager: [") .append(", life: "); if (!BaseAllocator.DEBUG) { sb.append("]\n"); } else { Preconditions.checkArgument(buffers != null, "IdentityHashMap of buffers must not be null"); synchronized (buffers) { sb.append("] holds ").append(buffers.size()).append(" buffers. \n"); for (ArrowBuf buf : buffers.keySet()) { buf.print(sb, indent + 2, verbosity); sb.append('\n'); } } } } /** * Get the {@link AllocationManager} used by this BufferLedger. * * @return The AllocationManager used by this BufferLedger. */ public AllocationManager getAllocationManager() { return allocationManager; } } apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferManager.java000066400000000000000000000034331515451615200333240ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * Manages a list of {@link ArrowBuf}s that can be reallocated as needed. Upon re-allocation the old * buffer will be freed. Managing a list of these buffers prevents some parts of the system from * needing to define a correct location to place the final call to free them. */ public interface BufferManager extends AutoCloseable { /** * Replace an old buffer with a new version at least of the provided size. Does not copy data. * * @param old Old Buffer that the user is no longer going to use. * @param newSize Size of new replacement buffer. * @return A new version of the buffer. */ ArrowBuf replace(ArrowBuf old, long newSize); /** * Get a managed buffer of indeterminate size. * * @return A buffer. */ ArrowBuf getManagedBuffer(); /** * Get a managed buffer of at least a certain size. * * @param size The desired size * @return A buffer */ ArrowBuf getManagedBuffer(long size); @Override void close(); } CheckAllocator.java000066400000000000000000000111421515451615200334130ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.io.IOException; import java.net.URL; import java.util.Enumeration; import java.util.LinkedHashSet; import java.util.Set; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * Static method to ensure we have a RootAllocator on the classpath and report which one is used. */ final class CheckAllocator { private static final Logger logger = LoggerFactory.getLogger(CheckAllocator.class); // unique package names needed by JPMS module naming private static final String ALLOCATOR_PATH_CORE = "org/apache/arrow/memory/DefaultAllocationManagerFactory.class"; private static final String ALLOCATOR_PATH_UNSAFE = "org/apache/arrow/memory/unsafe/DefaultAllocationManagerFactory.class"; private static final String ALLOCATOR_PATH_NETTY = "org/apache/arrow/memory/netty/DefaultAllocationManagerFactory.class"; private CheckAllocator() {} static String check() { Set urls = scanClasspath(); URL rootAllocator = assertOnlyOne(urls); reportResult(rootAllocator); if (rootAllocator.getPath().contains("memory-core") || rootAllocator.getPath().contains("/org/apache/arrow/memory/core/")) { return "org.apache.arrow.memory.DefaultAllocationManagerFactory"; } else if (rootAllocator.getPath().contains("memory-unsafe") || rootAllocator.getPath().contains("/org/apache/arrow/memory/unsafe/")) { return "org.apache.arrow.memory.unsafe.DefaultAllocationManagerFactory"; } else if (rootAllocator.getPath().contains("memory-netty") || rootAllocator.getPath().contains("/org/apache/arrow/memory/netty/")) { return "org.apache.arrow.memory.netty.DefaultAllocationManagerFactory"; } else { throw new IllegalStateException( "Unknown allocation manager type to infer. Current: " + rootAllocator.getPath()); } } @SuppressWarnings("URLEqualsHashCode") private static Set scanClasspath() { // LinkedHashSet appropriate here because it preserves insertion order // during iteration Set allocatorPathSet = new LinkedHashSet<>(); try { ClassLoader allocatorClassLoader = CheckAllocator.class.getClassLoader(); Enumeration paths; if (allocatorClassLoader == null) { paths = ClassLoader.getSystemResources(ALLOCATOR_PATH_CORE); if (!paths.hasMoreElements()) { paths = ClassLoader.getSystemResources(ALLOCATOR_PATH_UNSAFE); } if (!paths.hasMoreElements()) { paths = ClassLoader.getSystemResources(ALLOCATOR_PATH_NETTY); } } else { paths = allocatorClassLoader.getResources(ALLOCATOR_PATH_CORE); if (!paths.hasMoreElements()) { paths = allocatorClassLoader.getResources(ALLOCATOR_PATH_UNSAFE); } if (!paths.hasMoreElements()) { paths = allocatorClassLoader.getResources(ALLOCATOR_PATH_NETTY); } } while (paths.hasMoreElements()) { URL path = paths.nextElement(); allocatorPathSet.add(path); } } catch (IOException ioe) { logger.error("Error getting resources from path", ioe); } return allocatorPathSet; } private static void reportResult(URL rootAllocator) { String path = rootAllocator.getPath(); String subPath = path.substring(path.indexOf("memory")); logger.info("Using DefaultAllocationManager at {}", subPath); } private static URL assertOnlyOne(Set urls) { if (urls.size() > 1) { logger.warn("More than one DefaultAllocationManager on classpath. Choosing first found"); } if (urls.isEmpty()) { throw new RuntimeException( "No DefaultAllocationManager found on classpath. Can't allocate Arrow buffers." + " Please consider adding arrow-memory-netty or arrow-memory-unsafe as a dependency."); } return urls.iterator().next(); } } ChildAllocator.java000066400000000000000000000030621515451615200334230ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * Child allocator class. Only slightly different from the {@link RootAllocator}, in that these * can't be created directly, but must be obtained from {@link * BufferAllocator#newChildAllocator(String, AllocationListener, long, long)}. * *

    Child allocators can only be created by the root, or other children, so this class is package * private. */ class ChildAllocator extends BaseAllocator { /** * Constructor. * * @param parentAllocator parent allocator -- the one creating this child * @param name the name of this child allocator * @param config configuration of this child allocator */ ChildAllocator(BaseAllocator parentAllocator, String name, Config config) { super(parentAllocator, name, config); } } DefaultAllocationManagerOption.java000066400000000000000000000120701515451615200366140ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.reflect.Field; import org.apache.arrow.util.VisibleForTesting; import org.checkerframework.checker.nullness.qual.Nullable; /** A class for choosing the default allocation manager. */ public class DefaultAllocationManagerOption { /** The environmental variable to set the default allocation manager type. */ public static final String ALLOCATION_MANAGER_TYPE_ENV_NAME = "ARROW_ALLOCATION_MANAGER_TYPE"; /** The system property to set the default allocation manager type. */ public static final String ALLOCATION_MANAGER_TYPE_PROPERTY_NAME = "arrow.allocation.manager.type"; static final org.slf4j.Logger LOGGER = org.slf4j.LoggerFactory.getLogger(DefaultAllocationManagerOption.class); /** The default allocation manager factory. */ private static AllocationManager.@Nullable Factory DEFAULT_ALLOCATION_MANAGER_FACTORY = null; /** The allocation manager type. */ public enum AllocationManagerType { /** Netty based allocation manager. */ Netty, /** Unsafe based allocation manager. */ Unsafe, /** Unknown type. */ Unknown, } /** * Returns the default allocation manager type. * * @return the default allocation manager type. */ @SuppressWarnings("nullness:argument") // enum types valueOf are implicitly non-null @VisibleForTesting public static AllocationManagerType getDefaultAllocationManagerType() { AllocationManagerType ret = AllocationManagerType.Unknown; try { String envValue = System.getenv(ALLOCATION_MANAGER_TYPE_ENV_NAME); ret = AllocationManagerType.valueOf(envValue); } catch (IllegalArgumentException | NullPointerException e) { // ignore the exception, and make the allocation manager type remain unchanged } // system property takes precedence try { String propValue = System.getProperty(ALLOCATION_MANAGER_TYPE_PROPERTY_NAME); ret = AllocationManagerType.valueOf(propValue); } catch (IllegalArgumentException | NullPointerException e) { // ignore the exception, and make the allocation manager type remain unchanged } return ret; } static AllocationManager.Factory getDefaultAllocationManagerFactory() { if (DEFAULT_ALLOCATION_MANAGER_FACTORY != null) { return DEFAULT_ALLOCATION_MANAGER_FACTORY; } AllocationManagerType type = getDefaultAllocationManagerType(); switch (type) { case Netty: DEFAULT_ALLOCATION_MANAGER_FACTORY = getNettyFactory(); break; case Unsafe: DEFAULT_ALLOCATION_MANAGER_FACTORY = getUnsafeFactory(); break; case Unknown: LOGGER.info("allocation manager type not specified, using netty as the default type"); DEFAULT_ALLOCATION_MANAGER_FACTORY = getFactory(CheckAllocator.check()); break; default: throw new IllegalStateException("Unknown allocation manager type: " + type); } return DEFAULT_ALLOCATION_MANAGER_FACTORY; } @SuppressWarnings({"nullness:argument", "nullness:return"}) // incompatible argument for parameter obj of Field.get // Static member qualifying type may not be annotated private static AllocationManager.Factory getFactory(String clazzName) { try { Field field = Class.forName(clazzName).getDeclaredField("FACTORY"); field.setAccessible(true); return (AllocationManager.Factory) field.get(null); } catch (Exception e) { throw new RuntimeException("Unable to instantiate Allocation Manager for " + clazzName, e); } } private static AllocationManager.Factory getUnsafeFactory() { try { return getFactory("org.apache.arrow.memory.unsafe.UnsafeAllocationManager"); } catch (RuntimeException e) { throw new RuntimeException( "Please add arrow-memory-unsafe to your classpath," + " No DefaultAllocationManager found to instantiate an UnsafeAllocationManager", e); } } private static AllocationManager.Factory getNettyFactory() { try { return getFactory("org.apache.arrow.memory.netty.NettyAllocationManager"); } catch (RuntimeException e) { throw new RuntimeException( "Please add arrow-memory-netty to your classpath," + " No DefaultAllocationManager found to instantiate an NettyAllocationManager", e); } } } ForeignAllocation.java000066400000000000000000000033471515451615200341440ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * EXPERIMENTAL: a memory allocation that does not come from a BufferAllocator, but rather an * outside source (like JNI). * *

    To use this, subclass this class and implement {@link #release0()} to free the allocation. */ public abstract class ForeignAllocation { private final long memoryAddress; private final long size; /** * Create a new AllocationManager representing an imported buffer. * * @param size The buffer size. * @param memoryAddress The buffer address. */ protected ForeignAllocation(long size, long memoryAddress) { this.memoryAddress = memoryAddress; this.size = size; } /** Get the size of this allocation. */ public long getSize() { return size; } /** Get the address of this allocation. */ protected long memoryAddress() { return memoryAddress; } /** Free this allocation. Will only be called once. */ protected abstract void release0(); } ForeignAllocationManager.java000066400000000000000000000026211515451615200354310ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** An AllocationManager wrapping a ForeignAllocation. */ class ForeignAllocationManager extends AllocationManager { private final ForeignAllocation allocation; protected ForeignAllocationManager( BufferAllocator accountingAllocator, ForeignAllocation allocation) { super(accountingAllocator); this.allocation = allocation; } @Override public long getSize() { return allocation.getSize(); } @Override protected long memoryAddress() { return allocation.memoryAddress(); } @Override protected void release0() { allocation.release0(); } } LowCostIdentityHashMap.java000066400000000000000000000231121515451615200351030ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.util.Preconditions; import org.apache.arrow.util.VisibleForTesting; import org.checkerframework.checker.initialization.qual.Initialized; import org.checkerframework.checker.initialization.qual.UnderInitialization; import org.checkerframework.checker.nullness.qual.Nullable; /** * Highly specialized IdentityHashMap that implements only partial Map APIs. It incurs low initial * cost (just two elements by default). It assumes Value includes the Key - * Implements @ValueWithKeyIncluded iface that provides "getKey" method. * * @param Key type * @param Value type */ public class LowCostIdentityHashMap> { /* * The internal data structure to hold values. */ private @Nullable Object[] elementData; // elementData[index] = null; /* Actual number of values. */ private int size; /* * maximum number of elements that can be put in this map before having to * rehash. */ private int threshold; private static final int DEFAULT_MIN_SIZE = 1; /* Default load factor of 0.75; */ private static final int LOAD_FACTOR = 7500; /** Creates a Map with default expected maximum size. */ public LowCostIdentityHashMap() { this(DEFAULT_MIN_SIZE); } /** * Creates a Map with the specified maximum size parameter. * * @param maxSize The estimated maximum number of entries that will be put in this map. */ public LowCostIdentityHashMap(int maxSize) { if (maxSize >= 0) { this.size = 0; threshold = getThreshold(maxSize); elementData = newElementArrayUnderInitialized(computeElementArraySize()); } else { throw new IllegalArgumentException(); } } private int getThreshold(@UnderInitialization LowCostIdentityHashMap this, int maxSize) { // assign the threshold to maxSize initially, this will change to a // higher value if rehashing occurs. return maxSize > 2 ? maxSize : 2; } private int computeElementArraySize(@UnderInitialization LowCostIdentityHashMap this) { int arraySize = (int) (((long) threshold * 10000) / LOAD_FACTOR); // ensure arraySize is positive, the above cast from long to int type // leads to overflow and negative arraySize if threshold is too big return arraySize < 0 ? -arraySize : arraySize; } /** * Create a new element array. * * @param s the number of elements * @return Reference to the element array */ private Object[] newElementArrayInitialized( @Initialized LowCostIdentityHashMap this, int s) { return new Object[s]; } /** * Create a new element array. * * @param s the number of elements * @return Reference to the element array */ private Object[] newElementArrayUnderInitialized( @UnderInitialization LowCostIdentityHashMap this, int s) { return new Object[s]; } /** * Removes all elements from this map, leaving it empty. * * @see #isEmpty() * @see #size() */ public void clear() { size = 0; for (int i = 0; i < elementData.length; i++) { elementData[i] = null; } } /** * Returns whether this map contains the specified key. * * @param key the key to search for. * @return {@code true} if this map contains the specified key, {@code false} otherwise. */ public boolean containsKey(K key) { Preconditions.checkNotNull(key); int index = findIndex(key, elementData); return (elementData[index] == null) ? false : ((V) elementData[index]).getKey() == key; } /** * Returns whether this map contains the specified value. * * @param value the value to search for. * @return {@code true} if this map contains the specified value, {@code false} otherwise. */ public boolean containsValue(V value) { Preconditions.checkNotNull(value); for (int i = 0; i < elementData.length; i++) { if (elementData[i] == value) { return true; } } return false; } /** * Returns the value of the mapping with the specified key. * * @param key the key. * @return the value of the mapping with the specified key. */ public @Nullable V get(K key) { Preconditions.checkNotNull(key); int index = findIndex(key, elementData); return (elementData[index] == null) ? null : (((V) elementData[index]).getKey() == key) ? (V) elementData[index] : null; } /** * Returns the index where the key is found at, or the index of the next empty spot if the key is * not found in this table. */ @VisibleForTesting int findIndex(@Nullable Object key, @Nullable Object[] array) { int length = array.length; int index = getModuloHash(key, length); int last = (index + length - 1) % length; while (index != last) { if ((array[index] == null) || ((V) array[index]).getKey() == key) { /* * Found the key, or the next empty spot (which means key is not * in the table) */ break; } index = (index + 1) % length; } return index; } @VisibleForTesting static int getModuloHash(@Nullable Object key, int length) { return ((System.identityHashCode(key) & 0x7FFFFFFF) % length); } /** * Maps the specified key to the specified value. * * @param value the value. * @return the value of any previous mapping with the specified key or {@code null} if there was * no such mapping. */ public V put(V value) { Preconditions.checkNotNull(value); K key = value.getKey(); Preconditions.checkNotNull(key); int index = findIndex(key, elementData); // if the key doesn't exist in the table if (elementData[index] == null || ((V) elementData[index]).getKey() != key) { if (++size > threshold) { rehash(); index = findIndex(key, elementData); } // insert the key and assign the value to null initially elementData[index] = null; } // insert value to where it needs to go, return the old value Object result = elementData[index]; elementData[index] = value; return (V) result; } @VisibleForTesting void rehash() { int newlength = elementData.length * 15 / 10; if (newlength == 0) { newlength = 1; } @Nullable Object[] newData = newElementArrayInitialized(newlength); for (int i = 0; i < elementData.length; i++) { Object key = (elementData[i] == null) ? null : ((V) elementData[i]).getKey(); if (key != null) { // if not empty int index = findIndex(key, newData); newData[index] = elementData[i]; } } elementData = newData; computeMaxSize(); } private void computeMaxSize() { threshold = (int) ((long) elementData.length * LOAD_FACTOR / 10000); } /** * Removes the mapping with the specified key from this map. * * @param key the key of the mapping to remove. * @return the value of the removed mapping, or {@code null} if no mapping for the specified key * was found. */ public @Nullable V remove(K key) { Preconditions.checkNotNull(key); boolean hashedOk; int index; int next; int hash; Object result; Object object; index = next = findIndex(key, elementData); if (elementData[index] == null || ((V) elementData[index]).getKey() != key) { return null; } // store the value for this key result = elementData[index]; // clear value to allow movement of the rest of the elements elementData[index] = null; size--; // shift the following elements up if needed // until we reach an empty spot int length = elementData.length; while (true) { next = (next + 1) % length; object = elementData[next]; if (object == null) { break; } hash = getModuloHash(((V) object).getKey(), length); hashedOk = hash > index; if (next < index) { hashedOk = hashedOk || (hash <= next); } else { hashedOk = hashedOk && (hash <= next); } if (!hashedOk) { elementData[index] = object; index = next; elementData[index] = null; } } return (V) result; } /** * Returns whether this Map has no elements. * * @return {@code true} if this Map has no elements, {@code false} otherwise. * @see #size() */ public boolean isEmpty() { return size == 0; } /** * Returns the number of mappings in this Map. * * @return the number of mappings in this Map. */ public int size() { return size; } /** * Special API to return next value - substitute of regular Map.values.iterator().next(). * * @return next available value or null if none available */ public @Nullable V getNextValue() { for (int i = 0; i < elementData.length; i++) { if (elementData[i] != null) { return (V) elementData[i]; } } return null; } } OutOfMemoryException.java000066400000000000000000000042661515451615200346520ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.Optional; /** * Indicates memory could not be allocated for Arrow buffers. * *

    This is different from {@linkplain OutOfMemoryError} which indicates the JVM is out of memory. * This error indicates that static limit of one of Arrow's allocators (e.g. {@linkplain * BaseAllocator}) has been exceeded. */ public class OutOfMemoryException extends RuntimeException { static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(OutOfMemoryException.class); private static final long serialVersionUID = -6858052345185793382L; private Optional outcomeDetails = Optional.empty(); public OutOfMemoryException() { super(); } public OutOfMemoryException( String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) { super(message, cause, enableSuppression, writableStackTrace); } public OutOfMemoryException(String message, Throwable cause) { super(message, cause); } public OutOfMemoryException(String message) { super(message); } public OutOfMemoryException(String message, Optional details) { super(message); this.outcomeDetails = details; } public OutOfMemoryException(Throwable cause) { super(cause); } public Optional getOutcomeDetails() { return outcomeDetails; } } OwnershipTransferNOOP.java000066400000000000000000000023301515451615200347130ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** An {@link OwnershipTransferResult} indicating no transfer needed. */ public class OwnershipTransferNOOP implements OwnershipTransferResult { private final ArrowBuf buffer; OwnershipTransferNOOP(final ArrowBuf buf) { this.buffer = buf; } @Override public ArrowBuf getTransferredBuffer() { return buffer; } @Override public boolean getAllocationFit() { return true; } } OwnershipTransferResult.java000066400000000000000000000020201515451615200354120ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** The result of transferring an {@link ArrowBuf} between {@linkplain BufferAllocator}s. */ public interface OwnershipTransferResult { boolean getAllocationFit(); ArrowBuf getTransferredBuffer(); } ReferenceManager.java000066400000000000000000000147041515451615200337350ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * ReferenceManager is the reference count for one or more allocations. * *

    In order to integrate with the core {@link BufferAllocator} implementation, the allocation * itself should be represented by an {@link AllocationManager}, though this is not required by the * API. */ public interface ReferenceManager { /** * Return the reference count. * * @return reference count */ int getRefCount(); /** * Decrement this reference manager's reference count by 1 for the associated underlying memory. * If the reference count drops to 0, it implies that ArrowBufs managed by this reference manager * no longer need access to the underlying memory * * @return true if ref count has dropped to 0, false otherwise */ boolean release(); /** * Decrement this reference manager's reference count for the associated underlying memory. If the * reference count drops to 0, it implies that ArrowBufs managed by this reference manager no * longer need access to the underlying memory * * @param decrement the count to decrease the reference count by * @return the new reference count */ boolean release(int decrement); /** * Increment this reference manager's reference count by 1 for the associated underlying memory. */ void retain(); /** * Increment this reference manager's reference count by a given amount for the associated * underlying memory. * * @param increment the count to increase the reference count by */ void retain(int increment); /** * Create a new ArrowBuf that is associated with an alternative allocator for the purposes of * memory ownership and accounting. This has no impact on the reference counting for the current * ArrowBuf except in the situation where the passed in Allocator is the same as the current * buffer. This operation has no impact on the reference count of this ArrowBuf. The newly created * ArrowBuf with either have a reference count of 1 (in the case that this is the first time this * memory is being associated with the target allocator or in other words allocation manager * currently doesn't hold a mapping for the target allocator) or the current value of the * reference count for the target allocator-reference manager combination + 1 in the case that the * provided allocator already had an association to this underlying memory. * *

    The underlying allocation ({@link AllocationManager}) will not be copied. * * @param srcBuffer source ArrowBuf * @param targetAllocator The target allocator to create an association with. * @return A new ArrowBuf which shares the same underlying memory as this ArrowBuf. */ ArrowBuf retain(ArrowBuf srcBuffer, BufferAllocator targetAllocator); /** * Derive a new ArrowBuf from a given source ArrowBuf. The new derived ArrowBuf will share the * same reference count as rest of the ArrowBufs associated with this reference manager. * * @param sourceBuffer source ArrowBuf * @param index index (relative to source ArrowBuf) new ArrowBuf should be derived from * @param length length (bytes) of data in underlying memory that derived buffer will have access * to in underlying memory * @return derived buffer */ ArrowBuf deriveBuffer(ArrowBuf sourceBuffer, long index, long length); /** * Duplicate the memory accounting ownership of the backing allocation of the given ArrowBuf in * another allocator. This will generate a new ArrowBuf that carries an association with the same * underlying memory ({@link AllocationManager}s) as the given ArrowBuf. * * @param sourceBuffer source ArrowBuf * @param targetAllocator The target allocator to create an association with * @return {@link OwnershipTransferResult} with info on transfer result and new buffer */ OwnershipTransferResult transferOwnership(ArrowBuf sourceBuffer, BufferAllocator targetAllocator); /** * Get the buffer allocator associated with this reference manager. * * @return buffer allocator. */ BufferAllocator getAllocator(); /** * Total size (in bytes) of memory underlying this reference manager. * * @return Size (in bytes) of the memory chunk. */ long getSize(); /** * Get the total accounted size (in bytes). * * @return accounted size. */ long getAccountedSize(); String NO_OP_ERROR_MESSAGE = "Operation not supported on NO_OP Reference Manager"; // currently used for empty ArrowBufs ReferenceManager NO_OP = new ReferenceManager() { @Override public int getRefCount() { return 1; } @Override public boolean release() { return false; } @Override public boolean release(int decrement) { return false; } @Override public void retain() {} @Override public void retain(int increment) {} @Override public ArrowBuf retain(ArrowBuf srcBuffer, BufferAllocator targetAllocator) { return srcBuffer; } @Override public ArrowBuf deriveBuffer(ArrowBuf sourceBuffer, long index, long length) { return sourceBuffer; } @Override public OwnershipTransferResult transferOwnership( ArrowBuf sourceBuffer, BufferAllocator targetAllocator) { return new OwnershipTransferNOOP(sourceBuffer); } @Override public BufferAllocator getAllocator() { return new RootAllocator(0); } @Override public long getSize() { return 0L; } @Override public long getAccountedSize() { return 0L; } }; } ReusableBuffer.java000066400000000000000000000031321515451615200334310ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * A lightweight, automatically expanding container for holding byte data. * * @param The type of the underlying buffer. */ public interface ReusableBuffer { /** * Get the number of valid bytes in the data. * * @return the number of valid bytes in the data */ long getLength(); /** Get the buffer backing this ReusableBuffer. */ T getBuffer(); /** * Set the buffer to the contents of the given ArrowBuf. The internal buffer must resize if it * cannot fit the contents of the data. * * @param srcBytes the data to copy from * @param start the first position of the new data * @param len the number of bytes of the new data */ void set(ArrowBuf srcBytes, long start, long len); void set(byte[] srcBytes, long start, long len); } apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/RootAllocator.java000066400000000000000000000044371515451615200334110ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.memory.rounding.DefaultRoundingPolicy; import org.apache.arrow.memory.rounding.RoundingPolicy; import org.apache.arrow.util.VisibleForTesting; /** * A root allocator for using direct memory for Arrow Vectors/Arrays. Supports creating a tree of * descendant child allocators to facilitate better instrumentation of memory allocations. */ public class RootAllocator extends BaseAllocator { public RootAllocator() { this(AllocationListener.NOOP, Long.MAX_VALUE); } public RootAllocator(final long limit) { this(AllocationListener.NOOP, limit); } public RootAllocator(final AllocationListener listener, final long limit) { // todo fix DefaultRoundingPolicy when using Netty this(listener, limit, DefaultRoundingPolicy.DEFAULT_ROUNDING_POLICY); } /** * Constructor. * * @param listener the allocation listener * @param limit max allocation size in bytes * @param roundingPolicy the policy for rounding the buffer size */ public RootAllocator( final AllocationListener listener, final long limit, RoundingPolicy roundingPolicy) { this( configBuilder() .listener(listener) .maxAllocation(limit) .roundingPolicy(roundingPolicy) .build()); } public RootAllocator(Config config) { super(null, "ROOT", config); } /** Verify the accounting state of the allocation system. */ @VisibleForTesting public void verify() { verifyAllocator(); } } ValueWithKeyIncluded.java000066400000000000000000000020161515451615200345660ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * Helper interface to generify a value to be included in the map where key is part of the value. * * @param The type of the key. */ public interface ValueWithKeyIncluded { K getKey(); } apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/package-info.java000066400000000000000000000017741515451615200331520ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT 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 Allocation, Accounting and Management. See the Arrow Java documentation for details: Memory Management */ package org.apache.arrow.memory; apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/000077500000000000000000000000001515451615200315775ustar00rootroot00000000000000DefaultRoundingPolicy.java000066400000000000000000000077041515451615200366450ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.rounding; import org.apache.arrow.memory.util.CommonUtil; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * The default rounding policy. That is, if the requested size is within the chunk size, the rounded * size will be the next power of two. Otherwise, the rounded size will be identical to the * requested size. */ public class DefaultRoundingPolicy implements RoundingPolicy { private static final Logger logger = LoggerFactory.getLogger(DefaultRoundingPolicy.class); public final long chunkSize; /** * The variables here and the static block calculates the DEFAULT_CHUNK_SIZE. * *

    It was copied from {@link io.netty.buffer.PooledByteBufAllocator}. */ private static final long MIN_PAGE_SIZE = 4096; private static final long MAX_CHUNK_SIZE = ((long) Integer.MAX_VALUE + 1) / 2; private static final long DEFAULT_CHUNK_SIZE; static { long defaultPageSize = Long.getLong("org.apache.memory.allocator.pageSize", 8192); try { validateAndCalculatePageShifts(defaultPageSize); } catch (Throwable t) { defaultPageSize = 8192; } int defaultMaxOrder = Integer.getInteger("org.apache.memory.allocator.maxOrder", 11); try { validateAndCalculateChunkSize(defaultPageSize, defaultMaxOrder); } catch (Throwable t) { defaultMaxOrder = 11; } DEFAULT_CHUNK_SIZE = validateAndCalculateChunkSize(defaultPageSize, defaultMaxOrder); if (logger.isDebugEnabled()) { logger.debug("-Dorg.apache.memory.allocator.pageSize: {}", defaultPageSize); logger.debug("-Dorg.apache.memory.allocator.maxOrder: {}", defaultMaxOrder); } } private static long validateAndCalculatePageShifts(long pageSize) { if (pageSize < MIN_PAGE_SIZE) { throw new IllegalArgumentException( "pageSize: " + pageSize + " (expected: " + MIN_PAGE_SIZE + ")"); } if ((pageSize & pageSize - 1) != 0) { throw new IllegalArgumentException("pageSize: " + pageSize + " (expected: power of 2)"); } // Logarithm base 2. At this point we know that pageSize is a power of two. return Long.SIZE - 1L - Long.numberOfLeadingZeros(pageSize); } private static long validateAndCalculateChunkSize(long pageSize, int maxOrder) { if (maxOrder > 14) { throw new IllegalArgumentException("maxOrder: " + maxOrder + " (expected: 0-14)"); } // Ensure the resulting chunkSize does not overflow. long chunkSize = pageSize; for (long i = maxOrder; i > 0; i--) { if (chunkSize > MAX_CHUNK_SIZE / 2) { throw new IllegalArgumentException( String.format( "pageSize (%d) << maxOrder (%d) must not exceed %d", pageSize, maxOrder, MAX_CHUNK_SIZE)); } chunkSize <<= 1; } return chunkSize; } /** The singleton instance. */ public static final DefaultRoundingPolicy DEFAULT_ROUNDING_POLICY = new DefaultRoundingPolicy(DEFAULT_CHUNK_SIZE); private DefaultRoundingPolicy(long chunkSize) { this.chunkSize = chunkSize; } @Override public long getRoundedSize(long requestSize) { return requestSize < chunkSize ? CommonUtil.nextPowerOfTwo(requestSize) : requestSize; } } RoundingPolicy.java000066400000000000000000000021501515451615200353260ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.rounding; /** * The policy for rounding the buffer size, to improve performance and avoid memory fragmentation. * In particular, given a requested buffer size, the policy will determine the rounded buffer size. */ public interface RoundingPolicy { long getRoundedSize(long requestSize); } SegmentRoundingPolicy.java000066400000000000000000000055051515451615200366600ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.rounding; import com.google.errorprone.annotations.InlineMe; import org.apache.arrow.memory.util.LargeMemoryUtil; import org.apache.arrow.util.Preconditions; /** The rounding policy that each buffer size must a multiple of the segment size. */ public class SegmentRoundingPolicy implements RoundingPolicy { /** The minimal segment size. */ public static final long MIN_SEGMENT_SIZE = 1024L; /** * The segment size. It must be at least {@link SegmentRoundingPolicy#MIN_SEGMENT_SIZE}, and be a * power of 2. */ private long segmentSize; /** * Constructor for the segment rounding policy. * * @param segmentSize the segment size. * @throws IllegalArgumentException if the segment size is smaller than {@link * SegmentRoundingPolicy#MIN_SEGMENT_SIZE}, or is not a power of 2. * @deprecated use {@link SegmentRoundingPolicy#SegmentRoundingPolicy(long)} instead. */ @Deprecated(forRemoval = true) @InlineMe(replacement = "this((long) segmentSize)") public SegmentRoundingPolicy(int segmentSize) { this((long) segmentSize); } /** * Constructor for the segment rounding policy. * * @param segmentSize the segment size. * @throws IllegalArgumentException if the segment size is smaller than {@link * SegmentRoundingPolicy#MIN_SEGMENT_SIZE}, or is not a power of 2. */ public SegmentRoundingPolicy(long segmentSize) { Preconditions.checkArgument( segmentSize >= MIN_SEGMENT_SIZE, "The segment size cannot be smaller than %s", MIN_SEGMENT_SIZE); Preconditions.checkArgument( (segmentSize & (segmentSize - 1)) == 0, "The segment size must be a power of 2"); this.segmentSize = segmentSize; } @Override public long getRoundedSize(long requestSize) { return (requestSize + (segmentSize - 1)) / segmentSize * segmentSize; } @Deprecated(forRemoval = true) public int getSegmentSize() { return LargeMemoryUtil.checkedCastToInt(segmentSize); } public long getSegmentSizeAsLong() { return segmentSize; } } apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/util/000077500000000000000000000000001515451615200307275ustar00rootroot00000000000000ArrowBufPointer.java000066400000000000000000000131261515451615200346060ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.memory.util.hash.SimpleHasher; import org.apache.arrow.util.Preconditions; import org.checkerframework.checker.nullness.qual.Nullable; /** * Pointer to a memory region within an {@link ArrowBuf}. It will be used as the basis for * calculating hash code within a vector, and equality determination. */ public final class ArrowBufPointer implements Comparable { /** The hash code when the arrow buffer is null. */ public static final int NULL_HASH_CODE = 0; private @Nullable ArrowBuf buf; private long offset; private long length; private int hashCode = NULL_HASH_CODE; private final ArrowBufHasher hasher; /** A flag indicating if the underlying memory region has changed. */ private boolean hashCodeChanged = false; /** The default constructor. */ public ArrowBufPointer() { this(SimpleHasher.INSTANCE); } /** * Constructs an arrow buffer pointer with the specified hasher. * * @param hasher the hasher to use. */ public ArrowBufPointer(ArrowBufHasher hasher) { Preconditions.checkNotNull(hasher); this.hasher = hasher; this.buf = null; } /** * Constructs an Arrow buffer pointer. * * @param buf the underlying {@link ArrowBuf}, which can be null. * @param offset the start off set of the memory region pointed to. * @param length the length off set of the memory region pointed to. */ public ArrowBufPointer(ArrowBuf buf, long offset, long length) { this(buf, offset, length, SimpleHasher.INSTANCE); } /** * Constructs an Arrow buffer pointer. * * @param buf the underlying {@link ArrowBuf}, which can be null. * @param offset the start off set of the memory region pointed to. * @param length the length off set of the memory region pointed to. * @param hasher the hasher used to calculate the hash code. */ public ArrowBufPointer(ArrowBuf buf, long offset, long length, ArrowBufHasher hasher) { Preconditions.checkNotNull(hasher); this.hasher = hasher; set(buf, offset, length); } /** * Sets this pointer. * * @param buf the underlying {@link ArrowBuf}, which can be null. * @param offset the start off set of the memory region pointed to. * @param length the length off set of the memory region pointed to. */ public void set(ArrowBuf buf, long offset, long length) { this.buf = buf; this.offset = offset; this.length = length; hashCodeChanged = true; } /** * Gets the underlying buffer, or null if the underlying data is invalid or null. * * @return the underlying buffer, if any, or null if the underlying data is invalid or null. */ public @Nullable ArrowBuf getBuf() { return buf; } public long getOffset() { return offset; } public long getLength() { return length; } @Override public boolean equals(@Nullable Object o) { if (this == o) { return true; } if (o == null || getClass() != o.getClass()) { return false; } if (!hasher.equals(((ArrowBufPointer) o).hasher)) { // note that the hasher is incorporated in equality determination // this is to avoid problems in cases where two Arrow buffer pointers are not equal // while having different hashers and equal hash codes. return false; } ArrowBufPointer other = (ArrowBufPointer) o; if (buf == null || other.buf == null) { if (buf == null && other.buf == null) { return true; } else { return false; } } return ByteFunctionHelpers.equal( buf, offset, offset + length, other.buf, other.offset, other.offset + other.length) != 0; } @Override public int hashCode() { if (!hashCodeChanged) { return hashCode; } // re-compute the hash code if (buf == null) { hashCode = NULL_HASH_CODE; } else { hashCode = hasher.hashCode(buf, offset, length); } hashCodeChanged = false; return hashCode; } /** * Compare two arrow buffer pointers. The comparison is based on lexicographic order. * * @param that the other pointer to compare. * @return 0 if the two pointers are equal; a positive integer if this pointer is larger; a * negative integer if this pointer is smaller. */ @Override public int compareTo(ArrowBufPointer that) { if (this.buf == null || that.buf == null) { if (this.buf == null && that.buf == null) { return 0; } else { // null is smaller return this.buf == null ? -1 : 1; } } return ByteFunctionHelpers.compare( this.buf, this.offset, this.offset + this.length, that.buf, that.offset, that.offset + that.length); } } AssertionUtil.java000066400000000000000000000025311515451615200343210ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util; /** * Utility class to that provides {@link #ASSERT_ENABLED} constant to determine if assertions are * enabled. */ public class AssertionUtil { public static final boolean ASSERT_ENABLED; static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(AssertionUtil.class); static { boolean isAssertEnabled = false; assert isAssertEnabled = true; ASSERT_ENABLED = isAssertEnabled; } private AssertionUtil() {} public static boolean isAssertionsEnabled() { return ASSERT_ENABLED; } } AutoCloseableLock.java000066400000000000000000000023601515451615200350470ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util; import java.util.concurrent.locks.Lock; /** Simple wrapper class that allows Locks to be released via a try-with-resources block. */ public class AutoCloseableLock implements AutoCloseable { private final Lock lock; public AutoCloseableLock(Lock lock) { this.lock = lock; } public AutoCloseableLock open() { lock.lock(); return this; } @Override public void close() { lock.unlock(); } } ByteFunctionHelpers.java000066400000000000000000000234501515451615200354530ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util; import java.nio.ByteOrder; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BoundsChecking; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.memory.util.hash.SimpleHasher; /** Utility methods for memory comparison at a byte level. */ public class ByteFunctionHelpers { static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(ByteFunctionHelpers.class); private static final boolean LITTLE_ENDIAN = ByteOrder.nativeOrder() == ByteOrder.LITTLE_ENDIAN; private ByteFunctionHelpers() {} /** * Helper function to check for equality of bytes in two ArrowBufs. * * @param left Left ArrowBuf for comparison * @param lStart start offset in the buffer * @param lEnd end offset in the buffer * @param right Right ArrowBuf for comparison * @param rStart start offset in the buffer * @param rEnd end offset in the buffer * @return 1 if equals, 0 otherwise */ public static int equal( final ArrowBuf left, long lStart, long lEnd, final ArrowBuf right, long rStart, long rEnd) { if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { left.checkBytes(lStart, lEnd); right.checkBytes(rStart, rEnd); } return memEqual(left.memoryAddress(), lStart, lEnd, right.memoryAddress(), rStart, rEnd); } private static int memEqual( final long laddr, long lStart, long lEnd, final long raddr, long rStart, final long rEnd) { long n = lEnd - lStart; if (n == rEnd - rStart) { long lPos = laddr + lStart; long rPos = raddr + rStart; while (n > 63) { for (int x = 0; x < 8; x++) { long leftLong = MemoryUtil.getLong(lPos); long rightLong = MemoryUtil.getLong(rPos); if (leftLong != rightLong) { return 0; } lPos += 8; rPos += 8; } n -= 64; } while (n > 7) { long leftLong = MemoryUtil.getLong(lPos); long rightLong = MemoryUtil.getLong(rPos); if (leftLong != rightLong) { return 0; } lPos += 8; rPos += 8; n -= 8; } if (n > 3) { int leftInt = MemoryUtil.getInt(lPos); int rightInt = MemoryUtil.getInt(rPos); if (leftInt != rightInt) { return 0; } lPos += 4; rPos += 4; n -= 4; } while (n-- != 0) { byte leftByte = MemoryUtil.getByte(lPos); byte rightByte = MemoryUtil.getByte(rPos); if (leftByte != rightByte) { return 0; } lPos++; rPos++; } return 1; } else { return 0; } } /** * Helper function to compare a set of bytes in two ArrowBufs. * *

    Function will check data before completing in the case that * * @param left Left ArrowBuf to compare * @param lStart start offset in the buffer * @param lEnd end offset in the buffer * @param right Right ArrowBuf to compare * @param rStart start offset in the buffer * @param rEnd end offset in the buffer * @return 1 if left input is greater, -1 if left input is smaller, 0 otherwise */ public static int compare( final ArrowBuf left, long lStart, long lEnd, final ArrowBuf right, long rStart, long rEnd) { if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { left.checkBytes(lStart, lEnd); right.checkBytes(rStart, rEnd); } return memcmp(left.memoryAddress(), lStart, lEnd, right.memoryAddress(), rStart, rEnd); } private static int memcmp( final long laddr, long lStart, long lEnd, final long raddr, long rStart, final long rEnd) { long lLen = lEnd - lStart; long rLen = rEnd - rStart; long n = Math.min(rLen, lLen); long lPos = laddr + lStart; long rPos = raddr + rStart; while (n > 63) { for (int x = 0; x < 8; x++) { long leftLong = MemoryUtil.getLong(lPos); long rightLong = MemoryUtil.getLong(rPos); if (leftLong != rightLong) { if (LITTLE_ENDIAN) { return unsignedLongCompare(Long.reverseBytes(leftLong), Long.reverseBytes(rightLong)); } else { return unsignedLongCompare(leftLong, rightLong); } } lPos += 8; rPos += 8; } n -= 64; } while (n > 7) { long leftLong = MemoryUtil.getLong(lPos); long rightLong = MemoryUtil.getLong(rPos); if (leftLong != rightLong) { if (LITTLE_ENDIAN) { return unsignedLongCompare(Long.reverseBytes(leftLong), Long.reverseBytes(rightLong)); } else { return unsignedLongCompare(leftLong, rightLong); } } lPos += 8; rPos += 8; n -= 8; } if (n > 3) { int leftInt = MemoryUtil.getInt(lPos); int rightInt = MemoryUtil.getInt(rPos); if (leftInt != rightInt) { if (LITTLE_ENDIAN) { return unsignedIntCompare(Integer.reverseBytes(leftInt), Integer.reverseBytes(rightInt)); } else { return unsignedIntCompare(leftInt, rightInt); } } lPos += 4; rPos += 4; n -= 4; } while (n-- != 0) { byte leftByte = MemoryUtil.getByte(lPos); byte rightByte = MemoryUtil.getByte(rPos); if (leftByte != rightByte) { return ((leftByte & 0xFF) - (rightByte & 0xFF)) > 0 ? 1 : -1; } lPos++; rPos++; } if (lLen == rLen) { return 0; } return lLen > rLen ? 1 : -1; } /** * Helper function to compare a set of bytes in ArrowBuf to a ByteArray. * * @param left Left ArrowBuf for comparison purposes * @param lStart start offset in the buffer * @param lEnd end offset in the buffer * @param right second input to be compared * @param rStart start offset in the byte array * @param rEnd end offset in the byte array * @return 1 if left input is greater, -1 if left input is smaller, 0 otherwise */ public static int compare( final ArrowBuf left, int lStart, int lEnd, final byte[] right, int rStart, final int rEnd) { if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { left.checkBytes(lStart, lEnd); } return memcmp(left.memoryAddress(), lStart, lEnd, right, rStart, rEnd); } /** * Compares the two specified {@code long} values, treating them as unsigned values between {@code * 0} and {@code 2^64 - 1} inclusive. * * @param a the first unsigned {@code long} to compare * @param b the second unsigned {@code long} to compare * @return a negative value if {@code a} is less than {@code b}; a positive value if {@code a} is * greater than {@code b}; or zero if they are equal */ public static int unsignedLongCompare(long a, long b) { return Long.compare(a ^ Long.MIN_VALUE, b ^ Long.MIN_VALUE); } public static int unsignedIntCompare(int a, int b) { return Integer.compare(a ^ Integer.MIN_VALUE, b ^ Integer.MIN_VALUE); } private static int memcmp( final long laddr, int lStart, int lEnd, final byte[] right, int rStart, final int rEnd) { int lLen = lEnd - lStart; int rLen = rEnd - rStart; int n = Math.min(rLen, lLen); long lPos = laddr + lStart; int rPos = rStart; while (n > 7) { long leftLong = MemoryUtil.getLong(lPos); long rightLong = MemoryUtil.getLong(right, rPos); if (leftLong != rightLong) { if (LITTLE_ENDIAN) { return unsignedLongCompare(Long.reverseBytes(leftLong), Long.reverseBytes(rightLong)); } else { return unsignedLongCompare(leftLong, rightLong); } } lPos += 8; rPos += 8; n -= 8; } if (n > 3) { int leftInt = MemoryUtil.getInt(lPos); int rightInt = MemoryUtil.getInt(right, rPos); if (leftInt != rightInt) { if (LITTLE_ENDIAN) { return unsignedIntCompare(Integer.reverseBytes(leftInt), Integer.reverseBytes(rightInt)); } else { return unsignedIntCompare(leftInt, rightInt); } } lPos += 4; rPos += 4; n -= 4; } while (n-- != 0) { byte leftByte = MemoryUtil.getByte(lPos); byte rightByte = right[rPos]; if (leftByte != rightByte) { return ((leftByte & 0xFF) - (rightByte & 0xFF)) > 0 ? 1 : -1; } lPos++; rPos++; } if (lLen == rLen) { return 0; } return lLen > rLen ? 1 : -1; } /** Compute hashCode with the given {@link ArrowBuf} and start/end index. */ public static int hash(final ArrowBuf buf, long start, long end) { return hash(SimpleHasher.INSTANCE, buf, start, end); } /** * Compute hashCode with the given {@link ArrowBufHasher}, {@link ArrowBuf} and start/end index. */ public static final int hash(ArrowBufHasher hasher, final ArrowBuf buf, long start, long end) { if (hasher == null) { hasher = SimpleHasher.INSTANCE; } return hasher.hashCode(buf, start, end - start); } /** Generate a new hashCode with the given current hashCode and new hashCode. */ public static int combineHash(int currentHash, int newHash) { return currentHash * 31 + newHash; } } CommonUtil.java000066400000000000000000000043141515451615200336030ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util; import java.util.Arrays; /** Utilities and static methods needed for arrow-memory. */ public final class CommonUtil { private CommonUtil() {} /** * Rounds up the provided value to the nearest power of two. * * @param val An integer value. * @return The closest power of two of that value. */ public static int nextPowerOfTwo(int val) { if (val == 0 || val == 1) { return val + 1; } int highestBit = Integer.highestOneBit(val); if (highestBit == val) { return val; } else { return highestBit << 1; } } /** * Rounds up the provided value to the nearest power of two. * * @param val A long value. * @return The closest power of two of that value. */ public static long nextPowerOfTwo(long val) { if (val == 0 || val == 1) { return val + 1; } long highestBit = Long.highestOneBit(val); if (highestBit == val) { return val; } else { return highestBit << 1; } } /** * Specify an indentation amount when using a StringBuilder. * * @param sb StringBuilder to use * @param indent Indentation amount * @return the StringBuilder object with indentation applied */ public static StringBuilder indent(StringBuilder sb, int indent) { final char[] indentation = new char[indent * 2]; Arrays.fill(indentation, ' '); sb.append(indentation); return sb; } } apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/util/Float16.java000066400000000000000000000252771515451615200330230ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util; import org.apache.arrow.util.VisibleForTesting; /** * Lifted from Apache Parquet MR project: * https://github.com/apache/parquet-mr/blob/e87b80308869b77f914fcfd04364686e11158950/parquet-column/src/main/java/org/apache/parquet/schema/Float16.java * *

      * Changes made: *
    • Modify the data type input from Parquet-MR Binary (toFloat(Binary b)) to Arrow Java short * (toFloat(short b)) *
    • Expose NAN and POSITIVE_INFINITY variables *
    * * The class is a utility class to manipulate half-precision 16-bit IEEE 754 floating * point data types (also called fp16 or binary16). A half-precision float can be created from or * converted to single-precision floats, and is stored in a short data type. The IEEE 754 standard * specifies an float16 as having the following format: * *
      *
    • Sign bit: 1 bit *
    • Exponent width: 5 bits *
    • Significand: 10 bits *
    * *

    The format is laid out as follows: * * 

     * 1   11111   1111111111
 * ^   --^--   -----^----
 * sign  |          |_______ significand
 *       |
 *      -- exponent
 *
    * * Half-precision floating points can be useful to save memory and/or bandwidth at the expense of * range and precision when compared to single-precision floating points (float32). Ref: * https://android.googlesource.com/platform/libcore/+/master/luni/src/main/java/libcore/util/FP16.java */ public class Float16 { // Positive infinity of type half-precision float. public static final short POSITIVE_INFINITY = (short) 0x7c00; // A Not-a-Number representation of a half-precision float. public static final short NaN = (short) 0x7e00; // The bitmask to and a number with to obtain the sign bit. private static final int SIGN_MASK = 0x8000; // The offset to shift by to obtain the exponent bits. private static final int EXPONENT_SHIFT = 10; // The bitmask to and a number shifted by EXPONENT_SHIFT right, to obtain exponent bits. private static final int SHIFTED_EXPONENT_MASK = 0x1f; // The bitmask to and a number with to obtain significand bits. private static final int SIGNIFICAND_MASK = 0x3ff; // The offset of the exponent from the actual value. private static final int EXPONENT_BIAS = 15; // The offset to shift by to obtain the sign bit. private static final int SIGN_SHIFT = 15; // The bitmask to AND with to obtain exponent and significand bits. private static final int EXPONENT_SIGNIFICAND_MASK = 0x7fff; private static final int FP32_SIGN_SHIFT = 31; private static final int FP32_EXPONENT_SHIFT = 23; private static final int FP32_SHIFTED_EXPONENT_MASK = 0xff; private static final int FP32_SIGNIFICAND_MASK = 0x7fffff; private static final int FP32_EXPONENT_BIAS = 127; private static final int FP32_QNAN_MASK = 0x400000; private static final int FP32_DENORMAL_MAGIC = 126 << 23; private static final float FP32_DENORMAL_FLOAT = Float.intBitsToFloat(FP32_DENORMAL_MAGIC); /** * Returns true if the specified half-precision float value represents a Not-a-Number, false * otherwise. * * @param h A half-precision float value * @return True if the value is a NaN, false otherwise */ @VisibleForTesting public static boolean isNaN(short h) { return (h & EXPONENT_SIGNIFICAND_MASK) > POSITIVE_INFINITY; } /** * Compares the two specified half-precision float values. The following conditions apply during * the comparison: * *
      *
    • NaN is considered by this method to be equal to itself and greater than all other * half-precision float values (including {@code #POSITIVE_INFINITY}) *
    • POSITIVE_ZERO is considered by this method to be greater than NEGATIVE_ZERO. *
    * * @param x The first half-precision float value to compare. * @param y The second half-precision float value to compare * @return The value {@code 0} if {@code x} is numerically equal to {@code y}, a value less than * {@code 0} if {@code x} is numerically less than {@code y}, and a value greater than {@code * 0} if {@code x} is numerically greater than {@code y} */ @VisibleForTesting public static int compare(short x, short y) { boolean xIsNaN = isNaN(x); boolean yIsNaN = isNaN(y); if (!xIsNaN && !yIsNaN) { int first = ((x & SIGN_MASK) != 0 ? 0x8000 - (x & 0xffff) : x & 0xffff); int second = ((y & SIGN_MASK) != 0 ? 0x8000 - (y & 0xffff) : y & 0xffff); // Returns true if the first half-precision float value is less // (smaller toward negative infinity) than the second half-precision float value. if (first < second) { return -1; } // Returns true if the first half-precision float value is greater // (larger toward positive infinity) than the second half-precision float value. if (first > second) { return 1; } } // Collapse NaNs, akin to halfToIntBits(), but we want to keep // (signed) short value types to preserve the ordering of -0.0 // and +0.0 short xBits = xIsNaN ? NaN : x; short yBits = yIsNaN ? NaN : y; return (xBits == yBits ? 0 : (xBits < yBits ? -1 : 1)); } /** * Converts the specified half-precision float value into a single-precision float value. The * following special cases are handled: If the input is NaN, the returned value is Float NaN. If * the input is POSITIVE_INFINITY or NEGATIVE_INFINITY, the returned value is respectively Float * POSITIVE_INFINITY or Float NEGATIVE_INFINITY. If the input is 0 (positive or negative), the * returned value is +/-0.0f. Otherwise, the returned value is a normalized single-precision float * value. * * @param b The half-precision float value to convert to single-precision * @return A normalized single-precision float value */ @VisibleForTesting public static float toFloat(short b) { int bits = b & 0xffff; int s = bits & SIGN_MASK; int e = (bits >>> EXPONENT_SHIFT) & SHIFTED_EXPONENT_MASK; int m = bits & SIGNIFICAND_MASK; int outE = 0; int outM = 0; if (e == 0) { // Denormal or 0 if (m != 0) { // Convert denorm fp16 into normalized fp32 float o = Float.intBitsToFloat(FP32_DENORMAL_MAGIC + m); o -= FP32_DENORMAL_FLOAT; return s == 0 ? o : -o; } } else { outM = m << 13; if (e == 0x1f) { // Infinite or NaN outE = 0xff; if (outM != 0) { // SNaNs are quieted outM |= FP32_QNAN_MASK; } } else { outE = e - EXPONENT_BIAS + FP32_EXPONENT_BIAS; } } int out = (s << 16) | (outE << FP32_EXPONENT_SHIFT) | outM; return Float.intBitsToFloat(out); } /** * Converts the specified single-precision float value into a half-precision float value. The * following special cases are handled: * *

    If the input is NaN, the returned value is NaN. If the input is Float POSITIVE_INFINITY or * Float NEGATIVE_INFINITY, the returned value is respectively POSITIVE_INFINITY or * NEGATIVE_INFINITY. If the input is 0 (positive or negative), the returned value is * POSITIVE_ZERO or NEGATIVE_ZERO. If the input is a less than MIN_VALUE, the returned value is * flushed to POSITIVE_ZERO or NEGATIVE_ZERO. If the input is a less than MIN_NORMAL, the returned * value is a denorm half-precision float. Otherwise, the returned value is rounded to the nearest * representable half-precision float value. * * @param f The single-precision float value to convert to half-precision * @return A half-precision float value */ public static short toFloat16(float f) { int bits = Float.floatToRawIntBits(f); int s = (bits >>> FP32_SIGN_SHIFT); int e = (bits >>> FP32_EXPONENT_SHIFT) & FP32_SHIFTED_EXPONENT_MASK; int m = bits & FP32_SIGNIFICAND_MASK; int outE = 0; int outM = 0; if (e == 0xff) { // Infinite or NaN outE = 0x1f; outM = m != 0 ? 0x200 : 0; } else { e = e - FP32_EXPONENT_BIAS + EXPONENT_BIAS; if (e >= 0x1f) { // Overflow outE = 0x1f; } else if (e <= 0) { // Underflow if (e < -10) { // The absolute fp32 value is less than MIN_VALUE, flush to +/-0 } else { // The fp32 value is a normalized float less than MIN_NORMAL, // we convert to a denorm fp16 m = m | 0x800000; int shift = 14 - e; outM = m >> shift; int lowm = m & ((1 << shift) - 1); int hway = 1 << (shift - 1); // if above halfway or exactly halfway and outM is odd if (lowm + (outM & 1) > hway) { // Round to nearest even // Can overflow into exponent bit, which surprisingly is OK. // This increment relies on the +outM in the return statement below outM++; } } } else { outE = e; outM = m >> 13; // if above halfway or exactly halfway and outM is odd if ((m & 0x1fff) + (outM & 0x1) > 0x1000) { // Round to nearest even // Can overflow into exponent bit, which surprisingly is OK. // This increment relies on the +outM in the return statement below outM++; } } } // The outM is added here as the +1 increments for outM above can // cause an overflow in the exponent bit which is OK. return (short) ((s << SIGN_SHIFT) | (outE << EXPONENT_SHIFT) + outM); } /** * Returns a string representation of the specified half-precision float value. Calling this * method is equivalent to calling Float.toString(toFloat(h)). See {@link * Float#toString(float)} for more information on the format of the string representation. * * @param h A half-precision float value in binary little-endian format * @return A string representation of the specified value */ @VisibleForTesting public static String toFloatString(short h) { return Float.toString(Float16.toFloat(h)); } } HistoricalLog.java000066400000000000000000000152341515451615200342630ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util; import com.google.errorprone.annotations.FormatMethod; import com.google.errorprone.annotations.FormatString; import java.util.ArrayDeque; import java.util.Arrays; import java.util.Deque; import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; /** * Utility class that can be used to log activity within a class for later logging and debugging. * Supports recording events and recording the stack at the time they occur. */ public class HistoricalLog { private final Deque history = new ArrayDeque<>(); private final String idString; // the formatted id string private final int limit; // the limit on the number of events kept private @Nullable Event firstEvent; // the first stack trace recorded /** * Constructor. The format string will be formatted and have its arguments substituted at the time * this is called. * * @param idStringFormat {@link String#format} format string that can be used to identify this * object in a log. Including some kind of unique identifier that can be associated with the * object instance is best. * @param args for the format string, or nothing if none are required */ @FormatMethod public HistoricalLog(@FormatString final String idStringFormat, Object... args) { this(Integer.MAX_VALUE, idStringFormat, args); } /** * Constructor. The format string will be formatted and have its arguments substituted at the time * this is called. * *

    This form supports the specification of a limit that will limit the number of historical * entries kept (which keeps down the amount of memory used). With the limit, the first entry made * is always kept (under the assumption that this is the creation site of the object, which is * usually interesting), and then up to the limit number of entries are kept after that. Each time * a new entry is made, the oldest that is not the first is dropped. * * @param limit the maximum number of historical entries that will be kept, not including the * first entry made * @param idStringFormat {@link String#format} format string that can be used to identify this * object in a log. Including some kind of unique identifier that can be associated with the * object instance is best. * @param args for the format string, or nothing if none are required */ @FormatMethod public HistoricalLog(final int limit, @FormatString final String idStringFormat, Object... args) { this.limit = limit; this.idString = String.format(idStringFormat, args); this.firstEvent = null; } /** * Record an event. Automatically captures the stack trace at the time this is called. The format * string will be formatted and have its arguments substituted at the time this is called. * * @param noteFormat {@link String#format} format string that describes the event * @param args for the format string, or nothing if none are required */ @FormatMethod public synchronized void recordEvent(@FormatString final String noteFormat, Object... args) { final String note = String.format(noteFormat, args); final Event event = new Event(note); if (firstEvent == null) { firstEvent = event; } if (history.size() == limit) { history.removeFirst(); } history.add(event); } /** * Write the history of this object to the given {@link StringBuilder}. The history includes the * identifying string provided at construction time, and all the recorded events with their stack * traces. * * @param sb {@link StringBuilder} to write to * @param includeStackTrace whether to include the stacktrace of each event in the history */ public void buildHistory(final StringBuilder sb, boolean includeStackTrace) { buildHistory(sb, 0, includeStackTrace); } /** * Build the history and write it to sb. * * @param sb output * @param indent starting indent (usually "") * @param includeStackTrace whether to include the stacktrace of each event. */ public synchronized void buildHistory( final StringBuilder sb, int indent, boolean includeStackTrace) { final char[] indentation = new char[indent]; final char[] innerIndentation = new char[indent + 2]; Arrays.fill(indentation, ' '); Arrays.fill(innerIndentation, ' '); sb.append(indentation).append("event log for: ").append(idString).append('\n'); if (firstEvent != null) { long time = firstEvent.time; String note = firstEvent.note; final StackTrace stackTrace = firstEvent.stackTrace; sb.append(innerIndentation).append(time).append(' ').append(note).append('\n'); if (includeStackTrace) { stackTrace.writeToBuilder(sb, indent + 2); } for (final Event event : history) { if (event == firstEvent) { continue; } sb.append(innerIndentation) .append(" ") .append(event.time) .append(' ') .append(event.note) .append('\n'); if (includeStackTrace) { event.stackTrace.writeToBuilder(sb, indent + 2); sb.append('\n'); } } } } /** * Write the history of this object to the given {@link Logger}. The history includes the * identifying string provided at construction time, and all the recorded events with their stack * traces. * * @param logger {@link Logger} to write to */ public void logHistory(final Logger logger) { final StringBuilder sb = new StringBuilder(); buildHistory(sb, 0, true); logger.debug(sb.toString()); } private static class Event { private final String note; // the event text private final StackTrace stackTrace; // where the event occurred private final long time; public Event(final String note) { this.note = note; this.time = System.nanoTime(); stackTrace = new StackTrace(); } } } LargeMemoryUtil.java000066400000000000000000000026611515451615200346010ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util; import org.apache.arrow.memory.BoundsChecking; /** Contains utilities for dealing with a 64-bit address base. */ public final class LargeMemoryUtil { private LargeMemoryUtil() {} /** Casts length to an int, but raises an exception the value is outside the range of an int. */ public static int checkedCastToInt(long length) { if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { return Math.toIntExact(length); } return (int) length; } /** Returns a min(Integer.MAX_VALUE, length). */ public static int capAtMaxInt(long length) { return (int) Math.min(length, Integer.MAX_VALUE); } } MemoryUtil.java000066400000000000000000000240311515451615200336210ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util; 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; import java.security.AccessController; import java.security.PrivilegedAction; import org.checkerframework.checker.nullness.qual.Nullable; import sun.misc.Unsafe; /** Utilities for memory related operations. */ public class MemoryUtil { private static final org.slf4j.Logger logger = 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; /** The start offset of array data relative to the start address of the array object. */ private static final long BYTE_ARRAY_BASE_OFFSET; /** The offset of the address field with the {@link java.nio.ByteBuffer} object. */ private static final long BYTE_BUFFER_ADDRESS_OFFSET; /** If the native byte order is little-endian. */ public static final boolean LITTLE_ENDIAN = ByteOrder.nativeOrder() == ByteOrder.LITTLE_ENDIAN; // Java 1.8, 9, 11, 17, 21 becomes 1, 9, 11, 17, and 21. @SuppressWarnings("StringSplitter") private static final int majorVersion = Integer.parseInt(System.getProperty("java.specification.version").split("\\D+")[0]); static { try { // try to get the unsafe object final Object maybeUnsafe = AccessController.doPrivileged( new PrivilegedAction() { @Override @SuppressWarnings({"nullness:argument", "nullness:return"}) // incompatible argument for parameter obj of Field.get // incompatible types in return public Object run() { try { final Field unsafeField = Unsafe.class.getDeclaredField("theUnsafe"); unsafeField.setAccessible(true); return unsafeField.get(null); } catch (Throwable e) { return e; } } }); if (maybeUnsafe instanceof Throwable) { throw (Throwable) maybeUnsafe; } UNSAFE = (Unsafe) maybeUnsafe; // get the offset of the data inside a byte array object 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"); 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; final ByteBuffer direct = ByteBuffer.allocateDirect(1); try { final Object maybeDirectBufferConstructor = AccessController.doPrivileged( new PrivilegedAction() { @Override public Object run() { try { final Constructor constructor = (majorVersion >= 21) ? direct.getClass().getDeclaredConstructor(long.class, long.class) : direct.getClass().getDeclaredConstructor(long.class, int.class); constructor.setAccessible(true); logger.debug("Constructor for direct buffer found and made accessible"); return constructor; } catch (NoSuchMethodException e) { logger.debug("Cannot get constructor for direct buffer allocation", e); return e; } 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; } } }); if (maybeDirectBufferConstructor instanceof Constructor) { address = UNSAFE.allocateMemory(1); // try to use the constructor now try { ((Constructor) maybeDirectBufferConstructor).newInstance(address, 1); directBufferConstructor = (Constructor) maybeDirectBufferConstructor; logger.debug("direct buffer constructor: available"); } catch (InstantiationException | IllegalAccessException | InvocationTargetException e) { logger.warn("unable to instantiate a direct buffer via constructor", e); directBufferConstructor = null; } } else { logger.debug( "direct buffer constructor: unavailable", (Throwable) maybeDirectBufferConstructor); directBufferConstructor = null; } } finally { if (address != -1) { UNSAFE.freeMemory(address); } } DIRECT_BUFFER_CONSTRUCTOR = directBufferConstructor; } catch (Throwable e) { // This exception will get swallowed, but it's necessary for the static analysis that ensures // the static fields above get initialized final RuntimeException failure = new RuntimeException( "Failed to initialize MemoryUtil. You must start Java with " + "`--add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED` " + "(See https://arrow.apache.org/docs/java/install.html)", e); failure.printStackTrace(); throw failure; } } /** * Given a {@link ByteBuffer}, gets the address the underlying memory space. * * @param buf the byte buffer. * @return address of the underlying memory. */ public static long getByteBufferAddress(ByteBuffer buf) { 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() {} /** Create nio byte buffer. */ public static ByteBuffer directBuffer(long address, int capacity) { if (DIRECT_BUFFER_CONSTRUCTOR != null) { if (capacity < 0) { throw new IllegalArgumentException("Capacity is negative, has to be positive or 0"); } try { return (ByteBuffer) DIRECT_BUFFER_CONSTRUCTOR.newInstance(address, capacity); } catch (Throwable cause) { throw new Error(cause); } } throw new UnsupportedOperationException( "sun.misc.Unsafe or java.nio.DirectByteBuffer.(long, int) not available"); } @SuppressWarnings( "nullness:argument") // to handle null assignment on third party dependency: Unsafe private static void copyMemory( @Nullable Object srcBase, long srcOffset, @Nullable Object destBase, long destOffset, long bytes) { UNSAFE.copyMemory(srcBase, srcOffset, destBase, destOffset, bytes); } public static void copyMemory(long srcAddress, long destAddress, long bytes) { UNSAFE.copyMemory(srcAddress, destAddress, bytes); } public static void copyToMemory(byte[] src, long srcIndex, long destAddress, long bytes) { copyMemory(src, BYTE_ARRAY_BASE_OFFSET + srcIndex, null, destAddress, bytes); } public static void copyFromMemory(long srcAddress, byte[] dest, long destIndex, long bytes) { copyMemory(null, srcAddress, dest, BYTE_ARRAY_BASE_OFFSET + destIndex, bytes); } public static byte getByte(long address) { return UNSAFE.getByte(address); } public static void putByte(long address, byte value) { UNSAFE.putByte(address, value); } public static short getShort(long address) { return UNSAFE.getShort(address); } public static void putShort(long address, short value) { UNSAFE.putShort(address, value); } public static int getInt(long address) { return UNSAFE.getInt(address); } public static void putInt(long address, int value) { UNSAFE.putInt(address, value); } public static long getLong(long address) { return UNSAFE.getLong(address); } public static void putLong(long address, long value) { UNSAFE.putLong(address, value); } public static void setMemory(long address, long bytes, byte value) { UNSAFE.setMemory(address, bytes, value); } public static int getInt(byte[] bytes, int index) { return UNSAFE.getInt(bytes, BYTE_ARRAY_BASE_OFFSET + index); } public static long getLong(byte[] bytes, int index) { return UNSAFE.getLong(bytes, BYTE_ARRAY_BASE_OFFSET + index); } public static long allocateMemory(long bytes) { return UNSAFE.allocateMemory(bytes); } public static void freeMemory(long address) { UNSAFE.freeMemory(address); } } StackTrace.java000066400000000000000000000045451515451615200335470ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util; import java.util.Arrays; import org.checkerframework.checker.nullness.qual.Nullable; /** Convenient way of obtaining and manipulating stack traces for debugging. */ public class StackTrace { private final @Nullable StackTraceElement[] stackTraceElements; /** Constructor. Captures the current stack trace. */ public StackTrace() { final StackTraceElement[] stack = Thread.currentThread().getStackTrace(); // Skip first two elements to remove getStackTrace/StackTrace. stackTraceElements = Arrays.copyOfRange(stack, 2, stack.length); } /** * Write the stack trace to a StringBuilder. * * @param sb where to write it * @param indent how many double spaces to indent each line */ public void writeToBuilder(final StringBuilder sb, final int indent) { // create the indentation string final char[] indentation = new char[indent * 2]; Arrays.fill(indentation, ' '); // write the stack trace in standard Java format for (StackTraceElement ste : stackTraceElements) { if (ste != null) { sb.append(indentation) .append("at ") .append(ste.getClassName()) .append('.') .append(ste.getMethodName()) .append('(') .append(ste.getFileName()) .append(':') .append(Integer.toString(ste.getLineNumber())) .append(")\n"); } } } @Override public String toString() { final StringBuilder sb = new StringBuilder(); writeToBuilder(sb, 0); return sb.toString(); } } apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/util/hash/000077500000000000000000000000001515451615200316525ustar00rootroot00000000000000ArrowBufHasher.java000066400000000000000000000033061515451615200353220ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/util/hash/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.hash; import org.apache.arrow.memory.ArrowBuf; /** * Utility for calculating the hash code for a consecutive memory region. This class provides the * basic framework for efficiently calculating the hash code. * *

    A default light-weight implementation is given in {@link SimpleHasher}. */ public interface ArrowBufHasher { /** * Calculates the hash code for a memory region. * * @param address start address of the memory region. * @param length length of the memory region. * @return the hash code. */ int hashCode(long address, long length); /** * Calculates the hash code for a memory region. * * @param buf the buffer for the memory region. * @param offset offset within the buffer for the memory region. * @param length length of the memory region. * @return the hash code. */ int hashCode(ArrowBuf buf, long offset, long length); } MurmurHasher.java000066400000000000000000000117441515451615200350670ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/util/hash/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.hash; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.util.MemoryUtil; import org.checkerframework.checker.nullness.qual.Nullable; /** * Implementation of the Murmur hashing algorithm. Details of the algorithm can be found in * https://en.wikipedia.org/wiki/MurmurHash * *

    Murmur hashing is computationally expensive, as it involves several integer multiplications. * However, the produced hash codes have good quality in the sense that they are uniformly * distributed in the universe. * *

    Therefore, this algorithm is suitable for scenarios where uniform hashing is desired (e.g. in * an open addressing hash table/hash set). */ public class MurmurHasher implements ArrowBufHasher { private final int seed; /** Creates a default Murmur hasher, with seed 0. */ public MurmurHasher() { this(0); } /** * Creates a Murmur hasher. * * @param seed the seed for the hasher. */ public MurmurHasher(int seed) { this.seed = seed; } @Override public int hashCode(long address, long length) { return hashCode(address, length, seed); } @Override public int hashCode(ArrowBuf buf, long offset, long length) { buf.checkBytes(offset, offset + length); return hashCode(buf.memoryAddress() + offset, length); } /** * Calculates the hash code for a memory region. * * @param buf the buffer for the memory region. * @param offset offset within the buffer for the memory region. * @param length length of the memory region. * @param seed the seed. * @return the hash code. */ public static int hashCode(ArrowBuf buf, long offset, long length, int seed) { buf.checkBytes(offset, offset + length); return hashCode(buf.memoryAddress() + offset, length, seed); } /** * Calculates the hash code for a memory region. * * @param address start address of the memory region. * @param length length of the memory region. * @param seed the seed. * @return the hash code. */ public static int hashCode(long address, long length, int seed) { int index = 0; int hash = seed; while (index + 4 <= length) { int intValue = MemoryUtil.getInt(address + index); hash = combineHashCode(hash, intValue); index += 4; } if (index < length) { // process remaining data as a integer in little endian int intValue = 0; for (long i = length - 1; i >= index; i--) { intValue <<= 8; intValue |= (MemoryUtil.getByte(address + i) & 0x000000ff); index += 1; } hash = combineHashCode(hash, intValue); } return finalizeHashCode(hash, length); } /** * Combine the current hash code and a new int value to calculate a new hash code. * * @param currentHashCode the current hash code. * @param intValue the new int value. * @return the new hah code. */ public static int combineHashCode(int currentHashCode, int intValue) { int c1 = 0xcc9e2d51; int c2 = 0x1b873593; int r1 = 15; int r2 = 13; int m = 5; int n = 0xe6546b64; int k = intValue; k = k * c1; k = rotateLeft(k, r1); k = k * c2; int hash = currentHashCode; hash = hash ^ k; hash = rotateLeft(hash, r2); hash = hash * m + n; return hash; } /** * Finalizing the hash code. * * @param hashCode the current hash code. * @param length the length of the memory region. * @return the finalized hash code. */ public static int finalizeHashCode(int hashCode, long length) { hashCode = hashCode ^ (int) length; hashCode = hashCode ^ (hashCode >>> 16); hashCode = hashCode * 0x85ebca6b; hashCode = hashCode ^ (hashCode >>> 13); hashCode = hashCode * 0xc2b2ae35; hashCode = hashCode ^ (hashCode >>> 16); return hashCode; } private static int rotateLeft(int value, int count) { return (value << count) | (value >>> (32 - count)); } @Override public boolean equals(@Nullable Object o) { if (this == o) { return true; } if (!(o instanceof MurmurHasher)) { return false; } MurmurHasher that = (MurmurHasher) o; return seed == that.seed; } @Override public int hashCode() { return seed; } } SimpleHasher.java000066400000000000000000000073231515451615200350270ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/memory/util/hash/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.hash; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.util.MemoryUtil; import org.checkerframework.checker.nullness.qual.Nullable; /** * A simple hasher that calculates the hash code of integers as is, and does not perform any * finalization. So the computation is extremely efficient. * *

    This algorithm only provides the most basic semantics for the hash code. That is, if two * objects are equal, they must have equal hash code. However, the quality of the produced hash code * may not be good. In other words, the generated hash codes are far from being uniformly * distributed in the universe. * *

    Therefore, this algorithm is suitable only for scenarios where the most basic semantics of the * hash code is required (e.g. in scenarios that require fast and proactive data pruning) * *

    An object of this class is stateless, so it can be shared between threads. */ public class SimpleHasher implements ArrowBufHasher { public static SimpleHasher INSTANCE = new SimpleHasher(); protected SimpleHasher() {} /** * Calculates the hash code for a memory region. * * @param address start address of the memory region. * @param length length of the memory region. * @return the hash code. */ @Override public int hashCode(long address, long length) { int hashValue = 0; int index = 0; while (index + 8 <= length) { long longValue = MemoryUtil.getLong(address + index); int longHash = getLongHashCode(longValue); hashValue = combineHashCode(hashValue, longHash); index += 8; } if (index + 4 <= length) { int intValue = MemoryUtil.getInt(address + index); int intHash = intValue; hashValue = combineHashCode(hashValue, intHash); index += 4; } while (index < length) { byte byteValue = MemoryUtil.getByte(address + index); int byteHash = byteValue; hashValue = combineHashCode(hashValue, byteHash); index += 1; } return finalizeHashCode(hashValue); } /** * Calculates the hash code for a memory region. * * @param buf the buffer for the memory region. * @param offset offset within the buffer for the memory region. * @param length length of the memory region. * @return the hash code. */ @Override public int hashCode(ArrowBuf buf, long offset, long length) { buf.checkBytes(offset, offset + length); return hashCode(buf.memoryAddress() + offset, length); } protected int combineHashCode(int currentHashCode, int newHashCode) { return currentHashCode * 37 + newHashCode; } protected int getLongHashCode(long longValue) { return Long.hashCode(longValue); } protected int finalizeHashCode(int hashCode) { return hashCode; } @Override public int hashCode() { return 123; } @Override public boolean equals(@Nullable Object obj) { return obj != null && (obj instanceof SimpleHasher); } } apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/util/000077500000000000000000000000001515451615200274175ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java000066400000000000000000000200171515451615200331670ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.ArrayList; import java.util.Arrays; import java.util.Collection; 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 { // Utility class. Should not be instantiated private AutoCloseables() {} /** * Returns a new {@link AutoCloseable} that calls {@link #close(Iterable)} on autoCloseables * when close is called. */ public static AutoCloseable all( final @Nullable Collection autoCloseables) { return new AutoCloseable() { @Override public void close() throws Exception { AutoCloseables.close(autoCloseables); } }; } /** * Closes all autoCloseables if not null and suppresses exceptions by adding them to t. * * @param t the throwable to add suppressed exception to * @param autoCloseables the closeables to close */ public static void close(Throwable t, @Nullable AutoCloseable... autoCloseables) { if (autoCloseables == null) { return; } close(t, Arrays.asList(autoCloseables)); } /** * Closes all autoCloseables if not null and suppresses exceptions by adding them to t. * * @param t the throwable to add suppressed exception to * @param autoCloseables the closeables to close */ public static void close( Throwable t, @Nullable Iterable autoCloseables) { try { close(autoCloseables); } catch (Exception e) { t.addSuppressed(e); } } /** * Closes all autoCloseables if not null and suppresses subsequent exceptions if more than one. * * @param autoCloseables the closeables to close */ public static void close(@Nullable AutoCloseable... autoCloseables) throws Exception { if (autoCloseables == null) { return; } close(Arrays.asList(autoCloseables)); } /** * Closes all autoCloseables if not null and suppresses subsequent exceptions if more than one. * * @param ac the closeables to close */ 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) { return; } else if (ac instanceof AutoCloseable) { ((AutoCloseable) ac).close(); return; } Exception topLevelException = null; for (AutoCloseable closeable : ac) { try { if (closeable != null) { closeable.close(); } } catch (Exception e) { if (topLevelException == null) { topLevelException = e; } else if (e != topLevelException) { topLevelException.addSuppressed(e); } } } if (topLevelException != null) { throw topLevelException; } } /** Calls {@link #close(Iterable)} on the flattened list of closeables. */ @SafeVarargs public static void close(@Nullable Iterable... closeables) throws Exception { if (closeables == null) { return; } close(flatten(closeables)); } @SafeVarargs private static Iterable flatten( Iterable... closeables) { return new Iterable() { // Cast from Iterable to Iterable is safe in this // context // since there's no modification of the original collection @SuppressWarnings("unchecked") @Override public Iterator iterator() { return Arrays.stream(closeables) .flatMap( (Iterable i) -> 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(@Nullable AutoCloseable... ac) { if (ac == null || ac.length == 0) { return Collections.emptyList(); } else { final List nonNullAc = new ArrayList<>(); for (AutoCloseable autoCloseable : ac) { if (autoCloseable != null) { nonNullAc.add(autoCloseable); } } return nonNullAc; } } /** A closeable wrapper that will close the underlying closeables if a commit does not occur. */ public static class RollbackCloseable implements AutoCloseable { private boolean commit = false; private final List closeables; public RollbackCloseable(@Nullable AutoCloseable... closeables) { this.closeables = closeables == null ? new ArrayList<>() : new ArrayList<>(Arrays.asList(closeables)); } public T add(T t) { closeables.add(t); return t; } /** Add all of list to the rollback 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(@Nullable Iterable list) { if (list == null) { return; } for (AutoCloseable ac : list) { closeables.add(ac); } } public void commit() { commit = true; } @Override public void close() throws Exception { if (!commit) { AutoCloseables.close(closeables); } } } /** Creates an {@link RollbackCloseable} from the given closeables. */ public static RollbackCloseable rollbackable(@Nullable AutoCloseable... closeables) { return new RollbackCloseable(closeables); } /** * close() an {@link java.lang.AutoCloseable} without throwing a (checked) {@link * java.lang.Exception}. This wraps the close() call with a try-catch that will rethrow an * Exception wrapped with a {@link java.lang.RuntimeException}, providing a way to call close() * without having to do the try-catch everywhere or propagate the Exception. * * @param autoCloseable the AutoCloseable to close; may be null * @throws RuntimeException if an Exception occurs; the Exception is wrapped by the * RuntimeException */ public static void closeNoChecked(final @Nullable AutoCloseable autoCloseable) { if (autoCloseable != null) { try { autoCloseable.close(); } catch (final Exception e) { throw new RuntimeException("Exception while closing: " + e.getMessage(), e); } } } private static final AutoCloseable noOpAutocloseable = new AutoCloseable() { @Override public void close() {} }; /** * Get an AutoCloseable that does nothing. * * @return A do-nothing autocloseable */ public static AutoCloseable noop() { return noOpAutocloseable; } } apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/util/Collections2.java000066400000000000000000000061431515451615200326260ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.ArrayList; import java.util.Arrays; import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Spliterator; import java.util.Spliterators; import java.util.stream.Collectors; import java.util.stream.StreamSupport; /** * Utility methods for manipulating {@link java.util.Collections} and their * subclasses/implementations. */ public final class Collections2 { private Collections2() {} /** Creates a {@link List} from the elements remaining in iterator. */ public static List toList(Iterator iterator) { List target = new ArrayList<>(); iterator.forEachRemaining(target::add); return target; } /** Converts the iterable into a new {@link List}. */ public static List toList(Iterable iterable) { if (iterable instanceof Collection) { // If iterable is a collection, take advantage of it for a more efficient copy return new ArrayList((Collection) iterable); } return toList(iterable.iterator()); } /** Converts the iterable into a new immutable {@link List}. */ public static List toImmutableList(Iterable iterable) { return Collections.unmodifiableList(toList(iterable)); } /** Copies the elements of map to a new unmodifiable map. */ public static Map immutableMapCopy(Map map) { return Collections.unmodifiableMap(new HashMap<>(map)); } /** Copies the elements of list to a new unmodifiable list. */ public static List immutableListCopy(List list) { return Collections.unmodifiableList(new ArrayList<>(list)); } /** Copies the values to a new unmodifiable list. */ public static List asImmutableList(V... values) { return Collections.unmodifiableList(Arrays.asList(values)); } /** * Creates a human readable string from the remaining elements in iterator. * *

    The output should be similar to {@code Arrays#toString(Object[])} */ public static String toString(Iterator iterator) { return StreamSupport.stream( Spliterators.spliteratorUnknownSize(iterator, Spliterator.ORDERED), false) .map(String::valueOf) .collect(Collectors.joining(", ", "[", "]")); } } apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/util/Preconditions.java000066400000000000000000001322661515451615200331140ustar00rootroot00000000000000/* * Copyright (C) 2007 The Guava 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. */ package org.apache.arrow.util; import org.checkerframework.dataflow.qual.AssertMethod; /** * Static convenience methods that help a method or constructor check whether it was invoked * correctly (whether its preconditions have been met). These methods generally accept a * {@code boolean} expression which is expected to be {@code true} (or in the case of {@code * checkNotNull}, an object reference which is expected to be non-null). When {@code false} (or * {@code null}) is passed instead, the {@code Preconditions} method throws an unchecked exception, * which helps the calling method communicate to its caller that that caller has made * a mistake. Example: * * 

    {@code
 * /**
 *  * Returns the positive square root of the given value.
 *  *
 *  * @throws IllegalArgumentException if the value is negative
 *  *}{@code /
 * public static double sqrt(double value) {
 *   Preconditions.checkArgument(value >= 0.0, "negative value: %s", value);
 *   // calculate the square root
 * }
 *
 * void exampleBadCaller() {
 *   double d = sqrt(-1.0);
 * }
 * }
    * *

    In this example, {@code checkArgument} throws an {@code IllegalArgumentException} to indicate * that {@code exampleBadCaller} made an error in its call to {@code sqrt}. * *

    Warning about performance

    * *

    The goal of this class is to improve readability of code, but in some circumstances this may * come at a significant performance cost. Remember that parameter values for message construction * must all be computed eagerly, and autoboxing and varargs array creation may happen as well, even * when the precondition check then succeeds (as it should almost always do in production). In some * circumstances these wasted CPU cycles and allocations can add up to a real problem. * Performance-sensitive precondition checks can always be converted to the customary form: * * 

    {@code
 * if (value < 0.0) {
 *   throw new IllegalArgumentException("negative value: " + value);
 * }
 * }
    * *

    Other types of preconditions

    * *

    Not every type of precondition failure is supported by these methods. Continue to throw * standard JDK exceptions such as {@link java.util.NoSuchElementException} or {@link * UnsupportedOperationException} in the situations they are intended for. * *

    Non-preconditions

    * *

    It is of course possible to use the methods of this class to check for invalid conditions * which are not the caller's fault. Doing so is not recommended because it is * misleading to future readers of the code and of stack traces. See Conditional failures * explained in the Guava User Guide for more advice. * *

    {@code java.util.Objects.requireNonNull()}

    * *

    Projects which use {@code com.google.common} should generally avoid the use of {@link * java.util.Objects#requireNonNull(Object)}. Instead, use whichever of {@link * #checkNotNull(Object)} or {@link Verify#verifyNotNull(Object)} is appropriate to the situation. * (The same goes for the message-accepting overloads.) * *

    Only {@code %s} is supported

    * *

    In {@code Preconditions} error message template strings, only the {@code "%s"} specifier is * supported, not the full range of {@link java.util.Formatter} specifiers. * *

    More information

    * *

    See the Guava User Guide on using {@code * Preconditions}. * * @author Kevin Bourrillion * @since 2.0 */ public final class Preconditions { private Preconditions() {} /** * Ensures the truth of an expression involving one or more parameters to the calling method. * * @param expression a boolean expression * @throws IllegalArgumentException if {@code expression} is false */ @AssertMethod public static void checkArgument(boolean expression) { if (!expression) { throw new IllegalArgumentException(); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * * @param expression a boolean expression * @param errorMessage the exception message to use if the check fails; will be converted to a * string using {@link String#valueOf(Object)} * @throws IllegalArgumentException if {@code expression} is false */ @AssertMethod public static void checkArgument(boolean expression, Object errorMessage) { if (!expression) { throw new IllegalArgumentException(String.valueOf(errorMessage)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * * @param expression a boolean expression * @param errorMessageTemplate a template for the exception message should the check fail. The * message is formed by replacing each {@code %s} placeholder in the template with an * argument. These are matched by position - the first {@code %s} gets {@code * errorMessageArgs[0]}, etc. Unmatched arguments will be appended to the formatted message in * square braces. Unmatched placeholders will be left as-is. * @param errorMessageArgs the arguments to be substituted into the message template. Arguments * are converted to strings using {@link String#valueOf(Object)}. * @throws IllegalArgumentException if {@code expression} is false * @throws NullPointerException if the check fails and either {@code errorMessageTemplate} or * {@code errorMessageArgs} is null (don't let this happen) */ public static void checkArgument( boolean expression, String errorMessageTemplate, Object... errorMessageArgs) { if (!expression) { throw new IllegalArgumentException(format(errorMessageTemplate, errorMessageArgs)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, char p1) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, int p1) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, long p1) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, Object p1) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, char p1, char p2) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, char p1, int p2) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, char p1, long p2) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, char p1, Object p2) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, int p1, char p2) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, int p1, int p2) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, int p1, long p2) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, int p1, Object p2) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, long p1, char p2) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, long p1, int p2) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, long p1, long p2) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, long p1, Object p2) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, Object p1, char p2) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, Object p1, int p2) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, Object p1, long p2) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument(boolean b, String errorMessageTemplate, Object p1, Object p2) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument( boolean b, String errorMessageTemplate, Object p1, Object p2, Object p3) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2, p3)); } } /** * Ensures the truth of an expression involving one or more parameters to the calling method. * *

    See {@link #checkArgument(boolean, String, Object...)} for details. */ public static void checkArgument( boolean b, String errorMessageTemplate, Object p1, Object p2, Object p3, Object p4) { if (!b) { throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2, p3, p4)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * * @param expression a boolean expression * @throws IllegalStateException if {@code expression} is false */ @AssertMethod public static void checkState(boolean expression) { if (!expression) { throw new IllegalStateException(); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * * @param expression a boolean expression * @param errorMessage the exception message to use if the check fails; will be converted to a * string using {@link String#valueOf(Object)} * @throws IllegalStateException if {@code expression} is false */ @AssertMethod public static void checkState(boolean expression, Object errorMessage) { if (!expression) { throw new IllegalStateException(String.valueOf(errorMessage)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * * @param expression a boolean expression * @param errorMessageTemplate a template for the exception message should the check fail. The * message is formed by replacing each {@code %s} placeholder in the template with an * argument. These are matched by position - the first {@code %s} gets {@code * errorMessageArgs[0]}, etc. Unmatched arguments will be appended to the formatted message in * square braces. Unmatched placeholders will be left as-is. * @param errorMessageArgs the arguments to be substituted into the message template. Arguments * are converted to strings using {@link String#valueOf(Object)}. * @throws IllegalStateException if {@code expression} is false * @throws NullPointerException if the check fails and either {@code errorMessageTemplate} or * {@code errorMessageArgs} is null (don't let this happen) */ public static void checkState( boolean expression, String errorMessageTemplate, Object... errorMessageArgs) { if (!expression) { throw new IllegalStateException(format(errorMessageTemplate, errorMessageArgs)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, char p1) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, int p1) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, long p1) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, Object p1) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, char p1, char p2) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, char p1, int p2) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, char p1, long p2) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, char p1, Object p2) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, int p1, char p2) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, int p1, int p2) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, int p1, long p2) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, int p1, Object p2) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, long p1, char p2) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, long p1, int p2) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, long p1, long p2) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, long p1, Object p2) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, Object p1, char p2) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, Object p1, int p2) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, Object p1, long p2) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState(boolean b, String errorMessageTemplate, Object p1, Object p2) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState( boolean b, String errorMessageTemplate, Object p1, Object p2, Object p3) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2, p3)); } } /** * Ensures the truth of an expression involving the state of the calling instance, but not * involving any parameters to the calling method. * *

    See {@link #checkState(boolean, String, Object...)} for details. */ public static void checkState( boolean b, String errorMessageTemplate, Object p1, Object p2, Object p3, Object p4) { if (!b) { throw new IllegalStateException(format(errorMessageTemplate, p1, p2, p3, p4)); } } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * * @param reference an object reference * @return the non-null reference that was validated * @throws NullPointerException if {@code reference} is null */ public static T checkNotNull(T reference) { if (reference == null) { throw new NullPointerException(); } return reference; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * * @param reference an object reference * @param errorMessage the exception message to use if the check fails; will be converted to a * string using {@link String#valueOf(Object)} * @return the non-null reference that was validated * @throws NullPointerException if {@code reference} is null */ public static T checkNotNull(T reference, Object errorMessage) { if (reference == null) { throw new NullPointerException(String.valueOf(errorMessage)); } return reference; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * * @param reference an object reference * @param errorMessageTemplate a template for the exception message should the check fail. The * message is formed by replacing each {@code %s} placeholder in the template with an * argument. These are matched by position - the first {@code %s} gets {@code * errorMessageArgs[0]}, etc. Unmatched arguments will be appended to the formatted message in * square braces. Unmatched placeholders will be left as-is. * @param errorMessageArgs the arguments to be substituted into the message template. Arguments * are converted to strings using {@link String#valueOf(Object)}. * @return the non-null reference that was validated * @throws NullPointerException if {@code reference} is null */ public static T checkNotNull( T reference, String errorMessageTemplate, Object... errorMessageArgs) { if (reference == null) { // If either of these parameters is null, the right thing happens anyway throw new NullPointerException(format(errorMessageTemplate, errorMessageArgs)); } return reference; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, char p1) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, int p1) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, long p1) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, Object p1) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, char p1, char p2) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, char p1, int p2) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, char p1, long p2) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, char p1, Object p2) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, int p1, char p2) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, int p1, int p2) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, int p1, long p2) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, int p1, Object p2) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, long p1, char p2) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, long p1, int p2) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, long p1, long p2) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, long p1, Object p2) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, Object p1, char p2) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, Object p1, int p2) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, Object p1, long p2) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull(T obj, String errorMessageTemplate, Object p1, Object p2) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull( T obj, String errorMessageTemplate, Object p1, Object p2, Object p3) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2, p3)); } return obj; } /** * Ensures that an object reference passed as a parameter to the calling method is not null. * *

    See {@link #checkNotNull(Object, String, Object...)} for details. */ public static T checkNotNull( T obj, String errorMessageTemplate, Object p1, Object p2, Object p3, Object p4) { if (obj == null) { throw new NullPointerException(format(errorMessageTemplate, p1, p2, p3, p4)); } return obj; } /* * All recent hotspots (as of 2009) *really* like to have the natural code * * if (guardExpression) { * throw new BadException(messageExpression); * } * * refactored so that messageExpression is moved to a separate String-returning method. * * if (guardExpression) { * throw new BadException(badMsg(...)); * } * * The alternative natural refactorings into void or Exception-returning methods are much slower. * This is a big deal - we're talking factors of 2-8 in microbenchmarks, not just 10-20%. (This is * a hotspot optimizer bug, which should be fixed, but that's a separate, big project). * * The coding pattern above is heavily used in java.util, e.g. in ArrayList. There is a * RangeCheckMicroBenchmark in the JDK that was used to test this. * * But the methods in this class want to throw different exceptions, depending on the args, so it * appears that this pattern is not directly applicable. But we can use the ridiculous, devious * trick of throwing an exception in the middle of the construction of another exception. Hotspot * is fine with that. */ /** * Ensures that {@code index} specifies a valid element in an array, list or string of size * {@code size}. An element index may range from zero, inclusive, to {@code size}, exclusive. * * @param index a user-supplied index identifying an element of an array, list or string * @param size the size of that array, list or string * @return the value of {@code index} * @throws IndexOutOfBoundsException if {@code index} is negative or is not less than {@code size} * @throws IllegalArgumentException if {@code size} is negative */ public static int checkElementIndex(int index, int size) { return checkElementIndex(index, size, "index"); } /** * Ensures that {@code index} specifies a valid element in an array, list or string of size * {@code size}. An element index may range from zero, inclusive, to {@code size}, exclusive. * * @param index a user-supplied index identifying an element of an array, list or string * @param size the size of that array, list or string * @param desc the text to use to describe this index in an error message * @return the value of {@code index} * @throws IndexOutOfBoundsException if {@code index} is negative or is not less than {@code size} * @throws IllegalArgumentException if {@code size} is negative */ public static int checkElementIndex(int index, int size, String desc) { // Carefully optimized for execution by hotspot (explanatory comment above) if (index < 0 || index >= size) { throw new IndexOutOfBoundsException(badElementIndex(index, size, desc)); } return index; } private static String badElementIndex(int index, int size, String desc) { if (index < 0) { return format("%s (%s) must not be negative", desc, index); } else if (size < 0) { throw new IllegalArgumentException("negative size: " + size); } else { // index >= size return format("%s (%s) must be less than size (%s)", desc, index, size); } } /** * Ensures that {@code index} specifies a valid position in an array, list or string of * size {@code size}. A position index may range from zero to {@code size}, inclusive. * * @param index a user-supplied index identifying a position in an array, list or string * @param size the size of that array, list or string * @return the value of {@code index} * @throws IndexOutOfBoundsException if {@code index} is negative or is greater than {@code size} * @throws IllegalArgumentException if {@code size} is negative */ public static long checkPositionIndex(long index, long size) { return checkPositionIndex(index, size, "index"); } /** * Ensures that {@code index} specifies a valid position in an array, list or string of * size {@code size}. A position index may range from zero to {@code size}, inclusive. * * @param index a user-supplied index identifying a position in an array, list or string * @param size the size of that array, list or string * @param desc the text to use to describe this index in an error message * @return the value of {@code index} * @throws IndexOutOfBoundsException if {@code index} is negative or is greater than {@code size} * @throws IllegalArgumentException if {@code size} is negative */ public static long checkPositionIndex(long index, long size, String desc) { // Carefully optimized for execution by hotspot (explanatory comment above) if (index < 0 || index > size) { throw new IndexOutOfBoundsException(badPositionIndex(index, size, desc)); } return index; } private static String badPositionIndex(long index, long size, String desc) { if (index < 0) { return format("%s (%s) must not be negative", desc, index); } else if (size < 0) { throw new IllegalArgumentException("negative size: " + size); } else { // index > size return format("%s (%s) must not be greater than size (%s)", desc, index, size); } } /** * Ensures that {@code start} and {@code end} specify a valid positions in an array, list * or string of size {@code size}, and are in order. A position index may range from zero to * {@code size}, inclusive. * * @param start a user-supplied index identifying a starting position in an array, list or string * @param end a user-supplied index identifying a ending position in an array, list or string * @param size the size of that array, list or string * @throws IndexOutOfBoundsException if either index is negative or is greater than {@code size}, * or if {@code end} is less than {@code start} * @throws IllegalArgumentException if {@code size} is negative */ public static void checkPositionIndexes(int start, int end, int size) { // Carefully optimized for execution by hotspot (explanatory comment above) if (start < 0 || end < start || end > size) { throw new IndexOutOfBoundsException(badPositionIndexes(start, end, size)); } } private static String badPositionIndexes(int start, int end, int size) { if (start < 0 || start > size) { return badPositionIndex(start, size, "start index"); } if (end < 0 || end > size) { return badPositionIndex(end, size, "end index"); } // end < start return format("end index (%s) must not be less than start index (%s)", end, start); } /** * Substitutes each {@code %s} in {@code template} with an argument. These are matched by * position: the first {@code %s} gets {@code args[0]}, etc. If there are more arguments than * placeholders, the unmatched arguments will be appended to the end of the formatted message in * square braces. * * @param template a non-null string containing 0 or more {@code %s} placeholders. * @param args the arguments to be substituted into the message template. Arguments are converted * to strings using {@link String#valueOf(Object)}. Arguments can be null. */ // Note that this is somewhat-improperly used from Verify.java as well. static String format(String template, Object... args) { template = String.valueOf(template); // null -> "null" // start substituting the arguments into the '%s' placeholders StringBuilder builder = new StringBuilder(template.length() + 16 * args.length); int templateStart = 0; int i = 0; while (i < args.length) { int placeholderStart = template.indexOf("%s", templateStart); if (placeholderStart == -1) { break; } builder.append(template, templateStart, placeholderStart); builder.append(args[i++]); templateStart = placeholderStart + 2; } builder.append(template, templateStart, template.length()); // if we run out of placeholders, append the extra args in square braces if (i < args.length) { builder.append(" ["); builder.append(args[i++]); while (i < args.length) { builder.append(", "); builder.append(args[i++]); } builder.append(']'); } return builder.toString(); } } VisibleForTesting.java000066400000000000000000000020121515451615200336000ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/main/java/org/apache/arrow/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * Annotation to indicate a class member or class is visible only for the purposes of testing and * otherwise should not be referenced by other classes. */ public @interface VisibleForTesting {} apache-arrow-java-19.0.0/memory/memory-core/src/test/000077500000000000000000000000001515451615200224125ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/000077500000000000000000000000001515451615200233335ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/000077500000000000000000000000001515451615200241225ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/000077500000000000000000000000001515451615200253435ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200264755ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/memory/000077500000000000000000000000001515451615200300055ustar00rootroot00000000000000CountingAllocationListener.java000066400000000000000000000060531515451615200360770ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.checkerframework.checker.nullness.qual.Nullable; /** * Counting allocation listener. It counts the number of times it has been invoked, and how much * memory allocation it has seen When set to 'expand on fail', it attempts to expand the associated * allocator's limit. */ final class CountingAllocationListener implements AllocationListener { private int numPreCalls; private int numCalls; private int numReleaseCalls; private int numChildren; private long totalMem; private long currentMem; private boolean expandOnFail; @Nullable BufferAllocator expandAlloc; long expandLimit; CountingAllocationListener() { this.numCalls = 0; this.numChildren = 0; this.totalMem = 0; this.currentMem = 0; this.expandOnFail = false; this.expandAlloc = null; this.expandLimit = 0; } @Override public void onPreAllocation(long size) { numPreCalls++; } @Override public void onAllocation(long size) { numCalls++; totalMem += size; currentMem += size; } @Override public boolean onFailedAllocation(long size, AllocationOutcome outcome) { if (expandOnFail) { if (expandAlloc == null) { throw new IllegalStateException( "expandAlloc must be non-null because this " + "listener is set to expand on failure."); } expandAlloc.setLimit(expandLimit); return true; } return false; } @Override public void onRelease(long size) { numReleaseCalls++; currentMem -= size; } @Override public void onChildAdded(BufferAllocator parentAllocator, BufferAllocator childAllocator) { ++numChildren; } @Override public void onChildRemoved(BufferAllocator parentAllocator, BufferAllocator childAllocator) { --numChildren; } void setExpandOnFail(BufferAllocator expandAlloc, long expandLimit) { this.expandOnFail = true; this.expandAlloc = expandAlloc; this.expandLimit = expandLimit; } int getNumPreCalls() { return numPreCalls; } int getNumReleaseCalls() { return numReleaseCalls; } int getNumCalls() { return numCalls; } int getNumChildren() { return numChildren; } long getTotalMem() { return totalMem; } long getCurrentMem() { return currentMem; } } DefaultAllocationManagerFactory.java000066400000000000000000000036421515451615200370130ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.memory.util.MemoryUtil; /** * The default Allocation Manager Factory for a module. * *

    This is only used by tests and contains only a simplistic allocator method. */ public class DefaultAllocationManagerFactory implements AllocationManager.Factory { public static final AllocationManager.Factory FACTORY = new DefaultAllocationManagerFactory(); private static final ArrowBuf EMPTY = new ArrowBuf(ReferenceManager.NO_OP, null, 0, MemoryUtil.allocateMemory(0)); @Override public AllocationManager create(BufferAllocator accountingAllocator, long size) { return new AllocationManager(accountingAllocator) { private final long allocatedSize = size; private final long address = MemoryUtil.allocateMemory(size); @Override public long getSize() { return allocatedSize; } @Override protected long memoryAddress() { return address; } @Override protected void release0() { MemoryUtil.freeMemory(address); } }; } @Override public ArrowBuf empty() { return EMPTY; } } TestAccountant.java000066400000000000000000000125401515451615200335320ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.fail; import org.junit.jupiter.api.Test; public class TestAccountant { @Test public void basic() { ensureAccurateReservations(null); } @Test public void nested() { final Accountant parent = new Accountant(null, "test", 0, Long.MAX_VALUE); ensureAccurateReservations(parent); assertEquals(0, parent.getAllocatedMemory()); assertEquals(parent.getLimit() - parent.getAllocatedMemory(), parent.getHeadroom()); } @Test public void multiThread() throws InterruptedException { final Accountant parent = new Accountant(null, "test", 0, Long.MAX_VALUE); final int numberOfThreads = 32; final int loops = 100; Thread[] threads = new Thread[numberOfThreads]; for (int i = 0; i < numberOfThreads; i++) { Thread t = new Thread() { @Override public void run() { try { for (int i = 0; i < loops; i++) { ensureAccurateReservations(parent); } } catch (Exception ex) { ex.printStackTrace(); fail(ex.getMessage()); } } }; threads[i] = t; t.start(); } for (Thread thread : threads) { thread.join(); } assertEquals(0, parent.getAllocatedMemory()); assertEquals(parent.getLimit() - parent.getAllocatedMemory(), parent.getHeadroom()); } private void ensureAccurateReservations(Accountant outsideParent) { final Accountant parent = new Accountant(outsideParent, "test", 0, 10); assertEquals(0, parent.getAllocatedMemory()); final Accountant child = new Accountant(parent, "test", 2, Long.MAX_VALUE); assertEquals(2, parent.getAllocatedMemory()); assertEquals(10, child.getHeadroom()); { AllocationOutcome first = child.allocateBytes(1); assertEquals(AllocationOutcome.Status.SUCCESS, first.getStatus()); } // child will have new allocation assertEquals(1, child.getAllocatedMemory()); // root has no change since within reservation assertEquals(2, parent.getAllocatedMemory()); { AllocationOutcome first = child.allocateBytes(1); assertEquals(AllocationOutcome.Status.SUCCESS, first.getStatus()); } // child will have new allocation assertEquals(2, child.getAllocatedMemory()); // root has no change since within reservation assertEquals(2, parent.getAllocatedMemory()); child.releaseBytes(1); // child will have new allocation assertEquals(1, child.getAllocatedMemory()); // root has no change since within reservation assertEquals(2, parent.getAllocatedMemory()); { AllocationOutcome first = child.allocateBytes(2); assertEquals(AllocationOutcome.Status.SUCCESS, first.getStatus()); } // child will have new allocation assertEquals(3, child.getAllocatedMemory()); // went beyond reservation, now in parent accountant assertEquals(3, parent.getAllocatedMemory()); assertEquals(7, child.getHeadroom()); assertEquals(7, parent.getHeadroom()); { AllocationOutcome first = child.allocateBytes(7); assertEquals(AllocationOutcome.Status.SUCCESS, first.getStatus()); } // child will have new allocation assertEquals(10, child.getAllocatedMemory()); // went beyond reservation, now in parent accountant assertEquals(10, parent.getAllocatedMemory()); child.releaseBytes(9); assertEquals(1, child.getAllocatedMemory()); assertEquals(9, child.getHeadroom()); // back to reservation size assertEquals(2, parent.getAllocatedMemory()); assertEquals(8, parent.getHeadroom()); AllocationOutcome first = child.allocateBytes(10); assertEquals(AllocationOutcome.Status.FAILED_PARENT, first.getStatus()); // unchanged assertEquals(1, child.getAllocatedMemory()); assertEquals(2, parent.getAllocatedMemory()); boolean withinLimit = child.forceAllocate(10); assertFalse(withinLimit); // at new limit assertEquals(child.getAllocatedMemory(), 11); assertEquals(parent.getAllocatedMemory(), 11); assertEquals(-1, child.getHeadroom()); assertEquals(-1, parent.getHeadroom()); child.releaseBytes(11); assertEquals(child.getAllocatedMemory(), 0); assertEquals(parent.getAllocatedMemory(), 2); assertEquals(10, child.getHeadroom()); assertEquals(8, parent.getHeadroom()); child.close(); parent.close(); } } TestAllocationManager.java000066400000000000000000000026521515451615200350160ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 static org.junit.jupiter.api.Assertions.assertEquals; import org.junit.jupiter.api.Test; /** Test cases for {@link AllocationManager}. */ public class TestAllocationManager { @Test public void testAllocationManagerType() { // test unknown allocation manager type System.clearProperty(DefaultAllocationManagerOption.ALLOCATION_MANAGER_TYPE_PROPERTY_NAME); DefaultAllocationManagerOption.AllocationManagerType mgrType = DefaultAllocationManagerOption.getDefaultAllocationManagerType(); assertEquals(DefaultAllocationManagerOption.AllocationManagerType.Unknown, mgrType); } } apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/memory/TestArrowBuf.java000066400000000000000000000153021515451615200332400ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 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.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import ch.qos.logback.classic.Level; import ch.qos.logback.classic.Logger; import java.lang.reflect.Field; import java.lang.reflect.Modifier; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.util.Arrays; import org.apache.arrow.memory.util.Float16; import org.junit.jupiter.api.Test; import org.slf4j.LoggerFactory; public class TestArrowBuf { @Test public void testSliceOutOfBoundsLength_RaisesIndexOutOfBoundsException() { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(2)) { assertEquals(2, buf.capacity()); assertThrows(IndexOutOfBoundsException.class, () -> buf.slice(0, 3)); } } @Test public void testSliceOutOfBoundsIndexPlusLength_RaisesIndexOutOfBoundsException() { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(2)) { assertEquals(2, buf.capacity()); assertThrows(IndexOutOfBoundsException.class, () -> buf.slice(1, 2)); } } @Test public void testSliceOutOfBoundsIndex_RaisesIndexOutOfBoundsException() { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(2)) { assertEquals(2, buf.capacity()); assertThrows(IndexOutOfBoundsException.class, () -> buf.slice(3, 0)); } } @Test public void testSliceWithinBoundsLength_ReturnsSlice() { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(2)) { assertEquals(2, buf.capacity()); assertEquals(1, buf.slice(1, 1).capacity()); assertEquals(2, buf.slice(0, 2).capacity()); } } @Test public void testSetBytesSliced() { int arrLength = 64; byte[] expected = new byte[arrLength]; for (int i = 0; i < expected.length; i++) { expected[i] = (byte) i; } ByteBuffer data = ByteBuffer.wrap(expected); try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(expected.length)) { buf.setBytes(0, data, 0, data.capacity()); byte[] actual = new byte[expected.length]; buf.getBytes(0, actual); assertArrayEquals(expected, actual); } } @Test public void testSetBytesUnsliced() { int arrLength = 64; byte[] arr = new byte[arrLength]; for (int i = 0; i < arr.length; i++) { arr[i] = (byte) i; } ByteBuffer data = ByteBuffer.wrap(arr); int from = 10; int to = arrLength; byte[] expected = Arrays.copyOfRange(arr, from, to); try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(expected.length)) { buf.setBytes(0, data, from, to - from); byte[] actual = new byte[expected.length]; buf.getBytes(0, actual); assertArrayEquals(expected, actual); } } /** ARROW-9221: guard against big-endian byte buffers. */ @Test public void testSetBytesBigEndian() { final byte[] expected = new byte[64]; for (int i = 0; i < expected.length; i++) { expected[i] = (byte) i; } // Only this code path is susceptible: others use unsafe or byte-by-byte copies, while this // override copies longs. final ByteBuffer data = ByteBuffer.wrap(expected).asReadOnlyBuffer(); assertFalse(data.hasArray()); assertFalse(data.isDirect()); assertEquals(ByteOrder.BIG_ENDIAN, data.order()); try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(expected.length)) { buf.setBytes(0, data); byte[] actual = new byte[expected.length]; buf.getBytes(0, actual); assertArrayEquals(expected, actual); } } @Test /* * Test that allocation history is not recorded even though * assertions are enabled in tests (GH-34338). */ public void testEnabledAssertion() { ((Logger) LoggerFactory.getLogger("org.apache.arrow")).setLevel(Level.TRACE); try (BufferAllocator allocator = new RootAllocator(128)) { allocator.buffer(2); Exception e = assertThrows(IllegalStateException.class, allocator::close); assertFalse(e.getMessage().contains("event log for:")); } finally { ((Logger) LoggerFactory.getLogger("org.apache.arrow")).setLevel(null); } } @Test public void testEnabledHistoricalLog() { ((Logger) LoggerFactory.getLogger("org.apache.arrow")).setLevel(Level.TRACE); try { Field fieldDebug = BaseAllocator.class.getField("DEBUG"); fieldDebug.setAccessible(true); Field modifiersDebug = Field.class.getDeclaredField("modifiers"); modifiersDebug.setAccessible(true); modifiersDebug.setInt(fieldDebug, fieldDebug.getModifiers() & ~Modifier.FINAL); fieldDebug.set(null, true); try (BufferAllocator allocator = new RootAllocator(128)) { allocator.buffer(2); Exception e = assertThrows(IllegalStateException.class, allocator::close); assertTrue( e.getMessage().contains("event log for:"), // JDK8, JDK11 "Exception had the following message: " + e.getMessage()); } finally { fieldDebug.set(null, false); } } catch (Exception e) { assertTrue( e.toString().contains("java.lang.NoSuchFieldException: modifiers"), "Exception had the following toString(): " + e); // JDK17+ } finally { ((Logger) LoggerFactory.getLogger("org.apache.arrow")).setLevel(null); } } @Test public void testArrowBufFloat16() { try (BufferAllocator allocator = new RootAllocator(); ArrowBuf buf = allocator.buffer(1024)) { buf.setShort(0, Float16.toFloat16(+32.875f)); assertEquals((short) 0x501c, buf.getShort(0)); } } } TestBaseAllocator.java000066400000000000000000001275061515451615200341570ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 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.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; import java.util.Arrays; import java.util.Collection; import java.util.Collections; import java.util.Iterator; import org.apache.arrow.memory.AllocationOutcomeDetails.Entry; import org.apache.arrow.memory.rounding.RoundingPolicy; import org.apache.arrow.memory.rounding.SegmentRoundingPolicy; import org.apache.arrow.memory.util.AssertionUtil; import org.apache.arrow.memory.util.MemoryUtil; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; public class TestBaseAllocator { private static final int MAX_ALLOCATION = 8 * 1024; /* // ---------------------------------------- DEBUG ----------------------------------- @After public void checkBuffers() { final int bufferCount = UnsafeDirectLittleEndian.getBufferCount(); if (bufferCount != 0) { UnsafeDirectLittleEndian.logBuffers(logger); UnsafeDirectLittleEndian.releaseBuffers(); } assertEquals(0, bufferCount); } // @AfterClass // public static void dumpBuffers() { // UnsafeDirectLittleEndian.logBuffers(logger); // } // ---------------------------------------- DEBUG ------------------------------------ */ @Test public void test_privateMax() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { final ArrowBuf arrowBuf1 = rootAllocator.buffer(MAX_ALLOCATION / 2); assertNotNull(arrowBuf1, "allocation failed"); try (final BufferAllocator childAllocator = rootAllocator.newChildAllocator("noLimits", 0, MAX_ALLOCATION)) { final ArrowBuf arrowBuf2 = childAllocator.buffer(MAX_ALLOCATION / 2); assertNotNull(arrowBuf2, "allocation failed"); arrowBuf2.getReferenceManager().release(); } arrowBuf1.getReferenceManager().release(); } } @Test public void testRootAllocator_closeWithOutstanding() throws Exception { assertThrows( IllegalStateException.class, () -> { try { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { final ArrowBuf arrowBuf = rootAllocator.buffer(512); assertNotNull(arrowBuf, "allocation failed"); } } finally { /* * We expect there to be one unreleased underlying buffer because we're closing * without releasing it. */ /* // ------------------------------- DEBUG --------------------------------- final int bufferCount = UnsafeDirectLittleEndian.getBufferCount(); UnsafeDirectLittleEndian.releaseBuffers(); assertEquals(1, bufferCount); // ------------------------------- DEBUG --------------------------------- */ } }); } @Test @Disabled public void testRootAllocator_getEmpty() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { final ArrowBuf arrowBuf = rootAllocator.buffer(0); assertNotNull(arrowBuf, "allocation failed"); assertEquals(0, arrowBuf.capacity(), "capacity was non-zero"); assertTrue(arrowBuf.memoryAddress() != 0, "address should be valid"); arrowBuf.getReferenceManager().release(); } } @Disabled // TODO(DRILL-2740) @Test public void testAllocator_unreleasedEmpty() throws Exception { assertThrows( IllegalStateException.class, () -> { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { @SuppressWarnings("unused") final ArrowBuf arrowBuf = rootAllocator.buffer(0); } }); } @Test public void testAllocator_transferOwnership() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { final BufferAllocator childAllocator1 = rootAllocator.newChildAllocator("changeOwnership1", 0, MAX_ALLOCATION); final BufferAllocator childAllocator2 = rootAllocator.newChildAllocator("changeOwnership2", 0, MAX_ALLOCATION); final ArrowBuf arrowBuf1 = childAllocator1.buffer(MAX_ALLOCATION / 4); rootAllocator.verify(); final ReferenceManager referenceManager = arrowBuf1.getReferenceManager(); OwnershipTransferResult transferOwnership = referenceManager.transferOwnership(arrowBuf1, childAllocator2); assertEquiv(arrowBuf1, transferOwnership.getTransferredBuffer()); final boolean allocationFit = transferOwnership.getAllocationFit(); rootAllocator.verify(); assertTrue(allocationFit); arrowBuf1.getReferenceManager().release(); childAllocator1.close(); rootAllocator.verify(); transferOwnership.getTransferredBuffer().getReferenceManager().release(); childAllocator2.close(); } } static boolean equalsIgnoreOrder(Collection c1, Collection c2) { return (c1.size() == c2.size() && c1.containsAll(c2)); } @Test public void testAllocator_getParentAndChild() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { assertNull(rootAllocator.getParentAllocator()); try (final BufferAllocator childAllocator1 = rootAllocator.newChildAllocator("child1", 0, MAX_ALLOCATION)) { assertEquals(childAllocator1.getParentAllocator(), rootAllocator); assertTrue( equalsIgnoreOrder(Arrays.asList(childAllocator1), rootAllocator.getChildAllocators())); try (final BufferAllocator childAllocator2 = rootAllocator.newChildAllocator("child2", 0, MAX_ALLOCATION)) { assertEquals(childAllocator2.getParentAllocator(), rootAllocator); assertTrue( equalsIgnoreOrder( Arrays.asList(childAllocator1, childAllocator2), rootAllocator.getChildAllocators())); try (final BufferAllocator grandChildAllocator = childAllocator1.newChildAllocator("grand-child", 0, MAX_ALLOCATION)) { assertEquals(grandChildAllocator.getParentAllocator(), childAllocator1); assertTrue( equalsIgnoreOrder( Arrays.asList(grandChildAllocator), childAllocator1.getChildAllocators())); } } } } } @Test public void testAllocator_childRemovedOnClose() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { try (final BufferAllocator childAllocator1 = rootAllocator.newChildAllocator("child1", 0, MAX_ALLOCATION)) { try (final BufferAllocator childAllocator2 = rootAllocator.newChildAllocator("child2", 0, MAX_ALLOCATION)) { // root has two child allocators assertTrue( equalsIgnoreOrder( Arrays.asList(childAllocator1, childAllocator2), rootAllocator.getChildAllocators())); try (final BufferAllocator grandChildAllocator = childAllocator1.newChildAllocator("grand-child", 0, MAX_ALLOCATION)) { // child1 has one allocator i.e grand-child assertTrue( equalsIgnoreOrder( Arrays.asList(grandChildAllocator), childAllocator1.getChildAllocators())); } // grand-child closed assertTrue( equalsIgnoreOrder(Collections.EMPTY_SET, childAllocator1.getChildAllocators())); } // root has only one child left assertTrue( equalsIgnoreOrder(Arrays.asList(childAllocator1), rootAllocator.getChildAllocators())); } // all child allocators closed. assertTrue(equalsIgnoreOrder(Collections.EMPTY_SET, rootAllocator.getChildAllocators())); } } @Test public void testAllocator_shareOwnership() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { final BufferAllocator childAllocator1 = rootAllocator.newChildAllocator("shareOwnership1", 0, MAX_ALLOCATION); final BufferAllocator childAllocator2 = rootAllocator.newChildAllocator("shareOwnership2", 0, MAX_ALLOCATION); final ArrowBuf arrowBuf1 = childAllocator1.buffer(MAX_ALLOCATION / 4); rootAllocator.verify(); // share ownership of buffer. final ArrowBuf arrowBuf2 = arrowBuf1.getReferenceManager().retain(arrowBuf1, childAllocator2); rootAllocator.verify(); assertNotNull(arrowBuf2); assertNotEquals(arrowBuf2, arrowBuf1); assertEquiv(arrowBuf1, arrowBuf2); // release original buffer (thus transferring ownership to allocator 2. (should leave // allocator 1 in empty state) arrowBuf1.getReferenceManager().release(); rootAllocator.verify(); childAllocator1.close(); rootAllocator.verify(); final BufferAllocator childAllocator3 = rootAllocator.newChildAllocator("shareOwnership3", 0, MAX_ALLOCATION); final ArrowBuf arrowBuf3 = arrowBuf1.getReferenceManager().retain(arrowBuf1, childAllocator3); assertNotNull(arrowBuf3); assertNotEquals(arrowBuf3, arrowBuf1); assertNotEquals(arrowBuf3, arrowBuf2); assertEquiv(arrowBuf1, arrowBuf3); rootAllocator.verify(); arrowBuf2.getReferenceManager().release(); rootAllocator.verify(); childAllocator2.close(); rootAllocator.verify(); arrowBuf3.getReferenceManager().release(); rootAllocator.verify(); childAllocator3.close(); } } @Test public void testRootAllocator_createChildAndUse() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { try (final BufferAllocator childAllocator = rootAllocator.newChildAllocator("createChildAndUse", 0, MAX_ALLOCATION)) { final ArrowBuf arrowBuf = childAllocator.buffer(512); assertNotNull(arrowBuf, "allocation failed"); arrowBuf.getReferenceManager().release(); } } } @Test public void testRootAllocator_createChildDontClose() throws Exception { assertThrows( IllegalStateException.class, () -> { try { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { final BufferAllocator childAllocator = rootAllocator.newChildAllocator("createChildDontClose", 0, MAX_ALLOCATION); final ArrowBuf arrowBuf = childAllocator.buffer(512); assertNotNull(arrowBuf, "allocation failed"); } } finally { /* * We expect one underlying buffer because we closed a child allocator without * releasing the buffer allocated from it. */ /* // ------------------------------- DEBUG --------------------------------- final int bufferCount = UnsafeDirectLittleEndian.getBufferCount(); UnsafeDirectLittleEndian.releaseBuffers(); assertEquals(1, bufferCount); // ------------------------------- DEBUG --------------------------------- */ } }); } @Test public void testSegmentAllocator() { RoundingPolicy policy = new SegmentRoundingPolicy(1024L); try (RootAllocator allocator = new RootAllocator(AllocationListener.NOOP, 1024 * 1024, policy)) { ArrowBuf buf = allocator.buffer(798); assertEquals(1024, buf.capacity()); buf.setInt(333, 959); assertEquals(959, buf.getInt(333)); buf.close(); buf = allocator.buffer(1025); assertEquals(2048, buf.capacity()); buf.setInt(193, 939); assertEquals(939, buf.getInt(193)); buf.close(); } } @Test public void testSegmentAllocator_childAllocator() { RoundingPolicy policy = new SegmentRoundingPolicy(1024L); try (RootAllocator allocator = new RootAllocator(AllocationListener.NOOP, 1024 * 1024, policy); BufferAllocator childAllocator = allocator.newChildAllocator("child", 0, 512 * 1024)) { assertEquals("child", childAllocator.getName()); ArrowBuf buf = childAllocator.buffer(798); assertEquals(1024, buf.capacity()); buf.setInt(333, 959); assertEquals(959, buf.getInt(333)); buf.close(); buf = childAllocator.buffer(1025); assertEquals(2048, buf.capacity()); buf.setInt(193, 939); assertEquals(939, buf.getInt(193)); buf.close(); } } @Test public void testSegmentAllocator_smallSegment() { IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () -> new SegmentRoundingPolicy(128L)); assertEquals("The segment size cannot be smaller than 1024", e.getMessage()); } @Test public void testSegmentAllocator_segmentSizeNotPowerOf2() { IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () -> new SegmentRoundingPolicy(4097L)); assertEquals("The segment size must be a power of 2", e.getMessage()); } @Test public void testCustomizedAllocationManager() { try (BaseAllocator allocator = createAllocatorWithCustomizedAllocationManager()) { final ArrowBuf arrowBuf1 = allocator.buffer(MAX_ALLOCATION); assertNotNull(arrowBuf1, "allocation failed"); arrowBuf1.setInt(0, 1); assertEquals(1, arrowBuf1.getInt(0)); try { allocator.buffer(1); fail("allocated memory beyond max allowed"); } catch (OutOfMemoryException e) { // expected } arrowBuf1.getReferenceManager().release(); try { arrowBuf1.getInt(0); fail("data read from released buffer"); } catch (RuntimeException e) { // expected } } } private BaseAllocator createAllocatorWithCustomizedAllocationManager() { return new RootAllocator( BaseAllocator.configBuilder() .maxAllocation(MAX_ALLOCATION) .allocationManagerFactory( new AllocationManager.Factory() { @Override public AllocationManager create( BufferAllocator accountingAllocator, long requestedSize) { return new AllocationManager(accountingAllocator) { private final long address = MemoryUtil.allocateMemory(requestedSize); @Override protected long memoryAddress() { return address; } @Override protected void release0() { MemoryUtil.setMemory(address, requestedSize, (byte) 0); MemoryUtil.freeMemory(address); } @Override public long getSize() { return requestedSize; } }; } @Override public ArrowBuf empty() { return null; } }) .build()); } @Test public void testRootAllocator_listeners() throws Exception { CountingAllocationListener listener1 = new CountingAllocationListener(); assertEquals(0, listener1.getNumPreCalls()); assertEquals(0, listener1.getNumCalls()); assertEquals(0, listener1.getNumReleaseCalls()); assertEquals(0, listener1.getNumChildren()); assertEquals(0, listener1.getTotalMem()); CountingAllocationListener listener2 = new CountingAllocationListener(); assertEquals(0, listener2.getNumPreCalls()); assertEquals(0, listener2.getNumCalls()); assertEquals(0, listener2.getNumReleaseCalls()); assertEquals(0, listener2.getNumChildren()); assertEquals(0, listener2.getTotalMem()); // root and first-level child share the first listener // second-level and third-level child share the second listener try (final RootAllocator rootAllocator = new RootAllocator(listener1, MAX_ALLOCATION)) { try (final BufferAllocator c1 = rootAllocator.newChildAllocator("c1", 0, MAX_ALLOCATION)) { assertEquals(1, listener1.getNumChildren()); final ArrowBuf buf1 = c1.buffer(16); assertNotNull(buf1, "allocation failed"); assertEquals(1, listener1.getNumPreCalls()); assertEquals(1, listener1.getNumCalls()); assertEquals(0, listener1.getNumReleaseCalls()); assertEquals(16, listener1.getTotalMem()); buf1.getReferenceManager().release(); try (final BufferAllocator c2 = c1.newChildAllocator("c2", listener2, 0, MAX_ALLOCATION)) { assertEquals( 2, listener1.getNumChildren()); // c1 got a new child, so listener1 is notified. assertEquals(0, listener2.getNumChildren()); final ArrowBuf buf2 = c2.buffer(32); assertNotNull(buf2, "allocation failed"); assertEquals(1, listener1.getNumCalls()); assertEquals(16, listener1.getTotalMem()); assertEquals(1, listener2.getNumPreCalls()); assertEquals(1, listener2.getNumCalls()); assertEquals(0, listener2.getNumReleaseCalls()); assertEquals(32, listener2.getTotalMem()); buf2.getReferenceManager().release(); try (final BufferAllocator c3 = c2.newChildAllocator("c3", 0, MAX_ALLOCATION)) { assertEquals(2, listener1.getNumChildren()); assertEquals(1, listener2.getNumChildren()); final ArrowBuf buf3 = c3.buffer(64); assertNotNull(buf3, "allocation failed"); assertEquals(1, listener1.getNumPreCalls()); assertEquals(1, listener1.getNumCalls()); assertEquals(1, listener1.getNumReleaseCalls()); assertEquals(16, listener1.getTotalMem()); assertEquals(2, listener2.getNumPreCalls()); assertEquals(2, listener2.getNumCalls()); assertEquals(1, listener2.getNumReleaseCalls()); assertEquals(32 + 64, listener2.getTotalMem()); buf3.getReferenceManager().release(); } assertEquals(2, listener1.getNumChildren()); assertEquals(0, listener2.getNumChildren()); // third-level child removed } assertEquals(1, listener1.getNumChildren()); // second-level child removed assertEquals(0, listener2.getNumChildren()); } assertEquals(0, listener1.getNumChildren()); // first-level child removed assertEquals(2, listener2.getNumReleaseCalls()); } } @Test public void testRootAllocator_listenerAllocationFail() throws Exception { CountingAllocationListener listener1 = new CountingAllocationListener(); assertEquals(0, listener1.getNumCalls()); assertEquals(0, listener1.getTotalMem()); // Test attempts to allocate too much from a child whose limit is set to half of the max // allocation. The listener's callback triggers, expanding the child allocator's limit, so then // the allocation succeeds. try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { try (final BufferAllocator c1 = rootAllocator.newChildAllocator("c1", listener1, 0, MAX_ALLOCATION / 2)) { try { c1.buffer(MAX_ALLOCATION); fail("allocated memory beyond max allowed"); } catch (OutOfMemoryException e) { // expected } assertEquals(0, listener1.getNumCalls()); assertEquals(0, listener1.getTotalMem()); listener1.setExpandOnFail(c1, MAX_ALLOCATION); ArrowBuf arrowBuf = c1.buffer(MAX_ALLOCATION); assertNotNull(arrowBuf, "allocation failed"); assertEquals(1, listener1.getNumCalls()); assertEquals(MAX_ALLOCATION, listener1.getTotalMem()); arrowBuf.getReferenceManager().release(); } } } private static void allocateAndFree(final BufferAllocator allocator) { final ArrowBuf arrowBuf = allocator.buffer(512); assertNotNull(arrowBuf, "allocation failed"); arrowBuf.getReferenceManager().release(); final ArrowBuf arrowBuf2 = allocator.buffer(MAX_ALLOCATION); assertNotNull(arrowBuf2, "allocation failed"); arrowBuf2.getReferenceManager().release(); final int nBufs = 8; final ArrowBuf[] arrowBufs = new ArrowBuf[nBufs]; for (int i = 0; i < arrowBufs.length; ++i) { ArrowBuf arrowBufi = allocator.buffer(MAX_ALLOCATION / nBufs); assertNotNull(arrowBufi, "allocation failed"); arrowBufs[i] = arrowBufi; } for (ArrowBuf arrowBufi : arrowBufs) { arrowBufi.getReferenceManager().release(); } } @Test public void testAllocator_manyAllocations() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { try (final BufferAllocator childAllocator = rootAllocator.newChildAllocator("manyAllocations", 0, MAX_ALLOCATION)) { allocateAndFree(childAllocator); } } } @Test public void testAllocator_overAllocate() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { try (final BufferAllocator childAllocator = rootAllocator.newChildAllocator("overAllocate", 0, MAX_ALLOCATION)) { allocateAndFree(childAllocator); try { childAllocator.buffer(MAX_ALLOCATION + 1); fail("allocated memory beyond max allowed"); } catch (OutOfMemoryException e) { // expected } } } } @Test public void testAllocator_overAllocateParent() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { try (final BufferAllocator childAllocator = rootAllocator.newChildAllocator("overAllocateParent", 0, MAX_ALLOCATION)) { final ArrowBuf arrowBuf1 = rootAllocator.buffer(MAX_ALLOCATION / 2); assertNotNull(arrowBuf1, "allocation failed"); final ArrowBuf arrowBuf2 = childAllocator.buffer(MAX_ALLOCATION / 2); assertNotNull(arrowBuf2, "allocation failed"); try { childAllocator.buffer(MAX_ALLOCATION / 4); fail("allocated memory beyond max allowed"); } catch (OutOfMemoryException e) { // expected } arrowBuf1.getReferenceManager().release(); arrowBuf2.getReferenceManager().release(); } } } @Test public void testAllocator_failureAtParentLimitOutcomeDetails() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { try (final BufferAllocator childAllocator = rootAllocator.newChildAllocator("child", 0, MAX_ALLOCATION / 2)) { try (final BufferAllocator grandChildAllocator = childAllocator.newChildAllocator("grandchild", MAX_ALLOCATION / 4, MAX_ALLOCATION)) { OutOfMemoryException e = assertThrows( OutOfMemoryException.class, () -> grandChildAllocator.buffer(MAX_ALLOCATION)); // expected assertTrue(e.getMessage().contains("Unable to allocate buffer")); assertTrue(e.getOutcomeDetails().isPresent(), "missing outcome details"); AllocationOutcomeDetails outcomeDetails = e.getOutcomeDetails().get(); assertEquals(outcomeDetails.getFailedAllocator(), childAllocator); // The order of allocators should be child to root (request propagates to parent if // child cannot satisfy the request). Iterator iterator = outcomeDetails.allocEntries.iterator(); AllocationOutcomeDetails.Entry first = iterator.next(); assertEquals(MAX_ALLOCATION / 4, first.getAllocatedSize()); assertEquals(MAX_ALLOCATION, first.getRequestedSize()); assertFalse(first.isAllocationFailed()); AllocationOutcomeDetails.Entry second = iterator.next(); assertEquals(MAX_ALLOCATION - MAX_ALLOCATION / 4, second.getRequestedSize()); assertEquals(0, second.getAllocatedSize()); assertTrue(second.isAllocationFailed()); assertFalse(iterator.hasNext()); } } } } @Test public void testAllocator_failureAtRootLimitOutcomeDetails() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { try (final BufferAllocator childAllocator = rootAllocator.newChildAllocator("child", MAX_ALLOCATION / 2, Long.MAX_VALUE)) { try (final BufferAllocator grandChildAllocator = childAllocator.newChildAllocator("grandchild", MAX_ALLOCATION / 4, Long.MAX_VALUE)) { OutOfMemoryException e = assertThrows( OutOfMemoryException.class, () -> grandChildAllocator.buffer(MAX_ALLOCATION * 2)); assertTrue(e.getMessage().contains("Unable to allocate buffer")); assertTrue(e.getOutcomeDetails().isPresent(), "missing outcome details"); AllocationOutcomeDetails outcomeDetails = e.getOutcomeDetails().get(); assertEquals(outcomeDetails.getFailedAllocator(), rootAllocator); // The order of allocators should be child to root (request propagates to parent if // child cannot satisfy the request). Iterator iterator = outcomeDetails.allocEntries.iterator(); AllocationOutcomeDetails.Entry first = iterator.next(); assertEquals(MAX_ALLOCATION / 4, first.getAllocatedSize()); assertEquals(2 * MAX_ALLOCATION, first.getRequestedSize()); assertFalse(first.isAllocationFailed()); AllocationOutcomeDetails.Entry second = iterator.next(); assertEquals(MAX_ALLOCATION / 4, second.getAllocatedSize()); assertEquals(2 * MAX_ALLOCATION - MAX_ALLOCATION / 4, second.getRequestedSize()); assertFalse(second.isAllocationFailed()); AllocationOutcomeDetails.Entry third = iterator.next(); assertEquals(0, third.getAllocatedSize()); assertTrue(third.isAllocationFailed()); assertFalse(iterator.hasNext()); } } } } private static void testAllocator_sliceUpBufferAndRelease( final RootAllocator rootAllocator, final BufferAllocator bufferAllocator) { final ArrowBuf arrowBuf1 = bufferAllocator.buffer(MAX_ALLOCATION / 2); rootAllocator.verify(); final ArrowBuf arrowBuf2 = arrowBuf1.slice(16, arrowBuf1.capacity() - 32); rootAllocator.verify(); final ArrowBuf arrowBuf3 = arrowBuf2.slice(16, arrowBuf2.capacity() - 32); rootAllocator.verify(); @SuppressWarnings("unused") final ArrowBuf arrowBuf4 = arrowBuf3.slice(16, arrowBuf3.capacity() - 32); rootAllocator.verify(); arrowBuf3 .getReferenceManager() .release(); // since they share refcounts, one is enough to release them all rootAllocator.verify(); } @Test public void testAllocator_createSlices() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { testAllocator_sliceUpBufferAndRelease(rootAllocator, rootAllocator); try (final BufferAllocator childAllocator = rootAllocator.newChildAllocator("createSlices", 0, MAX_ALLOCATION)) { testAllocator_sliceUpBufferAndRelease(rootAllocator, childAllocator); } rootAllocator.verify(); testAllocator_sliceUpBufferAndRelease(rootAllocator, rootAllocator); try (final BufferAllocator childAllocator = rootAllocator.newChildAllocator("createSlices", 0, MAX_ALLOCATION)) { try (final BufferAllocator childAllocator2 = childAllocator.newChildAllocator("createSlices", 0, MAX_ALLOCATION)) { final ArrowBuf arrowBuf1 = childAllocator2.buffer(MAX_ALLOCATION / 8); @SuppressWarnings("unused") final ArrowBuf arrowBuf2 = arrowBuf1.slice(MAX_ALLOCATION / 16, MAX_ALLOCATION / 16); testAllocator_sliceUpBufferAndRelease(rootAllocator, childAllocator); arrowBuf1.getReferenceManager().release(); rootAllocator.verify(); } rootAllocator.verify(); testAllocator_sliceUpBufferAndRelease(rootAllocator, childAllocator); } rootAllocator.verify(); } } @Test public void testAllocator_sliceRanges() throws Exception { // final AllocatorOwner allocatorOwner = new NamedOwner("sliceRanges"); try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { // Populate a buffer with byte values corresponding to their indices. final ArrowBuf arrowBuf = rootAllocator.buffer(256); assertEquals(256, arrowBuf.capacity()); assertEquals(0, arrowBuf.readerIndex()); assertEquals(0, arrowBuf.readableBytes()); assertEquals(0, arrowBuf.writerIndex()); assertEquals(256, arrowBuf.writableBytes()); final ArrowBuf slice3 = arrowBuf.slice(); assertEquals(0, slice3.readerIndex()); assertEquals(0, slice3.readableBytes()); assertEquals(0, slice3.writerIndex()); // assertEquals(256, slice3.capacity()); // assertEquals(256, slice3.writableBytes()); for (int i = 0; i < 256; ++i) { arrowBuf.writeByte(i); } assertEquals(0, arrowBuf.readerIndex()); assertEquals(256, arrowBuf.readableBytes()); assertEquals(256, arrowBuf.writerIndex()); assertEquals(0, arrowBuf.writableBytes()); final ArrowBuf slice1 = arrowBuf.slice(); assertEquals(0, slice1.readerIndex()); assertEquals(256, slice1.readableBytes()); for (int i = 0; i < 10; ++i) { assertEquals(i, slice1.readByte()); } assertEquals(256 - 10, slice1.readableBytes()); for (int i = 0; i < 256; ++i) { assertEquals((byte) i, slice1.getByte(i)); } final ArrowBuf slice2 = arrowBuf.slice(25, 25); assertEquals(0, slice2.readerIndex()); assertEquals(25, slice2.readableBytes()); for (int i = 25; i < 50; ++i) { assertEquals(i, slice2.readByte()); } /* for(int i = 256; i > 0; --i) { slice3.writeByte(i - 1); } for(int i = 0; i < 256; ++i) { assertEquals(255 - i, slice1.getByte(i)); } */ arrowBuf.getReferenceManager().release(); // all the derived buffers share this fate } } @Test public void testAllocator_slicesOfSlices() throws Exception { // final AllocatorOwner allocatorOwner = new NamedOwner("slicesOfSlices"); try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { // Populate a buffer with byte values corresponding to their indices. final ArrowBuf arrowBuf = rootAllocator.buffer(256); for (int i = 0; i < 256; ++i) { arrowBuf.writeByte(i); } // Slice it up. final ArrowBuf slice0 = arrowBuf.slice(0, arrowBuf.capacity()); for (int i = 0; i < 256; ++i) { assertEquals((byte) i, arrowBuf.getByte(i)); } final ArrowBuf slice10 = slice0.slice(10, arrowBuf.capacity() - 10); for (int i = 10; i < 256; ++i) { assertEquals((byte) i, slice10.getByte(i - 10)); } final ArrowBuf slice20 = slice10.slice(10, arrowBuf.capacity() - 20); for (int i = 20; i < 256; ++i) { assertEquals((byte) i, slice20.getByte(i - 20)); } final ArrowBuf slice30 = slice20.slice(10, arrowBuf.capacity() - 30); for (int i = 30; i < 256; ++i) { assertEquals((byte) i, slice30.getByte(i - 30)); } arrowBuf.getReferenceManager().release(); } } @Test public void testAllocator_transferSliced() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { final BufferAllocator childAllocator1 = rootAllocator.newChildAllocator("transferSliced1", 0, MAX_ALLOCATION); final BufferAllocator childAllocator2 = rootAllocator.newChildAllocator("transferSliced2", 0, MAX_ALLOCATION); final ArrowBuf arrowBuf1 = childAllocator1.buffer(MAX_ALLOCATION / 8); final ArrowBuf arrowBuf2 = childAllocator2.buffer(MAX_ALLOCATION / 8); final ArrowBuf arrowBuf1s = arrowBuf1.slice(0, arrowBuf1.capacity() / 2); final ArrowBuf arrowBuf2s = arrowBuf2.slice(0, arrowBuf2.capacity() / 2); rootAllocator.verify(); OwnershipTransferResult result1 = arrowBuf2s.getReferenceManager().transferOwnership(arrowBuf2s, childAllocator1); assertEquiv(arrowBuf2s, result1.getTransferredBuffer()); rootAllocator.verify(); OwnershipTransferResult result2 = arrowBuf1s.getReferenceManager().transferOwnership(arrowBuf1s, childAllocator2); assertEquiv(arrowBuf1s, result2.getTransferredBuffer()); rootAllocator.verify(); result1.getTransferredBuffer().getReferenceManager().release(); result2.getTransferredBuffer().getReferenceManager().release(); arrowBuf1s.getReferenceManager().release(); // releases arrowBuf1 arrowBuf2s.getReferenceManager().release(); // releases arrowBuf2 childAllocator1.close(); childAllocator2.close(); } } @Test public void testAllocator_shareSliced() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { final BufferAllocator childAllocator1 = rootAllocator.newChildAllocator("transferSliced", 0, MAX_ALLOCATION); final BufferAllocator childAllocator2 = rootAllocator.newChildAllocator("transferSliced", 0, MAX_ALLOCATION); final ArrowBuf arrowBuf1 = childAllocator1.buffer(MAX_ALLOCATION / 8); final ArrowBuf arrowBuf2 = childAllocator2.buffer(MAX_ALLOCATION / 8); final ArrowBuf arrowBuf1s = arrowBuf1.slice(0, arrowBuf1.capacity() / 2); final ArrowBuf arrowBuf2s = arrowBuf2.slice(0, arrowBuf2.capacity() / 2); rootAllocator.verify(); final ArrowBuf arrowBuf2s1 = arrowBuf2s.getReferenceManager().retain(arrowBuf2s, childAllocator1); assertEquiv(arrowBuf2s, arrowBuf2s1); final ArrowBuf arrowBuf1s2 = arrowBuf1s.getReferenceManager().retain(arrowBuf1s, childAllocator2); assertEquiv(arrowBuf1s, arrowBuf1s2); rootAllocator.verify(); arrowBuf1s.getReferenceManager().release(); // releases arrowBuf1 arrowBuf2s.getReferenceManager().release(); // releases arrowBuf2 rootAllocator.verify(); arrowBuf2s1.getReferenceManager().release(); // releases the shared arrowBuf2 slice arrowBuf1s2.getReferenceManager().release(); // releases the shared arrowBuf1 slice childAllocator1.close(); childAllocator2.close(); } } @Test public void testAllocator_transferShared() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { final BufferAllocator childAllocator1 = rootAllocator.newChildAllocator("transferShared1", 0, MAX_ALLOCATION); final BufferAllocator childAllocator2 = rootAllocator.newChildAllocator("transferShared2", 0, MAX_ALLOCATION); final BufferAllocator childAllocator3 = rootAllocator.newChildAllocator("transferShared3", 0, MAX_ALLOCATION); final ArrowBuf arrowBuf1 = childAllocator1.buffer(MAX_ALLOCATION / 8); boolean allocationFit; ArrowBuf arrowBuf2 = arrowBuf1.getReferenceManager().retain(arrowBuf1, childAllocator2); rootAllocator.verify(); assertNotNull(arrowBuf2); assertNotEquals(arrowBuf2, arrowBuf1); assertEquiv(arrowBuf1, arrowBuf2); final ReferenceManager refManager1 = arrowBuf1.getReferenceManager(); final OwnershipTransferResult result1 = refManager1.transferOwnership(arrowBuf1, childAllocator3); allocationFit = result1.getAllocationFit(); final ArrowBuf arrowBuf3 = result1.getTransferredBuffer(); assertTrue(allocationFit); assertEquiv(arrowBuf1, arrowBuf3); rootAllocator.verify(); // Since childAllocator3 now has childAllocator1's buffer, 1, can close arrowBuf1.getReferenceManager().release(); childAllocator1.close(); rootAllocator.verify(); arrowBuf2.getReferenceManager().release(); childAllocator2.close(); rootAllocator.verify(); final BufferAllocator childAllocator4 = rootAllocator.newChildAllocator("transferShared4", 0, MAX_ALLOCATION); final ReferenceManager refManager3 = arrowBuf3.getReferenceManager(); final OwnershipTransferResult result3 = refManager3.transferOwnership(arrowBuf3, childAllocator4); allocationFit = result3.getAllocationFit(); final ArrowBuf arrowBuf4 = result3.getTransferredBuffer(); assertTrue(allocationFit); assertEquiv(arrowBuf3, arrowBuf4); rootAllocator.verify(); arrowBuf3.getReferenceManager().release(); childAllocator3.close(); rootAllocator.verify(); arrowBuf4.getReferenceManager().release(); childAllocator4.close(); rootAllocator.verify(); } } @Test public void testAllocator_unclaimedReservation() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { try (final BufferAllocator childAllocator1 = rootAllocator.newChildAllocator("unclaimedReservation", 0, MAX_ALLOCATION)) { try (final AllocationReservation reservation = childAllocator1.newReservation()) { assertTrue(reservation.add(64L)); } rootAllocator.verify(); } } } @Test public void testAllocator_claimedReservation() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { try (final BufferAllocator childAllocator1 = rootAllocator.newChildAllocator("claimedReservation", 0, MAX_ALLOCATION)) { try (final AllocationReservation reservation = childAllocator1.newReservation()) { assertTrue(reservation.add(32L)); assertTrue(reservation.add(32L)); final ArrowBuf arrowBuf = reservation.allocateBuffer(); assertEquals(64, arrowBuf.capacity()); rootAllocator.verify(); arrowBuf.getReferenceManager().release(); rootAllocator.verify(); } rootAllocator.verify(); } } } @Test public void testInitReservationAndLimit() throws Exception { try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { try (final BufferAllocator childAllocator = rootAllocator.newChildAllocator("child", 2048, 4096)) { assertEquals(2048, childAllocator.getInitReservation()); assertEquals(4096, childAllocator.getLimit()); } } } @Test public void multiple() throws Exception { final String owner = "test"; try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { final int op = 100000; BufferAllocator frag1 = allocator.newChildAllocator(owner, 1500000, Long.MAX_VALUE); BufferAllocator frag2 = allocator.newChildAllocator(owner, 500000, Long.MAX_VALUE); allocator.verify(); BufferAllocator allocator11 = frag1.newChildAllocator(owner, op, Long.MAX_VALUE); ArrowBuf b11 = allocator11.buffer(1000000); allocator.verify(); BufferAllocator allocator12 = frag1.newChildAllocator(owner, op, Long.MAX_VALUE); ArrowBuf b12 = allocator12.buffer(500000); allocator.verify(); BufferAllocator allocator21 = frag1.newChildAllocator(owner, op, Long.MAX_VALUE); allocator.verify(); BufferAllocator allocator22 = frag2.newChildAllocator(owner, op, Long.MAX_VALUE); ArrowBuf b22 = allocator22.buffer(2000000); allocator.verify(); BufferAllocator frag3 = allocator.newChildAllocator(owner, 1000000, Long.MAX_VALUE); allocator.verify(); BufferAllocator allocator31 = frag3.newChildAllocator(owner, op, Long.MAX_VALUE); ArrowBuf b31a = allocator31.buffer(200000); allocator.verify(); // Previously running operator completes b22.getReferenceManager().release(); allocator.verify(); allocator22.close(); b31a.getReferenceManager().release(); allocator31.close(); b12.getReferenceManager().release(); allocator12.close(); allocator21.close(); b11.getReferenceManager().release(); allocator11.close(); frag1.close(); frag2.close(); frag3.close(); } } // This test needs to run in non-debug mode. So disabling the assertion status through class // loader for this. // The test passes if run individually with -Dtest=TestBaseAllocator#testMemoryLeakWithReservation // but fails generally since the assertion status cannot be changed once the class is initialized. // So setting the test to @ignore @Test @Disabled public void testMemoryLeakWithReservation() throws Exception { // disabling assertion status AssertionUtil.class .getClassLoader() .setClassAssertionStatus(AssertionUtil.class.getName(), false); try (RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { ChildAllocator childAllocator1 = (ChildAllocator) rootAllocator.newChildAllocator("child1", 1024, MAX_ALLOCATION); rootAllocator.verify(); ChildAllocator childAllocator2 = (ChildAllocator) childAllocator1.newChildAllocator("child2", 1024, MAX_ALLOCATION); rootAllocator.verify(); childAllocator2.buffer(256); Exception exception = assertThrows(IllegalStateException.class, childAllocator2::close); String exMessage = exception.getMessage(); assertTrue(exMessage.contains("Memory leaked: (256)")); exception = assertThrows(IllegalStateException.class, childAllocator1::close); exMessage = exception.getMessage(); assertTrue(exMessage.contains("Memory leaked: (256)")); } } @Test public void testOverlimit() { try (BufferAllocator allocator = new RootAllocator(1024)) { try (BufferAllocator child1 = allocator.newChildAllocator("ChildA", 0, 1024); BufferAllocator child2 = allocator.newChildAllocator("ChildB", 1024, 1024)) { assertThrows( OutOfMemoryException.class, () -> { ArrowBuf buf1 = child1.buffer(8); buf1.close(); }); assertEquals(0, child1.getAllocatedMemory()); assertEquals(0, child2.getAllocatedMemory()); assertEquals(1024, allocator.getAllocatedMemory()); } } } @Test public void testOverlimitOverflow() { // Regression test for https://github.com/apache/arrow/issues/35960 try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { try (BufferAllocator child1 = allocator.newChildAllocator("ChildA", 0, Long.MAX_VALUE); BufferAllocator child2 = allocator.newChildAllocator("ChildB", Long.MAX_VALUE, Long.MAX_VALUE)) { assertThrows( OutOfMemoryException.class, () -> { ArrowBuf buf1 = child1.buffer(1024); buf1.close(); }); assertEquals(0, child1.getAllocatedMemory()); assertEquals(0, child2.getAllocatedMemory()); assertEquals(Long.MAX_VALUE, allocator.getAllocatedMemory()); } } } public void assertEquiv(ArrowBuf origBuf, ArrowBuf newBuf) { assertEquals(origBuf.readerIndex(), newBuf.readerIndex()); assertEquals(origBuf.writerIndex(), newBuf.writerIndex()); } } TestBoundaryChecking.java000066400000000000000000000124601515451615200346530ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import java.lang.reflect.Field; import java.net.URLClassLoader; import org.junit.jupiter.api.Test; /** Test cases for evaluating the value of {@link BoundsChecking#BOUNDS_CHECKING_ENABLED}. */ public class TestBoundaryChecking { /** * Get a copy of the current class loader. * * @return the newly created class loader. */ private ClassLoader copyClassLoader() { ClassLoader curClassLoader = this.getClass().getClassLoader(); if (curClassLoader instanceof URLClassLoader) { // for Java 1.8 return new URLClassLoader(((URLClassLoader) curClassLoader).getURLs(), null); } // for Java 1.9 and Java 11. return null; } /** * Get the value of flag {@link BoundsChecking#BOUNDS_CHECKING_ENABLED}. * * @param classLoader the class loader from which to get the flag value. * @return value of the flag. */ private boolean getFlagValue(ClassLoader classLoader) throws Exception { Class clazz = classLoader.loadClass("org.apache.arrow.memory.BoundsChecking"); Field field = clazz.getField("BOUNDS_CHECKING_ENABLED"); return (Boolean) field.get(null); } /** * Ensure the flag for bounds checking is enabled by default. This will protect users from JVM * crashes. */ @Test public void testDefaultValue() throws Exception { ClassLoader classLoader = copyClassLoader(); if (classLoader != null) { boolean boundsCheckingEnabled = getFlagValue(classLoader); assertTrue(boundsCheckingEnabled); } } /** * Test setting the bounds checking flag by the old property. * * @throws Exception if loading class {@link BoundsChecking#BOUNDS_CHECKING_ENABLED} fails. */ @Test public void testEnableOldProperty() throws Exception { String savedOldProperty = System.getProperty("drill.enable_unsafe_memory_access"); System.setProperty("drill.enable_unsafe_memory_access", "true"); ClassLoader classLoader = copyClassLoader(); if (classLoader != null) { boolean boundsCheckingEnabled = getFlagValue(classLoader); assertFalse(boundsCheckingEnabled); } // restore system property if (savedOldProperty != null) { System.setProperty("drill.enable_unsafe_memory_access", savedOldProperty); } else { System.clearProperty("drill.enable_unsafe_memory_access"); } } /** * Test setting the bounds checking flag by the new property. * * @throws Exception if loading class {@link BoundsChecking#BOUNDS_CHECKING_ENABLED} fails. */ @Test public void testEnableNewProperty() throws Exception { String savedNewProperty = System.getProperty("arrow.enable_unsafe_memory_access"); System.setProperty("arrow.enable_unsafe_memory_access", "true"); ClassLoader classLoader = copyClassLoader(); if (classLoader != null) { boolean boundsCheckingEnabled = getFlagValue(classLoader); assertFalse(boundsCheckingEnabled); } // restore system property if (savedNewProperty != null) { System.setProperty("arrow.enable_unsafe_memory_access", savedNewProperty); } else { System.clearProperty("arrow.enable_unsafe_memory_access"); } } /** * Test setting the bounds checking flag by both old and new properties. In this case, the new * property should take precedence. * * @throws Exception if loading class {@link BoundsChecking#BOUNDS_CHECKING_ENABLED} fails. */ @Test public void testEnableBothProperties() throws Exception { String savedOldProperty = System.getProperty("drill.enable_unsafe_memory_access"); String savedNewProperty = System.getProperty("arrow.enable_unsafe_memory_access"); System.setProperty("drill.enable_unsafe_memory_access", "false"); System.setProperty("arrow.enable_unsafe_memory_access", "true"); // new property takes precedence. ClassLoader classLoader = copyClassLoader(); if (classLoader != null) { boolean boundsCheckingEnabled = getFlagValue(classLoader); assertFalse(boundsCheckingEnabled); } // restore system property if (savedOldProperty != null) { System.setProperty("drill.enable_unsafe_memory_access", savedOldProperty); } else { System.clearProperty("drill.enable_unsafe_memory_access"); } if (savedNewProperty != null) { System.setProperty("arrow.enable_unsafe_memory_access", savedNewProperty); } else { System.clearProperty("arrow.enable_unsafe_memory_access"); } } } TestForeignAllocation.java000066400000000000000000000125311515451615200350320ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 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.util.MemoryUtil; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestForeignAllocation { BufferAllocator allocator; @BeforeEach public void before() { allocator = new RootAllocator(); } @AfterEach public void after() { allocator.close(); } @Test public void wrapForeignAllocation() { final long bufferSize = 16; UnsafeForeignAllocation allocation = new UnsafeForeignAllocation(bufferSize); try { assertEquals(0, allocator.getAllocatedMemory()); ArrowBuf buf = allocator.wrapForeignAllocation(allocation); assertEquals(bufferSize, buf.capacity()); buf.close(); assertTrue(allocation.released); } finally { allocation.release0(); } assertEquals(0, allocator.getAllocatedMemory()); } @Test public void wrapForeignAllocationWithAllocationListener() { final long bufferSize = 16; final CountingAllocationListener listener = new CountingAllocationListener(); try (BufferAllocator listenedAllocator = allocator.newChildAllocator("child", listener, 0L, allocator.getLimit())) { UnsafeForeignAllocation allocation = new UnsafeForeignAllocation(bufferSize); try { assertEquals(0, listenedAllocator.getAllocatedMemory()); ArrowBuf buf = listenedAllocator.wrapForeignAllocation(allocation); assertEquals(bufferSize, buf.capacity()); assertEquals(16, listener.getCurrentMem()); buf.close(); assertEquals(0, listener.getCurrentMem()); assertTrue(allocation.released); } finally { allocation.release0(); } assertEquals(0, listenedAllocator.getAllocatedMemory()); } assertEquals(1, listener.getNumPreCalls()); assertEquals(1, listener.getNumCalls()); assertEquals(1, listener.getNumReleaseCalls()); assertEquals(16, listener.getTotalMem()); } @Test public void wrapForeignAllocationFailedWithAllocationListener() { assertThrows( OutOfMemoryException.class, () -> { final long bufferSize = 16; final long limit = bufferSize - 1; final CountingAllocationListener listener = new CountingAllocationListener(); try (BufferAllocator listenedAllocator = allocator.newChildAllocator("child", listener, 0L, limit)) { UnsafeForeignAllocation allocation = new UnsafeForeignAllocation(bufferSize); try { assertEquals(0, listenedAllocator.getAllocatedMemory()); ArrowBuf buf = listenedAllocator.wrapForeignAllocation(allocation); assertEquals(bufferSize, buf.capacity()); buf.close(); assertTrue(allocation.released); } finally { allocation.release0(); } } }); } @Test public void wrapForeignAllocationWithAllocationListenerReclaimingSpace() { final long bufferSize = 16; final long limit = 2 * bufferSize - 1; final List buffersToBeFreed = new ArrayList<>(); final AllocationListener listener = new AllocationListener() { @Override public boolean onFailedAllocation(long size, AllocationOutcome outcome) { buffersToBeFreed.forEach(ArrowBuf::close); return true; } }; try (BufferAllocator listenedAllocator = allocator.newChildAllocator("child", listener, 0L, limit)) { final ArrowBuf buffer1 = listenedAllocator.buffer(bufferSize); buffersToBeFreed.add(buffer1); UnsafeForeignAllocation allocation = new UnsafeForeignAllocation(bufferSize); try (final ArrowBuf buffer2 = listenedAllocator.wrapForeignAllocation(allocation)) { assertEquals(bufferSize, buffer2.capacity()); assertEquals( 0, buffer1.getReferenceManager().getRefCount()); // buffer1 was closed by listener } } } private static class UnsafeForeignAllocation extends ForeignAllocation { boolean released = false; public UnsafeForeignAllocation(long bufferSize) { super(bufferSize, MemoryUtil.allocateMemory(bufferSize)); } @Override protected void release0() { if (!released) { MemoryUtil.freeMemory(memoryAddress()); released = true; } } } } TestLowCostIdentityHashMap.java000066400000000000000000000124331515451615200360020ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 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.assertTrue; import org.junit.jupiter.api.Test; /** To test simplified implementation of IdentityHashMap. */ public class TestLowCostIdentityHashMap { @Test public void testIdentityHashMap() { LowCostIdentityHashMap hashMap = new LowCostIdentityHashMap<>(); StringWithKey obj1 = new StringWithKey("s1key", "s1value"); StringWithKey obj2 = new StringWithKey("s2key", "s2value"); StringWithKey obj3 = new StringWithKey("s3key", "s3value"); StringWithKey obj4 = new StringWithKey("s1key", "s4value"); StringWithKey obj5 = new StringWithKey("s5key", "s5value"); assertNull(hashMap.put(obj1)); assertNull(hashMap.put(obj2)); assertNull(hashMap.put(obj3)); assertEquals(obj1, hashMap.put(obj4)); assertNull(hashMap.put(obj5)); assertEquals(4, hashMap.size()); assertEquals(obj4, hashMap.get("s1key")); assertNull(hashMap.remove("abc")); assertEquals(obj3, hashMap.remove("s3key")); assertEquals(3, hashMap.size()); assertFalse(hashMap.isEmpty()); StringWithKey nextValue = hashMap.getNextValue(); assertNotNull(nextValue); assertTrue( (hashMap.get("s1key") == nextValue || hashMap.get("s2key") == nextValue || hashMap.get("s5key") == nextValue)); assertTrue(hashMap.containsValue(obj4)); assertTrue(hashMap.containsValue(obj2)); assertTrue(hashMap.containsValue(obj5)); assertEquals(obj4, hashMap.remove("s1key")); nextValue = hashMap.getNextValue(); assertNotNull(nextValue); assertTrue(hashMap.get("s2key") == nextValue || hashMap.get("s5key") == nextValue); assertEquals(2, hashMap.size()); assertEquals(obj2, hashMap.remove("s2key")); assertEquals(obj5, hashMap.remove("s5key")); assertEquals(0, hashMap.size()); assertTrue(hashMap.isEmpty()); } @Test public void testLargeMap() { LowCostIdentityHashMap hashMap = new LowCostIdentityHashMap<>(); String[] keys = new String[200]; for (int i = 0; i < 200; i++) { keys[i] = "s" + i + "key"; } for (int i = 0; i < 100; i++) { if (i % 5 == 0 && i != 0) { StringWithKey obj = new StringWithKey(keys[i - 5], "s" + i + "value"); StringWithKey retObj = hashMap.put(obj); assertNotNull(retObj); StringWithKey obj1 = new StringWithKey(keys[i], "s" + 2 * i + "value"); StringWithKey retObj1 = hashMap.put(obj1); assertNull(retObj1); } else { StringWithKey obj = new StringWithKey(keys[i], "s" + i + "value"); StringWithKey retObj = hashMap.put(obj); assertNull(retObj); } } assertEquals(100, hashMap.size()); for (int i = 0; i < 100; i++) { StringWithKey returnObj = hashMap.get(keys[i]); assertNotNull(returnObj); if (i == 95) { assertEquals("s190value", returnObj.getValue()); continue; } if (i % 5 == 0) { assertEquals("s" + (i + 5) + "value", returnObj.getValue()); } else { assertEquals("s" + i + "value", returnObj.getValue()); } } for (int i = 0; i < 100; i++) { if (i % 4 == 0) { StringWithKey returnObj = hashMap.remove(keys[i]); assertNotNull(returnObj); assertFalse(hashMap.containsKey(keys[i])); } StringWithKey obj = new StringWithKey(keys[100 + i], "s" + (100 + i) + "value"); StringWithKey retObj = hashMap.put(obj); assertNull(retObj); assertTrue(hashMap.containsKey(keys[100 + i])); } assertEquals(175, hashMap.size()); for (int i = 0; i < 100; i++) { StringWithKey retObj = hashMap.getNextValue(); assertNotNull(retObj); hashMap.remove(retObj.getKey()); } assertFalse(hashMap.isEmpty()); assertEquals(75, hashMap.size()); hashMap.clear(); assertTrue(hashMap.isEmpty()); } private static class StringWithKey implements ValueWithKeyIncluded { private final String myValue; private final String myKey; StringWithKey(String myKey, String myValue) { this.myKey = myKey; this.myValue = myValue; } @Override public String getKey() { return myKey; } String getValue() { return myValue; } } } apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/memory/TestOpens.java000066400000000000000000000040031515451615200325710ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 static org.junit.jupiter.api.Assertions.assertThrows; 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 { /** 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, () -> MemoryUtil.directBuffer(0, 10)); boolean found = false; while (e != null) { 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"); } } apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/memory/util/000077500000000000000000000000001515451615200307625ustar00rootroot00000000000000TestArrowBufPointer.java000066400000000000000000000150121515451615200354750ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/memory/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertNotSame; 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.memory.util.hash.ArrowBufHasher; import org.apache.arrow.memory.util.hash.SimpleHasher; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link ArrowBufPointer}. */ public class TestArrowBufPointer { private final int BUFFER_LENGTH = 1024; private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testArrowBufPointersEqual() { try (ArrowBuf buf1 = allocator.buffer(BUFFER_LENGTH); ArrowBuf buf2 = allocator.buffer(BUFFER_LENGTH)) { for (int i = 0; i < BUFFER_LENGTH / 4; i++) { buf1.setInt(i * 4L, i * 1234); buf2.setInt(i * 4L, i * 1234); } ArrowBufPointer ptr1 = new ArrowBufPointer(null, 0, 100); ArrowBufPointer ptr2 = new ArrowBufPointer(null, 100, 5032); assertEquals(ptr1, ptr2); for (int i = 0; i < BUFFER_LENGTH / 4; i++) { ptr1.set(buf1, i * 4L, 4); ptr2.set(buf2, i * 4L, 4); assertEquals(ptr1, ptr2); } } } @Test public void testArrowBufPointersHashCode() { final int vectorLength = 100; try (ArrowBuf buf1 = allocator.buffer(vectorLength * 4); ArrowBuf buf2 = allocator.buffer(vectorLength * 4)) { for (int i = 0; i < vectorLength; i++) { buf1.setInt(i * 4L, i); buf2.setInt(i * 4L, i); } CounterHasher hasher1 = new CounterHasher(); CounterHasher hasher2 = new CounterHasher(); ArrowBufPointer pointer1 = new ArrowBufPointer(hasher1); assertEquals(ArrowBufPointer.NULL_HASH_CODE, pointer1.hashCode()); ArrowBufPointer pointer2 = new ArrowBufPointer(hasher2); assertEquals(ArrowBufPointer.NULL_HASH_CODE, pointer2.hashCode()); for (int i = 0; i < vectorLength; i++) { pointer1.set(buf1, i * 4L, 4); pointer2.set(buf2, i * 4L, 4); assertEquals(pointer1.hashCode(), pointer2.hashCode()); // verify that the hash codes have been re-computed assertEquals(hasher1.counter, i + 1); assertEquals(hasher2.counter, i + 1); } } } @Test public void testNullPointersHashCode() { ArrowBufPointer pointer = new ArrowBufPointer(); assertEquals(ArrowBufPointer.NULL_HASH_CODE, pointer.hashCode()); pointer.set(null, 0, 0); assertEquals(ArrowBufPointer.NULL_HASH_CODE, pointer.hashCode()); } @Test public void testReuseHashCode() { try (ArrowBuf buf = allocator.buffer(10)) { buf.setInt(0, 10); buf.setInt(4, 20); CounterHasher hasher = new CounterHasher(); ArrowBufPointer pointer = new ArrowBufPointer(hasher); pointer.set(buf, 0, 4); int hashCode = pointer.hashCode(); // hash code computed assertEquals(1, hasher.counter); // no hash code re-compute assertEquals(hashCode, pointer.hashCode()); assertEquals(1, hasher.counter); // hash code re-computed pointer.set(buf, 4, 4); assertNotSame(hashCode, pointer.hashCode()); assertEquals(2, hasher.counter); } } @Test public void testHashersForEquality() { try (ArrowBuf buf = allocator.buffer(10)) { // pointer 1 uses the default hasher ArrowBufPointer pointer1 = new ArrowBufPointer(buf, 0, 10); // pointer 2 uses the counter hasher ArrowBufPointer pointer2 = new ArrowBufPointer(buf, 0, 10, new CounterHasher()); // the two pointers cannot be equal, since they have different hashers assertNotEquals(pointer1, pointer2); } } @Test public void testArrowBufPointersComparison() { final int vectorLength = 100; try (ArrowBuf buf1 = allocator.buffer(vectorLength); ArrowBuf buf2 = allocator.buffer(vectorLength)) { for (int i = 0; i < vectorLength; i++) { buf1.setByte(i, i); buf2.setByte(i, i); } ArrowBufPointer pointer1 = new ArrowBufPointer(); ArrowBufPointer pointer2 = new ArrowBufPointer(); pointer1.set(buf1, 0, 10); pointer2.set(buf2, 0, 10); assertEquals(0, pointer1.compareTo(pointer2)); pointer1.set(null, 0, 0); pointer2.set(null, 0, 0); assertEquals(0, pointer1.compareTo(pointer2)); pointer2.set(buf2, 0, 5); assertTrue(pointer1.compareTo(pointer2) < 0); pointer1.set(buf1, 0, 10); assertTrue(pointer1.compareTo(pointer2) > 0); pointer1.set(buf1, 1, 5); pointer2.set(buf2, 3, 8); assertTrue(pointer1.compareTo(pointer2) < 0); } } /** * Hasher with a counter that increments each time a hash code is calculated. This is to validate * that the hash code in {@link ArrowBufPointer} is reused. */ static class CounterHasher implements ArrowBufHasher { protected int counter = 0; @Override public int hashCode(long address, long length) { counter += 1; return SimpleHasher.INSTANCE.hashCode(address, length); } @Override public int hashCode(ArrowBuf buf, long offset, long length) { counter += 1; return SimpleHasher.INSTANCE.hashCode(buf, offset, length); } @Override public int hashCode() { return super.hashCode(); } @Override public boolean equals(Object o) { return o instanceof CounterHasher; } } } TestByteFunctionHelpers.java000066400000000000000000000115351515451615200363470ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/memory/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util; import static org.junit.jupiter.api.Assertions.assertEquals; import java.nio.charset.StandardCharsets; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestByteFunctionHelpers { private BufferAllocator allocator; private static final int SIZE = 100; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testEquals() { ArrowBuf buffer1 = allocator.buffer(SIZE); ArrowBuf buffer2 = allocator.buffer(SIZE); for (int i = 0; i < SIZE; i++) { buffer1.setByte(i, i); buffer2.setByte(i, i); } // test three cases, length>8, length>3, length<3 assertEquals(1, ByteFunctionHelpers.equal(buffer1, 0, SIZE - 1, buffer2, 0, SIZE - 1)); assertEquals(1, ByteFunctionHelpers.equal(buffer1, 0, 6, buffer2, 0, 6)); assertEquals(1, ByteFunctionHelpers.equal(buffer1, 0, 2, buffer2, 0, 2)); // change value at index1 buffer1.setByte(1, 10); assertEquals(0, ByteFunctionHelpers.equal(buffer1, 0, SIZE - 1, buffer2, 0, SIZE - 1)); assertEquals(0, ByteFunctionHelpers.equal(buffer1, 0, 6, buffer2, 0, 6)); assertEquals(0, ByteFunctionHelpers.equal(buffer1, 0, 2, buffer2, 0, 2)); buffer1.close(); buffer2.close(); } @Test public void testCompare() { ArrowBuf buffer1 = allocator.buffer(SIZE); ArrowBuf buffer2 = allocator.buffer(SIZE); for (int i = 0; i < SIZE; i++) { buffer1.setByte(i, i); buffer2.setByte(i, i); } // test three cases, length>8, length>3, length<3 assertEquals(0, ByteFunctionHelpers.compare(buffer1, 0, SIZE - 1, buffer2, 0, SIZE - 1)); assertEquals(0, ByteFunctionHelpers.compare(buffer1, 0, 6, buffer2, 0, 6)); assertEquals(0, ByteFunctionHelpers.compare(buffer1, 0, 2, buffer2, 0, 2)); // change value at index 1 buffer1.setByte(1, 0); assertEquals(-1, ByteFunctionHelpers.compare(buffer1, 0, SIZE - 1, buffer2, 0, SIZE - 1)); assertEquals(-1, ByteFunctionHelpers.compare(buffer1, 0, 6, buffer2, 0, 6)); assertEquals(-1, ByteFunctionHelpers.compare(buffer1, 0, 2, buffer2, 0, 2)); buffer1.close(); buffer2.close(); } @Test public void testStringCompare() { String[] leftStrings = {"cat", "cats", "catworld", "dogs", "bags"}; String[] rightStrings = {"dog", "dogs", "dogworld", "dog", "sgab"}; for (int i = 0; i < leftStrings.length; ++i) { String leftStr = leftStrings[i]; String rightStr = rightStrings[i]; ArrowBuf left = allocator.buffer(SIZE); left.setBytes(0, leftStr.getBytes(StandardCharsets.UTF_8)); ArrowBuf right = allocator.buffer(SIZE); right.setBytes(0, rightStr.getBytes(StandardCharsets.UTF_8)); assertEquals( leftStr.compareTo(rightStr) < 0 ? -1 : 1, ByteFunctionHelpers.compare(left, 0, leftStr.length(), right, 0, rightStr.length())); left.close(); right.close(); } } @Test public void testCompareWithByteArray() { ArrowBuf buffer1 = allocator.buffer(SIZE); byte[] buffer2 = new byte[SIZE]; for (int i = 0; i < SIZE; i++) { buffer1.setByte(i, i); buffer2[i] = (byte) i; } // test three cases, length>8, length>3, length<3 assertEquals(0, ByteFunctionHelpers.compare(buffer1, 0, SIZE - 1, buffer2, 0, SIZE - 1)); assertEquals(0, ByteFunctionHelpers.compare(buffer1, 0, 6, buffer2, 0, 6)); assertEquals(0, ByteFunctionHelpers.compare(buffer1, 0, 2, buffer2, 0, 2)); // change value at index 1 buffer1.setByte(1, 0); assertEquals(-1, ByteFunctionHelpers.compare(buffer1, 0, SIZE - 1, buffer2, 0, SIZE - 1)); assertEquals(-1, ByteFunctionHelpers.compare(buffer1, 0, 6, buffer2, 0, 6)); assertEquals(-1, ByteFunctionHelpers.compare(buffer1, 0, 2, buffer2, 0, 2)); buffer1.close(); } } TestLargeMemoryUtil.java000077500000000000000000000101751515451615200354760ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/memory/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertInstanceOf; import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Method; import java.net.URLClassLoader; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; public class TestLargeMemoryUtil { /** * Get a copy of the current class loader. * * @return the newly created class loader. */ private ClassLoader copyClassLoader() { ClassLoader curClassLoader = this.getClass().getClassLoader(); if (curClassLoader instanceof URLClassLoader) { // for Java 1.8 return new URLClassLoader(((URLClassLoader) curClassLoader).getURLs(), null); } // for Java 1.9 and Java 11. return null; } /** * Use the checkedCastToInt method from the current classloader. * * @param classLoader the class loader from which to call the method. * @return the return value of the method. */ private int checkedCastToInt(ClassLoader classLoader, long value) throws Exception { Class clazz = classLoader.loadClass("org.apache.arrow.memory.util.LargeMemoryUtil"); Method method = clazz.getMethod("checkedCastToInt", long.class); return (int) method.invoke(null, value); } private void checkExpectedOverflow(ClassLoader classLoader, long value) { InvocationTargetException ex = Assertions.assertThrows( InvocationTargetException.class, () -> checkedCastToInt(classLoader, value)); assertInstanceOf(ArithmeticException.class, ex.getCause()); assertEquals("integer overflow", ex.getCause().getMessage()); } @Test public void testEnableLargeMemoryUtilCheck() throws Exception { String savedNewProperty = System.getProperty("arrow.enable_unsafe_memory_access"); System.setProperty("arrow.enable_unsafe_memory_access", "false"); try { ClassLoader classLoader = copyClassLoader(); if (classLoader != null) { assertEquals(Integer.MAX_VALUE, checkedCastToInt(classLoader, Integer.MAX_VALUE)); checkExpectedOverflow(classLoader, Integer.MAX_VALUE + 1L); checkExpectedOverflow(classLoader, Integer.MIN_VALUE - 1L); } } finally { // restore system property if (savedNewProperty != null) { System.setProperty("arrow.enable_unsafe_memory_access", savedNewProperty); } else { System.clearProperty("arrow.enable_unsafe_memory_access"); } } } @Test public void testDisabledLargeMemoryUtilCheck() throws Exception { String savedNewProperty = System.getProperty("arrow.enable_unsafe_memory_access"); System.setProperty("arrow.enable_unsafe_memory_access", "true"); try { ClassLoader classLoader = copyClassLoader(); if (classLoader != null) { assertEquals(Integer.MAX_VALUE, checkedCastToInt(classLoader, Integer.MAX_VALUE)); assertEquals(Integer.MIN_VALUE, checkedCastToInt(classLoader, Integer.MAX_VALUE + 1L)); assertEquals(Integer.MAX_VALUE, checkedCastToInt(classLoader, Integer.MIN_VALUE - 1L)); } } finally { // restore system property if (savedNewProperty != null) { System.setProperty("arrow.enable_unsafe_memory_access", savedNewProperty); } else { System.clearProperty("arrow.enable_unsafe_memory_access"); } } } } apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/memory/util/hash/000077500000000000000000000000001515451615200317055ustar00rootroot00000000000000TestArrowBufHasher.java000066400000000000000000000113651515451615200362210ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/memory/util/hash/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.hash; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import java.nio.charset.StandardCharsets; 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.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; /** Test cases for {@link ArrowBufHasher} and its subclasses. */ public class TestArrowBufHasher { private final int BUFFER_LENGTH = 1024; private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @ParameterizedTest(name = "hasher = {0}") @MethodSource("getHasher") void testHasher(String name, ArrowBufHasher hasher) { try (ArrowBuf buf1 = allocator.buffer(BUFFER_LENGTH); ArrowBuf buf2 = allocator.buffer(BUFFER_LENGTH)) { // prepare data for (int i = 0; i < BUFFER_LENGTH / 4; i++) { buf1.setFloat(i * 4L, i / 10.0f); buf2.setFloat(i * 4L, i / 10.0f); } verifyHashCodesEqual(hasher, buf1, 0, 100, buf2, 0, 100); verifyHashCodesEqual(hasher, buf1, 1, 5, buf2, 1, 5); verifyHashCodesEqual(hasher, buf1, 10, 17, buf2, 10, 17); verifyHashCodesEqual(hasher, buf1, 33, 25, buf2, 33, 25); verifyHashCodesEqual(hasher, buf1, 22, 22, buf2, 22, 22); verifyHashCodesEqual(hasher, buf1, 123, 333, buf2, 123, 333); verifyHashCodesEqual(hasher, buf1, 374, 1, buf2, 374, 1); verifyHashCodesEqual(hasher, buf1, 11, 0, buf2, 11, 0); verifyHashCodesEqual(hasher, buf1, 75, 25, buf2, 75, 25); verifyHashCodesEqual(hasher, buf1, 0, 1024, buf2, 0, 1024); } } private void verifyHashCodesEqual( ArrowBufHasher hasher, ArrowBuf buf1, int offset1, int length1, ArrowBuf buf2, int offset2, int length2) { int hashCode1 = hasher.hashCode(buf1, offset1, length1); int hashCode2 = hasher.hashCode(buf2, offset2, length2); assertEquals(hashCode1, hashCode2); } @ParameterizedTest(name = "hasher = {0}") @MethodSource("getHasher") public void testHasherNegative(String name, ArrowBufHasher hasher) { try (ArrowBuf buf = allocator.buffer(BUFFER_LENGTH)) { // prepare data for (int i = 0; i < BUFFER_LENGTH / 4; i++) { buf.setFloat(i * 4L, i / 10.0f); } assertThrows(IllegalArgumentException.class, () -> hasher.hashCode(buf, 0, -1)); assertThrows(IndexOutOfBoundsException.class, () -> hasher.hashCode(buf, 0, 1028)); assertThrows(IndexOutOfBoundsException.class, () -> hasher.hashCode(buf, 500, 1000)); } } @ParameterizedTest(name = "hasher = {0}") @MethodSource("getHasher") public void testHasherLessThanInt(String name, ArrowBufHasher hasher) { try (ArrowBuf buf1 = allocator.buffer(4); ArrowBuf buf2 = allocator.buffer(4)) { buf1.writeBytes("foo1".getBytes(StandardCharsets.UTF_8)); buf2.writeBytes("bar2".getBytes(StandardCharsets.UTF_8)); for (int i = 1; i <= 4; i++) { verifyHashCodeNotEqual(hasher, buf1, i, buf2, i); } } } private void verifyHashCodeNotEqual( ArrowBufHasher hasher, ArrowBuf buf1, int length1, ArrowBuf buf2, int length2) { int hashCode1 = hasher.hashCode(buf1, 0, length1); int hashCode2 = hasher.hashCode(buf2, 0, length2); assertNotEquals(hashCode1, hashCode2); } private static Stream getHasher() { return Stream.of( Arguments.of(SimpleHasher.class.getSimpleName(), SimpleHasher.INSTANCE), Arguments.of(MurmurHasher.class.getSimpleName(), new MurmurHasher())); } } apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/util/000077500000000000000000000000001515451615200274525ustar00rootroot00000000000000TestAutoCloseables.java000066400000000000000000000211231515451615200340020ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } TestCollections2.java000066400000000000000000000046761515451615200334530ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.fail; 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; /** Tests for {@code Collections2} class. */ public class TestCollections2 { @Test public void testToImmutableListFromIterable() { final List source = new ArrayList<>(Arrays.asList("foo", "bar", "baz")); final List copy = Collections2.toImmutableList(source); assertEquals(source, copy); try { copy.add("unexpected"); fail("add operation should not be supported"); } catch (UnsupportedOperationException ignored) { } try { copy.set(0, "unexpected"); fail("set operation should not be supported"); } catch (UnsupportedOperationException ignored) { } try { copy.remove(0); fail("remove operation should not be supported"); } catch (UnsupportedOperationException ignored) { } source.set(1, "newvalue"); source.add("anothervalue"); assertEquals("bar", copy.get(1)); assertEquals(3, copy.size()); } @Test public void testStringFromEmptyIterator() { assertEquals("[]", Collections2.toString(Collections.emptyIterator())); } @Test public void testStringFromIterator() { Iterator iterator = Arrays.asList("foo", "bar", "baz").iterator(); iterator.next(); assertEquals("[bar, baz]", Collections2.toString(iterator)); assertFalse(iterator.hasNext()); } } apache-arrow-java-19.0.0/memory/memory-core/src/test/java/org/apache/arrow/util/TestStackTrace.java000066400000000000000000000031651515451615200332060ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.memory.util.StackTrace; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; public class TestStackTrace { /** Check that the stack trace includes the origin line. */ @Test public void testStackTraceComplete() { final String stackTrace = new StackTrace().toString(); Assertions.assertTrue(stackTrace.contains("TestStackTrace.testStackTraceComplete"), stackTrace); } /** Check that the stack trace doesn't include getStackTrace or StackTrace. */ @Test public void testStackTraceOmit() { final String stackTrace = new StackTrace().toString(); Assertions.assertFalse(stackTrace.contains("Thread.getStackTrace"), stackTrace); Assertions.assertFalse( stackTrace.contains("org.apache.arrow.memory.util.StackTrace"), stackTrace); } } apache-arrow-java-19.0.0/memory/memory-core/src/test/resources/000077500000000000000000000000001515451615200244245ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-core/src/test/resources/logback.xml000066400000000000000000000025221515451615200265510ustar00rootroot00000000000000 %d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/000077500000000000000000000000001515451615200234235ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/pom.xml000066400000000000000000000035441515451615200247460ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-memory 19.0.0 arrow-memory-netty-buffer-patch Arrow Memory - Netty Buffer Netty Buffer needed to patch that is consumed by Arrow Memory Netty org.apache.arrow arrow-memory-core io.netty netty-buffer io.netty netty-common org.slf4j slf4j-api apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/000077500000000000000000000000001515451615200242125ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/main/000077500000000000000000000000001515451615200251365ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/main/java/000077500000000000000000000000001515451615200260575ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/main/java/io/000077500000000000000000000000001515451615200264665ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/main/java/io/netty/000077500000000000000000000000001515451615200276315ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/main/java/io/netty/buffer/000077500000000000000000000000001515451615200311025ustar00rootroot00000000000000ExpandableByteBuf.java000066400000000000000000000035361515451615200352210ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/main/java/io/netty/buffer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package io.netty.buffer; import org.apache.arrow.memory.BufferAllocator; /** * Allows us to decorate ArrowBuf to make it expandable so that we can use them in the context of * the Netty framework (thus supporting RPC level memory accounting). */ public class ExpandableByteBuf extends MutableWrappedByteBuf { private final BufferAllocator allocator; public ExpandableByteBuf(ByteBuf buffer, BufferAllocator allocator) { super(buffer); this.allocator = allocator; } @Override public ByteBuf copy(int index, int length) { return new ExpandableByteBuf(buffer.copy(index, length), allocator); } @Override public ByteBuf capacity(int newCapacity) { if (newCapacity > capacity()) { ByteBuf newBuf = NettyArrowBuf.unwrapBuffer(allocator.buffer(newCapacity)); newBuf.writeBytes(buffer, 0, buffer.capacity()); newBuf.readerIndex(buffer.readerIndex()); newBuf.writerIndex(buffer.writerIndex()); buffer.release(); buffer = newBuf; return newBuf; } else { return super.capacity(newCapacity); } } } LargeBuffer.java000066400000000000000000000022371515451615200340560ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/main/java/io/netty/buffer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package io.netty.buffer; /** * A MutableWrappedByteBuf that also maintains a metric of the number of huge buffer bytes and * counts. */ public class LargeBuffer extends MutableWrappedByteBuf { public LargeBuffer(ByteBuf buffer) { super(buffer); } @Override public ByteBuf copy(int index, int length) { return new LargeBuffer(buffer.copy(index, length)); } } MutableWrappedByteBuf.java000066400000000000000000000227131515451615200360700ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/main/java/io/netty/buffer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package io.netty.buffer; import io.netty.util.ByteProcessor; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.nio.channels.FileChannel; import java.nio.channels.GatheringByteChannel; import java.nio.channels.ScatteringByteChannel; /** * This is basically a complete copy of netty's DuplicatedByteBuf. We copy because we want to * override some behaviors and make buffer mutable. */ abstract class MutableWrappedByteBuf extends AbstractByteBuf { ByteBuf buffer; public MutableWrappedByteBuf(ByteBuf buffer) { super(buffer.maxCapacity()); if (buffer instanceof MutableWrappedByteBuf) { this.buffer = ((MutableWrappedByteBuf) buffer).buffer; } else { this.buffer = buffer; } setIndex(buffer.readerIndex(), buffer.writerIndex()); } @Override public ByteBuffer nioBuffer(int index, int length) { return unwrap().nioBuffer(index, length); } @Override public ByteBuf unwrap() { return buffer; } @Override public ByteBufAllocator alloc() { return buffer.alloc(); } @Override public ByteOrder order() { return buffer.order(); } @Override public boolean isDirect() { return buffer.isDirect(); } @Override public int capacity() { return buffer.capacity(); } @Override public ByteBuf capacity(int newCapacity) { buffer.capacity(newCapacity); return this; } @Override public boolean hasArray() { return buffer.hasArray(); } @Override public byte[] array() { return buffer.array(); } @Override public int arrayOffset() { return buffer.arrayOffset(); } @Override public boolean hasMemoryAddress() { return buffer.hasMemoryAddress(); } @Override public long memoryAddress() { return buffer.memoryAddress(); } @Override public byte getByte(int index) { return _getByte(index); } @Override protected byte _getByte(int index) { return buffer.getByte(index); } @Override public short getShort(int index) { return _getShort(index); } @Override protected short _getShort(int index) { return buffer.getShort(index); } @Override public short getShortLE(int index) { return buffer.getShortLE(index); } @Override protected short _getShortLE(int index) { return buffer.getShortLE(index); } @Override public int getUnsignedMedium(int index) { return _getUnsignedMedium(index); } @Override protected int _getUnsignedMedium(int index) { return buffer.getUnsignedMedium(index); } @Override public int getUnsignedMediumLE(int index) { return buffer.getUnsignedMediumLE(index); } @Override protected int _getUnsignedMediumLE(int index) { return buffer.getUnsignedMediumLE(index); } @Override public int getInt(int index) { return _getInt(index); } @Override protected int _getInt(int index) { return buffer.getInt(index); } @Override public int getIntLE(int index) { return buffer.getIntLE(index); } @Override protected int _getIntLE(int index) { return buffer.getIntLE(index); } @Override public long getLong(int index) { return _getLong(index); } @Override protected long _getLong(int index) { return buffer.getLong(index); } @Override public long getLongLE(int index) { return buffer.getLongLE(index); } @Override protected long _getLongLE(int index) { return buffer.getLongLE(index); } @Override public abstract ByteBuf copy(int index, int length); @Override public ByteBuf slice(int index, int length) { return new SlicedByteBuf(this, index, length); } @Override public ByteBuf getBytes(int index, ByteBuf dst, int dstIndex, int length) { buffer.getBytes(index, dst, dstIndex, length); return this; } @Override public ByteBuf getBytes(int index, byte[] dst, int dstIndex, int length) { buffer.getBytes(index, dst, dstIndex, length); return this; } @Override public ByteBuf getBytes(int index, ByteBuffer dst) { buffer.getBytes(index, dst); return this; } @Override public ByteBuf setByte(int index, int value) { _setByte(index, value); return this; } @Override protected void _setByte(int index, int value) { buffer.setByte(index, value); } @Override public ByteBuf setShort(int index, int value) { _setShort(index, value); return this; } @Override protected void _setShort(int index, int value) { buffer.setShort(index, value); } @Override public ByteBuf setShortLE(int index, int value) { buffer.setShortLE(index, value); return this; } @Override protected void _setShortLE(int index, int value) { buffer.setShortLE(index, value); } @Override public ByteBuf setMedium(int index, int value) { _setMedium(index, value); return this; } @Override protected void _setMedium(int index, int value) { buffer.setMedium(index, value); } @Override public ByteBuf setMediumLE(int index, int value) { buffer.setMediumLE(index, value); return this; } @Override protected void _setMediumLE(int index, int value) { buffer.setMediumLE(index, value); } @Override public ByteBuf setInt(int index, int value) { _setInt(index, value); return this; } @Override protected void _setInt(int index, int value) { buffer.setInt(index, value); } @Override public ByteBuf setIntLE(int index, int value) { buffer.setIntLE(index, value); return this; } @Override protected void _setIntLE(int index, int value) { buffer.setIntLE(index, value); } @Override public ByteBuf setLong(int index, long value) { _setLong(index, value); return this; } @Override protected void _setLong(int index, long value) { buffer.setLong(index, value); } @Override public ByteBuf setLongLE(int index, long value) { buffer.setLongLE(index, value); return this; } @Override protected void _setLongLE(int index, long value) { buffer.setLongLE(index, value); } @Override public ByteBuf setBytes(int index, byte[] src, int srcIndex, int length) { buffer.setBytes(index, src, srcIndex, length); return this; } @Override public ByteBuf setBytes(int index, ByteBuf src, int srcIndex, int length) { buffer.setBytes(index, src, srcIndex, length); return this; } @Override public ByteBuf setBytes(int index, ByteBuffer src) { buffer.setBytes(index, src); return this; } @Override public int setBytes(int index, FileChannel in, long position, int length) throws IOException { return buffer.setBytes(index, in, position, length); } @Override public ByteBuf getBytes(int index, OutputStream out, int length) throws IOException { buffer.getBytes(index, out, length); return this; } @Override public int getBytes(int index, GatheringByteChannel out, int length) throws IOException { return buffer.getBytes(index, out, length); } @Override public int setBytes(int index, InputStream in, int length) throws IOException { return buffer.setBytes(index, in, length); } @Override public int setBytes(int index, ScatteringByteChannel in, int length) throws IOException { return buffer.setBytes(index, in, length); } @Override public int getBytes(int index, FileChannel out, long position, int length) throws IOException { return buffer.getBytes(index, out, position, length); } @Override public int nioBufferCount() { return buffer.nioBufferCount(); } @Override public ByteBuffer[] nioBuffers(int index, int length) { return buffer.nioBuffers(index, length); } @Override public ByteBuffer internalNioBuffer(int index, int length) { return nioBuffer(index, length); } @Override public int forEachByte(int index, int length, ByteProcessor processor) { return buffer.forEachByte(index, length, processor); } @Override public int forEachByteDesc(int index, int length, ByteProcessor processor) { return buffer.forEachByteDesc(index, length, processor); } @Override public final int refCnt() { return unwrap().refCnt(); } @Override public final ByteBuf touch() { unwrap().touch(); return this; } @Override public final ByteBuf touch(Object hint) { unwrap().touch(hint); return this; } @Override public final ByteBuf retain() { unwrap().retain(); return this; } @Override public final ByteBuf retain(int increment) { unwrap().retain(increment); return this; } @Override public boolean release() { return release(1); } @Override public boolean release(int decrement) { boolean released = unwrap().release(decrement); return released; } } NettyArrowBuf.java000066400000000000000000000433051515451615200344460ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/main/java/io/netty/buffer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package io.netty.buffer; import io.netty.util.internal.PlatformDependent; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.nio.channels.FileChannel; import java.nio.channels.GatheringByteChannel; import java.nio.channels.ScatteringByteChannel; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BoundsChecking; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.patch.ArrowByteBufAllocator; import org.apache.arrow.memory.util.LargeMemoryUtil; import org.apache.arrow.util.Preconditions; import org.apache.arrow.util.VisibleForTesting; /** Netty specific wrapper over ArrowBuf for use in Netty framework. */ public class NettyArrowBuf extends AbstractByteBuf implements AutoCloseable { private final ArrowBuf arrowBuf; private final ArrowByteBufAllocator arrowByteBufAllocator; private long length; private final long address; /** * Constructs a new instance. * * @param arrowBuf The buffer to wrap. * @param bufferAllocator The allocator for the buffer. * @param length The length of this buffer. * @deprecated Use {@link #NettyArrowBuf(ArrowBuf, BufferAllocator, long)} instead. */ @Deprecated(forRemoval = true) public NettyArrowBuf( final ArrowBuf arrowBuf, final BufferAllocator bufferAllocator, final int length) { this(arrowBuf, bufferAllocator, (long) length); } /** * Constructs a new instance. * * @param arrowBuf The buffer to wrap. * @param bufferAllocator The allocator for the buffer. * @param length The length of this buffer. */ public NettyArrowBuf( final ArrowBuf arrowBuf, final BufferAllocator bufferAllocator, final long length) { super((int) length); this.arrowBuf = arrowBuf; this.arrowByteBufAllocator = new ArrowByteBufAllocator(bufferAllocator); this.length = length; this.address = arrowBuf.memoryAddress(); } @Override public ByteBuf copy() { throw new UnsupportedOperationException(); } @Override public ByteBuf copy(int index, int length) { throw new UnsupportedOperationException(); } @Override public ByteBuf retain() { arrowBuf.getReferenceManager().retain(); return this; } public ArrowBuf arrowBuf() { return arrowBuf; } @Override public ByteBuf retain(final int increment) { arrowBuf.getReferenceManager().retain(increment); return this; } @Override public boolean isDirect() { return true; } @Override public synchronized ByteBuf capacity(int newCapacity) { if (newCapacity == length) { return this; } Preconditions.checkArgument(newCapacity >= 0); if (newCapacity < length) { length = newCapacity; return this; } throw new UnsupportedOperationException( "Buffers don't support resizing that increases the size."); } @Override public ByteBuf unwrap() { // According to Netty's ByteBuf interface, unwrap() should return null if the buffer cannot be // unwrapped // https://github.com/netty/netty/blob/9fe796e10a433b6cd20ad78b2c39cd56b86ccd2e/buffer/src/main/java/io/netty/buffer/ByteBuf.java#L305 // Throwing here breaks toString() in AbstractByteBuf // Since toString() is used to build debug / error messages, this can cause strange behavior return null; } @Override public int refCnt() { return arrowBuf.getReferenceManager().getRefCount(); } @Override public ArrowByteBufAllocator alloc() { return arrowByteBufAllocator; } @Override public boolean hasArray() { return false; } @Override public byte[] array() { throw new UnsupportedOperationException("Operation not supported on direct buffer"); } @Override public int arrayOffset() { throw new UnsupportedOperationException("Operation not supported on direct buffer"); } @Override public boolean hasMemoryAddress() { return true; } @Override public long memoryAddress() { return this.address; } @Override public ByteBuf touch() { return this; } @Override public ByteBuf touch(Object hint) { return this; } @Override public int capacity() { return (int) Math.min(Integer.MAX_VALUE, arrowBuf.capacity()); } @Override public NettyArrowBuf slice() { return unwrapBuffer(arrowBuf.slice(readerIndex, writerIndex - readerIndex)); } @Override public NettyArrowBuf slice(int index, int length) { return unwrapBuffer(arrowBuf.slice(index, length)); } @Override public void close() { arrowBuf.close(); } @Override public boolean release() { return arrowBuf.getReferenceManager().release(); } @Override public boolean release(int decrement) { return arrowBuf.getReferenceManager().release(decrement); } @Override public NettyArrowBuf readerIndex(int readerIndex) { super.readerIndex(readerIndex); return this; } @Override public NettyArrowBuf writerIndex(int writerIndex) { super.writerIndex(writerIndex); return this; } @Override public int nioBufferCount() { return 1; } @Override public ByteBuffer internalNioBuffer(int index, int length) { ByteBuffer nioBuf = getDirectBuffer(index); // Follows convention from other ByteBuf implementations. return (ByteBuffer) nioBuf.clear().limit(length); } @Override public ByteBuffer[] nioBuffers() { return new ByteBuffer[] {nioBuffer()}; } @Override public ByteBuffer[] nioBuffers(int index, int length) { return new ByteBuffer[] {nioBuffer(index, length)}; } @Override public ByteBuffer nioBuffer() { return nioBuffer(readerIndex(), readableBytes()); } /** * Returns a buffer that is zero positioned but points to a slice of the original buffer starting * at given index. */ @Override public ByteBuffer nioBuffer(int index, int length) { chk(index, length); final ByteBuffer buffer = getDirectBuffer(index); buffer.limit(length); return buffer; } /** * Returns a buffer that is zero positioned but points to a slice of the original buffer starting * at given index. */ public ByteBuffer nioBuffer(long index, int length) { chk(index, length); final ByteBuffer buffer = getDirectBuffer(index); buffer.limit(length); return buffer; } /** * Get this ArrowBuf as a direct {@link ByteBuffer}. * * @return ByteBuffer */ private ByteBuffer getDirectBuffer(long index) { return PlatformDependent.directBuffer( addr(index), LargeMemoryUtil.checkedCastToInt(length - index)); } @Override public ByteBuf getBytes(int index, ByteBuffer dst) { arrowBuf.getBytes(index, dst); return this; } @Override public ByteBuf setBytes(int index, ByteBuffer src) { arrowBuf.setBytes(index, src); return this; } @Override public ByteBuf getBytes(int index, byte[] dst, int dstIndex, int length) { arrowBuf.getBytes(index, dst, dstIndex, length); return this; } @Override public ByteBuf setBytes(int index, byte[] src, int srcIndex, int length) { arrowBuf.setBytes(index, src, srcIndex, length); return this; } /** * Determine if the requested {@code index} and {@code length} will fit within {@code capacity}. * * @param index The starting index. * @param length The length which will be utilized (starting from {@code index}). * @param capacity The capacity that {@code index + length} is allowed to be within. * @return {@code true} if the requested {@code index} and {@code length} will fit within {@code * capacity}. {@code false} if this would result in an index out of bounds exception. */ private static boolean isOutOfBounds(int index, int length, int capacity) { return (index | length | (index + length) | (capacity - (index + length))) < 0; } @Override public ByteBuf getBytes(int index, ByteBuf dst, int dstIndex, int length) { chk(index, length); Preconditions.checkArgument(dst != null, "Expecting valid dst ByteBuffer"); if (isOutOfBounds(dstIndex, length, dst.capacity())) { throw new IndexOutOfBoundsException("dstIndex: " + dstIndex + " length: " + length); } else { final long srcAddress = addr(index); if (dst.hasMemoryAddress()) { final long dstAddress = dst.memoryAddress() + (long) dstIndex; PlatformDependent.copyMemory(srcAddress, dstAddress, (long) length); } else if (dst.hasArray()) { dstIndex += dst.arrayOffset(); PlatformDependent.copyMemory(srcAddress, dst.array(), dstIndex, (long) length); } else { dst.setBytes(dstIndex, this, index, length); } } return this; } @Override public ByteBuf setBytes(int index, ByteBuf src, int srcIndex, int length) { chk(index, length); Preconditions.checkArgument(src != null, "Expecting valid src ByteBuffer"); if (isOutOfBounds(srcIndex, length, src.capacity())) { throw new IndexOutOfBoundsException("srcIndex: " + srcIndex + " length: " + length); } else { if (length != 0) { final long dstAddress = addr(index); if (src.hasMemoryAddress()) { final long srcAddress = src.memoryAddress() + (long) srcIndex; PlatformDependent.copyMemory(srcAddress, dstAddress, (long) length); } else if (src.hasArray()) { srcIndex += src.arrayOffset(); PlatformDependent.copyMemory(src.array(), srcIndex, dstAddress, (long) length); } else { src.getBytes(srcIndex, this, index, length); } } } return this; } @Override public ByteBuf getBytes(int index, OutputStream out, int length) throws IOException { arrowBuf.getBytes(index, out, length); return this; } @Override public int setBytes(int index, InputStream in, int length) throws IOException { return arrowBuf.setBytes(index, in, length); } @Override public int getBytes(int index, GatheringByteChannel out, int length) throws IOException { Preconditions.checkArgument(out != null, "expecting valid gathering byte channel"); chk(index, length); if (length == 0) { return 0; } else { final ByteBuffer tmpBuf = getDirectBuffer(index); tmpBuf.clear().limit(length); return out.write(tmpBuf); } } @Override public int getBytes(int index, FileChannel out, long position, int length) throws IOException { chk(index, length); if (length == 0) { return 0; } else { final ByteBuffer tmpBuf = getDirectBuffer(index); tmpBuf.clear().limit(length); return out.write(tmpBuf, position); } } @Override public int setBytes(int index, ScatteringByteChannel in, int length) throws IOException { return (int) in.read(nioBuffers(index, length)); } @Override public int setBytes(int index, FileChannel in, long position, int length) throws IOException { return (int) in.read(nioBuffers(index, length)); } @Override public ByteOrder order() { return ByteOrder.LITTLE_ENDIAN; } @Override public ByteBuf order(ByteOrder endianness) { return this; } @Override protected int _getUnsignedMedium(int index) { return getUnsignedMedium(index); } @Override protected int _getUnsignedMediumLE(int index) { this.chk(index, 3); long addr = this.addr(index); return (PlatformDependent.getByte(addr) & 255) | (Short.reverseBytes(PlatformDependent.getShort(addr + 1L)) & '\uffff') << 8; } /*-------------------------------------------------* | | | get() APIs | | | *-------------------------------------------------*/ @Override protected byte _getByte(int index) { return getByte(index); } @Override public byte getByte(int index) { return arrowBuf.getByte(index); } @Override protected short _getShortLE(int index) { short s = getShort(index); return Short.reverseBytes(s); } @Override protected short _getShort(int index) { return getShort(index); } @Override public short getShort(int index) { return arrowBuf.getShort(index); } @Override protected int _getIntLE(int index) { int value = getInt(index); return Integer.reverseBytes(value); } @Override protected int _getInt(int index) { return getInt(index); } @Override public int getInt(int index) { return arrowBuf.getInt(index); } @Override protected long _getLongLE(int index) { long value = getLong(index); return Long.reverseBytes(value); } @Override protected long _getLong(int index) { return getLong(index); } @Override public long getLong(int index) { return arrowBuf.getLong(index); } /*-------------------------------------------------* | | | set() APIs | | | *-------------------------------------------------*/ @Override protected void _setByte(int index, int value) { setByte(index, value); } @Override public NettyArrowBuf setByte(int index, int value) { arrowBuf.setByte(index, value); return this; } @Override protected void _setShortLE(int index, int value) { this.chk(index, 2); PlatformDependent.putShort(this.addr(index), Short.reverseBytes((short) value)); } @Override protected void _setShort(int index, int value) { setShort(index, value); } @Override public NettyArrowBuf setShort(int index, int value) { arrowBuf.setShort(index, value); return this; } private long addr(long index) { return address + index; } /** * Helper function to do bounds checking at a particular index for particular length of data. * * @param index index (0 based relative to this ArrowBuf) * @param fieldLength provided length of data for get/set */ private void chk(long index, long fieldLength) { if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { // check reference count ensureAccessible(); // check bounds if (fieldLength < 0) { throw new IllegalArgumentException("length: " + fieldLength + " (expected: >= 0)"); } if (index < 0 || index > capacity() - fieldLength) { throw new IndexOutOfBoundsException( String.format( "index: %d, length: %d (expected: range(0, %d))", index, fieldLength, capacity())); } } } @Override protected void _setMedium(int index, int value) { setMedium(index, value); } @Override protected void _setMediumLE(int index, int value) { this.chk(index, 3); long addr = this.addr(index); PlatformDependent.putByte(addr, (byte) value); PlatformDependent.putShort(addr + 1L, Short.reverseBytes((short) (value >>> 8))); } @Override public NettyArrowBuf setMedium(int index, int value) { chk(index, 3); final long addr = addr(index); // we need to store 3 bytes starting from least significant byte // and ignoring the most significant byte // since arrow memory format is little endian, we will // first store the first 2 bytes followed by third byte // example: if the 4 byte int value is ABCD where A is MSB // D is LSB then we effectively want to store DCB in increasing // address to get Little Endian byte order // (short)value will give us CD and PlatformDependent.putShort() // will store them in LE order as DC starting at address addr // in order to get B, we do ABCD >>> 16 = 00AB => (byte)AB which // gives B. We store this at address addr + 2. So finally we get // DCB PlatformDependent.putShort(addr, (short) value); PlatformDependent.putByte(addr + 2, (byte) (value >>> 16)); return this; } @Override @VisibleForTesting protected void _setInt(int index, int value) { setInt(index, value); } @Override @VisibleForTesting protected void _setIntLE(int index, int value) { this.chk(index, 4); PlatformDependent.putInt(this.addr(index), Integer.reverseBytes(value)); } @Override public NettyArrowBuf setInt(int index, int value) { arrowBuf.setInt(index, value); return this; } @Override protected void _setLong(int index, long value) { setLong(index, value); } @Override public void _setLongLE(int index, long value) { this.chk(index, 8); PlatformDependent.putLong(this.addr(index), Long.reverseBytes(value)); } @Override public NettyArrowBuf setLong(int index, long value) { arrowBuf.setLong(index, value); return this; } /** unwrap arrow buffer into a netty buffer. */ public static NettyArrowBuf unwrapBuffer(ArrowBuf buf) { final NettyArrowBuf nettyArrowBuf = new NettyArrowBuf( buf, buf.getReferenceManager().getAllocator(), LargeMemoryUtil.checkedCastToInt(buf.capacity())); nettyArrowBuf.readerIndex(LargeMemoryUtil.checkedCastToInt(buf.readerIndex())); nettyArrowBuf.writerIndex(LargeMemoryUtil.checkedCastToInt(buf.writerIndex())); return nettyArrowBuf; } } PooledByteBufAllocatorL.java000066400000000000000000000221221515451615200363450ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/main/java/io/netty/buffer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package io.netty.buffer; import io.netty.util.internal.OutOfDirectMemoryError; import io.netty.util.internal.StringUtil; import java.lang.reflect.Field; import java.nio.ByteBuffer; import java.util.concurrent.atomic.AtomicLong; import org.apache.arrow.memory.OutOfMemoryException; import org.apache.arrow.memory.util.AssertionUtil; import org.apache.arrow.memory.util.LargeMemoryUtil; /** * The base allocator that we use for all of Arrow's memory management. Returns * UnsafeDirectLittleEndian buffers. */ public class PooledByteBufAllocatorL { private static final org.slf4j.Logger memoryLogger = org.slf4j.LoggerFactory.getLogger("arrow.allocator"); private static final int MEMORY_LOGGER_FREQUENCY_SECONDS = 60; public final UnsafeDirectLittleEndian empty; private final AtomicLong hugeBufferSize = new AtomicLong(0); private final AtomicLong hugeBufferCount = new AtomicLong(0); private final AtomicLong normalBufferSize = new AtomicLong(0); private final AtomicLong normalBufferCount = new AtomicLong(0); private final InnerAllocator allocator; public PooledByteBufAllocatorL() { allocator = new InnerAllocator(); empty = new UnsafeDirectLittleEndian(new DuplicatedByteBuf(Unpooled.EMPTY_BUFFER)); } /** Returns a {@linkplain UnsafeDirectLittleEndian} of the given size. */ public UnsafeDirectLittleEndian allocate(long size) { try { return allocator.directBuffer(LargeMemoryUtil.checkedCastToInt(size), Integer.MAX_VALUE); } catch (OutOfMemoryError e) { /* * OutOfDirectMemoryError is thrown by Netty when we exceed the direct memory limit defined by * -XX:MaxDirectMemorySize. OutOfMemoryError with "Direct buffer memory" message is thrown by * java.nio.Bits when we exceed the direct memory limit. This should never be hit in practice * as Netty is expected to throw an OutOfDirectMemoryError first. */ if (e instanceof OutOfDirectMemoryError || "Direct buffer memory".equals(e.getMessage())) { throw new OutOfMemoryException("Failure allocating buffer.", e); } throw e; } } public int getChunkSize() { return allocator.chunkSize(); } public long getHugeBufferSize() { return hugeBufferSize.get(); } public long getHugeBufferCount() { return hugeBufferCount.get(); } public long getNormalBufferSize() { return normalBufferSize.get(); } public long getNormalBufferCount() { return normalBufferSize.get(); } private static class AccountedUnsafeDirectLittleEndian extends UnsafeDirectLittleEndian { private final long initialCapacity; private final AtomicLong count; private final AtomicLong size; private AccountedUnsafeDirectLittleEndian(LargeBuffer buf, AtomicLong count, AtomicLong size) { super(buf); this.initialCapacity = buf.capacity(); this.count = count; this.size = size; } private AccountedUnsafeDirectLittleEndian( PooledUnsafeDirectByteBuf buf, AtomicLong count, AtomicLong size) { super(buf); this.initialCapacity = buf.capacity(); this.count = count; this.size = size; } @Override public ByteBuf copy() { throw new UnsupportedOperationException("copy method is not supported"); } @Override public ByteBuf copy(int index, int length) { throw new UnsupportedOperationException("copy method is not supported"); } @Override public boolean release(int decrement) { boolean released = super.release(decrement); if (released) { count.decrementAndGet(); size.addAndGet(-initialCapacity); } return released; } } private class InnerAllocator extends PooledByteBufAllocator { private final PoolArena[] directArenas; private final MemoryStatusThread statusThread; public InnerAllocator() { super(true); try { Field f = PooledByteBufAllocator.class.getDeclaredField("directArenas"); f.setAccessible(true); this.directArenas = (PoolArena[]) f.get(this); } catch (Exception e) { throw new RuntimeException( "Failure while initializing allocator. Unable to retrieve direct arenas field.", e); } if (memoryLogger.isTraceEnabled()) { statusThread = new MemoryStatusThread(this); statusThread.start(); } else { statusThread = null; } } private UnsafeDirectLittleEndian newDirectBufferL(int initialCapacity, int maxCapacity) { PoolThreadCache cache = threadCache(); PoolArena directArena = cache.directArena; if (directArena != null) { if (initialCapacity > chunkSize()) { // This is beyond chunk size so we'll allocate separately. ByteBuf buf = UnpooledByteBufAllocator.DEFAULT.directBuffer(initialCapacity, maxCapacity); hugeBufferSize.addAndGet(buf.capacity()); hugeBufferCount.incrementAndGet(); // logger.debug("Allocating huge buffer of size {}", initialCapacity, new Exception()); return new AccountedUnsafeDirectLittleEndian( new LargeBuffer(buf), hugeBufferCount, hugeBufferSize); } else { // within chunk, use arena. ByteBuf buf = directArena.allocate(cache, initialCapacity, maxCapacity); if (!(buf instanceof PooledUnsafeDirectByteBuf)) { fail(); } if (!AssertionUtil.ASSERT_ENABLED) { return new UnsafeDirectLittleEndian((PooledUnsafeDirectByteBuf) buf); } normalBufferSize.addAndGet(buf.capacity()); normalBufferCount.incrementAndGet(); return new AccountedUnsafeDirectLittleEndian( (PooledUnsafeDirectByteBuf) buf, normalBufferCount, normalBufferSize); } } else { throw fail(); } } private UnsupportedOperationException fail() { return new UnsupportedOperationException( "Arrow requires that the JVM used supports access sun.misc.Unsafe. This platform " + "didn't provide that functionality."); } @Override public UnsafeDirectLittleEndian directBuffer(int initialCapacity, int maxCapacity) { if (initialCapacity == 0 && maxCapacity == 0) { newDirectBuffer(initialCapacity, maxCapacity); } validate(initialCapacity, maxCapacity); return newDirectBufferL(initialCapacity, maxCapacity); } @Override public ByteBuf heapBuffer(int initialCapacity, int maxCapacity) { throw new UnsupportedOperationException("Arrow doesn't support using heap buffers."); } private void validate(int initialCapacity, int maxCapacity) { if (initialCapacity < 0) { throw new IllegalArgumentException( "initialCapacity: " + initialCapacity + " (expected: 0+)"); } if (initialCapacity > maxCapacity) { throw new IllegalArgumentException( String.format( "initialCapacity: %d (expected: not greater than maxCapacity(%d)", initialCapacity, maxCapacity)); } } @Override public String toString() { StringBuilder buf = new StringBuilder(); buf.append(directArenas.length); buf.append(" direct arena(s):"); buf.append(StringUtil.NEWLINE); for (PoolArena a : directArenas) { buf.append(a); } buf.append("Large buffers outstanding: "); buf.append(hugeBufferCount.get()); buf.append(" totaling "); buf.append(hugeBufferSize.get()); buf.append(" bytes."); buf.append('\n'); buf.append("Normal buffers outstanding: "); buf.append(normalBufferCount.get()); buf.append(" totaling "); buf.append(normalBufferSize.get()); buf.append(" bytes."); return buf.toString(); } private class MemoryStatusThread extends Thread { private final InnerAllocator allocator; public MemoryStatusThread(InnerAllocator allocator) { super("allocation.logger"); this.setDaemon(true); this.allocator = allocator; } @Override public void run() { while (true) { memoryLogger.trace("Memory Usage: \n{}", allocator); try { Thread.sleep(MEMORY_LOGGER_FREQUENCY_SECONDS * 1000); } catch (InterruptedException e) { return; } } } } } } UnsafeDirectLittleEndian.java000066400000000000000000000137511515451615200365460ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/main/java/io/netty/buffer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package io.netty.buffer; import io.netty.util.internal.PlatformDependent; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; import java.nio.ByteOrder; import java.util.concurrent.atomic.AtomicLong; /** * The underlying class we use for little-endian access to memory. Is used underneath ArrowBufs to * abstract away the Netty classes and underlying Netty memory management. */ public class UnsafeDirectLittleEndian extends WrappedByteBuf { private static final AtomicLong ID_GENERATOR = new AtomicLong(0); public final long id = ID_GENERATOR.incrementAndGet(); private final AbstractByteBuf wrapped; private final long memoryAddress; /** * Constructs a new instance. * * @param buf The buffer to wrap */ public UnsafeDirectLittleEndian(AbstractByteBuf buf) { super(buf); this.wrapped = buf; this.memoryAddress = buf.memoryAddress(); } private long addr(int index) { return memoryAddress + index; } @Override public long getLong(int index) { // wrapped.checkIndex(index, 8); long v = PlatformDependent.getLong(addr(index)); return v; } @Override public float getFloat(int index) { return Float.intBitsToFloat(getInt(index)); } @Override public ByteBuf slice() { return slice(this.readerIndex(), readableBytes()); } @Override public ByteBuf slice(int index, int length) { return new SlicedByteBuf(this, index, length); } @Override public ByteBuf order(ByteOrder endianness) { return this; } @Override public double getDouble(int index) { return Double.longBitsToDouble(getLong(index)); } @Override public char getChar(int index) { return (char) getShort(index); } @Override public long getUnsignedInt(int index) { return getInt(index) & 0xFFFFFFFFL; } @Override public int getInt(int index) { int v = PlatformDependent.getInt(addr(index)); return v; } @Override public int getUnsignedShort(int index) { return getShort(index) & 0xFFFF; } @Override public short getShort(int index) { short v = PlatformDependent.getShort(addr(index)); return v; } @Override public ByteBuf setShort(int index, int value) { wrapped.checkIndex(index, 2); setShort_(index, value); return this; } @Override public ByteBuf setInt(int index, int value) { wrapped.checkIndex(index, 4); setInt_(index, value); return this; } @Override public ByteBuf setLong(int index, long value) { wrapped.checkIndex(index, 8); setLong_(index, value); return this; } @Override public ByteBuf setChar(int index, int value) { setShort(index, value); return this; } @Override public ByteBuf setFloat(int index, float value) { setInt(index, Float.floatToRawIntBits(value)); return this; } @Override public ByteBuf setDouble(int index, double value) { setLong(index, Double.doubleToRawLongBits(value)); return this; } @Override public ByteBuf writeShort(int value) { wrapped.ensureWritable(2); setShort_(wrapped.writerIndex, value); wrapped.writerIndex += 2; return this; } @Override public ByteBuf writeInt(int value) { wrapped.ensureWritable(4); setInt_(wrapped.writerIndex, value); wrapped.writerIndex += 4; return this; } @Override public ByteBuf writeLong(long value) { wrapped.ensureWritable(8); setLong_(wrapped.writerIndex, value); wrapped.writerIndex += 8; return this; } @Override public ByteBuf writeChar(int value) { writeShort(value); return this; } @Override public ByteBuf writeFloat(float value) { writeInt(Float.floatToRawIntBits(value)); return this; } @Override public ByteBuf writeDouble(double value) { writeLong(Double.doubleToRawLongBits(value)); return this; } private void setShort_(int index, int value) { PlatformDependent.putShort(addr(index), (short) value); } private void setInt_(int index, int value) { PlatformDependent.putInt(addr(index), value); } private void setLong_(int index, long value) { PlatformDependent.putLong(addr(index), value); } @Override public byte getByte(int index) { return PlatformDependent.getByte(addr(index)); } @Override public ByteBuf setByte(int index, int value) { PlatformDependent.putByte(addr(index), (byte) value); return this; } @Override public boolean release() { return release(1); } @Override public int setBytes(int index, InputStream in, int length) throws IOException { wrapped.checkIndex(index, length); byte[] tmp = new byte[length]; int readBytes = in.read(tmp); if (readBytes > 0) { PlatformDependent.copyMemory(tmp, 0, addr(index), readBytes); } return readBytes; } @Override public ByteBuf getBytes(int index, OutputStream out, int length) throws IOException { wrapped.checkIndex(index, length); if (length != 0) { byte[] tmp = new byte[length]; PlatformDependent.copyMemory(addr(index), tmp, 0, length); out.write(tmp); } return this; } @Override public int hashCode() { return System.identityHashCode(this); } @Override public boolean equals(Object obj) { return this == obj; } } apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/main/java/org/000077500000000000000000000000001515451615200266465ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/main/java/org/apache/000077500000000000000000000000001515451615200300675ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200312215ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/main/java/org/apache/arrow/memory/000077500000000000000000000000001515451615200325315ustar00rootroot00000000000000patch/000077500000000000000000000000001515451615200335515ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/main/java/org/apache/arrow/memoryArrowByteBufAllocator.java000066400000000000000000000104541515451615200406340ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/main/java/org/apache/arrow/memory/patch/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.patch; import io.netty.buffer.AbstractByteBufAllocator; import io.netty.buffer.ByteBuf; import io.netty.buffer.CompositeByteBuf; import io.netty.buffer.ExpandableByteBuf; import io.netty.buffer.NettyArrowBuf; import org.apache.arrow.memory.BufferAllocator; /** * An implementation of ByteBufAllocator that wraps a Arrow BufferAllocator. This allows the RPC * layer to be accounted and managed using Arrow's BufferAllocator infrastructure. The only thin * different from a typical BufferAllocator is the signature and the fact that this Allocator * returns ExpandableByteBufs which enable otherwise non-expandable ArrowBufs to be expandable. * * @deprecated This class may be removed in a future release. */ @Deprecated public class ArrowByteBufAllocator extends AbstractByteBufAllocator { private static final int DEFAULT_BUFFER_SIZE = 4096; private static final int DEFAULT_MAX_COMPOSITE_COMPONENTS = 16; private final BufferAllocator allocator; public ArrowByteBufAllocator(BufferAllocator allocator) { this.allocator = allocator; } public BufferAllocator unwrap() { return allocator; } @Override public ByteBuf buffer() { return buffer(DEFAULT_BUFFER_SIZE); } @Override public ByteBuf buffer(int initialCapacity) { return new ExpandableByteBuf( NettyArrowBuf.unwrapBuffer(allocator.buffer(initialCapacity)), allocator); } @Override public ByteBuf buffer(int initialCapacity, int maxCapacity) { return buffer(initialCapacity); } @Override public ByteBuf ioBuffer() { return buffer(); } @Override public ByteBuf ioBuffer(int initialCapacity) { return buffer(initialCapacity); } @Override public ByteBuf ioBuffer(int initialCapacity, int maxCapacity) { return buffer(initialCapacity); } @Override public ByteBuf directBuffer() { return buffer(); } @Override public ByteBuf directBuffer(int initialCapacity) { return NettyArrowBuf.unwrapBuffer(allocator.buffer(initialCapacity)); } @Override public ByteBuf directBuffer(int initialCapacity, int maxCapacity) { return buffer(initialCapacity, maxCapacity); } @Override public CompositeByteBuf compositeBuffer() { return compositeBuffer(DEFAULT_MAX_COMPOSITE_COMPONENTS); } @Override public CompositeByteBuf compositeBuffer(int maxNumComponents) { return new CompositeByteBuf(this, true, maxNumComponents); } @Override public CompositeByteBuf compositeDirectBuffer() { return compositeBuffer(); } @Override public CompositeByteBuf compositeDirectBuffer(int maxNumComponents) { return compositeBuffer(maxNumComponents); } @Override public boolean isDirectBufferPooled() { return false; } @Override public ByteBuf heapBuffer() { throw fail(); } @Override public ByteBuf heapBuffer(int initialCapacity) { throw fail(); } @Override public ByteBuf heapBuffer(int initialCapacity, int maxCapacity) { throw fail(); } @Override public CompositeByteBuf compositeHeapBuffer() { throw fail(); } @Override public CompositeByteBuf compositeHeapBuffer(int maxNumComponents) { throw fail(); } @Override protected ByteBuf newHeapBuffer(int initialCapacity, int maxCapacity) { throw fail(); } @Override protected ByteBuf newDirectBuffer(int initialCapacity, int maxCapacity) { return buffer(initialCapacity, maxCapacity); } private RuntimeException fail() { throw new UnsupportedOperationException("Allocator doesn't support heap-based memory."); } } apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/test/000077500000000000000000000000001515451615200251715ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/test/java/000077500000000000000000000000001515451615200261125ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/test/java/io/000077500000000000000000000000001515451615200265215ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/test/java/io/netty/000077500000000000000000000000001515451615200276645ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/test/java/io/netty/buffer/000077500000000000000000000000001515451615200311355ustar00rootroot00000000000000TestUnsafeDirectLittleEndian.java000066400000000000000000000064661515451615200374460ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty-buffer-patch/src/test/java/io/netty/buffer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package io.netty.buffer; import static org.junit.jupiter.api.Assertions.assertEquals; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.charset.StandardCharsets; import org.junit.jupiter.api.Test; public class TestUnsafeDirectLittleEndian { @Test @SuppressWarnings("CatchAndPrintStackTrace") public void testPrimitiveGetSet() { ByteBuf byteBuf = Unpooled.directBuffer(64); UnsafeDirectLittleEndian unsafeDirect = getUnsafeDirectLittleEndian(byteBuf); assertEquals(Byte.MAX_VALUE, unsafeDirect.getByte(0)); assertEquals(-1, unsafeDirect.getByte(1)); assertEquals(Short.MAX_VALUE, unsafeDirect.getShort(2)); assertEquals(-2, unsafeDirect.getShort(4)); assertEquals((char) 65534, unsafeDirect.getChar(4)); assertEquals(Integer.MAX_VALUE, unsafeDirect.getInt(8)); assertEquals(-66052, unsafeDirect.getInt(12)); assertEquals(4294901244L, unsafeDirect.getUnsignedInt(12)); assertEquals(Long.MAX_VALUE, unsafeDirect.getLong(16)); assertEquals(-4295098372L, unsafeDirect.getLong(24)); assertEquals(1.23F, unsafeDirect.getFloat(32), 0.0); assertEquals(-1.23F, unsafeDirect.getFloat(36), 0.0); assertEquals(1.234567D, unsafeDirect.getDouble(40), 0.0); assertEquals(-1.234567D, unsafeDirect.getDouble(48), 0.0); byte[] inBytes = "1234567".getBytes(StandardCharsets.UTF_8); try (ByteArrayInputStream bais = new ByteArrayInputStream(inBytes); ByteArrayOutputStream baos = new ByteArrayOutputStream()) { assertEquals(5, unsafeDirect.setBytes(56, bais, 5)); unsafeDirect.getBytes(56, baos, 5); assertEquals("12345", new String(baos.toByteArray(), StandardCharsets.UTF_8)); } catch (IOException e) { e.printStackTrace(); } } private static UnsafeDirectLittleEndian getUnsafeDirectLittleEndian(ByteBuf byteBuf) { UnsafeDirectLittleEndian unsafeDirect = new UnsafeDirectLittleEndian(new LargeBuffer(byteBuf)); unsafeDirect.setByte(0, Byte.MAX_VALUE); unsafeDirect.setByte(1, -1); // 0xFF unsafeDirect.setShort(2, Short.MAX_VALUE); unsafeDirect.setShort(4, -2); // 0xFFFE unsafeDirect.setInt(8, Integer.MAX_VALUE); unsafeDirect.setInt(12, -66052); // 0xFFFE FDFC unsafeDirect.setLong(16, Long.MAX_VALUE); unsafeDirect.setLong(24, -4295098372L); // 0xFFFF FFFE FFFD FFFC unsafeDirect.setFloat(32, 1.23F); unsafeDirect.setFloat(36, -1.23F); unsafeDirect.setDouble(40, 1.234567D); unsafeDirect.setDouble(48, -1.234567D); return unsafeDirect; } } apache-arrow-java-19.0.0/memory/memory-netty/000077500000000000000000000000001515451615200210575ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/pom.xml000066400000000000000000000063351515451615200224030ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-memory 19.0.0 arrow-memory-netty Arrow Memory - Netty Netty allocator and utils for allocating memory in Arrow org.apache.arrow arrow-memory-core org.apache.arrow arrow-memory-netty-buffer-patch ${project.version} io.netty netty-buffer provided io.netty netty-common org.slf4j slf4j-api test ch.qos.logback logback-core test org.immutables value-annotations integration-tests org.apache.maven.plugins maven-failsafe-plugin default-it integration-test verify apache-arrow-java-19.0.0/memory/memory-netty/src/000077500000000000000000000000001515451615200216465ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/main/000077500000000000000000000000001515451615200225725ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/main/java/000077500000000000000000000000001515451615200235135ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/main/java/org/000077500000000000000000000000001515451615200243025ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/main/java/org/apache/000077500000000000000000000000001515451615200255235ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200266555ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/main/java/org/apache/arrow/memory/000077500000000000000000000000001515451615200301655ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/main/java/org/apache/arrow/memory/netty/000077500000000000000000000000001515451615200313305ustar00rootroot00000000000000DefaultAllocationManagerFactory.java000066400000000000000000000026461515451615200403410ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/main/java/org/apache/arrow/memory/netty/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.netty; import org.apache.arrow.memory.AllocationManager; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; /** The default Allocation Manager Factory for a module. */ public class DefaultAllocationManagerFactory implements AllocationManager.Factory { public static final AllocationManager.Factory FACTORY = NettyAllocationManager.FACTORY; @Override public AllocationManager create(BufferAllocator accountingAllocator, long size) { return FACTORY.create(accountingAllocator, size); } @Override public ArrowBuf empty() { return FACTORY.empty(); } } NettyAllocationManager.java000066400000000000000000000104011515451615200365140ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/main/java/org/apache/arrow/memory/netty/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.netty; import io.netty.buffer.PooledByteBufAllocatorL; import io.netty.buffer.UnsafeDirectLittleEndian; import io.netty.util.internal.PlatformDependent; import org.apache.arrow.memory.AllocationManager; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReferenceManager; /** * The default implementation of {@link AllocationManager}. The implementation is responsible for * managing when memory is allocated and returned to the Netty-based PooledByteBufAllocatorL. */ public class NettyAllocationManager extends AllocationManager { public static final AllocationManager.Factory FACTORY = new AllocationManager.Factory() { @Override public AllocationManager create(BufferAllocator accountingAllocator, long size) { return new NettyAllocationManager(accountingAllocator, size); } @Override public ArrowBuf empty() { return EMPTY_BUFFER; } }; /** * The default cut-off value for switching allocation strategies. If the request size is not * greater than the cut-off value, we will allocate memory by {@link PooledByteBufAllocatorL} * APIs, otherwise, we will use {@link PlatformDependent} APIs. */ public static final int DEFAULT_ALLOCATION_CUTOFF_VALUE = Integer.MAX_VALUE; private static final PooledByteBufAllocatorL INNER_ALLOCATOR = new PooledByteBufAllocatorL(); static final UnsafeDirectLittleEndian EMPTY = INNER_ALLOCATOR.empty; static final ArrowBuf EMPTY_BUFFER = new ArrowBuf(ReferenceManager.NO_OP, null, 0, NettyAllocationManager.EMPTY.memoryAddress()); static final long CHUNK_SIZE = INNER_ALLOCATOR.getChunkSize(); private final long allocatedSize; private final UnsafeDirectLittleEndian memoryChunk; private final long allocatedAddress; /** The cut-off value for switching allocation strategies. */ NettyAllocationManager( BufferAllocator accountingAllocator, long requestedSize, int allocationCutOffValue) { super(accountingAllocator); if (requestedSize > allocationCutOffValue) { this.memoryChunk = null; this.allocatedAddress = PlatformDependent.allocateMemory(requestedSize); this.allocatedSize = requestedSize; } else { this.memoryChunk = INNER_ALLOCATOR.allocate(requestedSize); this.allocatedAddress = memoryChunk.memoryAddress(); this.allocatedSize = memoryChunk.capacity(); } } NettyAllocationManager(BufferAllocator accountingAllocator, long requestedSize) { this(accountingAllocator, requestedSize, DEFAULT_ALLOCATION_CUTOFF_VALUE); } /** * Get the underlying memory chunk managed by this AllocationManager. * * @return the underlying memory chunk if the request size is not greater than the cutoff value * provided in the constructor , or null otherwise. * @deprecated this method will be removed in a future release. */ @Deprecated UnsafeDirectLittleEndian getMemoryChunk() { return memoryChunk; } @Override protected long memoryAddress() { return allocatedAddress; } @Override protected void release0() { if (memoryChunk == null) { PlatformDependent.freeMemory(allocatedAddress); } else { memoryChunk.release(); } } /** * Returns the underlying memory chunk size managed. * *

    NettyAllocationManager rounds requested size up to the next power of two. */ @Override public long getSize() { return allocatedSize; } } apache-arrow-java-19.0.0/memory/memory-netty/src/test/000077500000000000000000000000001515451615200226255ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/test/java/000077500000000000000000000000001515451615200235465ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/test/java/io/000077500000000000000000000000001515451615200241555ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/test/java/io/netty/000077500000000000000000000000001515451615200253205ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/test/java/io/netty/buffer/000077500000000000000000000000001515451615200265715ustar00rootroot00000000000000TestExpandableByteBuf.java000066400000000000000000000111051515451615200335370ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/test/java/io/netty/buffer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package io.netty.buffer; import static org.junit.jupiter.api.Assertions.assertEquals; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.Test; public class TestExpandableByteBuf { @Test public void testCapacity() { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(20)) { NettyArrowBuf srcByteBuf = NettyArrowBuf.unwrapBuffer(buf); ExpandableByteBuf expandableByteBuf = new ExpandableByteBuf(srcByteBuf, allocator); ByteBuf newByteBuf = expandableByteBuf.capacity(31); int capacity = newByteBuf.capacity(); assertEquals(32, capacity); } } @Test public void testCapacity1() { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(20)) { NettyArrowBuf srcByteBuf = NettyArrowBuf.unwrapBuffer(buf); ExpandableByteBuf expandableByteBuf = new ExpandableByteBuf(srcByteBuf, allocator); ByteBuf newByteBuf = expandableByteBuf.capacity(32); int capacity = newByteBuf.capacity(); assertEquals(32, capacity); } } @Test public void testSetAndGetIntValues() { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(20)) { NettyArrowBuf srcByteBuf = NettyArrowBuf.unwrapBuffer(buf); ExpandableByteBuf expandableByteBuf = new ExpandableByteBuf(srcByteBuf, allocator); int[] intVals = new int[] { Integer.MIN_VALUE, Short.MIN_VALUE - 1, Short.MIN_VALUE, 0, Short.MAX_VALUE, Short.MAX_VALUE + 1, Integer.MAX_VALUE }; for (int intValue : intVals) { expandableByteBuf.setInt(0, intValue); assertEquals(expandableByteBuf.getInt(0), intValue); assertEquals(expandableByteBuf.getIntLE(0), Integer.reverseBytes(intValue)); } } } @Test public void testSetAndGetLongValues() { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(20)) { NettyArrowBuf srcByteBuf = NettyArrowBuf.unwrapBuffer(buf); ExpandableByteBuf expandableByteBuf = new ExpandableByteBuf(srcByteBuf, allocator); long[] longVals = new long[] {Long.MIN_VALUE, 0, Long.MAX_VALUE}; for (long longValue : longVals) { expandableByteBuf.setLong(0, longValue); assertEquals(expandableByteBuf.getLong(0), longValue); assertEquals(expandableByteBuf.getLongLE(0), Long.reverseBytes(longValue)); } } } @Test public void testSetAndGetShortValues() { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(20)) { NettyArrowBuf srcByteBuf = NettyArrowBuf.unwrapBuffer(buf); ExpandableByteBuf expandableByteBuf = new ExpandableByteBuf(srcByteBuf, allocator); short[] shortVals = new short[] {Short.MIN_VALUE, 0, Short.MAX_VALUE}; for (short shortValue : shortVals) { expandableByteBuf.setShort(0, shortValue); assertEquals(expandableByteBuf.getShort(0), shortValue); assertEquals(expandableByteBuf.getShortLE(0), Short.reverseBytes(shortValue)); } } } @Test public void testSetAndGetByteValues() { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(20)) { NettyArrowBuf srcByteBuf = NettyArrowBuf.unwrapBuffer(buf); ExpandableByteBuf expandableByteBuf = new ExpandableByteBuf(srcByteBuf, allocator); byte[] byteVals = new byte[] {Byte.MIN_VALUE, 0, Byte.MAX_VALUE}; for (short byteValue : byteVals) { expandableByteBuf.setByte(0, byteValue); assertEquals(expandableByteBuf.getByte(0), byteValue); } } } } apache-arrow-java-19.0.0/memory/memory-netty/src/test/java/io/netty/buffer/TestNettyArrowBuf.java000066400000000000000000000127511515451615200330550ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package io.netty.buffer; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNull; import java.nio.ByteBuffer; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.memory.patch.ArrowByteBufAllocator; import org.junit.jupiter.api.Test; public class TestNettyArrowBuf { @Test public void testSliceWithoutArgs() { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(20)) { NettyArrowBuf nettyBuf = NettyArrowBuf.unwrapBuffer(buf); nettyBuf.writerIndex(20); nettyBuf.readerIndex(10); NettyArrowBuf slicedBuffer = nettyBuf.slice(); int readableBytes = slicedBuffer.readableBytes(); assertEquals(10, readableBytes); } } @Test public void testNioBuffer() { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(20)) { NettyArrowBuf nettyBuf = NettyArrowBuf.unwrapBuffer(buf); ByteBuffer byteBuffer = nettyBuf.nioBuffer(4, 6); // Nio Buffers should always be 0 indexed assertEquals(0, byteBuffer.position()); assertEquals(6, byteBuffer.limit()); // Underlying buffer has size 32 excluding 4 should have capacity of 28. assertEquals(28, byteBuffer.capacity()); } } @Test public void testInternalNioBuffer() { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(20)) { NettyArrowBuf nettyBuf = NettyArrowBuf.unwrapBuffer(buf); ByteBuffer byteBuffer = nettyBuf.internalNioBuffer(4, 6); assertEquals(0, byteBuffer.position()); assertEquals(6, byteBuffer.limit()); // Underlying buffer has size 32 excluding 4 should have capacity of 28. assertEquals(28, byteBuffer.capacity()); } } @Test public void testSetLEValues() { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(20)) { NettyArrowBuf nettyBuf = NettyArrowBuf.unwrapBuffer(buf); int[] intVals = new int[] { Integer.MIN_VALUE, Short.MIN_VALUE - 1, Short.MIN_VALUE, 0, Short.MAX_VALUE, Short.MAX_VALUE + 1, Integer.MAX_VALUE }; for (int intValue : intVals) { nettyBuf._setInt(0, intValue); assertEquals(nettyBuf._getIntLE(0), Integer.reverseBytes(intValue)); } long[] longVals = new long[] {Long.MIN_VALUE, 0, Long.MAX_VALUE}; for (long longValue : longVals) { nettyBuf._setLong(0, longValue); assertEquals(nettyBuf._getLongLE(0), Long.reverseBytes(longValue)); } short[] shortVals = new short[] {Short.MIN_VALUE, 0, Short.MAX_VALUE}; for (short shortValue : shortVals) { nettyBuf._setShort(0, shortValue); assertEquals(nettyBuf._getShortLE(0), Short.reverseBytes(shortValue)); } } } @Test public void testSetCompositeBuffer() { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(20); NettyArrowBuf buf2 = NettyArrowBuf.unwrapBuffer(allocator.buffer(20))) { CompositeByteBuf byteBufs = new CompositeByteBuf(new ArrowByteBufAllocator(allocator), true, 1); int expected = 4; buf2.setInt(0, expected); buf2.writerIndex(4); byteBufs.addComponent(true, buf2); NettyArrowBuf.unwrapBuffer(buf).setBytes(0, byteBufs, 4); int actual = buf.getInt(0); assertEquals(expected, actual); } } @Test public void testGetCompositeBuffer() { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(20)) { CompositeByteBuf byteBufs = new CompositeByteBuf(new ArrowByteBufAllocator(allocator), true, 1); int expected = 4; buf.setInt(0, expected); NettyArrowBuf buf2 = NettyArrowBuf.unwrapBuffer(allocator.buffer(20)); // composite buffers are a bit weird, need to jump hoops // to set capacity. byteBufs.addComponent(true, buf2); byteBufs.capacity(20); NettyArrowBuf.unwrapBuffer(buf).getBytes(0, byteBufs, 4); int actual = byteBufs.getInt(0); assertEquals(expected, actual); byteBufs.component(0).release(); } } @Test public void testUnwrapReturnsNull() { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(20)) { NettyArrowBuf nettyBuf = NettyArrowBuf.unwrapBuffer(buf); // NettyArrowBuf cannot be unwrapped, so unwrap() should return null per the Netty ByteBuf API assertNull(nettyBuf.unwrap()); } } } apache-arrow-java-19.0.0/memory/memory-netty/src/test/java/org/000077500000000000000000000000001515451615200243355ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/test/java/org/apache/000077500000000000000000000000001515451615200255565ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200267105ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/test/java/org/apache/arrow/memory/000077500000000000000000000000001515451615200302205ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/test/java/org/apache/arrow/memory/netty/000077500000000000000000000000001515451615200313635ustar00rootroot00000000000000ITTestLargeArrowBuf.java000066400000000000000000000050071515451615200357500ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/test/java/org/apache/arrow/memory/netty/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.netty; import static org.junit.jupiter.api.Assertions.assertEquals; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.Test; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * Integration test for large (more than 2GB) {@link org.apache.arrow.memory.ArrowBuf}. To run this * test, please make sure there is at least 4GB memory in the system. */ public class ITTestLargeArrowBuf { private static final Logger logger = LoggerFactory.getLogger(ITTestLargeArrowBuf.class); private void run(long bufSize) { try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); ArrowBuf largeBuf = allocator.buffer(bufSize)) { assertEquals(bufSize, largeBuf.capacity()); logger.trace("Successfully allocated a buffer with capacity {}", largeBuf.capacity()); for (long i = 0; i < bufSize / 8; i++) { largeBuf.setLong(i * 8, i); if ((i + 1) % 10000 == 0) { logger.trace("Successfully written {} long words", i + 1); } } logger.trace("Successfully written {} long words", bufSize / 8); for (long i = 0; i < bufSize / 8; i++) { long val = largeBuf.getLong(i * 8); assertEquals(i, val); if ((i + 1) % 10000 == 0) { logger.trace("Successfully read {} long words", i + 1); } } logger.trace("Successfully read {} long words", bufSize / 8); } logger.trace("Successfully released the large buffer."); } @Test public void testLargeArrowBuf() { run(4 * 1024 * 1024 * 1024L); } @Test public void testMaxIntArrowBuf() { run(Integer.MAX_VALUE); } } TestAllocationManagerNetty.java000066400000000000000000000030611515451615200374130ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/test/java/org/apache/arrow/memory/netty/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.netty; import static org.junit.jupiter.api.Assertions.assertEquals; import org.apache.arrow.memory.AllocationManager; import org.apache.arrow.memory.DefaultAllocationManagerOption; import org.junit.jupiter.api.Test; /** Test cases for {@link AllocationManager}. */ public class TestAllocationManagerNetty { @Test public void testAllocationManagerType() { // test netty allocation manager type System.setProperty( DefaultAllocationManagerOption.ALLOCATION_MANAGER_TYPE_PROPERTY_NAME, "Netty"); DefaultAllocationManagerOption.AllocationManagerType mgrType = DefaultAllocationManagerOption.getDefaultAllocationManagerType(); assertEquals(DefaultAllocationManagerOption.AllocationManagerType.Netty, mgrType); } } TestEmptyArrowBuf.java000066400000000000000000000065551515451615200355700ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/test/java/org/apache/arrow/memory/netty/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.netty; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import io.netty.buffer.PooledByteBufAllocatorL; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.ReferenceManager; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; public class TestEmptyArrowBuf { private static final int MAX_ALLOCATION = 8 * 1024; private static RootAllocator allocator; @BeforeAll public static void beforeClass() { allocator = new RootAllocator(MAX_ALLOCATION); } /** Ensure the allocator is closed. */ @AfterAll public static void afterClass() { if (allocator != null) { allocator.close(); } } @Test public void testZeroBuf() { // Exercise the historical log inside the empty ArrowBuf. This is initialized statically, and // there is a circular // dependency between ArrowBuf and BaseAllocator, so if the initialization happens in the wrong // order, the // historical log will be null even though RootAllocator.DEBUG is true. allocator.getEmpty().print(new StringBuilder(), 0, RootAllocator.Verbosity.LOG_WITH_STACKTRACE); } @Test public void testEmptyArrowBuf() { ArrowBuf buf = new ArrowBuf( ReferenceManager.NO_OP, null, 1024, new PooledByteBufAllocatorL().empty.memoryAddress()); buf.getReferenceManager().retain(); buf.getReferenceManager().retain(8); assertEquals(1024, buf.capacity()); assertEquals(1, buf.getReferenceManager().getRefCount()); assertEquals(0, buf.getActualMemoryConsumed()); for (int i = 0; i < 10; i++) { buf.setByte(i, i); } assertEquals(0, buf.getActualMemoryConsumed()); assertEquals(0, buf.getReferenceManager().getSize()); assertEquals(0, buf.getReferenceManager().getAccountedSize()); assertFalse(buf.getReferenceManager().release()); assertFalse(buf.getReferenceManager().release(2)); assertEquals(0, buf.getReferenceManager().getAllocator().getLimit()); assertEquals( buf, buf.getReferenceManager().transferOwnership(buf, allocator).getTransferredBuffer()); assertEquals(0, buf.readerIndex()); assertEquals(0, buf.writerIndex()); assertEquals(1, buf.refCnt()); ArrowBuf derive = buf.getReferenceManager().deriveBuffer(buf, 0, 100); assertEquals(derive, buf); assertEquals(1, buf.refCnt()); assertEquals(1, derive.refCnt()); buf.close(); } } TestEndianness.java000066400000000000000000000033411515451615200350770ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/test/java/org/apache/arrow/memory/netty/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.netty; import static org.junit.jupiter.api.Assertions.assertEquals; import io.netty.buffer.ByteBuf; import io.netty.buffer.NettyArrowBuf; import java.nio.ByteOrder; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.Test; public class TestEndianness { @Test public void testNativeEndian() { final BufferAllocator a = new RootAllocator(10000); final ByteBuf b = NettyArrowBuf.unwrapBuffer(a.buffer(4)); b.setInt(0, 35); if (ByteOrder.nativeOrder() == ByteOrder.LITTLE_ENDIAN) { assertEquals(35, b.getByte(0)); assertEquals(0, b.getByte(1)); assertEquals(0, b.getByte(2)); assertEquals(0, b.getByte(3)); } else { assertEquals(0, b.getByte(0)); assertEquals(0, b.getByte(1)); assertEquals(0, b.getByte(2)); assertEquals(35, b.getByte(3)); } b.release(); a.close(); } } TestNettyAllocationManager.java000066400000000000000000000102571515451615200374200ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/test/java/org/apache/arrow/memory/netty/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.netty; 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.assertNull; import org.apache.arrow.memory.AllocationManager; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.BufferLedger; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.Test; /** Test cases for {@link NettyAllocationManager}. */ public class TestNettyAllocationManager { static int CUSTOMIZED_ALLOCATION_CUTOFF_VALUE = 1024; private RootAllocator createCustomizedAllocator() { return new RootAllocator( RootAllocator.configBuilder() .allocationManagerFactory( new AllocationManager.Factory() { @Override public AllocationManager create(BufferAllocator accountingAllocator, long size) { return new NettyAllocationManager( accountingAllocator, size, CUSTOMIZED_ALLOCATION_CUTOFF_VALUE); } @Override public ArrowBuf empty() { return null; } }) .build()); } private void readWriteArrowBuf(ArrowBuf buffer) { // write buffer for (long i = 0; i < buffer.capacity() / 8; i++) { buffer.setLong(i * 8, i); } // read buffer for (long i = 0; i < buffer.capacity() / 8; i++) { long val = buffer.getLong(i * 8); assertEquals(i, val); } } /** Test the allocation strategy for small buffers.. */ @Test public void testSmallBufferAllocation() { final long bufSize = CUSTOMIZED_ALLOCATION_CUTOFF_VALUE - 512L; try (RootAllocator allocator = createCustomizedAllocator(); ArrowBuf buffer = allocator.buffer(bufSize)) { assertInstanceOf(BufferLedger.class, buffer.getReferenceManager()); BufferLedger bufferLedger = (BufferLedger) buffer.getReferenceManager(); // make sure we are using netty allocation manager AllocationManager allocMgr = bufferLedger.getAllocationManager(); assertInstanceOf(NettyAllocationManager.class, allocMgr); NettyAllocationManager nettyMgr = (NettyAllocationManager) allocMgr; // for the small buffer allocation strategy, the chunk is not null assertNotNull(nettyMgr.getMemoryChunk()); readWriteArrowBuf(buffer); } } /** Test the allocation strategy for large buffers.. */ @Test public void testLargeBufferAllocation() { final long bufSize = CUSTOMIZED_ALLOCATION_CUTOFF_VALUE + 1024L; try (RootAllocator allocator = createCustomizedAllocator(); ArrowBuf buffer = allocator.buffer(bufSize)) { assertInstanceOf(BufferLedger.class, buffer.getReferenceManager()); BufferLedger bufferLedger = (BufferLedger) buffer.getReferenceManager(); // make sure we are using netty allocation manager AllocationManager allocMgr = bufferLedger.getAllocationManager(); assertInstanceOf(NettyAllocationManager.class, allocMgr); NettyAllocationManager nettyMgr = (NettyAllocationManager) allocMgr; // for the large buffer allocation strategy, the chunk is null assertNull(nettyMgr.getMemoryChunk()); readWriteArrowBuf(buffer); } } } TestNettyAllocator.java000066400000000000000000000067001515451615200357560ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/test/java/org/apache/arrow/memory/netty/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.netty; import static org.junit.jupiter.api.Assertions.assertTrue; import ch.qos.logback.classic.Level; import ch.qos.logback.classic.Logger; import ch.qos.logback.classic.spi.ILoggingEvent; import ch.qos.logback.core.read.ListAppender; import io.netty.buffer.PooledByteBufAllocatorL; import java.util.Collections; import java.util.stream.Collectors; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.ReferenceManager; import org.junit.jupiter.api.Test; import org.slf4j.LoggerFactory; /** Test netty allocators. */ public class TestNettyAllocator { @Test @SuppressWarnings("SynchronizeOnNonFinalField") public void testMemoryUsage() { ListAppender memoryLogsAppender = new ListAppender<>(); memoryLogsAppender.list = Collections.synchronizedList(memoryLogsAppender.list); Logger logger = (Logger) LoggerFactory.getLogger("arrow.allocator"); try { logger.setLevel(Level.TRACE); logger.addAppender(memoryLogsAppender); memoryLogsAppender.start(); try (ArrowBuf buf = new ArrowBuf( ReferenceManager.NO_OP, null, 1024, new PooledByteBufAllocatorL().empty.memoryAddress())) { buf.memoryAddress(); } boolean result = false; long startTime = System.currentTimeMillis(); while ((System.currentTimeMillis() - startTime) < 10000) { // 10 seconds maximum for time to read logs // Lock on the list backing the appender since a background thread might try to add more // logs // while stream() is iterating over list elements. This would throw a flakey // ConcurrentModificationException. synchronized (memoryLogsAppender.list) { result = memoryLogsAppender.list.stream() .anyMatch( log -> log.toString().contains("Memory Usage: \n") && log.toString().contains("Large buffers outstanding: ") && log.toString().contains("Normal buffers outstanding: ") && log.getLevel().equals(Level.TRACE)); } if (result) { break; } } synchronized (memoryLogsAppender.list) { assertTrue( result, "Log messages are:\n" + memoryLogsAppender.list.stream() .map(ILoggingEvent::toString) .collect(Collectors.joining("\n"))); } } finally { memoryLogsAppender.stop(); logger.detachAppender(memoryLogsAppender); logger.setLevel(null); } } } apache-arrow-java-19.0.0/memory/memory-netty/src/test/resources/000077500000000000000000000000001515451615200246375ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-netty/src/test/resources/logback.xml000066400000000000000000000025221515451615200267640ustar00rootroot00000000000000 %d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n apache-arrow-java-19.0.0/memory/memory-unsafe/000077500000000000000000000000001515451615200211755ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/pom.xml000066400000000000000000000032061515451615200225130ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-memory 19.0.0 arrow-memory-unsafe Arrow Memory - Unsafe Allocator and utils for allocating memory in Arrow based on sun.misc.Unsafe org.apache.arrow arrow-memory-core org.immutables value-annotations apache-arrow-java-19.0.0/memory/memory-unsafe/src/000077500000000000000000000000001515451615200217645ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/main/000077500000000000000000000000001515451615200227105ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/main/java/000077500000000000000000000000001515451615200236315ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/main/java/module-info.java000066400000000000000000000017061515451615200267160ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ module org.apache.arrow.memory.unsafe { exports org.apache.arrow.memory.unsafe to org.apache.arrow.memory.core; requires org.apache.arrow.memory.core; } apache-arrow-java-19.0.0/memory/memory-unsafe/src/main/java/org/000077500000000000000000000000001515451615200244205ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/main/java/org/apache/000077500000000000000000000000001515451615200256415ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200267735ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/000077500000000000000000000000001515451615200303035ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/unsafe/000077500000000000000000000000001515451615200315645ustar00rootroot00000000000000DefaultAllocationManagerFactory.java000066400000000000000000000027001515451615200405640ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/unsafe/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.unsafe; import org.apache.arrow.memory.AllocationManager; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; /** The default Allocation Manager Factory for a module. */ public class DefaultAllocationManagerFactory implements AllocationManager.Factory { public static final AllocationManager.Factory FACTORY = UnsafeAllocationManager.FACTORY; @Override public AllocationManager create(BufferAllocator accountingAllocator, long size) { return FACTORY.create(accountingAllocator, size); } @Override public ArrowBuf empty() { return UnsafeAllocationManager.FACTORY.empty(); } } UnsafeAllocationManager.java000066400000000000000000000042701515451615200370750ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/unsafe/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.unsafe; import org.apache.arrow.memory.AllocationManager; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReferenceManager; import org.apache.arrow.memory.util.MemoryUtil; /** Allocation manager based on unsafe API. */ public final class UnsafeAllocationManager extends AllocationManager { private static final ArrowBuf EMPTY = new ArrowBuf(ReferenceManager.NO_OP, null, 0, MemoryUtil.allocateMemory(0)); public static final AllocationManager.Factory FACTORY = new Factory() { @Override public AllocationManager create(BufferAllocator accountingAllocator, long size) { return new UnsafeAllocationManager(accountingAllocator, size); } @Override public ArrowBuf empty() { return EMPTY; } }; private final long allocatedSize; private final long allocatedAddress; UnsafeAllocationManager(BufferAllocator accountingAllocator, long requestedSize) { super(accountingAllocator); allocatedAddress = MemoryUtil.allocateMemory(requestedSize); allocatedSize = requestedSize; } @Override public long getSize() { return allocatedSize; } @Override protected long memoryAddress() { return allocatedAddress; } @Override protected void release0() { MemoryUtil.freeMemory(allocatedAddress); } } apache-arrow-java-19.0.0/memory/memory-unsafe/src/test/000077500000000000000000000000001515451615200227435ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/test/java/000077500000000000000000000000001515451615200236645ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/test/java/org/000077500000000000000000000000001515451615200244535ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/test/java/org/apache/000077500000000000000000000000001515451615200256745ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200270265ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/test/java/org/apache/arrow/memory/000077500000000000000000000000001515451615200303365ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/test/java/org/apache/arrow/memory/unsafe/000077500000000000000000000000001515451615200316175ustar00rootroot00000000000000TestAllocationManagerUnsafe.java000066400000000000000000000030671515451615200377730ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/test/java/org/apache/arrow/memory/unsafe/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.unsafe; import static org.junit.jupiter.api.Assertions.assertEquals; import org.apache.arrow.memory.AllocationManager; import org.apache.arrow.memory.DefaultAllocationManagerOption; import org.junit.jupiter.api.Test; /** Test cases for {@link AllocationManager}. */ public class TestAllocationManagerUnsafe { @Test public void testAllocationManagerType() { // test unsafe allocation manager type System.setProperty( DefaultAllocationManagerOption.ALLOCATION_MANAGER_TYPE_PROPERTY_NAME, "Unsafe"); DefaultAllocationManagerOption.AllocationManagerType mgrType = DefaultAllocationManagerOption.getDefaultAllocationManagerType(); assertEquals(DefaultAllocationManagerOption.AllocationManagerType.Unsafe, mgrType); } } TestUnsafeAllocationManager.java000066400000000000000000000051501515451615200377660ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/test/java/org/apache/arrow/memory/unsafe/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.unsafe; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertInstanceOf; import org.apache.arrow.memory.AllocationManager; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.BufferLedger; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.Test; /** Test cases for {@link UnsafeAllocationManager}. */ public class TestUnsafeAllocationManager { private BufferAllocator createUnsafeAllocator() { return new RootAllocator( RootAllocator.configBuilder() .allocationManagerFactory(UnsafeAllocationManager.FACTORY) .build()); } private void readWriteArrowBuf(ArrowBuf buffer) { // write buffer for (long i = 0; i < buffer.capacity() / 8; i++) { buffer.setLong(i * 8, i); } // read buffer for (long i = 0; i < buffer.capacity() / 8; i++) { long val = buffer.getLong(i * 8); assertEquals(i, val); } } /** Test the memory allocation for {@link UnsafeAllocationManager}. */ @Test public void testBufferAllocation() { final long bufSize = 4096L; try (BufferAllocator allocator = createUnsafeAllocator(); ArrowBuf buffer = allocator.buffer(bufSize)) { assertInstanceOf(BufferLedger.class, buffer.getReferenceManager()); BufferLedger bufferLedger = (BufferLedger) buffer.getReferenceManager(); // make sure we are using unsafe allocation manager AllocationManager allocMgr = bufferLedger.getAllocationManager(); assertInstanceOf(UnsafeAllocationManager.class, allocMgr); UnsafeAllocationManager unsafeMgr = (UnsafeAllocationManager) allocMgr; assertEquals(bufSize, unsafeMgr.getSize()); readWriteArrowBuf(buffer); } } } apache-arrow-java-19.0.0/memory/memory-unsafe/src/test/resources/000077500000000000000000000000001515451615200247555ustar00rootroot00000000000000apache-arrow-java-19.0.0/memory/memory-unsafe/src/test/resources/logback.xml000066400000000000000000000025221515451615200271020ustar00rootroot00000000000000 %d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n apache-arrow-java-19.0.0/memory/pom.xml000066400000000000000000000026741515451615200177340ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-java-root 19.0.0 arrow-memory pom Arrow Memory memory-core memory-unsafe memory-netty-buffer-patch memory-netty apache-arrow-java-19.0.0/performance/000077500000000000000000000000001515451615200173775ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/pom.xml000066400000000000000000000143311515451615200207160ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-java-root 19.0.0 arrow-performance jar Arrow Performance Benchmarks JMH Performance benchmarks for other Arrow libraries. 1.37 benchmarks true .* 1 5 5 jmh-result.json json org.openjdk.jmh jmh-core ${jmh.version} org.apache.arrow arrow-vector ${arrow.vector.classifier} org.apache.arrow arrow-memory-core org.apache.arrow arrow-memory-netty runtime org.apache.avro avro ${dep.avro.version} org.apache.arrow arrow-avro com.h2database h2 2.3.232 runtime org.apache.arrow arrow-jdbc org.apache.arrow arrow-algorithm org.apache.maven.plugins maven-compiler-plugin org.openjdk.jmh jmh-generator-annprocess ${jmh.version} org.apache.maven.plugins maven-shade-plugin shade package ${uberjar.name} false org.openjdk.jmh.Main *:* META-INF/*.SF META-INF/*.DSA META-INF/*.RSA org.codehaus.mojo exec-maven-plugin ${skip.perf.benchmarks} test java -classpath org.openjdk.jmh.Main ${benchmark.filter} -f ${benchmark.forks} -jvmArgs ${benchmark.jvmargs} -wi ${benchmark.warmups} -i ${benchmark.runs} ${benchmark.list} -rff ${benchmark.resultfile} -rf ${benchmark.resultformat} run-java-benchmarks exec integration-test apache-arrow-java-19.0.0/performance/src/000077500000000000000000000000001515451615200201665ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/000077500000000000000000000000001515451615200211125ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/000077500000000000000000000000001515451615200220335ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/000077500000000000000000000000001515451615200226225ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/000077500000000000000000000000001515451615200240435ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200251755ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/adapter/000077500000000000000000000000001515451615200266155ustar00rootroot00000000000000AvroAdapterBenchmarks.java000066400000000000000000000117471515451615200336210ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/adapter/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.util.concurrent.TimeUnit; import org.apache.arrow.adapter.avro.AvroToArrow; import org.apache.arrow.adapter.avro.AvroToArrowConfig; import org.apache.arrow.adapter.avro.AvroToArrowConfigBuilder; import org.apache.arrow.adapter.avro.AvroToArrowVectorIterator; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.avro.Schema; import org.apache.avro.generic.GenericData; 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.DatumWriter; import org.apache.avro.io.Decoder; import org.apache.avro.io.DecoderFactory; import org.apache.avro.io.EncoderFactory; 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.runner.Runner; import org.openjdk.jmh.runner.RunnerException; import org.openjdk.jmh.runner.options.Options; import org.openjdk.jmh.runner.options.OptionsBuilder; /** Benchmarks for avro adapter. */ @State(Scope.Benchmark) public class AvroAdapterBenchmarks { private final int valueCount = 3000; private AvroToArrowConfig config; private Schema schema; private BinaryDecoder decoder; /** Setup benchmarks. */ @Setup public void prepare() throws Exception { BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); config = new AvroToArrowConfigBuilder(allocator).build(); String schemaStr = "{\n" + " \"namespace\": \"org.apache.arrow.avro\",\n" + " \"type\": \"record\",\n" + " \"name\": \"testBenchmark\",\n" + " \"fields\": [\n" + " {\"name\": \"f0\", \"type\": \"string\"},\n" + " {\"name\": \"f1\", \"type\": \"int\"},\n" + " {\"name\": \"f2\", \"type\": \"long\"},\n" + " {\"name\": \"f3\", \"type\": \"boolean\"},\n" + " {\"name\": \"f4\", \"type\": \"float\"}\n" + " ]\n" + "}"; schema = new Schema.Parser().parse(schemaStr); ByteArrayOutputStream out = new ByteArrayOutputStream(); BinaryEncoder encoder = new EncoderFactory().directBinaryEncoder(out, null); DatumWriter writer = new GenericDatumWriter(schema); for (int i = 0; i < valueCount; i++) { GenericRecord record = new GenericData.Record(schema); record.put(0, "test" + i); record.put(1, i); record.put(2, i + 1L); record.put(3, i % 2 == 0); record.put(4, i + 0.1f); writer.write(record, encoder); } decoder = new DecoderFactory().directBinaryDecoder(new ByteArrayInputStream(out.toByteArray()), null); } /** Tear down benchmarks. */ @TearDown public void tearDown() { config.getAllocator().close(); } /** * Test {@link AvroToArrow#avroToArrowIterator(Schema, Decoder, AvroToArrowConfig)}. * * @return useless. To avoid DCE by JIT. */ @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.NANOSECONDS) public int testAvroToArrow() throws Exception { decoder.inputStream().reset(); int sum = 0; try (AvroToArrowVectorIterator iter = AvroToArrow.avroToArrowIterator(schema, decoder, config)) { while (iter.hasNext()) { VectorSchemaRoot root = iter.next(); IntVector intVector = (IntVector) root.getVector("f1"); for (int i = 0; i < intVector.getValueCount(); i++) { sum += intVector.get(i); } root.close(); } } return sum; } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder().include(AvroAdapterBenchmarks.class.getSimpleName()).forks(1).build(); new Runner(opt).run(); } } apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/adapter/jdbc/000077500000000000000000000000001515451615200275175ustar00rootroot00000000000000JdbcAdapterBenchmarks.java000066400000000000000000000256621515451615200344570ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/adapter/jdbc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.jdbc; import java.sql.Connection; import java.sql.DriverManager; import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.Statement; import java.util.concurrent.TimeUnit; import org.apache.arrow.adapter.jdbc.consumer.BigIntConsumer; import org.apache.arrow.adapter.jdbc.consumer.BitConsumer; import org.apache.arrow.adapter.jdbc.consumer.IntConsumer; import org.apache.arrow.adapter.jdbc.consumer.JdbcConsumer; import org.apache.arrow.adapter.jdbc.consumer.VarCharConsumer; 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.IntVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; 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.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 Jdbc adapter. */ public class JdbcAdapterBenchmarks { // checkstyle:off: MissingJavadocMethod private static final int VALUE_COUNT = 3000; private static final String CREATE_STATEMENT = "CREATE TABLE test_table (f0 INT, f1 LONG, f2 VARCHAR, f3 BOOLEAN);"; private static final String INSERT_STATEMENT = "INSERT INTO test_table (f0, f1, f2, f3) VALUES (?, ?, ?, ?);"; private static final String QUERY = "SELECT f0, f1, f2, f3 FROM test_table;"; private static final String DROP_STATEMENT = "DROP TABLE test_table;"; private static final String URL = "jdbc:h2:mem:JdbcAdapterBenchmarks"; private static final String DRIVER = "org.h2.Driver"; /** State object for the jdbc e2e benchmark. */ @State(Scope.Benchmark) public static class JdbcState { private Connection conn = null; private ResultSet resultSet = null; private BufferAllocator allocator; private Statement statement; private JdbcToArrowConfig config; @Setup(Level.Trial) public void prepareState() throws Exception { allocator = new RootAllocator(Integer.MAX_VALUE); config = new JdbcToArrowConfigBuilder().setAllocator(allocator).setTargetBatchSize(1024).build(); Class.forName(DRIVER); conn = DriverManager.getConnection(URL); try (Statement stmt = conn.createStatement()) { stmt.executeUpdate(CREATE_STATEMENT); } for (int i = 0; i < VALUE_COUNT; i++) { // Insert data try (PreparedStatement stmt = conn.prepareStatement(INSERT_STATEMENT)) { stmt.setInt(1, i); stmt.setLong(2, i); stmt.setString(3, "test" + i); stmt.setBoolean(4, i % 2 == 0); stmt.executeUpdate(); } } } @Setup(Level.Invocation) public void prepareInvoke() throws Exception { statement = conn.createStatement(); resultSet = statement.executeQuery(QUERY); } @TearDown(Level.Invocation) public void tearDownInvoke() throws Exception { resultSet.close(); statement.close(); } @TearDown(Level.Trial) public void tearDownState() throws Exception { try (Statement stmt = conn.createStatement()) { stmt.executeUpdate(DROP_STATEMENT); } allocator.close(); } } /** State object for the consume benchmark. */ @State(Scope.Benchmark) public static class ConsumeState { private static final boolean NULLABLE = true; private Connection conn = null; private ResultSet resultSet = null; private BufferAllocator allocator; private Statement statement; private IntVector intVector; private BigIntVector longVector; private VarCharVector varCharVector; private BitVector bitVector; private JdbcConsumer intConsumer; private JdbcConsumer longConsumer; private JdbcConsumer varCharConsumer; private JdbcConsumer bitConsumer; private JdbcToArrowConfig config; @Setup(Level.Trial) public void prepare() throws Exception { allocator = new RootAllocator(Integer.MAX_VALUE); config = new JdbcToArrowConfigBuilder().setAllocator(allocator).setTargetBatchSize(1024).build(); Class.forName(DRIVER); conn = DriverManager.getConnection(URL); try (Statement stmt = conn.createStatement()) { stmt.executeUpdate(CREATE_STATEMENT); } for (int i = 0; i < VALUE_COUNT; i++) { // Insert data try (PreparedStatement stmt = conn.prepareStatement(INSERT_STATEMENT)) { stmt.setInt(1, i); stmt.setLong(2, i); stmt.setString(3, "test" + i); stmt.setBoolean(4, i % 2 == 0); stmt.executeUpdate(); } } statement = conn.createStatement(); resultSet = statement.executeQuery(QUERY); resultSet.next(); intVector = new IntVector("", allocator); intVector.allocateNew(VALUE_COUNT); intConsumer = IntConsumer.createConsumer(intVector, 1, NULLABLE); longVector = new BigIntVector("", allocator); longVector.allocateNew(VALUE_COUNT); longConsumer = BigIntConsumer.createConsumer(longVector, 2, NULLABLE); varCharVector = new VarCharVector("", allocator); varCharVector.allocateNew(VALUE_COUNT); varCharConsumer = VarCharConsumer.createConsumer(varCharVector, 3, NULLABLE); bitVector = new BitVector("", allocator); bitVector.allocateNew(VALUE_COUNT); bitConsumer = BitConsumer.createConsumer(bitVector, 4, NULLABLE); } @TearDown(Level.Trial) public void tearDown() throws Exception { try (Statement stmt = conn.createStatement()) { stmt.executeUpdate(DROP_STATEMENT); } resultSet.close(); statement.close(); conn.close(); intVector.close(); intConsumer.close(); longVector.close(); longConsumer.close(); varCharVector.close(); varCharConsumer.close(); bitVector.close(); bitConsumer.close(); allocator.close(); } } /** State object for the jdbc row consume benchmark. */ @State(Scope.Benchmark) public static class RowConsumeState { private Connection conn = null; private ResultSet resultSet = null; private BufferAllocator allocator; private Statement statement; private JdbcToArrowConfig config; private ArrowVectorIterator iter; private VectorSchemaRoot root; @Setup(Level.Trial) public void prepareState() throws Exception { allocator = new RootAllocator(Integer.MAX_VALUE); config = new JdbcToArrowConfigBuilder() .setAllocator(allocator) .setTargetBatchSize(VALUE_COUNT) .build(); Class.forName(DRIVER); conn = DriverManager.getConnection(URL); try (Statement stmt = conn.createStatement()) { stmt.executeUpdate(CREATE_STATEMENT); } for (int i = 0; i < VALUE_COUNT; i++) { // Insert data try (PreparedStatement stmt = conn.prepareStatement(INSERT_STATEMENT)) { stmt.setInt(1, i); stmt.setLong(2, i); stmt.setString(3, "test" + i); stmt.setBoolean(4, i % 2 == 0); stmt.executeUpdate(); } } } @Setup(Level.Invocation) public void prepareInvoke() throws Exception { statement = conn.createStatement(); resultSet = statement.executeQuery(QUERY); iter = JdbcToArrow.sqlToArrowVectorIterator(resultSet, config); root = iter.next(); iter.compositeConsumer.resetVectorSchemaRoot(root); } @TearDown(Level.Invocation) public void tearDownInvoke() throws Exception { resultSet.close(); statement.close(); iter.close(); } @TearDown(Level.Trial) public void tearDownState() throws Exception { try (Statement stmt = conn.createStatement()) { stmt.executeUpdate(DROP_STATEMENT); } allocator.close(); } } /** * Test {@link JdbcToArrow#sqlToArrowVectorIterator(ResultSet, JdbcToArrowConfig)}. * * @return useless. To avoid DCE by JIT. */ @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public int testJdbcToArrow(JdbcState state) throws Exception { int valueCount = 0; try (ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(state.resultSet, state.config)) { while (iter.hasNext()) { VectorSchemaRoot root = iter.next(); IntVector intVector = (IntVector) root.getFieldVectors().get(0); valueCount += intVector.getValueCount(); root.close(); } } return valueCount; } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public void consumeBenchmark(ConsumeState state) throws Exception { state.intConsumer.resetValueVector(state.intVector); state.longConsumer.resetValueVector(state.longVector); state.varCharConsumer.resetValueVector(state.varCharVector); state.bitConsumer.resetValueVector(state.bitVector); for (int i = 0; i < VALUE_COUNT; i++) { state.intConsumer.consume(state.resultSet); state.longConsumer.consume(state.resultSet); state.varCharConsumer.consume(state.resultSet); state.bitConsumer.consume(state.resultSet); } } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public void consumeRowsBenchmark(RowConsumeState state) throws Exception { for (int i = 0; i < VALUE_COUNT; i++) { state.iter.compositeConsumer.consume(state.resultSet); } } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder().include(JdbcAdapterBenchmarks.class.getSimpleName()).forks(1).build(); new Runner(opt).run(); } // checkstyle:on: MissingJavadocMethod } apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/algorithm/000077500000000000000000000000001515451615200271635ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/algorithm/search/000077500000000000000000000000001515451615200304305ustar00rootroot00000000000000ParallelSearcherBenchmarks.java000066400000000000000000000071621515451615200364310ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/algorithm/search/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.algorithm.search; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; 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 ParallelSearcher}. */ public class ParallelSearcherBenchmarks { // checkstyle:off: MissingJavadocMethod private static final int VECTOR_LENGTH = 1024 * 1024; /** State object for the benchmarks. */ @State(Scope.Benchmark) public static class SearchState { @Param({"1", "2", "5", "10", "20", "50", "100"}) int numThreads; BufferAllocator allocator; ExecutorService threadPool; IntVector targetVector; IntVector keyVector; ParallelSearcher searcher; @Setup(Level.Trial) public void prepare() { allocator = new RootAllocator(Integer.MAX_VALUE); targetVector = new IntVector("target vector", allocator); targetVector.allocateNew(VECTOR_LENGTH); keyVector = new IntVector("key vector", allocator); keyVector.allocateNew(1); threadPool = Executors.newFixedThreadPool(numThreads); for (int i = 0; i < VECTOR_LENGTH; i++) { targetVector.set(i, i); } targetVector.setValueCount(VECTOR_LENGTH); keyVector.set(0, VECTOR_LENGTH / 3); keyVector.setValueCount(1); } @Setup(Level.Invocation) public void prepareInvoke() { searcher = new ParallelSearcher<>(targetVector, threadPool, numThreads); } @TearDown(Level.Trial) public void tearDownState() { targetVector.close(); keyVector.close(); allocator.close(); threadPool.shutdown(); } } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public void searchBenchmark(SearchState state) throws Exception { state.searcher.search(state.keyVector, 0); } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder() .include(ParallelSearcherBenchmarks.class.getSimpleName()) .forks(1) .build(); new Runner(opt).run(); } // checkstyle:on: MissingJavadocMethod } apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/memory/000077500000000000000000000000001515451615200265055ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/memory/AllocatorBenchmarks.java000066400000000000000000000056421515451615200332750ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.concurrent.TimeUnit; import org.apache.arrow.memory.rounding.RoundingPolicy; import org.apache.arrow.memory.rounding.SegmentRoundingPolicy; 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.runner.Runner; import org.openjdk.jmh.runner.RunnerException; import org.openjdk.jmh.runner.options.Options; import org.openjdk.jmh.runner.options.OptionsBuilder; /** Benchmarks for allocators. */ public class AllocatorBenchmarks { /** Benchmark for the default allocator. */ @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public void defaultAllocatorBenchmark() { final int bufferSize = 1024; final int numBuffers = 1024; try (RootAllocator allocator = new RootAllocator(numBuffers * bufferSize)) { ArrowBuf[] buffers = new ArrowBuf[numBuffers]; for (int i = 0; i < numBuffers; i++) { buffers[i] = allocator.buffer(bufferSize); } for (int i = 0; i < numBuffers; i++) { buffers[i].close(); } } } /** Benchmark for allocator with segment rounding policy. */ @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public void segmentRoundingPolicyBenchmark() { final long bufferSize = 1024L; final int numBuffers = 1024; final long segmentSize = 1024L; RoundingPolicy policy = new SegmentRoundingPolicy(segmentSize); try (RootAllocator allocator = new RootAllocator(AllocationListener.NOOP, bufferSize * numBuffers, policy)) { ArrowBuf[] buffers = new ArrowBuf[numBuffers]; for (int i = 0; i < numBuffers; i++) { buffers[i] = allocator.buffer(bufferSize); } for (int i = 0; i < numBuffers; i++) { buffers[i].close(); } } } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder().include(AllocatorBenchmarks.class.getSimpleName()).forks(1).build(); new Runner(opt).run(); } } apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/memory/ArrowBufBenchmarks.java000066400000000000000000000044771515451615200331110ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.concurrent.TimeUnit; 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.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 ArrowBuf}. */ @State(Scope.Benchmark) public class ArrowBufBenchmarks { private static final int BUFFER_CAPACITY = 1024 * 1024; private static final int ALLOCATOR_CAPACITY = 1024 * 1024; private BufferAllocator allocator; private ArrowBuf buffer; /** Setup benchmarks. */ @Setup public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); buffer = allocator.buffer(BUFFER_CAPACITY); } /** Tear down benchmarks. */ @TearDown public void tearDown() { buffer.close(); allocator.close(); } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.NANOSECONDS) public void setZero() { buffer.setZero(0, BUFFER_CAPACITY); } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder().include(ArrowBufBenchmarks.class.getSimpleName()).forks(1).build(); new Runner(opt).run(); } } MemoryFootprintBenchmarks.java000066400000000000000000000167651515451615200344630ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/memory/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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(); } } apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/memory/util/000077500000000000000000000000001515451615200274625ustar00rootroot00000000000000ArrowBufPointerBenchmarks.java000066400000000000000000000060601515451615200353360ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/memory/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util; 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.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.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 ArrowBufPointer}. */ @State(Scope.Benchmark) public class ArrowBufPointerBenchmarks { private static final int BUFFER_CAPACITY = 1000; private static final int ALLOCATOR_CAPACITY = 1024 * 1024; private BufferAllocator allocator; private ArrowBuf buffer1; private ArrowBuf buffer2; private ArrowBufPointer pointer1; private ArrowBufPointer pointer2; /** Setup benchmarks. */ @Setup public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); buffer1 = allocator.buffer(BUFFER_CAPACITY); buffer2 = allocator.buffer(BUFFER_CAPACITY); for (int i = 0; i < BUFFER_CAPACITY; i++) { buffer1.setByte(i, i); buffer2.setByte(i, i); } // make the last bytes different buffer1.setByte(BUFFER_CAPACITY - 1, 12); buffer1.setByte(BUFFER_CAPACITY - 1, 123); pointer1 = new ArrowBufPointer(buffer1, 0, BUFFER_CAPACITY); pointer2 = new ArrowBufPointer(buffer2, 0, BUFFER_CAPACITY); } /** Tear down benchmarks. */ @TearDown public void tearDown() { buffer1.close(); buffer2.close(); allocator.close(); } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.NANOSECONDS) public int compareBenchmark() { return pointer1.compareTo(pointer2); } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder() .include(ArrowBufPointerBenchmarks.class.getSimpleName()) .forks(1) .build(); new Runner(opt).run(); } } ByteFunctionHelpersBenchmarks.java000066400000000000000000000105161515451615200362030ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/memory/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util; 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.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.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 ByteFunctionHelpers}. */ public class ByteFunctionHelpersBenchmarks { private static final int ALLOCATOR_CAPACITY = 1024 * 1024; /** * State object for the {@link ByteFunctionHelpersBenchmarks#arrowBufEquals(ArrowEqualState)} * benchmark. */ @State(Scope.Benchmark) public static class ArrowEqualState { // checkstyle:off: MissingJavadocMethod private static final int BUFFER_CAPACITY = 7; private BufferAllocator allocator; private ArrowBuf buffer1; private ArrowBuf buffer2; @Setup(Level.Trial) public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); buffer1 = allocator.buffer(BUFFER_CAPACITY); buffer2 = allocator.buffer(BUFFER_CAPACITY); for (int i = 0; i < BUFFER_CAPACITY; i++) { buffer1.setByte(i, i); buffer2.setByte(i, i); } } @TearDown(Level.Trial) public void tearDown() { buffer1.close(); buffer2.close(); allocator.close(); } } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.NANOSECONDS) public void arrowBufEquals(ArrowEqualState state) { ByteFunctionHelpers.equal( state.buffer1, 0, ArrowEqualState.BUFFER_CAPACITY - 1, state.buffer2, 0, ArrowEqualState.BUFFER_CAPACITY - 1); } /** * State object for the {@link * ByteFunctionHelpersBenchmarks#arrowBufArrayEquals(ArrowArrayEqualState)} benchmark. */ @State(Scope.Benchmark) public static class ArrowArrayEqualState { private static final int BUFFER_CAPACITY = 1024; private BufferAllocator allocator; private ArrowBuf buffer1; private byte[] buffer2; @Setup(Level.Trial) public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); buffer1 = allocator.buffer(BUFFER_CAPACITY); buffer2 = new byte[BUFFER_CAPACITY]; for (int i = 0; i < BUFFER_CAPACITY; i++) { buffer1.setByte(i, i); buffer2[i] = (byte) i; } } @TearDown(Level.Trial) public void tearDown() { buffer1.close(); allocator.close(); } } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.NANOSECONDS) public int arrowBufArrayEquals(ArrowArrayEqualState state) { return ByteFunctionHelpers.compare( state.buffer1, 0, ArrowArrayEqualState.BUFFER_CAPACITY, state.buffer2, 0, ArrowArrayEqualState.BUFFER_CAPACITY); } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder() .include(ByteFunctionHelpersBenchmarks.class.getSimpleName()) .forks(1) .build(); new Runner(opt).run(); } // checkstyle:on: MissingJavadocMethod } apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/000077500000000000000000000000001515451615200264775ustar00rootroot00000000000000BaseValueVectorBenchmarks.java000066400000000000000000000054051515451615200343170ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.BufferAllocator; import org.apache.arrow.memory.RootAllocator; 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.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 BaseValueVector}. */ @State(Scope.Benchmark) public class BaseValueVectorBenchmarks { private static final int VECTOR_LENGTH = 1024; private static final int ALLOCATOR_CAPACITY = 1024 * 1024; private BufferAllocator allocator; private IntVector vector; /** Setup benchmarks. */ @Setup public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); vector = new IntVector("vector", allocator); vector.allocateNew(VECTOR_LENGTH); } /** Tear down benchmarks. */ @TearDown public void tearDown() { vector.close(); allocator.close(); } /** * Test {@link BaseValueVector#computeCombinedBufferSize(int, int)}. * * @return useless. To avoid DCE by JIT. */ @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.NANOSECONDS) public int testComputeCombinedBufferSize() { int totalSize = 0; for (int i = 0; i < VECTOR_LENGTH; i++) { totalSize += vector.computeCombinedBufferSize(i, 4); } return totalSize; } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder() .include(BaseValueVectorBenchmarks.class.getSimpleName()) .forks(1) .build(); new Runner(opt).run(); } } BitVectorHelperBenchmarks.java000066400000000000000000000153161515451615200343300ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message.ArrowFieldNode; 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.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 BitVectorHelper}. */ public class BitVectorHelperBenchmarks { // checkstyle:off: MissingJavadocMethod /** State object for general benchmarks. */ @State(Scope.Benchmark) public static class BenchmarkState { private static final int VALIDITY_BUFFER_CAPACITY = 1024; private static final int ALLOCATOR_CAPACITY = 1024 * 1024; private BufferAllocator allocator; private ArrowBuf validityBuffer; private ArrowBuf oneBitValidityBuffer; /** Setup benchmarks. */ @Setup(Level.Trial) public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); validityBuffer = allocator.buffer(VALIDITY_BUFFER_CAPACITY / 8); for (int i = 0; i < VALIDITY_BUFFER_CAPACITY; i++) { if (i % 7 == 0) { BitVectorHelper.setBit(validityBuffer, i); } else { BitVectorHelper.unsetBit(validityBuffer, i); } } // only one 1 bit in the middle of the buffer oneBitValidityBuffer = allocator.buffer(VALIDITY_BUFFER_CAPACITY / 8); oneBitValidityBuffer.setZero(0, VALIDITY_BUFFER_CAPACITY / 8); BitVectorHelper.setBit(oneBitValidityBuffer, VALIDITY_BUFFER_CAPACITY / 2); } /** Tear down benchmarks. */ @TearDown(Level.Trial) public void tearDown() { validityBuffer.close(); oneBitValidityBuffer.close(); allocator.close(); } } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.NANOSECONDS) public int getNullCountBenchmark(BenchmarkState state) { return BitVectorHelper.getNullCount( state.validityBuffer, BenchmarkState.VALIDITY_BUFFER_CAPACITY); } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.NANOSECONDS) public boolean allBitsNullBenchmark(BenchmarkState state) { return BitVectorHelper.checkAllBitsEqualTo( state.oneBitValidityBuffer, BenchmarkState.VALIDITY_BUFFER_CAPACITY, true); } /** State object for {@link #loadValidityBufferAllOne(NonNullableValidityBufferState)}.. */ @State(Scope.Benchmark) public static class NonNullableValidityBufferState { private static final int VALIDITY_BUFFER_CAPACITY = 1024; private static final int ALLOCATOR_CAPACITY = 1024 * 1024; private BufferAllocator allocator; private ArrowBuf validityBuffer; private ArrowBuf loadResult; private ArrowFieldNode fieldNode; /** Setup benchmarks. */ @Setup(Level.Trial) public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); validityBuffer = allocator.buffer(VALIDITY_BUFFER_CAPACITY / 8); for (int i = 0; i < VALIDITY_BUFFER_CAPACITY; i++) { BitVectorHelper.setBit(validityBuffer, i); } fieldNode = new ArrowFieldNode(VALIDITY_BUFFER_CAPACITY, 0); } @TearDown(Level.Invocation) public void tearDownInvoke() { loadResult.close(); } /** Tear down benchmarks. */ @TearDown(Level.Trial) public void tearDown() { validityBuffer.close(); allocator.close(); } } /** * Benchmark for {@link BitVectorHelper#loadValidityBuffer(ArrowFieldNode, ArrowBuf, * BufferAllocator)} when all elements are not null. */ @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.NANOSECONDS) public void loadValidityBufferAllOne(NonNullableValidityBufferState state) { state.loadResult = BitVectorHelper.loadValidityBuffer(state.fieldNode, state.validityBuffer, state.allocator); } /** State object for {@link #setValidityBitBenchmark(ClearBitStateState)}. */ @State(Scope.Benchmark) public static class ClearBitStateState { private static final int VALIDITY_BUFFER_CAPACITY = 1024; private static final int ALLOCATOR_CAPACITY = 1024 * 1024; private BufferAllocator allocator; private ArrowBuf validityBuffer; private int bitToSet = 0; /** Setup benchmarks. */ @Setup(Level.Trial) public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); validityBuffer = allocator.buffer(VALIDITY_BUFFER_CAPACITY / 8); } /** Tear down benchmarks. */ @TearDown(Level.Trial) public void tearDown() { validityBuffer.close(); allocator.close(); } } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public void setValidityBitBenchmark(ClearBitStateState state) { for (int i = 0; i < ClearBitStateState.VALIDITY_BUFFER_CAPACITY; i++) { BitVectorHelper.setValidityBit(state.validityBuffer, i, state.bitToSet); } } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public void setValidityBitToZeroBenchmark(ClearBitStateState state) { for (int i = 0; i < ClearBitStateState.VALIDITY_BUFFER_CAPACITY; i++) { BitVectorHelper.unsetBit(state.validityBuffer, i); } } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder() .include(BitVectorHelperBenchmarks.class.getSimpleName()) .forks(1) .build(); new Runner(opt).run(); } // checkstyle:on: MissingJavadocMethod } DecimalVectorBenchmarks.java000066400000000000000000000070021515451615200340010ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.math.BigDecimal; 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.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.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 DecimalVector}. */ @State(Scope.Benchmark) public class DecimalVectorBenchmarks { private static final int VECTOR_LENGTH = 1024; private static final int ALLOCATOR_CAPACITY = 1024 * 1024; private BufferAllocator allocator; private DecimalVector vector; private ArrowBuf fromBuf; byte[] fromByteArray; /** Setup benchmarks. */ @Setup public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); vector = new DecimalVector("vector", allocator, 38, 16); vector.allocateNew(VECTOR_LENGTH); fromBuf = allocator.buffer(VECTOR_LENGTH * DecimalVector.TYPE_WIDTH); for (int i = 0; i < VECTOR_LENGTH; i++) { byte[] bytes = BigDecimal.valueOf(i).unscaledValue().toByteArray(); fromBuf.setBytes(i * DecimalVector.TYPE_WIDTH, bytes); } fromByteArray = new byte[DecimalVector.TYPE_WIDTH]; fromBuf.getBytes(0, fromByteArray); } /** Tear down benchmarks. */ @TearDown public void tearDown() { fromBuf.close(); vector.close(); allocator.close(); } /** Test writing on {@link DecimalVector} from arrow buf. */ @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public void setBigEndianArrowBufBenchmark() { int offset = 0; for (int i = 0; i < VECTOR_LENGTH; i++) { vector.setBigEndianSafe(i, offset, fromBuf, DecimalVector.TYPE_WIDTH); offset += 8; } } /** Test writing on {@link DecimalVector} from byte array. */ @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public void setBigEndianByteArrayBenchmark() { for (int i = 0; i < VECTOR_LENGTH; i++) { vector.setBigEndian(i, fromByteArray); } } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder() .include(DecimalVectorBenchmarks.class.getSimpleName()) .forks(1) .build(); new Runner(opt).run(); } } apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/Float8Benchmarks.java000066400000000000000000000072351515451615200325040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.BoundsChecking; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; 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.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 Float8Vector}. */ @State(Scope.Benchmark) public class Float8Benchmarks { // checkstyle:off: MissingJavadocMethod private static final int VECTOR_LENGTH = 1024; private static final int ALLOCATOR_CAPACITY = 1024 * 1024; private BufferAllocator allocator; private Float8Vector vector; private Float8Vector fromVector; /** Setup benchmarks. */ @Setup public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); vector = new Float8Vector("vector", allocator); vector.allocateNew(VECTOR_LENGTH); fromVector = new Float8Vector("vector", allocator); fromVector.allocateNew(VECTOR_LENGTH); for (int i = 0; i < VECTOR_LENGTH; i++) { if (i % 3 == 0) { fromVector.setNull(i); } else { fromVector.set(i, i * i); } } fromVector.setValueCount(VECTOR_LENGTH); } /** Tear down benchmarks. */ @TearDown public void tearDown() { vector.close(); fromVector.close(); allocator.close(); } /** * Test reading/writing on {@link Float8Vector}. The performance of this benchmark is influenced * by the states of two flags: 1. The flag for boundary checking. For details, please see {@link * BoundsChecking}. 2. The flag for null checking in get methods. For details, please see {@link * NullCheckingForGet}. * * @return useless. To avoid DCE by JIT. */ @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public double readWriteBenchmark() { double sum = 0; for (int i = 0; i < VECTOR_LENGTH; i++) { vector.set(i, i + 10.0); sum += vector.get(i); } return sum; } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public void copyFromBenchmark() { for (int i = 0; i < VECTOR_LENGTH; i++) { vector.copyFrom(i, i, (Float8Vector) fromVector); } } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder().include(Float8Benchmarks.class.getSimpleName()).forks(1).build(); new Runner(opt).run(); } // checkstyle:on: MissingJavadocMethod } FloatingPointBenchmarks.java000066400000000000000000000103761515451615200340450ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.compare.ApproxEqualsVisitor; import org.apache.arrow.vector.compare.Range; 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.runner.Runner; import org.openjdk.jmh.runner.RunnerException; import org.openjdk.jmh.runner.options.Options; import org.openjdk.jmh.runner.options.OptionsBuilder; /** Benchmarks for floating point vectors. */ @State(Scope.Benchmark) public class FloatingPointBenchmarks { // checkstyle:off: MissingJavadocMethod private static final int VECTOR_LENGTH = 1024; private static final int ALLOCATOR_CAPACITY = 1024 * 1024; private BufferAllocator allocator; private Float4Vector floatVector1; private Float4Vector floatVector2; private Float8Vector doubleVector1; private Float8Vector doubleVector2; private ApproxEqualsVisitor floatVisitor; private ApproxEqualsVisitor doubleVisitor; private Range range; /** Setup benchmarks. */ @Setup public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); floatVector1 = new Float4Vector("vector", allocator); floatVector2 = new Float4Vector("vector", allocator); doubleVector1 = new Float8Vector("vector", allocator); doubleVector2 = new Float8Vector("vector", allocator); floatVector1.allocateNew(VECTOR_LENGTH); floatVector2.allocateNew(VECTOR_LENGTH); doubleVector1.allocateNew(VECTOR_LENGTH); doubleVector2.allocateNew(VECTOR_LENGTH); for (int i = 0; i < VECTOR_LENGTH; i++) { if (i % 3 == 0) { floatVector1.setNull(i); floatVector2.setNull(i); doubleVector1.setNull(i); doubleVector2.setNull(i); } else { floatVector1.set(i, i * i); floatVector2.set(i, i * i); doubleVector1.set(i, i * i); doubleVector2.set(i, i * i); } } floatVector1.setValueCount(VECTOR_LENGTH); floatVector2.setValueCount(VECTOR_LENGTH); doubleVector1.setValueCount(VECTOR_LENGTH); doubleVector2.setValueCount(VECTOR_LENGTH); floatVisitor = new ApproxEqualsVisitor(floatVector1, floatVector2, 0.01f, 0.01); doubleVisitor = new ApproxEqualsVisitor(doubleVector1, doubleVector2, 0.01f, 0.01); range = new Range(0, 0, VECTOR_LENGTH); } /** Tear down benchmarks. */ @TearDown public void tearDown() { floatVector1.close(); floatVector2.close(); doubleVector1.close(); doubleVector2.close(); allocator.close(); } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public int approxEqualsBenchmark() { boolean floatResult = floatVisitor.visit(floatVector1, range); boolean doubleResult = doubleVisitor.visit(doubleVector1, range); return (floatResult ? 1 : 0) + (doubleResult ? 1 : 0); } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder() .include(FloatingPointBenchmarks.class.getSimpleName()) .forks(1) .build(); new Runner(opt).run(); } // checkstyle:on: MissingJavadocMethod } apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/IntBenchmarks.java000066400000000000000000000065171515451615200321030ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.complex.impl.IntWriterImpl; import org.apache.arrow.vector.holders.NullableIntHolder; 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.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 IntVector}. */ @State(Scope.Benchmark) public class IntBenchmarks { // checkstyle:off: MissingJavadocMethod private static final int VECTOR_LENGTH = 1024; private static final int ALLOCATOR_CAPACITY = 1024 * 1024; private BufferAllocator allocator; private IntVector vector; @Setup public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); vector = new IntVector("vector", allocator); vector.allocateNew(VECTOR_LENGTH); vector.setValueCount(VECTOR_LENGTH); } @TearDown public void tearDown() { vector.close(); allocator.close(); } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public void setWithValueHolder() { for (int i = 0; i < VECTOR_LENGTH; i++) { NullableIntHolder holder = new NullableIntHolder(); holder.isSet = i % 3 == 0 ? 0 : 1; if (holder.isSet == 1) { holder.value = i; } vector.setSafe(i, holder); } } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public void setIntDirectly() { for (int i = 0; i < VECTOR_LENGTH; i++) { vector.setSafe(i, i % 3 == 0 ? 0 : 1, i); } } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public void setWithWriter() { IntWriterImpl writer = new IntWriterImpl(vector); for (int i = 0; i < VECTOR_LENGTH; i++) { if (i % 3 != 0) { writer.writeInt(i); } } } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder().include(IntBenchmarks.class.getSimpleName()).forks(1).build(); new Runner(opt).run(); } // checkstyle:on: MissingJavadocMethod } apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/UuidVectorBenchmarks.java000066400000000000000000000101221515451615200334250ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 } apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/VarCharBenchmarks.java000066400000000000000000000060351515451615200326720ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.BufferAllocator; import org.apache.arrow.memory.RootAllocator; 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.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 VarCharVector}. */ @State(Scope.Benchmark) public class VarCharBenchmarks { // checkstyle:off: MissingJavadocMethod private static final int VECTOR_LENGTH = 1024; private static final int ALLOCATOR_CAPACITY = 1024 * 1024; private BufferAllocator allocator; private VarCharVector vector; private VarCharVector fromVector; /** Setup benchmarks. */ @Setup public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); vector = new VarCharVector("vector", allocator); vector.allocateNew(ALLOCATOR_CAPACITY / 4, VECTOR_LENGTH); fromVector = new VarCharVector("vector", allocator); fromVector.allocateNew(ALLOCATOR_CAPACITY / 4, VECTOR_LENGTH); for (int i = 0; i < VECTOR_LENGTH; i++) { if (i % 3 == 0) { fromVector.setNull(i); } else { fromVector.set(i, String.valueOf(i * 1000).getBytes()); } } fromVector.setValueCount(VECTOR_LENGTH); } /** Tear down benchmarks. */ @TearDown public void tearDown() { vector.close(); fromVector.close(); allocator.close(); } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public void copyFromBenchmark() { for (int i = 0; i < VECTOR_LENGTH; i++) { vector.copyFrom(i, i, fromVector); } } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder().include(VarCharBenchmarks.class.getSimpleName()).forks(1).build(); new Runner(opt).run(); } // checkstyle:on: MissingJavadocMethod } VariableWidthVectorBenchmarkConstants.java000066400000000000000000000020151515451615200367010ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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(); } VariableWidthVectorBenchmarks.java000066400000000000000000000101271515451615200351720ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 VariableWidthVectorBenchmarks { // 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 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(VariableWidthVectorBenchmarks.class.getSimpleName()) .forks(1) .build(); new Runner(opt).run(); } // checkstyle:on: MissingJavadocMethod } VariableWidthVectorInlineValueBenchmarks.java000066400000000000000000000101641515451615200373270ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 } VariableWidthViewVectorBenchmarks.java000066400000000000000000000101411515451615200360210ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 } VariableWidthViewVectorInlineValueBenchmarks.java000066400000000000000000000102231515451615200401560ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 } VectorLoaderBenchmark.java000066400000000000000000000065611515451615200334770ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; 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.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 VectorLoader}. */ public class VectorLoaderBenchmark { // checkstyle:off: MissingJavadocMethod private static final int ALLOCATOR_CAPACITY = 1024 * 1024; private static final int VECTOR_COUNT = 10; /** State for vector load benchmark. */ @State(Scope.Benchmark) public static class LoadState { private BufferAllocator allocator; private VarCharVector[] vectors; private ArrowRecordBatch recordBatch; private VectorSchemaRoot root; private VectorLoader loader; /** Setup benchmarks. */ @Setup(Level.Trial) public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); } @Setup(Level.Invocation) public void prepareInvoke() { vectors = new VarCharVector[VECTOR_COUNT]; for (int i = 0; i < VECTOR_COUNT; i++) { vectors[i] = new VarCharVector("vector", allocator); vectors[i].allocateNew(100, 10); } root = VectorSchemaRoot.of(vectors); VectorUnloader unloader = new VectorUnloader(root); recordBatch = unloader.getRecordBatch(); loader = new VectorLoader(root); } @TearDown(Level.Invocation) public void tearDownInvoke() { recordBatch.close(); root.close(); } /** Tear down benchmarks. */ @TearDown(Level.Trial) public void tearDown() { allocator.close(); } } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public void loadBenchmark(LoadState state) { state.loader.load(state.recordBatch); } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder().include(VectorLoaderBenchmark.class.getSimpleName()).forks(1).build(); new Runner(opt).run(); } // checkstyle:on: MissingJavadocMethod } VectorUnloaderBenchmark.java000066400000000000000000000063221515451615200340350ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; 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.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 VectorUnloader}. */ @State(Scope.Benchmark) public class VectorUnloaderBenchmark { // checkstyle:off: MissingJavadocMethod private static final int ALLOCATOR_CAPACITY = 1024 * 1024; private static final int VECTOR_COUNT = 10; private BufferAllocator allocator; private VarCharVector[] vectors; private VectorUnloader unloader; private ArrowRecordBatch recordBatch; /** Setup benchmarks. */ @Setup(Level.Trial) public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); } @Setup(Level.Invocation) public void prepareInvoke() { vectors = new VarCharVector[VECTOR_COUNT]; for (int i = 0; i < VECTOR_COUNT; i++) { vectors[i] = new VarCharVector("vector", allocator); vectors[i].allocateNew(100, 10); } unloader = new VectorUnloader(VectorSchemaRoot.of(vectors)); } @TearDown(Level.Invocation) public void tearDownInvoke() { if (recordBatch != null) { recordBatch.close(); } for (int i = 0; i < VECTOR_COUNT; i++) { vectors[i].close(); } } /** Tear down benchmarks. */ @TearDown(Level.Trial) public void tearDown() { allocator.close(); } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public void unloadBenchmark() { recordBatch = unloader.getRecordBatch(); } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder() .include(VectorUnloaderBenchmark.class.getSimpleName()) .forks(1) .build(); new Runner(opt).run(); } // checkstyle:on: MissingJavadocMethod } apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/dictionary/000077500000000000000000000000001515451615200306445ustar00rootroot00000000000000DictionaryEncoderBenchmarks.java000066400000000000000000000106001515451615200370300ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dictionary; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; import java.util.Random; import java.util.concurrent.TimeUnit; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.types.pojo.DictionaryEncoding; 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.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 DictionaryEncoder}. */ @State(Scope.Benchmark) public class DictionaryEncoderBenchmarks { private BufferAllocator allocator; private static final int DATA_SIZE = 1000; private static final int KEY_SIZE = 100; private static final int KEY_LENGTH = 10; private List keys = new ArrayList<>(); private VarCharVector vector; private VarCharVector dictionaryVector; /** Setup benchmarks. */ @Setup public void prepare() { for (int i = 0; i < KEY_SIZE; i++) { keys.add(generateUniqueKey(KEY_LENGTH)); } allocator = new RootAllocator(10 * 1024 * 1024); vector = new VarCharVector("vector", allocator); dictionaryVector = new VarCharVector("dict", allocator); vector.allocateNew(10240, DATA_SIZE); vector.setValueCount(DATA_SIZE); for (int i = 0; i < DATA_SIZE; i++) { byte[] value = keys.get(generateRandomIndex(KEY_SIZE)).getBytes(StandardCharsets.UTF_8); vector.setSafe(i, value, 0, value.length); } dictionaryVector.allocateNew(1024, 100); dictionaryVector.setValueCount(100); for (int i = 0; i < KEY_SIZE; i++) { byte[] value = keys.get(i).getBytes(StandardCharsets.UTF_8); dictionaryVector.setSafe(i, value, 0, value.length); } } /** Tear down benchmarks. */ @TearDown public void tearDown() { vector.close(); dictionaryVector.close(); keys.clear(); allocator.close(); } /** * Test encode for {@link DictionaryEncoder}. * * @return useless. To avoid DCE by JIT. */ @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.NANOSECONDS) public int testEncode() { Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); final ValueVector encoded = DictionaryEncoder.encode(vector, dictionary); encoded.close(); return 0; } private int generateRandomIndex(int max) { Random random = new Random(); return random.nextInt(max); } private String generateUniqueKey(int length) { String str = "abcdefghijklmnopqrstuvwxyz"; Random random = new Random(); StringBuffer sb = new StringBuffer(); for (int i = 0; i < length; i++) { int number = random.nextInt(26); sb.append(str.charAt(number)); } if (keys.contains(sb.toString())) { return generateUniqueKey(length); } return sb.toString(); } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder() .include(DictionaryEncoderBenchmarks.class.getSimpleName()) .forks(1) .build(); new Runner(opt).run(); } } apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/ipc/000077500000000000000000000000001515451615200272525ustar00rootroot00000000000000WriteChannelBenchmark.java000066400000000000000000000054271515451615200342440ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/ipc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.channels.Channels; import java.util.concurrent.TimeUnit; 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 WriteChannel}. */ public class WriteChannelBenchmark { // checkstyle:off: MissingJavadocMethod /** State object for align benchmark. */ @State(Scope.Benchmark) public static class AlignState { private ByteArrayOutputStream baos; private WriteChannel writeChannel; @Param({"1", "2", "3", "4", "5", "6", "7"}) public int alignSize; @Setup(Level.Invocation) public void prepareInvoke() throws IOException { baos = new ByteArrayOutputStream(8); writeChannel = new WriteChannel(Channels.newChannel(baos)); writeChannel.write(new byte[8 - alignSize]); } @TearDown(Level.Invocation) public void tearDownInvoke() throws IOException { writeChannel.close(); baos.close(); } } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.NANOSECONDS) public void alignBenchmark(AlignState state) throws IOException { state.writeChannel.align(); } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder().include(WriteChannelBenchmark.class.getSimpleName()).forks(1).build(); new Runner(opt).run(); } // checkstyle:on: MissingJavadocMethod } apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/ipc/message/000077500000000000000000000000001515451615200306765ustar00rootroot00000000000000ArrowRecordBatchBenchmarks.java000066400000000000000000000061021515451615200366520ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/ipc/message/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message; import java.util.ArrayList; import java.util.List; import java.util.concurrent.TimeUnit; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.VarCharVector; 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.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 ArrowRecordBatch}. */ @State(Scope.Benchmark) public class ArrowRecordBatchBenchmarks { // 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 BufferAllocator allocator; private VarCharVector vector; private List nodes; /** Setup benchmarks. */ @Setup public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); vector = new VarCharVector("vector", allocator); vector.allocateNew(VECTOR_CAPACITY, VECTOR_LENGTH); nodes = new ArrayList<>(); nodes.add(new ArrowFieldNode(VECTOR_LENGTH, 0)); nodes.add(new ArrowFieldNode(VECTOR_LENGTH, 0)); } /** Tear down benchmarks. */ @TearDown public void tearDown() { vector.close(); allocator.close(); } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.NANOSECONDS) public long createAndGetLength() { try (ArrowRecordBatch batch = new ArrowRecordBatch(VECTOR_LENGTH, nodes, vector.getFieldBuffers())) { return batch.computeBodyLength(); } } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder() .include(ArrowRecordBatchBenchmarks.class.getSimpleName()) .forks(1) .build(); new Runner(opt).run(); } // checkstyle:on: MissingJavadocMethod } apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/util/000077500000000000000000000000001515451615200274545ustar00rootroot00000000000000TransferPairBenchmarks.java000066400000000000000000000075761515451615200346550ustar00rootroot00000000000000apache-arrow-java-19.0.0/performance/src/main/java/org/apache/arrow/vector/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.nio.charset.StandardCharsets; import java.util.concurrent.TimeUnit; 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.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.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 TransferPair}. */ @State(Scope.Benchmark) public class TransferPairBenchmarks { // checkstyle:off: MissingJavadocMethod private static final int VECTOR_LENGTH = 1024; private static final int ALLOCATOR_CAPACITY = 1024 * 1024; private BufferAllocator allocator; private IntVector intVector; private VarCharVector varCharVector; /** Setup benchmarks. */ @Setup public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); intVector = new IntVector("intVector", allocator); varCharVector = new VarCharVector("varcharVector", allocator); intVector.allocateNew(VECTOR_LENGTH); varCharVector.allocateNew(VECTOR_LENGTH); for (int i = 0; i < VECTOR_LENGTH; i++) { if (i % 3 == 0) { intVector.setNull(i); varCharVector.setNull(i); } else { intVector.setSafe(i, i * i); varCharVector.setSafe(i, ("teststring" + i).getBytes(StandardCharsets.UTF_8)); } } intVector.setValueCount(VECTOR_LENGTH); varCharVector.setValueCount(VECTOR_LENGTH); } /** Tear down benchmarks. */ @TearDown public void tearDown() { intVector.close(); varCharVector.close(); allocator.close(); } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public int splitAndTransferIntVector() { IntVector toVector = new IntVector("intVector", allocator); toVector.setValueCount(VECTOR_LENGTH); TransferPair transferPair = intVector.makeTransferPair(toVector); transferPair.splitAndTransfer(0, VECTOR_LENGTH); toVector.close(); return 0; } @Benchmark @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public int splitAndTransferVarcharVector() { VarCharVector toVector = new VarCharVector("varcharVector", allocator); toVector.setValueCount(VECTOR_LENGTH); TransferPair transferPair = varCharVector.makeTransferPair(toVector); transferPair.splitAndTransfer(0, VECTOR_LENGTH); toVector.close(); return 0; } public static void main(String[] args) throws RunnerException { Options opt = new OptionsBuilder().include(TransferPairBenchmarks.class.getSimpleName()).forks(1).build(); new Runner(opt).run(); } // checkstyle:on: MissingJavadocMethod } apache-arrow-java-19.0.0/pom.xml000066400000000000000000001377601515451615200164310ustar00rootroot00000000000000 4.0.0 org.apache apache 35 org.apache.arrow arrow-java-root 19.0.0 pom Apache Arrow Java Root POM Apache Arrow is open source, in-memory columnar data structures and low-overhead messaging https://arrow.apache.org/ Developer List dev-subscribe@arrow.apache.org dev-unsubscribe@arrow.apache.org dev@arrow.apache.org https://lists.apache.org/list.html?dev@arrow.apache.org Commits List commits-subscribe@arrow.apache.org commits-unsubscribe@arrow.apache.org commits@arrow.apache.org https://lists.apache.org/list.html?commits@arrow.apache.org Issues List issues-subscribe@arrow.apache.org issues-unsubscribe@arrow.apache.org https://lists.apache.org/list.html?issues@arrow.apache.org GitHub List github-subscribe@arrow.apache.org github-unsubscribe@arrow.apache.org https://lists.apache.org/list.html?github@arrow.apache.org bom format memory arrow-variant vector tools adapter/jdbc flight performance algorithm adapter/avro compression scm:git:https://github.com/apache/arrow-java.git scm:git:https://github.com/apache/arrow-java.git v19.0.0 https://github.com/apache/arrow-java/tree/${project.scm.tag} GitHub https://github.com/apache/arrow-java/issues 1773307790 ${project.build.directory}/generated-sources 1.9.0 5.12.2 2.0.17 33.4.8-jre 4.2.9.Final 1.79.0 4.33.4 2.21.0 3.4.3 25.2.10 1.12.1 1.17.0 5.21.0 2 10.23.0 true 2.42.0 3.53.1 1.5.32 none -Xdoclint:none --add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED 11 11 11 11 3.2.2 org.apache.arrow arrow-bom ${project.version} pom import org.checkerframework checker-qual ${checker.framework.version} provided com.google.flatbuffers flatbuffers-java ${dep.fbs.version} com.google.errorprone error_prone_annotations ${error_prone_core.version} provided org.slf4j slf4j-api ${dep.slf4j.version} org.slf4j slf4j-jdk14 ${dep.slf4j.version} javax.annotation javax.annotation-api 1.3.2 org.assertj assertj-core 3.27.7 test org.immutables value-annotations 2.12.1 provided org.hamcrest hamcrest 3.0 com.fasterxml.jackson jackson-bom ${dep.jackson-bom.version} pom import com.google.guava guava-bom ${dep.guava-bom.version} pom import io.netty netty-bom ${dep.netty-bom.version} pom import io.grpc grpc-bom ${dep.grpc-bom.version} pom import com.google.protobuf protobuf-bom ${dep.protobuf-bom.version} pom import org.mockito mockito-bom ${dep.mockito-bom.version} pom import ch.qos.logback logback-classic ${logback.version} ch.qos.logback logback-core ${logback.version} org.slf4j jul-to-slf4j ${dep.slf4j.version} test org.slf4j jcl-over-slf4j ${dep.slf4j.version} test org.slf4j log4j-over-slf4j ${dep.slf4j.version} test org.junit.jupiter junit-jupiter-engine ${dep.junit.jupiter.version} test org.junit.jupiter junit-jupiter-api ${dep.junit.jupiter.version} test org.junit.jupiter junit-jupiter-params ${dep.junit.jupiter.version} test ch.qos.logback logback-classic ${logback.version} test de.huxhorn.lilith de.huxhorn.lilith.logback.appender.multiplex-classic 8.3.0 test org.apache.maven.plugins maven-compiler-plugin true false org.immutables value 2.12.1 maven-surefire-plugin ${surefire.add-opens.argLine} true true ${forkCount} true ${project.build.directory} true UTC 2097151 false maven-failsafe-plugin ${surefire.add-opens.argLine} ${project.build.directory} true UTC org.jacoco jacoco-maven-plugin 0.8.14 org.eclipse.m2e lifecycle-mapping 1.0.0 org.apache.maven.plugins maven-antrun-plugin [1.6,) run org.apache.maven.plugins maven-enforcer-plugin [1.2,) enforce org.apache.maven.plugins maven-remote-resources-plugin [1.1,) process org.apache.maven.plugins maven-checkstyle-plugin [0,) check org.apache.drill.tools drill-fmpp-maven-plugin [1.0,) generate false true org.apache.maven.plugins maven-javadoc-plugin **/module-info.java arrow-memory-netty-buffer-patch,arrow-memory-netty,flight-sql-jdbc-core,flight-integration-tests,arrow-performance com.gradle develocity-maven-extension arrow-git.properties properties/flight.properties org.apache.maven.plugins maven-surefire-plugin testingData ${arrow.test.dataRoot} org.xolstice.maven.plugins protobuf-maven-plugin 0.6.1 com.google.protobuf:protoc:${dep.protobuf-bom.version}:exe:${os.detected.classifier} grpc-java io.grpc:protoc-gen-grpc-java:${dep.grpc-bom.version}:exe:${os.detected.classifier} @generated=omit com.diffplug.spotless spotless-maven-plugin 2.44.4 org.codehaus.mojo build-helper-maven-plugin 3.6.1 org.codehaus.mojo properties-maven-plugin 1.2.1 org.codehaus.mojo exec-maven-plugin 3.6.3 org.codehaus.mojo versions-maven-plugin 2.21.0 pl.project13.maven git-commit-id-plugin 4.9.10 org.cyclonedx cyclonedx-maven-plugin 2.9.1 org.apache.drill.tools drill-fmpp-maven-plugin 1.22.0 org.apache.maven.plugins maven-jar-plugin **/logging.properties **/logback-test.xml **/logback.out.xml **/logback.xml org.apache.arrow ${username} https://arrow.apache.org/ test-jar true org.apache.maven.plugins maven-compiler-plugin 2048m true maven-enforcer-plugin avoid_bad_dependencies enforce verify commons-logging javax.servlet:servlet-api org.mortbay.jetty:servlet-api org.mortbay.jetty:servlet-api-2.5 log4j:log4j org.immutables:value org.checkerframework:checker junit:junit:4.* org.junit.vintage:junit-vintage-engine pl.project13.maven git-commit-id-plugin dd.MM.yyyy '@' HH:mm:ss z false false true false false false 7 -dirty true for-jars revision true target/classes/arrow-git.properties for-source-tarball revision false ./arrow-git.properties org.apache.maven.plugins maven-checkstyle-plugin **/module-info.java dev/checkstyle/checkstyle.xml dev/license/asf-java.license dev/checkstyle/suppressions.xml true UTF-8 true ${checkstyle.failOnViolation} ${checkstyle.failOnViolation} warning xml ${project.build.directory}/test/checkstyle-errors.xml false com.puppycrawl.tools checkstyle ${checkstyle.version} org.slf4j jcl-over-slf4j ${dep.slf4j.version} validate check validate org.apache.maven.plugins maven-dependency-plugin analyze analyze-only verify true true javax.annotation:javax.annotation-api:* org.apache.hadoop:hadoop-client-api org.cyclonedx cyclonedx-maven-plugin makeBom package org.apache.maven.plugins maven-project-info-reports-plugin org.apache.maven.plugins maven-site-plugin com.diffplug.spotless spotless-maven-plugin ${maven.multiModuleProjectDirectory}/dev/license/asf-xml.license (<configuration|<project) 1.17.0 ${maven.multiModuleProjectDirectory}/dev/license/asf-java.license package **/Preconditions.java **/IntObjectMap.java **/IntObjectHashMap.java spotless-check check kr.motd.maven os-maven-plugin 1.7.1 org.apache.maven.plugins maven-javadoc-plugin **/module-info.java arrow-memory-netty-buffer-patch,arrow-memory-netty,flight-sql-jdbc-core,flight-integration-tests,arrow-performance javadoc aggregate aggregate false org.apache.maven.plugins maven-project-info-reports-plugin org.apache.maven.plugins maven-site-plugin arrow-c-data c arrow-jni adapter/orc gandiva dataset c shade-flatbuffers shade-format-flatbuffers error-prone [17,) !m2e.version org.apache.maven.plugins maven-compiler-plugin -XDcompilePolicy=simple -Xplugin:ErrorProne -XepExcludedPaths:.*/(target/generated-source|format/src/main/java/org/apache/arrow/flatbuf)/.* -J--add-exports=jdk.compiler/com.sun.tools.javac.api=ALL-UNNAMED -J--add-exports=jdk.compiler/com.sun.tools.javac.file=ALL-UNNAMED -J--add-exports=jdk.compiler/com.sun.tools.javac.main=ALL-UNNAMED -J--add-exports=jdk.compiler/com.sun.tools.javac.model=ALL-UNNAMED -J--add-exports=jdk.compiler/com.sun.tools.javac.parser=ALL-UNNAMED -J--add-exports=jdk.compiler/com.sun.tools.javac.processing=ALL-UNNAMED -J--add-exports=jdk.compiler/com.sun.tools.javac.tree=ALL-UNNAMED -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 com.google.errorprone error_prone_core ${error_prone_core.version} code-coverage org.jacoco jacoco-maven-plugin default-prepare-agent prepare-agent default-prepare-agent-integration prepare-agent-integration org.jacoco jacoco-maven-plugin report false windows [17,] windows maven-surefire-plugin false generate-libs-cdata-all-os java-dist org.codehaus.mojo exec-maven-plugin cdata-cmake exec generate-resources cmake -S java -B java-jni -DARROW_JAVA_JNI_ENABLE_C=ON -DARROW_JAVA_JNI_ENABLE_DEFAULT=OFF -DBUILD_TESTING=OFF -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=${arrow.c.jni.dist.dir} ../ cdata-build exec generate-resources cmake --build java-jni --target install --config Release ../ generate-libs-jni-macos-linux java-dist false ON ON ON ON OFF ON ON ON org.codehaus.mojo exec-maven-plugin jni-cpp-cmake exec generate-resources cmake -S cpp -B cpp-jni -DARROW_BUILD_SHARED=OFF -DARROW_CSV=${ARROW_DATASET} -DARROW_DATASET=${ARROW_DATASET} -DARROW_DEPENDENCY_SOURCE=BUNDLED -DARROW_DEPENDENCY_USE_SHARED=OFF -DARROW_FILESYSTEM=ON -DARROW_GANDIVA=${ARROW_GANDIVA} -DARROW_GANDIVA_STATIC_LIBSTDCPP=ON -DARROW_JSON=${ARROW_DATASET} -DARROW_ORC=${ARROW_ORC} -DARROW_PARQUET=${ARROW_PARQUET} -DARROW_S3=ON -DARROW_SUBSTRAIT=${ARROW_DATASET} -DARROW_USE_CCACHE=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=java-dist -DCMAKE_UNITY_BUILD=ON ../ ${cpp.dependencies.builded} jni-cpp-build exec generate-resources cmake --build cpp-jni --target install --config Release ../ ${cpp.dependencies.builded} jni-cmake exec generate-resources cmake -S java -B java-jni -DARROW_JAVA_JNI_ENABLE_C=${ARROW_JAVA_JNI_ENABLE_C} -DARROW_JAVA_JNI_ENABLE_DATASET=${ARROW_JAVA_JNI_ENABLE_DATASET} -DARROW_JAVA_JNI_ENABLE_GANDIVA=${ARROW_JAVA_JNI_ENABLE_GANDIVA} -DARROW_JAVA_JNI_ENABLE_ORC=${ARROW_JAVA_JNI_ENABLE_ORC} -DARROW_JAVA_JNI_ENABLE_DEFAULT=ON -DBUILD_TESTING=OFF -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=${arrow.dataset.jni.dist.dir} -DCMAKE_PREFIX_PATH=${project.basedir}/../java-dist/lib/${os.detected.arch}/cmake -DProtobuf_USE_STATIC_LIBS=ON -DProtobuf_ROOT=${project.basedir}/../cpp-jni/protobuf_ep-install ../ jni-build exec generate-resources cmake --build java-jni --target install --config Release ../ generate-libs-jni-windows java-dist false ON OFF ON ON OFF ON OFF ON org.codehaus.mojo exec-maven-plugin jni-cpp-cmake exec generate-resources cmake -S cpp -B cpp-jni -DARROW_BUILD_SHARED=OFF -DARROW_CSV=${ARROW_DATASET} -DARROW_DATASET=${ARROW_DATASET} -DARROW_DEPENDENCY_USE_SHARED=OFF -DARROW_FILESYSTEM=ON -DARROW_GANDIVA=${ARROW_GANDIVA} -DARROW_JSON=${ARROW_DATASET} -DARROW_ORC=${ARROW_ORC} -DARROW_PARQUET=${ARROW_PARQUET} -DARROW_S3=ON -DARROW_SUBSTRAIT=${ARROW_DATASET} -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 ../ ${cpp.dependencies.builded} jni-cpp-build exec generate-resources ninja install ../cpp-jni ${cpp.dependencies.builded} jni-cmake exec generate-resources cmake -S java -B java-jni -DARROW_JAVA_JNI_ENABLE_C=${ARROW_JAVA_JNI_ENABLE_C} -DARROW_JAVA_JNI_ENABLE_DATASET=${ARROW_JAVA_JNI_ENABLE_DATASET} -DARROW_JAVA_JNI_ENABLE_GANDIVA=${ARROW_JAVA_JNI_ENABLE_GANDIVA} -DARROW_JAVA_JNI_ENABLE_ORC=${ARROW_JAVA_JNI_ENABLE_ORC} -DARROW_JAVA_JNI_ENABLE_DEFAULT=ON -DBUILD_TESTING=OFF -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=${arrow.dataset.jni.dist.dir} -DCMAKE_PREFIX_PATH=${project.basedir}/../java-dist/lib/${os.detected.arch}/cmake ../ jni-build exec generate-resources cmake --build java-jni --target install --config Release ../ cross-jdk-testing arrow.test.jdk-version maven-enforcer-plugin check-jdk-version-property enforce validate arrow.test.jdk-version "JDK version used for test must be specified." ^\d{2,} "JDK version used for test must be 11, 17, 21, ..." org.apache.maven.plugins maven-surefire-plugin ${arrow.test.jdk-version} org.apache.maven.plugins maven-failsafe-plugin ${arrow.test.jdk-version} apache-arrow-java-19.0.0/testing/000077500000000000000000000000001515451615200165535ustar00rootroot00000000000000apache-arrow-java-19.0.0/tools/000077500000000000000000000000001515451615200162365ustar00rootroot00000000000000apache-arrow-java-19.0.0/tools/pom.xml000066400000000000000000000134631515451615200175620ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-java-root 19.0.0 arrow-tools Arrow Tools Java applications for working with Arrow ValueVectors. org.apache.arrow arrow-memory-core org.apache.arrow arrow-vector ${arrow.vector.classifier} org.apache.arrow arrow-compression org.immutables value-annotations com.google.guava guava test commons-cli commons-cli 1.11.0 ch.qos.logback logback-classic test com.fasterxml.jackson.core jackson-core com.fasterxml.jackson.core jackson-databind org.slf4j slf4j-api org.apache.arrow arrow-memory-netty runtime org.apache.arrow arrow-vector ${project.version} tests test-jar test org.apache.maven.plugins maven-shade-plugin make-assembly shade package false true jar-with-dependencies **/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 org.apache.maven.plugins maven-dependency-plugin analyze verify com.fasterxml.jackson.core:* apache-arrow-java-19.0.0/tools/src/000077500000000000000000000000001515451615200170255ustar00rootroot00000000000000apache-arrow-java-19.0.0/tools/src/main/000077500000000000000000000000001515451615200177515ustar00rootroot00000000000000apache-arrow-java-19.0.0/tools/src/main/java/000077500000000000000000000000001515451615200206725ustar00rootroot00000000000000apache-arrow-java-19.0.0/tools/src/main/java/module-info.java000066400000000000000000000021011515451615200237450ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ module org.apache.arrow.tools { exports org.apache.arrow.tools; requires com.fasterxml.jackson.databind; requires org.apache.arrow.compression; requires org.apache.arrow.memory.core; requires org.apache.arrow.vector; requires org.apache.commons.cli; requires org.slf4j; } apache-arrow-java-19.0.0/tools/src/main/java/org/000077500000000000000000000000001515451615200214615ustar00rootroot00000000000000apache-arrow-java-19.0.0/tools/src/main/java/org/apache/000077500000000000000000000000001515451615200227025ustar00rootroot00000000000000apache-arrow-java-19.0.0/tools/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200240345ustar00rootroot00000000000000apache-arrow-java-19.0.0/tools/src/main/java/org/apache/arrow/tools/000077500000000000000000000000001515451615200251745ustar00rootroot00000000000000apache-arrow-java-19.0.0/tools/src/main/java/org/apache/arrow/tools/EchoServer.java000066400000000000000000000107321515451615200301070ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.tools; import java.io.IOException; import java.net.ServerSocket; import java.net.Socket; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; 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.ArrowStreamWriter; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** Simple server that echoes back data received. */ public class EchoServer { private static final Logger LOGGER = LoggerFactory.getLogger(EchoServer.class); private final ServerSocket serverSocket; private boolean closed = false; /** Constructs a new instance that binds to the given port. */ public EchoServer(int port) throws IOException { LOGGER.debug("Starting echo server."); serverSocket = new ServerSocket(port); LOGGER.debug("Running echo server on port: " + port()); } /** Main method to run the server, the first argument is an optional port number. */ public static void main(String[] args) throws Exception { int port; if (args.length > 0) { port = Integer.parseInt(args[0]); } else { port = 8080; } new EchoServer(port).run(); } public int port() { return serverSocket.getLocalPort(); } /** Starts the main server event loop. */ public void run() throws IOException { try { Socket clientSocket = null; ClientConnection client = null; while (!closed) { LOGGER.debug("Waiting to accept new client connection."); clientSocket = serverSocket.accept(); LOGGER.debug("Accepted new client connection."); client = new ClientConnection(clientSocket); try { client.run(); } catch (IOException e) { LOGGER.warn("Error handling client connection.", e); } LOGGER.debug("Closed connection with client"); } } catch (java.net.SocketException ex) { if (!closed) { throw ex; } } finally { serverSocket.close(); LOGGER.debug("Server closed."); } } public void close() throws IOException { closed = true; serverSocket.close(); } /** Handler for each client connection to the server. */ public static class ClientConnection implements AutoCloseable { public final Socket socket; public ClientConnection(Socket socket) { this.socket = socket; } /** Reads a record batch off the socket and writes it back out. */ public void run() throws IOException { // Read the entire input stream and write it back try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { ArrowStreamReader reader = new ArrowStreamReader(socket.getInputStream(), allocator); VectorSchemaRoot root = reader.getVectorSchemaRoot(); // load the first batch before instantiating the writer so that we have any dictionaries reader.loadNextBatch(); ArrowStreamWriter writer = new ArrowStreamWriter(root, reader, socket.getOutputStream()); writer.start(); int echoed = 0; while (true) { int rowCount = reader.getVectorSchemaRoot().getRowCount(); if (rowCount == 0) { break; } else { writer.writeBatch(); echoed += rowCount; reader.loadNextBatch(); } } writer.end(); Preconditions.checkState(reader.bytesRead() == writer.bytesWritten()); LOGGER.debug(String.format("Echoed %d records", echoed)); reader.close(false); writer.close(); } } @Override public void close() throws IOException { socket.close(); } } } apache-arrow-java-19.0.0/tools/src/main/java/org/apache/arrow/tools/FileRoundtrip.java000066400000000000000000000107431515451615200306320ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.tools; import java.io.File; import java.io.FileInputStream; import java.io.FileOutputStream; import java.io.IOException; import java.io.PrintStream; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.ArrowFileReader; import org.apache.arrow.vector.ipc.ArrowFileWriter; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.commons.cli.CommandLine; import org.apache.commons.cli.CommandLineParser; import org.apache.commons.cli.Options; import org.apache.commons.cli.ParseException; import org.apache.commons.cli.PosixParser; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** Application that verifies data can be round-tripped through a file. */ public class FileRoundtrip { private static final Logger LOGGER = LoggerFactory.getLogger(FileRoundtrip.class); private final Options options; private final PrintStream err; FileRoundtrip(PrintStream err) { this.err = err; this.options = new Options(); this.options.addOption("i", "in", true, "input file"); this.options.addOption("o", "out", true, "output file"); } public static void main(String[] args) { System.exit(new FileRoundtrip(System.err).run(args)); } private File validateFile(String type, String fileName) throws IOException { if (fileName == null) { throw new IllegalArgumentException("missing " + type + " file parameter"); } File f = new File(fileName); if (type.equals("input")) { if (!f.exists() || f.isDirectory()) { throw new IllegalArgumentException(type + " file not found: " + f.getAbsolutePath()); } } else if (type.equals("output")) { File parentDir = f.getParentFile(); if (parentDir != null && !parentDir.exists()) { if (!parentDir.mkdirs()) { throw new IOException( "Failed to create parent directory: " + parentDir.getAbsolutePath()); } } } return f; } int run(String[] args) { try { CommandLineParser parser = new PosixParser(); CommandLine cmd = parser.parse(options, args, false); String inFileName = cmd.getOptionValue("in"); String outFileName = cmd.getOptionValue("out"); File inFile = validateFile("input", inFileName); File outFile = validateFile("output", outFileName); try (BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); FileInputStream fileInputStream = new FileInputStream(inFile); ArrowFileReader arrowReader = new ArrowFileReader(fileInputStream.getChannel(), allocator)) { VectorSchemaRoot root = arrowReader.getVectorSchemaRoot(); Schema schema = root.getSchema(); LOGGER.debug("Input file size: " + inFile.length()); LOGGER.debug("Found schema: " + schema); try (FileOutputStream fileOutputStream = new FileOutputStream(outFile); ArrowFileWriter arrowWriter = new ArrowFileWriter(root, arrowReader, fileOutputStream.getChannel())) { arrowWriter.start(); while (true) { if (!arrowReader.loadNextBatch()) { break; } else { arrowWriter.writeBatch(); } } arrowWriter.end(); } LOGGER.debug("Output file size: " + outFile.length()); } } catch (ParseException e) { return fatalError("Invalid parameters", e); } catch (IOException e) { return fatalError("Error accessing files", e); } return 0; } private int fatalError(String message, Throwable e) { err.println(message); LOGGER.error(message, e); return 1; } } apache-arrow-java-19.0.0/tools/src/main/java/org/apache/arrow/tools/FileToStream.java000066400000000000000000000054571515451615200304100ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.tools; import java.io.File; import java.io.FileInputStream; import java.io.FileOutputStream; import java.io.IOException; import java.io.OutputStream; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.ArrowFileReader; import org.apache.arrow.vector.ipc.ArrowStreamWriter; /** * Converts an Arrow file to an Arrow stream. The file should be specified as the first argument and * the output is written to standard out. */ public class FileToStream { private FileToStream() {} /** Reads an Arrow file from in and writes it back to out. */ public static void convert(FileInputStream in, OutputStream out) throws IOException { BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); try (ArrowFileReader reader = new ArrowFileReader(in.getChannel(), allocator)) { VectorSchemaRoot root = reader.getVectorSchemaRoot(); // load the first batch before instantiating the writer so that we have any dictionaries // only writeBatches if we loaded one in the first place. boolean writeBatches = reader.loadNextBatch(); try (ArrowStreamWriter writer = new ArrowStreamWriter(root, reader, out)) { writer.start(); while (writeBatches) { writer.writeBatch(); if (!reader.loadNextBatch()) { break; } } writer.end(); } } } /** * Main method. The first arg is the file path. The second, optional argument, is an output file * location (defaults to standard out). */ public static void main(String[] args) throws IOException { if (args.length != 1 && args.length != 2) { System.err.println("Usage: FileToStream [output file]"); System.exit(1); } FileInputStream in = new FileInputStream(new File(args[0])); OutputStream out = args.length == 1 ? System.out : new FileOutputStream(new File(args[1])); convert(in, out); } } apache-arrow-java-19.0.0/tools/src/main/java/org/apache/arrow/tools/Integration.java000066400000000000000000000233141515451615200303250ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.tools; import java.io.File; import java.io.FileInputStream; import java.io.FileOutputStream; import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; import java.util.Iterator; import java.util.List; import org.apache.arrow.compression.CommonsCompressionFactory; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.ArrowFileReader; import org.apache.arrow.vector.ipc.ArrowFileWriter; import org.apache.arrow.vector.ipc.JsonFileReader; import org.apache.arrow.vector.ipc.JsonFileWriter; import org.apache.arrow.vector.ipc.message.ArrowBlock; import org.apache.arrow.vector.types.pojo.DictionaryEncoding; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.Validator; import org.apache.commons.cli.CommandLine; import org.apache.commons.cli.CommandLineParser; import org.apache.commons.cli.Options; import org.apache.commons.cli.ParseException; import org.apache.commons.cli.PosixParser; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** Application for cross language integration testing. */ public class Integration { private static final Logger LOGGER = LoggerFactory.getLogger(Integration.class); private final Options options; Integration() { this.options = new Options(); this.options.addOption("a", "arrow", true, "arrow file"); this.options.addOption("j", "json", true, "json file"); this.options.addOption( "c", "command", true, "command to execute: " + Arrays.toString(Command.values())); } /** Main method. */ public static void main(String[] args) { try { new Integration().run(args); } catch (ParseException e) { fatalError("Invalid parameters", e); } catch (IOException e) { fatalError("Error accessing files", e); } catch (RuntimeException e) { fatalError("Incompatible files", e); } } private static void fatalError(String message, Throwable e) { System.err.println(message); System.err.println(e.getMessage()); LOGGER.error(message, e); System.exit(1); } private File validateFile(String type, String fileName, boolean shouldExist) { if (fileName == null) { throw new IllegalArgumentException("missing " + type + " file parameter"); } File f = new File(fileName); if (shouldExist && (!f.exists() || f.isDirectory())) { throw new IllegalArgumentException(type + " file not found: " + f.getAbsolutePath()); } if (!shouldExist && f.exists()) { throw new IllegalArgumentException(type + " file already exists: " + f.getAbsolutePath()); } return f; } static void extractDictionaryEncodings(List fields, List encodings) { for (Field field : fields) { DictionaryEncoding encoding = field.getDictionary(); if (encoding != null) { encodings.add(encoding); } extractDictionaryEncodings(field.getChildren(), encodings); } } void run(String[] args) throws ParseException, IOException { CommandLineParser parser = new PosixParser(); CommandLine cmd = parser.parse(options, args, false); Command command = toCommand(cmd.getOptionValue("command")); File arrowFile = validateFile("arrow", cmd.getOptionValue("arrow"), command.arrowExists); File jsonFile = validateFile("json", cmd.getOptionValue("json"), command.jsonExists); command.execute(arrowFile, jsonFile); } private Command toCommand(String commandName) { try { return Command.valueOf(commandName); } catch (IllegalArgumentException e) { throw new IllegalArgumentException( "Unknown command: " + commandName + " expected one of " + Arrays.toString(Command.values())); } } /** Commands (actions) the application can perform. */ enum Command { ARROW_TO_JSON(true, false) { @Override public void execute(File arrowFile, File jsonFile) throws IOException { try (BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); FileInputStream fileInputStream = new FileInputStream(arrowFile); ArrowFileReader arrowReader = new ArrowFileReader(fileInputStream.getChannel(), allocator)) { VectorSchemaRoot root = arrowReader.getVectorSchemaRoot(); Schema schema = root.getSchema(); LOGGER.debug("Input file size: " + arrowFile.length()); LOGGER.debug("Found schema: " + schema); try (JsonFileWriter writer = new JsonFileWriter(jsonFile, JsonFileWriter.config().pretty(true))) { writer.start(schema, arrowReader); for (ArrowBlock rbBlock : arrowReader.getRecordBlocks()) { if (!arrowReader.loadRecordBatch(rbBlock)) { throw new IOException("Expected to load record batch"); } writer.write(root); } } LOGGER.debug("Output file size: " + jsonFile.length()); } } }, JSON_TO_ARROW(false, true) { @Override public void execute(File arrowFile, File jsonFile) throws IOException { try (BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); JsonFileReader reader = new JsonFileReader(jsonFile, allocator)) { Schema schema = reader.start(); LOGGER.debug("Input file size: " + jsonFile.length()); LOGGER.debug("Found schema: " + schema); try (FileOutputStream fileOutputStream = new FileOutputStream(arrowFile); VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); // TODO json dictionaries ArrowFileWriter arrowWriter = new ArrowFileWriter(root, reader, fileOutputStream.getChannel())) { arrowWriter.start(); while (reader.read(root)) { arrowWriter.writeBatch(); } arrowWriter.end(); } LOGGER.debug("Output file size: " + arrowFile.length()); } } }, VALIDATE(true, true) { @Override public void execute(File arrowFile, File jsonFile) throws IOException { try (BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); JsonFileReader jsonReader = new JsonFileReader(jsonFile, allocator); FileInputStream fileInputStream = new FileInputStream(arrowFile); ArrowFileReader arrowReader = new ArrowFileReader( fileInputStream.getChannel(), allocator, CommonsCompressionFactory.INSTANCE)) { Schema jsonSchema = jsonReader.start(); VectorSchemaRoot arrowRoot = arrowReader.getVectorSchemaRoot(); Schema arrowSchema = arrowRoot.getSchema(); LOGGER.debug("Arrow Input file size: " + arrowFile.length()); LOGGER.debug("ARROW schema: " + arrowSchema); LOGGER.debug("JSON Input file size: " + jsonFile.length()); LOGGER.debug("JSON schema: " + jsonSchema); Validator.compareSchemas(jsonSchema, arrowSchema); List recordBatches = arrowReader.getRecordBlocks(); Iterator iterator = recordBatches.iterator(); VectorSchemaRoot jsonRoot; int totalBatches = 0; while ((jsonRoot = jsonReader.read()) != null && iterator.hasNext()) { ArrowBlock rbBlock = iterator.next(); if (!arrowReader.loadRecordBatch(rbBlock)) { throw new IOException("Expected to load record batch"); } Validator.compareVectorSchemaRoot(arrowRoot, jsonRoot); jsonRoot.close(); totalBatches++; } // Validate Dictionaries after ArrowFileReader has read batches List encodingsJson = new ArrayList<>(); extractDictionaryEncodings(jsonSchema.getFields(), encodingsJson); List encodingsArrow = new ArrayList<>(); extractDictionaryEncodings(arrowSchema.getFields(), encodingsArrow); Validator.compareDictionaries(encodingsJson, encodingsArrow, jsonReader, arrowReader); boolean hasMoreJSON = jsonRoot != null; boolean hasMoreArrow = iterator.hasNext(); if (hasMoreJSON || hasMoreArrow) { throw new IllegalArgumentException( "Unexpected RecordBatches. Total: " + totalBatches + " J:" + hasMoreJSON + " " + "A:" + hasMoreArrow); } } } }; public final boolean arrowExists; public final boolean jsonExists; Command(boolean arrowExists, boolean jsonExists) { this.arrowExists = arrowExists; this.jsonExists = jsonExists; } public abstract void execute(File arrowFile, File jsonFile) throws IOException; } } apache-arrow-java-19.0.0/tools/src/main/java/org/apache/arrow/tools/StreamToFile.java000066400000000000000000000054531515451615200304040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.tools; import java.io.File; import java.io.FileInputStream; import java.io.FileOutputStream; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; import java.nio.channels.Channels; import org.apache.arrow.compression.CommonsCompressionFactory; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.ArrowFileWriter; import org.apache.arrow.vector.ipc.ArrowStreamReader; /** Converts an Arrow stream to an Arrow file. */ public class StreamToFile { /** Reads an Arrow stream from in and writes it to out. */ public static void convert(InputStream in, OutputStream out) throws IOException { BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); try (ArrowStreamReader reader = new ArrowStreamReader(in, allocator, CommonsCompressionFactory.INSTANCE)) { VectorSchemaRoot root = reader.getVectorSchemaRoot(); // load the first batch before instantiating the writer so that we have any dictionaries. // Only writeBatches if we load the first one. boolean writeBatches = reader.loadNextBatch(); try (ArrowFileWriter writer = new ArrowFileWriter(root, reader, Channels.newChannel(out))) { writer.start(); while (writeBatches) { writer.writeBatch(); if (!reader.loadNextBatch()) { break; } } writer.end(); } } } /** * Main method. Defaults to reading from standard in and standard out. If there are two arguments * the first is interpreted as the input file path, the second is the output file path. */ public static void main(String[] args) throws IOException { InputStream in = System.in; OutputStream out = System.out; if (args.length == 2) { in = new FileInputStream(new File(args[0])); out = new FileOutputStream(new File(args[1])); } convert(in, out); } } apache-arrow-java-19.0.0/tools/src/shade/000077500000000000000000000000001515451615200201115ustar00rootroot00000000000000apache-arrow-java-19.0.0/tools/src/shade/LICENSE.txt000066400000000000000000000406421515451615200217420ustar00rootroot00000000000000 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. apache-arrow-java-19.0.0/tools/src/shade/NOTICE.txt000066400000000000000000000231051515451615200216340ustar00rootroot00000000000000Apache 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 apache-arrow-java-19.0.0/tools/src/test/000077500000000000000000000000001515451615200200045ustar00rootroot00000000000000apache-arrow-java-19.0.0/tools/src/test/java/000077500000000000000000000000001515451615200207255ustar00rootroot00000000000000apache-arrow-java-19.0.0/tools/src/test/java/org/000077500000000000000000000000001515451615200215145ustar00rootroot00000000000000apache-arrow-java-19.0.0/tools/src/test/java/org/apache/000077500000000000000000000000001515451615200227355ustar00rootroot00000000000000apache-arrow-java-19.0.0/tools/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200240675ustar00rootroot00000000000000apache-arrow-java-19.0.0/tools/src/test/java/org/apache/arrow/tools/000077500000000000000000000000001515451615200252275ustar00rootroot00000000000000apache-arrow-java-19.0.0/tools/src/test/java/org/apache/arrow/tools/ArrowFileTestFixtures.java000066400000000000000000000172261515451615200323660ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.tools; 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.io.IOException; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ViewVarBinaryVector; import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.complex.NonNullableStructVector; import org.apache.arrow.vector.complex.impl.ComplexWriterImpl; import org.apache.arrow.vector.complex.writer.BaseWriter.ComplexWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter; import org.apache.arrow.vector.complex.writer.BigIntWriter; import org.apache.arrow.vector.complex.writer.IntWriter; import org.apache.arrow.vector.complex.writer.ViewVarBinaryWriter; import org.apache.arrow.vector.complex.writer.ViewVarCharWriter; import org.apache.arrow.vector.ipc.ArrowFileReader; import org.apache.arrow.vector.ipc.ArrowFileWriter; import org.apache.arrow.vector.ipc.message.ArrowBlock; import org.apache.arrow.vector.util.Text; public class ArrowFileTestFixtures { static final int COUNT = 10; static void writeData(int count, NonNullableStructVector parent) { ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); IntWriter intWriter = rootWriter.integer("int"); BigIntWriter bigIntWriter = rootWriter.bigInt("bigInt"); for (int i = 0; i < count; i++) { intWriter.setPosition(i); intWriter.writeInt(i); bigIntWriter.setPosition(i); bigIntWriter.writeBigInt(i); } writer.setValueCount(count); } private static String generateString(int length) { StringBuilder stringBuilder = new StringBuilder(length); for (int i = 0; i < length; i++) { stringBuilder.append(i); } return stringBuilder.toString(); } private static byte[] generateBytes(int length) { byte[] bytes = new byte[length]; for (int i = 0; i < length; i++) { bytes[i] = (byte) i; } return bytes; } static void writeVariableWidthViewData(int count, NonNullableStructVector parent) { ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); ViewVarCharWriter viewVarCharWriter = rootWriter.viewVarChar("viewVarChar"); ViewVarBinaryWriter viewVarBinaryWriter = rootWriter.viewVarBinary("viewVarBinary"); IntWriter intWriter = rootWriter.integer("int"); BigIntWriter bigIntWriter = rootWriter.bigInt("bigInt"); for (int i = 0; i < count; i++) { viewVarCharWriter.setPosition(i); viewVarCharWriter.writeViewVarChar(generateString(i)); viewVarBinaryWriter.setPosition(i); viewVarBinaryWriter.writeViewVarBinary(generateBytes(i)); intWriter.setPosition(i); intWriter.writeInt(i); bigIntWriter.setPosition(i); bigIntWriter.writeBigInt(i); } writer.setValueCount(count); } static void validateOutput(File testOutFile, BufferAllocator allocator) throws Exception { // read try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE); FileInputStream fileInputStream = new FileInputStream(testOutFile); ArrowFileReader arrowReader = new ArrowFileReader(fileInputStream.getChannel(), readerAllocator)) { VectorSchemaRoot root = arrowReader.getVectorSchemaRoot(); for (ArrowBlock rbBlock : arrowReader.getRecordBlocks()) { if (!arrowReader.loadRecordBatch(rbBlock)) { throw new IOException("Expected to read record batch"); } validateContent(COUNT, root); } } } static void validateVariadicOutput(File testOutFile, BufferAllocator allocator, int count) throws Exception { // read try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE); FileInputStream fileInputStream = new FileInputStream(testOutFile); ArrowFileReader arrowReader = new ArrowFileReader(fileInputStream.getChannel(), readerAllocator)) { VectorSchemaRoot root = arrowReader.getVectorSchemaRoot(); for (ArrowBlock rbBlock : arrowReader.getRecordBlocks()) { if (!arrowReader.loadRecordBatch(rbBlock)) { throw new IOException("Expected to read record batch"); } validateVariadicContent(count, root); } } } static void validateContent(int count, VectorSchemaRoot root) { assertEquals(count, root.getRowCount()); for (int i = 0; i < count; i++) { assertEquals(i, root.getVector("int").getObject(i)); assertEquals(Long.valueOf(i), root.getVector("bigInt").getObject(i)); } } static void validateVariadicContent(int count, VectorSchemaRoot root) { assertEquals(count, root.getRowCount()); ViewVarCharVector viewVarCharVector = (ViewVarCharVector) root.getVector("viewVarChar"); ViewVarBinaryVector viewVarBinaryVector = (ViewVarBinaryVector) root.getVector("viewVarBinary"); IntVector intVector = (IntVector) root.getVector("int"); BigIntVector bigIntVector = (BigIntVector) root.getVector("bigInt"); for (int i = 0; i < count; i++) { assertEquals(new Text(generateString(i)), viewVarCharVector.getObject(i)); assertArrayEquals(generateBytes(i), viewVarBinaryVector.get(i)); assertEquals(i, intVector.getObject(i)); assertEquals(Long.valueOf(i), bigIntVector.getObject(i)); } } static void write(FieldVector parent, File file) throws IOException { VectorSchemaRoot root = new VectorSchemaRoot(parent); try (FileOutputStream fileOutputStream = new FileOutputStream(file); ArrowFileWriter arrowWriter = new ArrowFileWriter(root, null, fileOutputStream.getChannel())) { arrowWriter.writeBatch(); } } static void writeInput(File testInFile, BufferAllocator allocator) throws IOException { try (BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); NonNullableStructVector parent = NonNullableStructVector.empty("parent", vectorAllocator)) { writeData(COUNT, parent); write(parent.getChild("root"), testInFile); } } static void writeVariableWidthViewInput(File testInFile, BufferAllocator allocator, int count) throws IOException { try (BufferAllocator vectorAllocator = allocator.newChildAllocator("original view vectors", 0, Integer.MAX_VALUE); NonNullableStructVector parent = NonNullableStructVector.empty("parent", vectorAllocator)) { writeVariableWidthViewData(count, parent); write(parent.getChild("root"), testInFile); } } } apache-arrow-java-19.0.0/tools/src/test/java/org/apache/arrow/tools/EchoServerTest.java000066400000000000000000000277201515451615200310070ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.tools; import static java.util.Arrays.asList; import static org.apache.arrow.vector.types.Types.MinorType.TINYINT; import static org.apache.arrow.vector.types.Types.MinorType.VARCHAR; 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.assertTrue; import com.google.common.collect.ImmutableList; import java.io.IOException; import java.net.Socket; import java.net.UnknownHostException; import java.nio.charset.StandardCharsets; 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.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.TinyIntVector; 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.impl.UnionListWriter; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider; import org.apache.arrow.vector.ipc.ArrowStreamReader; import org.apache.arrow.vector.ipc.ArrowStreamWriter; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.ArrowType.Int; 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.types.pojo.Schema; import org.apache.arrow.vector.util.Text; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; public class EchoServerTest { private static EchoServer server; private static int serverPort; private static Thread serverThread; @BeforeAll public static void startEchoServer() throws IOException { server = new EchoServer(0); serverPort = server.port(); serverThread = new Thread() { @Override public void run() { try { server.run(); } catch (IOException e) { e.printStackTrace(); } } }; serverThread.start(); } @AfterAll public static void stopEchoServer() throws IOException, InterruptedException { server.close(); serverThread.join(); } private void testEchoServer(int serverPort, Field field, TinyIntVector vector, int batches) throws UnknownHostException, IOException { VectorSchemaRoot root = new VectorSchemaRoot(asList(field), asList((FieldVector) vector), 0); try (BufferAllocator alloc = new RootAllocator(Long.MAX_VALUE); Socket socket = new Socket("localhost", serverPort); ArrowStreamWriter writer = new ArrowStreamWriter(root, null, socket.getOutputStream()); ArrowStreamReader reader = new ArrowStreamReader(socket.getInputStream(), alloc)) { writer.start(); for (int i = 0; i < batches; i++) { vector.allocateNew(16); for (int j = 0; j < 8; j++) { vector.set(j, j + i); vector.set(j + 8, 0, (byte) (j + i)); } vector.setValueCount(16); root.setRowCount(16); writer.writeBatch(); } writer.end(); assertEquals(new Schema(asList(field)), reader.getVectorSchemaRoot().getSchema()); TinyIntVector readVector = (TinyIntVector) reader.getVectorSchemaRoot().getFieldVectors().get(0); for (int i = 0; i < batches; i++) { assertTrue(reader.loadNextBatch()); assertEquals(16, reader.getVectorSchemaRoot().getRowCount()); assertEquals(16, readVector.getValueCount()); for (int j = 0; j < 8; j++) { assertEquals(j + i, readVector.get(j)); assertTrue(readVector.isNull(j + 8)); } } assertFalse(reader.loadNextBatch()); assertEquals(0, reader.getVectorSchemaRoot().getRowCount()); assertEquals(reader.bytesRead(), writer.bytesWritten()); } } @Test public void basicTest() throws InterruptedException, IOException { BufferAllocator alloc = new RootAllocator(Long.MAX_VALUE); Field field = new Field( "testField", new FieldType(true, new ArrowType.Int(8, true), null, null), Collections.emptyList()); TinyIntVector vector = new TinyIntVector("testField", FieldType.nullable(TINYINT.getType()), alloc); // Try an empty stream, just the header. testEchoServer(serverPort, field, vector, 0); // Try with one batch. testEchoServer(serverPort, field, vector, 1); // Try with a few testEchoServer(serverPort, field, vector, 10); } @Test public void testFlatDictionary() throws IOException { DictionaryEncoding writeEncoding = new DictionaryEncoding(1L, false, null); try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); IntVector writeVector = new IntVector( "varchar", new FieldType(true, MinorType.INT.getType(), writeEncoding, null), allocator); VarCharVector writeDictionaryVector = new VarCharVector("dict", FieldType.nullable(VARCHAR.getType()), allocator)) { ValueVectorDataPopulator.setVector(writeVector, 0, 1, null, 2, 1, 2); ValueVectorDataPopulator.setVector( writeDictionaryVector, "foo".getBytes(StandardCharsets.UTF_8), "bar".getBytes(StandardCharsets.UTF_8), "baz".getBytes(StandardCharsets.UTF_8)); List fields = ImmutableList.of(writeVector.getField()); List vectors = ImmutableList.of((FieldVector) writeVector); VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, 6); DictionaryProvider writeProvider = new MapDictionaryProvider(new Dictionary(writeDictionaryVector, writeEncoding)); try (Socket socket = new Socket("localhost", serverPort); ArrowStreamWriter writer = new ArrowStreamWriter(root, writeProvider, socket.getOutputStream()); ArrowStreamReader reader = new ArrowStreamReader(socket.getInputStream(), allocator)) { writer.start(); writer.writeBatch(); writer.end(); reader.loadNextBatch(); VectorSchemaRoot readerRoot = reader.getVectorSchemaRoot(); assertEquals(6, readerRoot.getRowCount()); FieldVector readVector = readerRoot.getFieldVectors().get(0); assertNotNull(readVector); DictionaryEncoding readEncoding = readVector.getField().getDictionary(); assertNotNull(readEncoding); assertEquals(1L, readEncoding.getId()); assertEquals(6, readVector.getValueCount()); assertEquals(0, readVector.getObject(0)); assertEquals(1, readVector.getObject(1)); assertEquals(null, readVector.getObject(2)); assertEquals(2, readVector.getObject(3)); assertEquals(1, readVector.getObject(4)); assertEquals(2, readVector.getObject(5)); Dictionary dictionary = reader.lookup(1L); assertNotNull(dictionary); VarCharVector dictionaryVector = ((VarCharVector) dictionary.getVector()); assertEquals(3, dictionaryVector.getValueCount()); assertEquals(new Text("foo"), dictionaryVector.getObject(0)); assertEquals(new Text("bar"), dictionaryVector.getObject(1)); assertEquals(new Text("baz"), dictionaryVector.getObject(2)); } } } @Test public void testNestedDictionary() throws IOException { DictionaryEncoding writeEncoding = new DictionaryEncoding(2L, false, null); try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); VarCharVector writeDictionaryVector = new VarCharVector("dictionary", FieldType.nullable(VARCHAR.getType()), allocator); ListVector writeVector = ListVector.empty("list", allocator)) { // data being written: // [['foo', 'bar'], ['foo'], ['bar']] -> [[0, 1], [0], [1]] writeDictionaryVector.allocateNew(); writeDictionaryVector.set(0, "foo".getBytes(StandardCharsets.UTF_8)); writeDictionaryVector.set(1, "bar".getBytes(StandardCharsets.UTF_8)); writeDictionaryVector.setValueCount(2); writeVector.addOrGetVector(new FieldType(true, MinorType.INT.getType(), writeEncoding, null)); writeVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(writeVector); listWriter.startList(); listWriter.writeInt(0); listWriter.writeInt(1); listWriter.endList(); listWriter.startList(); listWriter.writeInt(0); listWriter.endList(); listWriter.startList(); listWriter.writeInt(1); listWriter.endList(); listWriter.setValueCount(3); List fields = ImmutableList.of(writeVector.getField()); List vectors = ImmutableList.of((FieldVector) writeVector); VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, 3); DictionaryProvider writeProvider = new MapDictionaryProvider(new Dictionary(writeDictionaryVector, writeEncoding)); try (Socket socket = new Socket("localhost", serverPort); ArrowStreamWriter writer = new ArrowStreamWriter(root, writeProvider, socket.getOutputStream()); ArrowStreamReader reader = new ArrowStreamReader(socket.getInputStream(), allocator)) { writer.start(); writer.writeBatch(); writer.end(); reader.loadNextBatch(); VectorSchemaRoot readerRoot = reader.getVectorSchemaRoot(); assertEquals(3, readerRoot.getRowCount()); ListVector readVector = (ListVector) readerRoot.getFieldVectors().get(0); assertNotNull(readVector); assertNull(readVector.getField().getDictionary()); DictionaryEncoding readEncoding = readVector.getField().getChildren().get(0).getDictionary(); assertNotNull(readEncoding); assertEquals(2L, readEncoding.getId()); Field nestedField = readVector.getField().getChildren().get(0); DictionaryEncoding encoding = nestedField.getDictionary(); assertNotNull(encoding); assertEquals(2L, encoding.getId()); assertEquals(new Int(32, true), encoding.getIndexType()); assertEquals(3, readVector.getValueCount()); assertEquals(Arrays.asList(0, 1), readVector.getObject(0)); assertEquals(Arrays.asList(0), readVector.getObject(1)); assertEquals(Arrays.asList(1), readVector.getObject(2)); Dictionary readDictionary = reader.lookup(2L); assertNotNull(readDictionary); VarCharVector dictionaryVector = ((VarCharVector) readDictionary.getVector()); assertEquals(2, dictionaryVector.getValueCount()); assertEquals(new Text("foo"), dictionaryVector.getObject(0)); assertEquals(new Text("bar"), dictionaryVector.getObject(1)); } } } } apache-arrow-java-19.0.0/tools/src/test/java/org/apache/arrow/tools/TestFileRoundtrip.java000066400000000000000000000074501515451615200315260ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.tools; import static org.apache.arrow.tools.ArrowFileTestFixtures.validateOutput; import static org.apache.arrow.tools.ArrowFileTestFixtures.writeInput; 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.io.File; import java.io.FileOutputStream; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.ipc.InvalidArrowFileException; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; public class TestFileRoundtrip { @TempDir public File testFolder; @TempDir public File testAnotherFolder; private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Integer.MAX_VALUE); } @AfterEach public void tearDown() { allocator.close(); } @Test public void test() throws Exception { File testInFile = new File(testFolder, "testIn.arrow"); File testOutFile = new File(testFolder, "testOut.arrow"); writeInput(testInFile, allocator); String[] args = {"-i", testInFile.getAbsolutePath(), "-o", testOutFile.getAbsolutePath()}; int result = new FileRoundtrip(System.err).run(args); assertEquals(0, result); validateOutput(testOutFile, allocator); } @Test public void testDiffFolder() throws Exception { File testInFile = new File(testFolder, "testIn.arrow"); File testOutFile = new File(testAnotherFolder, "testOut.arrow"); writeInput(testInFile, allocator); String[] args = {"-i", testInFile.getAbsolutePath(), "-o", testOutFile.getAbsolutePath()}; int result = new FileRoundtrip(System.err).run(args); assertEquals(0, result); validateOutput(testOutFile, allocator); } @Test public void testNotFoundInput() { File testInFile = new File(testFolder, "testIn.arrow"); File testOutFile = new File(testFolder, "testOut.arrow"); String[] args = {"-i", testInFile.getAbsolutePath(), "-o", testOutFile.getAbsolutePath()}; Exception exception = assertThrows( IllegalArgumentException.class, () -> { new FileRoundtrip(System.err).run(args); }); assertTrue(exception.getMessage().contains("input file not found")); } @Test public void testSmallSizeInput() throws Exception { File testInFile = new File(testFolder, "testIn.arrow"); File testOutFile = new File(testFolder, "testOut.arrow"); // create an empty file new FileOutputStream(testInFile).close(); String[] args = {"-i", testInFile.getAbsolutePath(), "-o", testOutFile.getAbsolutePath()}; Exception exception = assertThrows( InvalidArrowFileException.class, () -> { new FileRoundtrip(System.err).run(args); }); assertEquals("file too small: 0", exception.getMessage()); } } apache-arrow-java-19.0.0/tools/src/test/java/org/apache/arrow/tools/TestIntegration.java000066400000000000000000000314201515451615200312150ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.tools; import static org.apache.arrow.tools.ArrowFileTestFixtures.validateOutput; import static org.apache.arrow.tools.ArrowFileTestFixtures.validateVariadicOutput; import static org.apache.arrow.tools.ArrowFileTestFixtures.write; import static org.apache.arrow.tools.ArrowFileTestFixtures.writeData; import static org.apache.arrow.tools.ArrowFileTestFixtures.writeInput; import static org.apache.arrow.tools.ArrowFileTestFixtures.writeVariableWidthViewInput; import static org.junit.jupiter.api.Assertions.assertEquals; 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 com.fasterxml.jackson.core.util.DefaultPrettyPrinter; import com.fasterxml.jackson.core.util.DefaultPrettyPrinter.NopIndenter; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.SerializationFeature; import java.io.BufferedReader; import java.io.File; import java.io.IOException; import java.io.StringReader; import java.net.URL; import java.util.Map; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.tools.Integration.Command; import org.apache.arrow.vector.complex.NonNullableStructVector; import org.apache.arrow.vector.complex.impl.ComplexWriterImpl; import org.apache.arrow.vector.complex.writer.BaseWriter.ComplexWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter; import org.apache.arrow.vector.complex.writer.BigIntWriter; import org.apache.arrow.vector.complex.writer.Float8Writer; import org.apache.arrow.vector.complex.writer.IntWriter; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; public class TestIntegration { @TempDir public File testFolder; private BufferAllocator allocator; private ObjectMapper om = new ObjectMapper(); { DefaultPrettyPrinter prettyPrinter = new DefaultPrettyPrinter(); prettyPrinter.indentArraysWith(NopIndenter.instance); om.setDefaultPrettyPrinter(prettyPrinter); om.enable(SerializationFeature.INDENT_OUTPUT); om.enable(SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS); } static void writeInputFloat(File testInFile, BufferAllocator allocator, double... f) throws IOException { try (BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); NonNullableStructVector parent = NonNullableStructVector.empty("parent", vectorAllocator)) { ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); Float8Writer floatWriter = rootWriter.float8("float"); for (int i = 0; i < f.length; i++) { floatWriter.setPosition(i); floatWriter.writeFloat8(f[i]); } writer.setValueCount(f.length); write(parent.getChild("root"), testInFile); } } static void writeInput2(File testInFile, BufferAllocator allocator) throws IOException { int count = ArrowFileTestFixtures.COUNT; try (BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); NonNullableStructVector parent = NonNullableStructVector.empty("parent", vectorAllocator)) { writeData(count, parent); ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); IntWriter intWriter = rootWriter.integer("int"); BigIntWriter bigIntWriter = rootWriter.bigInt("bigInt"); intWriter.setPosition(5); intWriter.writeInt(999); bigIntWriter.setPosition(4); bigIntWriter.writeBigInt(777L); writer.setValueCount(count); write(parent.getChild("root"), testInFile); } } @BeforeEach public void init() { allocator = new RootAllocator(Integer.MAX_VALUE); } @AfterEach public void tearDown() { allocator.close(); } @Test public void testValid() throws Exception { File testInFile = new File(testFolder, "testIn.arrow"); File testJSONFile = new File(testFolder, "testOut.json"); testJSONFile.delete(); File testOutFile = new File(testFolder, "testOut.arrow"); testOutFile.delete(); // generate an arrow file writeInput(testInFile, allocator); Integration integration = new Integration(); // convert it to json String[] args1 = { "-arrow", testInFile.getAbsolutePath(), "-json", testJSONFile.getAbsolutePath(), "-command", Command.ARROW_TO_JSON.name() }; integration.run(args1); // convert back to arrow String[] args2 = { "-arrow", testOutFile.getAbsolutePath(), "-json", testJSONFile.getAbsolutePath(), "-command", Command.JSON_TO_ARROW.name() }; integration.run(args2); // check it is the same validateOutput(testOutFile, allocator); // validate arrow against json String[] args3 = { "-arrow", testInFile.getAbsolutePath(), "-json", testJSONFile.getAbsolutePath(), "-command", Command.VALIDATE.name() }; integration.run(args3); } @Test public void testJSONRoundTripWithVariableWidth() throws Exception { URL resource = getClass().getResource("/integration_json_simple.json"); assertNotNull(resource); File testJSONFile = new File(resource.getFile()).getCanonicalFile(); File testOutFile = new File(testFolder, "testOut.arrow"); File testRoundTripJSONFile = new File(testFolder, "testOut.json"); testOutFile.delete(); testRoundTripJSONFile.delete(); Integration integration = new Integration(); // convert to arrow String[] args1 = { "-arrow", testOutFile.getAbsolutePath(), "-json", testJSONFile.getAbsolutePath(), "-command", Command.JSON_TO_ARROW.name() }; integration.run(args1); // convert back to json String[] args2 = { "-arrow", testOutFile.getAbsolutePath(), "-json", testRoundTripJSONFile.getAbsolutePath(), "-command", Command.ARROW_TO_JSON.name() }; integration.run(args2); BufferedReader orig = readNormalized(testJSONFile); BufferedReader rt = readNormalized(testRoundTripJSONFile); String i; String o; int j = 0; while ((i = orig.readLine()) != null && (o = rt.readLine()) != null) { assertEquals(i, o, "line: " + j); ++j; } } @Test public void testJSONRoundTripWithStruct() throws Exception { URL resource = getClass().getResource("/integration_json_struct.json"); assertNotNull(resource); File testJSONFile = new File(resource.getFile()).getCanonicalFile(); File testOutFile = new File(testFolder, "testOutStruct.arrow"); File testRoundTripJSONFile = new File(testFolder, "testOutStruct.json"); testOutFile.delete(); testRoundTripJSONFile.delete(); Integration integration = new Integration(); // convert to arrow String[] args1 = { "-arrow", testOutFile.getAbsolutePath(), "-json", testJSONFile.getAbsolutePath(), "-command", Command.JSON_TO_ARROW.name() }; integration.run(args1); // convert back to json String[] args2 = { "-arrow", testOutFile.getAbsolutePath(), "-json", testRoundTripJSONFile.getAbsolutePath(), "-command", Command.ARROW_TO_JSON.name() }; integration.run(args2); BufferedReader orig = readNormalized(testJSONFile); BufferedReader rt = readNormalized(testRoundTripJSONFile); String i; String o; int j = 0; while ((i = orig.readLine()) != null && (o = rt.readLine()) != null) { assertEquals(i, o, "line: " + j); ++j; } } private BufferedReader readNormalized(File f) throws IOException { Map tree = om.readValue(f.getCanonicalFile(), Map.class); String normalized = om.writeValueAsString(tree); return new BufferedReader(new StringReader(normalized)); } /** The test should not be sensitive to small variations in float representation. */ @Test public void testFloat() throws Exception { File testValidInFile = new File(testFolder, "testValidFloatIn.arrow"); File testInvalidInFile = new File(testFolder, "testAlsoValidFloatIn.arrow"); File testJSONFile = new File(testFolder, "testValidOut.json"); testJSONFile.delete(); // generate an arrow file writeInputFloat(testValidInFile, allocator, 912.4140000000002, 912.414); // generate a different arrow file writeInputFloat(testInvalidInFile, allocator, 912.414, 912.4140000000002); Integration integration = new Integration(); // convert the "valid" file to json String[] args1 = { "-arrow", testValidInFile.getAbsolutePath(), "-json", testJSONFile.getAbsolutePath(), "-command", Command.ARROW_TO_JSON.name() }; integration.run(args1); // compare the "invalid" file to the "valid" json String[] args3 = { "-arrow", testInvalidInFile.getAbsolutePath(), "-json", testJSONFile.getAbsolutePath(), "-command", Command.VALIDATE.name() }; // this should fail integration.run(args3); } @Test public void testInvalid() throws Exception { File testValidInFile = new File(testFolder, "testValidIn.arrow"); File testInvalidInFile = new File(testFolder, "testInvalidIn.arrow"); File testJSONFile = new File(testFolder, "testInvalidOut.json"); testJSONFile.delete(); // generate an arrow file writeInput(testValidInFile, allocator); // generate a different arrow file writeInput2(testInvalidInFile, allocator); Integration integration = new Integration(); // convert the "valid" file to json String[] args1 = { "-arrow", testValidInFile.getAbsolutePath(), "-json", testJSONFile.getAbsolutePath(), "-command", Command.ARROW_TO_JSON.name() }; integration.run(args1); // compare the "invalid" file to the "valid" json String[] args3 = { "-arrow", testInvalidInFile.getAbsolutePath(), "-json", testJSONFile.getAbsolutePath(), "-command", Command.VALIDATE.name() }; // this should fail IllegalArgumentException e = assertThrows( IllegalArgumentException.class, () -> { integration.run(args3); }); assertTrue(e.getMessage().contains("Different values in column"), e.getMessage()); assertTrue(e.getMessage().contains("999"), e.getMessage()); } @Test public void testValidateVariableWidthView() throws Exception { final int valueCount = 256; final int multiplier = 6; for (int i = 1; i < multiplier; i++) { File testInFile = new File(testFolder, "testIn.arrow"); File testJSONFile = new File(testFolder, "testOut.json"); testJSONFile.delete(); File testOutFile = new File(testFolder, "testOut.arrow"); testOutFile.delete(); writeVariableWidthViewInput(testInFile, allocator, multiplier * valueCount); Integration integration = new Integration(); // convert it to json String[] args1 = { "-arrow", testInFile.getAbsolutePath(), "-json", testJSONFile.getAbsolutePath(), "-command", Command.ARROW_TO_JSON.name() }; integration.run(args1); // convert back to arrow String[] args2 = { "-arrow", testOutFile.getAbsolutePath(), "-json", testJSONFile.getAbsolutePath(), "-command", Command.JSON_TO_ARROW.name() }; integration.run(args2); // check it is the same validateVariadicOutput(testOutFile, allocator, multiplier * valueCount); // validate arrow against json String[] args3 = { "-arrow", testInFile.getAbsolutePath(), "-json", testJSONFile.getAbsolutePath(), "-command", Command.VALIDATE.name() }; integration.run(args3); } } } apache-arrow-java-19.0.0/tools/src/test/resources/000077500000000000000000000000001515451615200220165ustar00rootroot00000000000000apache-arrow-java-19.0.0/tools/src/test/resources/integration_json_simple.json000066400000000000000000000041121515451615200276340ustar00rootroot00000000000000{ "schema": { "fields": [ { "name": "foo", "type": {"name": "int", "isSigned": true, "bitWidth": 32}, "nullable": true, "children": [] }, { "name": "bar", "type": {"name": "floatingpoint", "precision": "DOUBLE"}, "nullable": true, "children": [] }, { "name": "baz", "type": {"name": "utf8"}, "nullable": true, "children": [] } ] }, "batches": [ { "count": 5, "columns": [ { "name": "foo", "count": 5, "VALIDITY": [1, 0, 1, 1, 1], "DATA": [1, 2, 3, 4, 5] }, { "name": "bar", "count": 5, "VALIDITY": [1, 0, 0, 1, 1], "DATA": [1.0, 2.0, 3.0, 4.0, 5.0] }, { "name": "baz", "count": 5, "VALIDITY": [1, 0, 0, 1, 1], "OFFSET": [0, 2, 2, 2, 5, 9], "DATA": ["aa", "", "", "bbb", "cccc"] } ] }, { "count": 5, "columns": [ { "name": "foo", "count": 5, "VALIDITY": [1, 1, 1, 1, 1], "DATA": [1, 2, 3, 4, 5] }, { "name": "bar", "count": 5, "VALIDITY": [1, 1, 1, 1, 1], "DATA": [1.0, 2.0, 3.0, 4.0, 5.0] }, { "name": "baz", "count": 5, "VALIDITY": [1, 1, 1, 1, 1], "OFFSET": [0, 2, 3, 4, 7, 11], "DATA": ["aa", "b", "c", "ddd", "eeee"] } ] }, { "count": 5, "columns": [ { "name": "foo", "count": 5, "VALIDITY": [0, 0, 0, 0, 0], "DATA": [1, 2, 3, 4, 5] }, { "name": "bar", "count": 5, "VALIDITY": [0, 0, 0, 0, 0], "DATA": [1.0, 2.0, 3.0, 4.0, 5.0] }, { "name": "baz", "count": 5, "VALIDITY": [0, 0, 0, 0, 0], "OFFSET": [0, 0, 0, 0, 0, 0], "DATA": ["", "", "", "", ""] } ] } ] } apache-arrow-java-19.0.0/tools/src/test/resources/integration_json_struct.json000066400000000000000000000075761515451615200277100ustar00rootroot00000000000000{ "schema": { "fields": [ { "name": "struct_nullable", "type": { "name": "struct" }, "nullable": true, "children": [ { "name": "f1", "type": { "name": "int", "isSigned": true, "bitWidth": 32 }, "nullable": true, "children": [] }, { "name": "f2", "type": { "name": "utf8" }, "nullable": true, "children": [] } ] } ] }, "batches": [ { "count": 7, "columns": [ { "name": "struct_nullable", "count": 7, "VALIDITY": [ 0, 1, 1, 1, 0, 1, 0 ], "children": [ { "name": "f1", "count": 7, "VALIDITY": [ 1, 0, 1, 1, 1, 0, 0 ], "DATA": [ 1402032511, 290876774, 137773603, 410361374, 1959836418, 1995074679, -163525262 ] }, { "name": "f2", "count": 7, "VALIDITY": [ 0, 1, 1, 1, 0, 1, 0 ], "OFFSET": [ 0, 0, 7, 14, 21, 21, 28, 28 ], "DATA": [ "", "MhRNxD4", "3F9HBxK", "aVd88fp", "", "3loZrRf", "" ] } ] } ] }, { "count": 10, "columns": [ { "name": "struct_nullable", "count": 10, "VALIDITY": [ 0, 1, 1, 0, 1, 0, 0, 1, 1, 1 ], "children": [ { "name": "f1", "count": 10, "VALIDITY": [ 0, 0, 0, 0, 0, 0, 1, 0, 0, 0 ], "DATA": [ -2041500147, 1715692943, -35444996, 1425496657, 112765084, 1760754983, 413888857, 2039738337, -1924327700, 670528518 ] }, { "name": "f2", "count": 10, "VALIDITY": [ 1, 0, 0, 1, 1, 1, 1, 1, 1, 0 ], "OFFSET": [ 0, 7, 7, 7, 14, 21, 28, 35, 42, 49, 49 ], "DATA": [ "AS5oARE", "", "", "JGdagcX", "78SLiRw", "vbGf7OY", "5uh5fTs", "0ilsf82", "LjS9MbU", "" ] } ] } ] } ] } apache-arrow-java-19.0.0/tools/src/test/resources/logback.xml000066400000000000000000000025211515451615200241420ustar00rootroot00000000000000 %d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n apache-arrow-java-19.0.0/vector/000077500000000000000000000000001515451615200164005ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/pom.xml000066400000000000000000000230551515451615200177220ustar00rootroot00000000000000 4.0.0 org.apache.arrow arrow-java-root 19.0.0 arrow-vector Arrow Vectors An off-heap reference implementation for Arrow columnar data format. org.apache.arrow arrow-format org.apache.arrow arrow-memory-core org.immutables value-annotations com.fasterxml.jackson.core jackson-core com.fasterxml.jackson.core jackson-annotations com.fasterxml.jackson.core jackson-databind com.fasterxml.jackson.datatype jackson-datatype-jsr310 commons-codec commons-codec 1.21.0 org.apache.arrow arrow-memory-netty test org.apache.arrow arrow-memory-unsafe test com.google.flatbuffers flatbuffers-java ${dep.fbs.version} org.slf4j slf4j-api org.hamcrest hamcrest test org.apache.maven.plugins maven-surefire-plugin default-test test org.apache.arrow:arrow-memory-unsafe run-unsafe test test org.apache.arrow:arrow-memory-netty netty org.apache.maven.plugins maven-compiler-plugin -Werror org.apache.drill.tools drill-fmpp-maven-plugin generate-fmpp generate generate-sources src/main/codegen/config.fmpp ${project.build.directory}/generated-sources/fmpp src/main/codegen/templates org.apache.maven.plugins maven-shade-plugin shade package org.apache.arrow:arrow-format com.google.flatbuffers:* false true shade-format-flatbuffers com.google.flatbuffers 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 integration-tests org.apache.maven.plugins maven-failsafe-plugin 3600 false default-it integration-test verify apache-arrow-java-19.0.0/vector/src/000077500000000000000000000000001515451615200171675ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/000077500000000000000000000000001515451615200201135ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/codegen/000077500000000000000000000000001515451615200215175ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/codegen/config.fmpp000066400000000000000000000017761515451615200236630ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. data: { # TODO: Rename to ~valueVectorModesAndTypes for clarity. vv: tdd(../data/ValueVectorTypes.tdd), arrowTypes: tdd(../data/ArrowTypes.tdd) } freemarkerLinks: { includes: includes/ } apache-arrow-java-19.0.0/vector/src/main/codegen/data/000077500000000000000000000000001515451615200224305ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/codegen/data/ArrowTypes.tdd000066400000000000000000000066221515451615200252520ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. { types: [ { name: "Null", fields: [], complex: false }, { name: "Struct_", fields: [], complex: true }, { name: "List", fields: [], complex: true }, { name: "LargeList", fields: [], complex: true }, { name: "FixedSizeList", fields: [{name: "listSize", type: int}], complex: true }, { name: "Union", fields: [{name: "mode", type: short, valueType: UnionMode}, {name: "typeIds", type: "int[]"}], complex: true }, { name: "Map", fields: [{name: "keysSorted", type: boolean}], complex: true }, { name: "Int", fields: [{name: "bitWidth", type: int}, {name: "isSigned", type: boolean}], complex: false }, { name: "FloatingPoint", fields: [{name: precision, type: short, valueType: FloatingPointPrecision}], complex: false }, { name: "Utf8", fields: [], complex: false }, { name: "Utf8View", fields: [], complex: false }, { name: "LargeUtf8", fields: [], complex: false }, { name: "Binary", fields: [], complex: false }, { name: "BinaryView", fields: [], complex: false }, { name: "LargeBinary", fields: [], complex: false }, { name: "FixedSizeBinary", fields: [{name: "byteWidth", type: int}], complex: false } { name: "Bool", fields: [], complex: false }, { name: "Decimal", fields: [{name: "precision", type: int}, {name: "scale", type: int}, {name: "bitWidth", type: int}], complex: false }, { name: "Date", fields: [{name: "unit", type: short, valueType: DateUnit}] complex: false }, { name: "Time", fields: [{name: "unit", type: short, valueType: TimeUnit}, {name: "bitWidth", type: int}], complex: false }, { name: "Timestamp", fields: [{name: "unit", type: short, valueType: TimeUnit}, {name: "timezone", type: String}] complex: false }, { name: "Interval", fields: [{name: "unit", type: short, valueType: IntervalUnit}], complex: false }, { name: "Duration", fields: [{name: "unit", type: short, valueType: TimeUnit}], complex: false }, { name: "ListView", fields: [], complex: true }, { name: "LargeListView", fields: [], complex: true }, { name: "RunEndEncoded", fields: [], complex: true } ] } apache-arrow-java-19.0.0/vector/src/main/codegen/data/ValueVectorTypes.tdd000066400000000000000000000204061515451615200264130ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. { modes: [ {name: "Optional", prefix: "Nullable"}, {name: "Required", prefix: ""} ], types: [ { major: "Fixed", width: 1, javaType: "byte", boxedType: "Byte", fields: [{name: "value", type: "byte"}], minor: [ { class: "TinyInt", valueHolder: "IntHolder" }, { class: "UInt1", valueHolder: "UInt1Holder" } ] }, { major: "Fixed", width: 2, javaType: "char", boxedType: "Character", fields: [{name: "value", type: "char"}], minor: [ { class: "UInt2", valueHolder: "UInt2Holder"} ] }, { major: "Fixed", width: 2, javaType: "short", boxedType: "Short", fields: [{name: "value", type: "short"}], minor: [ { class: "SmallInt", valueHolder: "Int2Holder"}, ] }, { major: "Fixed", width: 2, javaType: "short", boxedType: "Short", fields: [{name: "value", type: "short"}], minor: [ { class: "Float2", valueHolder: "Int2Holder"}, ] }, { major: "Fixed", width: 4, javaType: "int", boxedType: "Integer", fields: [{name: "value", type: "int"}], minor: [ { class: "Int", valueHolder: "IntHolder"}, { class: "UInt4", valueHolder: "UInt4Holder" }, { class: "Float4", javaType: "float" , boxedType: "Float", fields: [{name: "value", type: "float"}]}, { class: "DateDay" }, { class: "IntervalYear", javaType: "int", friendlyType: "Period" }, { class: "TimeSec" }, { class: "TimeMilli", javaType: "int", friendlyType: "LocalDateTime" } ] }, { major: "Fixed", width: 8, javaType: "long", boxedType: "Long", fields: [{name: "value", type: "long"}], minor: [ { class: "BigInt"}, { class: "UInt8" }, { class: "Float8", javaType: "double", boxedType: "Double", fields: [{name: "value", type: "double"}] }, { class: "DateMilli", javaType: "long", friendlyType: "LocalDateTime" }, { class: "Duration", javaType: "long", friendlyType: "Duration", arrowType: "org.apache.arrow.vector.types.pojo.ArrowType.Duration", typeParams: [ {name: "unit", type: "org.apache.arrow.vector.types.TimeUnit"} ], arrowTypeConstructorParams: ["unit"]} { class: "TimeStampSec", javaType: "long", boxedType: "Long", friendlyType: "LocalDateTime" }, { class: "TimeStampMilli", javaType: "long", boxedType: "Long", friendlyType: "LocalDateTime" }, { class: "TimeStampMicro", javaType: "long", boxedType: "Long", friendlyType: "LocalDateTime" }, { class: "TimeStampNano", javaType: "long", boxedType: "Long", friendlyType: "LocalDateTime" }, { class: "TimeStampSecTZ", javaType: "long", boxedType: "Long", typeParams: [ {name: "timezone", type: "String"} ], arrowType: "org.apache.arrow.vector.types.pojo.ArrowType.Timestamp", arrowTypeConstructorParams: ["org.apache.arrow.vector.types.TimeUnit.SECOND", "timezone"] }, { class: "TimeStampMilliTZ", javaType: "long", boxedType: "Long", typeParams: [ {name: "timezone", type: "String"} ], arrowType: "org.apache.arrow.vector.types.pojo.ArrowType.Timestamp", arrowTypeConstructorParams: ["org.apache.arrow.vector.types.TimeUnit.MILLISECOND", "timezone"] }, { class: "TimeStampMicroTZ", javaType: "long", boxedType: "Long", typeParams: [ {name: "timezone", type: "String"} ], arrowType: "org.apache.arrow.vector.types.pojo.ArrowType.Timestamp", arrowTypeConstructorParams: ["org.apache.arrow.vector.types.TimeUnit.MICROSECOND", "timezone"] }, { class: "TimeStampNanoTZ", javaType: "long", boxedType: "Long", typeParams: [ {name: "timezone", type: "String"} ], arrowType: "org.apache.arrow.vector.types.pojo.ArrowType.Timestamp", arrowTypeConstructorParams: ["org.apache.arrow.vector.types.TimeUnit.NANOSECOND", "timezone"] }, { class: "TimeMicro" }, { class: "TimeNano" } ] }, { major: "Fixed", width: 8, javaType: "ArrowBuf", boxedType: "ArrowBuf", minor: [ { class: "IntervalDay", millisecondsOffset: 4, friendlyType: "Duration", fields: [ {name: "days", type:"int"}, {name: "milliseconds", type:"int"}] } ] }, { major: "Fixed", width: 16, javaType: "ArrowBuf", boxedType: "ArrowBuf", minor: [ { class: "IntervalMonthDayNano", daysOffset: 4, nanosecondsOffset: 8, friendlyType: "PeriodDuration", fields: [ {name: "months", type:"int"}, {name: "days", type:"int"}, {name: "nanoseconds", type:"long"}] } ] }, { major: "Fixed", width: 32, javaType: "ArrowBuf", boxedType: "ArrowBuf", minor: [ { class: "Decimal256", maxPrecisionDigits: 76, nDecimalDigits: 4, friendlyType: "BigDecimal", typeParams: [ {name: "scale", type: "int"}, { name: "precision", type: "int"}], arrowType: "org.apache.arrow.vector.types.pojo.ArrowType.Decimal", fields: [{name: "start", type: "long"}, {name: "buffer", type: "ArrowBuf"}] } ] }, { major: "Fixed", width: 16, javaType: "ArrowBuf", boxedType: "ArrowBuf", minor: [ { class: "Decimal", maxPrecisionDigits: 38, nDecimalDigits: 4, friendlyType: "BigDecimal", typeParams: [ {name: "scale", type: "int"}, { name: "precision", type: "int"}], arrowType: "org.apache.arrow.vector.types.pojo.ArrowType.Decimal", fields: [{name: "start", type: "long"}, {name: "buffer", type: "ArrowBuf"}] } ] }, { major: "Fixed", width: -1, javaType: "byte[]", boxedType: "ArrowBuf", minor: [ { class: "FixedSizeBinary", typeParams: [ {name: "byteWidth", type: "int"} ], arrowType: "org.apache.arrow.vector.types.pojo.ArrowType.FixedSizeBinary", friendlyType: "byte[]", fields: [{name: "buffer", type: "ArrowBuf"}], } ] }, { major: "VarLen", width: 4, javaType: "int", boxedType: "ArrowBuf", fields: [{name: "start", type: "int"}, {name: "end", type: "int"}, {name: "buffer", type: "ArrowBuf"}], minor: [ { class: "VarBinary" , friendlyType: "byte[]" }, { class: "VarChar" , friendlyType: "Text" }, { class: "ViewVarBinary" , friendlyType: "byte[]" }, { class: "ViewVarChar" , friendlyType: "Text" } ] }, { major: "VarLen", width: 8, javaType: "long", boxedType: "ArrowBuf", fields: [{name: "start", type: "long"}, {name: "end", type: "long"}, {name: "buffer", type: "ArrowBuf"}], minor: [ { class: "LargeVarChar" , friendlyType: "Text" } { class: "LargeVarBinary" , friendlyType: "byte[]" } ] }, { major: "Bit", width: 1, javaType: "int", boxedType: "Integer", minor: [ { class: "Bit" , friendlyType: "Boolean", fields: [{name: "value", type: "int"}] } ] } ] } apache-arrow-java-19.0.0/vector/src/main/codegen/includes/000077500000000000000000000000001515451615200233255ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/codegen/includes/license.ftl000066400000000000000000000014411515451615200254560ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ apache-arrow-java-19.0.0/vector/src/main/codegen/includes/vv_imports.ftl000066400000000000000000000044111515451615200262440ustar00rootroot00000000000000<#-- ~ Licensed to the Apache Software Foundation (ASF) under one or more ~ contributor license agreements. See the NOTICE file distributed with ~ this work for additional information regarding copyright ownership. ~ The ASF licenses this file to You under the Apache License, Version 2.0 ~ (the "License"); you may not use this file except in compliance with ~ the License. You may obtain a copy of the License at ~ ~ http://www.apache.org/licenses/LICENSE-2.0 ~ ~ Unless required by applicable law or agreed to in writing, software ~ distributed under the License is distributed on an "AS IS" BASIS, ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. ~ See the License for the specific language governing permissions and ~ limitations under the License. --> import static org.apache.arrow.util.Preconditions.checkArgument; import static org.apache.arrow.util.Preconditions.checkState; import com.google.flatbuffers.FlatBufferBuilder; import org.apache.arrow.memory.*; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.Types.*; import org.apache.arrow.vector.types.pojo.*; import org.apache.arrow.vector.types.pojo.ArrowType.*; import org.apache.arrow.vector.types.*; import org.apache.arrow.vector.*; import org.apache.arrow.vector.holders.*; import org.apache.arrow.vector.util.*; 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; import org.apache.arrow.vector.util.JsonStringArrayList; import java.util.Arrays; import java.util.Random; import java.util.List; import java.io.Closeable; import java.io.InputStream; import java.io.InputStreamReader; import java.nio.ByteBuffer; import java.math.BigDecimal; import java.math.BigInteger; import java.time.Duration; import java.time.LocalDateTime; import java.time.Period; import java.time.ZonedDateTime; apache-arrow-java-19.0.0/vector/src/main/codegen/templates/000077500000000000000000000000001515451615200235155ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/codegen/templates/AbstractFieldReader.java000066400000000000000000000077261515451615200302260ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/AbstractFieldReader.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; <#include "/@includes/vv_imports.ftl" /> /** * Source code generated using FreeMarker template ${.template_name} */ @SuppressWarnings("unused") public abstract class AbstractFieldReader extends AbstractBaseReader implements FieldReader{ protected AbstractFieldReader(){ super(); } /** * Returns true if the current value of the reader is not null * @return whether the current value is set */ public boolean isSet() { return true; } @Override public Field getField() { fail("getField"); return null; } <#list ["Object", "BigDecimal", "Short", "Integer", "Long", "Boolean", "LocalDateTime", "Duration", "Period", "Double", "Float", "Character", "Text", "String", "Byte", "byte[]", "PeriodDuration"] as friendlyType> <#assign safeType=friendlyType /> <#if safeType=="byte[]"><#assign safeType="ByteArray" /> public ${friendlyType} read${safeType}(int arrayIndex) { fail("read${safeType}(int arrayIndex)"); return null; } public ${friendlyType} read${safeType}() { fail("read${safeType}()"); return null; } public void copyAsValue(StructWriter writer) { fail("CopyAsValue StructWriter"); } public void copyAsField(String name, StructWriter writer) { fail("CopyAsField StructWriter"); } public void copyAsField(String name, ListWriter writer) { fail("CopyAsFieldList"); } public void copyAsField(String name, MapWriter writer) { fail("CopyAsFieldMap"); } <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> <#assign boxedType = (minor.boxedType!type.boxedType) /> public void read(${name}Holder holder) { fail("${name}"); } public void read(Nullable${name}Holder holder) { fail("${name}"); } public void read(int arrayIndex, ${name}Holder holder) { fail("Repeated${name}"); } public void read(int arrayIndex, Nullable${name}Holder holder) { fail("Repeated${name}"); } public void copyAsValue(${name}Writer writer) { fail("CopyAsValue${name}"); } public void copyAsField(String name, ${name}Writer writer) { fail("CopyAsField${name}"); } 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; } public FieldReader reader() { fail("reader()"); return null; } public int size() { fail("size()"); return -1; } 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())); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/AbstractFieldWriter.java000066400000000000000000000215241515451615200302700ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/AbstractFieldWriter.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; <#include "/@includes/vv_imports.ftl" /> /* * This class is generated using freemarker and the ${.template_name} template. * Note that changes to the AbstractFieldWriter template should also get reflected in the * AbstractPromotableFieldWriter, ComplexWriters, UnionFixedSizeListWriter, UnionListWriter * and UnionWriter templates and the PromotableWriter concrete code. */ @SuppressWarnings("unused") abstract class AbstractFieldWriter extends AbstractBaseWriter implements FieldWriter { protected boolean addVectorAsNullable = true; /** * Set flag to control the FieldType.nullable property when a writer creates a new vector. * If true then vectors created will be nullable, this is the default behavior. If false then * vectors created will be non-nullable. * * @param nullable Whether or not to create nullable vectors (default behavior is true) */ public void setAddVectorAsNullable(boolean nullable) { addVectorAsNullable = nullable; } @Override public void start() { throw new IllegalStateException(String.format("You tried to start when you are using a ValueWriter of type %s.", this.getClass().getSimpleName())); } @Override public void end() { throw new IllegalStateException(String.format("You tried to end when you are using a ValueWriter of type %s.", this.getClass().getSimpleName())); } @Override public void startList() { throw new IllegalStateException(String.format("You tried to start a list when you are using a ValueWriter of type %s.", this.getClass().getSimpleName())); } @Override public void endList() { throw new IllegalStateException(String.format("You tried to end a list when you are using a ValueWriter of type %s.", this.getClass().getSimpleName())); } @Override public void startListView() { throw new IllegalStateException(String.format("You tried to start a list view when you are using a ValueWriter of type %s.", this.getClass().getSimpleName())); } @Override public void endListView() { throw new IllegalStateException(String.format("You tried to end a list view when you are using a ValueWriter of type %s.", this.getClass().getSimpleName())); } @Override public void startMap() { throw new IllegalStateException(String.format("You tried to start a map when you are using a ValueWriter of type %s.", this.getClass().getSimpleName())); } @Override public void endMap() { throw new IllegalStateException(String.format("You tried to end a map when you are using a ValueWriter of type %s.", this.getClass().getSimpleName())); } @Override public void startEntry() { throw new IllegalStateException(String.format("You tried to start a map entry when you are using a ValueWriter of type %s.", this.getClass().getSimpleName())); } @Override public MapWriter key() { throw new IllegalStateException(String.format("You tried to start a map key when you are using a ValueWriter of type %s.", this.getClass().getSimpleName())); } @Override public MapWriter value() { throw new IllegalStateException(String.format("You tried to start a map value when you are using a ValueWriter of type %s.", this.getClass().getSimpleName())); } @Override 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) /> @Override public void write(${name}Holder holder) { fail("${name}"); } public void write${minor.class}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, ) { fail("${name}"); } <#if minor.class?starts_with("Decimal")> public void write${minor.class}(${friendlyType} value) { fail("${name}"); } public void write${minor.class}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, , ArrowType arrowType) { fail("${name}"); } public void writeBigEndianBytesTo${minor.class}(byte[] value) { fail("${name}"); } public void writeBigEndianBytesTo${minor.class}(byte[] value, ArrowType arrowType) { fail("${name}"); } <#if minor.class?ends_with("VarBinary")> public void write${minor.class}(byte[] value) { fail("${name}"); } public void write${minor.class}(byte[] value, int offset, int length) { fail("${name}"); } public void write${minor.class}(ByteBuffer value) { fail("${name}"); } public void write${minor.class}(ByteBuffer value, int offset, int length) { fail("${name}"); } <#if minor.class?ends_with("VarChar")> public void write${minor.class}(${friendlyType} value) { fail("${name}"); } public void write${minor.class}(String value) { fail("${name}"); } public void writeNull() { fail("${name}"); } /** * This implementation returns {@code false}. *

    * Must be overridden by struct writers. *

    */ @Override public boolean isEmptyStruct() { return false; } @Override public StructWriter struct() { fail("Struct"); return null; } @Override public ListWriter list() { fail("List"); return null; } @Override public ListWriter listView() { fail("ListView"); return null; } @Override public MapWriter map() { fail("Map"); return null; } @Override public StructWriter struct(String name) { fail("Struct"); return null; } @Override public ListWriter list(String name) { fail("List"); return null; } @Override public ListWriter listView(String name) { fail("ListView"); return null; } @Override public MapWriter map(String name) { fail("Map"); return null; } @Override public MapWriter map(boolean keysSorted) { fail("Map"); return null; } @Override 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" /> <#assign upperName = minor.class?upper_case /> <#assign capName = minor.class?cap_first /> <#if minor.typeParams?? > @Override public ${capName}Writer ${lowerName}(String name<#list minor.typeParams as typeParam>, ${typeParam.type} ${typeParam.name}) { fail("${capName}(" + <#list minor.typeParams as typeParam>"${typeParam.name}: " + ${typeParam.name} + ", " + ")"); return null; } @Override public ${capName}Writer ${lowerName}(String name) { fail("${capName}"); return null; } @Override public ${capName}Writer ${lowerName}() { fail("${capName}"); return null; } public void copyReader(FieldReader reader) { fail("Copy FieldReader"); } public void copyReaderToField(String name, FieldReader reader) { fail("Copy FieldReader to STring"); } private void fail(String name) { throw new IllegalArgumentException(String.format("You tried to write a %s type when you are using a ValueWriter of type %s.", name, this.getClass().getSimpleName())); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java000066400000000000000000000267601515451615200323240ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/AbstractPromotableFieldWriter.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; <#include "/@includes/vv_imports.ftl" /> <#function is_timestamp_tz type> <#return type?starts_with("TimeStamp") && type?ends_with("TZ")> /* * A FieldWriter which delegates calls to another FieldWriter. The delegate FieldWriter can be promoted to a new type * when necessary. Classes that extend this class are responsible for handling promotion. * * This class is generated using freemarker and the ${.template_name} template. * */ @SuppressWarnings("unused") abstract class AbstractPromotableFieldWriter extends AbstractFieldWriter { /** * Retrieve the FieldWriter, promoting if it is not a FieldWriter of the specified type * @param type the type of the values we want to write * @return the corresponding field writer */ protected FieldWriter getWriter(MinorType type) { return getWriter(type, null); } abstract protected FieldWriter getWriter(MinorType type, ArrowType arrowType); /** * @return the current FieldWriter */ abstract protected FieldWriter getWriter(); @Override public void start() { getWriter(MinorType.STRUCT).start(); } @Override public void end() { getWriter(MinorType.STRUCT).end(); setPosition(idx() + 1); } @Override public void startList() { getWriter(MinorType.LIST).startList(); } @Override public void endList() { getWriter(MinorType.LIST).endList(); setPosition(idx() + 1); } @Override public void startListView() { getWriter(MinorType.LISTVIEW).startListView(); } @Override public void endListView() { getWriter(MinorType.LISTVIEW).endListView(); setPosition(idx() + 1); } @Override public void startMap() { getWriter(MinorType.MAP).startMap(); } @Override public void endMap() { getWriter(MinorType.MAP).endMap(); setPosition(idx() + 1); } @Override public void startEntry() { getWriter(MinorType.MAP).startEntry(); } @Override public MapWriter key() { return getWriter(MinorType.MAP).key(); } @Override public MapWriter value() { return getWriter(MinorType.MAP).value(); } @Override public void endEntry() { getWriter(MinorType.MAP).endEntry(); } <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> <#assign fields = minor.fields!type.fields /> <#if minor.class == "Decimal"> @Override public void write(DecimalHolder holder) { getWriter(MinorType.DECIMAL).write(holder); } public void writeDecimal(int start, ArrowBuf buffer, ArrowType arrowType) { getWriter(MinorType.DECIMAL).writeDecimal(start, buffer, arrowType); } public void writeDecimal(int start, ArrowBuf buffer) { getWriter(MinorType.DECIMAL).writeDecimal(start, buffer); } public void writeBigEndianBytesToDecimal(byte[] value, ArrowType arrowType) { getWriter(MinorType.DECIMAL).writeBigEndianBytesToDecimal(value, arrowType); } public void writeBigEndianBytesToDecimal(byte[] value) { getWriter(MinorType.DECIMAL).writeBigEndianBytesToDecimal(value); } <#elseif minor.class == "Decimal256"> @Override public void write(Decimal256Holder holder) { getWriter(MinorType.DECIMAL256).write(holder); } public void writeDecimal256(long start, ArrowBuf buffer, ArrowType arrowType) { getWriter(MinorType.DECIMAL256).writeDecimal256(start, buffer, arrowType); } public void writeDecimal256(long start, ArrowBuf buffer) { getWriter(MinorType.DECIMAL256).writeDecimal256(start, buffer); } public void writeBigEndianBytesToDecimal256(byte[] value, ArrowType arrowType) { getWriter(MinorType.DECIMAL256).writeBigEndianBytesToDecimal256(value, arrowType); } public void writeBigEndianBytesToDecimal256(byte[] value) { getWriter(MinorType.DECIMAL256).writeBigEndianBytesToDecimal256(value); } <#elseif is_timestamp_tz(minor.class)> @Override public void write(${name}Holder holder) { ArrowType.Timestamp arrowTypeWithoutTz = (ArrowType.Timestamp) MinorType.${name?upper_case?remove_ending("TZ")}.getType(); // Take the holder.timezone similar to how PromotableWriter.java:write(DecimalHolder) takes the scale from the holder. ArrowType.Timestamp arrowType = new ArrowType.Timestamp(arrowTypeWithoutTz.getUnit(), holder.timezone); getWriter(MinorType.${name?upper_case}, arrowType).write(holder); } /** * @deprecated * The holder version should be used instead otherwise the timezone will default to UTC. * @see #write(${name}Holder) */ @Deprecated @Override public void write${minor.class}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, ) { ArrowType.Timestamp arrowTypeWithoutTz = (ArrowType.Timestamp) MinorType.${name?upper_case?remove_ending("TZ")}.getType(); // Assumes UTC if no timezone is provided ArrowType.Timestamp arrowType = new ArrowType.Timestamp(arrowTypeWithoutTz.getUnit(), "UTC"); getWriter(MinorType.${name?upper_case}, arrowType).write${minor.class}(<#list fields as field>${field.name}<#if field_has_next>, ); } <#elseif minor.class == "Duration"> @Override public void write(${name}Holder holder) { ArrowType.Duration arrowType = new ArrowType.Duration(holder.unit); getWriter(MinorType.${name?upper_case}, arrowType).write(holder); } /** * @deprecated * If you experience errors with using this version of the method, switch to the holder version. * The errors occur when using an untyped or unioned PromotableWriter, because this version of the * method does not have enough information to infer the ArrowType. * @see #write(${name}Holder) */ @Deprecated @Override public void write${minor.class}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, ) { getWriter(MinorType.${name?upper_case}).write${minor.class}(<#list fields as field>${field.name}<#if field_has_next>, ); } <#elseif minor.class == "FixedSizeBinary"> @Override public void write(${name}Holder holder) { ArrowType.FixedSizeBinary arrowType = new ArrowType.FixedSizeBinary(holder.byteWidth); getWriter(MinorType.${name?upper_case}, arrowType).write(holder); } /** * @deprecated * If you experience errors with using this version of the method, switch to the holder version. * The errors occur when using an untyped or unioned PromotableWriter, because this version of the * method does not have enough information to infer the ArrowType. * @see #write(${name}Holder) */ @Deprecated @Override public void write${minor.class}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, ) { getWriter(MinorType.${name?upper_case}).write${minor.class}(<#list fields as field>${field.name}<#if field_has_next>, ); } <#else> @Override public void write(${name}Holder holder) { getWriter(MinorType.${name?upper_case}).write(holder); } public void write${minor.class}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, ) { getWriter(MinorType.${name?upper_case}).write${minor.class}(<#list fields as field>${field.name}<#if field_has_next>, ); } <#if minor.class?ends_with("VarBinary")> @Override public void write${minor.class}(byte[] value) { getWriter(MinorType.${name?upper_case}).write${minor.class}(value); } @Override public void write${minor.class}(byte[] value, int offset, int length) { getWriter(MinorType.${name?upper_case}).write${minor.class}(value, offset, length); } @Override public void write${minor.class}(ByteBuffer value) { getWriter(MinorType.${name?upper_case}).write${minor.class}(value); } @Override public void write${minor.class}(ByteBuffer value, int offset, int length) { getWriter(MinorType.${name?upper_case}).write${minor.class}(value, offset, length); } <#elseif minor.class?ends_with("VarChar")> @Override public void write${minor.class}(Text value) { getWriter(MinorType.${name?upper_case}).write${minor.class}(value); } @Override public void write${minor.class}(String value) { getWriter(MinorType.${name?upper_case}).write${minor.class}(value); } public void writeNull() { } @Override public StructWriter struct() { return getWriter(MinorType.LIST).struct(); } @Override public ListWriter list() { return getWriter(MinorType.LIST).list(); } @Override public ListWriter listView() { return getWriter(MinorType.LISTVIEW).listView(); } @Override public MapWriter map() { return getWriter(MinorType.LIST).map(); } @Override 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); } @Override public ListWriter list(String name) { return getWriter(MinorType.STRUCT).list(name); } @Override public ListWriter listView(String name) { return getWriter(MinorType.STRUCT).listView(name); } @Override public MapWriter map(String name) { return getWriter(MinorType.STRUCT).map(name); } @Override 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" /> <#assign upperName = minor.class?upper_case /> <#assign capName = minor.class?cap_first /> <#if minor.typeParams?? > @Override public ${capName}Writer ${lowerName}(String name<#list minor.typeParams as typeParam>, ${typeParam.type} ${typeParam.name}) { return getWriter(MinorType.STRUCT).${lowerName}(name<#list minor.typeParams as typeParam>, ${typeParam.name}); } @Override public ${capName}Writer ${lowerName}(String name) { return getWriter(MinorType.STRUCT).${lowerName}(name); } @Override public ${capName}Writer ${lowerName}() { return getWriter(MinorType.LIST).${lowerName}(); } public void copyReader(FieldReader reader) { getWriter().copyReader(reader); } public void copyReaderToField(String name, FieldReader reader) { getWriter().copyReaderToField(name, reader); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/ArrowType.java000066400000000000000000000302311515451615200263130ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/types/pojo/ArrowType.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.types.pojo; import com.google.flatbuffers.FlatBufferBuilder; import java.util.Objects; 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; import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.annotation.JsonSubTypes; import com.fasterxml.jackson.annotation.JsonTypeInfo; /** * Arrow types * Source code generated using FreeMarker template ${.template_name} **/ @JsonTypeInfo( use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.PROPERTY, property = "name") @JsonSubTypes({ <#list arrowTypes.types as type> @JsonSubTypes.Type(value = ArrowType.${type.name?remove_ending("_")}.class, name = "${type.name?remove_ending("_")?lower_case}"), }) public abstract class ArrowType { public static abstract class PrimitiveType extends ArrowType { private PrimitiveType() { } @Override public boolean isComplex() { return false; } } public static abstract class ComplexType extends ArrowType { private ComplexType() { } @Override public boolean isComplex() { return true; } } public static enum ArrowTypeID { <#list arrowTypes.types as type> <#assign name = type.name> ${name?remove_ending("_")}(Type.${name}), NONE(Type.NONE); private final byte flatbufType; public byte getFlatbufID() { return this.flatbufType; } private ArrowTypeID(byte flatbufType) { this.flatbufType = flatbufType; } } @JsonIgnore public abstract ArrowTypeID getTypeID(); @JsonIgnore public abstract boolean isComplex(); public abstract int getType(FlatBufferBuilder builder); public abstract T accept(ArrowTypeVisitor visitor); /** * to visit the ArrowTypes * * type.accept(new ArrowTypeVisitor<Type>() { * ... * }); * */ public static interface ArrowTypeVisitor { <#list arrowTypes.types as type> T visit(${type.name?remove_ending("_")} type); default T visit(ExtensionType type) { return type.storageType().accept(this); } } /** * to visit the Complex ArrowTypes and bundle Primitive ones in one case */ public static abstract class ComplexTypeVisitor implements ArrowTypeVisitor { public T visit(PrimitiveType type) { throw new UnsupportedOperationException("Unexpected Primitive type: " + type); } <#list arrowTypes.types as type> <#if !type.complex> public final T visit(${type.name?remove_ending("_")} type) { return visit((PrimitiveType) type); } } /** * to visit the Primitive ArrowTypes and bundle Complex ones under one case */ public static abstract class PrimitiveTypeVisitor implements ArrowTypeVisitor { public T visit(ComplexType type) { throw new UnsupportedOperationException("Unexpected Complex type: " + type); } <#list arrowTypes.types as type> <#if type.complex> public final T visit(${type.name?remove_ending("_")} type) { return visit((ComplexType) type); } } <#list arrowTypes.types as type> <#assign name = type.name?remove_ending("_")> <#assign fields = type.fields> public static class ${name} extends <#if type.complex>ComplexType<#else>PrimitiveType { public static final ArrowTypeID TYPE_TYPE = ArrowTypeID.${name}; <#if type.fields?size == 0> public static final ${name} INSTANCE = new ${name}(); <#else> <#list fields as field> <#assign fieldType = field.valueType!field.type> ${fieldType} ${field.name}; <#if type.name == "Decimal"> // Needed to support golden file integration tests. @JsonCreator public static Decimal createDecimal( @JsonProperty("precision") int precision, @JsonProperty("scale") int scale, @JsonProperty("bitWidth") Integer bitWidth) { return new Decimal(precision, scale, bitWidth == null ? 128 : bitWidth); } /** * Construct Decimal with 128 bits. * * This is kept mainly for the sake of backward compatibility. * Please use {@link org.apache.arrow.vector.types.pojo.ArrowType.Decimal#Decimal(int, int, int)} instead. * * @deprecated This API will be removed in a future release. */ @Deprecated public Decimal(int precision, int scale) { this(precision, scale, 128); } <#else> @JsonCreator public ${type.name}( <#list type.fields as field> <#assign fieldType = field.valueType!field.type> @JsonProperty("${field.name}") ${fieldType} ${field.name}<#if field_has_next>, ) { <#list type.fields as field> this.${field.name} = ${field.name}; } <#list fields as field> <#assign fieldType = field.valueType!field.type> public ${fieldType} get${field.name?cap_first}() { return ${field.name}; } @Override public ArrowTypeID getTypeID() { return TYPE_TYPE; } @Override public int getType(FlatBufferBuilder builder) { <#list type.fields as field> <#if field.type == "String"> int ${field.name} = this.${field.name} == null ? -1 : builder.createString(this.${field.name}); <#if field.type == "int[]"> int ${field.name} = this.${field.name} == null ? -1 : org.apache.arrow.flatbuf.${type.name}.create${field.name?cap_first}Vector(builder, this.${field.name}); org.apache.arrow.flatbuf.${type.name}.start${type.name}(builder); <#list type.fields as field> <#if field.type == "String" || field.type == "int[]"> if (this.${field.name} != null) { org.apache.arrow.flatbuf.${type.name}.add${field.name?cap_first}(builder, ${field.name}); } <#else> org.apache.arrow.flatbuf.${type.name}.add${field.name?cap_first}(builder, this.${field.name}<#if field.valueType??>.getFlatbufID()); return org.apache.arrow.flatbuf.${type.name}.end${type.name}(builder); } public String toString() { return "${name}" <#if fields?size != 0> + "(" <#list fields as field> + <#if field.type == "int[]">java.util.Arrays.toString(${field.name})<#else>${field.name}<#if field_has_next> + ", " + ")" ; } @Override public int hashCode() { return java.util.Arrays.deepHashCode(new Object[] {<#list type.fields as field>${field.name}<#if field_has_next>, }); } @Override public boolean equals(Object obj) { if (!(obj instanceof ${name})) { return false; } <#if type.fields?size == 0> return true; <#else> ${type.name} that = (${type.name}) obj; return <#list type.fields as field>Objects.deepEquals(this.${field.name}, that.${field.name}) <#if field_has_next>&&<#else>; } @Override public T accept(ArrowTypeVisitor visitor) { return visitor.visit(this); } } /** * A user-defined data type that wraps an underlying storage type. */ public abstract static class ExtensionType extends ComplexType { /** The on-wire type for this user-defined type. */ public abstract ArrowType storageType(); /** The name of this user-defined type. Used to identify the type during serialization. */ public abstract String extensionName(); /** Check equality of this type to another user-defined type. */ public abstract boolean extensionEquals(ExtensionType other); /** Save any metadata for this type. */ public abstract String serialize(); /** Given saved metadata and the underlying storage type, construct a new instance of the user type. */ public abstract ArrowType deserialize(ArrowType storageType, String serializedData); /** Construct a vector for the user type. */ public abstract FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator); /** The field metadata key storing the name of the extension type. */ public static final String EXTENSION_METADATA_KEY_NAME = "ARROW:extension:name"; /** The field metadata key storing metadata for the extension type. */ public static final String EXTENSION_METADATA_KEY_METADATA = "ARROW:extension:metadata"; @Override public ArrowTypeID getTypeID() { return storageType().getTypeID(); } @Override public int getType(FlatBufferBuilder builder) { return storageType().getType(builder); } public String toString() { return "ExtensionType(" + extensionName() + ", " + storageType().toString() + ")"; } @Override public int hashCode() { return java.util.Arrays.deepHashCode(new Object[] {storageType(), extensionName()}); } @Override public boolean equals(Object obj) { if (!(obj instanceof ExtensionType)) { return false; } return this.extensionEquals((ExtensionType) obj); } @Override 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; public static org.apache.arrow.vector.types.pojo.ArrowType getTypeForField(org.apache.arrow.flatbuf.Field field) { switch(field.typeType()) { <#list arrowTypes.types as type> <#assign name = type.name?remove_ending("_")> <#assign nameLower = type.name?lower_case> <#assign fields = type.fields> case Type.${type.name}: { org.apache.arrow.flatbuf.${type.name} ${nameLower}Type = (org.apache.arrow.flatbuf.${type.name}) field.type(new org.apache.arrow.flatbuf.${type.name}()); <#list type.fields as field> <#if field.type == "int[]"> ${field.type} ${field.name} = new int[${nameLower}Type.${field.name}Length()]; for (int i = 0; i< ${field.name}.length; ++i) { ${field.name}[i] = ${nameLower}Type.${field.name}(i); } <#else> ${field.type} ${field.name} = ${nameLower}Type.${field.name}(); <#if type.name == "Decimal"> if (bitWidth != defaultDecimalBitWidth && bitWidth != 256) { throw new IllegalArgumentException("Library only supports 128-bit and 256-bit decimal values"); } return new ${name}(<#list type.fields as field><#if field.valueType??>${field.valueType}.fromFlatbufID(${field.name})<#else>${field.name}<#if field_has_next>, ); } default: throw new UnsupportedOperationException("Unsupported type: " + field.typeType()); } } public static Int getInt(org.apache.arrow.flatbuf.Field field) { org.apache.arrow.flatbuf.Int intType = (org.apache.arrow.flatbuf.Int) field.type(new org.apache.arrow.flatbuf.Int()); return new Int(intType.bitWidth(), intType.isSigned()); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/BaseReader.java000066400000000000000000000050121515451615200263530ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/reader/BaseReader.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.reader; <#include "/@includes/vv_imports.ftl" /> /** * Source code generated using FreeMarker template ${.template_name} */ @SuppressWarnings("unused") public interface BaseReader extends Positionable{ Field getField(); MinorType getMinorType(); void reset(); void read(UnionHolder holder); void read(int index, UnionHolder holder); void copyAsValue(UnionWriter writer); void read(DenseUnionHolder holder); void read(int index, DenseUnionHolder holder); void copyAsValue(DenseUnionWriter writer); boolean isSet(); public interface StructReader extends BaseReader, Iterable{ FieldReader reader(String name); } public interface RepeatedStructReader extends StructReader{ boolean next(); int size(); void copyAsValue(StructWriter writer); } public interface ListReader extends BaseReader{ FieldReader reader(); } public interface RepeatedListReader extends ListReader{ boolean next(); int size(); void copyAsValue(ListWriter writer); } public interface MapReader extends BaseReader{ FieldReader reader(); } public interface RepeatedMapReader extends MapReader{ boolean next(); int size(); void copyAsValue(MapWriter writer); } public interface ScalarReader extends <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> ${name}Reader, ExtensionReader, BaseReader {} interface ComplexReader{ StructReader rootAsStruct(); ListReader rootAsList(); boolean rootIsStruct(); boolean ok(); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/BaseWriter.java000066400000000000000000000123041515451615200264270ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/writer/BaseWriter.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.writer; <#include "/@includes/vv_imports.ftl" /> /* * File generated from ${.template_name} using FreeMarker. */ @SuppressWarnings("unused") public interface BaseWriter extends AutoCloseable, Positionable { int getValueCapacity(); void writeNull(); public interface StructWriter extends BaseWriter { Field getField(); /** * Whether this writer is a struct writer and is empty (has no children). * *

    * Intended only for use in determining whether to add dummy vector to * avoid empty (zero-column) schema, as in JsonReader. *

    * @return whether the struct is empty */ boolean isEmptyStruct(); <#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 /> <#assign capName = minor.class?cap_first /> <#if minor.typeParams?? > ${capName}Writer ${lowerName}(String name<#list minor.typeParams as typeParam>, ${typeParam.type} ${typeParam.name}); ${capName}Writer ${lowerName}(String name); 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); MapWriter map(String name, boolean keysSorted); void start(); void end(); } public interface ListWriter extends BaseWriter { void startList(); void endList(); void startListView(); void endListView(); StructWriter struct(); ListWriter list(); 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> <#assign lowerName = minor.class?uncap_first /> <#if lowerName == "int" ><#assign lowerName = "integer" /> <#assign upperName = minor.class?upper_case /> <#assign capName = minor.class?cap_first /> ${capName}Writer ${lowerName}(); } public interface MapWriter extends ListWriter { void startMap(); void endMap(); void startEntry(); void endEntry(); MapWriter key(); 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 {} public interface ComplexWriter { void allocate(); void clear(); void copyReader(FieldReader reader); StructWriter rootAsStruct(); ListWriter rootAsList(); ListWriter rootAsListView(); MapWriter rootAsMap(boolean keysSorted); void setPosition(int index); void setValueCount(int count); void reset(); } public interface StructOrListWriter { void start(); void end(); StructOrListWriter struct(String name); /** * @deprecated use {@link #listOfStruct()} instead. */ @Deprecated StructOrListWriter listoftstruct(String name); StructOrListWriter listOfStruct(String name); StructOrListWriter list(String name); boolean isStructWriter(); boolean isListWriter(); VarCharWriter varChar(String name); IntWriter integer(String name); BigIntWriter bigInt(String name); Float4Writer float4(String name); Float8Writer float8(String name); BitWriter bit(String name); VarBinaryWriter binary(String name); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/CaseSensitiveStructWriters.java000066400000000000000000000034551515451615200317210ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ <@pp.dropOutputFile /> <#list ["Nullable", "Single"] as mode> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/${mode}CaseSensitiveStructWriter.java" /> <#assign index = "idx()"> <#if mode == "Single"> <#assign containerClass = "NonNullableStructVector" /> <#else> <#assign containerClass = "StructVector" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; <#include "/@includes/vv_imports.ftl" /> /* * This class is generated using FreeMarker and the ${.template_name} template. */ @SuppressWarnings("unused") public class ${mode}CaseSensitiveStructWriter extends ${mode}StructWriter { public ${mode}CaseSensitiveStructWriter(${containerClass} container) { super(container); } @Override protected String handleCase(final String input){ return input; } @Override protected NullableStructWriterFactory getNullableStructWriterFactory() { return NullableStructWriterFactory.getNullableCaseSensitiveStructWriterFactoryInstance(); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/ComplexCopier.java000066400000000000000000000204461515451615200271370ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ import org.apache.arrow.vector.complex.impl.UnionMapReader; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.complex.writer.FieldWriter; import org.apache.arrow.vector.types.Types; <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/ComplexCopier.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; <#include "/@includes/vv_imports.ftl" /> /* * This class is generated using freemarker and the ${.template_name} template. */ @SuppressWarnings("unused") public class ComplexCopier { /** * Do a deep copy of the value in input into output * @param input field to read from * @param output field to write to */ public static void copy(FieldReader reader, FieldWriter writer) { final MinorType mt = reader.getMinorType(); switch (mt) { case LIST: case LISTVIEW: case LARGELIST: case LARGELISTVIEW: case FIXED_SIZE_LIST: if (reader.isSet()) { writer.startList(); while (reader.next()) { FieldReader childReader = reader.reader(); FieldWriter childWriter = getListWriterForReader(childReader, writer); if (childReader.isSet()) { copy(childReader, childWriter); } else { childWriter.writeNull(); } } writer.endList(); } else { writer.writeNull(); } break; case MAP: if (reader.isSet()) { UnionMapReader mapReader = (UnionMapReader) reader; writer.startMap(); while (mapReader.next()) { FieldReader structReader = reader.reader(); if (structReader.isSet()) { writer.startEntry(); copy(mapReader.key(), getMapWriterForReader(mapReader.key(), writer.key())); copy(mapReader.value(), getMapWriterForReader(mapReader.value(), writer.value())); writer.endEntry(); } else { writer.writeNull(); } } writer.endMap(); } else { writer.writeNull(); } break; case STRUCT: if (reader.isSet()) { writer.start(); for(String name : reader){ FieldReader childReader = reader.reader(name); if (childReader.getMinorType() != Types.MinorType.NULL) { FieldWriter childWriter = getStructWriterForReader(childReader, writer, name); if (childReader.isSet()) { copy(childReader, childWriter); } else { childWriter.writeNull(); } } } writer.end(); } else { 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/> <#if !minor.typeParams?? || minor.class?starts_with("Decimal") > case ${name?upper_case}: if (reader.isSet()) { Nullable${name}Holder ${uncappedName}Holder = new Nullable${name}Holder(); reader.read(${uncappedName}Holder); if (${uncappedName}Holder.isSet == 1) { writer.write${name}(<#list fields as field>${uncappedName}Holder.${field.name}<#if field_has_next>, <#if minor.class?starts_with("Decimal")>, new ArrowType.Decimal(${uncappedName}Holder.precision, ${uncappedName}Holder.scale, ${name}Holder.WIDTH * 8)); } } else { writer.writeNull(); } break; } } private static FieldWriter getStructWriterForReader(FieldReader reader, StructWriter writer, String name) { switch (reader.getMinorType()) { <#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.typeParams??> case ${name?upper_case}: return (FieldWriter) writer.<#if name == "Int">integer<#else>${uncappedName}(name); <#if minor.class?starts_with("Decimal")> case ${name?upper_case}: if (reader.getField().getType() instanceof ArrowType.Decimal) { ArrowType.Decimal type = (ArrowType.Decimal) reader.getField().getType(); return (FieldWriter) writer.${uncappedName}(name, type.getScale(), type.getPrecision()); } else { return (FieldWriter) writer.${uncappedName}(name); } case STRUCT: return (FieldWriter) writer.struct(name); case FIXED_SIZE_LIST: case LIST: return (FieldWriter) writer.list(name); case MAP: 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()); } } private static FieldWriter getListWriterForReader(FieldReader reader, ListWriter writer) { switch (reader.getMinorType()) { <#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.typeParams?? || minor.class?starts_with("Decimal") > case ${name?upper_case}: return (FieldWriter) writer.<#if name == "Int">integer<#else>${uncappedName}(); case STRUCT: return (FieldWriter) writer.struct(); case FIXED_SIZE_LIST: case LIST: case MAP: case NULL: 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()); } } private static FieldWriter getMapWriterForReader(FieldReader reader, MapWriter writer) { switch (reader.getMinorType()) { <#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.typeParams?? || minor.class?starts_with("Decimal") > case ${name?upper_case}: return (FieldWriter) writer.<#if name == "Int">integer<#else>${uncappedName}(); case STRUCT: return (FieldWriter) writer.struct(); case FIXED_SIZE_LIST: case LIST: case NULL: return (FieldWriter) writer.list(); case LISTVIEW: 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()); } } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/ComplexReaders.java000066400000000000000000000103671515451615200273040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ import java.lang.Override; import java.util.List; import org.apache.arrow.record.TransferPair; import org.apache.arrow.vector.complex.IndexHolder; import org.apache.arrow.vector.complex.writer.IntervalWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter; <@pp.dropOutputFile /> <#list vv.types as type> <#list type.minor as minor> <#list [""] as mode> <#assign lowerName = minor.class?uncap_first /> <#if lowerName == "int" ><#assign lowerName = "integer" /> <#assign name = minor.class?cap_first /> <#assign javaType = (minor.javaType!type.javaType) /> <#assign friendlyType = (minor.friendlyType!minor.boxedType!type.boxedType) /> <#assign safeType=friendlyType /> <#if safeType=="byte[]"><#assign safeType="ByteArray" /> <#assign hasFriendly = minor.friendlyType!"no" == "no" /> <#list ["Nullable"] as nullMode> <#if mode == "" > <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/${name}ReaderImpl.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; <#include "/@includes/vv_imports.ftl" /> /** * Source code generated using FreeMarker template ${.template_name} */ @SuppressWarnings("unused") public class ${name}ReaderImpl extends AbstractFieldReader { private final ${name}Vector vector; public ${name}ReaderImpl(${name}Vector vector){ super(); this.vector = vector; } public MinorType getMinorType(){ return vector.getMinorType(); } public Field getField(){ return vector.getField(); } public boolean isSet(){ return !vector.isNull(idx()); } public void copyAsValue(${minor.class?cap_first}Writer writer){ ${minor.class?cap_first}WriterImpl impl = (${minor.class?cap_first}WriterImpl) writer; impl.vector.copyFromSafe(idx(), impl.idx(), vector); } public void copyAsField(String name, StructWriter writer){ ${minor.class?cap_first}WriterImpl impl = (${minor.class?cap_first}WriterImpl) writer.${lowerName}(name); impl.vector.copyFromSafe(idx(), impl.idx(), vector); } <#if nullMode != "Nullable"> public void read(${minor.class?cap_first}Holder h){ vector.get(idx(), h); } public void read(Nullable${minor.class?cap_first}Holder h){ vector.get(idx(), h); } public ${friendlyType} read${safeType}(){ return vector.getObject(idx()); } <#if minor.class == "TimeStampSec" || minor.class == "TimeStampMilli" || minor.class == "TimeStampMicro" || minor.class == "TimeStampNano"> @Override public ${minor.boxedType} read${minor.boxedType}(){ return vector.get(idx()); } public void copyValue(FieldWriter w){ } public Object readObject(){ return (Object)vector.getObject(idx()); } } <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/reader/${name}Reader.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.reader; <#include "/@includes/vv_imports.ftl" /> /** * Source code generated using FreeMarker template ${.template_name} */ @SuppressWarnings("unused") public interface ${name}Reader extends BaseReader{ public void read(${minor.class?cap_first}Holder h); public void read(Nullable${minor.class?cap_first}Holder h); public Object readObject(); // read friendly type public ${friendlyType} read${safeType}(); public boolean isSet(); public void copyAsValue(${minor.class}Writer writer); public void copyAsField(String name, ${minor.class}Writer writer); } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/ComplexWriters.java000066400000000000000000000205621515451615200273540ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ <@pp.dropOutputFile /> <#list vv.types as type> <#list type.minor as minor> <#list ["Nullable"] as mode> <#assign name = minor.class?cap_first /> <#assign eName = name /> <#assign javaType = (minor.javaType!type.javaType) /> <#assign fields = minor.fields!type.fields /> <#assign friendlyType = (minor.friendlyType!minor.boxedType!type.boxedType) /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/${eName}WriterImpl.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; <#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 on the ${.template_name} template. */ @SuppressWarnings("unused") public class ${eName}WriterImpl extends AbstractFieldWriter { final ${name}Vector vector; <#if minor.class?ends_with("VarChar")> private final Text textBuffer = new Text(); public ${eName}WriterImpl(${name}Vector vector) { this.vector = vector; } @Override public Field getField() { return vector.getField(); } @Override public int getValueCapacity() { return vector.getValueCapacity(); } @Override public void allocate() { vector.allocateNew(); } @Override public void close() { vector.close(); } @Override public void clear() { vector.clear(); } @Override protected int idx() { return super.idx(); } <#if mode == "Repeated"> public void write(${minor.class?cap_first}Holder h) { mutator.addSafe(idx(), h); vector.setValueCount(idx()+1); } public void write(${minor.class?cap_first}Holder h) { mutator.addSafe(idx(), h); vector.setValueCount(idx()+1); } public void write${minor.class}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, ) { mutator.addSafe(idx(), <#list fields as field>${field.name}<#if field_has_next>, ); vector.setValueCount(idx()+1); } public void setPosition(int idx) { super.setPosition(idx); mutator.startNewValue(idx); } <#else> <#if !minor.class?starts_with("Decimal")> public void write(${minor.class}Holder h) { vector.setSafe(idx(), h); vector.setValueCount(idx()+1); } public void write(Nullable${minor.class}Holder h) { vector.setSafe(idx(), h); vector.setValueCount(idx()+1); } public void write${minor.class}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, ) { vector.setSafe(idx(), 1<#list fields as field><#if field.include!true >, ${field.name}); vector.setValueCount(idx()+1); } <#if minor.class?ends_with("VarChar")> @Override public void write${minor.class}(${friendlyType} value) { vector.setSafe(idx(), value); vector.setValueCount(idx()+1); } @Override public void write${minor.class}(String value) { textBuffer.set(value); vector.setSafe(idx(), textBuffer); vector.setValueCount(idx()+1); } <#if minor.class?starts_with("Decimal")> public void write(${minor.class}Holder h){ DecimalUtility.checkPrecisionAndScale(h.precision, h.scale, vector.getPrecision(), vector.getScale()); vector.setSafe(idx(), h); vector.setValueCount(idx() + 1); } public void write(Nullable${minor.class}Holder h){ if (h.isSet == 1) { DecimalUtility.checkPrecisionAndScale(h.precision, h.scale, vector.getPrecision(), vector.getScale()); } vector.setSafe(idx(), h); vector.setValueCount(idx() + 1); } public void write${minor.class}(long start, ArrowBuf buffer){ vector.setSafe(idx(), 1, start, buffer); vector.setValueCount(idx() + 1); } public void write${minor.class}(long start, ArrowBuf buffer, ArrowType arrowType){ DecimalUtility.checkPrecisionAndScale(((ArrowType.Decimal) arrowType).getPrecision(), ((ArrowType.Decimal) arrowType).getScale(), vector.getPrecision(), vector.getScale()); vector.setSafe(idx(), 1, start, buffer); vector.setValueCount(idx() + 1); } public void write${minor.class}(BigDecimal value){ // vector.setSafe already does precision and scale checking vector.setSafe(idx(), value); vector.setValueCount(idx() + 1); } public void writeBigEndianBytesTo${minor.class}(byte[] value, ArrowType arrowType){ DecimalUtility.checkPrecisionAndScale(((ArrowType.Decimal) arrowType).getPrecision(), ((ArrowType.Decimal) arrowType).getScale(), vector.getPrecision(), vector.getScale()); vector.setBigEndianSafe(idx(), value); vector.setValueCount(idx() + 1); } public void writeBigEndianBytesTo${minor.class}(byte[] value){ vector.setBigEndianSafe(idx(), value); vector.setValueCount(idx() + 1); } public void writeNull() { vector.setNull(idx()); vector.setValueCount(idx()+1); } <#if minor.class?ends_with("VarBinary")> public void write${minor.class}(byte[] value) { vector.setSafe(idx(), value); vector.setValueCount(idx() + 1); } public void write${minor.class}(byte[] value, int offset, int length) { vector.setSafe(idx(), value, offset, length); vector.setValueCount(idx() + 1); } public void write${minor.class}(ByteBuffer value) { vector.setSafe(idx(), value, 0, value.remaining()); vector.setValueCount(idx() + 1); } public void write${minor.class}(ByteBuffer value, int offset, int length) { vector.setSafe(idx(), value, offset, length); vector.setValueCount(idx() + 1); } } <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/writer/${eName}Writer.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.writer; <#include "/@includes/vv_imports.ftl" /> /* * This class is generated using FreeMarker on the ${.template_name} template. */ @SuppressWarnings("unused") public interface ${eName}Writer extends BaseWriter { public void write(${minor.class}Holder h); <#if minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> /** * @deprecated * The holder version should be used instead because the plain value version does not contain enough information * to fully specify this field type. * @see #write(${minor.class}Holder) */ @Deprecated public void write${minor.class}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, ); <#if minor.class?starts_with("Decimal")> public void write${minor.class}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, , ArrowType arrowType); public void write${minor.class}(${friendlyType} value); public void writeBigEndianBytesTo${minor.class}(byte[] value, ArrowType arrowType); /** * @deprecated * Use either the version that additionally takes in an ArrowType or use the holder version. * This version does not contain enough information to fully specify this field type. * @see #writeBigEndianBytesTo${minor.class}(byte[], ArrowType) * @see #write(${minor.class}Holder) */ @Deprecated public void writeBigEndianBytesTo${minor.class}(byte[] value); <#if minor.class?ends_with("VarBinary")> public void write${minor.class}(byte[] value); public void write${minor.class}(byte[] value, int offset, int length); public void write${minor.class}(ByteBuffer value); public void write${minor.class}(ByteBuffer value, int offset, int length); <#if minor.class?ends_with("VarChar")> public void write${minor.class}(${friendlyType} value); public void write${minor.class}(String value); } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/DenseUnionReader.java000066400000000000000000000153261515451615200275610ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ import org.apache.arrow.vector.complex.impl.UnionListReader; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/DenseUnionReader.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; <#include "/@includes/vv_imports.ftl" /> /** * Source code generated using FreeMarker template ${.template_name} */ @SuppressWarnings("unused") public class DenseUnionReader extends AbstractFieldReader { private BaseReader[] readers = new BaseReader[Byte.MAX_VALUE + 1]; public DenseUnionVector data; public DenseUnionReader(DenseUnionVector data) { this.data = data; } public MinorType getMinorType() { byte typeId = data.getTypeId(idx()); return data.getVectorByType(typeId).getMinorType(); } public byte getTypeId() { return data.getTypeId(idx()); } @Override public Field getField() { return data.getField(); } public boolean isSet(){ return !data.isNull(idx()); } public void read(DenseUnionHolder holder) { holder.reader = this; holder.isSet = this.isSet() ? 1 : 0; holder.typeId = getTypeId(); } public void read(int index, UnionHolder holder) { byte typeId = data.getTypeId(index); getList(typeId).read(index, holder); } private FieldReader getReaderForIndex(int index) { byte typeId = data.getTypeId(index); MinorType minorType = data.getVectorByType(typeId).getMinorType(); FieldReader reader = (FieldReader) readers[typeId]; if (reader != null) { return reader; } switch (minorType) { case NULL: reader = NullReader.INSTANCE; break; case STRUCT: reader = (FieldReader) getStruct(typeId); break; case LIST: reader = (FieldReader) getList(typeId); break; <#list vv.types as type> <#list type.minor as minor> <#assign name = minor.class?cap_first /> <#assign uncappedName = name?uncap_first/> <#if !minor.typeParams?? || minor.class?starts_with("Decimal")> case ${name?upper_case}: reader = (FieldReader) get${name}(typeId); break; default: throw new UnsupportedOperationException("Unsupported type: " + MinorType.values()[typeId]); } return reader; } private SingleStructReaderImpl structReader; private StructReader getStruct(byte typeId) { StructReader structReader = (StructReader) readers[typeId]; if (structReader == null) { structReader = (SingleStructReaderImpl) data.getVectorByType(typeId).getReader(); structReader.setPosition(idx()); readers[typeId] = structReader; } return structReader; } private UnionListReader listReader; private FieldReader getList(byte typeId) { UnionListReader listReader = (UnionListReader) readers[typeId]; if (listReader == null) { listReader = new UnionListReader((ListVector) data.getVectorByType(typeId)); listReader.setPosition(idx()); readers[typeId] = listReader; } return listReader; } private UnionMapReader mapReader; private FieldReader getMap(byte typeId) { UnionMapReader mapReader = (UnionMapReader) readers[typeId]; if (mapReader == null) { mapReader = new UnionMapReader((MapVector) data.getVectorByType(typeId)); mapReader.setPosition(idx()); readers[typeId] = mapReader; } return mapReader; } @Override public java.util.Iterator iterator() { throw new UnsupportedOperationException(); } @Override public void copyAsValue(UnionWriter writer) { writer.data.copyFrom(idx(), writer.idx(), data); } <#list ["Object", "BigDecimal", "Short", "Integer", "Long", "Boolean", "LocalDateTime", "Duration", "Period", "Double", "Float", "Character", "Text", "Byte", "byte[]", "PeriodDuration"] as friendlyType> <#assign safeType=friendlyType /> <#if safeType=="byte[]"><#assign safeType="ByteArray" /> @Override public ${friendlyType} read${safeType}() { return getReaderForIndex(idx()).read${safeType}(); } public int size() { return getReaderForIndex(idx()).size(); } <#list vv.types as type> <#list type.minor as minor> <#assign name = minor.class?cap_first /> <#assign uncappedName = name?uncap_first/> <#assign boxedType = (minor.boxedType!type.boxedType) /> <#assign javaType = (minor.javaType!type.javaType) /> <#assign friendlyType = (minor.friendlyType!minor.boxedType!type.boxedType) /> <#assign safeType=friendlyType /> <#if safeType=="byte[]"><#assign safeType="ByteArray" /> <#if !minor.typeParams?? || minor.class?starts_with("Decimal")> private ${name}ReaderImpl get${name}(byte typeId) { ${name}ReaderImpl reader = (${name}ReaderImpl) readers[typeId]; if (reader == null) { reader = new ${name}ReaderImpl((${name}Vector) data.getVectorByType(typeId)); reader.setPosition(idx()); readers[typeId] = reader; } return reader; } public void read(Nullable${name}Holder holder){ getReaderForIndex(idx()).read(holder); } public void copyAsValue(${name}Writer writer){ getReaderForIndex(idx()).copyAsValue(writer); } @Override public void copyAsValue(ListWriter writer) { ComplexCopier.copy(this, (FieldWriter) writer); } @Override public void setPosition(int index) { super.setPosition(index); byte typeId = data.getTypeId(index); if (readers[typeId] != null) { int offset = data.getOffset(index); readers[typeId].setPosition(offset); } } public FieldReader reader(byte typeId, String name){ return getStruct(typeId).reader(name); } public FieldReader reader(byte typeId) { return getList(typeId).reader(); } public boolean next() { return getReaderForIndex(idx()).next(); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/DenseUnionVector.java000066400000000000000000000771501515451615200276240ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReferenceManager; import org.apache.arrow.memory.util.CommonUtil; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BaseValueVector; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueIterableVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.AbstractStructVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.NonNullableStructVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.UnionMode; import org.apache.arrow.vector.compare.RangeEqualsVisitor; 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.CallBack; import org.apache.arrow.vector.util.DataSizeRoundingUtil; import org.apache.arrow.vector.util.TransferPair; import java.util.Arrays; import java.util.stream.Collectors; <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/DenseUnionVector.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex; <#include "/@includes/vv_imports.ftl" /> import java.util.ArrayList; import java.util.Collections; import java.util.Iterator; import org.apache.arrow.memory.util.CommonUtil; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.memory.util.hash.SimpleHasher; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.impl.ComplexCopier; import org.apache.arrow.vector.util.CallBack; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.BaseValueVector; import org.apache.arrow.vector.ValueIterableVector; import org.apache.arrow.vector.util.OversizedAllocationException; import org.apache.arrow.util.Preconditions; import static org.apache.arrow.vector.types.UnionMode.Dense; /* * This class is generated using freemarker and the ${.template_name} template. */ @SuppressWarnings("unused") /** * A vector which can hold values of different types. It does so by using a StructVector which contains a vector for each * primitive type that is stored. StructVector is used in order to take advantage of its serialization/deserialization methods, * as well as the addOrGet method. * * For performance reasons, DenseUnionVector stores a cached reference to each subtype vector, to avoid having to do the struct lookup * each time the vector is accessed. * Source code generated using FreeMarker template ${.template_name} */ public class DenseUnionVector extends AbstractContainerVector implements FieldVector, ValueIterableVector { int valueCount; NonNullableStructVector internalStruct; private ArrowBuf typeBuffer; private ArrowBuf offsetBuffer; /** * The key is type Id, and the value is vector. */ private ValueVector[] childVectors = new ValueVector[Byte.MAX_VALUE + 1]; /** * The index is the type id, and the value is the type field. */ private Field[] typeFields = new Field[Byte.MAX_VALUE + 1]; /** * The index is the index into the typeFields array, and the value is the logical field id. */ private byte[] typeMapFields = new byte[Byte.MAX_VALUE + 1]; /** * The next type id to allocate. */ private byte nextTypeId = 0; private FieldReader reader; private long typeBufferAllocationSizeInBytes; private long offsetBufferAllocationSizeInBytes; private final FieldType fieldType; public static final byte TYPE_WIDTH = 1; public static final byte OFFSET_WIDTH = 4; private static final FieldType INTERNAL_STRUCT_TYPE = new FieldType(/*nullable*/ false, ArrowType.Struct.INSTANCE, /*dictionary*/ null, /*metadata*/ null); public static DenseUnionVector empty(String name, BufferAllocator allocator) { FieldType fieldType = FieldType.notNullable(new ArrowType.Union( UnionMode.Dense, null)); return new DenseUnionVector(name, allocator, fieldType, null); } public DenseUnionVector(String name, BufferAllocator allocator, FieldType fieldType, CallBack callBack) { super(name, allocator, callBack); this.fieldType = fieldType; this.internalStruct = new NonNullableStructVector( "internal", allocator, INTERNAL_STRUCT_TYPE, callBack, AbstractStructVector.ConflictPolicy.CONFLICT_REPLACE, false); this.typeBuffer = allocator.getEmpty(); this.typeBufferAllocationSizeInBytes = BaseValueVector.INITIAL_VALUE_ALLOCATION * TYPE_WIDTH; this.offsetBuffer = allocator.getEmpty(); this.offsetBufferAllocationSizeInBytes = BaseValueVector.INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH; } public BufferAllocator getAllocator() { return allocator; } @Override public MinorType getMinorType() { return MinorType.DENSEUNION; } @Override public void initializeChildrenFromFields(List children) { for (Field field : children) { byte typeId = registerNewTypeId(field); FieldVector vector = (FieldVector) internalStruct.add(field.getName(), field.getFieldType()); vector.initializeChildrenFromFields(field.getChildren()); childVectors[typeId] = vector; } } @Override public List getChildrenFromFields() { return internalStruct.getChildrenFromFields(); } @Override public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers) { if (ownBuffers.size() != 2) { throw new IllegalArgumentException("Illegal buffer count for dense union with type " + getField().getFieldType() + ", expected " + 2 + ", got: " + ownBuffers.size()); } ArrowBuf buffer = ownBuffers.get(0); typeBuffer.getReferenceManager().release(); typeBuffer = buffer.getReferenceManager().retain(buffer, allocator); typeBufferAllocationSizeInBytes = typeBuffer.capacity(); buffer = ownBuffers.get(1); offsetBuffer.getReferenceManager().release(); offsetBuffer = buffer.getReferenceManager().retain(buffer, allocator); offsetBufferAllocationSizeInBytes = offsetBuffer.capacity(); this.valueCount = fieldNode.getLength(); } @Override public List getFieldBuffers() { List result = new ArrayList<>(2); setReaderAndWriterIndex(); result.add(typeBuffer); result.add(offsetBuffer); return result; } private void setReaderAndWriterIndex() { typeBuffer.readerIndex(0); typeBuffer.writerIndex(valueCount * TYPE_WIDTH); offsetBuffer.readerIndex(0); offsetBuffer.writerIndex((long) valueCount * OFFSET_WIDTH); } /** * Get the inner vectors. * * @deprecated This API will be removed as the current implementations no longer support inner vectors. * * @return the inner vectors for this field as defined by the TypeLayout */ @Override @Deprecated public List getFieldInnerVectors() { throw new UnsupportedOperationException("There are no inner vectors. Use geFieldBuffers"); } private String fieldName(byte typeId, MinorType type) { return type.name().toLowerCase() + typeId; } private FieldType fieldType(MinorType type) { return FieldType.nullable(type.getType()); } public synchronized byte registerNewTypeId(Field field) { if (nextTypeId == typeFields.length) { throw new IllegalStateException("Dense union vector support at most " + typeFields.length + " relative types. Please use union of union instead"); } byte typeId = nextTypeId; if (this.fieldType != null) { int[] typeIds = ((ArrowType.Union) this.fieldType.getType()).getTypeIds(); if (typeIds != null) { int thisTypeId = typeIds[nextTypeId]; if (thisTypeId > Byte.MAX_VALUE) { throw new IllegalStateException("Dense union vector types must be bytes. " + thisTypeId + " is too large"); } typeId = (byte) thisTypeId; } } typeFields[typeId] = field; typeMapFields[nextTypeId] = typeId; this.nextTypeId += 1; return typeId; } private T addOrGet(byte typeId, MinorType minorType, Class c) { return internalStruct.addOrGet(fieldName(typeId, minorType), fieldType(minorType), c); } private T addOrGet(byte typeId, MinorType minorType, ArrowType arrowType, Class c) { return internalStruct.addOrGet(fieldName(typeId, minorType), FieldType.nullable(arrowType), c); } @Override public long getOffsetBufferAddress() { return offsetBuffer.memoryAddress(); } @Override public long getDataBufferAddress() { throw new UnsupportedOperationException(); } @Override public long getValidityBufferAddress() { throw new UnsupportedOperationException(); } @Override public ArrowBuf getValidityBuffer() { throw new UnsupportedOperationException(); } @Override public ArrowBuf getOffsetBuffer() { return offsetBuffer; } public ArrowBuf getTypeBuffer() { return typeBuffer; } @Override public ArrowBuf getDataBuffer() { throw new UnsupportedOperationException(); } public StructVector getStruct(byte typeId) { StructVector structVector = typeId < 0 ? null : (StructVector) childVectors[typeId]; if (structVector == null) { int vectorCount = internalStruct.size(); structVector = addOrGet(typeId, MinorType.STRUCT, StructVector.class); if (internalStruct.size() > vectorCount) { structVector.allocateNew(); childVectors[typeId] = structVector; if (callBack != null) { callBack.doWork(); } } } return structVector; } <#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/> <#assign lowerCaseName = name?lower_case/> <#if !minor.typeParams?? || minor.class?starts_with("Decimal")> public ${name}Vector get${name}Vector(byte typeId<#if minor.class?starts_with("Decimal")>, ArrowType arrowType) { ValueVector vector = typeId < 0 ? null : childVectors[typeId]; if (vector == null) { int vectorCount = internalStruct.size(); vector = addOrGet(typeId, MinorType.${name?upper_case}<#if minor.class?starts_with("Decimal")>, arrowType, ${name}Vector.class); childVectors[typeId] = vector; if (internalStruct.size() > vectorCount) { vector.allocateNew(); if (callBack != null) { callBack.doWork(); } } } return (${name}Vector) vector; } public ListVector getList(byte typeId) { ListVector listVector = typeId < 0 ? null : (ListVector) childVectors[typeId]; if (listVector == null) { int vectorCount = internalStruct.size(); listVector = addOrGet(typeId, MinorType.LIST, ListVector.class); if (internalStruct.size() > vectorCount) { listVector.allocateNew(); childVectors[typeId] = listVector; if (callBack != null) { callBack.doWork(); } } } return listVector; } public MapVector getMap(byte typeId) { MapVector mapVector = typeId < 0 ? null : (MapVector) childVectors[typeId]; if (mapVector == null) { int vectorCount = internalStruct.size(); mapVector = addOrGet(typeId, MinorType.MAP, MapVector.class); if (internalStruct.size() > vectorCount) { mapVector.allocateNew(); childVectors[typeId] = mapVector; if (callBack != null) { callBack.doWork(); } } } return mapVector; } public byte getTypeId(int index) { return typeBuffer.getByte(index * TYPE_WIDTH); } public ValueVector getVectorByType(byte typeId) { return typeId < 0 ? null : childVectors[typeId]; } @Override public void allocateNew() throws OutOfMemoryException { /* new allocation -- clear the current buffers */ clear(); internalStruct.allocateNew(); try { allocateTypeBuffer(); allocateOffsetBuffer(); } catch (Exception e) { clear(); throw e; } } @Override public boolean allocateNewSafe() { /* new allocation -- clear the current buffers */ clear(); boolean safe = internalStruct.allocateNewSafe(); if (!safe) { return false; } try { allocateTypeBuffer(); allocateOffsetBuffer(); } catch (Exception e) { clear(); return false; } return true; } private void allocateTypeBuffer() { typeBuffer = allocator.buffer(typeBufferAllocationSizeInBytes); typeBuffer.readerIndex(0); setNegative(0, typeBuffer.capacity()); } private void allocateOffsetBuffer() { offsetBuffer = allocator.buffer(offsetBufferAllocationSizeInBytes); offsetBuffer.readerIndex(0); offsetBuffer.setZero(0, offsetBuffer.capacity()); } @Override public void reAlloc() { internalStruct.reAlloc(); reallocTypeBuffer(); reallocOffsetBuffer(); } public int getOffset(int index) { return offsetBuffer.getInt((long) index * OFFSET_WIDTH); } private void reallocTypeBuffer() { final long currentBufferCapacity = typeBuffer.capacity(); long newAllocationSize = currentBufferCapacity * 2; if (newAllocationSize == 0) { if (typeBufferAllocationSizeInBytes > 0) { newAllocationSize = typeBufferAllocationSizeInBytes; } else { newAllocationSize = BaseValueVector.INITIAL_VALUE_ALLOCATION * TYPE_WIDTH * 2; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); assert newAllocationSize >= 1; if (newAllocationSize > BaseValueVector.MAX_ALLOCATION_SIZE) { throw new OversizedAllocationException("Unable to expand the buffer"); } final ArrowBuf newBuf = allocator.buffer((int)newAllocationSize); newBuf.setBytes(0, typeBuffer, 0, currentBufferCapacity); typeBuffer.getReferenceManager().release(1); typeBuffer = newBuf; typeBufferAllocationSizeInBytes = (int)newAllocationSize; setNegative(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); } private void reallocOffsetBuffer() { final long currentBufferCapacity = offsetBuffer.capacity(); long newAllocationSize = currentBufferCapacity * 2; if (newAllocationSize == 0) { if (offsetBufferAllocationSizeInBytes > 0) { newAllocationSize = offsetBufferAllocationSizeInBytes; } else { newAllocationSize = BaseValueVector.INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH * 2; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); assert newAllocationSize >= 1; if (newAllocationSize > BaseValueVector.MAX_ALLOCATION_SIZE) { throw new OversizedAllocationException("Unable to expand the buffer"); } final ArrowBuf newBuf = allocator.buffer((int) newAllocationSize); newBuf.setBytes(0, offsetBuffer, 0, currentBufferCapacity); newBuf.setZero(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); offsetBuffer.getReferenceManager().release(1); offsetBuffer = newBuf; offsetBufferAllocationSizeInBytes = (int) newAllocationSize; } @Override public void setInitialCapacity(int numRecords) { } @Override public int getValueCapacity() { long capacity = getTypeBufferValueCapacity(); long offsetCapacity = getOffsetBufferValueCapacity(); if (offsetCapacity < capacity) { capacity = offsetCapacity; } long structCapacity = internalStruct.getValueCapacity(); if (structCapacity < capacity) { structCapacity = capacity; } return (int) capacity; } @Override public void close() { clear(); } @Override public void clear() { valueCount = 0; typeBuffer.getReferenceManager().release(); typeBuffer = allocator.getEmpty(); offsetBuffer.getReferenceManager().release(); offsetBuffer = allocator.getEmpty(); internalStruct.clear(); } @Override public void reset() { valueCount = 0; setNegative(0, typeBuffer.capacity()); offsetBuffer.setZero(0, offsetBuffer.capacity()); internalStruct.reset(); } @Override public Field getField() { int childCount = (int) Arrays.stream(typeFields).filter(field -> field != null).count(); List childFields = new ArrayList<>(childCount); int[] typeIds = new int[childCount]; for (int i = 0; i < typeFields.length; i++) { if (typeFields[i] != null) { int curIdx = childFields.size(); typeIds[curIdx] = i; childFields.add(typeFields[i]); } } FieldType fieldType; if (this.fieldType == null) { fieldType = FieldType.nullable(new ArrowType.Union(Dense, typeIds)); } else { final UnionMode mode = UnionMode.Dense; fieldType = new FieldType(this.fieldType.isNullable(), new ArrowType.Union(mode, typeIds), this.fieldType.getDictionary(), this.fieldType.getMetadata()); } return new Field(name, fieldType, childFields); } @Override public TransferPair getTransferPair(BufferAllocator allocator) { return getTransferPair(name, allocator); } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return getTransferPair(ref, allocator, null); } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return new org.apache.arrow.vector.complex.DenseUnionVector.TransferImpl(ref, allocator, callBack); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return getTransferPair(field, allocator, null); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { return new org.apache.arrow.vector.complex.DenseUnionVector.TransferImpl(field, allocator, callBack); } @Override public TransferPair makeTransferPair(ValueVector target) { return new TransferImpl((DenseUnionVector) target); } @Override public void copyFrom(int inIndex, int outIndex, ValueVector from) { Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); DenseUnionVector fromCast = (DenseUnionVector) from; int inOffset = fromCast.offsetBuffer.getInt((long) inIndex * OFFSET_WIDTH); fromCast.getReader().setPosition(inOffset); int outOffset = offsetBuffer.getInt((long) outIndex * OFFSET_WIDTH); getWriter().setPosition(outOffset); ComplexCopier.copy(fromCast.reader, writer); } @Override public void copyFromSafe(int inIndex, int outIndex, ValueVector from) { copyFrom(inIndex, outIndex, from); } public FieldVector addVector(byte typeId, FieldVector v) { final String name = v.getName().isEmpty() ? fieldName(typeId, v.getMinorType()) : v.getName(); Preconditions.checkState(internalStruct.getChild(name) == null, String.format("%s vector already exists", name)); final FieldVector newVector = internalStruct.addOrGet(name, v.getField().getFieldType(), v.getClass()); v.makeTransferPair(newVector).transfer(); internalStruct.putChild(name, newVector); childVectors[typeId] = newVector; if (callBack != null) { callBack.doWork(); } return newVector; } private class TransferImpl implements TransferPair { private final TransferPair[] internalTransferPairs = new TransferPair[nextTypeId]; private final DenseUnionVector to; public TransferImpl(String name, BufferAllocator allocator, CallBack callBack) { to = new DenseUnionVector(name, allocator, null, callBack); internalStruct.makeTransferPair(to.internalStruct); createTransferPairs(); } public TransferImpl(Field field, BufferAllocator allocator, CallBack callBack) { to = new DenseUnionVector(field.getName(), allocator, null, callBack); internalStruct.makeTransferPair(to.internalStruct); createTransferPairs(); } public TransferImpl(DenseUnionVector to) { this.to = to; internalStruct.makeTransferPair(to.internalStruct); createTransferPairs(); } private void createTransferPairs() { for (int i = 0; i < nextTypeId; i++) { ValueVector srcVec = internalStruct.getVectorById(i); ValueVector dstVec = to.internalStruct.getVectorById(i); to.typeFields[i] = typeFields[i]; to.typeMapFields[i] = typeMapFields[i]; to.childVectors[i] = dstVec; internalTransferPairs[i] = srcVec.makeTransferPair(dstVec); } } @Override public void transfer() { to.clear(); ReferenceManager refManager = typeBuffer.getReferenceManager(); to.typeBuffer = refManager.transferOwnership(typeBuffer, to.allocator).getTransferredBuffer(); refManager = offsetBuffer.getReferenceManager(); to.offsetBuffer = refManager.transferOwnership(offsetBuffer, to.allocator).getTransferredBuffer(); for (int i = 0; i < nextTypeId; i++) { if (internalTransferPairs[i] != null) { internalTransferPairs[i].transfer(); to.childVectors[i] = internalTransferPairs[i].getTo(); } } to.valueCount = valueCount; clear(); } @Override public void splitAndTransfer(int startIndex, int length) { to.clear(); // transfer type buffer int startPoint = startIndex * TYPE_WIDTH; int sliceLength = length * TYPE_WIDTH; ArrowBuf slicedBuffer = typeBuffer.slice(startPoint, sliceLength); ReferenceManager refManager = slicedBuffer.getReferenceManager(); to.typeBuffer = refManager.transferOwnership(slicedBuffer, to.allocator).getTransferredBuffer(); // transfer offset buffer while (to.offsetBuffer.capacity() < (long) length * OFFSET_WIDTH) { to.reallocOffsetBuffer(); } int [] typeCounts = new int[nextTypeId]; int [] typeStarts = new int[nextTypeId]; for (int i = 0; i < typeCounts.length; i++) { typeCounts[i] = 0; typeStarts[i] = -1; } for (int i = startIndex; i < startIndex + length; i++) { byte typeId = typeBuffer.getByte(i); if (typeId >= 0) { to.offsetBuffer.setInt((long) (i - startIndex) * OFFSET_WIDTH, typeCounts[typeId]); typeCounts[typeId] += 1; if (typeStarts[typeId] == -1) { typeStarts[typeId] = offsetBuffer.getInt((long) i * OFFSET_WIDTH); } } } // transfer vector values for (int i = 0; i < nextTypeId; i++) { if (typeCounts[i] > 0 && typeStarts[i] != -1) { internalTransferPairs[i].splitAndTransfer(typeStarts[i], typeCounts[i]); to.childVectors[i] = internalTransferPairs[i].getTo(); } } to.setValueCount(length); } @Override public ValueVector getTo() { return to; } @Override public void copyValueSafe(int from, int to) { this.to.copyFrom(from, to, DenseUnionVector.this); } } @Override public FieldReader getReader() { if (reader == null) { reader = new DenseUnionReader(this); } return reader; } public FieldWriter getWriter() { if (writer == null) { writer = new DenseUnionWriter(this); } return writer; } @Override public int getBufferSize() { return this.getBufferSizeFor(this.valueCount); } @Override public int getBufferSizeFor(final int count) { if (count == 0) { return 0; } int[] counts = new int[Byte.MAX_VALUE + 1]; for (int i = 0; i < count; i++) { byte typeId = getTypeId(i); if (typeId != -1) { counts[typeId] += 1; } } long childBytes = 0; for (int typeId = 0; typeId < childVectors.length; typeId++) { ValueVector childVector = childVectors[typeId]; if (childVector != null) { childBytes += childVector.getBufferSizeFor(counts[typeId]); } } return (int) (count * TYPE_WIDTH + (long) count * OFFSET_WIDTH + childBytes); } @Override public ArrowBuf[] getBuffers(boolean clear) { List list = new java.util.ArrayList<>(); setReaderAndWriterIndex(); if (getBufferSize() != 0) { list.add(typeBuffer); list.add(offsetBuffer); list.addAll(java.util.Arrays.asList(internalStruct.getBuffers(clear))); } if (clear) { valueCount = 0; typeBuffer.getReferenceManager().retain(); typeBuffer.close(); typeBuffer = allocator.getEmpty(); offsetBuffer.getReferenceManager().retain(); offsetBuffer.close(); offsetBuffer = allocator.getEmpty(); } return list.toArray(new ArrowBuf[list.size()]); } @Override public Iterator iterator() { return internalStruct.iterator(); } private ValueVector getVector(int index) { byte typeId = typeBuffer.getByte(index * TYPE_WIDTH); return getVectorByType(typeId); } public Object getObject(int index) { ValueVector vector = getVector(index); if (vector != null) { int offset = offsetBuffer.getInt((long) index * OFFSET_WIDTH); return vector.isNull(offset) ? null : vector.getObject(offset); } return null; } public void get(int index, DenseUnionHolder holder) { FieldReader reader = new DenseUnionReader(DenseUnionVector.this); reader.setPosition(index); holder.reader = reader; } public int getValueCount() { return valueCount; } /** * IMPORTANT: Union types always return non null as there is no validity buffer. * * To check validity correctly you must check the underlying vector. */ public boolean isNull(int index) { return false; } @Override public int getNullCount() { return 0; } public int isSet(int index) { return isNull(index) ? 0 : 1; } DenseUnionWriter writer; public void setValueCount(int valueCount) { this.valueCount = valueCount; while (valueCount > getTypeBufferValueCapacity()) { reallocTypeBuffer(); reallocOffsetBuffer(); } setChildVectorValueCounts(); } private void setChildVectorValueCounts() { int [] counts = new int[Byte.MAX_VALUE + 1]; for (int i = 0; i < this.valueCount; i++) { byte typeId = getTypeId(i); if (typeId != -1) { counts[typeId] += 1; } } for (int i = 0; i < nextTypeId; i++) { childVectors[typeMapFields[i]].setValueCount(counts[typeMapFields[i]]); } } public void setSafe(int index, DenseUnionHolder holder) { FieldReader reader = holder.reader; if (writer == null) { writer = new DenseUnionWriter(DenseUnionVector.this); } int offset = offsetBuffer.getInt((long) index * OFFSET_WIDTH); MinorType type = reader.getMinorType(); writer.setPosition(offset); byte typeId = holder.typeId; switch (type) { <#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.typeParams?? || minor.class?starts_with("Decimal")> case ${name?upper_case}: Nullable${name}Holder ${uncappedName}Holder = new Nullable${name}Holder(); reader.read(${uncappedName}Holder); setSafe(index, ${uncappedName}Holder); break; case STRUCT: case LIST: { setTypeId(index, typeId); ComplexCopier.copy(reader, writer); break; } default: throw new UnsupportedOperationException(); } } <#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.typeParams?? || minor.class?starts_with("Decimal")> public void setSafe(int index, Nullable${name}Holder holder) { while (index >= getOffsetBufferValueCapacity()) { reallocOffsetBuffer(); } byte typeId = getTypeId(index); ${name}Vector vector = get${name}Vector(typeId<#if minor.class?starts_with("Decimal")>, new ArrowType.Decimal(holder.precision, holder.scale, holder.WIDTH * 8)); int offset = vector.getValueCount(); vector.setValueCount(offset + 1); vector.setSafe(offset, holder); offsetBuffer.setInt((long) index * OFFSET_WIDTH, offset); } public void setTypeId(int index, byte typeId) { while (index >= getTypeBufferValueCapacity()) { reallocTypeBuffer(); } typeBuffer.setByte(index * TYPE_WIDTH , typeId); } private int getTypeBufferValueCapacity() { return (int) typeBuffer.capacity() / TYPE_WIDTH; } public void setOffset(int index, int offset) { while (index >= getOffsetBufferValueCapacity()) { reallocOffsetBuffer(); } offsetBuffer.setInt((long) index * OFFSET_WIDTH, offset); } private long getOffsetBufferValueCapacity() { return offsetBuffer.capacity() / OFFSET_WIDTH; } @Override public int hashCode(int index, ArrowBufHasher hasher) { if (isNull(index)) { return 0; } int offset = offsetBuffer.getInt((long) index * OFFSET_WIDTH); return getVector(index).hashCode(offset, hasher); } @Override public int hashCode(int index) { return hashCode(index, SimpleHasher.INSTANCE); } @Override public OUT accept(VectorVisitor visitor, IN value) { return visitor.visit(this, value); } @Override public String getName() { return name; } private void setNegative(long start, long end) { for (long i = start;i < end; i++) { typeBuffer.setByte(i, -1); } } @Override public T addOrGet(String name, FieldType fieldType, Class clazz) { return internalStruct.addOrGet(name, fieldType, clazz); } @Override public T getChild(String name, Class clazz) { return internalStruct.getChild(name, clazz); } @Override public VectorWithOrdinal getChildVectorWithOrdinal(String name) { return internalStruct.getChildVectorWithOrdinal(name); } @Override public int size() { return internalStruct.size(); } @Override public void setInitialCapacity(int valueCount, double density) { for (final ValueVector vector : internalStruct) { if (vector instanceof DensityAwareVector) { ((DensityAwareVector) vector).setInitialCapacity(valueCount, density); } else { vector.setInitialCapacity(valueCount); } } } /** * Set the element at the given index to null. For DenseUnionVector, it throws an UnsupportedOperationException * as nulls are not supported at the top level and isNull() always returns false. * * @param index position of element * @throws UnsupportedOperationException whenever invoked */ @Override public void setNull(int index) { throw new UnsupportedOperationException("The method setNull() is not supported on DenseUnionVector."); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/DenseUnionWriter.java000066400000000000000000000236621515451615200276350ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.impl.NullableStructWriterFactory; import org.apache.arrow.vector.types.Types; <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/DenseUnionWriter.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; <#include "/@includes/vv_imports.ftl" /> import org.apache.arrow.vector.complex.writer.BaseWriter; import org.apache.arrow.vector.types.Types.MinorType; /* * This class is generated using freemarker and the ${.template_name} template. */ @SuppressWarnings("unused") public class DenseUnionWriter extends AbstractFieldWriter implements FieldWriter { DenseUnionVector data; private BaseWriter[] writers = new BaseWriter[Byte.MAX_VALUE + 1]; private final NullableStructWriterFactory nullableStructWriterFactory; public DenseUnionWriter(DenseUnionVector vector) { this(vector, NullableStructWriterFactory.getNullableStructWriterFactoryInstance()); } public DenseUnionWriter(DenseUnionVector vector, NullableStructWriterFactory nullableStructWriterFactory) { data = vector; this.nullableStructWriterFactory = nullableStructWriterFactory; } @Override public void setPosition(int index) { super.setPosition(index); for (BaseWriter writer : writers) { if (writer != null) { writer.setPosition(index); } } } @Override public void start() { byte typeId = data.getTypeId(idx()); getStructWriter((byte) idx()).start(); } @Override public void end() { byte typeId = data.getTypeId(idx()); getStructWriter(typeId).end(); } @Override public void startList() { byte typeId = data.getTypeId(idx()); getListWriter(typeId).startList(); } @Override public void endList() { byte typeId = data.getTypeId(idx()); getListWriter(typeId).endList(); } @Override public void startListView() { byte typeId = data.getTypeId(idx()); getListViewWriter(typeId).startList(); } @Override public void endListView() { byte typeId = data.getTypeId(idx()); getListViewWriter(typeId).endList(); } private StructWriter getStructWriter(byte typeId) { StructWriter structWriter = (StructWriter) writers[typeId]; if (structWriter == null) { structWriter = nullableStructWriterFactory.build((StructVector) data.getVectorByType(typeId)); writers[typeId] = structWriter; } return structWriter; } public StructWriter asStruct(byte typeId) { data.setTypeId(idx(), typeId); return getStructWriter(typeId); } private ListWriter getListWriter(byte typeId) { ListWriter listWriter = (ListWriter) writers[typeId]; if (listWriter == null) { listWriter = new UnionListWriter((ListVector) data.getVectorByType(typeId), nullableStructWriterFactory); writers[typeId] = listWriter; } return listWriter; } private ListWriter getListViewWriter(byte typeId) { ListWriter listWriter = (ListWriter) writers[typeId]; if (listWriter == null) { listWriter = new UnionListViewWriter((ListViewVector) data.getVectorByType(typeId), nullableStructWriterFactory); writers[typeId] = listWriter; } return listWriter; } public ListWriter asList(byte typeId) { data.setTypeId(idx(), typeId); return getListWriter(typeId); } private MapWriter getMapWriter(byte typeId) { MapWriter mapWriter = (MapWriter) writers[typeId]; if (mapWriter == null) { mapWriter = new UnionMapWriter((MapVector) data.getVectorByType(typeId)); writers[typeId] = mapWriter; } return mapWriter; } public MapWriter asMap(byte typeId) { data.setTypeId(idx(), typeId); return getMapWriter(typeId); } BaseWriter getWriter(byte typeId) { MinorType minorType = data.getVectorByType(typeId).getMinorType(); switch (minorType) { case STRUCT: return getStructWriter(typeId); case LIST: return getListWriter(typeId); case MAP: return getMapWriter(typeId); <#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.typeParams?? || minor.class?starts_with("Decimal")> case ${name?upper_case}: return get${name}Writer(typeId); default: throw new UnsupportedOperationException("Unknown type: " + minorType); } } <#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.typeParams?? || minor.class?starts_with("Decimal")> private ${name}Writer get${name}Writer(byte typeId) { ${name}Writer writer = (${name}Writer) writers[typeId]; if (writer == null) { writer = new ${name}WriterImpl((${name}Vector) data.getVectorByType(typeId)); writers[typeId] = writer; } return writer; } public ${name}Writer as${name}(byte typeId) { data.setTypeId(idx(), typeId); return get${name}Writer(typeId); } @Override public void write(${name}Holder holder) { throw new UnsupportedOperationException(); } public void write${minor.class}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, , byte typeId<#if minor.class?starts_with("Decimal")>, ArrowType arrowType) { data.setTypeId(idx(), typeId); get${name}Writer(typeId).setPosition(data.getOffset(idx())); get${name}Writer(typeId).write${name}(<#list fields as field>${field.name}<#if field_has_next>, <#if minor.class?starts_with("Decimal")>, arrowType); } public void writeNull() { } @Override public StructWriter struct() { byte typeId = data.getTypeId(idx()); data.setTypeId(idx(), typeId); getListWriter(typeId).setPosition(data.getOffset(idx())); return getListWriter(typeId).struct(); } @Override public ListWriter list() { byte typeId = data.getTypeId(idx()); data.setTypeId(idx(), typeId); getListWriter(typeId).setPosition(data.getOffset(idx())); return getListWriter(typeId).list(); } @Override public ListWriter list(String name) { byte typeId = data.getTypeId(idx()); data.setTypeId(idx(), typeId); getStructWriter(typeId).setPosition(data.getOffset(idx())); return getStructWriter(typeId).list(name); } @Override public MapWriter map() { byte typeId = data.getTypeId(idx()); data.setTypeId(idx(), typeId); getListWriter(typeId).setPosition(data.getOffset(idx())); return getMapWriter(typeId).map(); } @Override public MapWriter map(String name) { byte typeId = data.getTypeId(idx()); data.setTypeId(idx(), typeId); getStructWriter(typeId).setPosition(data.getOffset(idx())); return getStructWriter(typeId).map(name); } @Override public MapWriter map(String name, boolean keysSorted) { byte typeId = data.getTypeId(idx()); data.setTypeId(idx(), typeId); getStructWriter(typeId).setPosition(data.getOffset(idx())); return getStructWriter(typeId).map(name, keysSorted); } @Override public StructWriter struct(String name) { byte typeId = data.getTypeId(idx()); data.setTypeId(idx(), typeId); getStructWriter(typeId).setPosition(data.getOffset(idx())); return getStructWriter(typeId).struct(name); } <#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 /> <#assign capName = minor.class?cap_first /> <#if !minor.typeParams?? || minor.class?starts_with("Decimal") > @Override public ${capName}Writer ${lowerName}(String name) { byte typeId = data.getTypeId(idx()); data.setTypeId(idx(), typeId); getStructWriter(typeId).setPosition(data.getOffset(idx())); return getStructWriter(typeId).${lowerName}(name); } @Override public ${capName}Writer ${lowerName}() { byte typeId = data.getTypeId(idx()); data.setTypeId(idx(), typeId); getListWriter(typeId).setPosition(data.getOffset(idx())); return getListWriter(typeId).${lowerName}(); } <#if minor.class?starts_with("Decimal")> public ${capName}Writer ${lowerName}(String name<#list minor.typeParams as typeParam>, ${typeParam.type} ${typeParam.name}) { byte typeId = data.getTypeId(idx()); data.setTypeId(idx(), typeId); getStructWriter(typeId).setPosition(data.getOffset(idx())); return getStructWriter(typeId).${lowerName}(name<#list minor.typeParams as typeParam>, ${typeParam.name}); } @Override public void allocate() { data.allocateNew(); } @Override public void clear() { data.clear(); } @Override public void close() throws Exception { data.close(); } @Override public Field getField() { return data.getField(); } @Override public int getValueCapacity() { return data.getValueCapacity(); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/HolderReaderImpl.java000066400000000000000000000134121515451615200275430ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ <@pp.dropOutputFile /> <#list vv.types as type> <#list type.minor as minor> <#list ["", "Nullable"] as holderMode> <#assign nullMode = holderMode /> <#assign lowerName = minor.class?uncap_first /> <#if lowerName == "int" ><#assign lowerName = "integer" /> <#assign name = minor.class?cap_first /> <#assign javaType = (minor.javaType!type.javaType) /> <#assign friendlyType = (minor.friendlyType!minor.boxedType!type.boxedType) /> <#assign safeType=friendlyType /> <#if safeType=="byte[]"><#assign safeType="ByteArray" /> <#assign fields = (minor.fields!type.fields) + minor.typeParams![]/> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/${holderMode}${name}HolderReaderImpl.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; <#include "/@includes/vv_imports.ftl" /> // Source code generated using FreeMarker template ${.template_name} @SuppressWarnings("unused") public class ${holderMode}${name}HolderReaderImpl extends AbstractFieldReader { private ${nullMode}${name}Holder holder; public ${holderMode}${name}HolderReaderImpl(${holderMode}${name}Holder 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 next on a single value reader."); } @Override public MinorType getMinorType() { return MinorType.${name?upper_case}; } @Override public boolean isSet() { <#if holderMode == "Nullable"> return this.holder.isSet == 1; <#else> return true; } @Override public void read(${name}Holder h) { <#list fields as field> h.${field.name} = holder.${field.name}; } @Override public void read(Nullable${name}Holder h) { <#list fields as field> h.${field.name} = holder.${field.name}; h.isSet = isSet() ? 1 : 0; } // read friendly type @Override public ${friendlyType} read${safeType}() { <#if nullMode == "Nullable"> if (!isSet()) { return null; } <#if type.major == "VarLen"> <#if type.width == 4> int length = holder.end - holder.start; <#elseif type.width == 8> int length = (int) (holder.end - holder.start); byte[] value = new byte [length]; holder.buffer.getBytes(holder.start, value, 0, length); <#if minor.class == "VarBinary" || minor.class == "LargeVarBinary" || minor.class == "ViewVarBinary"> return value; <#elseif minor.class == "VarChar" || minor.class == "LargeVarChar" || minor.class == "ViewVarChar"> Text text = new Text(); text.set(value); return text; <#elseif minor.class == "IntervalDay"> return Duration.ofDays(holder.days).plusMillis(holder.milliseconds); <#elseif minor.class == "IntervalYear"> return Period.ofMonths(holder.value); <#elseif minor.class == "IntervalMonthDayNano"> return new PeriodDuration(Period.ofMonths(holder.months).plusDays(holder.days), Duration.ofNanos(holder.nanoseconds)); <#elseif minor.class == "Duration"> return DurationVector.toDuration(holder.value, holder.unit); <#elseif minor.class == "Bit" > return new Boolean(holder.value != 0); <#elseif minor.class == "Decimal"> byte[] bytes = new byte[${type.width}]; holder.buffer.getBytes(holder.start, bytes, 0, ${type.width}); ${friendlyType} value = new BigDecimal(new BigInteger(bytes), holder.scale); return value; <#elseif minor.class == "Decimal256"> byte[] bytes = new byte[${type.width}]; holder.buffer.getBytes(holder.start, bytes, 0, ${type.width}); ${friendlyType} value = new BigDecimal(new BigInteger(bytes), holder.scale); return value; <#elseif minor.class == "FixedSizeBinary"> byte[] value = new byte [holder.byteWidth]; holder.buffer.getBytes(0, value, 0, holder.byteWidth); return value; <#elseif minor.class == "TimeStampSec"> final long millis = java.util.concurrent.TimeUnit.SECONDS.toMillis(holder.value); return DateUtility.getLocalDateTimeFromEpochMilli(millis); <#elseif minor.class == "TimeStampMilli" || minor.class == "DateMilli" || minor.class == "TimeMilli"> return DateUtility.getLocalDateTimeFromEpochMilli(holder.value); <#elseif minor.class == "TimeStampMicro"> return DateUtility.getLocalDateTimeFromEpochMicro(holder.value); <#elseif minor.class == "TimeStampNano"> return DateUtility.getLocalDateTimeFromEpochNano(holder.value); <#else> ${friendlyType} value = new ${friendlyType}(this.holder.value); return value; } @Override public Object readObject() { return read${safeType}(); } <#if nullMode != "Nullable"> public void copyAsValue(${minor.class?cap_first}Writer writer){ writer.write(holder); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/NullReader.java000066400000000000000000000077121515451615200264240ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ import org.apache.arrow.vector.types.pojo.ArrowType.Null; import org.apache.arrow.vector.types.pojo.Field; <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/NullReader.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; <#include "/@includes/vv_imports.ftl" /> /** * Source code generated using FreeMarker template ${.template_name} */ @SuppressWarnings("unused") public class NullReader extends AbstractBaseReader implements FieldReader{ public static final NullReader INSTANCE = new NullReader(); public static final NullReader EMPTY_LIST_INSTANCE = new NullReader(MinorType.NULL); public static final NullReader EMPTY_STRUCT_INSTANCE = new NullReader(MinorType.STRUCT); private MinorType type; private NullReader(){ super(); type = MinorType.NULL; } private NullReader(MinorType type){ super(); this.type = type; } @Override public MinorType getMinorType() { return type; } @Override public Field getField() { return new Field("", FieldType.nullable(new Null()), null); } public void copyAsValue(StructWriter writer) {} public void copyAsValue(ListWriter writer) {} public void copyAsValue(UnionWriter writer) {} <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> public void read(${name}Holder holder){ throw new UnsupportedOperationException("NullReader cannot write into non-nullable holder"); } public void read(Nullable${name}Holder holder){ holder.isSet = 0; } public void read(int arrayIndex, ${name}Holder holder){ throw new ArrayIndexOutOfBoundsException(); } public void copyAsValue(${minor.class}Writer writer){} public void copyAsField(String name, ${minor.class}Writer writer){} public void read(int arrayIndex, Nullable${name}Holder holder){ throw new ArrayIndexOutOfBoundsException(); } public void read(ExtensionHolder holder) { holder.isSet = 0; } public int size(){ return 0; } public boolean isSet(){ return false; } public boolean next(){ return false; } public RepeatedStructReader struct(){ return this; } public RepeatedListReader list(){ return this; } public StructReader struct(String name){ return this; } public ListReader list(String name){ return this; } public FieldReader reader(String name){ return this; } public FieldReader reader(){ return this; } private void fail(String name){ throw new IllegalArgumentException(String.format("You tried to read a %s type when you are using a ValueReader of type %s.", name, this.getClass().getSimpleName())); } <#list ["Object", "BigDecimal", "Short", "Integer", "Long", "Boolean", "LocalDateTime", "Duration", "Period", "Double", "Float", "Character", "Text", "String", "Byte", "byte[]", "PeriodDuration"] as friendlyType> <#assign safeType=friendlyType /> <#if safeType=="byte[]"><#assign safeType="ByteArray" /> public ${friendlyType} read${safeType}(int arrayIndex){ return null; } public ${friendlyType} read${safeType}(){ return null; } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/PromotableViewWriter.java000066400000000000000000000137301515451615200305200ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/PromotableViewWriter.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; import java.util.Locale; <#include "/@includes/vv_imports.ftl" /> /** * This FieldWriter implementation delegates all FieldWriter API calls to an inner FieldWriter. This * inner field writer can start as a specific type, and this class will promote the writer to a * UnionWriter if a call is made that the specifically typed writer cannot handle. A new UnionVector * is created, wrapping the original vector, and replaces the original vector in the parent vector, * which can be either an AbstractStructVector or a ListViewVector. * *

    The writer used can either be for single elements (struct) or lists. */ public class PromotableViewWriter extends PromotableWriter { public PromotableViewWriter(ValueVector v, FixedSizeListVector fixedListVector) { super(v, fixedListVector); } public PromotableViewWriter(ValueVector v, FixedSizeListVector fixedListVector, NullableStructWriterFactory nullableStructWriterFactory) { super(v, fixedListVector, nullableStructWriterFactory); } public PromotableViewWriter(ValueVector v, LargeListVector largeListVector) { super(v, largeListVector); } public PromotableViewWriter(ValueVector v, LargeListVector largeListVector, NullableStructWriterFactory nullableStructWriterFactory) { super(v, largeListVector, nullableStructWriterFactory); } public PromotableViewWriter(ValueVector v, ListVector listVector) { super(v, listVector); } public PromotableViewWriter(ValueVector v, ListVector listVector, NullableStructWriterFactory nullableStructWriterFactory) { super(v, listVector, nullableStructWriterFactory); } public PromotableViewWriter(ValueVector v, ListViewVector listViewVector, NullableStructWriterFactory nullableStructWriterFactory) { super(v, listViewVector, nullableStructWriterFactory); } public PromotableViewWriter(ValueVector v, LargeListViewVector largeListViewVector) { super(v, largeListViewVector); } public PromotableViewWriter(ValueVector v, LargeListViewVector largeListViewVector, NullableStructWriterFactory nullableStructWriterFactory) { super(v, largeListViewVector, nullableStructWriterFactory); } public PromotableViewWriter(ValueVector v, AbstractStructVector parentContainer) { super(v, parentContainer); } public PromotableViewWriter(ValueVector v, AbstractStructVector parentContainer, NullableStructWriterFactory nullableStructWriterFactory) { super(v, parentContainer, nullableStructWriterFactory); } @Override protected FieldWriter getWriter(MinorType type, ArrowType arrowType) { if (state == State.UNION) { if (requiresArrowType(type)) { writer = ((UnionWriter) writer).toViewWriter(); ((UnionViewWriter) writer).getWriter(type, arrowType); } else { writer = ((UnionWriter) writer).toViewWriter(); ((UnionViewWriter) writer).getWriter(type); } } else if (state == State.UNTYPED) { if (type == null) { // ??? return null; } if (arrowType == null) { arrowType = type.getType(); } FieldType fieldType = new FieldType(addVectorAsNullable, arrowType, null, null); ValueVector v; if (listVector != null) { v = listVector.addOrGetVector(fieldType).getVector(); } else if (fixedListVector != null) { v = fixedListVector.addOrGetVector(fieldType).getVector(); } else if (listViewVector != null) { v = listViewVector.addOrGetVector(fieldType).getVector(); } else if (largeListVector != null) { v = largeListVector.addOrGetVector(fieldType).getVector(); } else { v = largeListViewVector.addOrGetVector(fieldType).getVector(); } v.allocateNew(); setWriter(v); writer.setPosition(position); } else if (type != this.type) { promoteToUnion(); if (requiresArrowType(type)) { writer = ((UnionWriter) writer).toViewWriter(); ((UnionViewWriter) writer).getWriter(type, arrowType); } else { writer = ((UnionWriter) writer).toViewWriter(); ((UnionViewWriter) writer).getWriter(type); } } return writer; } @Override public StructWriter struct() { return getWriter(MinorType.LISTVIEW).struct(); } <#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 /> <#assign capName = minor.class?cap_first /> @Override public ${capName}Writer ${lowerName}() { return getWriter(MinorType.LISTVIEW).${lowerName}(); } @Override public void allocate() { getWriter().allocate(); } @Override public void clear() { getWriter().clear(); } @Override public Field getField() { return getWriter().getField(); } @Override public int getValueCapacity() { return getWriter().getValueCapacity(); } @Override public void close() throws Exception { getWriter().close(); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/PromotableWriter.java000066400000000000000000000445631515451615200276750ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/PromotableWriter.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; import java.util.Locale; <#include "/@includes/vv_imports.ftl" /> /** * This FieldWriter implementation delegates all FieldWriter API calls to an inner FieldWriter. This * inner field writer can start as a specific type, and this class will promote the writer to a * UnionWriter if a call is made that the specifically typed writer cannot handle. A new UnionVector * is created, wrapping the original vector, and replaces the original vector in the parent vector, * which can be either an AbstractStructVector or a ListVector. * *

    The writer used can either be for single elements (struct) or lists. */ public class PromotableWriter extends AbstractPromotableFieldWriter { protected final AbstractStructVector parentContainer; protected final ListVector listVector; protected final ListViewVector listViewVector; protected final FixedSizeListVector fixedListVector; protected final LargeListVector largeListVector; protected final LargeListViewVector largeListViewVector; protected final NullableStructWriterFactory nullableStructWriterFactory; protected int position; protected static final int MAX_DECIMAL_PRECISION = 38; protected static final int MAX_DECIMAL256_PRECISION = 76; protected enum State { UNTYPED, SINGLE, UNION } protected MinorType type; protected ValueVector vector; protected UnionVector unionVector; protected State state; protected FieldWriter writer; /** * Constructs a new instance. * * @param v The vector to write. * @param parentContainer The parent container for the vector. */ public PromotableWriter(ValueVector v, AbstractStructVector parentContainer) { this(v, parentContainer, NullableStructWriterFactory.getNullableStructWriterFactoryInstance()); } /** * Constructs a new instance. * * @param v The vector to initialize the writer with. * @param parentContainer The parent container for the vector. * @param nullableStructWriterFactory The factory to create the delegate writer. */ public PromotableWriter( ValueVector v, AbstractStructVector parentContainer, NullableStructWriterFactory nullableStructWriterFactory) { this.parentContainer = parentContainer; this.listVector = null; this.listViewVector = null; this.fixedListVector = null; this.largeListVector = null; this.largeListViewVector = null; this.nullableStructWriterFactory = nullableStructWriterFactory; init(v); } /** * Constructs a new instance. * * @param v The vector to initialize the writer with. * @param listVector The vector that serves as a parent of v. */ public PromotableWriter(ValueVector v, ListVector listVector) { this(v, listVector, NullableStructWriterFactory.getNullableStructWriterFactoryInstance()); } /** * Constructs a new instance. * * @param v The vector to initialize the writer with. * @param fixedListVector The vector that serves as a parent of v. */ public PromotableWriter(ValueVector v, FixedSizeListVector fixedListVector) { this(v, fixedListVector, NullableStructWriterFactory.getNullableStructWriterFactoryInstance()); } /** * Constructs a new instance. * * @param v The vector to initialize the writer with. * @param largeListVector The vector that serves as a parent of v. */ public PromotableWriter(ValueVector v, LargeListVector largeListVector) { this(v, largeListVector, NullableStructWriterFactory.getNullableStructWriterFactoryInstance()); } /** * Constructs a new instance. * * @param v The vector to initialize the writer with. * @param listViewVector The vector that serves as a parent of v. */ public PromotableWriter(ValueVector v, ListViewVector listViewVector) { this(v, listViewVector, NullableStructWriterFactory.getNullableStructWriterFactoryInstance()); } /** * Constructs a new instance. * * @param v The vector to initialize the writer with. * @param largeListViewVector The vector that serves as a parent of v. */ public PromotableWriter(ValueVector v, LargeListViewVector largeListViewVector) { this(v, largeListViewVector, NullableStructWriterFactory.getNullableStructWriterFactoryInstance()); } /** * Constructs a new instance. * * @param v The vector to initialize the writer with. * @param listVector The vector that serves as a parent of v. * @param nullableStructWriterFactory The factory to create the delegate writer. */ public PromotableWriter( ValueVector v, ListVector listVector, NullableStructWriterFactory nullableStructWriterFactory) { this.listVector = listVector; this.listViewVector = null; this.parentContainer = null; this.fixedListVector = null; this.largeListVector = null; this.largeListViewVector = null; this.nullableStructWriterFactory = nullableStructWriterFactory; init(v); } /** * Constructs a new instance. * * @param v The vector to initialize the writer with. * @param listViewVector The vector that serves as a parent of v. * @param nullableStructWriterFactory The factory to create the delegate writer. */ public PromotableWriter( ValueVector v, ListViewVector listViewVector, NullableStructWriterFactory nullableStructWriterFactory) { this.listViewVector = listViewVector; this.listVector = null; this.parentContainer = null; this.fixedListVector = null; this.largeListVector = null; this.largeListViewVector = null; this.nullableStructWriterFactory = nullableStructWriterFactory; init(v); } /** * Constructs a new instance. * * @param v The vector to initialize the writer with. * @param fixedListVector The vector that serves as a parent of v. * @param nullableStructWriterFactory The factory to create the delegate writer. */ public PromotableWriter( ValueVector v, FixedSizeListVector fixedListVector, NullableStructWriterFactory nullableStructWriterFactory) { this.fixedListVector = fixedListVector; this.parentContainer = null; this.listVector = null; this.listViewVector = null; this.largeListVector = null; this.largeListViewVector = null; this.nullableStructWriterFactory = nullableStructWriterFactory; init(v); } /** * Constructs a new instance. * * @param v The vector to initialize the writer with. * @param largeListVector The vector that serves as a parent of v. * @param nullableStructWriterFactory The factory to create the delegate writer. */ public PromotableWriter( ValueVector v, LargeListVector largeListVector, NullableStructWriterFactory nullableStructWriterFactory) { this.largeListVector = largeListVector; this.fixedListVector = null; this.parentContainer = null; this.listVector = null; this.listViewVector = null; this.largeListViewVector = null; this.nullableStructWriterFactory = nullableStructWriterFactory; init(v); } /** * Constructs a new instance. * * @param v The vector to initialize the writer with. * @param largeListViewVector The vector that serves as a parent of v. * @param nullableStructWriterFactory The factory to create the delegate writer. */ public PromotableWriter( ValueVector v, LargeListViewVector largeListViewVector, NullableStructWriterFactory nullableStructWriterFactory) { this.largeListViewVector = largeListViewVector; this.fixedListVector = null; this.parentContainer = null; this.listVector = null; this.listViewVector = null; this.largeListVector = null; this.nullableStructWriterFactory = nullableStructWriterFactory; init(v); } private void init(ValueVector v) { if (v instanceof UnionVector) { state = State.UNION; unionVector = (UnionVector) v; writer = new UnionWriter(unionVector, nullableStructWriterFactory); } else if (v instanceof NullVector) { state = State.UNTYPED; } else { setWriter(v); } } @Override public void setAddVectorAsNullable(boolean nullable) { super.setAddVectorAsNullable(nullable); if (writer instanceof AbstractFieldWriter) { ((AbstractFieldWriter) writer).setAddVectorAsNullable(nullable); } } protected void setWriter(ValueVector v) { state = State.SINGLE; vector = v; type = v.getMinorType(); switch (type) { case STRUCT: writer = nullableStructWriterFactory.build((StructVector) vector); break; case LIST: writer = new UnionListWriter((ListVector) vector, nullableStructWriterFactory); break; case LISTVIEW: writer = new UnionListViewWriter((ListViewVector) vector, nullableStructWriterFactory); break; case MAP: writer = new UnionMapWriter((MapVector) vector); break; 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; } } @Override public void writeNull() { FieldWriter w = getWriter(); if (w != null) { w.writeNull(); } setPosition(idx() + 1); } @Override public void setPosition(int index) { super.setPosition(index); FieldWriter w = getWriter(); if (w == null) { position = index; } else { w.setPosition(index); } } protected boolean requiresArrowType(MinorType type) { return type == MinorType.DECIMAL || type == MinorType.MAP || type == MinorType.DURATION || type == MinorType.FIXEDSIZEBINARY || type == MinorType.EXTENSIONTYPE || (type.name().startsWith("TIMESTAMP") && type.name().endsWith("TZ")); } @Override protected FieldWriter getWriter(MinorType type, ArrowType arrowType) { if (state == State.UNION) { if (requiresArrowType(type)) { ((UnionWriter) writer).getWriter(type, arrowType); } else { ((UnionWriter) writer).getWriter(type); } } else if (state == State.UNTYPED) { if (type == null) { // ??? return null; } if (arrowType == null) { arrowType = type.getType(); } FieldType fieldType = new FieldType(addVectorAsNullable, arrowType, null, null); ValueVector v; if (listVector != null) { v = listVector.addOrGetVector(fieldType).getVector(); } else if (fixedListVector != null) { v = fixedListVector.addOrGetVector(fieldType).getVector(); } else if (listViewVector != null) { v = listViewVector.addOrGetVector(fieldType).getVector(); } else { v = largeListVector.addOrGetVector(fieldType).getVector(); } v.allocateNew(); setWriter(v); writer.setPosition(position); } else if (type != this.type) { promoteToUnion(); if (requiresArrowType(type)) { ((UnionWriter) writer).getWriter(type, arrowType); } else { ((UnionWriter) writer).getWriter(type); } } return writer; } @Override public boolean isEmptyStruct() { return writer.isEmptyStruct(); } @Override protected FieldWriter getWriter() { return writer; } protected FieldWriter promoteToUnion() { String name = vector.getField().getName(); TransferPair tp = vector.getTransferPair( vector.getMinorType().name().toLowerCase(Locale.ROOT), vector.getAllocator()); tp.transfer(); if (parentContainer != null) { // TODO allow dictionaries in complex types unionVector = parentContainer.addOrGetUnion(name); unionVector.allocateNew(); } else if (listVector != null) { unionVector = listVector.promoteToUnion(); } else if (fixedListVector != null) { unionVector = fixedListVector.promoteToUnion(); } else if (largeListVector != null) { unionVector = largeListVector.promoteToUnion(); } else if (listViewVector != null) { unionVector = listViewVector.promoteToUnion(); } unionVector.addVector((FieldVector) tp.getTo()); writer = new UnionWriter(unionVector, nullableStructWriterFactory); writer.setPosition(idx()); for (int i = 0; i <= idx(); i++) { unionVector.setType(i, vector.getMinorType()); } vector = null; state = State.UNION; return writer; } @Override public void write(DecimalHolder holder) { getWriter( MinorType.DECIMAL, new ArrowType.Decimal(MAX_DECIMAL_PRECISION, holder.scale, /*bitWidth=*/ 128)) .write(holder); } @Override public void writeDecimal(long start, ArrowBuf buffer, ArrowType arrowType) { getWriter( MinorType.DECIMAL, new ArrowType.Decimal( MAX_DECIMAL_PRECISION, ((ArrowType.Decimal) arrowType).getScale(), /*bitWidth=*/ 128)) .writeDecimal(start, buffer, arrowType); } @Override public void writeDecimal(BigDecimal value) { getWriter( MinorType.DECIMAL, new ArrowType.Decimal(MAX_DECIMAL_PRECISION, value.scale(), /*bitWidth=*/ 128)) .writeDecimal(value); } @Override public void writeBigEndianBytesToDecimal(byte[] value, ArrowType arrowType) { getWriter( MinorType.DECIMAL, new ArrowType.Decimal( MAX_DECIMAL_PRECISION, ((ArrowType.Decimal) arrowType).getScale(), /*bitWidth=*/ 128)) .writeBigEndianBytesToDecimal(value, arrowType); } @Override public void write(Decimal256Holder holder) { getWriter( MinorType.DECIMAL256, new ArrowType.Decimal(MAX_DECIMAL256_PRECISION, holder.scale, /*bitWidth=*/ 256)) .write(holder); } @Override public void writeDecimal256(long start, ArrowBuf buffer, ArrowType arrowType) { getWriter( MinorType.DECIMAL256, new ArrowType.Decimal( MAX_DECIMAL256_PRECISION, ((ArrowType.Decimal) arrowType).getScale(), /*bitWidth=*/ 256)) .writeDecimal256(start, buffer, arrowType); } @Override public void writeDecimal256(BigDecimal value) { getWriter( MinorType.DECIMAL256, new ArrowType.Decimal(MAX_DECIMAL256_PRECISION, value.scale(), /*bitWidth=*/ 256)) .writeDecimal256(value); } @Override public void writeBigEndianBytesToDecimal256(byte[] value, ArrowType arrowType) { getWriter( MinorType.DECIMAL256, new ArrowType.Decimal( MAX_DECIMAL256_PRECISION, ((ArrowType.Decimal) arrowType).getScale(), /*bitWidth=*/ 256)) .writeBigEndianBytesToDecimal256(value, arrowType); } @Override public void writeVarBinary(byte[] value) { getWriter(MinorType.VARBINARY).writeVarBinary(value); } @Override public void writeVarBinary(byte[] value, int offset, int length) { getWriter(MinorType.VARBINARY).writeVarBinary(value, offset, length); } @Override public void writeVarBinary(ByteBuffer value) { getWriter(MinorType.VARBINARY).writeVarBinary(value); } @Override public void writeVarBinary(ByteBuffer value, int offset, int length) { getWriter(MinorType.VARBINARY).writeVarBinary(value, offset, length); } @Override public void writeLargeVarBinary(byte[] value) { getWriter(MinorType.LARGEVARBINARY).writeLargeVarBinary(value); } @Override public void writeLargeVarBinary(byte[] value, int offset, int length) { getWriter(MinorType.LARGEVARBINARY).writeLargeVarBinary(value, offset, length); } @Override public void writeLargeVarBinary(ByteBuffer value) { getWriter(MinorType.LARGEVARBINARY).writeLargeVarBinary(value); } @Override public void writeLargeVarBinary(ByteBuffer value, int offset, int length) { getWriter(MinorType.LARGEVARBINARY).writeLargeVarBinary(value, offset, length); } @Override public void writeVarChar(Text value) { getWriter(MinorType.VARCHAR).writeVarChar(value); } @Override public void writeVarChar(String value) { getWriter(MinorType.VARCHAR).writeVarChar(value); } @Override public void writeLargeVarChar(Text value) { getWriter(MinorType.LARGEVARCHAR).writeLargeVarChar(value); } @Override 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(); } @Override public void clear() { getWriter().clear(); } @Override public Field getField() { return getWriter().getField(); } @Override public int getValueCapacity() { return getWriter().getValueCapacity(); } @Override public void close() throws Exception { getWriter().close(); } /** * Convert the writer to a PromotableViewWriter. * * @return The writer as a PromotableViewWriter. */ public PromotableViewWriter toViewWriter() { PromotableViewWriter promotableViewWriter = new PromotableViewWriter(unionVector, parentContainer, nullableStructWriterFactory); promotableViewWriter.position = position; promotableViewWriter.writer = writer; promotableViewWriter.state = state; promotableViewWriter.unionVector = unionVector; promotableViewWriter.type = MinorType.LISTVIEW; return promotableViewWriter; } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/StructWriters.java000066400000000000000000000332321515451615200272270ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ <@pp.dropOutputFile /> <#list ["Nullable", "Single"] as mode> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/${mode}StructWriter.java" /> <#assign index = "idx()"> <#if mode == "Single"> <#assign containerClass = "NonNullableStructVector" /> <#else> <#assign containerClass = "StructVector" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; <#include "/@includes/vv_imports.ftl" /> import java.util.Map; import java.util.HashMap; import org.apache.arrow.vector.holders.RepeatedStructHolder; import org.apache.arrow.vector.AllocationHelper; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.complex.writer.FieldWriter; <#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. */ @SuppressWarnings("unused") public class ${mode}StructWriter extends AbstractFieldWriter { protected final ${containerClass} container; private int initialCapacity; private final Map fields = new HashMap<>(); public ${mode}StructWriter(${containerClass} container) { <#if mode == "Single"> if (container instanceof StructVector) { throw new IllegalArgumentException("Invalid container: " + container); } this.container = container; this.initialCapacity = 0; for (Field child : container.getField().getChildren()) { MinorType minorType = Types.getMinorTypeForArrowType(child.getType()); addVectorAsNullable = child.isNullable(); switch (minorType) { case STRUCT: struct(child.getName()); break; case LIST: list(child.getName()); break; case LISTVIEW: listView(child.getName()); break; case MAP: { ArrowType.Map arrowType = (ArrowType.Map) child.getType(); map(child.getName(), arrowType.getKeysSorted()); break; } case DENSEUNION: { FieldType fieldType = new FieldType(addVectorAsNullable, MinorType.DENSEUNION.getType(), null, null); DenseUnionWriter writer = new DenseUnionWriter(container.addOrGet(child.getName(), fieldType, DenseUnionVector.class), getNullableStructWriterFactory()); 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()); fields.put(handleCase(child.getName()), writer); break; <#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 /> case ${upperName}: { <#if minor.typeParams?? > ${minor.arrowType} arrowType = (${minor.arrowType})child.getType(); ${lowerName}(child.getName()<#list minor.typeParams as typeParam>, arrowType.get${typeParam.name?cap_first}()); <#else> ${lowerName}(child.getName()); break; } default: throw new UnsupportedOperationException("Unknown type: " + minorType); } } } protected String handleCase(final String input) { return input.toLowerCase(); } protected NullableStructWriterFactory getNullableStructWriterFactory() { return NullableStructWriterFactory.getNullableStructWriterFactoryInstance(); } @Override public int getValueCapacity() { return container.getValueCapacity(); } public void setInitialCapacity(int initialCapacity) { this.initialCapacity = initialCapacity; container.setInitialCapacity(initialCapacity); } @Override public boolean isEmptyStruct() { return 0 == container.size(); } @Override public Field getField() { return container.getField(); } @Override public StructWriter struct(String name) { String finalName = handleCase(name); FieldWriter writer = fields.get(finalName); if(writer == null){ int vectorCount=container.size(); FieldType fieldType = new FieldType(addVectorAsNullable, MinorType.STRUCT.getType(), null, null); StructVector vector = container.addOrGet(name, fieldType, StructVector.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.STRUCT); } } 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(); container.close(); } @Override public void allocate() { container.allocateNew(); for(final FieldWriter w : fields.values()) { w.allocate(); } } @Override public void clear() { container.clear(); for(final FieldWriter w : fields.values()) { w.clear(); } } @Override public ListWriter list(String name) { String finalName = handleCase(name); FieldWriter writer = fields.get(finalName); int vectorCount = container.size(); if(writer == null) { FieldType fieldType = new FieldType(addVectorAsNullable, MinorType.LIST.getType(), null, null); writer = new PromotableWriter(container.addOrGet(name, fieldType, ListVector.class), container, getNullableStructWriterFactory()); if (container.size() > vectorCount) { writer.allocate(); } writer.setPosition(idx()); fields.put(finalName, writer); } else { if (writer instanceof PromotableWriter) { // ensure writers are initialized ((PromotableWriter)writer).getWriter(MinorType.LIST); } } return writer; } @Override public ListWriter listView(String name) { String finalName = handleCase(name); FieldWriter writer = fields.get(finalName); int vectorCount = container.size(); if(writer == null) { FieldType fieldType = new FieldType(addVectorAsNullable, MinorType.LISTVIEW.getType(), null, null); writer = new PromotableViewWriter(container.addOrGet(name, fieldType, ListViewVector.class), container, getNullableStructWriterFactory()); if (container.size() > vectorCount) { writer.allocate(); } writer.setPosition(idx()); fields.put(finalName, writer); } else { if (writer instanceof PromotableViewWriter) { // ensure writers are initialized ((PromotableViewWriter) writer).getWriter(MinorType.LISTVIEW); } else { writer = ((PromotableWriter) writer).toViewWriter(); ((PromotableViewWriter) writer).getWriter(MinorType.LISTVIEW); } } return writer; } @Override public MapWriter map(String name) { return map(name, false); } @Override public MapWriter map(String name, boolean keysSorted) { FieldWriter writer = fields.get(handleCase(name)); if(writer == null) { ValueVector vector; ValueVector currentVector = container.getChild(name); MapVector v = container.addOrGet(name, new FieldType(addVectorAsNullable, new ArrowType.Map(keysSorted) ,null, null), MapVector.class); writer = new PromotableWriter(v, container, getNullableStructWriterFactory()); vector = v; if (currentVector == null || currentVector != vector) { if(this.initialCapacity > 0) { vector.setInitialCapacity(this.initialCapacity); } vector.allocateNewSafe(); } writer.setPosition(idx()); fields.put(handleCase(name), writer); } else { if (writer instanceof PromotableWriter) { // ensure writers are initialized ((PromotableWriter)writer).getWriter(MinorType.MAP, new ArrowType.Map(keysSorted)); } } return writer; } public void setValueCount(int count) { container.setValueCount(count); } @Override public void setPosition(int index) { super.setPosition(index); for(final FieldWriter w: fields.values()) { w.setPosition(index); } } <#if mode="Nullable"> @Override public void writeNull() { container.setNull(idx()); setValueCount(idx()+1); super.setPosition(idx()+1); } @Override public void start() { <#if mode == "Single"> <#else> container.setIndexDefined(idx()); } @Override public void end() { setPosition(idx()+1); } <#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 /> <#assign capName = minor.class?cap_first /> <#assign vectName = capName /> <#if minor.typeParams?? > @Override public ${minor.class}Writer ${lowerName}(String name) { // returns existing writer final FieldWriter writer = fields.get(handleCase(name)); Preconditions.checkNotNull(writer); return writer; } @Override public ${minor.class}Writer ${lowerName}(String name<#list minor.typeParams as typeParam>, ${typeParam.type} ${typeParam.name}) { <#else> @Override public ${minor.class}Writer ${lowerName}(String name) { FieldWriter writer = fields.get(handleCase(name)); if(writer == null) { ValueVector vector; ValueVector currentVector = container.getChild(name); ${vectName}Vector v = container.addOrGet(name, new FieldType(addVectorAsNullable, <#if minor.typeParams??> <#if minor.arrowTypeConstructorParams??> <#assign constructorParams = minor.arrowTypeConstructorParams /> <#else> <#assign constructorParams = [] /> <#list minor.typeParams?reverse as typeParam> <#assign constructorParams = constructorParams + [ typeParam.name ] /> new ${minor.arrowType}(${constructorParams?join(", ")}<#if minor.class?starts_with("Decimal")>, ${vectName}Vector.TYPE_WIDTH * 8) <#else> MinorType.${upperName}.getType() ,null, null), ${vectName}Vector.class); writer = new PromotableWriter(v, container, getNullableStructWriterFactory()); vector = v; if (currentVector == null || currentVector != vector) { if(this.initialCapacity > 0) { vector.setInitialCapacity(this.initialCapacity); } vector.allocateNewSafe(); } writer.setPosition(idx()); fields.put(handleCase(name), writer); } else { if (writer instanceof PromotableWriter) { // ensure writers are initialized <#if minor.class?starts_with("Decimal")> ((PromotableWriter)writer).getWriter(MinorType.${upperName}<#if minor.class?starts_with("Decimal")>, new ${minor.arrowType}(precision, scale, ${vectName}Vector.TYPE_WIDTH * 8)); <#elseif is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> <#if minor.arrowTypeConstructorParams??> <#assign constructorParams = minor.arrowTypeConstructorParams /> <#else> <#assign constructorParams = [] /> <#list minor.typeParams?reverse as typeParam> <#assign constructorParams = constructorParams + [ typeParam.name ] /> ArrowType arrowType = new ${minor.arrowType}(${constructorParams?join(", ")}); ((PromotableWriter)writer).getWriter(MinorType.${upperName}, arrowType); <#else> ((PromotableWriter)writer).getWriter(MinorType.${upperName}); } } return writer; } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java000066400000000000000000000237721515451615200313270ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.complex.writer.Decimal256Writer; import org.apache.arrow.vector.complex.writer.DecimalWriter; import org.apache.arrow.vector.holders.Decimal256Holder; import org.apache.arrow.vector.holders.DecimalHolder; import java.lang.UnsupportedOperationException; import java.math.BigDecimal; <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/UnionFixedSizeListWriter.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; <#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. */ @SuppressWarnings("unused") public class UnionFixedSizeListWriter extends AbstractFieldWriter { protected FixedSizeListVector vector; protected PromotableWriter writer; private boolean inStruct = false; private String structName; private final int listSize; public UnionFixedSizeListWriter(FixedSizeListVector vector) { this(vector, NullableStructWriterFactory.getNullableStructWriterFactoryInstance()); } public UnionFixedSizeListWriter(FixedSizeListVector vector, NullableStructWriterFactory nullableStructWriterFactory) { this.vector = vector; this.writer = new PromotableWriter(vector.getDataVector(), vector, nullableStructWriterFactory); this.listSize = vector.getListSize(); } public UnionFixedSizeListWriter(FixedSizeListVector vector, AbstractFieldWriter parent) { this(vector); } @Override public void allocate() { vector.allocateNew(); } @Override public void clear() { vector.clear(); } @Override public Field getField() { return vector.getField(); } public void setValueCount(int count) { vector.setValueCount(count); } @Override public int getValueCapacity() { return vector.getValueCapacity(); } @Override public void close() throws Exception { vector.close(); writer.close(); } @Override public void setPosition(int index) { super.setPosition(index); } <#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 ${minor.class}Writer ${lowerName}() { return this; } <#if minor.typeParams?? > @Override 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 ${minor.class}Writer ${lowerName}(String name) { structName = name; return writer.${lowerName}(name); } @Override public StructWriter struct() { inStruct = true; return this; } @Override public ListWriter list() { return writer; } @Override public ListWriter list(String name) { ListWriter listWriter = writer.list(name); return listWriter; } @Override public StructWriter struct(String name) { StructWriter structWriter = writer.struct(name); return structWriter; } @Override public MapWriter map() { return writer; } @Override public MapWriter map(String name) { MapWriter mapWriter = writer.map(name); return mapWriter; } @Override public MapWriter map(boolean keysSorted) { writer.map(keysSorted); return writer; } @Override public MapWriter map(String name, boolean keysSorted) { MapWriter mapWriter = writer.map(name, keysSorted); return mapWriter; } @Override public void startList() { int start = vector.startNewValue(idx()); writer.setPosition(start); } @Override public void endList() { setPosition(idx() + 1); } @Override public void start() { writer.start(); } @Override public void end() { writer.end(); inStruct = false; } @Override 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.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 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); } <#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.write(holder); writer.setPosition(writer.idx()+1); } <#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.write${name}(start, buffer, arrowType); writer.setPosition(writer.idx()+1); } @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.write(holder); writer.setPosition(writer.idx()+1); } @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.write${name}(value); writer.setPosition(writer.idx()+1); } @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.writeBigEndianBytesTo${name}(value, arrowType); writer.setPosition(writer.idx() + 1); } <#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.write${name}(<#list fields as field>holder.${field.name}<#if field_has_next>, ); writer.setPosition(writer.idx()+1); } <#if minor.class?ends_with("VarBinary")> @Override public void write${minor.class}(byte[] value) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } writer.write${minor.class}(value); writer.setPosition(writer.idx() + 1); } @Override public void write${minor.class}(byte[] value, int offset, int length) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } writer.write${minor.class}(value, offset, length); writer.setPosition(writer.idx() + 1); } @Override public void write${minor.class}(ByteBuffer value) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } writer.write${minor.class}(value); writer.setPosition(writer.idx() + 1); } @Override public void write${minor.class}(ByteBuffer value, int offset, int length) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } writer.write${minor.class}(value, offset, length); writer.setPosition(writer.idx() + 1); } <#elseif minor.class?ends_with("VarChar")> @Override public void write${minor.class}(Text value) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } writer.write${minor.class}(value); writer.setPosition(writer.idx() + 1); } @Override public void write${minor.class}(String value) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } writer.write${minor.class}(value); writer.setPosition(writer.idx() + 1); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/UnionListWriter.java000066400000000000000000000274151515451615200275120ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.complex.writer.Decimal256Writer; import org.apache.arrow.vector.complex.writer.DecimalWriter; import org.apache.arrow.vector.holders.Decimal256Holder; import org.apache.arrow.vector.holders.DecimalHolder; import java.lang.UnsupportedOperationException; import java.math.BigDecimal; <@pp.dropOutputFile /> <#list ["List", "ListView", "LargeList", "LargeListView"] as listName> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/Union${listName}Writer.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; <#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. */ @SuppressWarnings("unused") public class Union${listName}Writer extends AbstractFieldWriter { protected ${listName}Vector vector; protected PromotableWriter writer; 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> private static final int OFFSET_WIDTH = 4; <#if listName == "ListView"> private static final long SIZE_WIDTH = 4; <#if listName == "LargeListView"> private static final long SIZE_WIDTH = 8; public Union${listName}Writer(${listName}Vector vector) { this(vector, NullableStructWriterFactory.getNullableStructWriterFactoryInstance()); } public Union${listName}Writer(${listName}Vector vector, NullableStructWriterFactory nullableStructWriterFactory) { this.vector = vector; <#if listName = "ListView" || listName = "LargeListView"> this.writer = new PromotableViewWriter(vector.getDataVector(), vector, nullableStructWriterFactory); <#else> this.writer = new PromotableWriter(vector.getDataVector(), vector, nullableStructWriterFactory); } public Union${listName}Writer(${listName}Vector vector, AbstractFieldWriter parent) { this(vector); } @Override public void allocate() { vector.allocateNew(); } @Override public void clear() { vector.clear(); } @Override public Field getField() { return vector.getField(); } public void setValueCount(int count) { vector.setValueCount(count); } @Override public int getValueCapacity() { return vector.getValueCapacity(); } @Override public void close() throws Exception { vector.close(); writer.close(); } @Override public void setPosition(int index) { super.setPosition(index); } <#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 ${minor.class}Writer ${lowerName}() { return this; } <#if minor.typeParams?? > @Override 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 ${minor.class}Writer ${lowerName}(String name) { structName = name; return writer.${lowerName}(name); } @Override public StructWriter struct() { inStruct = true; return this; } @Override public ListWriter list() { return writer; } @Override public ListWriter list(String name) { ListWriter listWriter = writer.list(name); return listWriter; } @Override public ListWriter listView() { return writer; } @Override public ListWriter listView(String name) { ListWriter listWriter = writer.listView(name); return listWriter; } @Override public StructWriter struct(String name) { StructWriter structWriter = writer.struct(name); return structWriter; } @Override public MapWriter map() { return writer; } @Override public MapWriter map(String name) { MapWriter mapWriter = writer.map(name); return mapWriter; } @Override public MapWriter map(boolean keysSorted) { writer.map(keysSorted); return writer; } @Override public MapWriter map(String name, boolean keysSorted) { MapWriter mapWriter = writer.map(name, 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() { vector.startNewValue(idx()); writer.setPosition(checkedCastToInt(vector.getOffsetBuffer().getLong((idx() + 1L) * OFFSET_WIDTH))); listStarted = true; } @Override public void endList() { vector.getOffsetBuffer().setLong((idx() + 1L) * OFFSET_WIDTH, writer.idx()); setPosition(idx() + 1); listStarted = false; } <#elseif listName == "ListView"> @Override public void startList() { vector.startNewValue(idx()); writer.setPosition(vector.getOffsetBuffer().getInt((idx()) * OFFSET_WIDTH)); listStarted = true; } @Override public void endList() { int sizeUptoIdx = 0; for (int i = 0; i < idx(); i++) { sizeUptoIdx += vector.getSizeBuffer().getInt(i * SIZE_WIDTH); } vector.getSizeBuffer().setInt(idx() * SIZE_WIDTH, writer.idx() - sizeUptoIdx); setPosition(idx() + 1); listStarted = false; } @Override public void startListView() { vector.startNewValue(idx()); writer.setPosition(vector.getOffsetBuffer().getInt((idx()) * OFFSET_WIDTH)); listStarted = true; } @Override public void endListView() { int sizeUptoIdx = 0; for (int i = 0; i < idx(); i++) { sizeUptoIdx += vector.getSizeBuffer().getInt(i * SIZE_WIDTH); } vector.getSizeBuffer().setInt(idx() * SIZE_WIDTH, writer.idx() - sizeUptoIdx); setPosition(idx() + 1); listStarted = false; } <#elseif listName == "LargeListView"> @Override public void startList() { vector.startNewValue(idx()); writer.setPosition(vector.getOffsetBuffer().getInt((idx()) * OFFSET_WIDTH)); listStarted = true; } @Override public void endList() { int sizeUptoIdx = 0; for (int i = 0; i < idx(); i++) { sizeUptoIdx += vector.getSizeBuffer().getInt(i * SIZE_WIDTH); } vector.getSizeBuffer().setInt(idx() * SIZE_WIDTH, writer.idx() - sizeUptoIdx); setPosition(idx() + 1); listStarted = false; } @Override public void startListView() { vector.startNewValue(idx()); writer.setPosition(checkedCastToInt(vector.getOffsetBuffer().getInt((idx()) * OFFSET_WIDTH))); listStarted = true; } @Override public void endListView() { int sizeUptoIdx = 0; for (int i = 0; i < idx(); i++) { sizeUptoIdx += vector.getSizeBuffer().getInt(i * SIZE_WIDTH); } vector.getSizeBuffer().setInt(idx() * SIZE_WIDTH, writer.idx() - sizeUptoIdx); setPosition(idx() + 1); listStarted = false; } <#else> @Override public void startList() { vector.startNewValue(idx()); writer.setPosition(vector.getOffsetBuffer().getInt((idx() + 1L) * OFFSET_WIDTH)); listStarted = true; } @Override public void endList() { vector.getOffsetBuffer().setInt((idx() + 1L) * OFFSET_WIDTH, writer.idx()); setPosition(idx() + 1); listStarted = false; } @Override public void start() { writer.start(); } @Override public void end() { writer.end(); inStruct = false; } @Override public void writeNull() { if (!listStarted){ vector.setNull(idx()); } else { writer.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 /> <#assign fields = minor.fields!type.fields /> <#assign uncappedName = name?uncap_first/> @Override public void write${name}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, ) { writer.write${name}(<#list fields as field>${field.name}<#if field_has_next>, ); writer.setPosition(writer.idx()+1); } <#if is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> @Override public void write(${name}Holder holder) { writer.write(holder); writer.setPosition(writer.idx()+1); } <#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); } @Override public void write(${name}Holder holder) { writer.write(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); } <#else> @Override public void write(${name}Holder holder) { writer.write${name}(<#list fields as field>holder.${field.name}<#if field_has_next>, ); writer.setPosition(writer.idx()+1); } <#if minor.class?ends_with("VarBinary")> @Override public void write${minor.class}(byte[] value) { writer.write${minor.class}(value); writer.setPosition(writer.idx() + 1); } @Override public void write${minor.class}(byte[] value, int offset, int length) { writer.write${minor.class}(value, offset, length); writer.setPosition(writer.idx() + 1); } @Override public void write${minor.class}(ByteBuffer value) { writer.write${minor.class}(value); writer.setPosition(writer.idx() + 1); } @Override public void write${minor.class}(ByteBuffer value, int offset, int length) { writer.write${minor.class}(value, offset, length); writer.setPosition(writer.idx() + 1); } <#elseif minor.class?ends_with("VarChar")> @Override public void write${minor.class}(Text value) { writer.write${minor.class}(value); writer.setPosition(writer.idx() + 1); } @Override public void write${minor.class}(String value) { writer.write${minor.class}(value); writer.setPosition(writer.idx() + 1); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/UnionMapWriter.java000066400000000000000000000172071515451615200273120ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.complex.writer.Decimal256Writer; import org.apache.arrow.vector.complex.writer.DecimalWriter; import org.apache.arrow.vector.holders.Decimal256Holder; import org.apache.arrow.vector.holders.DecimalHolder; import java.lang.UnsupportedOperationException; import java.math.BigDecimal; <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/UnionMapWriter.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; <#include "/@includes/vv_imports.ftl" /> /* * This class is generated using freemarker and the ${.template_name} template. */ /** *

    Writer for MapVectors. This extends UnionListWriter to simplify writing map entries to a list * of struct elements, with "key" and "value" fields. The procedure for writing a map begin with * {@link #startMap()} followed by {@link #startEntry()}. An entry is written by using the * {@link #key()} writer to write the key, then the {@link #value()} writer to write a value. After * writing the value, call {@link #endEntry()} to complete the entry. Each map can have 1 or more * entries. When done writing entries, call {@link #endMap()} to complete the map. * *

    NOTE: the MapVector can have NULL values by not writing to position. If a map is started with * {@link #startMap()}, then it must have a key written. The value of a map entry can be NULL by * not using the {@link #value()} writer. * *

    Example to write the following map to position 5 of a vector * 

    {@code
 *   // {
 *   //   1 -> 3,
 *   //   2 -> 4,
 *   //   3 -> NULL
 *   // }
 *
 *   UnionMapWriter writer = ...
 *
 *   writer.setPosition(5);
 *   writer.startMap();
 *   writer.startEntry();
 *   writer.key().integer().writeInt(1);
 *   writer.value().integer().writeInt(3);
 *   writer.endEntry();
 *   writer.startEntry();
 *   writer.key().integer().writeInt(2);
 *   writer.value().integer().writeInt(4);
 *   writer.endEntry();
 *   writer.startEntry();
 *   writer.key().integer().writeInt(3);
 *   writer.endEntry();
 *   writer.endMap();
 *
    *

    */ @SuppressWarnings("unused") public class UnionMapWriter extends UnionListWriter { /** * Current mode for writing map entries, set by calling {@link #key()} or {@link #value()} * and reset with a call to {@link #endEntry()}. With KEY mode, a struct writer with field * named "key" is returned. With VALUE mode, a struct writer with field named "value" is * returned. In OFF mode, the writer will behave like a standard UnionListWriter */ private enum MapWriteMode { OFF, KEY, VALUE, } private MapWriteMode mode = MapWriteMode.OFF; private StructWriter entryWriter; public UnionMapWriter(MapVector vector) { super(vector); entryWriter = struct(); } /** Start writing a map that consists of 1 or more entries. */ public void startMap() { startList(); } /** Complete the map. */ public void endMap() { endList(); } /** * Start a map entry that should be followed by calls to {@link #key()} and {@link #value()} * writers. Call {@link #endEntry()} to complete the entry. */ public void startEntry() { writer.setAddVectorAsNullable(false); entryWriter.start(); } /** Complete the map entry. */ public void endEntry() { entryWriter.end(); mode = MapWriteMode.OFF; writer.setAddVectorAsNullable(true); } /** Return the key writer that is used to write to the "key" field. */ public UnionMapWriter key() { writer.setAddVectorAsNullable(false); mode = MapWriteMode.KEY; return this; } /** Return the value writer that is used to write to the "value" field. */ public UnionMapWriter value() { writer.setAddVectorAsNullable(true); mode = MapWriteMode.VALUE; return this; } <#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?? > @Override public ${name}Writer ${uncappedName}() { switch (mode) { case KEY: return entryWriter.${uncappedName}(MapVector.KEY_NAME); case VALUE: return entryWriter.${uncappedName}(MapVector.VALUE_NAME); default: return this; } } @Override public DecimalWriter decimal() { switch (mode) { case KEY: return entryWriter.decimal(MapVector.KEY_NAME); case VALUE: return entryWriter.decimal(MapVector.VALUE_NAME); default: return this; } } @Override public Decimal256Writer decimal256() { switch (mode) { case KEY: return entryWriter.decimal256(MapVector.KEY_NAME); case VALUE: return entryWriter.decimal256(MapVector.VALUE_NAME); default: return this; } } @Override public StructWriter struct() { switch (mode) { case KEY: return entryWriter.struct(MapVector.KEY_NAME); case VALUE: return entryWriter.struct(MapVector.VALUE_NAME); default: return super.struct(); } } @Override public ListWriter list() { switch (mode) { case KEY: return entryWriter.list(MapVector.KEY_NAME); case VALUE: return entryWriter.list(MapVector.VALUE_NAME); default: return super.list(); } } @Override public MapWriter map(boolean keysSorted) { switch (mode) { case KEY: return entryWriter.map(MapVector.KEY_NAME, keysSorted); case VALUE: return entryWriter.map(MapVector.VALUE_NAME, keysSorted); default: return super.map(); } } @Override public MapWriter map() { switch (mode) { case KEY: return entryWriter.map(MapVector.KEY_NAME); case VALUE: return entryWriter.map(MapVector.VALUE_NAME); default: 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; } } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/UnionReader.java000066400000000000000000000172161515451615200266020ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/UnionReader.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex.impl; <#include "/@includes/vv_imports.ftl" /> <#function is_timestamp_tz type> <#return type?starts_with("TimeStamp") && type?ends_with("TZ")> /** * Source code generated using FreeMarker template ${.template_name} */ @SuppressWarnings("unused") public class UnionReader extends AbstractFieldReader { private static final int NUM_SUPPORTED_TYPES = 51; private BaseReader[] readers = new BaseReader[NUM_SUPPORTED_TYPES]; public UnionVector data; public UnionReader(UnionVector data) { this.data = data; } public MinorType getMinorType() { return TYPES[data.getTypeValue(idx())]; } private static MinorType[] TYPES = new MinorType[NUM_SUPPORTED_TYPES]; static { for (MinorType minorType : MinorType.values()) { TYPES[minorType.ordinal()] = minorType; } } @Override public Field getField() { return data.getField(); } public boolean isSet(){ return !data.isNull(idx()); } public void read(UnionHolder holder) { holder.reader = this; holder.isSet = this.isSet() ? 1 : 0; } public void read(int index, UnionHolder holder) { getList().read(index, 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) { return reader; } switch (MinorType.values()[typeValue]) { case NULL: return NullReader.INSTANCE; case STRUCT: return (FieldReader) getStruct(); case LIST: return (FieldReader) getList(); case LISTVIEW: return (FieldReader) getListView(); case MAP: return (FieldReader) getMap(); <#list vv.types as type> <#list type.minor as minor> <#assign name = minor.class?cap_first /> <#assign uncappedName = name?uncap_first/> <#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> case ${name?upper_case}: return (FieldReader) get${name}(); 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() { if (structReader == null) { structReader = (SingleStructReaderImpl) data.getStruct().getReader(); structReader.setPosition(idx()); readers[MinorType.STRUCT.ordinal()] = structReader; } return structReader; } private UnionListReader listReader; private FieldReader getList() { if (listReader == null) { listReader = new UnionListReader(data.getList()); listReader.setPosition(idx()); readers[MinorType.LIST.ordinal()] = listReader; } return listReader; } private UnionListViewReader listViewReader; private FieldReader getListView() { if (listViewReader == null) { listViewReader = new UnionListViewReader(data.getListView()); listViewReader.setPosition(idx()); readers[MinorType.LISTVIEW.ordinal()] = listViewReader; } return listViewReader; } private UnionMapReader mapReader; private FieldReader getMap() { if (mapReader == null) { mapReader = new UnionMapReader(data.getMap()); mapReader.setPosition(idx()); readers[MinorType.MAP.ordinal()] = mapReader; } return mapReader; } @Override public java.util.Iterator iterator() { return getStruct().iterator(); } @Override public void copyAsValue(UnionWriter writer) { writer.data.copyFrom(idx(), writer.idx(), data); } <#list ["Object", "BigDecimal", "Short", "Integer", "Long", "Boolean", "LocalDateTime", "Duration", "Period", "Double", "Float", "Character", "Text", "Byte", "byte[]", "PeriodDuration"] as friendlyType> <#assign safeType=friendlyType /> <#if safeType=="byte[]"><#assign safeType="ByteArray" /> @Override public ${friendlyType} read${safeType}() { return getReaderForIndex(idx()).read${safeType}(); } public int size() { return getReaderForIndex(idx()).size(); } <#list vv.types as type> <#list type.minor as minor> <#assign name = minor.class?cap_first /> <#assign uncappedName = name?uncap_first/> <#assign boxedType = (minor.boxedType!type.boxedType) /> <#assign javaType = (minor.javaType!type.javaType) /> <#assign friendlyType = (minor.friendlyType!minor.boxedType!type.boxedType) /> <#assign safeType=friendlyType /> <#if safeType=="byte[]"><#assign safeType="ByteArray" /> <#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> private ${name}ReaderImpl ${uncappedName}Reader; private ${name}ReaderImpl get${name}() { if (${uncappedName}Reader == null) { ${uncappedName}Reader = new ${name}ReaderImpl(data.get${name}Vector()); ${uncappedName}Reader.setPosition(idx()); readers[MinorType.${name?upper_case}.ordinal()] = ${uncappedName}Reader; } return ${uncappedName}Reader; } public void read(Nullable${name}Holder holder){ getReaderForIndex(idx()).read(holder); } public void copyAsValue(${name}Writer writer){ getReaderForIndex(idx()).copyAsValue(writer); } @Override public void copyAsValue(ListWriter writer) { ComplexCopier.copy(this, (FieldWriter) writer); } @Override public void setPosition(int index) { super.setPosition(index); for (BaseReader reader : readers) { if (reader != null) { reader.setPosition(index); } } } public FieldReader reader(String name){ return getStruct().reader(name); } public FieldReader reader() { return getList().reader(); } public boolean next() { return getReaderForIndex(idx()).next(); } public void read(ExtensionHolder holder){ getReaderForIndex(idx(), holder.type()).read(holder); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/UnionVector.java000066400000000000000000000745631515451615200266520ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReferenceManager; import org.apache.arrow.memory.util.CommonUtil; import org.apache.arrow.memory.util.hash.ArrowBufHasher; 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; import org.apache.arrow.vector.complex.NonNullableStructVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.UnionMode; import org.apache.arrow.vector.compare.RangeEqualsVisitor; 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.CallBack; import org.apache.arrow.vector.util.DataSizeRoundingUtil; <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/UnionVector.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.complex; <#include "/@includes/vv_imports.ftl" /> import java.util.ArrayList; import java.util.Collections; import java.util.Iterator; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.util.CommonUtil; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.impl.ComplexCopier; import org.apache.arrow.vector.util.CallBack; import org.apache.arrow.vector.util.ValueVectorUtility; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.memory.util.ArrowBufPointer; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.vector.BaseValueVector; import org.apache.arrow.vector.util.OversizedAllocationException; import org.apache.arrow.util.Preconditions; import static org.apache.arrow.vector.types.UnionMode.Sparse; import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import static org.apache.arrow.memory.util.LargeMemoryUtil.capAtMaxInt; <#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. */ @SuppressWarnings("unused") /** * A vector which can hold values of different types. It does so by using a StructVector which contains a vector for each * primitive type that is stored. StructVector is used in order to take advantage of its serialization/deserialization methods, * as well as the addOrGet method. * * For performance reasons, UnionVector stores a cached reference to each subtype vector, to avoid having to do the struct lookup * each time the vector is accessed. * Source code generated using FreeMarker template ${.template_name} */ public class UnionVector extends AbstractContainerVector implements FieldVector { int valueCount; NonNullableStructVector internalStruct; protected ArrowBuf typeBuffer; private StructVector structVector; private ListVector listVector; private ListViewVector listViewVector; private MapVector mapVector; private FieldReader reader; private int singleType = 0; private ValueVector singleVector; private int typeBufferAllocationSizeInBytes; private final FieldType fieldType; private final Field[] typeIds = new Field[Byte.MAX_VALUE + 1]; public static final byte TYPE_WIDTH = 1; private static final FieldType INTERNAL_STRUCT_TYPE = new FieldType(false /*nullable*/, ArrowType.Struct.INSTANCE, null /*dictionary*/, null /*metadata*/); public static UnionVector empty(String name, BufferAllocator allocator) { FieldType fieldType = FieldType.nullable(new ArrowType.Union( UnionMode.Sparse, null)); return new UnionVector(name, allocator, fieldType, null); } public UnionVector(String name, BufferAllocator allocator, FieldType fieldType, CallBack callBack) { super(name, allocator, callBack); this.fieldType = fieldType; this.internalStruct = new NonNullableStructVector( "internal", allocator, INTERNAL_STRUCT_TYPE, callBack, AbstractStructVector.ConflictPolicy.CONFLICT_REPLACE, false); this.typeBuffer = allocator.getEmpty(); this.typeBufferAllocationSizeInBytes = BaseValueVector.INITIAL_VALUE_ALLOCATION * TYPE_WIDTH; } public BufferAllocator getAllocator() { return allocator; } @Override public MinorType getMinorType() { return MinorType.UNION; } @Override public void initializeChildrenFromFields(List children) { int count = 0; for (Field child: children) { int typeId = Types.getMinorTypeForArrowType(child.getType()).ordinal(); if (this.fieldType != null) { int[] typeIds = ((ArrowType.Union)this.fieldType.getType()).getTypeIds(); if (typeIds != null) { typeId = typeIds[count++]; } } typeIds[typeId] = child; } internalStruct.initializeChildrenFromFields(children); } @Override public List getChildrenFromFields() { return internalStruct.getChildrenFromFields(); } @Override public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers) { if (ownBuffers.size() != 1) { throw new IllegalArgumentException("Illegal buffer count, expected 1, got: " + ownBuffers.size()); } ArrowBuf buffer = ownBuffers.get(0); typeBuffer.getReferenceManager().release(); typeBuffer = buffer.getReferenceManager().retain(buffer, allocator); typeBufferAllocationSizeInBytes = checkedCastToInt(typeBuffer.capacity()); this.valueCount = fieldNode.getLength(); } @Override public List getFieldBuffers() { List result = new ArrayList<>(1); setReaderAndWriterIndex(); result.add(typeBuffer); return result; } private void setReaderAndWriterIndex() { typeBuffer.readerIndex(0); typeBuffer.writerIndex(valueCount * TYPE_WIDTH); } /** * Get the inner vectors. * * @deprecated This API will be removed as the current implementations no longer support inner vectors. * * @return the inner vectors for this field as defined by the TypeLayout */ @Deprecated @Override public List getFieldInnerVectors() { throw new UnsupportedOperationException("There are no inner vectors. Use geFieldBuffers"); } private String fieldName(MinorType type) { return type.name().toLowerCase(); } private FieldType fieldType(MinorType type) { return FieldType.nullable(type.getType()); } private T addOrGet(Types.MinorType minorType, Class c) { return addOrGet(null, minorType, c); } private T addOrGet(String name, Types.MinorType minorType, ArrowType arrowType, Class c) { return internalStruct.addOrGet(name == null ? fieldName(minorType) : name, FieldType.nullable(arrowType), c); } private T addOrGet(String name, Types.MinorType minorType, Class c) { return internalStruct.addOrGet(name == null ? fieldName(minorType) : name, fieldType(minorType), c); } @Override public long getValidityBufferAddress() { throw new UnsupportedOperationException(); } public long getTypeBufferAddress() { return typeBuffer.memoryAddress(); } @Override public long getDataBufferAddress() { throw new UnsupportedOperationException(); } @Override public long getOffsetBufferAddress() { throw new UnsupportedOperationException(); } public ArrowBuf getTypeBuffer() { return typeBuffer; } @Override public ArrowBuf getValidityBuffer() { throw new UnsupportedOperationException(); } @Override public ArrowBuf getDataBuffer() { throw new UnsupportedOperationException(); } @Override public ArrowBuf getOffsetBuffer() { throw new UnsupportedOperationException(); } public StructVector getStruct() { if (structVector == null) { int vectorCount = internalStruct.size(); structVector = addOrGet(MinorType.STRUCT, StructVector.class); if (internalStruct.size() > vectorCount) { structVector.allocateNew(); if (callBack != null) { callBack.doWork(); } } } return structVector; } <#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/> <#assign lowerCaseName = name?lower_case/> <#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> private ${name}Vector ${uncappedName}Vector; <#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) { ${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; } public ${name}Vector get${name}Vector(ArrowType arrowType) { return get${name}Vector(null, arrowType); } public ${name}Vector get${name}Vector(String name, ArrowType arrowType) { if (${uncappedName}Vector == null) { int vectorCount = internalStruct.size(); ${uncappedName}Vector = addOrGet(name, MinorType.${name?upper_case}, arrowType, ${name}Vector.class); if (internalStruct.size() > vectorCount) { ${uncappedName}Vector.allocateNew(); if (callBack != null) { callBack.doWork(); } } } return ${uncappedName}Vector; } <#else> public ${name}Vector get${name}Vector() { return get${name}Vector(null); } public ${name}Vector get${name}Vector(String name) { if (${uncappedName}Vector == null) { int vectorCount = internalStruct.size(); ${uncappedName}Vector = addOrGet(name, MinorType.${name?upper_case}, ${name}Vector.class); if (internalStruct.size() > vectorCount) { ${uncappedName}Vector.allocateNew(); if (callBack != null) { callBack.doWork(); } } } return ${uncappedName}Vector; } public ListVector getList() { if (listVector == null) { int vectorCount = internalStruct.size(); listVector = addOrGet(MinorType.LIST, ListVector.class); if (internalStruct.size() > vectorCount) { listVector.allocateNew(); if (callBack != null) { callBack.doWork(); } } } return listVector; } public ListViewVector getListView() { if (listViewVector == null) { int vectorCount = internalStruct.size(); listViewVector = addOrGet(MinorType.LISTVIEW, ListViewVector.class); if (internalStruct.size() > vectorCount) { listViewVector.allocateNew(); if (callBack != null) { callBack.doWork(); } } } return listViewVector; } public MapVector getMap() { if (mapVector == null) { throw new IllegalArgumentException("No map present. Provide ArrowType argument to create a new vector"); } return mapVector; } public MapVector getMap(ArrowType arrowType) { return getMap(null, arrowType); } public MapVector getMap(String name, ArrowType arrowType) { if (mapVector == null) { int vectorCount = internalStruct.size(); mapVector = addOrGet(name, MinorType.MAP, arrowType, MapVector.class); if (internalStruct.size() > vectorCount) { mapVector.allocateNew(); if (callBack != null) { callBack.doWork(); } } } 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); } @Override public void allocateNew() throws OutOfMemoryException { /* new allocation -- clear the current buffers */ clear(); internalStruct.allocateNew(); try { allocateTypeBuffer(); } catch (Exception e) { clear(); throw e; } } @Override public boolean allocateNewSafe() { /* new allocation -- clear the current buffers */ clear(); boolean safe = internalStruct.allocateNewSafe(); if (!safe) { return false; } try { allocateTypeBuffer(); } catch (Exception e) { clear(); return false; } return true; } private void allocateTypeBuffer() { typeBuffer = allocator.buffer(typeBufferAllocationSizeInBytes); typeBuffer.readerIndex(0); typeBuffer.setZero(0, typeBuffer.capacity()); } @Override public void reAlloc() { internalStruct.reAlloc(); reallocTypeBuffer(); } private void reallocTypeBuffer() { final long currentBufferCapacity = typeBuffer.capacity(); long newAllocationSize = currentBufferCapacity * 2; if (newAllocationSize == 0) { if (typeBufferAllocationSizeInBytes > 0) { newAllocationSize = typeBufferAllocationSizeInBytes; } else { newAllocationSize = BaseValueVector.INITIAL_VALUE_ALLOCATION * TYPE_WIDTH * 2; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); assert newAllocationSize >= 1; if (newAllocationSize > BaseValueVector.MAX_ALLOCATION_SIZE) { throw new OversizedAllocationException("Unable to expand the buffer"); } final ArrowBuf newBuf = allocator.buffer(checkedCastToInt(newAllocationSize)); newBuf.setBytes(0, typeBuffer, 0, currentBufferCapacity); newBuf.setZero(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); typeBuffer.getReferenceManager().release(1); typeBuffer = newBuf; typeBufferAllocationSizeInBytes = (int)newAllocationSize; } @Override public void setInitialCapacity(int numRecords) { } @Override public int getValueCapacity() { return Math.min(getTypeBufferValueCapacity(), internalStruct.getValueCapacity()); } @Override public void close() { clear(); } @Override public void clear() { valueCount = 0; typeBuffer.getReferenceManager().release(); typeBuffer = allocator.getEmpty(); internalStruct.clear(); } @Override public void reset() { valueCount = 0; typeBuffer.setZero(0, typeBuffer.capacity()); internalStruct.reset(); } @Override public Field getField() { List childFields = new ArrayList<>(); List children = internalStruct.getChildren(); int[] typeIds = new int[children.size()]; for (ValueVector v : children) { typeIds[childFields.size()] = v.getMinorType().ordinal(); childFields.add(v.getField()); } FieldType fieldType; if (this.fieldType == null) { fieldType = FieldType.nullable(new ArrowType.Union(Sparse, typeIds)); } else { final UnionMode mode = ((ArrowType.Union)this.fieldType.getType()).getMode(); fieldType = new FieldType(this.fieldType.isNullable(), new ArrowType.Union(mode, typeIds), this.fieldType.getDictionary(), this.fieldType.getMetadata()); } return new Field(name, fieldType, childFields); } @Override public TransferPair getTransferPair(BufferAllocator allocator) { return getTransferPair(name, allocator); } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return getTransferPair(ref, allocator, null); } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return new org.apache.arrow.vector.complex.UnionVector.TransferImpl(ref, allocator, callBack); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return getTransferPair(field, allocator, null); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { return new org.apache.arrow.vector.complex.UnionVector.TransferImpl(field, allocator, callBack); } @Override public TransferPair makeTransferPair(ValueVector target) { return new TransferImpl((UnionVector) target); } @Override public void copyFrom(int inIndex, int outIndex, ValueVector from) { Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); UnionVector fromCast = (UnionVector) from; fromCast.getReader().setPosition(inIndex); getWriter().setPosition(outIndex); ComplexCopier.copy(fromCast.reader, writer); } @Override public void copyFromSafe(int inIndex, int outIndex, ValueVector from) { copyFrom(inIndex, outIndex, from); } public FieldVector addVector(FieldVector v) { final String name = v.getName().isEmpty() ? fieldName(v.getMinorType()) : v.getName(); Preconditions.checkState(internalStruct.getChild(name) == null, String.format("%s vector already exists", name)); final FieldVector newVector = internalStruct.addOrGet(name, v.getField().getFieldType(), v.getClass()); v.makeTransferPair(newVector).transfer(); internalStruct.putChild(name, newVector); if (callBack != null) { callBack.doWork(); } return newVector; } /** * Directly put a vector to internalStruct without creating a new one with same type. */ public void directAddVector(FieldVector v) { String name = fieldName(v.getMinorType()); Preconditions.checkState(internalStruct.getChild(name) == null, String.format("%s vector already exists", name)); internalStruct.putChild(name, v); if (callBack != null) { callBack.doWork(); } } private class TransferImpl implements TransferPair { private final TransferPair internalStructVectorTransferPair; private final UnionVector to; public TransferImpl(String name, BufferAllocator allocator, CallBack callBack) { to = new UnionVector(name, allocator, /* field type */ null, callBack); internalStructVectorTransferPair = internalStruct.makeTransferPair(to.internalStruct); } public TransferImpl(Field field, BufferAllocator allocator, CallBack callBack) { to = new UnionVector(field.getName(), allocator, null, callBack); internalStructVectorTransferPair = internalStruct.makeTransferPair(to.internalStruct); } public TransferImpl(UnionVector to) { this.to = to; internalStructVectorTransferPair = internalStruct.makeTransferPair(to.internalStruct); } @Override public void transfer() { to.clear(); ReferenceManager refManager = typeBuffer.getReferenceManager(); to.typeBuffer = refManager.transferOwnership(typeBuffer, to.allocator).getTransferredBuffer(); internalStructVectorTransferPair.transfer(); to.valueCount = valueCount; clear(); } @Override public void splitAndTransfer(int startIndex, int length) { Preconditions.checkArgument(startIndex >= 0 && length >= 0 && startIndex + length <= valueCount, "Invalid parameters startIndex: %s, length: %s for valueCount: %s", startIndex, length, valueCount); to.clear(); internalStructVectorTransferPair.splitAndTransfer(startIndex, length); final int startPoint = startIndex * TYPE_WIDTH; final int sliceLength = length * TYPE_WIDTH; final ArrowBuf slicedBuffer = typeBuffer.slice(startPoint, sliceLength); final ReferenceManager refManager = slicedBuffer.getReferenceManager(); to.typeBuffer = refManager.transferOwnership(slicedBuffer, to.allocator).getTransferredBuffer(); to.setValueCount(length); } @Override public ValueVector getTo() { return to; } @Override public void copyValueSafe(int from, int to) { this.to.copyFrom(from, to, UnionVector.this); } } @Override public FieldReader getReader() { if (reader == null) { reader = new UnionReader(this); } return reader; } public FieldWriter getWriter() { if (writer == null) { writer = new UnionWriter(this); } return writer; } @Override public int getBufferSize() { if (valueCount == 0) { return 0; } return (valueCount * TYPE_WIDTH) + internalStruct.getBufferSize(); } @Override public int getBufferSizeFor(final int valueCount) { if (valueCount == 0) { return 0; } long bufferSize = 0; for (final ValueVector v : (Iterable) this) { bufferSize += v.getBufferSizeFor(valueCount); } return (int) bufferSize + (valueCount * TYPE_WIDTH); } @Override public ArrowBuf[] getBuffers(boolean clear) { List list = new java.util.ArrayList<>(); setReaderAndWriterIndex(); if (getBufferSize() != 0) { list.add(typeBuffer); list.addAll(java.util.Arrays.asList(internalStruct.getBuffers(clear))); } if (clear) { valueCount = 0; typeBuffer.getReferenceManager().retain(); typeBuffer.getReferenceManager().release(); typeBuffer = allocator.getEmpty(); } return list.toArray(new ArrowBuf[list.size()]); } @Override public Iterator iterator() { return internalStruct.iterator(); } public ValueVector getVector(int index) { return getVector(index, null); } public ValueVector getVector(int index, ArrowType arrowType) { int type = typeBuffer.getByte(index * TYPE_WIDTH); return getVectorByType(type, arrowType); } public ValueVector getVectorByType(int typeId) { return getVectorByType(typeId, null); } public ValueVector getVectorByType(int typeId, ArrowType arrowType) { Field type = typeIds[typeId]; Types.MinorType minorType; String name = null; if (type == null) { minorType = Types.MinorType.values()[typeId]; } else { minorType = Types.getMinorTypeForArrowType(type.getType()); name = type.getName(); } switch (minorType) { case NULL: return null; <#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.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> case ${name?upper_case}: return get${name}Vector(name<#if minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary">, arrowType); case STRUCT: return getStruct(); case LIST: return getList(); case LISTVIEW: return getListView(); case MAP: return getMap(name, arrowType); case EXTENSIONTYPE: return getExtension(arrowType); default: throw new UnsupportedOperationException("Cannot support type: " + MinorType.values()[typeId]); } } public Object getObject(int index) { ValueVector vector = getVector(index); if (vector != null) { return vector.isNull(index) ? null : vector.getObject(index); } return null; } public byte[] get(int index) { return null; } public void get(int index, ComplexHolder holder) { } public void get(int index, UnionHolder holder) { FieldReader reader = new UnionReader(UnionVector.this); reader.setPosition(index); holder.reader = reader; } public int getValueCount() { return valueCount; } /** * IMPORTANT: Union types always return non null as there is no validity buffer. * * To check validity correctly you must check the underlying vector. */ public boolean isNull(int index) { return false; } @Override public int getNullCount() { return 0; } public int isSet(int index) { return isNull(index) ? 0 : 1; } UnionWriter writer; public void setValueCount(int valueCount) { this.valueCount = valueCount; while (valueCount > getTypeBufferValueCapacity()) { reallocTypeBuffer(); } internalStruct.setValueCount(valueCount); } public void setSafe(int index, UnionHolder holder) { setSafe(index, holder, null); } public void setSafe(int index, UnionHolder holder, ArrowType arrowType) { FieldReader reader = holder.reader; if (writer == null) { writer = new UnionWriter(UnionVector.this); } writer.setPosition(index); MinorType type = reader.getMinorType(); switch (type) { <#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.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> case ${name?upper_case}: Nullable${name}Holder ${uncappedName}Holder = new Nullable${name}Holder(); reader.read(${uncappedName}Holder); <#if minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> setSafe(index, ${uncappedName}Holder, arrowType); <#else> setSafe(index, ${uncappedName}Holder); break; case STRUCT: { ComplexCopier.copy(reader, writer); break; } case LIST: { ComplexCopier.copy(reader, writer); break; } default: throw new UnsupportedOperationException(); } } <#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.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> <#if minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> public void setSafe(int index, Nullable${name}Holder holder, ArrowType arrowType) { setType(index, MinorType.${name?upper_case}); get${name}Vector(null, arrowType).setSafe(index, holder); } <#else> public void setSafe(int index, Nullable${name}Holder holder) { setType(index, MinorType.${name?upper_case}); get${name}Vector(null).setSafe(index, holder); } public void setType(int index, MinorType type) { while (index >= getTypeBufferValueCapacity()) { reallocTypeBuffer(); } typeBuffer.setByte(index * TYPE_WIDTH , (byte) type.ordinal()); } private int getTypeBufferValueCapacity() { return capAtMaxInt(typeBuffer.capacity() / TYPE_WIDTH); } @Override public int hashCode(int index) { return hashCode(index, null); } @Override public int hashCode(int index, ArrowBufHasher hasher) { ValueVector vec = getVector(index); if (vec == null) { return ArrowBufPointer.NULL_HASH_CODE; } return vec.hashCode(index, hasher); } @Override public OUT accept(VectorVisitor visitor, IN value) { return visitor.visit(this, value); } @Override public String getName() { return name; } @Override public String toString() { return ValueVectorUtility.getToString(this, 0, getValueCount()); } @Override public T addOrGet(String name, FieldType fieldType, Class clazz) { return internalStruct.addOrGet(name, fieldType, clazz); } @Override public T getChild(String name, Class clazz) { return internalStruct.getChild(name, clazz); } @Override public VectorWithOrdinal getChildVectorWithOrdinal(String name) { return internalStruct.getChildVectorWithOrdinal(name); } @Override public int size() { return internalStruct.size(); } @Override public void setInitialCapacity(int valueCount, double density) { for (final ValueVector vector : internalStruct) { if (vector instanceof DensityAwareVector) { ((DensityAwareVector) vector).setInitialCapacity(valueCount, density); } else { vector.setInitialCapacity(valueCount); } } } /** * Set the element at the given index to null. For UnionVector, it throws an UnsupportedOperationException * as nulls are not supported at the top level and isNull() always returns false. * * @param index position of element * @throws UnsupportedOperationException whenever invoked */ @Override public void setNull(int index) { throw new UnsupportedOperationException("The method setNull() is not supported on UnionVector."); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/UnionViewWriter.java000066400000000000000000000217171515451615200275100ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/UnionViewWriter.java" /> package org.apache.arrow.vector.complex.impl; <#include "/@includes/vv_imports.ftl" /> import org.apache.arrow.vector.complex.writer.BaseWriter; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.complex.impl.NullableStructWriterFactory; import org.apache.arrow.vector.types.Types; <#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. */ @SuppressWarnings("unused") public class UnionViewWriter extends UnionWriter { public UnionViewWriter(UnionVector vector) { this(vector, NullableStructWriterFactory.getNullableStructWriterFactoryInstance()); } public UnionViewWriter(UnionVector vector, NullableStructWriterFactory nullableStructWriterFactory) { super(vector, nullableStructWriterFactory); } @Override public StructWriter struct() { data.setType(idx(), MinorType.LISTVIEW); getListWriter().setPosition(idx()); return getListWriter().struct(); } <#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/> <#assign friendlyType = (minor.friendlyType!minor.boxedType!type.boxedType) /> <#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> private ${name}Writer ${name?uncap_first}Writer; <#if minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> private ${name}Writer get${name}Writer(ArrowType arrowType) { if (${uncappedName}Writer == null) { ${uncappedName}Writer = new ${name}WriterImpl(data.get${name}Vector(arrowType)); ${uncappedName}Writer.setPosition(idx()); writers.add(${uncappedName}Writer); } return ${uncappedName}Writer; } public ${name}Writer as${name}(ArrowType arrowType) { data.setType(idx(), MinorType.${name?upper_case}); return get${name}Writer(arrowType); } <#else> private ${name}Writer get${name}Writer() { if (${uncappedName}Writer == null) { ${uncappedName}Writer = new ${name}WriterImpl(data.get${name}Vector()); ${uncappedName}Writer.setPosition(idx()); writers.add(${uncappedName}Writer); } return ${uncappedName}Writer; } public ${name}Writer as${name}() { data.setType(idx(), MinorType.${name?upper_case}); return get${name}Writer(); } @Override public void write(${name}Holder holder) { data.setType(idx(), MinorType.${name?upper_case}); <#if minor.class?starts_with("Decimal")> ArrowType arrowType = new ArrowType.Decimal(holder.precision, holder.scale, ${name}Holder.WIDTH * 8); get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).write${name}(<#list fields as field>holder.${field.name}<#if field_has_next>, , arrowType); <#elseif is_timestamp_tz(minor.class)> ArrowType.Timestamp arrowTypeWithoutTz = (ArrowType.Timestamp) MinorType.${name?upper_case?remove_ending("TZ")}.getType(); ArrowType arrowType = new ArrowType.Timestamp(arrowTypeWithoutTz.getUnit(), holder.timezone); get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).write(holder); <#elseif minor.class == "Duration"> ArrowType arrowType = new ArrowType.Duration(holder.unit); get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).write(holder); <#elseif minor.class == "FixedSizeBinary"> ArrowType arrowType = new ArrowType.FixedSizeBinary(holder.byteWidth); get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).write(holder); <#else> get${name}Writer().setPosition(idx()); get${name}Writer().write${name}(<#list fields as field>holder.${field.name}<#if field_has_next>, ); } public void write${minor.class}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, <#if minor.class?starts_with("Decimal")>, ArrowType arrowType) { data.setType(idx(), MinorType.${name?upper_case}); <#if minor.class?starts_with("Decimal")> get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).write${name}(<#list fields as field>${field.name}<#if field_has_next>, , arrowType); <#elseif is_timestamp_tz(minor.class)> ArrowType.Timestamp arrowTypeWithoutTz = (ArrowType.Timestamp) MinorType.${name?upper_case?remove_ending("TZ")}.getType(); ArrowType arrowType = new ArrowType.Timestamp(arrowTypeWithoutTz.getUnit(), "UTC"); get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).write${name}(<#list fields as field>${field.name}<#if field_has_next>, ); <#elseif minor.class == "Duration" || minor.class == "FixedSizeBinary"> // This is expected to throw. There's nothing more that we can do here since we can't infer any // sort of default unit for the Duration or a default width for the FixedSizeBinary types. ArrowType arrowType = MinorType.${name?upper_case}.getType(); get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).write${name}(<#list fields as field>${field.name}<#if field_has_next>, ); <#else> get${name}Writer().setPosition(idx()); get${name}Writer().write${name}(<#list fields as field>${field.name}<#if field_has_next>, ); } <#if minor.class?starts_with("Decimal")> public void write${name}(${friendlyType} value) { data.setType(idx(), MinorType.${name?upper_case}); ArrowType arrowType = new ArrowType.Decimal(value.precision(), value.scale(), ${name}Vector.TYPE_WIDTH * 8); get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).write${name}(value); } public void writeBigEndianBytesTo${name}(byte[] value, ArrowType arrowType) { data.setType(idx(), MinorType.${name?upper_case}); get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).writeBigEndianBytesTo${name}(value, arrowType); } <#elseif minor.class?ends_with("VarBinary")> @Override public void write${minor.class}(byte[] value) { get${name}Writer().setPosition(idx()); get${name}Writer().write${minor.class}(value); } @Override public void write${minor.class}(byte[] value, int offset, int length) { get${name}Writer().setPosition(idx()); get${name}Writer().write${minor.class}(value, offset, length); } @Override public void write${minor.class}(ByteBuffer value) { get${name}Writer().setPosition(idx()); get${name}Writer().write${minor.class}(value); } @Override public void write${minor.class}(ByteBuffer value, int offset, int length) { get${name}Writer().setPosition(idx()); get${name}Writer().write${minor.class}(value, offset, length); } <#elseif minor.class?ends_with("VarChar")> @Override public void write${minor.class}(${friendlyType} value) { get${name}Writer().setPosition(idx()); get${name}Writer().write${minor.class}(value); } @Override public void write${minor.class}(String value) { get${name}Writer().setPosition(idx()); get${name}Writer().write${minor.class}(value); } <#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 /> <#assign capName = minor.class?cap_first /> <#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> @Override public ${capName}Writer ${lowerName}() { data.setType(idx(), MinorType.LISTVIEW); getListViewWriter().setPosition(idx()); return getListViewWriter().${lowerName}(); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/UnionWriter.java000066400000000000000000000440311515451615200266470ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.complex.impl.NullableStructWriterFactory; import org.apache.arrow.vector.types.Types; <@pp.dropOutputFile /> <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/UnionWriter.java" /> <#include "/@includes/license.ftl" /> 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; <#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. */ @SuppressWarnings("unused") public class UnionWriter extends AbstractFieldWriter implements FieldWriter { protected UnionVector data; protected StructWriter structWriter; protected UnionListWriter listWriter; protected UnionListViewWriter listViewWriter; protected UnionMapWriter mapWriter; protected List writers = new java.util.ArrayList<>(); protected final NullableStructWriterFactory nullableStructWriterFactory; public UnionWriter(UnionVector vector) { this(vector, NullableStructWriterFactory.getNullableStructWriterFactoryInstance()); } public UnionWriter(UnionVector vector, NullableStructWriterFactory nullableStructWriterFactory) { data = vector; this.nullableStructWriterFactory = nullableStructWriterFactory; } /** * Convert the UnionWriter to a UnionViewWriter. * * @return the converted UnionViewWriter */ public UnionViewWriter toViewWriter() { UnionViewWriter unionViewWriter = new UnionViewWriter(data, nullableStructWriterFactory); unionViewWriter.structWriter = structWriter; unionViewWriter.listWriter = listWriter; unionViewWriter.listViewWriter = listViewWriter; unionViewWriter.mapWriter = mapWriter; unionViewWriter.writers = writers; unionViewWriter.setPosition(this.getPosition()); return unionViewWriter; } @Override public void setPosition(int index) { super.setPosition(index); for (BaseWriter writer : writers) { writer.setPosition(index); } } @Override public void start() { data.setType(idx(), MinorType.STRUCT); getStructWriter().start(); } @Override public void end() { getStructWriter().end(); } @Override public void startList() { getListWriter().startList(); data.setType(idx(), MinorType.LIST); } @Override public void endList() { getListWriter().endList(); } @Override public void startListView() { getListViewWriter().startListView(); data.setType(idx(), MinorType.LISTVIEW); } @Override public void endListView() { getListViewWriter().endListView(); } @Override public void startMap() { getMapWriter().startMap(); data.setType(idx(), MinorType.MAP); } @Override public void endMap() { getMapWriter().endMap(); } @Override public void startEntry() { getMapWriter().startEntry(); } @Override public MapWriter key() { return getMapWriter().key(); } @Override public MapWriter value() { return getMapWriter().value(); } @Override public void endEntry() { getMapWriter().endEntry(); } private StructWriter getStructWriter() { if (structWriter == null) { structWriter = nullableStructWriterFactory.build(data.getStruct()); structWriter.setPosition(idx()); writers.add(structWriter); } return structWriter; } public StructWriter asStruct() { data.setType(idx(), MinorType.STRUCT); return getStructWriter(); } protected ListWriter getListWriter() { if (listWriter == null) { listWriter = new UnionListWriter(data.getList(), nullableStructWriterFactory); listWriter.setPosition(idx()); writers.add(listWriter); } return listWriter; } protected ListWriter getListViewWriter() { if (listViewWriter == null) { listViewWriter = new UnionListViewWriter(data.getListView(), nullableStructWriterFactory); listViewWriter.setPosition(idx()); writers.add(listViewWriter); } return listViewWriter; } public ListWriter asList() { data.setType(idx(), MinorType.LIST); return getListWriter(); } public ListWriter asListView() { data.setType(idx(), MinorType.LISTVIEW); return getListViewWriter(); } private MapWriter getMapWriter() { if (mapWriter == null) { mapWriter = new UnionMapWriter(data.getMap(new ArrowType.Map(false))); mapWriter.setPosition(idx()); writers.add(mapWriter); } return mapWriter; } private MapWriter getMapWriter(ArrowType arrowType) { if (mapWriter == null) { mapWriter = new UnionMapWriter(data.getMap(arrowType)); mapWriter.setPosition(idx()); writers.add(mapWriter); } return mapWriter; } public MapWriter asMap(ArrowType arrowType) { data.setType(idx(), MinorType.MAP); 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); } BaseWriter getWriter(MinorType minorType, ArrowType arrowType) { switch (minorType) { case STRUCT: return getStructWriter(); case LIST: return getListWriter(); case LISTVIEW: 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 /> <#assign fields = minor.fields!type.fields /> <#assign uncappedName = name?uncap_first/> <#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> case ${name?upper_case}: <#if minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> return get${name}Writer(arrowType); <#else> return get${name}Writer(); default: throw new UnsupportedOperationException("Unknown type: " + minorType); } } <#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/> <#assign friendlyType = (minor.friendlyType!minor.boxedType!type.boxedType) /> <#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> private ${name}Writer ${name?uncap_first}Writer; <#if minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> private ${name}Writer get${name}Writer(ArrowType arrowType) { if (${uncappedName}Writer == null) { ${uncappedName}Writer = new ${name}WriterImpl(data.get${name}Vector(arrowType)); ${uncappedName}Writer.setPosition(idx()); writers.add(${uncappedName}Writer); } return ${uncappedName}Writer; } public ${name}Writer as${name}(ArrowType arrowType) { data.setType(idx(), MinorType.${name?upper_case}); return get${name}Writer(arrowType); } <#else> private ${name}Writer get${name}Writer() { if (${uncappedName}Writer == null) { ${uncappedName}Writer = new ${name}WriterImpl(data.get${name}Vector()); ${uncappedName}Writer.setPosition(idx()); writers.add(${uncappedName}Writer); } return ${uncappedName}Writer; } public ${name}Writer as${name}() { data.setType(idx(), MinorType.${name?upper_case}); return get${name}Writer(); } @Override public void write(${name}Holder holder) { data.setType(idx(), MinorType.${name?upper_case}); <#if minor.class?starts_with("Decimal")> ArrowType arrowType = new ArrowType.Decimal(holder.precision, holder.scale, ${name}Holder.WIDTH * 8); get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).write${name}(<#list fields as field>holder.${field.name}<#if field_has_next>, , arrowType); <#elseif is_timestamp_tz(minor.class)> ArrowType.Timestamp arrowTypeWithoutTz = (ArrowType.Timestamp) MinorType.${name?upper_case?remove_ending("TZ")}.getType(); ArrowType arrowType = new ArrowType.Timestamp(arrowTypeWithoutTz.getUnit(), holder.timezone); get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).write(holder); <#elseif minor.class == "Duration"> ArrowType arrowType = new ArrowType.Duration(holder.unit); get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).write(holder); <#elseif minor.class == "FixedSizeBinary"> ArrowType arrowType = new ArrowType.FixedSizeBinary(holder.byteWidth); get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).write(holder); <#else> get${name}Writer().setPosition(idx()); get${name}Writer().write${name}(<#list fields as field>holder.${field.name}<#if field_has_next>, ); } public void write${minor.class}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, <#if minor.class?starts_with("Decimal")>, ArrowType arrowType) { data.setType(idx(), MinorType.${name?upper_case}); <#if minor.class?starts_with("Decimal")> get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).write${name}(<#list fields as field>${field.name}<#if field_has_next>, , arrowType); <#elseif is_timestamp_tz(minor.class)> ArrowType.Timestamp arrowTypeWithoutTz = (ArrowType.Timestamp) MinorType.${name?upper_case?remove_ending("TZ")}.getType(); ArrowType arrowType = new ArrowType.Timestamp(arrowTypeWithoutTz.getUnit(), "UTC"); get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).write${name}(<#list fields as field>${field.name}<#if field_has_next>, ); <#elseif minor.class == "Duration" || minor.class == "FixedSizeBinary"> // This is expected to throw. There's nothing more that we can do here since we can't infer any // sort of default unit for the Duration or a default width for the FixedSizeBinary types. ArrowType arrowType = MinorType.${name?upper_case}.getType(); get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).write${name}(<#list fields as field>${field.name}<#if field_has_next>, ); <#else> get${name}Writer().setPosition(idx()); get${name}Writer().write${name}(<#list fields as field>${field.name}<#if field_has_next>, ); } <#if minor.class?starts_with("Decimal")> public void write${name}(${friendlyType} value) { data.setType(idx(), MinorType.${name?upper_case}); ArrowType arrowType = new ArrowType.Decimal(value.precision(), value.scale(), ${name}Vector.TYPE_WIDTH * 8); get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).write${name}(value); } public void writeBigEndianBytesTo${name}(byte[] value, ArrowType arrowType) { data.setType(idx(), MinorType.${name?upper_case}); get${name}Writer(arrowType).setPosition(idx()); get${name}Writer(arrowType).writeBigEndianBytesTo${name}(value, arrowType); } <#elseif minor.class?ends_with("VarBinary")> @Override public void write${minor.class}(byte[] value) { get${name}Writer().setPosition(idx()); get${name}Writer().write${minor.class}(value); } @Override public void write${minor.class}(byte[] value, int offset, int length) { get${name}Writer().setPosition(idx()); get${name}Writer().write${minor.class}(value, offset, length); } @Override public void write${minor.class}(ByteBuffer value) { get${name}Writer().setPosition(idx()); get${name}Writer().write${minor.class}(value); } @Override public void write${minor.class}(ByteBuffer value, int offset, int length) { get${name}Writer().setPosition(idx()); get${name}Writer().write${minor.class}(value, offset, length); } <#elseif minor.class?ends_with("VarChar")> @Override public void write${minor.class}(${friendlyType} value) { get${name}Writer().setPosition(idx()); get${name}Writer().write${minor.class}(value); } @Override public void write${minor.class}(String value) { get${name}Writer().setPosition(idx()); get${name}Writer().write${minor.class}(value); } public void writeNull() { } @Override public StructWriter struct() { data.setType(idx(), MinorType.LIST); getListWriter().setPosition(idx()); return getListWriter().struct(); } @Override public ListWriter list() { data.setType(idx(), MinorType.LIST); getListWriter().setPosition(idx()); return getListWriter().list(); } @Override public ListWriter list(String name) { data.setType(idx(), MinorType.STRUCT); getStructWriter().setPosition(idx()); return getStructWriter().list(name); } @Override public ListWriter listView() { data.setType(idx(), MinorType.LISTVIEW); getListViewWriter().setPosition(idx()); return getListViewWriter().listView(); } @Override public ListWriter listView(String name) { data.setType(idx(), MinorType.STRUCT); getStructWriter().setPosition(idx()); return getStructWriter().listView(name); } @Override public StructWriter struct(String name) { data.setType(idx(), MinorType.STRUCT); getStructWriter().setPosition(idx()); return getStructWriter().struct(name); } @Override public MapWriter map() { data.setType(idx(), MinorType.MAP); getListWriter().setPosition(idx()); return getListWriter().map(); } @Override public MapWriter map(boolean keysSorted) { data.setType(idx(), MinorType.MAP); getListWriter().setPosition(idx()); return getListWriter().map(keysSorted); } @Override public MapWriter map(String name) { data.setType(idx(), MinorType.MAP); getStructWriter().setPosition(idx()); return getStructWriter().map(name); } @Override public MapWriter map(String name, boolean keysSorted) { data.setType(idx(), MinorType.MAP); getStructWriter().setPosition(idx()); 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" /> <#assign upperName = minor.class?upper_case /> <#assign capName = minor.class?cap_first /> <#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> @Override public ${capName}Writer ${lowerName}(String name) { data.setType(idx(), MinorType.STRUCT); getStructWriter().setPosition(idx()); return getStructWriter().${lowerName}(name); } @Override public ${capName}Writer ${lowerName}() { data.setType(idx(), MinorType.LIST); getListWriter().setPosition(idx()); return getListWriter().${lowerName}(); } <#if minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> @Override public ${capName}Writer ${lowerName}(String name<#list minor.typeParams as typeParam>, ${typeParam.type} ${typeParam.name}) { data.setType(idx(), MinorType.STRUCT); getStructWriter().setPosition(idx()); return getStructWriter().${lowerName}(name<#list minor.typeParams as typeParam>, ${typeParam.name}); } @Override public void allocate() { data.allocateNew(); } @Override public void clear() { data.clear(); } @Override public void close() throws Exception { data.close(); } @Override public Field getField() { return data.getField(); } @Override public int getValueCapacity() { return data.getValueCapacity(); } } apache-arrow-java-19.0.0/vector/src/main/codegen/templates/ValueHolders.java000066400000000000000000000046611515451615200267640ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ <@pp.dropOutputFile /> <#list vv.modes as mode> <#list vv.types as type> <#list type.minor as minor> <#assign className="${mode.prefix}${minor.class}Holder" /> <@pp.changeOutputFile name="/org/apache/arrow/vector/holders/${className}.java" /> <#include "/@includes/license.ftl" /> package org.apache.arrow.vector.holders; <#include "/@includes/vv_imports.ftl" /> /** * Source code generated using FreeMarker template ${.template_name} */ public final class ${className} implements ValueHolder{ <#if mode.name == "Repeated"> /** The first index (inclusive) into the Vector. **/ public int start; /** The last index (exclusive) into the Vector. **/ public int end; /** The Vector holding the actual values. **/ public ${minor.class}Vector vector; <#else> public static final int WIDTH = ${type.width}; <#if mode.name == "Optional">public int isSet; <#else>public final int isSet = 1; <#assign fields = (minor.fields!type.fields) + (minor.typeParams![]) /> <#list fields as field> public ${field.type} ${field.name}; /** * Reason for not supporting the operation is that ValueHolders are potential scalar * replacements and hence we don't want any methods to be invoked on them. */ public int hashCode(){ throw new UnsupportedOperationException(); } /** * Reason for not supporting the operation is that ValueHolders are potential scalar * replacements and hence we don't want any methods to be invoked on them. */ public String toString(){ throw new UnsupportedOperationException(); } } apache-arrow-java-19.0.0/vector/src/main/java/000077500000000000000000000000001515451615200210345ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/module-info.java000066400000000000000000000041571515451615200241240ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ module org.apache.arrow.vector { exports org.apache.arrow.vector; exports org.apache.arrow.vector.compare; exports org.apache.arrow.vector.compare.util; exports org.apache.arrow.vector.complex; exports org.apache.arrow.vector.complex.impl; exports org.apache.arrow.vector.complex.reader; exports org.apache.arrow.vector.complex.writer; exports org.apache.arrow.vector.compression; exports org.apache.arrow.vector.dictionary; exports org.apache.arrow.vector.extension; exports org.apache.arrow.vector.holders; exports org.apache.arrow.vector.ipc; exports org.apache.arrow.vector.ipc.message; exports org.apache.arrow.vector.table; exports org.apache.arrow.vector.types; exports org.apache.arrow.vector.types.pojo; exports org.apache.arrow.vector.util; exports org.apache.arrow.vector.validate; opens org.apache.arrow.vector.types.pojo to com.fasterxml.jackson.databind; requires com.fasterxml.jackson.annotation; requires com.fasterxml.jackson.core; requires com.fasterxml.jackson.databind; requires com.fasterxml.jackson.datatype.jsr310; requires flatbuffers.java; requires jdk.unsupported; requires org.apache.arrow.format; requires org.apache.arrow.memory.core; requires org.apache.commons.codec; requires org.slf4j; uses org.apache.arrow.vector.compression.CompressionCodec.Factory; } apache-arrow-java-19.0.0/vector/src/main/java/org/000077500000000000000000000000001515451615200216235ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/000077500000000000000000000000001515451615200230445ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/000077500000000000000000000000001515451615200241765ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/000077500000000000000000000000001515451615200255005ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/AddOrGetResult.java000066400000000000000000000027041515451615200311760ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.util.Preconditions; /** * Tuple class containing a vector and whether it was created. * * @param The type of vector the result is for. */ public class AddOrGetResult { private final V vector; private final boolean created; /** Constructs a new object. */ public AddOrGetResult(V vector, boolean created) { this.vector = Preconditions.checkNotNull(vector); this.created = created; } /** Returns the vector. */ public V getVector() { return vector; } /** Returns whether the vector is created. */ public boolean isCreated() { return created; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/AllocationHelper.java000066400000000000000000000077601515451615200316020ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.complex.RepeatedFixedWidthVectorLike; import org.apache.arrow.vector.complex.RepeatedVariableWidthVectorLike; /** Helper utility methods for allocating storage for Vectors. */ public class AllocationHelper { private AllocationHelper() {} /** * Allocates the vector. * * @param v The vector to allocate. * @param valueCount Number of values to allocate. * @param bytesPerValue bytes per value. * @throws org.apache.arrow.memory.OutOfMemoryException if it can't allocate the memory. */ public static void allocate(ValueVector v, int valueCount, int bytesPerValue) { allocate(v, valueCount, bytesPerValue, 5); } /** * Allocates memory for a vector assuming given number of values and their width. * * @param v The vector the allocate. * @param valueCount The number of elements to allocate. * @param bytesPerValue The bytes per value to use for allocating underlying storage * @param childValCount If v is a repeated vector, this is number of child elements * to allocate. * @throws org.apache.arrow.memory.OutOfMemoryException if it can't allocate the memory. */ public static void allocatePrecomputedChildCount( ValueVector v, int valueCount, int bytesPerValue, int childValCount) { if (v instanceof FixedWidthVector) { ((FixedWidthVector) v).allocateNew(valueCount); } else if (v instanceof VariableWidthVector) { ((VariableWidthVector) v).allocateNew(valueCount * bytesPerValue, valueCount); } else if (v instanceof RepeatedFixedWidthVectorLike) { ((RepeatedFixedWidthVectorLike) v).allocateNew(valueCount, childValCount); } else if (v instanceof RepeatedVariableWidthVectorLike) { ((RepeatedVariableWidthVectorLike) v) .allocateNew(childValCount * bytesPerValue, valueCount, childValCount); } else { v.allocateNew(); } } /** * Allocates memory for a vector assuming given number of values and their width. * * @param v The vector the allocate. * @param valueCount The number of elements to allocate. * @param bytesPerValue The bytes per value to use for allocating underlying storage * @param repeatedPerTop If v is a repeated vector, this is assumed number of * elements per child. * @throws org.apache.arrow.memory.OutOfMemoryException if it can't allocate the memory */ public static void allocate( ValueVector v, int valueCount, int bytesPerValue, int repeatedPerTop) { allocatePrecomputedChildCount(v, valueCount, bytesPerValue, repeatedPerTop * valueCount); } /** * Allocates the exact amount if v is fixed width, otherwise falls back to dynamic allocation. * * @param v value vector we are trying to allocate * @param valueCount size we are trying to allocate * @throws org.apache.arrow.memory.OutOfMemoryException if it can't allocate the memory */ public static void allocateNew(ValueVector v, int valueCount) { if (v instanceof FixedWidthVector) { ((FixedWidthVector) v).allocateNew(valueCount); } else if (v instanceof VariableWidthVector) { ((VariableWidthVector) v).allocateNew(valueCount); } else { v.allocateNew(); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java000066400000000000000000000714341515451615200323710ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.memory.util.LargeMemoryUtil.capAtMaxInt; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Collections; import java.util.List; 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.MemoryUtil; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.util.CallBack; import org.apache.arrow.vector.util.OversizedAllocationException; import org.apache.arrow.vector.util.TransferPair; /** * BaseFixedWidthVector provides an abstract interface for implementing vectors of fixed width * values. The vectors are nullable implying that zero or more elements in the vector could be NULL. */ public abstract class BaseFixedWidthVector extends BaseValueVector implements FixedWidthVector, FieldVector, VectorDefinitionSetter { private final int typeWidth; protected int lastValueCapacity; protected int actualValueCapacity; protected final Field field; private int allocationMonitor; protected ArrowBuf valueBuffer; /** * Constructs a new instance. * * @param field field materialized by this vector * @param allocator The allocator to use for allocating memory for the vector. * @param typeWidth The width in bytes of the type. */ public BaseFixedWidthVector(Field field, final BufferAllocator allocator, final int typeWidth) { super(allocator); this.typeWidth = typeWidth; this.field = field; valueCount = 0; allocationMonitor = 0; validityBuffer = allocator.getEmpty(); valueBuffer = allocator.getEmpty(); lastValueCapacity = INITIAL_VALUE_ALLOCATION; refreshValueCapacity(); } @Override public int getTypeWidth() { return typeWidth; } @Override public String getName() { return field.getName(); } /* TODO: * see if getNullCount() can be made faster -- O(1) */ /* TODO: * Once the entire hierarchy has been refactored, move common functions * like getNullCount() to top level * base class BaseValueVector. * * Along with this, some class members (validityBuffer) can also be * abstracted out to top level base class. * * Right now BaseValueVector is the top level base class for other * vector types in ValueVector hierarchy (non-nullable) and those * vectors have not yet been refactored/removed so moving things to * the top class as of now is not a good idea. */ /** * Get the memory address of buffer that manages the validity (NULL or NON-NULL nature) of * elements in the vector. * * @return starting address of the buffer */ @Override public long getValidityBufferAddress() { return validityBuffer.memoryAddress(); } /** * Get the memory address of buffer that stores the data for elements in the vector. * * @return starting address of the buffer */ @Override public long getDataBufferAddress() { return valueBuffer.memoryAddress(); } /** * Get the memory address of buffer that stores the offsets for elements in the vector. This * operation is not supported for fixed-width vectors. * * @return starting address of the buffer * @throws UnsupportedOperationException for fixed width vectors */ @Override public long getOffsetBufferAddress() { throw new UnsupportedOperationException("not supported for fixed-width vectors"); } /** * 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. * * @return buffer */ @Override public ArrowBuf getValidityBuffer() { return validityBuffer; } /** * Get the buffer that stores the data for elements in the vector. * * @return buffer */ @Override public ArrowBuf getDataBuffer() { return valueBuffer; } /** * buffer that stores the offsets for elements in the vector. This operation is not supported for * fixed-width vectors. * * @return buffer * @throws UnsupportedOperationException for fixed width vectors */ @Override public ArrowBuf getOffsetBuffer() { throw new UnsupportedOperationException("not supported for fixed-width vectors"); } /** * Sets the desired value capacity for the vector. This function doesn't allocate any memory for * the vector. * * @param valueCount desired number of elements in the vector */ @Override public void setInitialCapacity(int valueCount) { computeAndCheckBufferSize(valueCount); lastValueCapacity = valueCount; } /** * Get the current value capacity for the vector. * * @return number of elements that vector can hold. */ @Override public int getValueCapacity() { return actualValueCapacity; } /** Call this if you change the capacity of valueBuffer or validityBuffer. */ protected void refreshValueCapacity() { actualValueCapacity = Math.min(getValueBufferValueCapacity(), getValidityBufferValueCapacity()); } protected int getValueBufferValueCapacity() { return capAtMaxInt(valueBuffer.capacity() / typeWidth); } protected int getValidityBufferValueCapacity() { return capAtMaxInt(validityBuffer.capacity() * 8); } /** zero out the vector and the data in associated buffers. */ @Override public void zeroVector() { initValidityBuffer(); initValueBuffer(); } /* zero out the validity buffer */ private void initValidityBuffer() { validityBuffer.setZero(0, validityBuffer.capacity()); } /* zero out the data buffer */ private void initValueBuffer() { valueBuffer.setZero(0, valueBuffer.capacity()); } /** * Reset the vector to initial state. Same as {@link #zeroVector()}. Note that this method doesn't * release any memory. */ @Override public void reset() { valueCount = 0; zeroVector(); } /** Close the vector and release the associated buffers. */ @Override public void close() { clear(); } /** Same as {@link #close()}. */ @Override public void clear() { valueCount = 0; validityBuffer = releaseBuffer(validityBuffer); valueBuffer = releaseBuffer(valueBuffer); refreshValueCapacity(); } /* used to step down the memory allocation */ protected void incrementAllocationMonitor() { if (allocationMonitor < 0) { allocationMonitor = 0; } allocationMonitor++; } /* used to step up the memory allocation */ protected void decrementAllocationMonitor() { if (allocationMonitor > 0) { allocationMonitor = 0; } allocationMonitor--; } /** Same as {@link #allocateNewSafe()}. */ @Override public void allocateNew() { allocateNew(lastValueCapacity); } /** * Allocate memory for the vector. We internally use a default value count of 4096 to allocate * memory for at least these many elements in the vector. See {@link #allocateNew(int)} for * allocating memory for specific number of elements in the vector. * * @return false if memory allocation fails, true otherwise. */ @Override public boolean allocateNewSafe() { try { allocateNew(lastValueCapacity); return true; } catch (Exception e) { return false; } } /** * Allocate memory for the vector to support storing at least the provided number of elements in * the vector. This method must be called prior to using the ValueVector. * * @param valueCount the desired number of elements in the vector * @throws org.apache.arrow.memory.OutOfMemoryException on error */ @Override public void allocateNew(int valueCount) { computeAndCheckBufferSize(valueCount); /* we are doing a new allocation -- release the current buffers */ clear(); try { allocateBytes(valueCount); } catch (Exception e) { clear(); throw e; } } /* * Compute the buffer size required for 'valueCount', and check if it's within bounds. */ private long computeAndCheckBufferSize(int valueCount) { final long size = computeCombinedBufferSize(valueCount, typeWidth); if (size > MAX_ALLOCATION_SIZE) { throw new OversizedAllocationException( "Memory required for vector capacity " + valueCount + " is (" + size + "), which is more than max allowed (" + MAX_ALLOCATION_SIZE + ")"); } return size; } /** * Actual memory allocation is done by this function. All the calculations and knowledge about * what size to allocate is upto the callers of this method. Callers appropriately handle errors * if memory allocation fails here. Callers should also take care of determining that desired size * is within the bounds of max allocation allowed and any other error conditions. */ private void allocateBytes(int valueCount) { DataAndValidityBuffers buffers = allocFixedDataAndValidityBufs(valueCount, typeWidth); valueBuffer = buffers.getDataBuf(); validityBuffer = buffers.getValidityBuf(); zeroVector(); refreshValueCapacity(); lastValueCapacity = getValueCapacity(); } /** * During splitAndTransfer, if we splitting from a random position within a byte, we can't just * 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. */ @Override protected void allocateValidityBuffer(final long validityBufferSize) { super.allocateValidityBuffer(validityBufferSize); refreshValueCapacity(); } /** * Get the potential buffer size for a particular number of records. * * @param count desired number of elements in the vector * @return estimated size of underlying buffers if the vector holds a given number of elements */ @Override public int getBufferSizeFor(final int count) { if (count == 0) { return 0; } return (count * typeWidth) + BitVectorHelper.getValidityBufferSizeFromCount(count); } /** * Get the size (number of bytes) of underlying buffers used by this vector. * * @return size of underlying buffers. */ @Override public int getBufferSize() { if (valueCount == 0) { return 0; } return (valueCount * typeWidth) + BitVectorHelper.getValidityBufferSizeFromCount(valueCount); } /** * Get information about how this field is materialized. * * @return the field corresponding to this vector */ @Override public Field getField() { return field; } /** * Return the underlying buffers associated with this vector. Note that this doesn't impact the * reference counts for this buffer so it only should be used for in-context access. Also note * that this buffer changes regularly thus external classes shouldn't hold a reference to it * (unless they change it). * * @param clear Whether to clear vector before returning; the buffers will still be refcounted but * the returned array will be the only reference to them * @return The underlying {@link ArrowBuf buffers} that is used by this vector instance. */ @Override public ArrowBuf[] getBuffers(boolean clear) { final ArrowBuf[] buffers; setReaderAndWriterIndex(); if (getBufferSize() == 0) { buffers = new ArrowBuf[0]; } else { buffers = new ArrowBuf[2]; buffers[0] = validityBuffer; buffers[1] = valueBuffer; } if (clear) { for (final ArrowBuf buffer : buffers) { buffer.getReferenceManager().retain(1); } clear(); } return buffers; } /** * Resize the vector to increase the capacity. The internal behavior is to double the current * value capacity. */ @Override public void reAlloc() { int targetValueCount = getValueCapacity() * 2; if (targetValueCount == 0) { if (lastValueCapacity > 0) { targetValueCount = lastValueCapacity; } else { targetValueCount = INITIAL_VALUE_ALLOCATION * 2; } } computeAndCheckBufferSize(targetValueCount); DataAndValidityBuffers buffers = allocFixedDataAndValidityBufs(targetValueCount, typeWidth); final ArrowBuf newValueBuffer = buffers.getDataBuf(); newValueBuffer.setBytes(0, valueBuffer, 0, valueBuffer.capacity()); newValueBuffer.setZero( valueBuffer.capacity(), newValueBuffer.capacity() - valueBuffer.capacity()); valueBuffer.getReferenceManager().release(); valueBuffer = newValueBuffer; final ArrowBuf newValidityBuffer = buffers.getValidityBuf(); newValidityBuffer.setBytes(0, validityBuffer, 0, validityBuffer.capacity()); newValidityBuffer.setZero( validityBuffer.capacity(), newValidityBuffer.capacity() - validityBuffer.capacity()); validityBuffer.getReferenceManager().release(); validityBuffer = newValidityBuffer; refreshValueCapacity(); lastValueCapacity = getValueCapacity(); } /** * Get the inner vectors. * * @return the inner vectors for this field as defined by the TypeLayout * @deprecated This API will be removed as the current implementations no longer support inner * vectors. */ @Deprecated @Override public List getFieldInnerVectors() { throw new UnsupportedOperationException("There are no inner vectors. Use getFieldBuffers"); } /** * Initialize the children in schema for this Field. This operation is a NO-OP for scalar types * since they don't have any children. * * @param children the schema * @throws IllegalArgumentException if children is a non-empty list for scalar types. */ @Override public void initializeChildrenFromFields(List children) { if (!children.isEmpty()) { throw new IllegalArgumentException("primitive type vector cannot have children"); } } /** * Get the inner child vectors. * * @return list of child vectors for complex types, empty list for scalar vector types */ @Override public List getChildrenFromFields() { return Collections.emptyList(); } /** * Load the buffers of this vector with provided source buffers. The caller manages the source * buffers and populates them before invoking this method. * * @param fieldNode the fieldNode indicating the value count * @param ownBuffers the buffers for this Field (own buffers only, children not included) */ @Override public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers) { if (ownBuffers.size() != 2) { throw new IllegalArgumentException( "Illegal buffer count, expected " + 2 + ", got: " + ownBuffers.size()); } ArrowBuf bitBuffer = ownBuffers.get(0); ArrowBuf dataBuffer = ownBuffers.get(1); validityBuffer.getReferenceManager().release(); validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); valueBuffer.getReferenceManager().release(); valueBuffer = dataBuffer.getReferenceManager().retain(dataBuffer, allocator); refreshValueCapacity(); valueCount = fieldNode.getLength(); } /** * Get the buffers belonging to this vector. * * @return the inner buffers. */ @Override public List getFieldBuffers() { List result = new ArrayList<>(2); setReaderAndWriterIndex(); result.add(validityBuffer); result.add(valueBuffer); return result; } /** Set the reader and writer indexes for the inner buffers. */ private void setReaderAndWriterIndex() { validityBuffer.readerIndex(0); valueBuffer.readerIndex(0); if (valueCount == 0) { validityBuffer.writerIndex(0); valueBuffer.writerIndex(0); } else { validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); if (typeWidth == 0) { /* specialized handling for BitVector */ valueBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); } else { valueBuffer.writerIndex((long) valueCount * typeWidth); } } } /** Validate the scalar values held by this vector. */ public void validateScalars() { // No validation by default. } /** * Construct a transfer pair of this vector and another vector of same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @param callBack not used * @return TransferPair */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return getTransferPair(ref, allocator); } /** * Construct a transfer pair of this vector and another vector of same type. * * @param field The field materialized by this vector. * @param allocator allocator for the target vector * @param callBack not used * @return TransferPair */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { return getTransferPair(field, allocator); } /** * Construct a transfer pair of this vector and another vector of same type. * * @param allocator allocator for the target vector * @return TransferPair */ @Override public TransferPair getTransferPair(BufferAllocator allocator) { return getTransferPair(getName(), allocator); } /** * Construct a transfer pair of this vector and another vector of same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return TransferPair */ @Override public abstract TransferPair getTransferPair(String ref, BufferAllocator allocator); /** * Construct a transfer pair of this vector and another vector of same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return TransferPair */ @Override public abstract TransferPair getTransferPair(Field field, BufferAllocator allocator); /** * Transfer this vector's data to another vector. The memory associated with this vector is * transferred to the allocator of target vector for accounting and management purposes. * * @param target destination vector for transfer */ public void transferTo(BaseFixedWidthVector target) { compareTypes(target, "transferTo"); target.clear(); target.validityBuffer = transferBuffer(validityBuffer, target.allocator); target.valueBuffer = transferBuffer(valueBuffer, target.allocator); target.valueCount = valueCount; target.refreshValueCapacity(); clear(); } /** * Slice this vector at desired index and length and transfer the corresponding data to the target * vector. * * @param startIndex start position of the split in source vector. * @param length length of the split. * @param target destination vector */ public void splitAndTransferTo(int startIndex, int length, BaseFixedWidthVector target) { Preconditions.checkArgument( startIndex >= 0 && length >= 0 && startIndex + length <= valueCount, "Invalid parameters startIndex: %s, length: %s for valueCount: %s", startIndex, length, valueCount); compareTypes(target, "splitAndTransferTo"); target.clear(); splitAndTransferValidityBuffer(startIndex, length, target); splitAndTransferValueBuffer(startIndex, length, target); target.setValueCount(length); } /** Data buffer can always be split and transferred using slicing. */ private void splitAndTransferValueBuffer( int startIndex, int length, BaseFixedWidthVector target) { final int startPoint = startIndex * typeWidth; final int sliceLength = length * typeWidth; final ArrowBuf slicedBuffer = valueBuffer.slice(startPoint, sliceLength); target.valueBuffer = transferBuffer(slicedBuffer, target.allocator); target.refreshValueCapacity(); } @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(); } /*----------------------------------------------------------------* | | | common getters and setters | | | *----------------------------------------------------------------*/ /** * Get the number of elements that are null in the vector. * * @return the number of null elements. */ @Override public int getNullCount() { return BitVectorHelper.getNullCount(validityBuffer, valueCount); } /** * Get the value count of vector. This will always be zero unless {@link #setValueCount(int)} has * been called prior to calling this. * * @return valueCount for the vector */ @Override public int getValueCount() { return valueCount; } /** * Set value count for the vector. * * @param valueCount value count to set */ @Override public void setValueCount(int valueCount) { this.valueCount = valueCount; final int currentValueCapacity = getValueCapacity(); while (valueCount > getValueCapacity()) { reAlloc(); } /* * We are trying to understand the pattern of memory allocation. * If initially, the user did vector.allocateNew(), we would have * allocated memory of default size (4096 * type width). * Later on user invokes setValueCount(count). * * If the existing value capacity is twice as large as the * valueCount, we know that we over-provisioned memory in the * first place when default memory allocation was done because user * really needs a much less value count in the vector. * * We record this by bumping up the allocationMonitor. If this pattern * happens for certain number of times and allocationMonitor * reaches the threshold (internal hardcoded) value, subsequent * call to allocateNew() will take care of stepping down the * default memory allocation size. * * Another case would be under-provisioning the initial memory and * thus going through a lot of realloc(). Here the goal is to * see if we can minimize the number of reallocations. Again the * state is recorded in allocationMonitor by decrementing it * (negative value). If a threshold is hit, realloc will try to * allocate more memory in order to possibly avoid a future realloc. * This case is also applicable to setSafe() methods which can trigger * a realloc() and thus we record the state there as well. */ if (valueCount > 0) { if (currentValueCapacity >= (valueCount * 2)) { incrementAllocationMonitor(); } else if (currentValueCapacity <= (valueCount / 2)) { decrementAllocationMonitor(); } } setReaderAndWriterIndex(); } /** * Check if the given index is within the current value capacity of the vector. * * @param index position to check * @return true if index is within the current value capacity */ public boolean isSafe(int index) { return index < getValueCapacity(); } /** * Check if element at given index is null. * * @param index position of element * @return true if element at given index is null, false otherwise */ @Override public boolean isNull(int index) { return (isSet(index) == 0); } /** * Same as {@link #isNull(int)}. * * @param index position of element * @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; } /** * Mark the particular position in the vector as non-null. * * @param index position of the element. */ @Override public void setIndexDefined(int index) { handleSafe(index); BitVectorHelper.setBit(validityBuffer, index); } public void set(int index, byte[] value, int start, int length) { throw new UnsupportedOperationException(); } public void setSafe(int index, byte[] value, int start, int length) { throw new UnsupportedOperationException(); } public void set(int index, ByteBuffer value, int start, int length) { throw new UnsupportedOperationException(); } public void setSafe(int index, ByteBuffer value, int start, int length) { throw new UnsupportedOperationException(); } /*----------------------------------------------------------------* | | | helper methods for setters | | | *----------------------------------------------------------------*/ protected void handleSafe(int index) { while (index >= getValueCapacity()) { decrementAllocationMonitor(); reAlloc(); } } /** * Copy a cell value from a particular index in source vector to a particular position in this * vector. The source vector should be of the same type as this one. * * @param fromIndex position to copy from in source vector * @param thisIndex position to copy to in this vector * @param from source vector */ @Override public void copyFrom(int fromIndex, int thisIndex, ValueVector from) { Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); if (from.isNull(fromIndex)) { BitVectorHelper.unsetBit(this.getValidityBuffer(), thisIndex); } else { BitVectorHelper.setBit(this.getValidityBuffer(), thisIndex); MemoryUtil.copyMemory( from.getDataBuffer().memoryAddress() + (long) fromIndex * typeWidth, this.getDataBuffer().memoryAddress() + (long) thisIndex * typeWidth, typeWidth); } } /** * Same as {@link #copyFrom(int, int, ValueVector)} except that it handles the case when the * capacity of the vector needs to be expanded before copy. * * @param fromIndex position to copy from in source vector * @param thisIndex position to copy to in this vector * @param from source vector */ @Override public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) { Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); handleSafe(thisIndex); copyFrom(fromIndex, thisIndex, from); } /** * Set the element at the given index to null. * * @param index position of element */ @Override public void setNull(int index) { handleSafe(index); // not really needed to set the bit to 0 as long as // the buffer always starts from 0. BitVectorHelper.unsetBit(validityBuffer, index); } @Override public ArrowBufPointer getDataPointer(int index) { return getDataPointer(index, new ArrowBufPointer()); } @Override public ArrowBufPointer getDataPointer(int index, ArrowBufPointer reuse) { if (isNull(index)) { reuse.set(null, 0, 0); } else { reuse.set(valueBuffer, (long) index * typeWidth, typeWidth); } return reuse; } @Override public int hashCode(int index) { return hashCode(index, null); } @Override public int hashCode(int index, ArrowBufHasher hasher) { if (isNull(index)) { return ArrowBufPointer.NULL_HASH_CODE; } long start = (long) typeWidth * index; long end = (long) typeWidth * (index + 1); return ByteFunctionHelpers.hash(hasher, this.getDataBuffer(), start, end); } @Override public OUT accept(VectorVisitor visitor, IN value) { return visitor.visit(this, value); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/BaseIntVector.java000066400000000000000000000031441515451615200310550ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** Interface for all int type vectors. */ public interface BaseIntVector extends FieldVector { /** * Sets the value at index, note this value may need to be truncated. Note this is safe version * (i.e. call setSafe(int, ...) method in vector) */ void setWithPossibleTruncate(int index, long value); /** * Sets the value at index, note this value may need to be truncated. Note this is unsafe version * (i.e. call set(int, ...) method in vector) */ void setUnsafeWithPossibleTruncate(int index, long value); /** * Gets the value at index. This value may have been extended to long and will throw {@link * NullPointerException} if the value is null. Note null check could be turned off via {@link * NullCheckingForGet}. */ long getValueAsLong(int index); } BaseLargeVariableWidthVector.java000066400000000000000000001405111515451615200337440ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.memory.util.LargeMemoryUtil.capAtMaxInt; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Collections; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.OutOfMemoryException; 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.hash.ArrowBufHasher; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.util.CallBack; import org.apache.arrow.vector.util.OversizedAllocationException; import org.apache.arrow.vector.util.TransferPair; /** * BaseLargeVariableWidthVector is a base class providing functionality for large strings/large * bytes types. */ public abstract class BaseLargeVariableWidthVector extends BaseValueVector implements VariableWidthFieldVector { private static final int DEFAULT_RECORD_BYTE_COUNT = 12; private static final int INITIAL_BYTE_COUNT = INITIAL_VALUE_ALLOCATION * DEFAULT_RECORD_BYTE_COUNT; private int lastValueCapacity; private long lastValueAllocationSizeInBytes; /* 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 valueBuffer; protected ArrowBuf offsetBuffer; protected int lastSet; protected final Field field; /** * Constructs a new instance. * * @param field The field materialized by this vector. * @param allocator The allocator to use for creating/resizing buffers */ public BaseLargeVariableWidthVector(Field field, final BufferAllocator allocator) { super(allocator); this.field = field; lastValueAllocationSizeInBytes = INITIAL_BYTE_COUNT; // -1 because we require one extra slot for the offset array. lastValueCapacity = INITIAL_VALUE_ALLOCATION - 1; valueCount = 0; lastSet = -1; offsetBuffer = allocator.getEmpty(); validityBuffer = allocator.getEmpty(); valueBuffer = allocator.getEmpty(); } @Override public String getName() { return field.getName(); } /** * 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. * * @return buffer */ @Override public ArrowBuf getValidityBuffer() { return validityBuffer; } /** * Get the buffer that stores the data for elements in the vector. * * @return buffer */ @Override public ArrowBuf getDataBuffer() { return valueBuffer; } /** * buffer that stores the offsets for elements in the vector. This operation is not supported for * fixed-width vectors. * * @return buffer */ @Override public ArrowBuf getOffsetBuffer() { return offsetBuffer; } /** * Get the memory address of buffer that stores the offsets for elements in the vector. * * @return starting address of the buffer */ @Override public long getOffsetBufferAddress() { return offsetBuffer.memoryAddress(); } /** * Get the memory address of buffer that manages the validity (NULL or NON-NULL nature) of * elements in the vector. * * @return starting address of the buffer */ @Override public long getValidityBufferAddress() { return validityBuffer.memoryAddress(); } /** * Get the memory address of buffer that stores the data for elements in the vector. * * @return starting address of the buffer */ @Override public long getDataBufferAddress() { return valueBuffer.memoryAddress(); } /** * Sets the desired value capacity for the vector. This function doesn't allocate any memory for * the vector. * * @param valueCount desired number of elements in the vector */ @Override public void setInitialCapacity(int valueCount) { final long size = (long) valueCount * DEFAULT_RECORD_BYTE_COUNT; checkDataBufferSize(size); computeAndCheckOffsetsBufferSize(valueCount); lastValueAllocationSizeInBytes = size; lastValueCapacity = valueCount; } /** * Sets the desired value capacity for the vector. This function doesn't allocate any memory for * the vector. * * @param valueCount desired number of elements in the vector * @param density average number of bytes per variable width element */ @Override public void setInitialCapacity(int valueCount, double density) { long size = Math.max((long) (valueCount * density), 1L); checkDataBufferSize(size); computeAndCheckOffsetsBufferSize(valueCount); lastValueAllocationSizeInBytes = size; lastValueCapacity = valueCount; } /** * Get the density of this ListVector. * * @return density */ public double getDensity() { if (valueCount == 0) { return 0.0D; } final long startOffset = getStartOffset(0); final long endOffset = getStartOffset(valueCount); final double totalListSize = endOffset - startOffset; return totalListSize / valueCount; } /** * Get the current capacity which does not exceed either validity buffer or offset buffer. Note: * Here the `getValueCapacity` has no relationship with the value buffer. * * @return number of elements that vector can hold. */ @Override public int getValueCapacity() { final long offsetValueCapacity = Math.max(getOffsetBufferValueCapacity() - 1, 0); return capAtMaxInt(Math.min(offsetValueCapacity, getValidityBufferValueCapacity())); } private long getValidityBufferValueCapacity() { return validityBuffer.capacity() * 8; } private long getOffsetBufferValueCapacity() { return offsetBuffer.capacity() / OFFSET_WIDTH; } /** zero out the vector and the data in associated buffers. */ public void zeroVector() { initValidityBuffer(); initOffsetBuffer(); valueBuffer.setZero(0, valueBuffer.capacity()); } /* zero out the validity buffer */ private void initValidityBuffer() { validityBuffer.setZero(0, validityBuffer.capacity()); } /* zero out the offset buffer */ private void initOffsetBuffer() { offsetBuffer.setZero(0, offsetBuffer.capacity()); } /** * Reset the vector to initial state. Same as {@link #zeroVector()}. Note that this method doesn't * release any memory. */ @Override public void reset() { zeroVector(); lastSet = -1; valueCount = 0; } /** Close the vector and release the associated buffers. */ @Override public void close() { clear(); } /** Same as {@link #close()}. */ @Override public void clear() { validityBuffer = releaseBuffer(validityBuffer); valueBuffer = releaseBuffer(valueBuffer); offsetBuffer = releaseBuffer(offsetBuffer); lastSet = -1; valueCount = 0; } /** * Get the inner vectors. * * @return the inner vectors for this field as defined by the TypeLayout * @deprecated This API will be removed as the current implementations no longer support inner * vectors. */ @Override @Deprecated public List getFieldInnerVectors() { throw new UnsupportedOperationException("There are no inner vectors. Use getFieldBuffers"); } /** * Initialize the children in schema for this Field. This operation is a NO-OP for scalar types * since they don't have any children. * * @param children the schema * @throws IllegalArgumentException if children is a non-empty list for scalar types. */ @Override public void initializeChildrenFromFields(List children) { if (!children.isEmpty()) { throw new IllegalArgumentException("primitive type vector cannot have children"); } } /** * Get the inner child vectors. * * @return list of child vectors for complex types, empty list for scalar vector types */ @Override public List getChildrenFromFields() { return Collections.emptyList(); } /** * Load the buffers of this vector with provided source buffers. The caller manages the source * buffers and populates them before invoking this method. * * @param fieldNode the fieldNode indicating the value count * @param ownBuffers the buffers for this Field (own buffers only, children not included) */ @Override public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers) { ArrowBuf bitBuffer = ownBuffers.get(0); ArrowBuf offBuffer = ownBuffers.get(1); ArrowBuf dataBuffer = ownBuffers.get(2); validityBuffer.getReferenceManager().release(); validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); offsetBuffer.getReferenceManager().release(); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); valueBuffer.getReferenceManager().release(); valueBuffer = dataBuffer.getReferenceManager().retain(dataBuffer, allocator); lastSet = fieldNode.getLength() - 1; valueCount = fieldNode.getLength(); } /** * Get the buffers belonging to this vector. * * @return the inner buffers. */ @Override public List getFieldBuffers() { // before flight/IPC, we must bring the vector to a consistent state. // this is because, it is possible that the offset buffers of some trailing values // are not updated. this may cause some data in the data buffer being lost. // for details, please see TestValueVector#testUnloadVariableWidthVector. fillHoles(valueCount); List result = new ArrayList<>(3); setReaderAndWriterIndex(); result.add(validityBuffer); result.add(offsetBuffer); result.add(valueBuffer); return result; } /** * Export the buffers of the fields for C Data Interface. This method traverse the buffers and * export buffer and buffer's memory address into a list of buffers and a pointer to the list of * buffers. */ @Override public void exportCDataBuffers(List buffers, ArrowBuf buffersPtr, long nullValue) { // before flight/IPC, we must bring the vector to a consistent state. // this is because, it is possible that the offset buffers of some trailing values // are not updated. this may cause some data in the data buffer being lost. // for details, please see TestValueVector#testUnloadVariableWidthVector. fillHoles(valueCount); exportBuffer(validityBuffer, buffers, buffersPtr, nullValue, true); if (offsetBuffer.capacity() == 0) { // Empty offset buffer is allowed for historical reason. // To export it through C Data interface, we need to allocate a buffer with one offset. // We set `retain = false` to explicitly not increase the ref count for the exported buffer. // The ref count of the newly created buffer (i.e., 1) already represents the usage // at imported side. exportBuffer(allocateOffsetBuffer(OFFSET_WIDTH), buffers, buffersPtr, nullValue, false); } else { exportBuffer(offsetBuffer, buffers, buffersPtr, nullValue, true); } exportBuffer(valueBuffer, buffers, buffersPtr, nullValue, true); } /** Set the reader and writer indexes for the inner buffers. */ private void setReaderAndWriterIndex() { validityBuffer.readerIndex(0); offsetBuffer.readerIndex(0); valueBuffer.readerIndex(0); if (valueCount == 0) { validityBuffer.writerIndex(0); valueBuffer.writerIndex(0); } else { final long lastDataOffset = getStartOffset(valueCount); 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()}. */ @Override public void allocateNew() { allocateNew(lastValueAllocationSizeInBytes, lastValueCapacity); } /** * Allocate memory for the vector. We internally use a default value count of 4096 to allocate * memory for at least these many elements in the vector. See {@link #allocateNew(long, int)} for * allocating memory for specific number of elements in the vector. * * @return false if memory allocation fails, true otherwise. */ @Override public boolean allocateNewSafe() { try { allocateNew(lastValueAllocationSizeInBytes, lastValueCapacity); return true; } catch (Exception e) { return false; } } /** * Allocate memory for the vector to support storing at least the provided number of elements in * the vector. This method must be called prior to using the ValueVector. * * @param totalBytes desired total memory capacity * @param valueCount the desired number of elements in the vector * @throws org.apache.arrow.memory.OutOfMemoryException if memory allocation fails */ @Override public void allocateNew(long totalBytes, int valueCount) { assert totalBytes >= 0; checkDataBufferSize(totalBytes); computeAndCheckOffsetsBufferSize(valueCount); /* we are doing a new allocation -- release the current buffers */ clear(); try { allocateBytes(totalBytes, valueCount); } catch (Exception e) { clear(); throw e; } } @Override public void allocateNew(int valueCount) { allocateNew(lastValueAllocationSizeInBytes, valueCount); } /* Check if the data buffer size is within bounds. */ private void checkDataBufferSize(long size) { if (size > MAX_ALLOCATION_SIZE || size < 0) { throw new OversizedAllocationException( "Memory required for vector " + " is (" + size + "), which is more than max allowed (" + MAX_ALLOCATION_SIZE + ")"); } } /** * Compute the buffer size required for 'valueCount' offsets and validity, and check if it's * within bounds. */ private long computeAndCheckOffsetsBufferSize(int valueCount) { /* to track the end offset of last data element in vector, we need * an additional slot in offset buffer. */ final long size = computeCombinedBufferSize(valueCount + 1, OFFSET_WIDTH); if (size > MAX_ALLOCATION_SIZE) { throw new OversizedAllocationException( "Memory required for vector capacity " + valueCount + " is (" + size + "), which is more than max allowed (" + MAX_ALLOCATION_SIZE + ")"); } return size; } /* allocate the inner buffers */ private void allocateBytes(final long valueBufferSize, final int valueCount) { /* allocate data buffer */ long curSize = valueBufferSize; valueBuffer = allocator.buffer(curSize); valueBuffer.readerIndex(0); /* allocate offset buffer and validity buffer */ DataAndValidityBuffers buffers = allocFixedDataAndValidityBufs(valueCount + 1, OFFSET_WIDTH); offsetBuffer = buffers.getDataBuf(); validityBuffer = buffers.getValidityBuf(); initOffsetBuffer(); initValidityBuffer(); lastValueCapacity = getValueCapacity(); lastValueAllocationSizeInBytes = capAtMaxInt(valueBuffer.capacity()); } /* allocate offset buffer */ private ArrowBuf allocateOffsetBuffer(final long size) { ArrowBuf offsetBuffer = allocator.buffer(size); offsetBuffer.readerIndex(0); offsetBuffer.setZero(0, offsetBuffer.capacity()); return offsetBuffer; } /* allocate validity buffer */ @Override protected void allocateValidityBuffer(final long size) { super.allocateValidityBuffer(size); } /** * Resize the vector to increase the capacity. The internal behavior is to double the current * value capacity. */ @Override public void reAlloc() { reallocDataBuffer(); reallocValidityAndOffsetBuffers(); } /** * Reallocate the data buffer. Data Buffer stores the actual data for LARGEVARCHAR or * LARGEVARBINARY elements in the vector. The behavior is to double the size of buffer. * * @throws OversizedAllocationException if the desired new size is more than max allowed * @throws OutOfMemoryException if the internal memory allocation fails */ public void reallocDataBuffer() { final long currentBufferCapacity = valueBuffer.capacity(); long newAllocationSize = currentBufferCapacity * 2; if (newAllocationSize == 0) { if (lastValueAllocationSizeInBytes > 0) { newAllocationSize = lastValueAllocationSizeInBytes; } else { newAllocationSize = INITIAL_BYTE_COUNT * 2; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); assert newAllocationSize >= 1; checkDataBufferSize(newAllocationSize); final ArrowBuf newBuf = allocator.buffer(newAllocationSize); newBuf.setBytes(0, valueBuffer, 0, currentBufferCapacity); valueBuffer.getReferenceManager().release(); valueBuffer = newBuf; lastValueAllocationSizeInBytes = valueBuffer.capacity(); } /** * Reallocate the validity and offset buffers for this vector. Validity buffer is used to track * the NULL or NON-NULL nature of elements in the vector and offset buffer is used to store the * lengths of variable width elements in the vector. * *

    Note that data buffer for variable length vectors moves independent of the companion * validity and offset buffers. This is in contrast to what we have for fixed width vectors. * *

    So even though we may have setup an initial capacity of 1024 elements in the vector, it is * quite possible that we need to reAlloc() the data buffer when we are setting the 5th element in * the vector simply because previous variable length elements have exhausted the buffer capacity. * However, we really don't need to reAlloc() validity and offset buffers until we try to set the * 1025th element This is why we do a separate check for safe methods to determine which buffer * needs reallocation. * * @throws OversizedAllocationException if the desired new size is more than max allowed * @throws OutOfMemoryException if the internal memory allocation fails */ public void reallocValidityAndOffsetBuffers() { int targetOffsetCount = capAtMaxInt((offsetBuffer.capacity() / OFFSET_WIDTH) * 2); if (targetOffsetCount == 0) { if (lastValueCapacity > 0) { targetOffsetCount = (lastValueCapacity + 1); } else { targetOffsetCount = 2 * (INITIAL_VALUE_ALLOCATION + 1); } } computeAndCheckOffsetsBufferSize(targetOffsetCount); DataAndValidityBuffers buffers = allocFixedDataAndValidityBufs(targetOffsetCount, OFFSET_WIDTH); final ArrowBuf newOffsetBuffer = buffers.getDataBuf(); newOffsetBuffer.setBytes(0, offsetBuffer, 0, offsetBuffer.capacity()); newOffsetBuffer.setZero( offsetBuffer.capacity(), newOffsetBuffer.capacity() - offsetBuffer.capacity()); offsetBuffer.getReferenceManager().release(); offsetBuffer = newOffsetBuffer; final ArrowBuf newValidityBuffer = buffers.getValidityBuf(); newValidityBuffer.setBytes(0, validityBuffer, 0, validityBuffer.capacity()); newValidityBuffer.setZero( validityBuffer.capacity(), newValidityBuffer.capacity() - validityBuffer.capacity()); validityBuffer.getReferenceManager().release(); validityBuffer = newValidityBuffer; lastValueCapacity = getValueCapacity(); } /** * Get the size (number of bytes) of underlying data buffer. * * @return number of bytes in the data buffer */ @Override public int getByteCapacity() { return capAtMaxInt(valueBuffer.capacity()); } @Override public int sizeOfValueBuffer() { if (valueCount == 0) { return 0; } return capAtMaxInt(getStartOffset(valueCount)); } /** * Get the size (number of bytes) of underlying buffers used by this vector. * * @return size of underlying buffers. */ @Override public int getBufferSize() { return getBufferSizeFor(this.valueCount); } /** * Get the potential buffer size for a particular number of records. * * @param valueCount desired number of elements in the vector * @return estimated size of underlying buffers if the vector holds a given number of elements */ @Override public int getBufferSizeFor(final int valueCount) { if (valueCount == 0) { return 0; } 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); return capAtMaxInt(validityBufferSize + offsetBufferSize + dataBufferSize); } /** * Get information about how this field is materialized. * * @return the field corresponding to this vector */ @Override public Field getField() { return field; } /** * Return the underlying buffers associated with this vector. Note that this doesn't impact the * reference counts for this buffer so it only should be used for in-context access. Also note * that this buffer changes regularly thus external classes shouldn't hold a reference to it * (unless they change it). * * @param clear Whether to clear vector before returning; the buffers will still be refcounted but * the returned array will be the only reference to them * @return The underlying {@link io.netty.buffer.ArrowBuf buffers} that is used by this vector * instance. */ @Override public ArrowBuf[] getBuffers(boolean clear) { final ArrowBuf[] buffers; setReaderAndWriterIndex(); if (getBufferSize() == 0) { buffers = new ArrowBuf[0]; } else { buffers = new ArrowBuf[3]; buffers[0] = validityBuffer; buffers[1] = offsetBuffer; buffers[2] = valueBuffer; } if (clear) { for (final ArrowBuf buffer : buffers) { buffer.getReferenceManager().retain(); } clear(); } return buffers; } /** Validate the scalar values held by this vector. */ public void validateScalars() { // No validation by default. } /** * Construct a transfer pair of this vector and another vector of same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @param callBack not used * @return TransferPair */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return getTransferPair(ref, allocator); } /** * Construct a transfer pair of this vector and another vector of same type. * * @param field The field materialized by this vector * @param allocator allocator for the target vector * @param callBack not used * @return TransferPair */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { return getTransferPair(field, allocator); } /** * Construct a transfer pair of this vector and another vector of same type. * * @param allocator allocator for the target vector * @return TransferPair */ @Override public TransferPair getTransferPair(BufferAllocator allocator) { return getTransferPair(getName(), allocator); } /** * Construct a transfer pair of this vector and another vector of same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return TransferPair */ @Override public abstract TransferPair getTransferPair(String ref, BufferAllocator allocator); /** * Construct a transfer pair of this vector and another vector of same type. * * @param field The field materialized by this vector * @param allocator allocator for the target vector * @return TransferPair */ @Override public abstract TransferPair getTransferPair(Field field, BufferAllocator allocator); /** * Transfer this vector's data to another vector. The memory associated with this vector is * transferred to the allocator of target vector for accounting and management purposes. * * @param target destination vector for transfer */ public void transferTo(BaseLargeVariableWidthVector target) { compareTypes(target, "transferTo"); target.clear(); target.validityBuffer = transferBuffer(validityBuffer, target.allocator); target.valueBuffer = transferBuffer(valueBuffer, target.allocator); target.offsetBuffer = transferBuffer(offsetBuffer, target.allocator); target.setLastSet(this.lastSet); if (this.valueCount > 0) { target.setValueCount(this.valueCount); } clear(); } /** * Slice this vector at desired index and length and transfer the corresponding data to the target * vector. * * @param startIndex start position of the split in source vector. * @param length length of the split. * @param target destination vector */ public void splitAndTransferTo(int startIndex, int length, BaseLargeVariableWidthVector target) { Preconditions.checkArgument( startIndex >= 0 && length >= 0 && startIndex + length <= valueCount, "Invalid parameters startIndex: %s, length: %s for valueCount: %s", startIndex, length, valueCount); compareTypes(target, "splitAndTransferTo"); target.clear(); if (length > 0) { splitAndTransferValidityBuffer(startIndex, length, target); splitAndTransferOffsetBuffer(startIndex, length, target); target.setLastSet(length - 1); target.setValueCount(length); } } /** * Transfer the offsets along with data. Unlike the data buffer, we cannot simply slice the offset * buffer for split and transfer. The reason is that offsets in the target vector have to be * adjusted and made relative to the staring offset in source vector from the start index of * split. This is why, we need to explicitly allocate the offset buffer and set the adjusted * offsets in the target vector. */ private void splitAndTransferOffsetBuffer( int startIndex, int length, BaseLargeVariableWidthVector target) { final long start = getStartOffset(startIndex); final long end = getStartOffset(startIndex + length); final long dataLength = end - start; target.offsetBuffer = target.allocateOffsetBuffer((long) (length + 1) * OFFSET_WIDTH); for (int i = 0; i < length + 1; i++) { final long relativeSourceOffset = getStartOffset(startIndex + i) - start; target.offsetBuffer.setLong((long) i * OFFSET_WIDTH, relativeSourceOffset); } final ArrowBuf slicedBuffer = valueBuffer.slice(start, dataLength); target.valueBuffer = transferBuffer(slicedBuffer, target.allocator); } @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(); } target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); target.validityBuffer.getReferenceManager().retain(); } /*----------------------------------------------------------------* | | | common getters and setters | | | *----------------------------------------------------------------*/ /** * Get the number of elements that are null in the vector. * * @return the number of null elements. */ @Override public int getNullCount() { return BitVectorHelper.getNullCount(validityBuffer, valueCount); } /** * Check if the given index is within the current value capacity of the vector. * * @param index position to check * @return true if index is within the current value capacity */ public boolean isSafe(int index) { return index < getValueCapacity(); } /** * Check if element at given index is null. * * @param index position of element * @return true if element at given index is null */ @Override public boolean isNull(int index) { return (isSet(index) == 0); } /** * Same as {@link #isNull(int)}. * * @param index position of element * @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; } /** * Get the value count of vector. This will always be zero unless setValueCount(int) has been * called prior to calling this. * * @return valueCount for the vector */ @Override public int getValueCount() { return valueCount; } /** * Sets the value count for the vector. * * @param valueCount value count */ @Override public void setValueCount(int valueCount) { assert valueCount >= 0; this.valueCount = valueCount; while (valueCount > getValueCapacity()) { reallocValidityAndOffsetBuffers(); } fillHoles(valueCount); lastSet = valueCount - 1; setReaderAndWriterIndex(); } /** * Create holes in the vector upto the given index (exclusive). Holes will be created from the * current last set position in the vector. * * @param index target index */ @Override public void fillEmpties(int index) { handleSafe(index, emptyByteArray.length); fillHoles(index); lastSet = index - 1; } /** * Set the index of last non-null element in the vector. It is important to call this method with * appropriate value before calling {@link #setValueCount(int)}. * * @param value desired index of last non-null element. */ @Override public void setLastSet(int value) { lastSet = value; } /** * Get the index of last non-null element in the vector. * * @return index of the last non-null element */ @Override public int getLastSet() { return lastSet; } /** * Mark the particular position in the vector as non-null. * * @param index position of the element. */ @Override public void setIndexDefined(int index) { // We need to check and realloc both validity and offset buffer while (index >= getValueCapacity()) { reallocValidityAndOffsetBuffers(); } BitVectorHelper.setBit(validityBuffer, index); } /** * Sets the value length for an element. * * @param index position of the element to set * @param length length of the element */ public void setValueLengthSafe(int index, int length) { assert index >= 0; handleSafe(index, length); fillHoles(index); final long startOffset = getStartOffset(index); offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, startOffset + length); lastSet = index; } /** * Get the variable length element at specified index as Text. * * @param index position of element to get * @return greater than 0 length for non-null element, 0 otherwise */ @Override public int getValueLength(int index) { assert index >= 0; if (isSet(index) == 0) { return 0; } final long startOffset = getStartOffset(index); final int dataLength = (int) (getEndOffset(index) - startOffset); return dataLength; } /** * Set the variable length element at the specified index to the supplied byte array. This is same * as using {@link #set(int, byte[], int, int)} with start as 0 and length as value.length * * @param index position of the element to set * @param value array of bytes to write */ @Override public void set(int index, byte[] value) { assert index >= 0; fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); setBytes(index, value, 0, value.length); lastSet = index; } /** * Same as {@link #set(int, byte[])} except that it handles the case where index and length of new * element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param value array of bytes to write */ @Override public void setSafe(int index, byte[] value) { assert index >= 0; handleSafe(index, value.length); fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); setBytes(index, value, 0, value.length); lastSet = index; } /** * Set the variable length element at the specified index to the supplied byte array. * * @param index position of the element to set * @param value array of bytes to write * @param start start index in array of bytes * @param length length of data in array of bytes */ @Override public void set(int index, byte[] value, int start, int length) { assert index >= 0; fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); setBytes(index, value, start, length); lastSet = index; } /** * Same as {@link #set(int, byte[], int, int)} except that it handles the case where index and * length of new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param value array of bytes to write * @param start start index in array of bytes * @param length length of data in array of bytes */ public void setSafe(int index, byte[] value, int start, int length) { assert index >= 0; handleSafe(index, length); fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); setBytes(index, value, start, length); lastSet = index; } /** * Set the variable length element at the specified index to the content in supplied ByteBuffer. * * @param index position of the element to set * @param value ByteBuffer with data * @param start start index in ByteBuffer * @param length length of data in ByteBuffer */ @Override public void set(int index, ByteBuffer value, int start, int length) { assert index >= 0; fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); final long startOffset = getStartOffset(index); offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, startOffset + length); valueBuffer.setBytes(startOffset, value, start, length); lastSet = index; } /** * Same as {@link #set(int, ByteBuffer, int, int)} except that it handles the case where index and * length of new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param value ByteBuffer with data * @param start start index in ByteBuffer * @param length length of data in ByteBuffer */ public void setSafe(int index, ByteBuffer value, int start, int length) { assert index >= 0; handleSafe(index, length); fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); final long startOffset = getStartOffset(index); offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, startOffset + length); valueBuffer.setBytes(startOffset, value, start, length); lastSet = index; } /** * Set the element at the given index to null. * * @param index position of element */ @Override public void setNull(int index) { // We need to check and realloc both validity and offset buffer while (index >= getValueCapacity()) { reallocValidityAndOffsetBuffers(); } BitVectorHelper.unsetBit(validityBuffer, index); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param start start position of data in buffer * @param end end position of data in buffer * @param buffer data buffer containing the variable width element to be stored in the vector */ public void set(int index, int isSet, long start, long end, ArrowBuf buffer) { assert index >= 0; final long dataLength = end - start; fillHoles(index); BitVectorHelper.setValidityBit(validityBuffer, index, isSet); final long startOffset = offsetBuffer.getLong((long) index * OFFSET_WIDTH); offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, end); valueBuffer.setBytes(startOffset, buffer, start, dataLength); lastSet = index; } /** * Same as {@link #set(int, int, long, long, ArrowBuf)} except that it handles the case when index * is greater than or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param start start position of data in buffer * @param end end position of data in buffer * @param buffer data buffer containing the variable width element to be stored in the vector */ public void setSafe(int index, int isSet, long start, long end, ArrowBuf buffer) { assert index >= 0; final long dataLength = end - start; handleSafe(index, (int) dataLength); fillHoles(index); BitVectorHelper.setValidityBit(validityBuffer, index, isSet); final long startOffset = offsetBuffer.getLong((long) index * OFFSET_WIDTH); offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, startOffset + dataLength); valueBuffer.setBytes(startOffset, buffer, start, dataLength); lastSet = index; } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param start start position of data in buffer * @param length length of data in buffer * @param buffer data buffer containing the variable width element to be stored in the vector */ public void set(int index, long start, int length, ArrowBuf buffer) { assert index >= 0; fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); final long startOffset = offsetBuffer.getLong((long) index * OFFSET_WIDTH); offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, startOffset + length); final ArrowBuf bb = buffer.slice(start, length); valueBuffer.setBytes(startOffset, bb); lastSet = index; } /** * Same as {@link #set(int, int, long, long, ArrowBuf)} except that it handles the case when index * is greater than or equal to current value capacity of the vector. * * @param index position of the new value * @param start start position of data in buffer * @param length length of data in buffer * @param buffer data buffer containing the variable width element to be stored in the vector */ public void setSafe(int index, long start, int length, ArrowBuf buffer) { assert index >= 0; handleSafe(index, length); fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); final long startOffset = offsetBuffer.getLong((long) index * OFFSET_WIDTH); offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, startOffset + length); final ArrowBuf bb = buffer.slice(start, length); valueBuffer.setBytes(startOffset, bb); lastSet = index; } /*----------------------------------------------------------------* | | | helper methods for setters | | | *----------------------------------------------------------------*/ protected final void fillHoles(int index) { for (int i = lastSet + 1; i < index; i++) { setBytes(i, emptyByteArray, 0, emptyByteArray.length); } lastSet = index - 1; } protected final void setBytes(int index, byte[] value, int start, int length) { /* end offset of current last element in the vector. this will * be the start offset of new element we are trying to store. */ final long startOffset = getStartOffset(index); /* set new end offset */ offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, startOffset + length); /* store the var length data in value buffer */ valueBuffer.setBytes(startOffset, value, start, length); } /** * Gets the starting offset of a record, given its index. * * @param index index of the record. * @return the starting offset of the record. */ protected final long getStartOffset(int index) { return offsetBuffer.getLong((long) index * OFFSET_WIDTH); } protected final void handleSafe(int index, int dataLength) { /* * IMPORTANT: * value buffer for variable length vectors moves independent * of the companion validity and offset buffers. This is in * contrast to what we have for fixed width vectors. * * Here there is no concept of getValueCapacity() in the * data stream. getValueCapacity() is applicable only to validity * and offset buffers. * * So even though we may have setup an initial capacity of 1024 * elements in the vector, it is quite possible * that we need to reAlloc() the data buffer when we are setting * the 5th element in the vector simply because previous * variable length elements have exhausted the buffer capacity. * However, we really don't need to reAlloc() validity and * offset buffers until we try to set the 1025th element * This is why we do a separate check for safe methods to * determine which buffer needs reallocation. */ while (index >= getValueCapacity()) { reallocValidityAndOffsetBuffers(); } final long startOffset = lastSet < 0 ? 0L : getStartOffset(lastSet + 1); while (valueBuffer.capacity() < (startOffset + dataLength)) { reallocDataBuffer(); } } /** * Method used by Json Writer to read a variable width element from the variable width vector and * write to Json. * *

    This method should not be used externally. * * @param data buffer storing the variable width vector elements * @param offset buffer storing the offsets of variable width vector elements * @param index position of the element in the vector * @return array of bytes */ public static byte[] get(final ArrowBuf data, final ArrowBuf offset, int index) { final long currentStartOffset = offset.getLong((long) index * OFFSET_WIDTH); final int dataLength = (int) (offset.getLong((long) (index + 1) * OFFSET_WIDTH) - currentStartOffset); final byte[] result = new byte[dataLength]; data.getBytes(currentStartOffset, result, 0, dataLength); return result; } /** * Method used by Json Reader to explicitly set the offsets of the variable width vector data. The * method takes care of allocating the memory for offsets if the caller hasn't done so. * *

    This method should not be used externally. * * @param buffer ArrowBuf to store offsets for variable width elements * @param allocator memory allocator * @param valueCount number of elements * @param index position of the element * @param value offset of the element * @return buffer holding the offsets */ public static ArrowBuf set( ArrowBuf buffer, BufferAllocator allocator, int valueCount, int index, long value) { if (buffer == null) { buffer = allocator.buffer((long) valueCount * OFFSET_WIDTH); } buffer.setLong((long) index * OFFSET_WIDTH, value); if (index == (valueCount - 1)) { buffer.writerIndex((long) valueCount * OFFSET_WIDTH); } return buffer; } /** * Copy a cell value from a particular index in source vector to a particular position in this * vector. * * @param fromIndex position to copy from in source vector * @param thisIndex position to copy to in this vector * @param from source vector */ @Override public void copyFrom(int fromIndex, int thisIndex, ValueVector from) { Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); if (from.isNull(fromIndex)) { fillHoles(thisIndex); BitVectorHelper.unsetBit(this.validityBuffer, thisIndex); final long copyStart = offsetBuffer.getLong((long) thisIndex * OFFSET_WIDTH); offsetBuffer.setLong((long) (thisIndex + 1) * OFFSET_WIDTH, copyStart); } else { final long start = from.getOffsetBuffer().getLong((long) fromIndex * OFFSET_WIDTH); final long end = from.getOffsetBuffer().getLong((long) (fromIndex + 1) * OFFSET_WIDTH); final long length = end - start; fillHoles(thisIndex); BitVectorHelper.setBit(this.validityBuffer, thisIndex); final long copyStart = getStartOffset(thisIndex); from.getDataBuffer().getBytes(start, this.valueBuffer, copyStart, (int) length); offsetBuffer.setLong((long) (thisIndex + 1) * OFFSET_WIDTH, copyStart + length); } lastSet = thisIndex; } /** * Same as {@link #copyFrom(int, int, ValueVector)} except that it handles the case when the * capacity of the vector needs to be expanded before copy. * * @param fromIndex position to copy from in source vector * @param thisIndex position to copy to in this vector * @param from source vector */ @Override public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) { Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); if (from.isNull(fromIndex)) { handleSafe(thisIndex, 0); fillHoles(thisIndex); BitVectorHelper.unsetBit(this.validityBuffer, thisIndex); final long copyStart = offsetBuffer.getLong((long) thisIndex * OFFSET_WIDTH); offsetBuffer.setLong((long) (thisIndex + 1) * OFFSET_WIDTH, copyStart); } else { final long start = from.getOffsetBuffer().getLong((long) fromIndex * OFFSET_WIDTH); final long end = from.getOffsetBuffer().getLong((long) (fromIndex + 1) * OFFSET_WIDTH); final int length = (int) (end - start); handleSafe(thisIndex, length); fillHoles(thisIndex); BitVectorHelper.setBit(this.validityBuffer, thisIndex); final long copyStart = getStartOffset(thisIndex); from.getDataBuffer().getBytes(start, this.valueBuffer, copyStart, length); offsetBuffer.setLong((long) (thisIndex + 1) * OFFSET_WIDTH, copyStart + length); } lastSet = thisIndex; } @Override public ArrowBufPointer getDataPointer(int index) { return getDataPointer(index, new ArrowBufPointer()); } @Override public ArrowBufPointer getDataPointer(int index, ArrowBufPointer reuse) { if (isNull(index)) { reuse.set(null, 0, 0); } else { long offset = getStartOffset(index); int length = (int) (getEndOffset(index) - offset); reuse.set(valueBuffer, offset, length); } return reuse; } @Override public int hashCode(int index) { return hashCode(index, null); } @Override public int hashCode(int index, ArrowBufHasher hasher) { if (isNull(index)) { return ArrowBufPointer.NULL_HASH_CODE; } final long start = getStartOffset(index); final long end = getEndOffset(index); return ByteFunctionHelpers.hash(hasher, this.getDataBuffer(), start, end); } @Override public OUT accept(VectorVisitor visitor, IN value) { return visitor.visit(this, value); } protected final long getEndOffset(int index) { return offsetBuffer.getLong((long) (index + 1) * OFFSET_WIDTH); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java000066400000000000000000000315531515451615200314040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.Collections; import java.util.Iterator; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReferenceManager; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.util.DataSizeRoundingUtil; import org.apache.arrow.vector.util.TransferPair; import org.apache.arrow.vector.util.ValueVectorUtility; /** * Base class for other Arrow Vector Types. Provides basic functionality around memory management. */ public abstract class BaseValueVector implements ValueVector { public static final String MAX_ALLOCATION_SIZE_PROPERTY = "arrow.vector.max_allocation_bytes"; public static final long MAX_ALLOCATION_SIZE = Long.getLong(MAX_ALLOCATION_SIZE_PROPERTY, Long.MAX_VALUE); /* * For all fixed width vectors, the value and validity buffers are sliced from a single buffer. * Similarly, for variable width vectors, the offsets and validity buffers are sliced from a * single buffer. To ensure the single buffer is power-of-2 size, the initial value allocation * should be less than power-of-2. For IntVectors, this comes to 3970*4 (15880) for the data * buffer and 504 bytes for the validity buffer, totalling to 16384 (2^16). */ public static final int INITIAL_VALUE_ALLOCATION = 3970; protected final BufferAllocator allocator; protected volatile FieldReader fieldReader; protected ArrowBuf validityBuffer; protected int valueCount; protected BaseValueVector(BufferAllocator allocator) { this.allocator = Preconditions.checkNotNull(allocator, "allocator cannot be null"); } @Override public abstract String getName(); /** Representation of vector suitable for debugging. */ @Override public String toString() { return ValueVectorUtility.getToString(this, 0, getValueCount()); } @Override public void clear() {} @Override public void close() { clear(); } @Override public TransferPair getTransferPair(BufferAllocator allocator) { return getTransferPair(getName(), allocator); } @Override public Iterator iterator() { return Collections.emptyIterator(); } /** * Checks to ensure that every buffer vv uses has a positive reference count, throws * if this precondition isn't met. Returns true otherwise. */ public static boolean checkBufRefs(final ValueVector vv) { for (final ArrowBuf buffer : vv.getBuffers(false)) { if (buffer.refCnt() <= 0) { throw new IllegalStateException("zero refcount"); } } return true; } @Override public BufferAllocator getAllocator() { return allocator; } void compareTypes(BaseValueVector target, String caller) { if (this.getMinorType() != target.getMinorType()) { throw new UnsupportedOperationException(caller + " should have vectors of exact same type"); } } protected ArrowBuf releaseBuffer(ArrowBuf buffer) { buffer.getReferenceManager().release(); buffer = allocator.getEmpty(); return buffer; } /** * 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); } /* round up bytes for the validity buffer for the given valueCount */ private static long roundUp8ForValidityBuffer(long valueCount) { return ((valueCount + 63) >> 6) << 3; } long computeCombinedBufferSize(int valueCount, int typeWidth) { Preconditions.checkArgument(valueCount >= 0, "valueCount must be >= 0"); Preconditions.checkArgument(typeWidth >= 0, "typeWidth must be >= 0"); // compute size of validity buffer. long bufferSize = roundUp8ForValidityBuffer(valueCount); // add the size of the value buffer. if (typeWidth == 0) { // for boolean type, value-buffer and validity-buffer are of same size. bufferSize *= 2; } else { bufferSize += DataSizeRoundingUtil.roundUpTo8Multiple((long) valueCount * typeWidth); } return allocator.getRoundingPolicy().getRoundedSize(bufferSize); } /** * Each vector has a different reader that implements the FieldReader interface. Overridden * methods must make sure to return the correct concrete reader implementation. * * @return Returns a lambda that initializes a reader when called. */ protected abstract FieldReader getReaderImpl(); /** * Default implementation to create a reader for the vector. Depends on the individual vector * class' implementation of {@link #getReaderImpl} to initialize the reader appropriately. * * @return Concrete instance of FieldReader by using double-checked locking. */ @Override public FieldReader getReader() { FieldReader reader = fieldReader; if (reader != null) { return reader; } synchronized (this) { if (fieldReader == null) { fieldReader = getReaderImpl(); } return fieldReader; } } /** Container for primitive vectors (1 for the validity bit-mask and one to hold the values). */ static class DataAndValidityBuffers { private ArrowBuf dataBuf; private ArrowBuf validityBuf; DataAndValidityBuffers(ArrowBuf dataBuf, ArrowBuf validityBuf) { this.dataBuf = dataBuf; this.validityBuf = validityBuf; } ArrowBuf getDataBuf() { return dataBuf; } ArrowBuf getValidityBuf() { return validityBuf; } } DataAndValidityBuffers allocFixedDataAndValidityBufs(int valueCount, int typeWidth) { long bufferSize = computeCombinedBufferSize(valueCount, typeWidth); assert bufferSize <= MAX_ALLOCATION_SIZE; long validityBufferSize; long dataBufferSize; if (typeWidth == 0) { validityBufferSize = dataBufferSize = bufferSize / 2; } else { // Due to the rounding policy, the bufferSize could be greater than the // requested size. Utilize the allocated buffer fully.; long actualCount = (long) ((bufferSize * 8.0) / (8 * typeWidth + 1)); do { validityBufferSize = roundUp8ForValidityBuffer(actualCount); dataBufferSize = DataSizeRoundingUtil.roundUpTo8Multiple(actualCount * typeWidth); if (validityBufferSize + dataBufferSize <= bufferSize) { break; } --actualCount; } while (true); } /* allocate combined buffer */ ArrowBuf combinedBuffer = allocator.buffer(bufferSize); /* slice into requested lengths */ ArrowBuf dataBuf = null; ArrowBuf validityBuf = null; long bufferOffset = 0; for (int numBuffers = 0; numBuffers < 2; ++numBuffers) { long len = (numBuffers == 0 ? dataBufferSize : validityBufferSize); ArrowBuf buf = combinedBuffer.slice(bufferOffset, len); buf.getReferenceManager().retain(); buf.readerIndex(0); buf.writerIndex(0); bufferOffset += len; if (numBuffers == 0) { dataBuf = buf; } else { validityBuf = buf; } } combinedBuffer.getReferenceManager().release(); return new DataAndValidityBuffers(dataBuf, validityBuf); } public static ArrowBuf transferBuffer( final ArrowBuf srcBuffer, final BufferAllocator targetAllocator) { final ReferenceManager referenceManager = srcBuffer.getReferenceManager(); return referenceManager.transferOwnership(srcBuffer, targetAllocator).getTransferredBuffer(); } @Override public void copyFrom(int fromIndex, int thisIndex, ValueVector from) { throw new UnsupportedOperationException(); } @Override 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()); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java000066400000000000000000001445211515451615200330550ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.memory.util.LargeMemoryUtil.capAtMaxInt; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Collections; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.OutOfMemoryException; 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.hash.ArrowBufHasher; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.util.CallBack; import org.apache.arrow.vector.util.OversizedAllocationException; import org.apache.arrow.vector.util.TransferPair; /** BaseVariableWidthVector is a base class providing functionality for strings/bytes types. */ public abstract class BaseVariableWidthVector extends BaseValueVector implements VariableWidthFieldVector { private static final int DEFAULT_RECORD_BYTE_COUNT = 8; private static final int INITIAL_BYTE_COUNT = INITIAL_VALUE_ALLOCATION * DEFAULT_RECORD_BYTE_COUNT; private static final int MAX_BUFFER_SIZE = (int) Math.min(MAX_ALLOCATION_SIZE, Integer.MAX_VALUE); private int lastValueCapacity; private long lastValueAllocationSizeInBytes; /* 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 valueBuffer; protected ArrowBuf offsetBuffer; protected int lastSet; protected final Field field; /** * Constructs a new instance. * * @param field The field materialized by this vector. * @param allocator The allocator to use for creating/resizing buffers */ public BaseVariableWidthVector(Field field, final BufferAllocator allocator) { super(allocator); this.field = field; lastValueAllocationSizeInBytes = INITIAL_BYTE_COUNT; // -1 because we require one extra slot for the offset array. lastValueCapacity = INITIAL_VALUE_ALLOCATION - 1; valueCount = 0; lastSet = -1; offsetBuffer = allocator.getEmpty(); validityBuffer = allocator.getEmpty(); valueBuffer = allocator.getEmpty(); } @Override public String getName() { return field.getName(); } /* TODO: * see if getNullCount() can be made faster -- O(1) */ /* TODO: * Once the entire hierarchy has been refactored, move common functions * like getNullCount() to top level * base class BaseValueVector. * * Along with this, some class members (validityBuffer) can also be * abstracted out to top level base class. * * Right now BaseValueVector is the top level base class for other * vector types in ValueVector hierarchy (non-nullable) and those * vectors have not yet been refactored/removed so moving things to * the top class as of now is not a good idea. */ /** * 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. * * @return buffer */ @Override public ArrowBuf getValidityBuffer() { return validityBuffer; } /** * Get the buffer that stores the data for elements in the vector. * * @return buffer */ @Override public ArrowBuf getDataBuffer() { return valueBuffer; } /** * buffer that stores the offsets for elements in the vector. This operation is not supported for * fixed-width vectors. * * @return buffer */ @Override public ArrowBuf getOffsetBuffer() { return offsetBuffer; } /** * Get the memory address of buffer that stores the offsets for elements in the vector. * * @return starting address of the buffer */ @Override public long getOffsetBufferAddress() { return offsetBuffer.memoryAddress(); } /** * Get the memory address of buffer that manages the validity (NULL or NON-NULL nature) of * elements in the vector. * * @return starting address of the buffer */ @Override public long getValidityBufferAddress() { return validityBuffer.memoryAddress(); } /** * Get the memory address of buffer that stores the data for elements in the vector. * * @return starting address of the buffer */ @Override public long getDataBufferAddress() { return valueBuffer.memoryAddress(); } /** * Sets the desired value capacity for the vector. This function doesn't allocate any memory for * the vector. * * @param valueCount desired number of elements in the vector */ @Override public void setInitialCapacity(int valueCount) { final long size = (long) valueCount * DEFAULT_RECORD_BYTE_COUNT; checkDataBufferSize(size); computeAndCheckOffsetsBufferSize(valueCount); lastValueAllocationSizeInBytes = (int) size; lastValueCapacity = valueCount; } /** * Sets the desired value capacity for the vector. This function doesn't allocate any memory for * the vector. * * @param valueCount desired number of elements in the vector * @param density average number of bytes per variable width element */ @Override public void setInitialCapacity(int valueCount, double density) { long size = Math.max((long) (valueCount * density), 1L); checkDataBufferSize(size); computeAndCheckOffsetsBufferSize(valueCount); lastValueAllocationSizeInBytes = (int) size; lastValueCapacity = valueCount; } /** * Get the density of this ListVector. * * @return density */ public double getDensity() { if (valueCount == 0) { return 0.0D; } final int startOffset = getStartOffset(0); final int endOffset = getStartOffset(valueCount); final double totalListSize = endOffset - startOffset; return totalListSize / valueCount; } /** * Get the current capacity which does not exceed either validity buffer or offset buffer. Note: * Here the `getValueCapacity` has no relationship with the value buffer. * * @return number of elements that vector can hold. */ @Override public int getValueCapacity() { final int offsetValueCapacity = Math.max(getOffsetBufferValueCapacity() - 1, 0); return Math.min(offsetValueCapacity, getValidityBufferValueCapacity()); } private int getValidityBufferValueCapacity() { return capAtMaxInt(validityBuffer.capacity() * 8); } private int getOffsetBufferValueCapacity() { return capAtMaxInt(offsetBuffer.capacity() / OFFSET_WIDTH); } /** zero out the vector and the data in associated buffers. */ public void zeroVector() { initValidityBuffer(); initOffsetBuffer(); valueBuffer.setZero(0, valueBuffer.capacity()); } /* zero out the validity buffer */ private void initValidityBuffer() { validityBuffer.setZero(0, validityBuffer.capacity()); } /* zero out the offset buffer */ private void initOffsetBuffer() { offsetBuffer.setZero(0, offsetBuffer.capacity()); } /** * Reset the vector to initial state. Same as {@link #zeroVector()}. Note that this method doesn't * release any memory. */ @Override public void reset() { zeroVector(); lastSet = -1; valueCount = 0; } /** Close the vector and release the associated buffers. */ @Override public void close() { clear(); } /** Same as {@link #close()}. */ @Override public void clear() { validityBuffer = releaseBuffer(validityBuffer); valueBuffer = releaseBuffer(valueBuffer); offsetBuffer = releaseBuffer(offsetBuffer); lastSet = -1; valueCount = 0; } /** * Get the inner vectors. * * @return the inner vectors for this field as defined by the TypeLayout * @deprecated This API will be removed as the current implementations no longer support inner * vectors. */ @Deprecated @Override public List getFieldInnerVectors() { throw new UnsupportedOperationException("There are no inner vectors. Use getFieldBuffers"); } /** * Initialize the children in schema for this Field. This operation is a NO-OP for scalar types * since they don't have any children. * * @param children the schema * @throws IllegalArgumentException if children is a non-empty list for scalar types. */ @Override public void initializeChildrenFromFields(List children) { if (!children.isEmpty()) { throw new IllegalArgumentException("primitive type vector cannot have children"); } } /** * Get the inner child vectors. * * @return list of child vectors for complex types, empty list for scalar vector types */ @Override public List getChildrenFromFields() { return Collections.emptyList(); } /** * Load the buffers of this vector with provided source buffers. The caller manages the source * buffers and populates them before invoking this method. * * @param fieldNode the fieldNode indicating the value count * @param ownBuffers the buffers for this Field (own buffers only, children not included) */ @Override public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers) { ArrowBuf bitBuffer = ownBuffers.get(0); ArrowBuf offBuffer = ownBuffers.get(1); ArrowBuf dataBuffer = ownBuffers.get(2); validityBuffer.getReferenceManager().release(); validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); offsetBuffer.getReferenceManager().release(); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); valueBuffer.getReferenceManager().release(); valueBuffer = dataBuffer.getReferenceManager().retain(dataBuffer, allocator); lastSet = fieldNode.getLength() - 1; valueCount = fieldNode.getLength(); } /** * Get the buffers belonging to this vector. * * @return the inner buffers. */ @Override public List getFieldBuffers() { // before flight/IPC, we must bring the vector to a consistent state. // this is because, it is possible that the offset buffers of some trailing values // are not updated. this may cause some data in the data buffer being lost. // for details, please see TestValueVector#testUnloadVariableWidthVector. fillHoles(valueCount); List result = new ArrayList<>(3); setReaderAndWriterIndex(); result.add(validityBuffer); result.add(offsetBuffer); result.add(valueBuffer); return result; } /** * Export the buffers of the fields for C Data Interface. This method traverse the buffers and * export buffer and buffer's memory address into a list of buffers and a pointer to the list of * buffers. */ @Override public void exportCDataBuffers(List buffers, ArrowBuf buffersPtr, long nullValue) { // before flight/IPC, we must bring the vector to a consistent state. // this is because, it is possible that the offset buffers of some trailing values // are not updated. this may cause some data in the data buffer being lost. // for details, please see TestValueVector#testUnloadVariableWidthVector. fillHoles(valueCount); exportBuffer(validityBuffer, buffers, buffersPtr, nullValue, true); if (offsetBuffer.capacity() == 0) { // Empty offset buffer is allowed for historical reason. // To export it through C Data interface, we need to allocate a buffer with one offset. // We set `retain = false` to explicitly not increase the ref count for the exported buffer. // The ref count of the newly created buffer (i.e., 1) already represents the usage // at imported side. exportBuffer(allocateOffsetBuffer(OFFSET_WIDTH), buffers, buffersPtr, nullValue, false); } else { exportBuffer(offsetBuffer, buffers, buffersPtr, nullValue, true); } exportBuffer(valueBuffer, buffers, buffersPtr, nullValue, true); } /** Set the reader and writer indexes for the inner buffers. */ private void setReaderAndWriterIndex() { validityBuffer.readerIndex(0); offsetBuffer.readerIndex(0); valueBuffer.readerIndex(0); if (valueCount == 0) { validityBuffer.writerIndex(0); valueBuffer.writerIndex(0); } else { final int lastDataOffset = getStartOffset(valueCount); 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()}. */ @Override public void allocateNew() { allocateNew(lastValueAllocationSizeInBytes, lastValueCapacity); } /** * Allocate memory for the vector. We internally use a default value count of 4096 to allocate * memory for at least these many elements in the vector. See {@link #allocateNew(long, int)} for * allocating memory for specific number of elements in the vector. * * @return false if memory allocation fails, true otherwise. */ @Override public boolean allocateNewSafe() { try { allocateNew(lastValueAllocationSizeInBytes, lastValueCapacity); return true; } catch (Exception e) { return false; } } /** * Allocate memory for the vector to support storing at least the provided number of elements in * the vector. This method must be called prior to using the ValueVector. * * @param totalBytes desired total memory capacity * @param valueCount the desired number of elements in the vector * @throws org.apache.arrow.memory.OutOfMemoryException if memory allocation fails */ @Override public void allocateNew(long totalBytes, int valueCount) { assert totalBytes >= 0; checkDataBufferSize(totalBytes); computeAndCheckOffsetsBufferSize(valueCount); /* we are doing a new allocation -- release the current buffers */ clear(); try { allocateBytes(totalBytes, valueCount); } catch (Exception e) { clear(); throw e; } } @Override public void allocateNew(int valueCount) { allocateNew(lastValueAllocationSizeInBytes, valueCount); } /* Check if the data buffer size is within bounds. */ private void checkDataBufferSize(long size) { if (size > MAX_BUFFER_SIZE || size < 0) { throw new OversizedAllocationException( "Memory required for vector " + "is (" + size + "), which is overflow or more than max allowed (" + MAX_BUFFER_SIZE + "). " + "You could consider using LargeVarCharVector/LargeVarBinaryVector for large strings/large bytes types"); } } /* * Compute the buffer size required for 'valueCount' offsets and validity, and check if it's * within bounds. */ private long computeAndCheckOffsetsBufferSize(int valueCount) { /* to track the end offset of last data element in vector, we need * an additional slot in offset buffer. */ final long size = computeCombinedBufferSize(valueCount + 1, OFFSET_WIDTH); if (size > MAX_BUFFER_SIZE) { throw new OversizedAllocationException( "Memory required for vector capacity " + valueCount + " is (" + size + "), which is more than max allowed (" + MAX_BUFFER_SIZE + ")"); } return size; } /* allocate the inner buffers */ private void allocateBytes(final long valueBufferSize, final int valueCount) { /* allocate data buffer */ long curSize = valueBufferSize; valueBuffer = allocator.buffer(curSize); valueBuffer.readerIndex(0); /* allocate offset buffer and validity buffer */ DataAndValidityBuffers buffers = allocFixedDataAndValidityBufs(valueCount + 1, OFFSET_WIDTH); offsetBuffer = buffers.getDataBuf(); validityBuffer = buffers.getValidityBuf(); initOffsetBuffer(); initValidityBuffer(); lastValueCapacity = getValueCapacity(); lastValueAllocationSizeInBytes = capAtMaxInt(valueBuffer.capacity()); } /* allocate offset buffer */ private ArrowBuf allocateOffsetBuffer(final long size) { final int curSize = (int) size; ArrowBuf offsetBuffer = allocator.buffer(curSize); offsetBuffer.readerIndex(0); offsetBuffer.setZero(0, offsetBuffer.capacity()); return offsetBuffer; } /* allocate validity buffer */ @Override protected void allocateValidityBuffer(final long size) { super.allocateValidityBuffer(size); } /** * Resize the vector to increase the capacity. The internal behavior is to double the current * value capacity. */ @Override public void reAlloc() { reallocDataBuffer(); reallocValidityAndOffsetBuffers(); } /** * Reallocate the data buffer. Data Buffer stores the actual data for VARCHAR or VARBINARY * elements in the vector. The behavior is to double the size of buffer. * * @throws OversizedAllocationException if the desired new size is more than max allowed * @throws OutOfMemoryException if the internal memory allocation fails */ public void reallocDataBuffer() { final long currentBufferCapacity = valueBuffer.capacity(); long newAllocationSize = currentBufferCapacity * 2; if (newAllocationSize == 0) { if (lastValueAllocationSizeInBytes > 0) { newAllocationSize = lastValueAllocationSizeInBytes; } else { newAllocationSize = INITIAL_BYTE_COUNT * 2L; } } reallocDataBuffer(newAllocationSize); } /** * Reallocate the data buffer to given size. Data Buffer stores the actual data for VARCHAR or * VARBINARY elements in the vector. The actual allocate size may be larger than the request one * because it will round up the provided value to the nearest power of two. * * @param desiredAllocSize the desired new allocation size * @throws OversizedAllocationException if the desired new size is more than max allowed * @throws OutOfMemoryException if the internal memory allocation fails */ public void reallocDataBuffer(long desiredAllocSize) { if (desiredAllocSize == 0) { return; } final long newAllocationSize = Math.min(CommonUtil.nextPowerOfTwo(desiredAllocSize), MAX_BUFFER_SIZE); assert newAllocationSize >= 1; if (newAllocationSize < desiredAllocSize) { checkDataBufferSize(desiredAllocSize); } final ArrowBuf newBuf = allocator.buffer(newAllocationSize); newBuf.setBytes(0, valueBuffer, 0, valueBuffer.capacity()); valueBuffer.getReferenceManager().release(); valueBuffer = newBuf; lastValueAllocationSizeInBytes = valueBuffer.capacity(); } /** * Reallocate the validity and offset buffers for this vector. Validity buffer is used to track * the NULL or NON-NULL nature of elements in the vector and offset buffer is used to store the * lengths of variable width elements in the vector. * *

    Note that data buffer for variable length vectors moves independent of the companion * validity and offset buffers. This is in contrast to what we have for fixed width vectors. * *

    So even though we may have setup an initial capacity of 1024 elements in the vector, it is * quite possible that we need to reAlloc() the data buffer when we are setting the 5th element in * the vector simply because previous variable length elements have exhausted the buffer capacity. * However, we really don't need to reAlloc() validity and offset buffers until we try to set the * 1025th element This is why we do a separate check for safe methods to determine which buffer * needs reallocation. * * @throws OversizedAllocationException if the desired new size is more than max allowed * @throws OutOfMemoryException if the internal memory allocation fails */ public void reallocValidityAndOffsetBuffers() { int targetOffsetCount = capAtMaxInt((offsetBuffer.capacity() / OFFSET_WIDTH) * 2); if (targetOffsetCount == 0) { if (lastValueCapacity > 0) { targetOffsetCount = (lastValueCapacity + 1); } else { targetOffsetCount = 2 * (INITIAL_VALUE_ALLOCATION + 1); } } computeAndCheckOffsetsBufferSize(targetOffsetCount); DataAndValidityBuffers buffers = allocFixedDataAndValidityBufs(targetOffsetCount, OFFSET_WIDTH); final ArrowBuf newOffsetBuffer = buffers.getDataBuf(); newOffsetBuffer.setBytes(0, offsetBuffer, 0, offsetBuffer.capacity()); newOffsetBuffer.setZero( offsetBuffer.capacity(), newOffsetBuffer.capacity() - offsetBuffer.capacity()); offsetBuffer.getReferenceManager().release(); offsetBuffer = newOffsetBuffer; final ArrowBuf newValidityBuffer = buffers.getValidityBuf(); newValidityBuffer.setBytes(0, validityBuffer, 0, validityBuffer.capacity()); newValidityBuffer.setZero( validityBuffer.capacity(), newValidityBuffer.capacity() - validityBuffer.capacity()); validityBuffer.getReferenceManager().release(); validityBuffer = newValidityBuffer; lastValueCapacity = getValueCapacity(); } /** * Get the size (number of bytes) of underlying data buffer. * * @return number of bytes in the data buffer */ @Override public int getByteCapacity() { return capAtMaxInt(valueBuffer.capacity()); } @Override public int sizeOfValueBuffer() { if (valueCount == 0) { return 0; } return offsetBuffer.getInt((long) valueCount * OFFSET_WIDTH); } /** * Get the size (number of bytes) of underlying buffers used by this vector. * * @return size of underlying buffers. */ @Override public int getBufferSize() { return getBufferSizeFor(this.valueCount); } /** * Get the potential buffer size for a particular number of records. * * @param valueCount desired number of elements in the vector * @return estimated size of underlying buffers if the vector holds a given number of elements */ @Override public int getBufferSizeFor(final int valueCount) { if (valueCount == 0) { return 0; } 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); return validityBufferSize + offsetBufferSize + dataBufferSize; } /** * Get information about how this field is materialized. * * @return the field corresponding to this vector */ @Override public Field getField() { return field; } /** * Return the underlying buffers associated with this vector. Note that this doesn't impact the * reference counts for this buffer so it only should be used for in-context access. Also note * that this buffer changes regularly thus external classes shouldn't hold a reference to it * (unless they change it). * * @param clear Whether to clear vector before returning; the buffers will still be refcounted but * the returned array will be the only reference to them * @return The underlying {@link ArrowBuf buffers} that is used by this vector instance. */ @Override public ArrowBuf[] getBuffers(boolean clear) { final ArrowBuf[] buffers; setReaderAndWriterIndex(); if (getBufferSize() == 0) { buffers = new ArrowBuf[0]; } else { buffers = new ArrowBuf[3]; buffers[0] = validityBuffer; buffers[1] = offsetBuffer; buffers[2] = valueBuffer; } if (clear) { for (final ArrowBuf buffer : buffers) { buffer.getReferenceManager().retain(); } clear(); } return buffers; } /** Validate the scalar values held by this vector. */ public void validateScalars() { // No validation by default. } /** * Construct a transfer pair of this vector and another vector of same type. * * @param field The field materialized by this vector. * @param allocator allocator for the target vector * @param callBack not used * @return TransferPair */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { return getTransferPair(field, allocator); } /** * Construct a transfer pair of this vector and another vector of same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @param callBack not used * @return TransferPair */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return getTransferPair(ref, allocator); } /** * Construct a transfer pair of this vector and another vector of same type. * * @param allocator allocator for the target vector * @return TransferPair */ @Override public TransferPair getTransferPair(BufferAllocator allocator) { return getTransferPair(getName(), allocator); } /** * Construct a transfer pair of this vector and another vector of same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return TransferPair */ @Override public abstract TransferPair getTransferPair(String ref, BufferAllocator allocator); /** * Construct a transfer pair of this vector and another vector of same type. * * @param field The field materialized by this vector. * @param allocator allocator for the target vector * @return TransferPair */ @Override public abstract TransferPair getTransferPair(Field field, BufferAllocator allocator); /** * Transfer this vector'data to another vector. The memory associated with this vector is * transferred to the allocator of target vector for accounting and management purposes. * * @param target destination vector for transfer */ public void transferTo(BaseVariableWidthVector target) { compareTypes(target, "transferTo"); target.clear(); target.validityBuffer = transferBuffer(validityBuffer, target.allocator); target.valueBuffer = transferBuffer(valueBuffer, target.allocator); target.offsetBuffer = transferBuffer(offsetBuffer, target.allocator); target.setLastSet(this.lastSet); if (this.valueCount > 0) { target.setValueCount(this.valueCount); } clear(); } /** * Slice this vector at desired index and length and transfer the corresponding data to the target * vector. * * @param startIndex start position of the split in source vector. * @param length length of the split. * @param target destination vector */ public void splitAndTransferTo(int startIndex, int length, BaseVariableWidthVector target) { Preconditions.checkArgument( startIndex >= 0 && length >= 0 && startIndex + length <= valueCount, "Invalid parameters startIndex: %s, length: %s for valueCount: %s", startIndex, length, valueCount); compareTypes(target, "splitAndTransferTo"); target.clear(); if (length > 0) { splitAndTransferValidityBuffer(startIndex, length, target); splitAndTransferOffsetBuffer(startIndex, length, target); target.setLastSet(length - 1); target.setValueCount(length); } } /** * Transfer the offsets along with data. Unlike the data buffer, we cannot simply slice the offset * buffer for split and transfer. The reason is that offsets in the target vector have to be * adjusted and made relative to the staring offset in source vector from the start index of * split. This is why, we need to explicitly allocate the offset buffer and set the adjusted * offsets in the target vector. */ private void splitAndTransferOffsetBuffer( int startIndex, int length, BaseVariableWidthVector target) { final int start = getStartOffset(startIndex); final int end = getStartOffset(startIndex + length); final int dataLength = end - start; if (start == 0) { final ArrowBuf slicedOffsetBuffer = offsetBuffer.slice( startIndex * ((long) OFFSET_WIDTH), (1 + length) * ((long) OFFSET_WIDTH)); target.offsetBuffer = transferBuffer(slicedOffsetBuffer, target.allocator); } else { target.offsetBuffer = target.allocateOffsetBuffer((long) (length + 1) * OFFSET_WIDTH); for (int i = 0; i < length + 1; i++) { final int relativeSourceOffset = getStartOffset(startIndex + i) - start; target.offsetBuffer.setInt((long) i * OFFSET_WIDTH, relativeSourceOffset); } } final ArrowBuf slicedBuffer = valueBuffer.slice(start, dataLength); target.valueBuffer = transferBuffer(slicedBuffer, target.allocator); } @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(); } final ArrowBuf slicedValidityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); target.validityBuffer = transferBuffer(slicedValidityBuffer, target.allocator); } /*----------------------------------------------------------------* | | | common getters and setters | | | *----------------------------------------------------------------*/ /** * Get the number of elements that are null in the vector. * * @return the number of null elements. */ @Override public int getNullCount() { return BitVectorHelper.getNullCount(validityBuffer, valueCount); } /** * Check if the given index is within the current value capacity of the vector. * * @param index position to check * @return true if index is within the current value capacity */ public boolean isSafe(int index) { return index < getValueCapacity(); } /** * Check if element at given index is null. * * @param index position of element * @return true if element at given index is null */ @Override public boolean isNull(int index) { return (isSet(index) == 0); } /** * Same as {@link #isNull(int)}. * * @param index position of element * @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; } /** * Get the value count of vector. This will always be zero unless setValueCount(int) has been * called prior to calling this. * * @return valueCount for the vector */ @Override public int getValueCount() { return valueCount; } /** * Sets the value count for the vector. * * @param valueCount value count */ @Override public void setValueCount(int valueCount) { assert valueCount >= 0; this.valueCount = valueCount; while (valueCount > getValueCapacity()) { reallocValidityAndOffsetBuffers(); } fillHoles(valueCount); lastSet = valueCount - 1; setReaderAndWriterIndex(); } /** * Create holes in the vector upto the given index (exclusive). Holes will be created from the * current last set position in the vector. * * @param index target index */ @Override public void fillEmpties(int index) { handleSafe(index, emptyByteArray.length); fillHoles(index); lastSet = index - 1; } /** * Set the index of last non-null element in the vector. It is important to call this method with * appropriate value before calling {@link #setValueCount(int)}. * * @param value desired index of last non-null element. */ @Override public void setLastSet(int value) { lastSet = value; } /** * Get the index of last non-null element in the vector. * * @return index of the last non-null element */ @Override public int getLastSet() { return lastSet; } /** * Get the starting position (offset) in the data stream for a given element in the vector. * * @param index position of the element in the vector * @return starting offset for the element */ public long getStartEnd(int index) { return offsetBuffer.getLong((long) index * OFFSET_WIDTH); } /** * Mark the particular position in the vector as non-null. * * @param index position of the element. */ @Override public void setIndexDefined(int index) { // We need to check and realloc both validity and offset buffer while (index >= getValueCapacity()) { reallocValidityAndOffsetBuffers(); } BitVectorHelper.setBit(validityBuffer, index); } /** * Sets the value length for an element. * * @param index position of the element to set * @param length length of the element */ @Override public void setValueLengthSafe(int index, int length) { assert index >= 0; handleSafe(index, length); fillHoles(index); final int startOffset = getStartOffset(index); offsetBuffer.setInt((index + 1) * ((long) OFFSET_WIDTH), startOffset + length); lastSet = index; } /** * Get the variable length element at specified index as Text. * * @param index position of element to get * @return greater than 0 length for non-null element, 0 otherwise */ @Override public int getValueLength(int index) { assert index >= 0; if (isSet(index) == 0) { return 0; } final int startOffset = getStartOffset(index); final int dataLength = getEndOffset(index) - startOffset; return dataLength; } /** * Set the variable length element at the specified index to the supplied byte array. This is same * as using {@link #set(int, byte[], int, int)} with start as 0 and length as value.length * * @param index position of the element to set * @param value array of bytes to write */ @Override public void set(int index, byte[] value) { assert index >= 0; fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); setBytes(index, value, 0, value.length); lastSet = index; } /** * Same as {@link #set(int, byte[])} except that it handles the case where index and length of new * element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param value array of bytes to write */ @Override public void setSafe(int index, byte[] value) { assert index >= 0; handleSafe(index, value.length); fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); setBytes(index, value, 0, value.length); lastSet = index; } /** * Set the variable length element at the specified index to the supplied byte array. * * @param index position of the element to set * @param value array of bytes to write * @param start start index in array of bytes * @param length length of data in array of bytes */ public void set(int index, byte[] value, int start, int length) { assert index >= 0; fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); setBytes(index, value, start, length); lastSet = index; } /** * Same as {@link #set(int, byte[], int, int)} except that it handles the case where index and * length of new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param value array of bytes to write * @param start start index in array of bytes * @param length length of data in array of bytes */ public void setSafe(int index, byte[] value, int start, int length) { assert index >= 0; handleSafe(index, length); fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); setBytes(index, value, start, length); lastSet = index; } /** * Set the variable length element at the specified index to the content in supplied ByteBuffer. * * @param index position of the element to set * @param value ByteBuffer with data * @param start start index in ByteBuffer * @param length length of data in ByteBuffer */ @Override public void set(int index, ByteBuffer value, int start, int length) { assert index >= 0; fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); final int startOffset = getStartOffset(index); offsetBuffer.setInt((index + 1) * ((long) OFFSET_WIDTH), startOffset + length); valueBuffer.setBytes(startOffset, value, start, length); lastSet = index; } /** * Same as {@link #set(int, ByteBuffer, int, int)} except that it handles the case where index and * length of new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param value ByteBuffer with data * @param start start index in ByteBuffer * @param length length of data in ByteBuffer */ public void setSafe(int index, ByteBuffer value, int start, int length) { assert index >= 0; handleSafe(index, length); fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); final int startOffset = getStartOffset(index); offsetBuffer.setInt((index + 1) * ((long) OFFSET_WIDTH), startOffset + length); valueBuffer.setBytes(startOffset, value, start, length); lastSet = index; } /** * Set the element at the given index to null. * * @param index position of element */ @Override public void setNull(int index) { // We need to check and realloc both validity and offset buffer while (index >= getValueCapacity()) { reallocValidityAndOffsetBuffers(); } BitVectorHelper.unsetBit(validityBuffer, index); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param start start position of data in buffer * @param end end position of data in buffer * @param buffer data buffer containing the variable width element to be stored in the vector */ public void set(int index, int isSet, int start, int end, ArrowBuf buffer) { assert index >= 0; final int dataLength = end - start; fillHoles(index); BitVectorHelper.setValidityBit(validityBuffer, index, isSet); final int startOffset = offsetBuffer.getInt((long) index * OFFSET_WIDTH); offsetBuffer.setInt((index + 1) * ((long) OFFSET_WIDTH), startOffset + dataLength); valueBuffer.setBytes(startOffset, buffer, start, dataLength); lastSet = index; } /** * Same as {@link #set(int, int, int, int, ArrowBuf)} except that it handles the case when index * is greater than or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param start start position of data in buffer * @param end end position of data in buffer * @param buffer data buffer containing the variable width element to be stored in the vector */ public void setSafe(int index, int isSet, int start, int end, ArrowBuf buffer) { assert index >= 0; final int dataLength = end - start; handleSafe(index, dataLength); fillHoles(index); BitVectorHelper.setValidityBit(validityBuffer, index, isSet); final int startOffset = offsetBuffer.getInt((long) index * OFFSET_WIDTH); offsetBuffer.setInt((long) (index + 1) * OFFSET_WIDTH, startOffset + dataLength); valueBuffer.setBytes(startOffset, buffer, start, dataLength); lastSet = index; } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param start start position of data in buffer * @param length length of data in buffer * @param buffer data buffer containing the variable width element to be stored in the vector */ public void set(int index, int start, int length, ArrowBuf buffer) { assert index >= 0; fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); final int startOffset = offsetBuffer.getInt((long) index * OFFSET_WIDTH); offsetBuffer.setInt((long) (index + 1) * OFFSET_WIDTH, startOffset + length); final ArrowBuf bb = buffer.slice(start, length); valueBuffer.setBytes(startOffset, bb); lastSet = index; } /** * Same as {@link #set(int, int, int, int, ArrowBuf)} except that it handles the case when index * is greater than or equal to current value capacity of the vector. * * @param index position of the new value * @param start start position of data in buffer * @param length length of data in buffer * @param buffer data buffer containing the variable width element to be stored in the vector */ public void setSafe(int index, int start, int length, ArrowBuf buffer) { assert index >= 0; handleSafe(index, length); fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); final int startOffset = getStartOffset(index); offsetBuffer.setInt((long) (index + 1) * OFFSET_WIDTH, startOffset + length); final ArrowBuf bb = buffer.slice(start, length); valueBuffer.setBytes(startOffset, bb); lastSet = index; } /*----------------------------------------------------------------* | | | helper methods for setters | | | *----------------------------------------------------------------*/ protected final void fillHoles(int index) { for (int i = lastSet + 1; i < index; i++) { setBytes(i, emptyByteArray, 0, emptyByteArray.length); } lastSet = index - 1; } protected final void setBytes(int index, byte[] value, int start, int length) { /* end offset of current last element in the vector. this will * be the start offset of new element we are trying to store. */ final int startOffset = getStartOffset(index); /* set new end offset */ offsetBuffer.setInt((long) (index + 1) * OFFSET_WIDTH, startOffset + length); /* store the var length data in value buffer */ valueBuffer.setBytes(startOffset, value, start, length); } public final int getStartOffset(int index) { return offsetBuffer.getInt((long) index * OFFSET_WIDTH); } protected final void handleSafe(int index, int dataLength) { /* * IMPORTANT: * value buffer for variable length vectors moves independent * of the companion validity and offset buffers. This is in * contrast to what we have for fixed width vectors. * * Here there is no concept of getValueCapacity() in the * data stream. getValueCapacity() is applicable only to validity * and offset buffers. * * So even though we may have setup an initial capacity of 1024 * elements in the vector, it is quite possible * that we need to reAlloc() the data buffer when we are setting * the 5th element in the vector simply because previous * variable length elements have exhausted the buffer capacity. * However, we really don't need to reAlloc() validity and * offset buffers until we try to set the 1025th element * This is why we do a separate check for safe methods to * determine which buffer needs reallocation. */ while (index >= getValueCapacity()) { reallocValidityAndOffsetBuffers(); } final long startOffset = lastSet < 0 ? 0 : getStartOffset(lastSet + 1); final long targetCapacity = startOffset + dataLength; if (valueBuffer.capacity() < targetCapacity) { reallocDataBuffer(targetCapacity); } } /** * Method used by Json Writer to read a variable width element from the variable width vector and * write to Json. * *

    This method should not be used externally. * * @param data buffer storing the variable width vector elements * @param offset buffer storing the offsets of variable width vector elements * @param index position of the element in the vector * @return array of bytes */ public static byte[] get(final ArrowBuf data, final ArrowBuf offset, int index) { final int currentStartOffset = offset.getInt((long) index * OFFSET_WIDTH); final int dataLength = offset.getInt((long) (index + 1) * OFFSET_WIDTH) - currentStartOffset; final byte[] result = new byte[dataLength]; data.getBytes(currentStartOffset, result, 0, dataLength); return result; } /** * Method used by Json Reader to explicitly set the offsets of the variable width vector data. The * method takes care of allocating the memory for offsets if the caller hasn't done so. * *

    This method should not be used externally. * * @param buffer ArrowBuf to store offsets for variable width elements * @param allocator memory allocator * @param valueCount number of elements * @param index position of the element * @param value offset of the element * @return buffer holding the offsets */ public static ArrowBuf set( ArrowBuf buffer, BufferAllocator allocator, int valueCount, int index, int value) { if (buffer == null) { buffer = allocator.buffer((long) valueCount * OFFSET_WIDTH); } buffer.setInt((long) index * OFFSET_WIDTH, value); if (index == (valueCount - 1)) { buffer.writerIndex((long) valueCount * OFFSET_WIDTH); } return buffer; } /** * Copy a cell value from a particular index in source vector to a particular position in this * vector. * * @param fromIndex position to copy from in source vector * @param thisIndex position to copy to in this vector * @param from source vector */ @Override public void copyFrom(int fromIndex, int thisIndex, ValueVector from) { Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); if (from.isNull(fromIndex)) { fillHoles(thisIndex); BitVectorHelper.unsetBit(this.validityBuffer, thisIndex); final int copyStart = offsetBuffer.getInt((long) thisIndex * OFFSET_WIDTH); offsetBuffer.setInt((long) (thisIndex + 1) * OFFSET_WIDTH, copyStart); } else { final int start = from.getOffsetBuffer().getInt((long) fromIndex * OFFSET_WIDTH); final int end = from.getOffsetBuffer().getInt((long) (fromIndex + 1) * OFFSET_WIDTH); final int length = end - start; fillHoles(thisIndex); BitVectorHelper.setBit(this.validityBuffer, thisIndex); final int copyStart = getStartOffset(thisIndex); from.getDataBuffer().getBytes(start, this.valueBuffer, copyStart, length); offsetBuffer.setInt((long) (thisIndex + 1) * OFFSET_WIDTH, copyStart + length); } lastSet = thisIndex; } /** * Same as {@link #copyFrom(int, int, ValueVector)} except that it handles the case when the * capacity of the vector needs to be expanded before copy. * * @param fromIndex position to copy from in source vector * @param thisIndex position to copy to in this vector * @param from source vector */ @Override public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) { Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); if (from.isNull(fromIndex)) { handleSafe(thisIndex, 0); fillHoles(thisIndex); BitVectorHelper.unsetBit(this.validityBuffer, thisIndex); final int copyStart = getStartOffset(thisIndex); offsetBuffer.setInt((long) (thisIndex + 1) * OFFSET_WIDTH, copyStart); } else { final int start = from.getOffsetBuffer().getInt((long) fromIndex * OFFSET_WIDTH); final int end = from.getOffsetBuffer().getInt((long) (fromIndex + 1) * OFFSET_WIDTH); final int length = end - start; handleSafe(thisIndex, length); fillHoles(thisIndex); BitVectorHelper.setBit(this.validityBuffer, thisIndex); final int copyStart = getStartOffset(thisIndex); from.getDataBuffer().getBytes(start, this.valueBuffer, copyStart, length); offsetBuffer.setInt((long) (thisIndex + 1) * OFFSET_WIDTH, copyStart + length); } lastSet = thisIndex; } @Override public ArrowBufPointer getDataPointer(int index) { return getDataPointer(index, new ArrowBufPointer()); } @Override public ArrowBufPointer getDataPointer(int index, ArrowBufPointer reuse) { if (isNull(index)) { reuse.set(null, 0, 0); } else { int offset = getStartOffset(index); int length = getEndOffset(index) - offset; reuse.set(valueBuffer, offset, length); } return reuse; } @Override public int hashCode(int index) { return hashCode(index, null); } @Override public int hashCode(int index, ArrowBufHasher hasher) { if (isNull(index)) { return ArrowBufPointer.NULL_HASH_CODE; } final int start = getStartOffset(index); final int end = getEndOffset(index); return ByteFunctionHelpers.hash(hasher, this.getDataBuffer(), start, end); } @Override public OUT accept(VectorVisitor visitor, IN value) { return visitor.visit(this, value); } /** Gets the ending offset of a record, given its index. */ public final int getEndOffset(int index) { return offsetBuffer.getInt((long) (index + 1) * OFFSET_WIDTH); } } BaseVariableWidthViewVector.java000066400000000000000000001567521515451615200336420ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.memory.util.LargeMemoryUtil.capAtMaxInt; import static org.apache.arrow.vector.util.DataSizeRoundingUtil.roundUpToMultipleOf16; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Collections; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.OutOfMemoryException; import org.apache.arrow.memory.ReusableBuffer; 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.hash.ArrowBufHasher; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.util.CallBack; import org.apache.arrow.vector.util.OversizedAllocationException; import org.apache.arrow.vector.util.TransferPair; /** * BaseVariableWidthViewVector is a base class providing functionality for strings/bytes types in * view format. */ public abstract class BaseVariableWidthViewVector extends BaseValueVector implements VariableWidthFieldVector { // A single element of a view comprises 16 bytes public static final int ELEMENT_SIZE = 16; public static final int INITIAL_VIEW_VALUE_ALLOCATION = 4096; private static final int INITIAL_BYTE_COUNT = INITIAL_VIEW_VALUE_ALLOCATION * ELEMENT_SIZE; private static final int MAX_BUFFER_SIZE = (int) Math.min(MAX_ALLOCATION_SIZE, Integer.MAX_VALUE); private int lastValueCapacity; private long lastValueAllocationSizeInBytes; /* * Variable Width View Vector comprises the following format * * Short strings, length <= 12 * | Bytes 0-3 | Bytes 4-15 | * |------------|---------------------------------------| * | length | data (padded with 0) | * |------------|---------------------------------------| * * Long strings, length > 12 * | Bytes 0-3 | Bytes 4-7 | Bytes 8-11 | Bytes 12-15 | * |------------|------------|------------|-------------| * | length | prefix | buf.index | offset | * |------------|------------|------------|-------------| * * */ // 12 byte unsigned int to track inline views public static final int INLINE_SIZE = 12; // The first 4 bytes of view are allocated for length public static final int LENGTH_WIDTH = 4; // The second 4 bytes of view are allocated for prefix width public static final int PREFIX_WIDTH = 4; // 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[] {}; // 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 lastSet; protected final Field field; /** * Constructs a new instance. * * @param field The field materialized by this vector * @param allocator The allocator to use for creating/resizing buffers */ public BaseVariableWidthViewVector(Field field, final BufferAllocator allocator) { super(allocator); this.field = field; lastValueAllocationSizeInBytes = INITIAL_BYTE_COUNT; lastValueCapacity = INITIAL_VIEW_VALUE_ALLOCATION; valueCount = 0; lastSet = -1; validityBuffer = allocator.getEmpty(); viewBuffer = allocator.getEmpty(); dataBuffers = new ArrayList<>(); } @Override public String getName() { return field.getName(); } /* TODO: * see if getNullCount() can be made faster -- O(1) */ /* TODO: * Once the entire hierarchy has been refactored, move common functions * like getNullCount() to top level * base class BaseValueVector. * * Along with this, some class members (validityBuffer) can also be * abstracted out to top level base class. * * Right now BaseValueVector is the top level base class for other * vector types in ValueVector hierarchy (non-nullable) and those * vectors have not yet been refactored/removed so moving things to * the top class as of now is not a good idea. */ /** * 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. * * @return buffer */ @Override public ArrowBuf getValidityBuffer() { return validityBuffer; } /** * Get the buffer that stores the data for elements in the vector. * * @return buffer */ @Override public ArrowBuf getDataBuffer() { return viewBuffer; } /** * Get the buffers that store the data for views in the vector. * * @return list of ArrowBuf */ public List getDataBuffers() { return dataBuffers; } /** * BaseVariableWidthViewVector doesn't support offset buffer. * * @return throws UnsupportedOperationException */ @Override public ArrowBuf getOffsetBuffer() { throw new UnsupportedOperationException( "Offset buffer is not supported in BaseVariableWidthViewVector"); } /** * BaseVariableWidthViewVector doesn't support offset buffer. * * @return throws UnsupportedOperationException */ @Override public long getOffsetBufferAddress() { throw new UnsupportedOperationException( "Offset buffer is not supported in BaseVariableWidthViewVector"); } /** * Get the memory address of buffer that manages the validity (NULL or NON-NULL nature) of * elements in the vector. * * @return starting address of the buffer */ @Override public long getValidityBufferAddress() { return validityBuffer.memoryAddress(); } /** * Get the memory address of buffer that stores the data for elements in the vector. * * @return starting address of the buffer */ @Override public long getDataBufferAddress() { return viewBuffer.memoryAddress(); } /** * Sets the desired value capacity for the vector. This function doesn't allocate any memory for * the vector. * * @param valueCount desired number of elements in the vector */ @Override public void setInitialCapacity(int valueCount) { final long size = (long) valueCount * ELEMENT_SIZE; checkDataBufferSize(size); lastValueAllocationSizeInBytes = (int) size; lastValueCapacity = valueCount; } /** * Sets the desired value capacity for the vector. This function doesn't allocate any memory for * the vector. * * @param valueCount desired number of elements in the vector * @param density average number of bytes per variable width view element */ @Override public void setInitialCapacity(int valueCount, double density) { final long size = (long) valueCount * ELEMENT_SIZE; initialDataBufferSize = (int) (valueCount * density); checkDataBufferSize(size); lastValueAllocationSizeInBytes = (int) size; lastValueCapacity = valueCount; } /** * Get the density of this ListVector. * * @return density */ public double getDensity() { if (valueCount == 0) { return 0.0D; } final double totalListSize = getTotalValueLengthUpToIndex(valueCount); return totalListSize / valueCount; } /** * Get the current capacity which does not exceed either validity buffer or value buffer. Note: * Here the `getValueCapacity` has a relationship with the value buffer. * * @return number of elements that vector can hold. */ @Override public int getValueCapacity() { final int validityCapacity = getValidityBufferValueCapacity(); final int valueBufferCapacity = Math.max(capAtMaxInt(viewBuffer.capacity() / ELEMENT_SIZE), 0); return Math.min(valueBufferCapacity, validityCapacity); } private int getValidityBufferValueCapacity() { return capAtMaxInt(validityBuffer.capacity() * 8); } /** zero out the vector and the data in associated buffers. */ public void zeroVector() { initValidityBuffer(); viewBuffer.setZero(0, viewBuffer.capacity()); clearDataBuffers(); } /* zero out the validity buffer */ private void initValidityBuffer() { validityBuffer.setZero(0, validityBuffer.capacity()); } /** Reset the vector to initial state. Note that this method doesn't release any memory. */ @Override public void reset() { zeroVector(); lastSet = -1; valueCount = 0; } /** Close the vector and release the associated buffers. */ @Override public void close() { clear(); } /** Same as {@link #close()}. */ @Override public void clear() { validityBuffer = releaseBuffer(validityBuffer); viewBuffer = releaseBuffer(viewBuffer); clearDataBuffers(); lastSet = -1; valueCount = 0; } /** Release the data buffers and clear the list. */ public void clearDataBuffers() { for (ArrowBuf buffer : dataBuffers) { releaseBuffer(buffer); } dataBuffers.clear(); } /** * Get the inner vectors. * * @return the inner vectors for this field as defined by the TypeLayout * @deprecated This API will be removed as the current implementations no longer support inner * vectors. */ @Deprecated @Override public List getFieldInnerVectors() { throw new UnsupportedOperationException("There are no inner vectors. Use getFieldBuffers"); } /** * Initialize the children in schema for this Field. This operation is a NO-OP for scalar types * since they don't have any children. * * @param children the schema * @throws IllegalArgumentException if children is a non-empty list for scalar types. */ @Override public void initializeChildrenFromFields(List children) { if (!children.isEmpty()) { throw new IllegalArgumentException("primitive type vector cannot have children"); } } /** * Get the inner child vectors. * * @return list of child vectors for complex types, empty list for scalar vector types */ @Override public List getChildrenFromFields() { return Collections.emptyList(); } /** * Load the buffers of this vector with provided source buffers. The caller manages the source * buffers and populates them before invoking this method. * * @param fieldNode the fieldNode indicating the value count * @param ownBuffers the buffers for this Field (own buffers only, children not included) */ @Override public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers) { ArrowBuf bitBuf = ownBuffers.get(0); ArrowBuf viewBuf = ownBuffers.get(1); List dataBufs = ownBuffers.subList(2, ownBuffers.size()); this.clear(); this.viewBuffer = viewBuf.getReferenceManager().retain(viewBuf, allocator); this.validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuf, allocator); for (ArrowBuf dataBuf : dataBufs) { this.dataBuffers.add(dataBuf.getReferenceManager().retain(dataBuf, allocator)); } lastSet = fieldNode.getLength() - 1; valueCount = fieldNode.getLength(); } /** * Get the buffers belonging to this vector. * * @return the inner buffers. */ @Override public List getFieldBuffers() { List result = new ArrayList<>(2 + dataBuffers.size()); setReaderAndWriterIndex(); result.add(validityBuffer); result.add(viewBuffer); // append data buffers result.addAll(dataBuffers); return result; } /** Set the reader and writer indexes for the inner buffers. */ private void setReaderAndWriterIndex() { validityBuffer.readerIndex(0); viewBuffer.readerIndex(0); if (valueCount == 0) { validityBuffer.writerIndex(0); viewBuffer.writerIndex(0); } else { validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); viewBuffer.writerIndex(valueCount * ELEMENT_SIZE); } } /** Same as {@link #allocateNewSafe()}. */ @Override public void allocateNew() { allocateNew(lastValueAllocationSizeInBytes, lastValueCapacity); } /** * Allocate memory for the vector. We internally use a default value count of 4096 to allocate * memory for at least these many elements in the vector. See {@link #allocateNew(long, int)} for * allocating memory for specific number of elements in the vector. * * @return false if memory allocation fails, true otherwise. */ @Override public boolean allocateNewSafe() { try { allocateNew(lastValueAllocationSizeInBytes, lastValueCapacity); return true; } catch (Exception e) { return false; } } /** * Allocate memory for the vector to support storing at least the provided number of elements in * the vector. This method must be called prior to using the ValueVector. * * @param totalBytes desired total memory capacity * @param valueCount the desired number of elements in the vector * @throws OutOfMemoryException if memory allocation fails */ @Override public void allocateNew(long totalBytes, int valueCount) { assert totalBytes >= 0; checkDataBufferSize(totalBytes); /* we are doing a new allocation -- release the current buffers */ clear(); try { allocateBytes(totalBytes, valueCount); } catch (Exception e) { clear(); throw e; } } @Override public void allocateNew(int valueCount) { allocateNew(lastValueAllocationSizeInBytes, valueCount); } /* Check if the data buffer size is within bounds. */ private void checkDataBufferSize(long size) { if (size > MAX_BUFFER_SIZE || size < 0) { throw new OversizedAllocationException( "Memory required for vector " + "is (" + size + "), which is overflow or more than max allowed (" + MAX_BUFFER_SIZE + "). " + "You could consider using LargeVarCharVector/LargeVarBinaryVector for large strings/large bytes types"); } } /* allocate the inner buffers */ private void allocateBytes(final long valueBufferSize, final int valueCount) { /* allocate data buffer */ viewBuffer = allocator.buffer(valueBufferSize); viewBuffer.readerIndex(0); validityBuffer = allocator.buffer((valueCount + 7) / 8); initValidityBuffer(); lastValueCapacity = getValueCapacity(); lastValueAllocationSizeInBytes = capAtMaxInt(viewBuffer.capacity()); } /** * Resize the vector to increase the capacity. The internal behavior is to double the current * value capacity. */ @Override public void reAlloc() { reallocViewBuffer(); reallocViewDataBuffer(); reallocValidityBuffer(); } /** * Reallocate the view buffer. View Buffer stores the views for VIEWVARCHAR or VIEWVARBINARY * elements in the vector. The behavior is to double the size of buffer. * * @throws OversizedAllocationException if the desired new size is more than max allowed * @throws OutOfMemoryException if the internal memory allocation fails */ public void reallocViewBuffer() { long currentViewBufferCapacity = viewBuffer.capacity(); long newAllocationSize = currentViewBufferCapacity * 2; if (newAllocationSize == 0) { if (lastValueAllocationSizeInBytes > 0) { newAllocationSize = lastValueAllocationSizeInBytes; } else { newAllocationSize = INITIAL_BYTE_COUNT * 2L; } } reallocViewBuffer(newAllocationSize); } /** Reallocate the data buffer associated with view buffer. */ public void reallocViewDataBuffer() { long currentDataBufferCapacity = 0; if (!dataBuffers.isEmpty()) { currentDataBufferCapacity = dataBuffers.get(dataBuffers.size() - 1).capacity(); } long newAllocationSize = currentDataBufferCapacity * 2; if (newAllocationSize == 0) { if (lastValueAllocationSizeInBytes > 0) { newAllocationSize = lastValueAllocationSizeInBytes; } else { newAllocationSize = INITIAL_BYTE_COUNT * 2L; } } reallocViewDataBuffer(newAllocationSize); } /** * Reallocate the view buffer to given size. View Buffer stores the views for VIEWVARCHAR or * VIEWVARBINARY elements in the vector. The actual allocated size may be larger than the request * one because it will round up the provided value to the nearest power of two. * * @param desiredAllocSize the desired new allocation size * @throws OversizedAllocationException if the desired new size is more than max allowed * @throws OutOfMemoryException if the internal memory allocation fails */ public void reallocViewBuffer(long desiredAllocSize) { if (desiredAllocSize == 0) { return; } long newAllocationSize = Math.min(CommonUtil.nextPowerOfTwo(desiredAllocSize), MAX_BUFFER_SIZE); assert newAllocationSize >= 1; // 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()); viewBuffer.getReferenceManager().release(); viewBuffer = newBuf; lastValueAllocationSizeInBytes = viewBuffer.capacity(); lastValueCapacity = getValueCapacity(); } /** * Reallocate the data buffer for views. * * @param desiredAllocSize allocation size in bytes */ public void reallocViewDataBuffer(long desiredAllocSize) { if (desiredAllocSize == 0) { return; } if (dataBuffers.isEmpty()) { return; } ArrowBuf currentBuf = dataBuffers.get(dataBuffers.size() - 1); if (currentBuf.capacity() - currentBuf.writerIndex() >= desiredAllocSize) { return; } final long newAllocationSize = Math.min(CommonUtil.nextPowerOfTwo(desiredAllocSize), MAX_BUFFER_SIZE); assert newAllocationSize >= 1; if (newAllocationSize < desiredAllocSize) { checkDataBufferSize(desiredAllocSize); } final ArrowBuf newBuf = allocator.buffer(newAllocationSize); dataBuffers.add(newBuf); } /** Reallocate Validity buffer. */ public void reallocValidityBuffer() { int targetValidityCount = capAtMaxInt((validityBuffer.capacity() * 8) * 2); if (targetValidityCount == 0) { if (lastValueCapacity > 0) { targetValidityCount = lastValueCapacity; } else { targetValidityCount = 2 * INITIAL_VALUE_ALLOCATION; } } long validityBufferSize = computeValidityBufferSize(targetValidityCount); final ArrowBuf newValidityBuffer = allocator.buffer(validityBufferSize); newValidityBuffer.setBytes(0, validityBuffer, 0, validityBuffer.capacity()); newValidityBuffer.setZero( validityBuffer.capacity(), newValidityBuffer.capacity() - validityBuffer.capacity()); validityBuffer.getReferenceManager().release(); validityBuffer = newValidityBuffer; lastValueCapacity = getValueCapacity(); } private long computeValidityBufferSize(int valueCount) { return (valueCount + 7) / 8; } /** * Get the size (number of bytes) of underlying view buffer. * * @return number of bytes in the view buffer */ @Override public int getByteCapacity() { return capAtMaxInt(viewBuffer.capacity()); } @Override public int sizeOfValueBuffer() { throw new UnsupportedOperationException( "sizeOfValueBuffer is not supported for BaseVariableWidthViewVector"); } /** * Get the size (number of bytes) of underlying elements in the view buffer. * * @return number of bytes used by data in the view buffer */ public int sizeOfViewBufferElements() { if (valueCount == 0) { return 0; } int totalSize = 0; for (int i = 0; i < valueCount; i++) { totalSize += getValueLength(i); } return totalSize; } /** * Get the size (number of bytes) of underlying buffers used by this vector. * * @return size of underlying buffers. */ @Override public int getBufferSize() { return getBufferSizeFor(this.valueCount); } /** * Get the potential buffer size for a particular number of records. * * @param valueCount desired number of elements in the vector * @return estimated size of underlying buffers if the vector holds a given number of elements */ @Override public int getBufferSizeFor(final int valueCount) { if (valueCount == 0) { return 0; } final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); final int viewBufferSize = valueCount * ELEMENT_SIZE; final int dataBufferSize = getDataBufferSize(); return validityBufferSize + viewBufferSize + dataBufferSize; } private int getDataBufferSize() { int dataBufferSize = 0; for (ArrowBuf buf : dataBuffers) { dataBufferSize += (int) buf.writerIndex(); } return dataBufferSize; } /** * Get information about how this field is materialized. * * @return the field corresponding to this vector */ @Override public Field getField() { return field; } /** * Return the underlying buffers associated with this vector. Note that this doesn't impact the * reference counts for this buffer, so it only should be used for in-context access. Also note * that this buffer changes regularly, thus external classes shouldn't hold a reference to it * (unless they change it). * * @param clear Whether to clear vector before returning, the buffers will still be refcounted but * the returned array will be the only reference to them * @return The underlying {@link ArrowBuf buffers} that is used by this vector instance. */ @Override public ArrowBuf[] getBuffers(boolean clear) { final ArrowBuf[] buffers; setReaderAndWriterIndex(); if (getBufferSize() == 0) { buffers = new ArrowBuf[0]; } else { final int dataBufferSize = dataBuffers.size(); // validity and view buffers final int fixedBufferSize = 2; buffers = new ArrowBuf[fixedBufferSize + dataBufferSize]; buffers[0] = validityBuffer; buffers[1] = viewBuffer; for (int i = fixedBufferSize; i < fixedBufferSize + dataBufferSize; i++) { buffers[i] = dataBuffers.get(i - fixedBufferSize); } } if (clear) { for (final ArrowBuf buffer : buffers) { buffer.getReferenceManager().retain(); } clear(); } return buffers; } /** Validate the scalar values held by this vector. */ public void validateScalars() { // No validation by default. } /** * Construct a transfer pair of this vector and another vector of the same type. * * @param field The field materialized by this vector. * @param allocator allocator for the target vector * @param callBack not used * @return TransferPair */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { return getTransferPair(field, allocator); } /** * Construct a transfer pair of this vector and another vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @param callBack not used * @return TransferPair */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return getTransferPair(ref, allocator); } /** * Construct a transfer pair of this vector and another vector of the same type. * * @param allocator allocator for the target vector * @return TransferPair */ @Override public TransferPair getTransferPair(BufferAllocator allocator) { return getTransferPair(getName(), allocator); } /** * Construct a transfer pair of this vector and another vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return TransferPair */ @Override public abstract TransferPair getTransferPair(String ref, BufferAllocator allocator); /** * Construct a transfer pair of this vector and another vector of the same type. * * @param field The field materialized by this vector. * @param allocator allocator for the target vector * @return TransferPair */ @Override public abstract TransferPair getTransferPair(Field field, BufferAllocator allocator); /** * Transfer this vector's data to another vector. The memory associated with this vector is * transferred to the allocator of target vector for accounting and management purposes. * * @param target destination vector for transfer */ public void transferTo(BaseVariableWidthViewVector target) { compareTypes(target, "transferTo"); target.clear(); target.validityBuffer = transferBuffer(validityBuffer, target.allocator); target.viewBuffer = transferBuffer(viewBuffer, target.allocator); target.dataBuffers = new ArrayList<>(dataBuffers.size()); for (int i = 0; i < dataBuffers.size(); i++) { target.dataBuffers.add(transferBuffer(dataBuffers.get(i), target.allocator)); } target.setLastSet(this.lastSet); if (this.valueCount > 0) { target.setValueCount(this.valueCount); } clear(); } /** * Slice this vector at desired index and length and transfer the corresponding data to the target * vector. * * @param startIndex start position of the split in source vector. * @param length length of the split. * @param target destination vector */ public void splitAndTransferTo(int startIndex, int length, BaseVariableWidthViewVector target) { Preconditions.checkArgument( startIndex >= 0 && length >= 0 && startIndex + length <= valueCount, "Invalid parameters startIndex: %s, length: %s for valueCount: %s", startIndex, length, valueCount); compareTypes(target, "splitAndTransferTo"); target.clear(); if (length > 0) { splitAndTransferValidityBuffer(startIndex, length, target); splitAndTransferViewBufferAndDataBuffer(startIndex, length, target); target.setLastSet(length - 1); target.setValueCount(length); } } /* allocate validity buffer */ @Override protected void allocateValidityBuffer(final long size) { super.allocateValidityBuffer(size); } @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(); } final ArrowBuf slicedValidityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); target.validityBuffer = transferBuffer(slicedValidityBuffer, target.allocator); } /** * In split and transfer, the view buffer and the data buffer will be allocated. Then the values * will be copied from the source vector to the target vector. Allocation and setting are * preferred over transfer since the buf index and buf offset needs to be overwritten when large * strings are added. * * @param startIndex starting index * @param length number of elements to be copied * @param target target vector */ private void splitAndTransferViewBufferAndDataBuffer( int startIndex, int length, BaseVariableWidthViewVector target) { if (length == 0) { return; } if (target.viewBuffer != null) { target.viewBuffer.getReferenceManager().release(); } // allocate target view buffer target.viewBuffer = target.allocator.buffer(length * ELEMENT_SIZE); for (int i = startIndex; i < startIndex + length; i++) { final int stringLength = getValueLength(i); // keeping track of writing index in the target view buffer int writePosition = (i - startIndex) * ELEMENT_SIZE; // keeping track of reading index in the source view buffer int readPosition = i * ELEMENT_SIZE; // set length target.viewBuffer.setInt(writePosition, stringLength); if (stringLength <= INLINE_SIZE) { // handle inline buffer writePosition += LENGTH_WIDTH; readPosition += LENGTH_WIDTH; // set data by copying the required portion from the source buffer target.viewBuffer.setBytes(writePosition, viewBuffer, readPosition, stringLength); } else { // handle non-inline buffer final int readBufIndex = viewBuffer.getInt(((long) i * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH); final int readBufOffset = viewBuffer.getInt( ((long) i * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH + BUF_INDEX_WIDTH); final ArrowBuf dataBuf = dataBuffers.get(readBufIndex); // allocate data buffer ArrowBuf currentDataBuf = target.allocateOrGetLastDataBuffer(stringLength); final long currentOffSet = currentDataBuf.writerIndex(); writePosition += LENGTH_WIDTH; readPosition += LENGTH_WIDTH; // set prefix target.viewBuffer.setBytes(writePosition, viewBuffer, readPosition, PREFIX_WIDTH); writePosition += PREFIX_WIDTH; // set buf id target.viewBuffer.setInt(writePosition, target.dataBuffers.size() - 1); writePosition += BUF_INDEX_WIDTH; // set offset target.viewBuffer.setInt(writePosition, (int) currentOffSet); currentDataBuf.setBytes(currentOffSet, dataBuf, readBufOffset, stringLength); currentDataBuf.writerIndex(currentOffSet + stringLength); } } } /*----------------------------------------------------------------* | | | common getters and setters | | | *----------------------------------------------------------------*/ /** * Get the number of elements that are null in the vector. * * @return the number of null elements. */ @Override public int getNullCount() { return BitVectorHelper.getNullCount(validityBuffer, valueCount); } /** * Check if the given index is within the current value capacity of the vector. * * @param index position to check * @return true if the index is within the current value capacity */ public boolean isSafe(int index) { return index < getValueCapacity(); } /** * Check if an element at given index is null. * * @param index position of an element * @return true if an element at given index is null */ @Override public boolean isNull(int index) { return (isSet(index) == 0); } /** * Same as {@link #isNull(int)}. * * @param index position of an element * @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; } /** * Get the value count of vector. This will always be zero unless setValueCount(int) has been * called prior to calling this. * * @return valueCount for the vector */ @Override public int getValueCount() { return valueCount; } /** * Sets the value count for the vector. * * @param valueCount value count */ @Override public void setValueCount(int valueCount) { assert valueCount >= 0; this.valueCount = valueCount; while (valueCount > getValueCapacity()) { reallocViewBuffer(); reallocValidityBuffer(); } lastSet = valueCount - 1; setReaderAndWriterIndex(); } /** * Create holes in the vector upto the given index (exclusive). Holes will be created from the * current last-set position in the vector. * * @param index target index */ @Override public void fillEmpties(int index) { handleSafe(index, EMPTY_BYTE_ARRAY.length); lastSet = index - 1; } /** * Set the index of the last non-null element in the vector. It is important to call this method * with appropriate value before calling {@link #setValueCount(int)}. * * @param value desired index of last non-null element. */ @Override public void setLastSet(int value) { lastSet = value; } /** * Get the index of the last non-null element in the vector. * * @return index of the last non-null element */ @Override public int getLastSet() { return lastSet; } /** * Mark the particular position in the vector as non-null. * * @param index position of the element. */ @Override public void setIndexDefined(int index) { // We need to check and reallocate the validity buffer while (index >= getValueCapacity()) { reallocValidityBuffer(); } BitVectorHelper.setBit(validityBuffer, index); } /** * Sets the value length for an element. * * @param index position of the element to set * @param length length of the element */ @Override public void setValueLengthSafe(int index, int length) { assert index >= 0; handleSafe(index, length); lastSet = index; } /** * Get the length of the element at specified index. * * @param index position of an element to get * @return greater than length 0 for a non-null element, 0 otherwise */ @Override public int getValueLength(int index) { assert index >= 0; if (index < 0 || index >= viewBuffer.capacity() / ELEMENT_SIZE) { throw new IndexOutOfBoundsException("Index out of bounds: " + index); } if (isSet(index) == 0) { return 0; } return viewBuffer.getInt(((long) index * ELEMENT_SIZE)); } /** * Set the variable length element at the specified index to the supplied byte array. This is same * as using {@link #set(int, byte[], int, int)} with start as Zero and length as #value.length * * @param index position of the element to set * @param value array of bytes to write */ public void set(int index, byte[] value) { assert index >= 0; BitVectorHelper.setBit(validityBuffer, index); setBytes(index, value, 0, value.length); lastSet = index; } /** * Same as {@link #set(int, byte[])} except that it handles the case where index and length of a * new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param value array of bytes to write */ @Override public void setSafe(int index, byte[] value) { assert index >= 0; // check if the current index can be populated handleSafe(index, value.length); BitVectorHelper.setBit(validityBuffer, index); setBytes(index, value, 0, value.length); lastSet = index; } /** * Set the variable length element at the specified index to the supplied byte array. * * @param index position of the element to set * @param value array of bytes to write * @param start start index in an array of bytes * @param length length of data in an array of bytes */ public void set(int index, byte[] value, int start, int length) { assert index >= 0; BitVectorHelper.setBit(validityBuffer, index); setBytes(index, value, start, length); lastSet = index; } /** * Same as {@link #set(int, byte[], int, int)} except that it handles the case where index and * length of a new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param value array of bytes to write * @param start start index in an array of bytes * @param length length of data in an array of bytes */ public void setSafe(int index, byte[] value, int start, int length) { assert index >= 0; handleSafe(index, length); BitVectorHelper.setBit(validityBuffer, index); setBytes(index, value, start, length); lastSet = index; } /** * Set the variable length element at the specified index to the content in supplied ByteBuffer. * * @param index position of the element to set * @param value ByteBuffer with data * @param start start index in ByteBuffer * @param length length of data in ByteBuffer */ public void set(int index, ByteBuffer value, int start, int length) { assert index >= 0; BitVectorHelper.setBit(validityBuffer, index); setBytes(index, value.array(), start, length); lastSet = index; } /** * Same as {@link #set(int, ByteBuffer, int, int)} except that it handles the case where index and * length of a new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param value ByteBuffer with data * @param start start index in ByteBuffer * @param length length of data in ByteBuffer */ public void setSafe(int index, ByteBuffer value, int start, int length) { assert index >= 0; handleSafe(index, length); BitVectorHelper.setBit(validityBuffer, index); setBytes(index, value.array(), start, length); lastSet = index; } /** * Set the element at the given index to null. * * @param index position of an element */ @Override public void setNull(int index) { handleSafe(index, 0); BitVectorHelper.unsetBit(validityBuffer, index); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet Zero for NULL value, 1 otherwise * @param start start position of data in buffer * @param end end position of data in buffer * @param buffer data buffer containing the variable width element to be stored in the vector */ public void set(int index, int isSet, int start, int end, ArrowBuf buffer) { assert index >= 0; final int dataLength = end - start; BitVectorHelper.setValidityBit(validityBuffer, index, isSet); setBytes(index, buffer, start, dataLength); lastSet = index; } /** * Same as {@link #set(int, int, int, int, ArrowBuf)} except that it handles the case when index * is greater than or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet Zero for NULL value, 1 otherwise * @param start start position of data in buffer * @param end end position of data in buffer * @param buffer data buffer containing the variable width element to be stored in the vector */ public void setSafe(int index, int isSet, int start, int end, ArrowBuf buffer) { assert index >= 0; final int dataLength = end - start; handleSafe(index, dataLength); BitVectorHelper.setValidityBit(validityBuffer, index, isSet); setBytes(index, buffer, start, dataLength); lastSet = index; } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param start start position of data in buffer * @param length length of data in buffer * @param buffer data buffer containing the variable width element to be stored in the vector */ public void set(int index, int start, int length, ArrowBuf buffer) { assert index >= 0; BitVectorHelper.setBit(validityBuffer, index); setBytes(index, buffer, start, length); lastSet = index; } /** * Same as {@link #set(int, int, int, int, ArrowBuf)} except that it handles the case when index * is greater than or equal to current value capacity of the vector. * * @param index position of the new value * @param start start position of data in buffer * @param length length of data in buffer * @param buffer data buffer containing the variable width element to be stored in the vector */ public void setSafe(int index, int start, int length, ArrowBuf buffer) { assert index >= 0; handleSafe(index, length); BitVectorHelper.setBit(validityBuffer, index); setBytes(index, buffer, start, length); lastSet = index; } /*----------------------------------------------------------------* | | | helper methods for setters | | | *----------------------------------------------------------------*/ protected ArrowBuf allocateOrGetLastDataBuffer(int length) { long dataBufferSize; if (initialDataBufferSize > 0) { dataBufferSize = Math.max(initialDataBufferSize, length); } else { dataBufferSize = Math.max(lastValueAllocationSizeInBytes, length); } if (dataBuffers.isEmpty() || dataBuffers.get(dataBuffers.size() - 1).capacity() - dataBuffers.get(dataBuffers.size() - 1).writerIndex() < length) { ArrowBuf newBuf = allocator.buffer(dataBufferSize); dataBuffers.add(newBuf); } return dataBuffers.get(dataBuffers.size() - 1); } /** * This method is used to create a view buffer for a variable width vector. It handles both inline * and data buffers. * *

    If the length of the value is less than or equal to {@link #INLINE_SIZE}, the value is * stored in the valueBuffer directly as an inline buffer. The valueBuffer stores the length of * the value followed by the value itself. If the length of the value is greater than {@link * #INLINE_SIZE}, a new buffer is allocated and added to dataBuffers to hold the value. The * viewBuffer in this case stores the length of the value, a prefix of the value, the index of the * new buffer in dataBuffers, and the offset of the value in the new buffer. * * @param index The index at which the new value will be inserted. * @param value The byte array that contains the data to be inserted. * @param start The start index in the byte array from where the data for the new value begins. * @param length The length of the data in the byte array that belongs to the new value. */ protected final void setBytes(int index, byte[] value, int start, int length) { int writePosition = index * 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); writePosition += LENGTH_WIDTH; // set data viewBuffer.setBytes(writePosition, value, start, length); } else { // allocate data buffer ArrowBuf currentBuf = allocateOrGetLastDataBuffer(length); // set length viewBuffer.setInt(writePosition, length); writePosition += LENGTH_WIDTH; // set prefix viewBuffer.setBytes(writePosition, value, start, PREFIX_WIDTH); writePosition += PREFIX_WIDTH; // set buf id viewBuffer.setInt(writePosition, dataBuffers.size() - 1); writePosition += BUF_INDEX_WIDTH; // set offset viewBuffer.setInt(writePosition, (int) currentBuf.writerIndex()); currentBuf.setBytes(currentBuf.writerIndex(), value, start, length); currentBuf.writerIndex(currentBuf.writerIndex() + length); } } /** * This method is used to create a view buffer for a variable width vector. Similar to {@link * #setBytes(int index, byte[] value, int start, int length)} * * @param index The index at which the new value will be inserted. * @param valueBuf The byte array that contains the data to be inserted. * @param start The start index in the byte array from where the data for the new value begins. * @param length The length of the data in the byte array that belongs to the new value. */ protected final void setBytes(int index, ArrowBuf valueBuf, int start, int length) { int writePosition = index * 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); writePosition += LENGTH_WIDTH; // set data viewBuffer.setBytes(writePosition, valueBuf, start, length); } else { // allocate data buffer ArrowBuf currentBuf = allocateOrGetLastDataBuffer(length); // set length viewBuffer.setInt(writePosition, length); writePosition += LENGTH_WIDTH; // set prefix viewBuffer.setBytes(writePosition, valueBuf, start, PREFIX_WIDTH); writePosition += PREFIX_WIDTH; // set buf id viewBuffer.setInt(writePosition, dataBuffers.size() - 1); writePosition += BUF_INDEX_WIDTH; // set offset viewBuffer.setInt(writePosition, (int) currentBuf.writerIndex()); currentBuf.setBytes(currentBuf.writerIndex(), valueBuf, start, length); currentBuf.writerIndex(currentBuf.writerIndex() + length); } } /** * Get the total length of the elements up to the given index. * * @param index The index of the element in the vector. * @return The total length up to the element at the given index. */ public final int getTotalValueLengthUpToIndex(int index) { int totalLength = 0; for (int i = 0; i < index - 1; i++) { totalLength += getValueLength(i); } return totalLength; } protected final void handleSafe(int index, int dataLength) { final long targetCapacity = roundUpToMultipleOf16((long) index * ELEMENT_SIZE + dataLength); if (viewBuffer.capacity() < targetCapacity) { reallocViewBuffer(targetCapacity); } while (index >= getValidityBufferValueCapacity()) { reallocValidityBuffer(); } } /** * Copy a cell value from a particular index in source vector to a particular position in this * vector. * * @param fromIndex position to copy from in source vector * @param thisIndex position to copy to in this vector * @param from source vector */ @Override public void copyFrom(int fromIndex, int thisIndex, ValueVector from) { Preconditions.checkArgument(getMinorType() == from.getMinorType()); if (from.isNull(fromIndex)) { BitVectorHelper.unsetBit(validityBuffer, thisIndex); } else { final int viewLength = from.getDataBuffer().getInt((long) fromIndex * ELEMENT_SIZE); copyFromNotNull(fromIndex, thisIndex, from, viewLength); } lastSet = thisIndex; } /** * Same as {@link #copyFrom(int, int, ValueVector)} except that it handles the case when the * capacity of the vector needs to be expanded before copy. * * @param fromIndex position to copy from in source vector * @param thisIndex position to copy to in this vector * @param from source vector */ @Override public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) { Preconditions.checkArgument(getMinorType() == from.getMinorType()); if (from.isNull(fromIndex)) { handleSafe(thisIndex, 0); BitVectorHelper.unsetBit(validityBuffer, thisIndex); } else { final int viewLength = from.getDataBuffer().getInt((long) fromIndex * ELEMENT_SIZE); handleSafe(thisIndex, viewLength); copyFromNotNull(fromIndex, thisIndex, from, viewLength); } lastSet = thisIndex; } private void copyFromNotNull(int fromIndex, int thisIndex, ValueVector from, int viewLength) { BitVectorHelper.setBit(validityBuffer, thisIndex); final int start = thisIndex * ELEMENT_SIZE; final int copyStart = fromIndex * ELEMENT_SIZE; if (viewLength > INLINE_SIZE) { final int bufIndex = from.getDataBuffer() .getInt(((long) fromIndex * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH); final int dataOffset = from.getDataBuffer() .getInt( ((long) fromIndex * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH + BUF_INDEX_WIDTH); final ArrowBuf dataBuf = ((BaseVariableWidthViewVector) from).dataBuffers.get(bufIndex); final ArrowBuf thisDataBuf = allocateOrGetLastDataBuffer(viewLength); viewBuffer.setBytes(start, from.getDataBuffer(), copyStart, LENGTH_WIDTH + PREFIX_WIDTH); int writePosition = start + LENGTH_WIDTH + PREFIX_WIDTH; // set buf id viewBuffer.setInt(writePosition, dataBuffers.size() - 1); writePosition += BUF_INDEX_WIDTH; // set offset viewBuffer.setInt(writePosition, (int) thisDataBuf.writerIndex()); thisDataBuf.setBytes(thisDataBuf.writerIndex(), dataBuf, dataOffset, viewLength); thisDataBuf.writerIndex(thisDataBuf.writerIndex() + viewLength); } else { from.getDataBuffer().getBytes(copyStart, viewBuffer, start, ELEMENT_SIZE); } } @Override public ArrowBufPointer getDataPointer(int index) { return getDataPointer(index, new ArrowBufPointer()); } @Override public ArrowBufPointer getDataPointer(int index, ArrowBufPointer reuse) { if (isNull(index)) { reuse.set(null, 0, 0); } else { int length = getValueLength(index); if (length < INLINE_SIZE) { int start = index * ELEMENT_SIZE + LENGTH_WIDTH; reuse.set(viewBuffer, start, length); } else { final int bufIndex = viewBuffer.getInt(((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH); ArrowBuf dataBuf = dataBuffers.get(bufIndex); reuse.set(dataBuf, 0, length); } } return reuse; } @Override public int hashCode(int index) { return hashCode(index, null); } @Override public int hashCode(int index, ArrowBufHasher hasher) { if (isNull(index)) { return ArrowBufPointer.NULL_HASH_CODE; } final int length = getValueLength(index); if (length < INLINE_SIZE) { int start = index * ELEMENT_SIZE + LENGTH_WIDTH; return ByteFunctionHelpers.hash(hasher, this.getDataBuffer(), start, start + length); } else { final int bufIndex = viewBuffer.getInt(((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH); final int dataOffset = viewBuffer.getInt( ((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH + BUF_INDEX_WIDTH); ArrowBuf dataBuf = dataBuffers.get(bufIndex); return ByteFunctionHelpers.hash(hasher, dataBuf, dataOffset, dataOffset + length); } } /** * Retrieves the data of a variable-width element at a given index in the vector. * *

    If the length of the data is greater than {@link #INLINE_SIZE}, the data is stored in an * inline buffer. The method retrieves the buffer index and data offset from the viewBuffer, and * then retrieves the data from the corresponding buffer in the dataBuffers list. * *

    If the length of the data is less than or equal to {@link #INLINE_SIZE}, the data is stored * directly in the viewBuffer. The method retrieves the data directly from the viewBuffer. * * @param index position of the element in the vector * @return byte array containing the data of the element */ protected byte[] getData(int index) { final int dataLength = getValueLength(index); byte[] result = new byte[dataLength]; if (dataLength > INLINE_SIZE) { // data is in the data buffer // get buffer index final int bufferIndex = viewBuffer.getInt(((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH); // get data offset final int dataOffset = viewBuffer.getInt( ((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH + BUF_INDEX_WIDTH); dataBuffers.get(bufferIndex).getBytes(dataOffset, result, 0, dataLength); } else { // data is in the view buffer viewBuffer.getBytes((long) index * ELEMENT_SIZE + BUF_INDEX_WIDTH, result, 0, dataLength); } return result; } protected void getData(int index, ReusableBuffer buffer) { final int dataLength = getValueLength(index); if (dataLength > INLINE_SIZE) { // data is in the data buffer // get buffer index final int bufferIndex = viewBuffer.getInt(((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH); // get data offset final int dataOffset = viewBuffer.getInt( ((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH + BUF_INDEX_WIDTH); ArrowBuf dataBuf = dataBuffers.get(bufferIndex); buffer.set(dataBuf, dataOffset, dataLength); } else { // data is in the value buffer buffer.set(viewBuffer, ((long) index * ELEMENT_SIZE) + BUF_INDEX_WIDTH, dataLength); } } @Override public OUT accept(VectorVisitor visitor, IN value) { return visitor.visit(this, value); } /** * Retrieves the export buffer count for the C Data Interface. * *

    For Variadic types, an additional buffer is kept to store the size of each variadic buffer * since that information cannot be retrieved in the C Data import. * *

    In the C Data Interface, the binary view import expects at least three buffers. The variadic * size buffer is merely allocated to determine the number of elements per each variadic buffer, * and it is not part of the imported data. * *

    The count is set to 3 + dataBuffers.size(). Three is formed by validity, view, and variadic * size buffer. * * @return the number of buffers to be exported */ @Override public int getExportedCDataBufferCount() { return 3 + dataBuffers.size(); } /** * Get the data buffer of the vector. Note that an additional buffer is appended to store the size * of each variadic buffer's size. * * @param buffers list of buffers to be exported * @param buffersPtr buffer to store the pointers to the exported buffers * @param nullValue null value */ @Override public void exportCDataBuffers(List buffers, ArrowBuf buffersPtr, long nullValue) { exportBuffer(validityBuffer, buffers, buffersPtr, nullValue, true); exportBuffer(viewBuffer, buffers, buffersPtr, nullValue, true); // allocating additional space to keep the number of variadic buffers ArrowBuf variadicSizeBuffer = allocator.buffer((long) Long.BYTES * dataBuffers.size()); // variadicSizeBuffer.setZero(0, variadicSizeBuffer.capacity()); // export data buffers for (int i = 0; i < dataBuffers.size(); i++) { ArrowBuf dataBuf = dataBuffers.get(i); // calculate sizes for variadic size buffer variadicSizeBuffer.setLong((long) i * Long.BYTES, dataBuf.capacity()); exportBuffer(dataBuf, buffers, buffersPtr, nullValue, true); } // export variadic size buffer exportBuffer(variadicSizeBuffer, buffers, buffersPtr, nullValue, false); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/BigIntVector.java000066400000000000000000000262571515451615200307160ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.BigIntReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.BigIntHolder; import org.apache.arrow.vector.holders.NullableBigIntHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * BigIntVector implements a fixed width vector (8 bytes) of integer values which could be null. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public final class BigIntVector extends BaseFixedWidthVector implements BaseIntVector, ValueIterableVector { public static final byte TYPE_WIDTH = 8; /** * Instantiate a BigIntVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public BigIntVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.BIGINT.getType()), allocator); } /** * Instantiate a BigIntVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public BigIntVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a BigIntVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public BigIntVector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new BigIntReaderImpl(BigIntVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.BIGINT; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public long get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getLong((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableBigIntHolder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getLong((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Long getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getLong((long) index * TYPE_WIDTH); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, long value) { valueBuffer.setLong((long) index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, long value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableBigIntHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, BigIntHolder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, long)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, long value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableBigIntHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableBigIntHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, BigIntHolder)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, BigIntHolder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void set(int index, int isSet, long value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, long)} except that it handles the case when index is greater than * or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setSafe(int index, int isSet, long value) { handleSafe(index); set(index, isSet, value); } /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static long get(final ArrowBuf buffer, final int index) { return buffer.getLong((long) index * TYPE_WIDTH); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising of this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((BigIntVector) to); } @Override public void setWithPossibleTruncate(int index, long value) { this.setSafe(index, value); } @Override public void setUnsafeWithPossibleTruncate(int index, long value) { this.set(index, value); } @Override public long getValueAsLong(int index) { return this.get(index); } private class TransferImpl implements TransferPair { BigIntVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new BigIntVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new BigIntVector(field, allocator); } public TransferImpl(BigIntVector to) { this.to = to; } @Override public BigIntVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, BigIntVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/BitVector.java000066400000000000000000000472371515451615200302610ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.memory.util.LargeMemoryUtil.capAtMaxInt; import static org.apache.arrow.vector.NullCheckingForGet.NULL_CHECKING_ENABLED; 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.hash.ArrowBufHasher; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.complex.impl.BitReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.BitHolder; import org.apache.arrow.vector.holders.NullableBitHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.OversizedAllocationException; import org.apache.arrow.vector.util.TransferPair; /** * BitVector implements a fixed width (1 bit) vector of boolean values which could be null. Each * value in the vector corresponds to a single bit in the underlying data stream backing the vector. */ public final class BitVector extends BaseFixedWidthVector implements ValueIterableVector { private static final int HASH_CODE_FOR_ZERO = 17; private static final int HASH_CODE_FOR_ONE = 19; /** * Instantiate a BitVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public BitVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.BIT.getType()), allocator); } /** * Instantiate a BitVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public BitVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a BitVector. This doesn't allocate any memory for the data in vector. * * @param field the Field materialized by this vector * @param allocator allocator for memory management. */ public BitVector(Field field, BufferAllocator allocator) { super(field, allocator, 0); } @Override protected FieldReader getReaderImpl() { return new BitReaderImpl(BitVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.BIT; } /** * Sets the desired value capacity for the vector. This function doesn't allocate any memory for * the vector. * * @param valueCount desired number of elements in the vector */ @Override public void setInitialCapacity(int 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"); } lastValueCapacity = valueCount; } @Override protected int getValueBufferValueCapacity() { return capAtMaxInt(valueBuffer.capacity() * 8); } /** * Get the potential buffer size for a particular number of records. * * @param count desired number of elements in the vector * @return estimated size of underlying buffers if the vector holds a given number of elements */ @Override public int getBufferSizeFor(final int count) { if (count == 0) { return 0; } return 2 * BitVectorHelper.getValidityBufferSizeFromCount(count); } /** * Get the size (number of bytes) of underlying buffers used by this vector. * * @return size of underlying buffers. */ @Override public int getBufferSize() { return getBufferSizeFor(valueCount); } /** * Slice this vector at desired index and length and transfer the corresponding data to the target * vector. * * @param startIndex start position of the split in source vector. * @param length length of the split. * @param target destination vector */ @Override public void splitAndTransferTo(int startIndex, int length, BaseFixedWidthVector target) { Preconditions.checkArgument( startIndex >= 0 && length >= 0 && startIndex + length <= valueCount, "Invalid parameters startIndex: %s, length: %s for valueCount: %s", startIndex, length, valueCount); compareTypes(target, "splitAndTransferTo"); target.clear(); target.validityBuffer = splitAndTransferBuffer(startIndex, length, validityBuffer, target.validityBuffer); target.valueBuffer = splitAndTransferBuffer(startIndex, length, valueBuffer, target.valueBuffer); target.refreshValueCapacity(); target.setValueCount(length); } private ArrowBuf splitAndTransferBuffer( int startIndex, int length, ArrowBuf sourceBuffer, ArrowBuf destBuffer) { int firstByteSource = BitVectorHelper.byteIndex(startIndex); int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1); int byteSizeTarget = BitVectorHelper.getValidityBufferSizeFromCount(length); int offset = startIndex % 8; if (length > 0) { if (offset == 0) { /* slice */ if (destBuffer != null) { destBuffer.getReferenceManager().release(); } destBuffer = sourceBuffer.slice(firstByteSource, byteSizeTarget); destBuffer.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. */ destBuffer = allocator.buffer(byteSizeTarget); destBuffer.readerIndex(0); destBuffer.setZero(0, destBuffer.capacity()); for (int i = 0; i < byteSizeTarget - 1; i++) { byte b1 = BitVectorHelper.getBitsFromCurrentByte(sourceBuffer, firstByteSource + i, offset); byte b2 = BitVectorHelper.getBitsFromNextByte(sourceBuffer, firstByteSource + i + 1, offset); destBuffer.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( sourceBuffer, firstByteSource + byteSizeTarget - 1, offset); byte b2 = BitVectorHelper.getBitsFromNextByte( sourceBuffer, firstByteSource + byteSizeTarget, offset); destBuffer.setByte(byteSizeTarget - 1, b1 + b2); } else { byte b1 = BitVectorHelper.getBitsFromCurrentByte( sourceBuffer, firstByteSource + byteSizeTarget - 1, offset); destBuffer.setByte(byteSizeTarget - 1, b1); } } } return destBuffer; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ private int getBit(int index) { final int byteIndex = index >> 3; final byte b = valueBuffer.getByte(byteIndex); final int bitIndex = index & 7; return (b >> bitIndex) & 0x01; } /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public int get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return getBit(index); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableBitHolder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = getBit(index); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Boolean getObject(int index) { if (isSet(index) == 0) { return null; } else { return getBit(index) != 0; } } /** * Copy a cell value from a particular index in source vector to a particular position in this * vector. * * @param fromIndex position to copy from in source vector * @param thisIndex position to copy to in this vector * @param from source vector */ @Override public void copyFrom(int fromIndex, int thisIndex, ValueVector from) { Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); boolean fromIsSet = BitVectorHelper.get(from.getValidityBuffer(), fromIndex) != 0; if (fromIsSet) { BitVectorHelper.setBit(validityBuffer, thisIndex); BitVectorHelper.setValidityBit(valueBuffer, thisIndex, ((BitVector) from).getBit(fromIndex)); } else { BitVectorHelper.unsetBit(validityBuffer, thisIndex); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, int value) { BitVectorHelper.setBit(validityBuffer, index); if (value != 0) { BitVectorHelper.setBit(valueBuffer, index); } else { BitVectorHelper.unsetBit(valueBuffer, index); } } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableBitHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); if (holder.value != 0) { BitVectorHelper.setBit(valueBuffer, index); } else { BitVectorHelper.unsetBit(valueBuffer, index); } } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, BitHolder holder) { BitVectorHelper.setBit(validityBuffer, index); if (holder.value != 0) { BitVectorHelper.setBit(valueBuffer, index); } else { BitVectorHelper.unsetBit(valueBuffer, index); } } /** * Same as {@link #set(int, int)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, int value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableBitHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableBitHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, BitHolder)} except that it handles the case when index is greater than * or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, BitHolder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void set(int index, int isSet, int value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, int)} except that it handles the case when index is greater than * or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setSafe(int index, int isSet, int value) { handleSafe(index); set(index, isSet, value); } /** * Set the element at the given index to one. * * @param index position of element */ public void setToOne(int index) { BitVectorHelper.setBit(validityBuffer, index); BitVectorHelper.setBit(valueBuffer, index); } /** * Same as {@link #setToOne(int)} except that it handles the case when index is greater than or * equal to current value capacity of the vector. * * @param index position of the element */ public void setSafeToOne(int index) { handleSafe(index); setToOne(index); } @Override public ArrowBufPointer getDataPointer(int index) { throw new UnsupportedOperationException(); } @Override public ArrowBufPointer getDataPointer(int index, ArrowBufPointer reuse) { throw new UnsupportedOperationException(); } @Override public int hashCode(int index) { if (isNull(index)) { return ArrowBufPointer.NULL_HASH_CODE; } else { if (get(index) == 0) { return HASH_CODE_FOR_ZERO; } else { return HASH_CODE_FOR_ONE; } } } @Override public int hashCode(int index, ArrowBufHasher hasher) { return hashCode(index); } /** * Set count bits to 1 in data starting at firstBitIndex. * * @param firstBitIndex the index of the first bit to set * @param count the number of bits to set */ public void setRangeToOne(int firstBitIndex, int count) { int startByteIndex = BitVectorHelper.byteIndex(firstBitIndex); final int lastBitIndex = firstBitIndex + count; final int endByteIndex = BitVectorHelper.byteIndex(lastBitIndex); final int startByteBitIndex = BitVectorHelper.bitIndex(firstBitIndex); final int endBytebitIndex = BitVectorHelper.bitIndex(lastBitIndex); if (count < 8 && startByteIndex == endByteIndex) { // handles the case where we don't have a first and a last byte byte bitMask = 0; for (int i = startByteBitIndex; i < endBytebitIndex; ++i) { bitMask |= (byte) (1L << i); } BitVectorHelper.setBitMaskedByte(validityBuffer, startByteIndex, bitMask); BitVectorHelper.setBitMaskedByte(valueBuffer, startByteIndex, bitMask); } else { // fill in first byte (if it's not full) if (startByteBitIndex != 0) { final byte bitMask = (byte) (0xFFL << startByteBitIndex); BitVectorHelper.setBitMaskedByte(validityBuffer, startByteIndex, bitMask); BitVectorHelper.setBitMaskedByte(valueBuffer, startByteIndex, bitMask); ++startByteIndex; } // fill in one full byte at a time validityBuffer.setOne(startByteIndex, endByteIndex - startByteIndex); valueBuffer.setOne(startByteIndex, endByteIndex - startByteIndex); // fill in the last byte (if it's not full) if (endBytebitIndex != 0) { final int byteIndex = BitVectorHelper.byteIndex(lastBitIndex - endBytebitIndex); final byte bitMask = (byte) (0xFFL >>> ((8 - endBytebitIndex) & 7)); BitVectorHelper.setBitMaskedByte(validityBuffer, byteIndex, bitMask); BitVectorHelper.setBitMaskedByte(valueBuffer, byteIndex, bitMask); } } } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((BitVector) to); } private class TransferImpl implements TransferPair { BitVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new BitVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new BitVector(field, allocator); } public TransferImpl(BitVector to) { this.to = to; } @Override public BitVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, BitVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java000066400000000000000000000401151515451615200314050ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.memory.util.LargeMemoryUtil.checkedCastToInt; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BoundsChecking; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.MemoryUtil; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.util.DataSizeRoundingUtil; /** * Helper class for performing generic operations on a bit vector buffer. External use of this class * is not recommended. */ public class BitVectorHelper { private BitVectorHelper() {} /** Get the index of byte corresponding to bit index in validity buffer. */ public static long byteIndex(long absoluteBitIndex) { return absoluteBitIndex >> 3; } /** Get the relative index of bit within the byte in validity buffer. */ public static int bitIndex(long absoluteBitIndex) { return checkedCastToInt(absoluteBitIndex & 7); } /** Get the index of byte corresponding to bit index in validity buffer. */ public static int byteIndex(int absoluteBitIndex) { return absoluteBitIndex >> 3; } /** Get the relative index of bit within the byte in validity buffer. */ public static int bitIndex(int absoluteBitIndex) { return absoluteBitIndex & 7; } /** * Set the bit at provided index to 1. * * @param validityBuffer validity buffer of the vector * @param index index to be set */ public static void setBit(ArrowBuf validityBuffer, long index) { // it can be observed that some logic is duplicate of the logic in setValidityBit. // this is because JIT cannot always remove the if branch in setValidityBit, // so we give a dedicated implementation for setting bits. final long byteIndex = byteIndex(index); final int bitIndex = bitIndex(index); // the byte is promoted to an int, because according to Java specification, // bytes will be promoted to ints automatically, upon expression evaluation. // by promoting it manually, we avoid the unnecessary conversions. int currentByte = validityBuffer.getByte(byteIndex); final int bitMask = 1 << bitIndex; currentByte |= bitMask; validityBuffer.setByte(byteIndex, currentByte); } /** * Set the bit at provided index to 0. * * @param validityBuffer validity buffer of the vector * @param index index to be set */ public static void unsetBit(ArrowBuf validityBuffer, int index) { // it can be observed that some logic is duplicate of the logic in setValidityBit. // this is because JIT cannot always remove the if branch in setValidityBit, // so we give a dedicated implementation for unsetting bits. final int byteIndex = byteIndex(index); final int bitIndex = bitIndex(index); // the byte is promoted to an int, because according to Java specification, // bytes will be promoted to ints automatically, upon expression evaluation. // by promoting it manually, we avoid the unnecessary conversions. int currentByte = validityBuffer.getByte(byteIndex); final int bitMask = 1 << bitIndex; currentByte &= ~bitMask; validityBuffer.setByte(byteIndex, currentByte); } /** * Set the bit at a given index to provided value (1 or 0). * * @param validityBuffer validity buffer of the vector * @param index index to be set * @param value value to set */ public static void setValidityBit(ArrowBuf validityBuffer, int index, int value) { final int byteIndex = byteIndex(index); final int bitIndex = bitIndex(index); // the byte is promoted to an int, because according to Java specification, // bytes will be promoted to ints automatically, upon expression evaluation. // by promoting it manually, we avoid the unnecessary conversions. int currentByte = validityBuffer.getByte(byteIndex); final int bitMask = 1 << bitIndex; if (value != 0) { currentByte |= bitMask; } else { currentByte &= ~bitMask; } validityBuffer.setByte(byteIndex, currentByte); } /** * Set the bit at a given index to provided value (1 or 0). Internally takes care of allocating * the buffer if the caller didn't do so. * * @param validityBuffer validity buffer of the vector * @param allocator allocator for the buffer * @param valueCount number of values to allocate/set * @param index index to be set * @param value value to set * @return ArrowBuf */ public static ArrowBuf setValidityBit( ArrowBuf validityBuffer, BufferAllocator allocator, int valueCount, int index, int value) { if (validityBuffer == null) { validityBuffer = allocator.buffer(getValidityBufferSizeFromCount(valueCount)); } setValidityBit(validityBuffer, index, value); if (index == (valueCount - 1)) { validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount)); } return validityBuffer; } /** * Check if a bit at a given index is set or not. * * @param buffer buffer to check * @param index index of the buffer * @return 1 if bit is set, 0 otherwise. */ public static int get(final ArrowBuf buffer, int index) { final int byteIndex = index >> 3; final byte b = buffer.getByte(byteIndex); final int bitIndex = index & 7; return (b >> bitIndex) & 0x01; } /** * 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 */ public static int getValidityBufferSizeFromCount(int valueCount) { return DataSizeRoundingUtil.divideBy8Ceil(valueCount); } /** * Given a validity buffer, find the number of bits that are not set. This is used to compute the * number of null elements in a nullable vector. * * @param validityBuffer validity buffer of the vector * @param valueCount number of values in the vector * @return number of bits not set. */ public static int getNullCount(final ArrowBuf validityBuffer, final int valueCount) { if (valueCount == 0) { return 0; } int count = 0; 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; int index = 0; while (index + 8 <= fullBytesCount) { long longValue = validityBuffer.getLong(index); count += Long.bitCount(longValue); index += 8; } if (index + 4 <= fullBytesCount) { int intValue = validityBuffer.getInt(index); count += Integer.bitCount(intValue); index += 4; } while (index < fullBytesCount) { byte byteValue = validityBuffer.getByte(index); count += Integer.bitCount(byteValue & 0xFF); index += 1; } // handling with the last bits if (remainder != 0) { byte byteValue = validityBuffer.getByte(sizeInBytes - 1); // making the remaining bits all 1s if it is not fully filled byte mask = (byte) (0xFF << remainder); byteValue = (byte) (byteValue | mask); count += Integer.bitCount(byteValue & 0xFF); } return 8 * sizeInBytes - count; } /** * Tests if all bits in a validity buffer are equal 0 or 1, according to the specified parameter. * * @param validityBuffer the validity buffer. * @param valueCount the bit count. * @param checkOneBits if set to true, the method checks if all bits are equal to 1; otherwise, it * checks if all bits are equal to 0. * @return true if all bits are 0 or 1 according to the parameter, and false otherwise. */ public static boolean checkAllBitsEqualTo( final ArrowBuf validityBuffer, final int valueCount, final boolean checkOneBits) { if (valueCount == 0) { return true; } final int sizeInBytes = getValidityBufferSizeFromCount(valueCount); // boundary check validityBuffer.checkBytes(0, sizeInBytes); // 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; // the integer number to compare against final int intToCompare = checkOneBits ? -1 : 0; int index = 0; while (index + 8 <= fullBytesCount) { long longValue = MemoryUtil.getLong(validityBuffer.memoryAddress() + index); if (longValue != (long) intToCompare) { return false; } index += 8; } if (index + 4 <= fullBytesCount) { int intValue = MemoryUtil.getInt(validityBuffer.memoryAddress() + index); if (intValue != intToCompare) { return false; } index += 4; } while (index < fullBytesCount) { byte byteValue = MemoryUtil.getByte(validityBuffer.memoryAddress() + index); if (byteValue != (byte) intToCompare) { return false; } index += 1; } // handling with the last bits if (remainder != 0) { byte byteValue = MemoryUtil.getByte(validityBuffer.memoryAddress() + sizeInBytes - 1); byte mask = (byte) ((1 << remainder) - 1); byteValue = (byte) (byteValue & mask); if (checkOneBits) { if ((mask & byteValue) != mask) { return false; } } else { if (byteValue != (byte) 0) { return false; } } } return true; } /** Returns the byte at index from data right-shifted by offset. */ public static byte getBitsFromCurrentByte( final ArrowBuf data, final int index, final int offset) { return (byte) ((data.getByte(index) & 0xFF) >>> offset); } /** Returns the byte at index from left-shifted by (8 - offset). */ public static byte getBitsFromNextByte(ArrowBuf data, int index, int offset) { return (byte) ((data.getByte(index) << (8 - offset))); } /** * Returns a new buffer if the source validity buffer is either all null or all not-null, * otherwise returns a buffer pointing to the same memory as source. * * @param fieldNode The fieldNode containing the null count * @param sourceValidityBuffer The source validity buffer that will have its position copied if * there is a mix of null and non-null values * @param allocator The allocator to use for creating a new buffer if necessary. * @return A new buffer that is either allocated or points to the same memory as * sourceValidityBuffer. */ public static ArrowBuf loadValidityBuffer( final ArrowFieldNode fieldNode, final ArrowBuf sourceValidityBuffer, final BufferAllocator allocator) { final int valueCount = fieldNode.getLength(); ArrowBuf newBuffer = null; // Create a new validity buffer iff both of the following are true: // - validity buffer is not present, that is, it is either null or empty (in the case of // IPC for instance). // - values are either all NULLs or all non-NULLs boolean isValidityBufferNull = sourceValidityBuffer == null || sourceValidityBuffer.capacity() == 0; if (isValidityBufferNull && (fieldNode.getNullCount() == 0 || fieldNode.getNullCount() == valueCount)) { newBuffer = allocator.buffer(getValidityBufferSizeFromCount(valueCount)); newBuffer.setZero(0, newBuffer.capacity()); if (fieldNode.getNullCount() != 0) { /* all NULLs */ return newBuffer; } /* all non-NULLs */ int fullBytesCount = valueCount / 8; newBuffer.setOne(0, fullBytesCount); int remainder = valueCount % 8; if (remainder > 0) { byte bitMask = (byte) (0xFFL >>> ((8 - remainder) & 7)); newBuffer.setByte(fullBytesCount, bitMask); } } else { /* mixed byte pattern -- create another ArrowBuf associated with the * target allocator */ newBuffer = sourceValidityBuffer.getReferenceManager().retain(sourceValidityBuffer, allocator); } return newBuffer; } /** * Set the byte of the given index in the data buffer by applying a bit mask to the current byte * at that index. * * @param data buffer to set * @param byteIndex byteIndex within the buffer * @param bitMask bit mask to be set */ static void setBitMaskedByte(ArrowBuf data, int byteIndex, byte bitMask) { byte currentByte = data.getByte(byteIndex); currentByte |= bitMask; data.setByte(byteIndex, currentByte); } /** * Concat two validity buffers. * * @param input1 the first validity buffer. * @param numBits1 the number of bits in the first validity buffer. * @param input2 the second validity buffer. * @param numBits2 the number of bits in the second validity buffer. * @param output the output validity buffer. It can be the same one as the first input. The caller * must make sure the output buffer has enough capacity. */ public static void concatBits( ArrowBuf input1, int numBits1, ArrowBuf input2, int numBits2, ArrowBuf output) { int numBytes1 = DataSizeRoundingUtil.divideBy8Ceil(numBits1); int numBytes2 = DataSizeRoundingUtil.divideBy8Ceil(numBits2); int numBytesOut = DataSizeRoundingUtil.divideBy8Ceil(numBits1 + numBits2); if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { output.checkBytes(0, numBytesOut); } // copy the first bit set if (input1 != output) { MemoryUtil.copyMemory(input1.memoryAddress(), output.memoryAddress(), numBytes1); } if (bitIndex(numBits1) == 0) { // The number of bits for the first bit set is a multiple of 8, so the boundary is at byte // boundary. // For this case, we have a shortcut to copy all bytes from the second set after the byte // boundary. MemoryUtil.copyMemory(input2.memoryAddress(), output.memoryAddress() + numBytes1, numBytes2); return; } // the number of bits to fill a full byte after the first input is processed int numBitsToFill = 8 - bitIndex(numBits1); // mask to clear high bits int mask = (1 << (8 - numBitsToFill)) - 1; int numFullBytes = numBits2 / 8; int prevByte = output.getByte(numBytes1 - 1) & mask; for (int i = 0; i < numFullBytes; i++) { int curByte = input2.getByte(i) & 0xff; // first fill the bits to a full byte int byteToFill = (curByte << (8 - numBitsToFill)) & 0xff; output.setByte(numBytes1 + i - 1, byteToFill | prevByte); // fill remaining bits in the current byte // note that it is also the previous byte for the next iteration prevByte = curByte >>> numBitsToFill; } int lastOutputByte = prevByte; // the number of extra bits for the second input, relative to full bytes int numTrailingBits = bitIndex(numBits2); if (numTrailingBits == 0) { output.setByte(numBytes1 + numFullBytes - 1, lastOutputByte); return; } // process remaining bits from input2 int remByte = input2.getByte(numBytes2 - 1) & 0xff; int byteToFill = remByte << (8 - numBitsToFill); lastOutputByte |= byteToFill; output.setByte(numBytes1 + numFullBytes - 1, lastOutputByte); if (numTrailingBits > numBitsToFill) { // clear all bits for the last byte before writing output.setByte(numBytes1 + numFullBytes, 0); // some remaining bits cannot be filled in the previous byte int leftByte = remByte >>> numBitsToFill; output.setByte(numBytes1 + numFullBytes, leftByte); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/BufferBacked.java000066400000000000000000000021301515451615200306420ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; /** Content is backed by a buffer and can be loaded/unloaded. */ public interface BufferBacked { void load(ArrowFieldNode fieldNode, ArrowBuf data); ArrowBuf unLoad(); } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/BufferLayout.java000066400000000000000000000122671515451615200307620ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.util.Preconditions; /** * Metadata class that captures the "type" of an Arrow buffer. (e.g. data buffers, offset buffers * for variable width types and validity buffers). */ public class BufferLayout { /** * Enumeration of the different logical types a buffer can have. Data buffer is common to most of * the layouts. Offset buffer is used for variable width types. Validity buffer is used for * nullable types. Type buffer is used for Union types. Size buffer is used for ListView and * LargeListView types. */ public enum BufferType { DATA("DATA"), OFFSET("OFFSET"), VALIDITY("VALIDITY"), TYPE("TYPE_ID"), SIZE("SIZE"), VIEWS("VIEWS"), VARIADIC_DATA_BUFFERS("VARIADIC_DATA_BUFFERS"); private final String name; BufferType(String name) { this.name = name; } public String getName() { return name; } } private static final BufferLayout VALIDITY_BUFFER = new BufferLayout(BufferType.VALIDITY, 1); private static final BufferLayout OFFSET_BUFFER = new BufferLayout(BufferType.OFFSET, 32); private static final BufferLayout LARGE_OFFSET_BUFFER = new BufferLayout(BufferType.OFFSET, 64); private static final BufferLayout TYPE_BUFFER = new BufferLayout(BufferType.TYPE, 32); private static final BufferLayout BIT_BUFFER = new BufferLayout(BufferType.DATA, 1); private static final BufferLayout VALUES_256 = new BufferLayout(BufferType.DATA, 256); private static final BufferLayout VALUES_128 = new BufferLayout(BufferType.DATA, 128); private static final BufferLayout VALUES_64 = new BufferLayout(BufferType.DATA, 64); private static final BufferLayout VALUES_32 = new BufferLayout(BufferType.DATA, 32); private static final BufferLayout VALUES_16 = new BufferLayout(BufferType.DATA, 16); private static final BufferLayout VALUES_8 = new BufferLayout(BufferType.DATA, 8); private static final BufferLayout LARGE_SIZE_BUFFER = new BufferLayout(BufferType.SIZE, 64); private static final BufferLayout SIZE_BUFFER = new BufferLayout(BufferType.SIZE, 32); private static final BufferLayout VIEW_BUFFER = new BufferLayout(BufferType.VIEWS, 16); public static BufferLayout typeBuffer() { return TYPE_BUFFER; } public static BufferLayout offsetBuffer() { return OFFSET_BUFFER; } public static BufferLayout largeOffsetBuffer() { return LARGE_OFFSET_BUFFER; } public static BufferLayout sizeBuffer() { return SIZE_BUFFER; } public static BufferLayout largeSizeBuffer() { return LARGE_SIZE_BUFFER; } /** * Returns a databuffer for the given bitwidth. Only supports powers of two between 8 and 128 * inclusive. */ public static BufferLayout dataBuffer(int typeBitWidth) { switch (typeBitWidth) { case 8: return VALUES_8; case 16: return VALUES_16; case 32: return VALUES_32; case 64: return VALUES_64; case 128: return VALUES_128; case 256: return VALUES_256; default: throw new IllegalArgumentException("only 8, 16, 32, 64, 128, or 256 bits supported"); } } public static BufferLayout booleanVector() { return BIT_BUFFER; } public static BufferLayout validityVector() { return VALIDITY_BUFFER; } public static BufferLayout byteVector() { return dataBuffer(8); } public static BufferLayout viewVector() { return VIEW_BUFFER; } private final short typeBitWidth; private final BufferType type; BufferLayout(BufferType type, int typeBitWidth) { super(); this.type = Preconditions.checkNotNull(type); this.typeBitWidth = (short) typeBitWidth; if (typeBitWidth <= 0) { throw new IllegalArgumentException("bitWidth invalid: " + typeBitWidth); } } public int getTypeBitWidth() { return typeBitWidth; } public BufferType getType() { return type; } @Override public String toString() { return String.format("%s(%s)", type, typeBitWidth); } @Override public int hashCode() { return 31 * (31 + type.hashCode()) + typeBitWidth; } @Override public boolean equals(Object obj) { if (this == obj) { return true; } if (obj == null) { return false; } if (!(obj instanceof BufferLayout)) { return false; } BufferLayout other = (BufferLayout) obj; return type.equals(other.type) && (typeBitWidth == other.typeBitWidth); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/DateDayVector.java000066400000000000000000000255671515451615200310600ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.DateDayReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.DateDayHolder; import org.apache.arrow.vector.holders.NullableDateDayHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * DateDayVector implements a fixed width (4 bytes) vector of date values which could be null. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public final class DateDayVector extends BaseFixedWidthVector implements ValueIterableVector { public static final byte TYPE_WIDTH = 4; /** * Instantiate a DateDayVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public DateDayVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.DATEDAY.getType()), allocator); } /** * Instantiate a DateDayVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public DateDayVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a DateDayVector. This doesn't allocate any memory for the data in vector. * * @param field Field materialized by this vector * @param allocator allocator for memory management. */ public DateDayVector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new DateDayReaderImpl(DateDayVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.DATEDAY; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public int get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getInt((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableDateDayHolder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getInt((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Integer getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getInt((long) index * TYPE_WIDTH); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, int value) { valueBuffer.setInt((long) index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, int value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableDateDayHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, DateDayHolder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, int)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, int value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableDateDayHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableDateDayHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, DateDayHolder)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, DateDayHolder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void set(int index, int isSet, int value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, int)} except that it handles the case when index is greater than * or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setSafe(int index, int isSet, int value) { handleSafe(index); set(index, isSet, value); } /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static int get(final ArrowBuf buffer, final int index) { return buffer.getInt((long) index * TYPE_WIDTH); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((DateDayVector) to); } private class TransferImpl implements TransferPair { DateDayVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new DateDayVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new DateDayVector(field, allocator); } public TransferImpl(DateDayVector to) { this.to = to; } @Override public DateDayVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, DateDayVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/DateMilliVector.java000066400000000000000000000261511515451615200313770ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import java.time.LocalDateTime; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.DateMilliReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.DateMilliHolder; import org.apache.arrow.vector.holders.NullableDateMilliHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.DateUtility; import org.apache.arrow.vector.util.TransferPair; /** * DateMilliVector implements a fixed width vector (8 bytes) of date values which could be null. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public final class DateMilliVector extends BaseFixedWidthVector implements ValueIterableVector { public static final byte TYPE_WIDTH = 8; /** * Instantiate a DateMilliVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public DateMilliVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.DATEMILLI.getType()), allocator); } /** * Instantiate a DateMilliVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public DateMilliVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a DateMilliVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public DateMilliVector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new DateMilliReaderImpl(DateMilliVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.DATEMILLI; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public long get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getLong((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableDateMilliHolder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getLong((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public LocalDateTime getObject(int index) { if (isSet(index) == 0) { return null; } else { final long millis = valueBuffer.getLong((long) index * TYPE_WIDTH); return DateUtility.getLocalDateTimeFromEpochMilli(millis); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, long value) { valueBuffer.setLong((long) index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, long value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableDateMilliHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, DateMilliHolder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, long)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, long value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableDateMilliHolder)} except that it handles the case when index * is greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableDateMilliHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, DateMilliHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, DateMilliHolder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void set(int index, int isSet, long value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, long)} except that it handles the case when index is greater than * or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setSafe(int index, int isSet, long value) { handleSafe(index); set(index, isSet, value); } /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static long get(final ArrowBuf buffer, final int index) { return buffer.getLong((long) index * TYPE_WIDTH); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((DateMilliVector) to); } private class TransferImpl implements TransferPair { DateMilliVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new DateMilliVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new DateMilliVector(field, allocator); } public TransferImpl(DateMilliVector to) { this.to = to; } @Override public DateMilliVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, DateMilliVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/Decimal256Vector.java000066400000000000000000000517251515451615200313330ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import java.math.BigDecimal; import java.nio.ByteOrder; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.MemoryUtil; import org.apache.arrow.vector.complex.impl.Decimal256ReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.Decimal256Holder; import org.apache.arrow.vector.holders.NullableDecimal256Holder; 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.DecimalUtility; import org.apache.arrow.vector.util.TransferPair; import org.apache.arrow.vector.validate.ValidateUtil; /** * Decimal256Vector implements a fixed width vector (32 bytes) of decimal values which could be * null. A validity buffer (bit vector) is maintained to track which elements in the vector are * null. */ public final class Decimal256Vector extends BaseFixedWidthVector implements ValueIterableVector { public static final int MAX_PRECISION = 76; public static final byte TYPE_WIDTH = 32; private static final boolean LITTLE_ENDIAN = ByteOrder.nativeOrder() == ByteOrder.LITTLE_ENDIAN; private final int precision; private final int scale; /** * Instantiate a Decimal256Vector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public Decimal256Vector(String name, BufferAllocator allocator, int precision, int scale) { this( name, FieldType.nullable(new ArrowType.Decimal(precision, scale, /* bitWidth= */ TYPE_WIDTH * 8)), allocator); } /** * Instantiate a Decimal256Vector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public Decimal256Vector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a Decimal256Vector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public Decimal256Vector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); ArrowType.Decimal arrowType = (ArrowType.Decimal) field.getFieldType().getType(); this.precision = arrowType.getPrecision(); this.scale = arrowType.getScale(); } @Override protected FieldReader getReaderImpl() { return new Decimal256ReaderImpl(Decimal256Vector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.DECIMAL256; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public ArrowBuf get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.slice((long) index * TYPE_WIDTH, TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableDecimal256Holder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.buffer = valueBuffer; holder.precision = precision; holder.scale = scale; holder.start = ((long) index) * TYPE_WIDTH; } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public BigDecimal getObject(int index) { if (isSet(index) == 0) { return null; } else { return getObjectNotNull(index); } } /** * Same as {@link #getObject(int)}, but does not check for null. * * @param index position of element * @return element at given index */ public BigDecimal getObjectNotNull(int index) { return DecimalUtility.getBigDecimalFromArrowBuf(valueBuffer, index, scale, TYPE_WIDTH); } /** Return precision for the decimal value. */ public int getPrecision() { return precision; } /** Return scale for the decimal value. */ public int getScale() { return scale; } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the element at the given index to the given value. * * @param index position of element * @param buffer ArrowBuf containing decimal value. */ public void set(int index, ArrowBuf buffer) { BitVectorHelper.setBit(validityBuffer, index); valueBuffer.setBytes((long) index * TYPE_WIDTH, buffer, 0, TYPE_WIDTH); } /** * Set the decimal element at given index to the provided array of bytes. Decimal256 is now * implemented as Native Endian. This API allows the user to pass a decimal value in the form of * byte array in BE byte order. * *

    Consumers of Arrow code can use this API instead of first swapping the source bytes (doing a * write and read) and then finally writing to ArrowBuf of decimal vector. * *

    This method takes care of adding the necessary padding if the length of byte array is less * than 32 (length of decimal type). * * @param index position of element * @param value array of bytes containing decimal in big endian byte order. */ public void setBigEndian(int index, byte[] value) { BitVectorHelper.setBit(validityBuffer, index); final int length = value.length; // do the bound check. valueBuffer.checkBytes((long) index * TYPE_WIDTH, (long) (index + 1) * TYPE_WIDTH); long outAddress = valueBuffer.memoryAddress() + (long) index * TYPE_WIDTH; if (length == 0) { MemoryUtil.setMemory(outAddress, Decimal256Vector.TYPE_WIDTH, (byte) 0); return; } if (LITTLE_ENDIAN) { // swap bytes to convert BE to LE for (int byteIdx = 0; byteIdx < length; ++byteIdx) { MemoryUtil.putByte(outAddress + byteIdx, value[length - 1 - byteIdx]); } if (length == TYPE_WIDTH) { return; } if (length < TYPE_WIDTH) { // sign extend final byte pad = (byte) (value[0] < 0 ? 0xFF : 0x00); MemoryUtil.setMemory(outAddress + length, Decimal256Vector.TYPE_WIDTH - length, pad); return; } } else { if (length <= TYPE_WIDTH) { // copy data from value to outAddress MemoryUtil.copyToMemory( value, 0, outAddress + Decimal256Vector.TYPE_WIDTH - length, length); // sign extend final byte pad = (byte) (value[0] < 0 ? 0xFF : 0x00); MemoryUtil.setMemory(outAddress, Decimal256Vector.TYPE_WIDTH - length, pad); return; } } throw new IllegalArgumentException( "Invalid decimal value length. Valid length in [1 - 32], got " + length); } /** * Set the element at the given index to the given value. * * @param index position of element * @param start start index of data in the buffer * @param buffer ArrowBuf containing decimal value. */ public void set(int index, long start, ArrowBuf buffer) { BitVectorHelper.setBit(validityBuffer, index); valueBuffer.setBytes((long) index * TYPE_WIDTH, buffer, start, TYPE_WIDTH); } /** * Sets the element at given index using the buffer whose size maybe <= 32 bytes. * * @param index index to write the decimal to * @param start start of value in the buffer * @param buffer contains the decimal in native endian bytes * @param length length of the value in the buffer */ public void setSafe(int index, long start, ArrowBuf buffer, int length) { handleSafe(index); BitVectorHelper.setBit(validityBuffer, index); // do the bound checks. buffer.checkBytes(start, start + length); valueBuffer.checkBytes((long) index * TYPE_WIDTH, (long) (index + 1) * TYPE_WIDTH); long inAddress = buffer.memoryAddress() + start; long outAddress = valueBuffer.memoryAddress() + (long) index * TYPE_WIDTH; if (LITTLE_ENDIAN) { MemoryUtil.copyMemory(inAddress, outAddress, length); // sign extend if (length < TYPE_WIDTH) { byte msb = MemoryUtil.getByte(inAddress + length - 1); final byte pad = (byte) (msb < 0 ? 0xFF : 0x00); MemoryUtil.setMemory(outAddress + length, Decimal256Vector.TYPE_WIDTH - length, pad); } } else { MemoryUtil.copyMemory(inAddress, outAddress + Decimal256Vector.TYPE_WIDTH - length, length); // sign extend if (length < TYPE_WIDTH) { byte msb = MemoryUtil.getByte(inAddress); final byte pad = (byte) (msb < 0 ? 0xFF : 0x00); MemoryUtil.setMemory(outAddress, Decimal256Vector.TYPE_WIDTH - length, pad); } } } /** * Sets the element at given index using the buffer whose size maybe <= 32 bytes. * * @param index index to write the decimal to * @param start start of value in the buffer * @param buffer contains the decimal in big endian bytes * @param length length of the value in the buffer */ public void setBigEndianSafe(int index, long start, ArrowBuf buffer, int length) { handleSafe(index); BitVectorHelper.setBit(validityBuffer, index); // do the bound checks. buffer.checkBytes(start, start + length); valueBuffer.checkBytes((long) index * TYPE_WIDTH, (long) (index + 1) * TYPE_WIDTH); // not using buffer.getByte() to avoid boundary checks for every byte. long inAddress = buffer.memoryAddress() + start; long outAddress = valueBuffer.memoryAddress() + (long) index * TYPE_WIDTH; if (LITTLE_ENDIAN) { // swap bytes to convert BE to LE for (int byteIdx = 0; byteIdx < length; ++byteIdx) { byte val = MemoryUtil.getByte((inAddress + length - 1) - byteIdx); MemoryUtil.putByte(outAddress + byteIdx, val); } // sign extend if (length < 32) { byte msb = MemoryUtil.getByte(inAddress); final byte pad = (byte) (msb < 0 ? 0xFF : 0x00); MemoryUtil.setMemory(outAddress + length, Decimal256Vector.TYPE_WIDTH - length, pad); } } else { MemoryUtil.copyMemory(inAddress, outAddress + Decimal256Vector.TYPE_WIDTH - length, length); // sign extend if (length < TYPE_WIDTH) { byte msb = MemoryUtil.getByte(inAddress); final byte pad = (byte) (msb < 0 ? 0xFF : 0x00); MemoryUtil.setMemory(outAddress, Decimal256Vector.TYPE_WIDTH - length, pad); } } } /** * Set the element at the given index to the given value. * * @param index position of element * @param value BigDecimal containing decimal value. */ public void set(int index, BigDecimal value) { BitVectorHelper.setBit(validityBuffer, index); DecimalUtility.checkPrecisionAndScale(value, precision, scale); DecimalUtility.writeBigDecimalToArrowBuf(value, valueBuffer, index, TYPE_WIDTH); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value long value. */ public void set(int index, long value) { BitVectorHelper.setBit(validityBuffer, index); DecimalUtility.writeLongToArrowBuf(value, valueBuffer, index, TYPE_WIDTH); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableDecimal256Holder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); valueBuffer.setBytes((long) index * TYPE_WIDTH, holder.buffer, holder.start, TYPE_WIDTH); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, Decimal256Holder holder) { BitVectorHelper.setBit(validityBuffer, index); valueBuffer.setBytes((long) index * TYPE_WIDTH, holder.buffer, holder.start, TYPE_WIDTH); } /** * Same as {@link #set(int, ArrowBuf)} except that it handles the case when index is greater than * or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param buffer ArrowBuf containing decimal value. */ public void setSafe(int index, ArrowBuf buffer) { handleSafe(index); set(index, buffer); } /** * Same as {@link #setBigEndian(int, byte[])} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. */ public void setBigEndianSafe(int index, byte[] value) { handleSafe(index); setBigEndian(index, value); } /** * Same as {@link #set(int, int, ArrowBuf)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param start start index of data in the buffer * @param buffer ArrowBuf containing decimal value. */ public void setSafe(int index, long start, ArrowBuf buffer) { handleSafe(index); set(index, start, buffer); } /** * Same as {@link #set(int, BigDecimal)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value BigDecimal containing decimal value. */ public void setSafe(int index, BigDecimal value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, long)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value long value. */ public void setSafe(int index, long value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableDecimalHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableDecimal256Holder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, Decimal256Holder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, Decimal256Holder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param start start position of the value in the buffer * @param buffer buffer containing the value to be stored in the vector */ public void set(int index, int isSet, long start, ArrowBuf buffer) { if (isSet > 0) { set(index, start, buffer); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #setSafe(int, int, int, ArrowBuf)} except that it handles the case when the * position of new value is beyond the current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param start start position of the value in the buffer * @param buffer buffer containing the value to be stored in the vector */ public void setSafe(int index, int isSet, long start, ArrowBuf buffer) { handleSafe(index); set(index, isSet, start, buffer); } @Override public void validateScalars() { for (int i = 0; i < getValueCount(); ++i) { BigDecimal value = getObject(i); if (value != null) { ValidateUtil.validateOrThrow( DecimalUtility.checkPrecisionAndScaleNoThrow(value, getPrecision(), getScale()), "Invalid value for Decimal256Vector at position " + i + ". Value does not fit in precision " + getPrecision() + " and scale " + getScale() + "."); } } } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((Decimal256Vector) to); } private class TransferImpl implements TransferPair { Decimal256Vector to; public TransferImpl(String ref, BufferAllocator allocator) { to = (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) { to = new Decimal256Vector(field, allocator); } public TransferImpl(Decimal256Vector to) { this.to = to; } @Override public Decimal256Vector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, Decimal256Vector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/DecimalVector.java000066400000000000000000000514371515451615200310760ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import java.math.BigDecimal; import java.nio.ByteOrder; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.MemoryUtil; import org.apache.arrow.vector.complex.impl.DecimalReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.DecimalHolder; import org.apache.arrow.vector.holders.NullableDecimalHolder; 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.DecimalUtility; import org.apache.arrow.vector.util.TransferPair; import org.apache.arrow.vector.validate.ValidateUtil; /** * DecimalVector implements a fixed width vector (16 bytes) of decimal values which could be null. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public final class DecimalVector extends BaseFixedWidthVector implements ValueIterableVector { public static final int MAX_PRECISION = 38; public static final byte TYPE_WIDTH = 16; private static final boolean LITTLE_ENDIAN = ByteOrder.nativeOrder() == ByteOrder.LITTLE_ENDIAN; private final int precision; private final int scale; /** * Instantiate a DecimalVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public DecimalVector(String name, BufferAllocator allocator, int precision, int scale) { this( name, FieldType.nullable(new ArrowType.Decimal(precision, scale, TYPE_WIDTH * 8)), allocator); } /** * Instantiate a DecimalVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public DecimalVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a DecimalVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public DecimalVector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); ArrowType.Decimal arrowType = (ArrowType.Decimal) field.getFieldType().getType(); this.precision = arrowType.getPrecision(); this.scale = arrowType.getScale(); } @Override protected FieldReader getReaderImpl() { return new DecimalReaderImpl(DecimalVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.DECIMAL; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public ArrowBuf get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.slice((long) index * TYPE_WIDTH, TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableDecimalHolder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.buffer = valueBuffer; holder.precision = precision; holder.scale = scale; holder.start = (long) index * TYPE_WIDTH; } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public BigDecimal getObject(int index) { if (isSet(index) == 0) { return null; } else { return getObjectNotNull(index); } } /** * Same as {@link #getObject(int)} but does not check for null. * * @param index position of element * @return element at given index */ public BigDecimal getObjectNotNull(int index) { return DecimalUtility.getBigDecimalFromArrowBuf(valueBuffer, index, scale, TYPE_WIDTH); } /** Return precision for the decimal value. */ public int getPrecision() { return precision; } /** Return scale for the decimal value. */ public int getScale() { return scale; } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the element at the given index to the given value. * * @param index position of element * @param buffer ArrowBuf containing decimal value. */ public void set(int index, ArrowBuf buffer) { BitVectorHelper.setBit(validityBuffer, index); valueBuffer.setBytes((long) index * TYPE_WIDTH, buffer, 0, TYPE_WIDTH); } /** * Set the decimal element at given index to the provided array of bytes. Decimal is now * implemented as Native Endian. This API allows the user to pass a decimal value in the form of * byte array in BE byte order. * *

    Consumers of Arrow code can use this API instead of first swapping the source bytes (doing a * write and read) and then finally writing to ArrowBuf of decimal vector. * *

    This method takes care of adding the necessary padding if the length of byte array is less * than 16 (length of decimal type). * * @param index position of element * @param value array of bytes containing decimal in big endian byte order. */ public void setBigEndian(int index, byte[] value) { BitVectorHelper.setBit(validityBuffer, index); final int length = value.length; // do the bound check. valueBuffer.checkBytes((long) index * TYPE_WIDTH, (long) (index + 1) * TYPE_WIDTH); long outAddress = valueBuffer.memoryAddress() + (long) index * TYPE_WIDTH; if (length == 0) { MemoryUtil.setMemory(outAddress, DecimalVector.TYPE_WIDTH, (byte) 0); return; } if (LITTLE_ENDIAN) { // swap bytes to convert BE to LE for (int byteIdx = 0; byteIdx < length; ++byteIdx) { MemoryUtil.putByte(outAddress + byteIdx, value[length - 1 - byteIdx]); } if (length == TYPE_WIDTH) { return; } if (length < TYPE_WIDTH) { // sign extend final byte pad = (byte) (value[0] < 0 ? 0xFF : 0x00); MemoryUtil.setMemory(outAddress + length, DecimalVector.TYPE_WIDTH - length, pad); return; } } else { if (length <= TYPE_WIDTH) { // copy data from value to outAddress MemoryUtil.copyToMemory(value, 0, outAddress + DecimalVector.TYPE_WIDTH - length, length); // sign extend final byte pad = (byte) (value[0] < 0 ? 0xFF : 0x00); MemoryUtil.setMemory(outAddress, DecimalVector.TYPE_WIDTH - length, pad); return; } } throw new IllegalArgumentException( "Invalid decimal value length. Valid length in [1 - 16], got " + length); } /** * Set the element at the given index to the given value. * * @param index position of element * @param start start index of data in the buffer * @param buffer ArrowBuf containing decimal value. */ public void set(int index, long start, ArrowBuf buffer) { BitVectorHelper.setBit(validityBuffer, index); valueBuffer.setBytes((long) index * TYPE_WIDTH, buffer, start, TYPE_WIDTH); } /** * Sets the element at given index using the buffer whose size maybe <= 16 bytes. * * @param index index to write the decimal to * @param start start of value in the buffer * @param buffer contains the decimal in native endian bytes * @param length length of the value in the buffer */ public void setSafe(int index, long start, ArrowBuf buffer, int length) { handleSafe(index); BitVectorHelper.setBit(validityBuffer, index); // do the bound checks. buffer.checkBytes(start, start + length); valueBuffer.checkBytes((long) index * TYPE_WIDTH, (long) (index + 1) * TYPE_WIDTH); long inAddress = buffer.memoryAddress() + start; long outAddress = valueBuffer.memoryAddress() + (long) index * TYPE_WIDTH; if (LITTLE_ENDIAN) { MemoryUtil.copyMemory(inAddress, outAddress, length); // sign extend if (length < TYPE_WIDTH) { byte msb = MemoryUtil.getByte(inAddress + length - 1); final byte pad = (byte) (msb < 0 ? 0xFF : 0x00); MemoryUtil.setMemory(outAddress + length, DecimalVector.TYPE_WIDTH - length, pad); } } else { MemoryUtil.copyMemory(inAddress, outAddress + DecimalVector.TYPE_WIDTH - length, length); // sign extend if (length < TYPE_WIDTH) { byte msb = MemoryUtil.getByte(inAddress); final byte pad = (byte) (msb < 0 ? 0xFF : 0x00); MemoryUtil.setMemory(outAddress, DecimalVector.TYPE_WIDTH - length, pad); } } } /** * Sets the element at given index using the buffer whose size maybe <= 16 bytes. * * @param index index to write the decimal to * @param start start of value in the buffer * @param buffer contains the decimal in big endian bytes * @param length length of the value in the buffer */ public void setBigEndianSafe(int index, long start, ArrowBuf buffer, int length) { handleSafe(index); BitVectorHelper.setBit(validityBuffer, index); // do the bound checks. buffer.checkBytes(start, start + length); valueBuffer.checkBytes((long) index * TYPE_WIDTH, (long) (index + 1) * TYPE_WIDTH); // not using buffer.getByte() to avoid boundary checks for every byte. long inAddress = buffer.memoryAddress() + start; long outAddress = valueBuffer.memoryAddress() + (long) index * TYPE_WIDTH; if (LITTLE_ENDIAN) { // swap bytes to convert BE to LE for (int byteIdx = 0; byteIdx < length; ++byteIdx) { byte val = MemoryUtil.getByte((inAddress + length - 1) - byteIdx); MemoryUtil.putByte(outAddress + byteIdx, val); } // sign extend if (length < TYPE_WIDTH) { byte msb = MemoryUtil.getByte(inAddress); final byte pad = (byte) (msb < 0 ? 0xFF : 0x00); MemoryUtil.setMemory(outAddress + length, DecimalVector.TYPE_WIDTH - length, pad); } } else { MemoryUtil.copyMemory(inAddress, outAddress + DecimalVector.TYPE_WIDTH - length, length); // sign extend if (length < TYPE_WIDTH) { byte msb = MemoryUtil.getByte(inAddress); final byte pad = (byte) (msb < 0 ? 0xFF : 0x00); MemoryUtil.setMemory(outAddress, DecimalVector.TYPE_WIDTH - length, pad); } } } /** * Set the element at the given index to the given value. * * @param index position of element * @param value BigDecimal containing decimal value. */ public void set(int index, BigDecimal value) { BitVectorHelper.setBit(validityBuffer, index); DecimalUtility.checkPrecisionAndScale(value, precision, scale); DecimalUtility.writeBigDecimalToArrowBuf(value, valueBuffer, index, TYPE_WIDTH); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value long value. */ public void set(int index, long value) { BitVectorHelper.setBit(validityBuffer, index); DecimalUtility.writeLongToArrowBuf(value, valueBuffer, index, TYPE_WIDTH); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableDecimalHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); valueBuffer.setBytes((long) index * TYPE_WIDTH, holder.buffer, holder.start, TYPE_WIDTH); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, DecimalHolder holder) { BitVectorHelper.setBit(validityBuffer, index); valueBuffer.setBytes((long) index * TYPE_WIDTH, holder.buffer, holder.start, TYPE_WIDTH); } /** * Same as {@link #set(int, ArrowBuf)} except that it handles the case when index is greater than * or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param buffer ArrowBuf containing decimal value. */ public void setSafe(int index, ArrowBuf buffer) { handleSafe(index); set(index, buffer); } /** * Same as {@link #setBigEndian(int, byte[])} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. */ public void setBigEndianSafe(int index, byte[] value) { handleSafe(index); setBigEndian(index, value); } /** * Same as {@link #set(int, long, ArrowBuf)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param start start index of data in the buffer * @param buffer ArrowBuf containing decimal value. */ public void setSafe(int index, long start, ArrowBuf buffer) { handleSafe(index); set(index, start, buffer); } /** * Same as {@link #set(int, BigDecimal)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value BigDecimal containing decimal value. */ public void setSafe(int index, BigDecimal value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, long)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value long value. */ public void setSafe(int index, long value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableDecimalHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableDecimalHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, DecimalHolder)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, DecimalHolder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param start start position of the value in the buffer * @param buffer buffer containing the value to be stored in the vector */ public void set(int index, int isSet, long start, ArrowBuf buffer) { if (isSet > 0) { set(index, start, buffer); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, long, ArrowBuf)} except that it handles the case when the * position of new value is beyond the current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param start start position of the value in the buffer * @param buffer buffer containing the value to be stored in the vector */ public void setSafe(int index, int isSet, long start, ArrowBuf buffer) { handleSafe(index); set(index, isSet, start, buffer); } @Override public void validateScalars() { for (int i = 0; i < getValueCount(); ++i) { BigDecimal value = getObject(i); if (value != null) { ValidateUtil.validateOrThrow( DecimalUtility.checkPrecisionAndScaleNoThrow(value, getPrecision(), getScale()), "Invalid value for DecimalVector at position " + i + ". Value does not fit in precision " + getPrecision() + " and scale " + getScale() + "."); } } } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((DecimalVector) to); } private class TransferImpl implements TransferPair { DecimalVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = (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) { to = new DecimalVector(field, allocator); } public TransferImpl(DecimalVector to) { this.to = to; } @Override public DecimalVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, DecimalVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/DensityAwareVector.java000066400000000000000000000044711515451615200321330ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * Vector that support density aware initial capacity settings. We use this for ListVector and * VarCharVector as of now to control the memory allocated. * *

    For ListVector, we have been using a multiplier of 5 to compute the initial capacity of the * inner data vector. For deeply nested lists and lists with lots of NULL values, this is * over-allocation upfront. So density helps to be conservative when computing the value capacity of * the inner vector. * *

    For example, a density value of 10 implies each position in the list vector has a list of 10 * values. So we will provision an initial capacity of (valuecount * 10) for the inner vector. A * density value of 0.1 implies out of 10 positions in the list vector, 1 position has a list of * size 1 and remaining positions are null (no lists) or empty lists. This helps in tightly * controlling the memory we provision for inner data vector. * *

    Similar analogy is applicable for VarCharVector where the capacity of the data buffer can be * controlled using density multiplier instead of default multiplier of 8 (default size of average * varchar length). * *

    Also from container vectors, we propagate the density down the inner vectors so that they can * use it appropriately. */ public interface DensityAwareVector { /** * Set value with density. * * @param valueCount the number of values in this vector * @param density the density of the vector */ void setInitialCapacity(int valueCount, double density); } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/DurationVector.java000066400000000000000000000325031515451615200313160ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.concurrent.TimeUnit.MICROSECONDS; import static org.apache.arrow.vector.NullCheckingForGet.NULL_CHECKING_ENABLED; import java.time.Duration; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.DurationReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.DurationHolder; import org.apache.arrow.vector.holders.NullableDurationHolder; 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.TransferPair; /** * DurationVector implements a fixed width vector (8 bytes) of a configurable TimeUnit granularity * duration values which could be null. A validity buffer (bit vector) is maintained to track which * elements in the vector are null. */ public final class DurationVector extends BaseFixedWidthVector implements ValueIterableVector { public static final byte TYPE_WIDTH = 8; private final TimeUnit unit; /** * Instantiate a DurationVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public DurationVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a DurationVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public DurationVector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); this.unit = ((ArrowType.Duration) field.getFieldType().getType()).getUnit(); } @Override protected FieldReader getReaderImpl() { return new DurationReaderImpl(DurationVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link MinorType} */ @Override public MinorType getMinorType() { return MinorType.DURATION; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static long get(final ArrowBuf buffer, final int index) { return buffer.getLong((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public ArrowBuf get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { return null; } return valueBuffer.slice((long) index * TYPE_WIDTH, TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableDurationHolder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = get(valueBuffer, index); holder.unit = this.unit; } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Duration getObject(int index) { if (isSet(index) == 0) { return null; } else { return getObjectNotNull(index); } } /** * Same as {@link #getObject(int)} but does not check for null. * * @param index position of element * @return element at given index */ public Duration getObjectNotNull(int index) { final long value = get(valueBuffer, index); return toDuration(value, unit); } /** Converts the given value and unit to the appropriate {@link Duration}. */ public static Duration toDuration(long value, TimeUnit unit) { switch (unit) { case SECOND: return Duration.ofSeconds(value); case MILLISECOND: return Duration.ofMillis(value); case NANOSECOND: return Duration.ofNanos(value); case MICROSECOND: return Duration.ofNanos(MICROSECONDS.toNanos(value)); default: throw new IllegalArgumentException("Unknown timeunit: " + unit); } } /** * Get the Interval value at a given index as a {@link StringBuilder} object. * * @param index position of the element * @return String Builder object with Interval in java.time.Duration format. */ public StringBuilder getAsStringBuilder(int index) { if (isSet(index) == 0) { return null; } else { return getAsStringBuilderHelper(index); } } private StringBuilder getAsStringBuilderHelper(int index) { return new StringBuilder(getObject(index).toString()); } /** Gets the time unit of the duration. */ public TimeUnit getUnit() { return unit; } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, ArrowBuf value) { BitVectorHelper.setBit(validityBuffer, index); valueBuffer.setBytes((long) index * TYPE_WIDTH, value, 0, TYPE_WIDTH); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value The duration value (in the timeunit associated with this vector) */ public void set(int index, long value) { final long offsetIndex = (long) index * TYPE_WIDTH; BitVectorHelper.setBit(validityBuffer, index); valueBuffer.setLong(offsetIndex, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableDurationHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (!this.unit.equals(holder.unit)) { throw new IllegalArgumentException( String.format("holder.unit: %s not equal to vector unit: %s", holder.unit, this.unit)); } else if (holder.isSet > 0) { set(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, DurationHolder holder) { if (!this.unit.equals(holder.unit)) { throw new IllegalArgumentException( String.format("holder.unit: %s not equal to vector unit: %s", holder.unit, this.unit)); } set(index, holder.value); } /** * Same as {@link #set(int, ArrowBuf)} except that it handles the case when index is greater than * or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, ArrowBuf value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, long)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value duration in the time unit this vector was constructed with */ public void setSafe(int index, long value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableDurationHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableDurationHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, DurationHolder)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, DurationHolder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value The duration value (in the TimeUnit associated with this vector). */ public void set(int index, int isSet, long value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, long)} except that it handles the case when index is greater than * or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value The duration value (in the timeunit associated with this vector) */ public void setSafe(int index, int isSet, long value) { handleSafe(index); set(index, isSet, value); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising of this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((DurationVector) to); } private class TransferImpl implements TransferPair { DurationVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new DurationVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new DurationVector(field, allocator); } public TransferImpl(DurationVector to) { this.to = to; } @Override public DurationVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, DurationVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ElementAddressableVector.java000066400000000000000000000032221515451615200332500ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.memory.util.ArrowBufPointer; /** * Vector for which each data element resides in a continuous memory region, so it can be pointed to * by an {@link org.apache.arrow.memory.util.ArrowBufPointer}. */ public interface ElementAddressableVector extends ValueVector { /** * Gets the pointer for the data at the given index. * * @param index the index for the data. * @return the pointer to the data. */ ArrowBufPointer getDataPointer(int index); /** * Gets the pointer for the data at the given index. * * @param index the index for the data. * @param reuse the data pointer to fill, this avoids creating a new pointer object. * @return the pointer to the data, it should be the same one as the input parameter */ ArrowBufPointer getDataPointer(int index, ArrowBufPointer reuse); } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ExtensionTypeVector.java000066400000000000000000000170751515451615200323560ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.Iterator; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.OutOfMemoryException; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.util.CallBack; import org.apache.arrow.vector.util.TransferPair; /** * A vector that wraps an underlying vector, used to help implement extension types. * * @param The wrapped vector type. */ public abstract class ExtensionTypeVector extends BaseValueVector implements FieldVector { private final T underlyingVector; private final String name; /** * Instantiate an extension type vector. * * @param name name of the vector * @param allocator allocator for memory management * @param underlyingVector underlying filed vector */ public ExtensionTypeVector(String name, BufferAllocator allocator, T underlyingVector) { super(allocator); Preconditions.checkNotNull(underlyingVector, "underlyingVector cannot be null."); this.name = name; this.underlyingVector = underlyingVector; } /** * Instantiate an extension type vector. * * @param field field materialized by this vector. * @param allocator allocator for memory management * @param underlyingVector underlying filed vector */ public ExtensionTypeVector(Field field, BufferAllocator allocator, T underlyingVector) { this(field.getName(), allocator, underlyingVector); } @Override public String getName() { return name; } /** Get the underlying vector. */ public T getUnderlyingVector() { return underlyingVector; } @Override public void allocateNew() throws OutOfMemoryException { this.underlyingVector.allocateNew(); } @Override public boolean allocateNewSafe() { return this.underlyingVector.allocateNewSafe(); } @Override public void reAlloc() { this.underlyingVector.reAlloc(); } @Override public void setInitialCapacity(int numRecords) { this.underlyingVector.setInitialCapacity(numRecords); } @Override public int getValueCapacity() { return this.underlyingVector.getValueCapacity(); } @Override public void reset() { this.underlyingVector.reset(); } @Override public Field getField() { return this.underlyingVector.getField(); } @Override public MinorType getMinorType() { return MinorType.EXTENSIONTYPE; } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return underlyingVector.getTransferPair(ref, allocator); } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return underlyingVector.getTransferPair(ref, allocator, callBack); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return underlyingVector.getTransferPair(field, allocator); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { return underlyingVector.getTransferPair(field, allocator, callBack); } @Override public TransferPair makeTransferPair(ValueVector target) { return underlyingVector.makeTransferPair(target); } @Override protected FieldReader getReaderImpl() { return underlyingVector.getReader(); } @Override public int getBufferSize() { return underlyingVector.getBufferSize(); } @Override public int getBufferSizeFor(int valueCount) { return underlyingVector.getBufferSizeFor(valueCount); } @Override public ArrowBuf[] getBuffers(boolean clear) { return underlyingVector.getBuffers(clear); } @Override public ArrowBuf getValidityBuffer() { return underlyingVector.getValidityBuffer(); } @Override public ArrowBuf getDataBuffer() { return underlyingVector.getDataBuffer(); } @Override public ArrowBuf getOffsetBuffer() { return underlyingVector.getOffsetBuffer(); } @Override public int getValueCount() { return underlyingVector.getValueCount(); } @Override public void setValueCount(int valueCount) { underlyingVector.setValueCount(valueCount); } /** * Get the extension object at the specified index. * *

    Generally, this should access the underlying vector and construct the corresponding Java * object from the raw data. */ @Override public abstract Object getObject(int index); @Override public int getNullCount() { return underlyingVector.getNullCount(); } @Override public boolean isNull(int index) { return underlyingVector.isNull(index); } @Override public void setNull(int index) { underlyingVector.setNull(index); } @Override public void initializeChildrenFromFields(List children) { underlyingVector.initializeChildrenFromFields(children); } @Override public List getChildrenFromFields() { return underlyingVector.getChildrenFromFields(); } @Override public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers) { underlyingVector.loadFieldBuffers(fieldNode, ownBuffers); } @Override public List getFieldBuffers() { return underlyingVector.getFieldBuffers(); } /** * Get the inner vectors. * * @return the inner vectors for this field as defined by the TypeLayout * @deprecated This API will be removed as the current implementations no longer support inner * vectors. */ @Deprecated @Override public List getFieldInnerVectors() { return underlyingVector.getFieldInnerVectors(); } @Override public long getValidityBufferAddress() { return underlyingVector.getValidityBufferAddress(); } @Override public long getDataBufferAddress() { return underlyingVector.getDataBufferAddress(); } @Override public long getOffsetBufferAddress() { return underlyingVector.getOffsetBufferAddress(); } @Override public void clear() { underlyingVector.clear(); } @Override public void close() { underlyingVector.close(); } @Override public TransferPair getTransferPair(BufferAllocator allocator) { return underlyingVector.getTransferPair(allocator); } @Override public Iterator iterator() { return underlyingVector.iterator(); } @Override public BufferAllocator getAllocator() { return underlyingVector.getAllocator(); } @Override public OUT accept(VectorVisitor visitor, IN value) { return visitor.visit(this, value); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/FieldVector.java000066400000000000000000000111531515451615200305520ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.types.pojo.Field; /** * A vector corresponding to a Field in the schema. It has inner vectors backed by buffers * (validity, offsets, data, ...) */ public interface FieldVector extends ValueVector { /** * Initializes the child vectors to be later loaded with loadBuffers. * * @param children the schema */ void initializeChildrenFromFields(List children); /** * The returned list is the same size as the list passed to initializeChildrenFromFields. * * @return the children according to schema (empty for primitive types) */ List getChildrenFromFields(); /** * Loads data in the vectors. (ownBuffers must be the same size as getFieldVectors()) * * @param fieldNode the fieldNode * @param ownBuffers the buffers for this Field (own buffers only, children not included) */ void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers); /** * Get the buffers of the fields, (same size as getFieldVectors() since it is their content). * * @return the buffers containing the data for this vector (ready for reading) */ List getFieldBuffers(); /** * Retrieves the export buffer count for the C Data Interface. * * @return the number of variadic buffers */ default int getExportedCDataBufferCount() { return getFieldBuffers().size(); } /** * Export a given buffer and its memory address into a list of buffers and a pointer to the list * of buffers. * * @param buffer the buffer to export * @param buffers the list of buffers * @param buffersPtr the pointer to the list of buffers * @param nullValue the null value to use for null buffer * @param retain whether to retain the buffer when exporting */ default void exportBuffer( ArrowBuf buffer, List buffers, ArrowBuf buffersPtr, long nullValue, boolean retain) { if (buffer != null) { if (retain) { buffer.getReferenceManager().retain(); } buffersPtr.writeLong(buffer.memoryAddress()); } else { buffersPtr.writeLong(nullValue); } buffers.add(buffer); } /** * Export the buffers of the fields for C Data Interface. This method traverse the buffers and * export buffer and buffer's memory address into a list of buffers and a pointer to the list of * buffers. * *

    By default, when exporting a buffer, it will increase ref count for exported buffer that * counts the usage at imported side. */ default void exportCDataBuffers(List buffers, ArrowBuf buffersPtr, long nullValue) { List fieldBuffers = getFieldBuffers(); for (ArrowBuf arrowBuf : fieldBuffers) { exportBuffer(arrowBuf, buffers, buffersPtr, nullValue, true); } } /** * Get the inner vectors. * * @return the inner vectors for this field as defined by the TypeLayout * @deprecated This API will be removed as the current implementations no longer support inner * vectors. */ @Deprecated List getFieldInnerVectors(); /** * Gets the starting address of the underlying buffer associated with validity vector. * * @return buffer address */ long getValidityBufferAddress(); /** * Gets the starting address of the underlying buffer associated with data vector. * * @return buffer address */ long getDataBufferAddress(); /** * Gets the starting address of the underlying buffer associated with offset vector. * * @return buffer address */ long getOffsetBufferAddress(); /** * Set the element at the given index to null. * * @param index the value to change */ void setNull(int index); } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/FixedSizeBinaryVector.java000066400000000000000000000335651515451615200326010ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReusableBuffer; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.complex.impl.FixedSizeBinaryReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.FixedSizeBinaryHolder; import org.apache.arrow.vector.holders.NullableFixedSizeBinaryHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType.FixedSizeBinary; 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.validate.ValidateUtil; /** * FixedSizeBinaryVector implements a fixed width vector of binary values which could be null. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public class FixedSizeBinaryVector extends BaseFixedWidthVector implements ValueIterableVector { private final int byteWidth; /** * Instantiate a FixedSizeBinaryVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. * @param byteWidth byte width of the binary values */ public FixedSizeBinaryVector(String name, BufferAllocator allocator, int byteWidth) { this(name, FieldType.nullable(new FixedSizeBinary(byteWidth)), allocator); } /** * Instantiate a FixedSizeBinaryVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public FixedSizeBinaryVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a FixedSizeBinaryVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public FixedSizeBinaryVector(Field field, BufferAllocator allocator) { super(field, allocator, ((FixedSizeBinary) field.getFieldType().getType()).getByteWidth()); byteWidth = ((FixedSizeBinary) field.getFieldType().getType()).getByteWidth(); } @Override protected FieldReader getReaderImpl() { return new FixedSizeBinaryReaderImpl(FixedSizeBinaryVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.FIXEDSIZEBINARY; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public byte[] get(int index) { assert index >= 0; if (NULL_CHECKING_ENABLED && isSet(index) == 0) { return null; } final byte[] dst = new byte[byteWidth]; valueBuffer.getBytes((long) index * byteWidth, dst, 0, byteWidth); return dst; } /** * Read the value at the given position to the given output buffer. The caller is responsible for * checking for nullity first. * * @param index position of element. * @param buffer the buffer to write into. */ public void read(int index, ReusableBuffer buffer) { final int startOffset = index * byteWidth; buffer.set(valueBuffer, startOffset, byteWidth); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element * @param holder nullable holder to carry the buffer */ public void get(int index, NullableFixedSizeBinaryHolder holder) { assert index >= 0; if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.buffer = valueBuffer.slice((long) index * byteWidth, byteWidth); holder.byteWidth = byteWidth; } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public byte[] getObject(int index) { if (isSet(index) == 0) { return null; } return get(valueBuffer, index, byteWidth); } public int getByteWidth() { return byteWidth; } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** Sets the value at index to the provided one. */ public void set(int index, byte[] value) { assert index >= 0; Preconditions.checkNotNull(value, "expecting a valid byte array"); assert byteWidth <= value.length; BitVectorHelper.setBit(validityBuffer, index); valueBuffer.setBytes((long) index * byteWidth, value, 0, byteWidth); } /** * Same as {@link #set(int, byte[])} but reallocates if index is larger than * capacity. */ public void setSafe(int index, byte[] value) { handleSafe(index); set(index, value); } /** Sets the value if isSet is positive, otherwise sets the index to null/invalid. */ public void set(int index, int isSet, byte[] value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } public void setSafe(int index, int isSet, byte[] value) { handleSafe(index); set(index, isSet, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param buffer ArrowBuf containing binary value. */ public void set(int index, ArrowBuf buffer) { assert index >= 0; assert byteWidth <= buffer.capacity(); BitVectorHelper.setBit(validityBuffer, index); valueBuffer.setBytes((long) index * byteWidth, buffer, 0, byteWidth); } /** * Same as {@link #set(int, ArrowBuf)} except that it handles the case when index is greater than * or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param buffer ArrowBuf containing binary value. */ public void setSafe(int index, ArrowBuf buffer) { handleSafe(index); set(index, buffer); } /** * Set the element at the given index to the given value. * * @param index position of element * @param buffer ArrowBuf containing binary value. */ public void set(int index, int isSet, ArrowBuf buffer) { if (isSet > 0) { set(index, buffer); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, ArrowBuf)} except that it handles the case when index is greater than * or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param buffer ArrowBuf containing binary value. */ public void setSafe(int index, int isSet, ArrowBuf buffer) { handleSafe(index); set(index, isSet, buffer); } /** * Set the variable length element at the specified index to the data buffer supplied in the * holder. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void set(int index, FixedSizeBinaryHolder holder) { if (this.byteWidth != holder.byteWidth) { throw new IllegalArgumentException( String.format( "holder.byteWidth: %d not equal to vector byteWidth: %d", holder.byteWidth, this.byteWidth)); } set(index, holder.buffer); } /** * Same as {@link #set(int, FixedSizeBinaryHolder)} except that it handles the case where index * and length of new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void setSafe(int index, FixedSizeBinaryHolder holder) { handleSafe(index); set(index, holder); } /** * Set the variable length element at the specified index to the data buffer supplied in the * holder. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void set(int index, NullableFixedSizeBinaryHolder holder) { if (holder.isSet < 0) { throw new IllegalArgumentException("holder has a negative isSet value"); } else if (this.byteWidth != holder.byteWidth) { throw new IllegalArgumentException( String.format( "holder.byteWidth: %d not equal to vector byteWidth: %d", holder.byteWidth, this.byteWidth)); } else if (holder.isSet > 0) { set(index, holder.buffer); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, NullableFixedSizeBinaryHolder)} except that it handles the case where * index and length of new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void setSafe(int index, NullableFixedSizeBinaryHolder holder) { handleSafe(index); set(index, holder); } /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static byte[] get(final ArrowBuf buffer, final int index, final int byteWidth) { final byte[] dst = new byte[byteWidth]; buffer.getBytes((long) index * byteWidth, dst, 0, byteWidth); return dst; } @Override public void validateScalars() { for (int i = 0; i < getValueCount(); ++i) { byte[] value = get(i); if (value != null) { ValidateUtil.validateOrThrow( value.length == byteWidth, "Invalid value for FixedSizeBinaryVector at position " + i + ". The length was " + value.length + " but the length of each element should be " + byteWidth + "."); } } } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((FixedSizeBinaryVector) to); } private class TransferImpl implements TransferPair { FixedSizeBinaryVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new FixedSizeBinaryVector(ref, allocator, FixedSizeBinaryVector.this.byteWidth); } public TransferImpl(Field field, BufferAllocator allocator) { to = new FixedSizeBinaryVector(field, allocator); } public TransferImpl(FixedSizeBinaryVector to) { this.to = to; } @Override public FixedSizeBinaryVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, FixedSizeBinaryVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/FixedWidthVector.java000066400000000000000000000025441515451615200315720ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * Interface for all fixed width {@link ElementAddressableVector} (e.g. integer, fixed size binary, * etc). */ public interface FixedWidthVector extends ElementAddressableVector { /** * Allocate a new memory space for this vector. Must be called prior to using the ValueVector. * * @param valueCount Number of values in the vector. */ void allocateNew(int valueCount); /** Zero out the underlying buffer backing this vector. */ void zeroVector(); /** Get the width of the type in bytes. */ int getTypeWidth(); } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/Float2Vector.java000066400000000000000000000323111515451615200306550ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.Float16; import org.apache.arrow.vector.complex.impl.Float2ReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.Float2Holder; import org.apache.arrow.vector.holders.NullableFloat2Holder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * Float2Vector implements a fixed width (2 bytes) vector of short values which could be null. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public final class Float2Vector extends BaseFixedWidthVector implements FloatingPointVector, ValueIterableVector { public static final byte TYPE_WIDTH = 2; /** * Instantiate a Float2Vector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public Float2Vector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.FLOAT2.getType()), allocator); } /** * Instantiate a Float2Vector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public Float2Vector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a Float2Vector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public Float2Vector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new Float2ReaderImpl(Float2Vector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link MinorType} */ @Override public MinorType getMinorType() { return MinorType.FLOAT2; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public short get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getShort((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableFloat2Holder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getShort((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Short getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getShort((long) index * TYPE_WIDTH); } } /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ static short get(final ArrowBuf buffer, final int index) { return buffer.getShort((long) index * TYPE_WIDTH); } @Override public double getValueAsDouble(int index) { return getValueAsFloat(index); } public float getValueAsFloat(int index) { return Float16.toFloat(this.get(index)); } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, short value) { valueBuffer.setShort((long) index * TYPE_WIDTH, value); } private void setValue(int index, float value) { valueBuffer.setShort((long) index * TYPE_WIDTH, Float16.toFloat16(value)); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, short value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void setWithPossibleTruncate(int index, float value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableFloat2Holder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, Float2Holder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, short)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, short value) { handleSafe(index); set(index, value); } /** * Same as {@link #setWithPossibleTruncate(int, float)} except that it handles the case when index * is greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafeWithPossibleTruncate(int index, float value) { handleSafe(index); setWithPossibleTruncate(index, value); } /** * Same as {@link #set(int, NullableFloat2Holder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableFloat2Holder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, Float2Holder)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, Float2Holder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void set(int index, int isSet, short value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setWithPossibleTruncate(int index, int isSet, float value) { if (isSet > 0) { setWithPossibleTruncate(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, short)} except that it handles the case when index is greater * than or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setSafe(int index, int isSet, short value) { handleSafe(index); set(index, isSet, value); } /** * Same as {@link #set(int, int, short)} except that it handles the case when index is greater * than or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setSafeWithPossibleTruncate(int index, int isSet, float value) { handleSafe(index); setWithPossibleTruncate(index, isSet, value); } @Override public void setWithPossibleTruncate(int index, double value) { throw new UnsupportedOperationException( "The operation for double data types is not supported."); } @Override public void setSafeWithPossibleTruncate(int index, double value) { throw new UnsupportedOperationException( "The operation for double data types is not supported."); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((Float2Vector) to); } private class TransferImpl implements TransferPair { Float2Vector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new Float2Vector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new Float2Vector(field, allocator); } public TransferImpl(Float2Vector to) { this.to = to; } @Override public Float2Vector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, Float2Vector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/Float4Vector.java000066400000000000000000000263071515451615200306670ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.Float4ReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.Float4Holder; import org.apache.arrow.vector.holders.NullableFloat4Holder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * Float4Vector implements a fixed width vector (4 bytes) of float values which could be null. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public final class Float4Vector extends BaseFixedWidthVector implements FloatingPointVector, ValueIterableVector { public static final byte TYPE_WIDTH = 4; /** * Instantiate a Float4Vector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public Float4Vector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.FLOAT4.getType()), allocator); } /** * Instantiate a Float4Vector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public Float4Vector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a Float4Vector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public Float4Vector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new Float4ReaderImpl(Float4Vector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.FLOAT4; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public float get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getFloat((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableFloat4Holder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getFloat((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Float getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getFloat((long) index * TYPE_WIDTH); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, float value) { valueBuffer.setFloat((long) index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, float value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableFloat4Holder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, Float4Holder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, float)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, float value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableFloat4Holder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableFloat4Holder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, Float4Holder)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, Float4Holder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void set(int index, int isSet, float value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, float)} except that it handles the case when index is greater * than or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setSafe(int index, int isSet, float value) { handleSafe(index); set(index, isSet, value); } /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static float get(final ArrowBuf buffer, final int index) { return buffer.getFloat((long) index * TYPE_WIDTH); } @Override public void setWithPossibleTruncate(int index, double value) { set(index, (float) value); } @Override public void setSafeWithPossibleTruncate(int index, double value) { setSafe(index, (float) value); } @Override public double getValueAsDouble(int index) { return get(index); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((Float4Vector) to); } private class TransferImpl implements TransferPair { Float4Vector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new Float4Vector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new Float4Vector(field, allocator); } public TransferImpl(Float4Vector to) { this.to = to; } @Override public Float4Vector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, Float4Vector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/Float8Vector.java000066400000000000000000000263101515451615200306650ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.Float8ReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.Float8Holder; import org.apache.arrow.vector.holders.NullableFloat8Holder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * Float8Vector implements a fixed width vector (8 bytes) of double values which could be null. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public final class Float8Vector extends BaseFixedWidthVector implements FloatingPointVector, ValueIterableVector { public static final byte TYPE_WIDTH = 8; /** * Instantiate a Float8Vector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public Float8Vector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.FLOAT8.getType()), allocator); } /** * Instantiate a Float8Vector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public Float8Vector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a Float8Vector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public Float8Vector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new Float8ReaderImpl(Float8Vector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.FLOAT8; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public double get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getDouble((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableFloat8Holder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getDouble((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Double getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getDouble((long) index * TYPE_WIDTH); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, double value) { valueBuffer.setDouble((long) index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, double value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableFloat8Holder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, Float8Holder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, double)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, double value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableFloat8Holder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableFloat8Holder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, Float8Holder)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, Float8Holder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void set(int index, int isSet, double value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, double)} except that it handles the case when index is greater * than or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setSafe(int index, int isSet, double value) { handleSafe(index); set(index, isSet, value); } /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static double get(final ArrowBuf buffer, final int index) { return buffer.getDouble((long) index * TYPE_WIDTH); } @Override public void setWithPossibleTruncate(int index, double value) { set(index, value); } @Override public void setSafeWithPossibleTruncate(int index, double value) { setSafe(index, value); } @Override public double getValueAsDouble(int index) { return get(index); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((Float8Vector) to); } private class TransferImpl implements TransferPair { Float8Vector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new Float8Vector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new Float8Vector(field, allocator); } public TransferImpl(Float8Vector to) { this.to = to; } @Override public Float8Vector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, Float8Vector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/FloatingPointVector.java000066400000000000000000000032031515451615200323010ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** The interface for vectors with floating point values. */ public interface FloatingPointVector extends ValueVector { /** * Sets the value at the given index, note this value may be truncated internally. * * @param index the index to set. * @param value the value to set. */ void setWithPossibleTruncate(int index, double value); /** * Sets the value at the given index, note this value may be truncated internally. Any * expansion/reallocation is handled automatically. * * @param index the index to set. * @param value the value to set. */ void setSafeWithPossibleTruncate(int index, double value); /** * Gets the value at the given index. * * @param index the index to retrieve the value. * @return the value at the index. */ double getValueAsDouble(int index); } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/GenerateSampleData.java000066400000000000000000000311631515451615200320350ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.math.BigDecimal; import java.nio.charset.Charset; /** * Helper class to generate test data for Nullable fixed and variable width scalar vectors. Previous * implementations of java vector classes provided generateTestData(now deprecated) API to populate * the vector with sample data. This class should be used for that purpose. */ public class GenerateSampleData { private GenerateSampleData() {} /** Populates vector with valueCount random values. */ public static void generateTestData(final ValueVector vector, final int valueCount) { if (vector instanceof IntVector) { writeIntData((IntVector) vector, valueCount); } else if (vector instanceof DecimalVector) { writeDecimalData((DecimalVector) vector, valueCount); } else if (vector instanceof BitVector) { writeBooleanData((BitVector) vector, valueCount); } else if (vector instanceof VarCharVector) { writeVarCharData((VarCharVector) vector, valueCount); } else if (vector instanceof VarBinaryVector) { writeVarBinaryData((VarBinaryVector) vector, valueCount); } else if (vector instanceof BigIntVector) { writeBigIntData((BigIntVector) vector, valueCount); } else if (vector instanceof Float4Vector) { writeFloatData((Float4Vector) vector, valueCount); } else if (vector instanceof Float8Vector) { writeDoubleData((Float8Vector) vector, valueCount); } else if (vector instanceof DateDayVector) { writeDateDayData((DateDayVector) vector, valueCount); } else if (vector instanceof DateMilliVector) { writeDateMilliData((DateMilliVector) vector, valueCount); } else if (vector instanceof IntervalDayVector) { writeIntervalDayData((IntervalDayVector) vector, valueCount); } else if (vector instanceof IntervalYearVector) { writeIntervalYearData((IntervalYearVector) vector, valueCount); } else if (vector instanceof SmallIntVector) { writeSmallIntData((SmallIntVector) vector, valueCount); } else if (vector instanceof TinyIntVector) { writeTinyIntData((TinyIntVector) vector, valueCount); } else if (vector instanceof TimeMicroVector) { writeTimeMicroData((TimeMicroVector) vector, valueCount); } else if (vector instanceof TimeMilliVector) { writeTimeMilliData((TimeMilliVector) vector, valueCount); } else if (vector instanceof TimeNanoVector) { writeTimeNanoData((TimeNanoVector) vector, valueCount); } else if (vector instanceof TimeSecVector) { writeTimeSecData((TimeSecVector) vector, valueCount); } else if (vector instanceof TimeStampSecVector) { writeTimeStampData((TimeStampSecVector) vector, valueCount); } else if (vector instanceof TimeStampMicroVector) { writeTimeStampData((TimeStampMicroVector) vector, valueCount); } else if (vector instanceof TimeStampMilliVector) { writeTimeStampData((TimeStampMilliVector) vector, valueCount); } else if (vector instanceof TimeStampNanoVector) { writeTimeStampData((TimeStampNanoVector) vector, valueCount); } else if (vector instanceof TimeStampSecTZVector) { writeTimeStampData((TimeStampSecTZVector) vector, valueCount); } else if (vector instanceof TimeStampMicroTZVector) { writeTimeStampData((TimeStampMicroTZVector) vector, valueCount); } else if (vector instanceof TimeStampMilliTZVector) { writeTimeStampData((TimeStampMilliTZVector) vector, valueCount); } else if (vector instanceof TimeStampNanoTZVector) { writeTimeStampData((TimeStampNanoTZVector) vector, valueCount); } else if (vector instanceof UInt1Vector) { writeUInt1Data((UInt1Vector) vector, valueCount); } else if (vector instanceof UInt2Vector) { writeUInt2Data((UInt2Vector) vector, valueCount); } else if (vector instanceof UInt4Vector) { writeUInt4Data((UInt4Vector) vector, valueCount); } else if (vector instanceof UInt8Vector) { writeUInt8Data((UInt8Vector) vector, valueCount); } } private static void writeTimeStampData(TimeStampVector vector, int valueCount) { final long even = 100000; final long odd = 200000; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeDecimalData(DecimalVector vector, int valueCount) { final BigDecimal even = new BigDecimal("0.0543278923"); final BigDecimal odd = new BigDecimal("2.0543278923"); for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeIntData(IntVector vector, int valueCount) { final int even = 1000; final int odd = 2000; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeBooleanData(BitVector vector, int valueCount) { final int even = 0; final int odd = 1; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeIntervalYearData(IntervalYearVector vector, int valueCount) { final int even = 1; final int odd = 2; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeIntervalDayData(IntervalDayVector vector, int valueCount) { for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, 1, 50); } else { vector.setSafe(i, 2, 100); } } vector.setValueCount(valueCount); } private static void writeTimeSecData(TimeSecVector vector, int valueCount) { final int even = 500; final int odd = 900; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeTimeMilliData(TimeMilliVector vector, int valueCount) { final int even = 1000; final int odd = 2000; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeTimeMicroData(TimeMicroVector vector, int valueCount) { final long even = 1000000000; final long odd = 2000000000; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeTimeNanoData(TimeNanoVector vector, int valueCount) { final long even = 1000000000; final long odd = 2000000000; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeDateDayData(DateDayVector vector, int valueCount) { final int even = 1000; final int odd = 2000; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeDateMilliData(DateMilliVector vector, int valueCount) { final long even = 1000000000; final long odd = 2000000000; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeSmallIntData(SmallIntVector vector, int valueCount) { final short even = 10; final short odd = 20; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeTinyIntData(TinyIntVector vector, int valueCount) { final byte even = 1; final byte odd = 2; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeUInt1Data(UInt1Vector vector, int valueCount) { final byte even = 1; final byte odd = 2; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeUInt2Data(UInt2Vector vector, int valueCount) { final short even = 10; final short odd = 20; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeUInt4Data(UInt4Vector vector, int valueCount) { final int even = 1000; final int odd = 2000; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeUInt8Data(UInt8Vector vector, int valueCount) { final long even = 1000000000; final long odd = 2000000000; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeBigIntData(BigIntVector vector, int valueCount) { final long even = 1000000000; final long odd = 2000000000; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeFloatData(Float4Vector vector, int valueCount) { final float even = 20.3f; final float odd = 40.2f; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeDoubleData(Float8Vector vector, int valueCount) { final double even = 20.2373; final double odd = 40.2378; for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeVarBinaryData(VarBinaryVector vector, int valueCount) { Charset utf8Charset = Charset.forName("UTF-8"); final byte[] even = "AAAAA1".getBytes(utf8Charset); final byte[] odd = "BBBBBBBBB2".getBytes(utf8Charset); for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } private static void writeVarCharData(VarCharVector vector, int valueCount) { Charset utf8Charset = Charset.forName("UTF-8"); final byte[] even = "AAAAA1".getBytes(utf8Charset); final byte[] odd = "BBBBBBBBB2".getBytes(utf8Charset); for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/IntVector.java000066400000000000000000000261311515451615200302630ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.IntReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.IntHolder; import org.apache.arrow.vector.holders.NullableIntHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * IntVector implements a fixed width (4 bytes) vector of integer values which could be null. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public final class IntVector extends BaseFixedWidthVector implements BaseIntVector, ValueIterableVector { public static final byte TYPE_WIDTH = 4; /** * Instantiate a IntVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public IntVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.INT.getType()), allocator); } /** * Instantiate a IntVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public IntVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a IntVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public IntVector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new IntReaderImpl(IntVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.INT; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public int get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getInt((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableIntHolder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getInt((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Integer getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getInt((long) index * TYPE_WIDTH); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, int value) { valueBuffer.setInt((long) index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, int value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableIntHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, IntHolder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, int)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, int value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableIntHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableIntHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, IntHolder)} except that it handles the case when index is greater than * or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, IntHolder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void set(int index, int isSet, int value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, int)} except that it handles the case when index is greater than * or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setSafe(int index, int isSet, int value) { handleSafe(index); set(index, isSet, value); } /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static int get(final ArrowBuf buffer, final int index) { return buffer.getInt((long) index * TYPE_WIDTH); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((IntVector) to); } @Override public void setWithPossibleTruncate(int index, long value) { this.setSafe(index, (int) value); } @Override public void setUnsafeWithPossibleTruncate(int index, long value) { this.set(index, (int) value); } @Override public long getValueAsLong(int index) { return this.get(index); } private class TransferImpl implements TransferPair { IntVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new IntVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new IntVector(field, allocator); } public TransferImpl(IntVector to) { this.to = to; } @Override public IntVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, IntVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/IntervalDayVector.java000066400000000000000000000350371515451615200317600ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import java.time.Duration; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.IntervalDayReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.IntervalDayHolder; import org.apache.arrow.vector.holders.NullableIntervalDayHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * IntervalDayVector implements a fixed width vector (8 bytes) of interval (days and milliseconds) * values which could be null. A validity buffer (bit vector) is maintained to track which elements * in the vector are null. */ public final class IntervalDayVector extends BaseFixedWidthVector implements ValueIterableVector { public static final byte TYPE_WIDTH = 8; private static final byte MILLISECOND_OFFSET = 4; /** * Instantiate a IntervalDayVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public IntervalDayVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.INTERVALDAY.getType()), allocator); } /** * Instantiate a IntervalDayVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public IntervalDayVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a IntervalDayVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public IntervalDayVector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new IntervalDayReaderImpl(IntervalDayVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.INTERVALDAY; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Given a data buffer, get the number of days stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return day value stored at the index. */ public static int getDays(final ArrowBuf buffer, final int index) { return buffer.getInt((long) index * TYPE_WIDTH); } /** * Given a data buffer, get the get the number of milliseconds stored at a particular position in * the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return milliseconds value stored at the index. */ public static int getMilliseconds(final ArrowBuf buffer, final int index) { return buffer.getInt((long) index * TYPE_WIDTH + MILLISECOND_OFFSET); } /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public ArrowBuf get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { return null; } return valueBuffer.slice((long) index * TYPE_WIDTH, TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableIntervalDayHolder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } final long startIndex = (long) index * TYPE_WIDTH; holder.isSet = 1; holder.days = valueBuffer.getInt(startIndex); holder.milliseconds = valueBuffer.getInt(startIndex + MILLISECOND_OFFSET); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Duration getObject(int index) { if (isSet(index) == 0) { return null; } else { return getObjectNotNull(index); } } /** * Same as {@link #getObject(int)} but does not check for null. * * @param index position of element * @return element at given index */ public Duration getObjectNotNull(int index) { final long startIndex = (long) index * TYPE_WIDTH; final int days = valueBuffer.getInt(startIndex); final int milliseconds = valueBuffer.getInt(startIndex + MILLISECOND_OFFSET); return Duration.ofDays(days).plusMillis(milliseconds); } /** * Get the Interval value at a given index as a {@link StringBuilder} object. * * @param index position of the element * @return String Builder object with Interval value as [days, hours, minutes, seconds, millis] */ public StringBuilder getAsStringBuilder(int index) { if (isSet(index) == 0) { return null; } else { return getAsStringBuilderHelper(index); } } private StringBuilder getAsStringBuilderHelper(int index) { final long startIndex = (long) index * TYPE_WIDTH; final int days = valueBuffer.getInt(startIndex); int millis = valueBuffer.getInt(startIndex + MILLISECOND_OFFSET); final int hours = millis / org.apache.arrow.vector.util.DateUtility.hoursToMillis; millis = millis % org.apache.arrow.vector.util.DateUtility.hoursToMillis; final int minutes = millis / org.apache.arrow.vector.util.DateUtility.minutesToMillis; millis = millis % org.apache.arrow.vector.util.DateUtility.minutesToMillis; final int seconds = millis / org.apache.arrow.vector.util.DateUtility.secondsToMillis; millis = millis % org.apache.arrow.vector.util.DateUtility.secondsToMillis; final String dayString = (Math.abs(days) == 1) ? " day " : " days "; return new StringBuilder() .append(days) .append(dayString) .append(hours) .append(":") .append(minutes) .append(":") .append(seconds) .append(".") .append(millis); } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, ArrowBuf value) { BitVectorHelper.setBit(validityBuffer, index); valueBuffer.setBytes((long) index * TYPE_WIDTH, value, 0, TYPE_WIDTH); } /** * Set the element at the given index to the given value. * * @param index position of element * @param days days for the interval * @param milliseconds milliseconds for the interval */ public void set(int index, int days, int milliseconds) { final long offsetIndex = (long) index * TYPE_WIDTH; BitVectorHelper.setBit(validityBuffer, index); valueBuffer.setInt(offsetIndex, days); valueBuffer.setInt((offsetIndex + MILLISECOND_OFFSET), milliseconds); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableIntervalDayHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { set(index, holder.days, holder.milliseconds); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, IntervalDayHolder holder) { set(index, holder.days, holder.milliseconds); } /** * Same as {@link #set(int, ArrowBuf)} except that it handles the case when index is greater than * or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, ArrowBuf value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, int, int)} except that it handles the case when index is greater than * or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param days days for the interval * @param milliseconds milliseconds for the interval */ public void setSafe(int index, int days, int milliseconds) { handleSafe(index); set(index, days, milliseconds); } /** * Same as {@link #set(int, NullableIntervalDayHolder)} except that it handles the case when index * is greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableIntervalDayHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, IntervalDayHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, IntervalDayHolder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param days days component of interval * @param milliseconds millisecond component of interval */ public void set(int index, int isSet, int days, int milliseconds) { if (isSet > 0) { set(index, days, milliseconds); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, int, int)} except that it handles the case when index is greater * than or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param days days component of interval * @param milliseconds millisecond component of interval */ public void setSafe(int index, int isSet, int days, int milliseconds) { handleSafe(index); set(index, isSet, days, milliseconds); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((IntervalDayVector) to); } private class TransferImpl implements TransferPair { IntervalDayVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new IntervalDayVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new IntervalDayVector(field, allocator); } public TransferImpl(IntervalDayVector to) { this.to = to; } @Override public IntervalDayVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, IntervalDayVector.this); } } } IntervalMonthDayNanoVector.java000066400000000000000000000354541515451615200335260ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import java.time.Duration; import java.time.Period; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.IntervalMonthDayNanoReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.IntervalMonthDayNanoHolder; import org.apache.arrow.vector.holders.NullableIntervalMonthDayNanoHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * IntervalMonthDayNanoVector implements a fixed width vector (16 bytes) of interval (month, days * and nanoseconds) values which could be null. A validity buffer (bit vector) is maintained to * track which elements in the vector are null. * *

    Month, day and nanoseconds are independent from one another and there is no specific limits * imposed on their values. */ public final class IntervalMonthDayNanoVector extends BaseFixedWidthVector implements ValueIterableVector { public static final byte TYPE_WIDTH = 16; private static final byte DAY_OFFSET = 4; private static final byte NANOSECOND_OFFSET = 8; /** * Instantiate a IntervalMonthDayNanoVector. This doesn't allocate any memory for the data in * vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public IntervalMonthDayNanoVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.INTERVALMONTHDAYNANO.getType()), allocator); } /** * Instantiate a IntervalMonthDayNanoVector. This doesn't allocate any memory for the data in * vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public IntervalMonthDayNanoVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a IntervalMonthDayNanoVector. This doesn't allocate any memory for the data in * vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public IntervalMonthDayNanoVector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new IntervalMonthDayNanoReaderImpl(IntervalMonthDayNanoVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.INTERVALMONTHDAYNANO; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Given a data buffer, get the number of months stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return day value stored at the index. */ public static int getMonths(final ArrowBuf buffer, final int index) { return buffer.getInt((long) index * TYPE_WIDTH); } /** * Given a data buffer, get the number of days stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return day value stored at the index. */ public static int getDays(final ArrowBuf buffer, final int index) { return buffer.getInt((long) index * TYPE_WIDTH + DAY_OFFSET); } /** * Given a data buffer, get the get the number of nanoseconds stored at a particular position in * the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return nanoseconds value stored at the index. */ public static long getNanoseconds(final ArrowBuf buffer, final int index) { return buffer.getLong((long) index * TYPE_WIDTH + NANOSECOND_OFFSET); } /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public ArrowBuf get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { return null; } return valueBuffer.slice((long) index * TYPE_WIDTH, TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableIntervalMonthDayNanoHolder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } final long startIndex = (long) index * TYPE_WIDTH; holder.isSet = 1; holder.months = valueBuffer.getInt(startIndex); holder.days = valueBuffer.getInt(startIndex + DAY_OFFSET); holder.nanoseconds = valueBuffer.getLong(startIndex + NANOSECOND_OFFSET); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public PeriodDuration getObject(int index) { if (isSet(index) == 0) { return null; } else { final long startIndex = (long) index * TYPE_WIDTH; final int months = valueBuffer.getInt(startIndex); final int days = valueBuffer.getInt(startIndex + DAY_OFFSET); final long nanoseconds = valueBuffer.getLong(startIndex + NANOSECOND_OFFSET); return new PeriodDuration( Period.ofMonths(months).plusDays(days), Duration.ofNanos(nanoseconds)); } } /** * Get the Interval value at a given index as a {@link StringBuilder} object. * * @param index position of the element * @return String Builder object with Interval value as */ public StringBuilder getAsStringBuilder(int index) { if (isSet(index) == 0) { return null; } else { return getAsStringBuilderHelper(index); } } private StringBuilder getAsStringBuilderHelper(int index) { return new StringBuilder().append(getObject(index).toString()).append(" "); } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, ArrowBuf value) { BitVectorHelper.setBit(validityBuffer, index); valueBuffer.setBytes((long) index * TYPE_WIDTH, value, 0, TYPE_WIDTH); } /** * Set the element at the given index to the given value. * * @param index position of element * @param months months component of interval * @param days days component of interval * @param nanoseconds nanosecond component of interval */ public void set(int index, int months, int days, long nanoseconds) { final long offsetIndex = (long) index * TYPE_WIDTH; BitVectorHelper.setBit(validityBuffer, index); valueBuffer.setInt(offsetIndex, months); valueBuffer.setInt(offsetIndex + DAY_OFFSET, days); valueBuffer.setLong((offsetIndex + NANOSECOND_OFFSET), nanoseconds); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableIntervalMonthDayNanoHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { set(index, holder.months, holder.days, holder.nanoseconds); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, IntervalMonthDayNanoHolder holder) { set(index, holder.months, holder.days, holder.nanoseconds); } /** * Same as {@link #set(int, ArrowBuf)} except that it handles the case when index is greater than * or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, ArrowBuf value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, int, int, long)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param months months for the interval * @param days days for the interval * @param nanoseconds nanoseconds for the interval */ public void setSafe(int index, int months, int days, long nanoseconds) { handleSafe(index); set(index, months, days, nanoseconds); } /** * Same as {@link #set(int, NullableIntervalMonthDayNanoHolder)} except that it handles the case * when index is greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableIntervalMonthDayNanoHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, IntervalMonthDayNanoHolder)} except that it handles the case when * index is greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, IntervalMonthDayNanoHolder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param months months component of interval * @param days days component of interval * @param nanoseconds nanosecond component of interval */ public void set(int index, int isSet, int months, int days, long nanoseconds) { if (isSet > 0) { set(index, months, days, nanoseconds); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, int, int, long)} except that it handles the case when index is * greater than or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param months months component of interval * @param days days component of interval * @param nanoseconds nanosecond component of interval */ public void setSafe(int index, int isSet, int months, int days, long nanoseconds) { handleSafe(index); set(index, isSet, months, days, nanoseconds); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising of this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((IntervalMonthDayNanoVector) to); } private class TransferImpl implements TransferPair { IntervalMonthDayNanoVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new IntervalMonthDayNanoVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new IntervalMonthDayNanoVector(field, allocator); } public TransferImpl(IntervalMonthDayNanoVector to) { this.to = to; } @Override public IntervalMonthDayNanoVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, IntervalMonthDayNanoVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/IntervalYearVector.java000066400000000000000000000301451515451615200321360ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import java.time.Period; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.IntervalYearReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.IntervalYearHolder; import org.apache.arrow.vector.holders.NullableIntervalYearHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * IntervalYearVector implements a fixed width (4 bytes) vector of interval (years and months) * values which could be null. A validity buffer (bit vector) is maintained to track which elements * in the vector are null. */ public final class IntervalYearVector extends BaseFixedWidthVector implements ValueIterableVector { public static final byte TYPE_WIDTH = 4; /** * Instantiate a IntervalYearVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public IntervalYearVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.INTERVALYEAR.getType()), allocator); } /** * Instantiate a IntervalYearVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public IntervalYearVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a IntervalYearVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public IntervalYearVector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new IntervalYearReaderImpl(IntervalYearVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.INTERVALYEAR; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static int getTotalMonths(final ArrowBuf buffer, final int index) { return buffer.getInt((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public int get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getInt((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableIntervalYearHolder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getInt((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Period getObject(int index) { if (isSet(index) == 0) { return null; } else { final int interval = valueBuffer.getInt((long) index * TYPE_WIDTH); // TODO: verify interval is in months return Period.ofMonths(interval); } } /** * Get the Interval value at a given index as a {@link StringBuilder} object. * * @param index position of the element * @return String Builder object with Interval value as [years, months] */ public StringBuilder getAsStringBuilder(int index) { if (isSet(index) == 0) { return null; } else { return getAsStringBuilderHelper(index); } } private StringBuilder getAsStringBuilderHelper(int index) { int value = valueBuffer.getInt((long) index * TYPE_WIDTH); final int years = (value / org.apache.arrow.vector.util.DateUtility.yearsToMonths); final int months = (value % org.apache.arrow.vector.util.DateUtility.yearsToMonths); final String yearString = (Math.abs(years) == 1) ? " year " : " years "; final String monthString = (Math.abs(months) == 1) ? " month " : " months "; return new StringBuilder().append(years).append(yearString).append(months).append(monthString); } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, int value) { valueBuffer.setInt((long) index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, int value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableIntervalYearHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, IntervalYearHolder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, int)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, int value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableIntervalYearHolder)} except that it handles the case when * index is greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableIntervalYearHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, IntervalYearHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, IntervalYearHolder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void set(int index, int isSet, int value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, int)} except that it handles the case when index is greater than * or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setSafe(int index, int isSet, int value) { handleSafe(index); set(index, isSet, value); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((IntervalYearVector) to); } private class TransferImpl implements TransferPair { IntervalYearVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new IntervalYearVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new IntervalYearVector(field, allocator); } public TransferImpl(IntervalYearVector to) { this.to = to; } @Override public IntervalYearVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, IntervalYearVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/LargeVarBinaryVector.java000066400000000000000000000263671515451615200324140ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReusableBuffer; import org.apache.arrow.vector.complex.impl.LargeVarBinaryReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.LargeVarBinaryHolder; import org.apache.arrow.vector.holders.NullableLargeVarBinaryHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * LargeVarBinaryVector implements a large variable width vector of binary values which could be * NULL. A validity buffer (bit vector) is maintained to track which elements in the vector are * null. The size of the underlying buffer can be over 2GB. */ public final class LargeVarBinaryVector extends BaseLargeVariableWidthVector implements ValueIterableVector { /** * Instantiate a LargeVarBinaryVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public LargeVarBinaryVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.LARGEVARBINARY.getType()), allocator); } /** * Instantiate a LargeVarBinaryVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public LargeVarBinaryVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a LargeVarBinaryVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public LargeVarBinaryVector(Field field, BufferAllocator allocator) { super(field, allocator); } @Override protected FieldReader getReaderImpl() { return new LargeVarBinaryReaderImpl(LargeVarBinaryVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.LARGEVARBINARY; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the variable length element at specified index as byte array. * * @param index position of element to get * @return array of bytes for non-null element, null otherwise */ public byte[] get(int index) { assert index >= 0; if (NULL_CHECKING_ENABLED && 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; } /** * Read the value at the given position to the given output buffer. The caller is responsible for * checking for nullity first. * * @param index position of element. * @param buffer the buffer to write into. */ @Override public void read(int index, ReusableBuffer buffer) { final long startOffset = getStartOffset(index); final long dataLength = getEndOffset(index) - startOffset; buffer.set(valueBuffer, startOffset, dataLength); } /** * Get the variable length element at specified index as Text. * * @param index position of element to get * @return byte array for non-null element, null otherwise */ @Override public byte[] getObject(int 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; } /** * Get the variable length element at specified index and sets the state in provided holder. * * @param index position of element to get * @param holder data holder to be populated by this function */ public void get(int index, NullableLargeVarBinaryHolder holder) { assert index >= 0; if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.start = getStartOffset(index); holder.end = getEndOffset(index); holder.buffer = valueBuffer; } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the variable length element at the specified index to the data buffer supplied in the * holder. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void set(int index, LargeVarBinaryHolder holder) { assert index >= 0; fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); final int dataLength = (int) (holder.end - holder.start); final long startOffset = getStartOffset(index); offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, startOffset + dataLength); valueBuffer.setBytes(startOffset, holder.buffer, holder.start, dataLength); lastSet = index; } /** * Same as {@link #set(int, LargeVarBinaryHolder)} except that it handles the case where index and * length of new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void setSafe(int index, LargeVarBinaryHolder holder) { assert index >= 0; final int dataLength = (int) (holder.end - holder.start); handleSafe(index, dataLength); fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); final long startOffset = getStartOffset(index); offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, startOffset + dataLength); valueBuffer.setBytes(startOffset, holder.buffer, holder.start, dataLength); lastSet = index; } /** * Set the variable length element at the specified index to the data buffer supplied in the * holder. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void set(int index, NullableLargeVarBinaryHolder holder) { assert index >= 0; fillHoles(index); BitVectorHelper.setValidityBit(validityBuffer, index, holder.isSet); final long startOffset = getStartOffset(index); if (holder.isSet != 0) { final int dataLength = (int) (holder.end - holder.start); offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, startOffset + dataLength); valueBuffer.setBytes(startOffset, holder.buffer, holder.start, dataLength); } else { offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, startOffset); } lastSet = index; } /** * Same as {@link #set(int, NullableLargeVarBinaryHolder)} except that it handles the case where * index and length of new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void setSafe(int index, NullableLargeVarBinaryHolder holder) { assert index >= 0; if (holder.isSet != 0) { final int dataLength = (int) (holder.end - holder.start); handleSafe(index, dataLength); fillHoles(index); final long startOffset = getStartOffset(index); offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, startOffset + dataLength); valueBuffer.setBytes(startOffset, holder.buffer, holder.start, dataLength); } else { fillEmpties(index + 1); } BitVectorHelper.setValidityBit(validityBuffer, index, holder.isSet); lastSet = index; } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising of this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((LargeVarBinaryVector) to); } private class TransferImpl implements TransferPair { LargeVarBinaryVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new LargeVarBinaryVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new LargeVarBinaryVector(field, allocator); } public TransferImpl(LargeVarBinaryVector to) { this.to = to; } @Override public LargeVarBinaryVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, LargeVarBinaryVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/LargeVarCharVector.java000066400000000000000000000303661515451615200320370ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReusableBuffer; import org.apache.arrow.vector.complex.impl.LargeVarCharReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.LargeVarCharHolder; import org.apache.arrow.vector.holders.NullableLargeVarCharHolder; import org.apache.arrow.vector.types.Types; 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; import org.apache.arrow.vector.validate.ValidateUtil; /** * LargeVarCharVector implements a variable width vector of VARCHAR values which could be NULL. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. * *

    The offset width of this vector is 8, so the underlying buffer can be larger than 2GB. */ public final class LargeVarCharVector extends BaseLargeVariableWidthVector implements ValueIterableVector { /** * Instantiate a LargeVarCharVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public LargeVarCharVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(Types.MinorType.LARGEVARCHAR.getType()), allocator); } /** * Instantiate a LargeVarCharVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public LargeVarCharVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a LargeVarCharVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public LargeVarCharVector(Field field, BufferAllocator allocator) { super(field, allocator); } @Override protected FieldReader getReaderImpl() { return new LargeVarCharReaderImpl(LargeVarCharVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public Types.MinorType getMinorType() { return Types.MinorType.LARGEVARCHAR; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the variable length element at specified index as byte array. * * @param index position of element to get * @return array of bytes for non-null element, null otherwise */ @Override public byte[] get(int index) { assert index >= 0; if (NULL_CHECKING_ENABLED && 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; } /** * Get the variable length element at specified index as Text. * * @param index position of element to get * @return Text object for non-null element, null otherwise */ @Override public Text getObject(int index) { assert index >= 0; if (isSet(index) == 0) { return null; } final Text result = new Text(); read(index, result); return result; } /** * Read the value at the given position to the given output buffer. The caller is responsible for * checking for nullity first. * * @param index position of element. * @param buffer the buffer to write into. */ @Override public void read(int index, ReusableBuffer buffer) { final long startOffset = getStartOffset(index); final long dataLength = getEndOffset(index) - startOffset; buffer.set(valueBuffer, startOffset, dataLength); } /** * Get the variable length element at specified index and sets the state in provided holder. * * @param index position of element to get * @param holder data holder to be populated by this function */ public void get(int index, NullableLargeVarCharHolder holder) { assert index >= 0; if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.start = getStartOffset(index); holder.end = getEndOffset(index); holder.buffer = valueBuffer; } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the variable length element at the specified index to the data buffer supplied in the * holder. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void set(int index, LargeVarCharHolder holder) { assert index >= 0; fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); final int dataLength = (int) (holder.end - holder.start); final long startOffset = getStartOffset(index); offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, startOffset + dataLength); valueBuffer.setBytes(startOffset, holder.buffer, holder.start, dataLength); lastSet = index; } /** * Same as {@link #set(int, LargeVarCharHolder)} except that it handles the case where index and * length of new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void setSafe(int index, LargeVarCharHolder holder) { assert index >= 0; final int dataLength = (int) (holder.end - holder.start); handleSafe(index, dataLength); fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); final long startOffset = getStartOffset(index); offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, startOffset + dataLength); valueBuffer.setBytes(startOffset, holder.buffer, holder.start, dataLength); lastSet = index; } /** * Set the variable length element at the specified index to the data buffer supplied in the * holder. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void set(int index, NullableLargeVarCharHolder holder) { assert index >= 0; fillHoles(index); BitVectorHelper.setValidityBit(validityBuffer, index, holder.isSet); final long startOffset = getStartOffset(index); if (holder.isSet != 0) { final int dataLength = (int) (holder.end - holder.start); offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, startOffset + dataLength); valueBuffer.setBytes(startOffset, holder.buffer, holder.start, dataLength); } else { offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, startOffset); } lastSet = index; } /** * Same as {@link #set(int, NullableLargeVarCharHolder)} except that it handles the case where * index and length of new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void setSafe(int index, NullableLargeVarCharHolder holder) { assert index >= 0; if (holder.isSet != 0) { final int dataLength = (int) (holder.end - holder.start); handleSafe(index, dataLength); fillHoles(index); final long startOffset = getStartOffset(index); offsetBuffer.setLong((long) (index + 1) * OFFSET_WIDTH, startOffset + dataLength); valueBuffer.setBytes(startOffset, holder.buffer, holder.start, dataLength); } else { fillHoles(index + 1); } BitVectorHelper.setValidityBit(validityBuffer, index, holder.isSet); lastSet = index; } /** * Set the variable length element at the specified index to the content in supplied Text. * * @param index position of the element to set * @param text Text object with data */ public void set(int index, Text text) { set(index, text.getBytes(), 0, (int) text.getLength()); } /** * Same as {@link #set(int, NullableLargeVarCharHolder)} except that it handles the case where * index and length of new element are beyond the existing capacity of the vector. * * @param index position of the element to set. * @param text Text object with data */ public void setSafe(int index, Text text) { setSafe(index, text.getBytes(), 0, (int) text.getLength()); } @Override public void validateScalars() { for (int i = 0; i < getValueCount(); ++i) { byte[] value = get(i); if (value != null) { ValidateUtil.validateOrThrow( Text.validateUTF8NoThrow(value), "Non-UTF-8 data in VarCharVector at position " + i + "."); } } } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new LargeVarCharVector.TransferImpl(ref, allocator); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new LargeVarCharVector.TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new LargeVarCharVector.TransferImpl((LargeVarCharVector) to); } private class TransferImpl implements TransferPair { LargeVarCharVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new LargeVarCharVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new LargeVarCharVector(field, allocator); } public TransferImpl(LargeVarCharVector to) { this.to = to; } @Override public LargeVarCharVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, LargeVarCharVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/NullCheckingForGet.java000066400000000000000000000061351515451615200320250ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * Configuration class to determine if null checking should be enabled or disabled for the "get" * methods. For example, the get method of class org.apache.arrow.vector.Float8Vector first checks * if the value at the given index is null, before retrieving the value. This configuration will * turn on and off such checks. * *

    Null checking is on by default. You can disable it by setting either the system property or * the environmental variable to "false". The system property is named * "arrow.enable_null_check_for_get" and the environmental variable is named * "ARROW_ENABLE_NULL_CHECK_FOR_GET". When both the system property and the environmental variable * are set, the system property takes precedence. * *

    Disabling null-checking in the "get" methods may lead to performance improvements. For * example, suppose we have the following micro-benchmark: * *

    * * 

    {@code
 * Float8Vector vector = ...
 *
 * public void test() {
 *   sum = 0;
 *   for (int i = 0; i < 1024; i++) {
 *     vector.set(i, i + 10.0);
 *     safeSum += vector.get(i);
 *   }
 * }
 *
 * }
    * *

    Performance evaluations of the micro-benchmark with the JMH framework reveal that, disabling * null checking has the following effects: 1. The amounts of byte code and assembly code generated * by JIT are both smaller. 2. The performance improves by about 30% (2.819 ± 0.005 us/op vs. 4.069 * ± 0.004 us/op). * *

    Therefore, for scenarios where the user can be sure that the null-checking is unnecessary, it * is beneficial to disable it with this configuration. */ public class NullCheckingForGet { /** The flag to indicate if null checking is enabled for "get" methods. */ public static final boolean NULL_CHECKING_ENABLED; static { String envProperty = System.getenv("ARROW_ENABLE_NULL_CHECK_FOR_GET"); String sysProperty = System.getProperty("arrow.enable_null_check_for_get"); // The system property has a higher priority than the environmental variable. String flagValue = sysProperty; if (flagValue == null) { flagValue = envProperty; } // The flag is set to false only if the system property/environmental // variable is explicitly set to "false". NULL_CHECKING_ENABLED = !"false".equals(flagValue); } private NullCheckingForGet() {} } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/NullVector.java000066400000000000000000000216511515451615200304450ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex.BaseRepeatedValueVector.DATA_VECTOR_NAME; import java.util.Collections; import java.util.Iterator; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.OutOfMemoryException; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.impl.NullReader; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; 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.util.CallBack; import org.apache.arrow.vector.util.TransferPair; /** A null type vector. */ public class NullVector implements FieldVector, ValueIterableVector { private int valueCount; protected Field field; /** * Instantiate a NullVector. * * @param name name of the vector */ public NullVector(String name) { this(name, FieldType.nullable(Types.MinorType.NULL.getType())); } /** * Instantiate a NullVector with the given number of values. * * @param name name of the vector * @param valueCount number of values (i.e., nulls) in this vector. */ public NullVector(String name, int valueCount) { this(new Field(name, FieldType.nullable(Types.MinorType.NULL.getType()), null), valueCount); } /** * Instantiate a NullVector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector. */ public NullVector(String name, FieldType fieldType) { this(new Field(name, fieldType, null)); } /** * Instantiate a NullVector. * * @param field field materialized by this vector. */ public NullVector(Field field) { this(field, 0); } /** * Instantiate a NullVector with the given number of values. * * @param field field materialized by this vector. * @param valueCount number of values (i.e., nulls) in this vector. */ public NullVector(Field field, int valueCount) { this.field = field; this.valueCount = valueCount; } @Deprecated public NullVector() { this(new Field(DATA_VECTOR_NAME, FieldType.nullable(new ArrowType.Null()), null)); } @Override public void close() {} @Override public void clear() {} @Override public void reset() {} @Override public Field getField() { return field; } @Override public Types.MinorType getMinorType() { return Types.MinorType.NULL; } @Override public TransferPair getTransferPair(BufferAllocator allocator) { return getTransferPair(getName(), allocator); } @Override public Iterator iterator() { return Collections.emptyIterator(); } @Override public int getBufferSize() { return 0; } @Override public int getBufferSizeFor(final int valueCount) { return 0; } @Override public ArrowBuf[] getBuffers(boolean clear) { return new ArrowBuf[0]; } @Override public void allocateNew() throws OutOfMemoryException { allocateNewSafe(); } @Override public boolean allocateNewSafe() { return true; } @Override public void reAlloc() {} /* * IMPORTANT NOTE * It's essential that NullVector (and ZeroVector) do not require BufferAllocator for any data storage. * However, some methods of the parent interface may require passing in a BufferAllocator, even if null. * * @return null */ @Override public BufferAllocator getAllocator() { return null; } @Override public void setInitialCapacity(int numRecords) {} @Override public int getValueCapacity() { return this.valueCount; } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field.getName()); } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return getTransferPair(ref, allocator); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { return getTransferPair(field, allocator); } @Override public TransferPair makeTransferPair(ValueVector target) { return new TransferImpl((NullVector) target); } @Override public FieldReader getReader() { return NullReader.INSTANCE; } @Override public void initializeChildrenFromFields(List children) { if (!children.isEmpty()) { throw new IllegalArgumentException("Null vector has no children"); } } @Override public List getChildrenFromFields() { return Collections.emptyList(); } @Override public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers) { Preconditions.checkArgument(ownBuffers.isEmpty(), "Null vector has no buffers"); valueCount = fieldNode.getLength(); } @Override public List getFieldBuffers() { return Collections.emptyList(); } /** * Get the inner vectors. * * @return the inner vectors for this field as defined by the TypeLayout * @deprecated This API will be removed as the current implementations no longer support inner * vectors. */ @Deprecated @Override public List getFieldInnerVectors() { return Collections.emptyList(); } @Override public long getValidityBufferAddress() { throw new UnsupportedOperationException(); } @Override public long getDataBufferAddress() { throw new UnsupportedOperationException(); } @Override public long getOffsetBufferAddress() { throw new UnsupportedOperationException(); } @Override public ArrowBuf getValidityBuffer() { throw new UnsupportedOperationException(); } @Override public ArrowBuf getDataBuffer() { throw new UnsupportedOperationException(); } @Override public ArrowBuf getOffsetBuffer() { throw new UnsupportedOperationException(); } @Override public int getValueCount() { return this.valueCount; } @Override public void setValueCount(int valueCount) { this.valueCount = valueCount; } @Override public Object getObject(int index) { return null; } @Override public int getNullCount() { return this.valueCount; } /** * Set the element at the given index to null. In a NullVector, this is a no-op. * * @param index position of element */ @Override public void setNull(int index) {} @Override public boolean isNull(int index) { return true; } @Override public int hashCode(int index) { return 31; } @Override public int hashCode(int index, ArrowBufHasher hasher) { return 31; } @Override public OUT accept(VectorVisitor visitor, IN value) { return visitor.visit(this, value); } @Override public void copyFrom(int fromIndex, int thisIndex, ValueVector from) { throw new UnsupportedOperationException(); } @Override public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) { throw new UnsupportedOperationException(); } @Override public String getName() { return this.getField().getName(); } private class TransferImpl implements TransferPair { NullVector to; public TransferImpl(String ref) { to = new NullVector(ref); } @Deprecated public TransferImpl() { to = new NullVector(); } public TransferImpl(NullVector to) { this.to = to; } @Override public NullVector getTo() { return to; } @Override public void transfer() { to.valueCount = valueCount; } @Override public void splitAndTransfer(int startIndex, int length) { to.valueCount = length; } @Override public void copyValueSafe(int fromIndex, int toIndex) { if (toIndex > to.valueCount) { to.valueCount = toIndex; } } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/PeriodDuration.java000066400000000000000000000076501515451615200313030ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.time.temporal.ChronoUnit.DAYS; import static java.time.temporal.ChronoUnit.MONTHS; import static java.time.temporal.ChronoUnit.NANOS; import static java.time.temporal.ChronoUnit.SECONDS; import static java.time.temporal.ChronoUnit.YEARS; import java.time.Duration; import java.time.Period; import java.time.temporal.ChronoUnit; import java.time.temporal.Temporal; import java.time.temporal.TemporalAmount; import java.time.temporal.TemporalUnit; import java.time.temporal.UnsupportedTemporalTypeException; import java.util.Arrays; import java.util.Collections; import java.util.List; import org.apache.arrow.util.Preconditions; /** Combination of Period and Duration for representing this interval type as a POJO. */ public class PeriodDuration implements TemporalAmount { private static final List SUPPORTED_UNITS = Collections.unmodifiableList( Arrays.asList(YEARS, MONTHS, DAYS, SECONDS, NANOS)); private final Period period; private final Duration duration; public PeriodDuration(Period period, Duration duration) { this.period = Preconditions.checkNotNull(period); this.duration = Preconditions.checkNotNull(duration); } public Period getPeriod() { return period; } public Duration getDuration() { return duration; } @Override public long get(TemporalUnit unit) { if (unit instanceof ChronoUnit) { switch ((ChronoUnit) unit) { case YEARS: return period.getYears(); case MONTHS: return period.getMonths(); case DAYS: return period.getDays(); case SECONDS: return duration.getSeconds(); case NANOS: return duration.getNano(); default: break; } } throw new UnsupportedTemporalTypeException("Unsupported TemporalUnit: " + unit); } @Override public List getUnits() { return SUPPORTED_UNITS; } @Override public Temporal addTo(Temporal temporal) { return temporal.plus(period).plus(duration); } @Override public Temporal subtractFrom(Temporal temporal) { return temporal.minus(period).minus(duration); } /** * Format this PeriodDuration as an ISO-8601 interval. * * @return An ISO-8601 formatted string representing the interval. */ public String toISO8601IntervalString() { if (duration.isZero()) { return period.toString(); } String durationString = duration.toString(); if (period.isZero()) { return durationString; } // Remove 'P' from duration string and concatenate to produce an ISO-8601 representation return period + durationString.substring(1); } @Override public String toString() { return period.toString() + " " + duration.toString(); } @Override public boolean equals(Object o) { if (!(o instanceof PeriodDuration)) { return false; } PeriodDuration other = (PeriodDuration) o; return this.period.equals(other.period) && this.duration.equals(other.duration); } @Override public int hashCode() { return this.period.hashCode() * 31 + this.duration.hashCode(); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/SchemaChangeCallBack.java000066400000000000000000000032701515451615200322300ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.util.CallBack; /** * Callback for when the Schema for the Vector changes (generally happens when a vector is promoted * to a union type from a single value type). */ public class SchemaChangeCallBack implements CallBack { private boolean schemaChanged = false; /** Constructs a schema-change callback with the schema-changed state set to {@code false}. */ public SchemaChangeCallBack() {} /** Sets the schema-changed state to {@code true}. */ @Override public void doWork() { schemaChanged = true; } /** * Returns the value of schema-changed state, resetting the schema-changed state * to {@code false}. * * @return the previous schema-changed state */ public boolean getSchemaChangedAndReset() { final boolean current = schemaChanged; schemaChanged = false; return current; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/SmallIntVector.java000066400000000000000000000277451515451615200312700ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.SmallIntReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableSmallIntHolder; import org.apache.arrow.vector.holders.SmallIntHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * SmallIntVector implements a fixed width (2 bytes) vector of short values which could be null. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public final class SmallIntVector extends BaseFixedWidthVector implements BaseIntVector, ValueIterableVector { public static final byte TYPE_WIDTH = 2; /** * Instantiate a SmallIntVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public SmallIntVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.SMALLINT.getType()), allocator); } /** * Instantiate a SmallIntVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public SmallIntVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a SmallIntVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public SmallIntVector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new SmallIntReaderImpl(SmallIntVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.SMALLINT; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public short get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getShort((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableSmallIntHolder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getShort((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Short getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getShort((long) index * TYPE_WIDTH); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, int value) { valueBuffer.setShort((long) index * TYPE_WIDTH, value); } private void setValue(int index, short value) { valueBuffer.setShort((long) index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, int value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, short value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableSmallIntHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, SmallIntHolder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, int)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, int value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, short)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, short value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableSmallIntHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableSmallIntHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, SmallIntHolder)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, SmallIntHolder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void set(int index, int isSet, short value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, short)} except that it handles the case when index is greater * than or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setSafe(int index, int isSet, short value) { handleSafe(index); set(index, isSet, value); } /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static short get(final ArrowBuf buffer, final int index) { return buffer.getShort((long) index * TYPE_WIDTH); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((SmallIntVector) to); } @Override public void setWithPossibleTruncate(int index, long value) { this.setSafe(index, (int) value); } @Override public void setUnsafeWithPossibleTruncate(int index, long value) { this.set(index, (int) value); } @Override public long getValueAsLong(int index) { return this.get(index); } private class TransferImpl implements TransferPair { SmallIntVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new SmallIntVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new SmallIntVector(field, allocator); } public TransferImpl(SmallIntVector to) { this.to = to; } @Override public SmallIntVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, SmallIntVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/TimeMicroVector.java000066400000000000000000000257161515451615200314310ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.TimeMicroReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableTimeMicroHolder; import org.apache.arrow.vector.holders.TimeMicroHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * TimeMicroVector implements a fixed width vector (8 bytes) of time (microsecond resolution) values * which could be null. A validity buffer (bit vector) is maintained to track which elements in the * vector are null. */ public final class TimeMicroVector extends BaseFixedWidthVector implements ValueIterableVector { public static final byte TYPE_WIDTH = 8; /** * Instantiate a TimeMicroVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public TimeMicroVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.TIMEMICRO.getType()), allocator); } /** * Instantiate a TimeMicroVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public TimeMicroVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a TimeMicroVector. This doesn't allocate any memory for the data in vector. * * @param field Field materialized by this vector * @param allocator allocator for memory management. */ public TimeMicroVector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new TimeMicroReaderImpl(TimeMicroVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.TIMEMICRO; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public long get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getLong((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableTimeMicroHolder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getLong((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Long getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getLong((long) index * TYPE_WIDTH); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, long value) { valueBuffer.setLong((long) index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, long value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableTimeMicroHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, TimeMicroHolder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, long)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, long value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableTimeMicroHolder)} except that it handles the case when index * is greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableTimeMicroHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, TimeMicroHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, TimeMicroHolder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void set(int index, int isSet, long value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, long)} except that it handles the case when index is greater than * or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setSafe(int index, int isSet, long value) { handleSafe(index); set(index, isSet, value); } /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static long get(final ArrowBuf buffer, int index) { return buffer.getLong((long) index * TYPE_WIDTH); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((TimeMicroVector) to); } private class TransferImpl implements TransferPair { TimeMicroVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new TimeMicroVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new TimeMicroVector(field, allocator); } public TransferImpl(TimeMicroVector to) { this.to = to; } @Override public TimeMicroVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, TimeMicroVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/TimeMilliVector.java000066400000000000000000000262371515451615200314250ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import java.time.LocalDateTime; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.TimeMilliReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableTimeMilliHolder; import org.apache.arrow.vector.holders.TimeMilliHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.DateUtility; import org.apache.arrow.vector.util.TransferPair; /** * TimeMilliVector implements a fixed width (4 bytes) vector of time (millisecond resolution) values * which could be null. A validity buffer (bit vector) is maintained to track which elements in the * vector are null. */ public final class TimeMilliVector extends BaseFixedWidthVector implements ValueIterableVector { public static final byte TYPE_WIDTH = 4; /** * Instantiate a TimeMilliVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public TimeMilliVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.TIMEMILLI.getType()), allocator); } /** * Instantiate a TimeMilliVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public TimeMilliVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a TimeMilliVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public TimeMilliVector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new TimeMilliReaderImpl(TimeMilliVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.TIMEMILLI; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public int get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getInt((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableTimeMilliHolder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getInt((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public LocalDateTime getObject(int index) { if (isSet(index) == 0) { return null; } final int millis = valueBuffer.getInt((long) index * TYPE_WIDTH); // TODO: this doesn't seem right, time not from epoch return DateUtility.getLocalDateTimeFromEpochMilli(millis); } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, int value) { valueBuffer.setInt((long) index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, int value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableTimeMilliHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, TimeMilliHolder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, int)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, int value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableTimeMilliHolder)} except that it handles the case when index * is greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableTimeMilliHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, TimeMilliHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, TimeMilliHolder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void set(int index, int isSet, int value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, int)} except that it handles the case when index is greater than * or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setSafe(int index, int isSet, int value) { handleSafe(index); set(index, isSet, value); } /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static int get(final ArrowBuf buffer, final int index) { return buffer.getInt((long) index * TYPE_WIDTH); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((TimeMilliVector) to); } private class TransferImpl implements TransferPair { TimeMilliVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new TimeMilliVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new TimeMilliVector(field, allocator); } public TransferImpl(TimeMilliVector to) { this.to = to; } @Override public TimeMilliVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, TimeMilliVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/TimeNanoVector.java000066400000000000000000000256661515451615200312570ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.TimeNanoReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableTimeNanoHolder; import org.apache.arrow.vector.holders.TimeNanoHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * TimeNanoVector implements a fixed width vector (8 bytes) of time (nanosecond resolution) values * which could be null. A validity buffer (bit vector) is maintained to track which elements in the * vector are null. */ public final class TimeNanoVector extends BaseFixedWidthVector implements ValueIterableVector { public static final byte TYPE_WIDTH = 8; /** * Instantiate a TimeNanoVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public TimeNanoVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.TIMENANO.getType()), allocator); } /** * Instantiate a TimeNanoVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public TimeNanoVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a TimeNanoVector. This doesn't allocate any memory for the data in vector. * * @param field Field materialized by this vector * @param allocator allocator for memory management. */ public TimeNanoVector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new TimeNanoReaderImpl(TimeNanoVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.TIMENANO; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public long get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getLong((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableTimeNanoHolder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getLong((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Long getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getLong((long) index * TYPE_WIDTH); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, long value) { valueBuffer.setLong((long) index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, long value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableTimeNanoHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, TimeNanoHolder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, long)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, long value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableTimeNanoHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableTimeNanoHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, TimeNanoHolder)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, TimeNanoHolder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void set(int index, int isSet, long value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, long)} except that it handles the case when index is greater than * or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setSafe(int index, int isSet, long value) { handleSafe(index); set(index, isSet, value); } /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static long get(final ArrowBuf buffer, final int index) { return buffer.getLong((long) index * TYPE_WIDTH); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((TimeNanoVector) to); } private class TransferImpl implements TransferPair { TimeNanoVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new TimeNanoVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new TimeNanoVector(field, allocator); } public TransferImpl(TimeNanoVector to) { this.to = to; } @Override public TimeNanoVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, TimeNanoVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/TimeSecVector.java000066400000000000000000000256211515451615200310650ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.TimeSecReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableTimeSecHolder; import org.apache.arrow.vector.holders.TimeSecHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * TimeSecVector implements a fixed width (4 bytes) vector of time (seconds resolution) values which * could be null. A validity buffer (bit vector) is maintained to track which elements in the vector * are null. */ public final class TimeSecVector extends BaseFixedWidthVector implements ValueIterableVector { public static final byte TYPE_WIDTH = 4; /** * Instantiate a TimeSecVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public TimeSecVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.TIMESEC.getType()), allocator); } /** * Instantiate a TimeSecVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public TimeSecVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a TimeSecVector. This doesn't allocate any memory for the data in vector. * * @param field Field materialized by this vector * @param allocator allocator for memory management. */ public TimeSecVector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new TimeSecReaderImpl(TimeSecVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.TIMESEC; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public int get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getInt((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableTimeSecHolder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getInt((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Integer getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getInt((long) index * TYPE_WIDTH); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, int value) { valueBuffer.setInt((long) index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, int value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableTimeSecHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, TimeSecHolder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, int)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, int value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableTimeSecHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableTimeSecHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, TimeSecHolder)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, TimeSecHolder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void set(int index, int isSet, int value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, int)} except that it handles the case when index is greater than * or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setSafe(int index, int isSet, int value) { handleSafe(index); set(index, isSet, value); } /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static int get(final ArrowBuf buffer, final int index) { return buffer.getInt((long) index * TYPE_WIDTH); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising of this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((TimeSecVector) to); } private class TransferImpl implements TransferPair { TimeSecVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new TimeSecVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new TimeSecVector(field, allocator); } public TransferImpl(TimeSecVector to) { this.to = to; } @Override public TimeSecVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, TimeSecVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/TimeStampMicroTZVector.java000066400000000000000000000222601515451615200327030ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.TimeStampMicroTZReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableTimeStampMicroTZHolder; import org.apache.arrow.vector.holders.TimeStampMicroTZHolder; 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.TransferPair; /** * TimeStampMicroTZVector implements a fixed width vector (8 bytes) of timestamp (microsecond * resolution) values which could be null. A validity buffer (bit vector) is maintained to track * which elements in the vector are null. */ public final class TimeStampMicroTZVector extends TimeStampVector implements ValueIterableVector { private final String timeZone; /** * Instantiate a TimeStampMicroTZVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public TimeStampMicroTZVector(String name, BufferAllocator allocator, String timeZone) { this( name, FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MICROSECOND, timeZone)), allocator); } /** * Instantiate a TimeStampMicroTZVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampMicroTZVector(String name, FieldType fieldType, BufferAllocator allocator) { super(name, fieldType, allocator); ArrowType.Timestamp arrowType = (ArrowType.Timestamp) fieldType.getType(); timeZone = arrowType.getTimezone(); } /** * Instantiate a TimeStampMicroTZVector. This doesn't allocate any memory for the data in vector. * * @param field Field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampMicroTZVector(Field field, BufferAllocator allocator) { super(field, allocator); ArrowType.Timestamp arrowType = (ArrowType.Timestamp) field.getFieldType().getType(); timeZone = arrowType.getTimezone(); } @Override protected FieldReader getReaderImpl() { return new TimeStampMicroTZReaderImpl(TimeStampMicroTZVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.TIMESTAMPMICROTZ; } /** * Get the time zone of the timestamps stored in this vector. * * @return the time zone of the timestamps stored in this vector. */ public String getTimeZone() { return this.timeZone; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableTimeStampMicroTZHolder holder) { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getLong((long) index * TYPE_WIDTH); holder.timezone = timeZone; } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Long getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getLong((long) index * TYPE_WIDTH); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableTimeStampMicroTZHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } 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 { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, TimeStampMicroTZHolder holder) { 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); } /** * Same as {@link #set(int, NullableTimeStampMicroTZHolder)} except that it handles the case when * index is greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableTimeStampMicroTZHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, TimeStampMicroTZHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, TimeStampMicroTZHolder holder) { handleSafe(index); set(index, holder); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { TimeStampMicroTZVector to = new TimeStampMicroTZVector(ref, field.getFieldType(), allocator); return new TransferImpl(to); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { TimeStampMicroTZVector to = new TimeStampMicroTZVector(field, allocator); return new TransferImpl(to); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((TimeStampMicroTZVector) to); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/TimeStampMicroVector.java000066400000000000000000000202101515451615200324160ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import java.time.LocalDateTime; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.TimeStampMicroReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableTimeStampMicroHolder; import org.apache.arrow.vector.holders.TimeStampMicroHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.DateUtility; import org.apache.arrow.vector.util.TransferPair; /** * TimeStampMicroVector implements a fixed width vector (8 bytes) of timestamp (microsecond * resolution) values which could be null. A validity buffer (bit vector) is maintained to track * which elements in the vector are null. */ public final class TimeStampMicroVector extends TimeStampVector implements ValueIterableVector { /** * Instantiate a TimeStampMicroVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public TimeStampMicroVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.TIMESTAMPMICRO.getType()), allocator); } /** * Instantiate a TimeStampMicroVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampMicroVector(String name, FieldType fieldType, BufferAllocator allocator) { super(name, fieldType, allocator); } /** * Instantiate a TimeStampMicroVector. This doesn't allocate any memory for the data in vector. * * @param field Field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampMicroVector(Field field, BufferAllocator allocator) { super(field, allocator); } @Override protected FieldReader getReaderImpl() { return new TimeStampMicroReaderImpl(TimeStampMicroVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.TIMESTAMPMICRO; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableTimeStampMicroHolder holder) { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getLong((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public LocalDateTime getObject(int index) { if (isSet(index) == 0) { return null; } else { final long micros = valueBuffer.getLong((long) index * TYPE_WIDTH); return DateUtility.getLocalDateTimeFromEpochMicro(micros); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableTimeStampMicroHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, TimeStampMicroHolder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, NullableTimeStampMicroHolder)} except that it handles the case when * index is greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableTimeStampMicroHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, TimeStampMicroHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, TimeStampMicroHolder holder) { handleSafe(index); set(index, holder); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { TimeStampMicroVector to = new TimeStampMicroVector(ref, field.getFieldType(), allocator); return new TransferImpl(to); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { TimeStampMicroVector to = new TimeStampMicroVector(field, allocator); return new TransferImpl(to); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((TimeStampMicroVector) to); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/TimeStampMilliTZVector.java000066400000000000000000000222601515451615200327000ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.TimeStampMilliTZReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableTimeStampMilliTZHolder; import org.apache.arrow.vector.holders.TimeStampMilliTZHolder; 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.TransferPair; /** * TimeStampMilliTZVector implements a fixed width vector (8 bytes) of timestamp (millisecond * resolution) values which could be null. A validity buffer (bit vector) is maintained to track * which elements in the vector are null. */ public final class TimeStampMilliTZVector extends TimeStampVector implements ValueIterableVector { private final String timeZone; /** * Instantiate a TimeStampMilliTZVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public TimeStampMilliTZVector(String name, BufferAllocator allocator, String timeZone) { this( name, FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, timeZone)), allocator); } /** * Instantiate a TimeStampMilliTZVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampMilliTZVector(String name, FieldType fieldType, BufferAllocator allocator) { super(name, fieldType, allocator); ArrowType.Timestamp arrowType = (ArrowType.Timestamp) fieldType.getType(); timeZone = arrowType.getTimezone(); } /** * Instantiate a TimeStampMilliTZVector. This doesn't allocate any memory for the data in vector. * * @param field Field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampMilliTZVector(Field field, BufferAllocator allocator) { super(field, allocator); ArrowType.Timestamp arrowType = (ArrowType.Timestamp) field.getFieldType().getType(); timeZone = arrowType.getTimezone(); } @Override protected FieldReader getReaderImpl() { return new TimeStampMilliTZReaderImpl(TimeStampMilliTZVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.TIMESTAMPMILLITZ; } /** * Get the time zone of the timestamps stored in this vector. * * @return the time zone of the timestamps stored in this vector. */ public String getTimeZone() { return this.timeZone; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableTimeStampMilliTZHolder holder) { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getLong((long) index * TYPE_WIDTH); holder.timezone = timeZone; } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Long getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getLong((long) index * TYPE_WIDTH); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableTimeStampMilliTZHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } 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 { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, TimeStampMilliTZHolder holder) { 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); } /** * Same as {@link #set(int, NullableTimeStampMilliTZHolder)} except that it handles the case when * index is greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableTimeStampMilliTZHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, TimeStampMilliTZHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, TimeStampMilliTZHolder holder) { handleSafe(index); set(index, holder); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { TimeStampMilliTZVector to = new TimeStampMilliTZVector(ref, field.getFieldType(), allocator); return new TransferImpl(to); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { TimeStampMilliTZVector to = new TimeStampMilliTZVector(field, allocator); return new TransferImpl(to); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((TimeStampMilliTZVector) to); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/TimeStampMilliVector.java000066400000000000000000000202101515451615200324130ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import java.time.LocalDateTime; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.TimeStampMilliReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableTimeStampMilliHolder; import org.apache.arrow.vector.holders.TimeStampMilliHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.DateUtility; import org.apache.arrow.vector.util.TransferPair; /** * TimeStampMilliVector implements a fixed width vector (8 bytes) of timestamp (millisecond * resolution) values which could be null. A validity buffer (bit vector) is maintained to track * which elements in the vector are null. */ public final class TimeStampMilliVector extends TimeStampVector implements ValueIterableVector { /** * Instantiate a TimeStampMilliVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public TimeStampMilliVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.TIMESTAMPMILLI.getType()), allocator); } /** * Instantiate a TimeStampMilliVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampMilliVector(String name, FieldType fieldType, BufferAllocator allocator) { super(name, fieldType, allocator); } /** * Instantiate a TimeStampMilliVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampMilliVector(Field field, BufferAllocator allocator) { super(field, allocator); } @Override protected FieldReader getReaderImpl() { return new TimeStampMilliReaderImpl(TimeStampMilliVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.TIMESTAMPMILLI; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableTimeStampMilliHolder holder) { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getLong((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public LocalDateTime getObject(int index) { if (isSet(index) == 0) { return null; } else { final long millis = valueBuffer.getLong((long) index * TYPE_WIDTH); return DateUtility.getLocalDateTimeFromEpochMilli(millis); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableTimeStampMilliHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, TimeStampMilliHolder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, NullableTimeStampMilliHolder)} except that it handles the case when * index is greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableTimeStampMilliHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, TimeStampMilliHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, TimeStampMilliHolder holder) { handleSafe(index); set(index, holder); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { TimeStampMilliVector to = new TimeStampMilliVector(ref, field.getFieldType(), allocator); return new TransferImpl(to); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { TimeStampMilliVector to = new TimeStampMilliVector(field, allocator); return new TransferImpl(to); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((TimeStampMilliVector) to); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/TimeStampNanoTZVector.java000066400000000000000000000222161515451615200325260ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.TimeStampNanoTZReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableTimeStampNanoTZHolder; 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.TransferPair; /** * TimeStampNanoTZVector implements a fixed width vector (8 bytes) of timestamp (nanosecond * resolution) values which could be null. A validity buffer (bit vector) is maintained to track * which elements in the vector are null. */ public final class TimeStampNanoTZVector extends TimeStampVector implements ValueIterableVector { private final String timeZone; /** * Instantiate a TimeStampNanoTZVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public TimeStampNanoTZVector(String name, BufferAllocator allocator, String timeZone) { this( name, FieldType.nullable(new ArrowType.Timestamp(TimeUnit.NANOSECOND, timeZone)), allocator); } /** * Instantiate a TimeStampNanoTZVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampNanoTZVector(String name, FieldType fieldType, BufferAllocator allocator) { super(name, fieldType, allocator); ArrowType.Timestamp arrowType = (ArrowType.Timestamp) fieldType.getType(); timeZone = arrowType.getTimezone(); } /** * Instantiate a TimeStampNanoTZVector. This doesn't allocate any memory for the data in vector. * * @param field Field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampNanoTZVector(Field field, BufferAllocator allocator) { super(field, allocator); ArrowType.Timestamp arrowType = (ArrowType.Timestamp) field.getFieldType().getType(); timeZone = arrowType.getTimezone(); } @Override protected FieldReader getReaderImpl() { return new TimeStampNanoTZReaderImpl(TimeStampNanoTZVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.TIMESTAMPNANOTZ; } /** * Get the time zone of the timestamps stored in this vector. * * @return the time zone of the timestamps stored in this vector. */ public String getTimeZone() { return this.timeZone; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableTimeStampNanoTZHolder holder) { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getLong((long) index * TYPE_WIDTH); holder.timezone = timeZone; } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Long getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getLong((long) index * TYPE_WIDTH); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableTimeStampNanoTZHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } 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 { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, TimeStampNanoTZHolder holder) { 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); } /** * Same as {@link #set(int, NullableTimeStampNanoTZHolder)} except that it handles the case when * index is greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableTimeStampNanoTZHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, TimeStampNanoTZHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, TimeStampNanoTZHolder holder) { handleSafe(index); set(index, holder); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { TimeStampNanoTZVector to = new TimeStampNanoTZVector(ref, field.getFieldType(), allocator); return new TransferImpl(to); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { TimeStampNanoTZVector to = new TimeStampNanoTZVector(field, allocator); return new TransferImpl(to); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((TimeStampNanoTZVector) to); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/TimeStampNanoVector.java000066400000000000000000000201511515451615200322440ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import java.time.LocalDateTime; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.TimeStampNanoReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableTimeStampNanoHolder; import org.apache.arrow.vector.holders.TimeStampNanoHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.DateUtility; import org.apache.arrow.vector.util.TransferPair; /** * TimeStampNanoVector implements a fixed width vector (8 bytes) of timestamp (nanosecond * resolution) values which could be null. A validity buffer (bit vector) is maintained to track * which elements in the vector are null. */ public final class TimeStampNanoVector extends TimeStampVector implements ValueIterableVector { /** * Instantiate a TimeStampNanoVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public TimeStampNanoVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.TIMESTAMPNANO.getType()), allocator); } /** * Instantiate a TimeStampNanoVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampNanoVector(String name, FieldType fieldType, BufferAllocator allocator) { super(name, fieldType, allocator); } /** * Instantiate a TimeStampNanoVector. This doesn't allocate any memory for the data in vector. * * @param field Field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampNanoVector(Field field, BufferAllocator allocator) { super(field, allocator); } @Override protected FieldReader getReaderImpl() { return new TimeStampNanoReaderImpl(TimeStampNanoVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.TIMESTAMPNANO; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableTimeStampNanoHolder holder) { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getLong((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public LocalDateTime getObject(int index) { if (isSet(index) == 0) { return null; } else { final long nanos = valueBuffer.getLong((long) index * TYPE_WIDTH); return DateUtility.getLocalDateTimeFromEpochNano(nanos); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableTimeStampNanoHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, TimeStampNanoHolder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, NullableTimeStampNanoHolder)} except that it handles the case when * index is greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableTimeStampNanoHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, TimeStampNanoHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, TimeStampNanoHolder holder) { handleSafe(index); set(index, holder); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { TimeStampNanoVector to = new TimeStampNanoVector(ref, field.getFieldType(), allocator); return new TransferImpl(to); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { TimeStampNanoVector to = new TimeStampNanoVector(field, allocator); return new TransferImpl(to); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((TimeStampNanoVector) to); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/TimeStampSecTZVector.java000066400000000000000000000221161515451615200323440ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.TimeStampSecTZReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableTimeStampSecTZHolder; import org.apache.arrow.vector.holders.TimeStampSecTZHolder; 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.TransferPair; /** * TimeStampSecTZVector implements a fixed width vector (8 bytes) of timestamp (seconds resolution) * values which could be null. A validity buffer (bit vector) is maintained to track which elements * in the vector are null. */ public final class TimeStampSecTZVector extends TimeStampVector implements ValueIterableVector { private final String timeZone; /** * Instantiate a TimeStampSecTZVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public TimeStampSecTZVector(String name, BufferAllocator allocator, String timeZone) { this(name, FieldType.nullable(new ArrowType.Timestamp(TimeUnit.SECOND, timeZone)), allocator); } /** * Instantiate a TimeStampSecTZVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampSecTZVector(String name, FieldType fieldType, BufferAllocator allocator) { super(name, fieldType, allocator); ArrowType.Timestamp arrowType = (ArrowType.Timestamp) fieldType.getType(); timeZone = arrowType.getTimezone(); } /** * Instantiate a TimeStampSecTZVector. This doesn't allocate any memory for the data in vector. * * @param field Field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampSecTZVector(Field field, BufferAllocator allocator) { super(field, allocator); ArrowType.Timestamp arrowType = (ArrowType.Timestamp) field.getFieldType().getType(); timeZone = arrowType.getTimezone(); } @Override protected FieldReader getReaderImpl() { return new TimeStampSecTZReaderImpl(TimeStampSecTZVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.TIMESTAMPSECTZ; } /** * Get the time zone of the timestamps stored in this vector. * * @return the time zone of the timestamps stored in this vector. */ public String getTimeZone() { return this.timeZone; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableTimeStampSecTZHolder holder) { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getLong((long) index * TYPE_WIDTH); holder.timezone = timeZone; } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ public Long getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getLong((long) index * TYPE_WIDTH); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableTimeStampSecTZHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } 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 { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, TimeStampSecTZHolder holder) { 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); } /** * Same as {@link #set(int, NullableTimeStampSecTZHolder)} except that it handles the case when * index is greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableTimeStampSecTZHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, TimeStampSecTZHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, TimeStampSecTZHolder holder) { handleSafe(index); set(index, holder); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising of this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { TimeStampSecTZVector to = new TimeStampSecTZVector(ref, field.getFieldType(), allocator); return new TransferImpl(to); } /** * Construct a TransferPair comprising of this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { TimeStampSecTZVector to = new TimeStampSecTZVector(field, allocator); return new TransferImpl(to); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((TimeStampSecTZVector) to); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/TimeStampSecVector.java000066400000000000000000000202341515451615200320650ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import java.time.LocalDateTime; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.TimeStampSecReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableTimeStampSecHolder; import org.apache.arrow.vector.holders.TimeStampSecHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.DateUtility; import org.apache.arrow.vector.util.TransferPair; /** * TimeStampSecVector implements a fixed width vector (8 bytes) of timestamp (seconds resolution) * values which could be null. A validity buffer (bit vector) is maintained to track which elements * in the vector are null. */ public final class TimeStampSecVector extends TimeStampVector implements ValueIterableVector { /** * Instantiate a TimeStampSecVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public TimeStampSecVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.TIMESTAMPSEC.getType()), allocator); } /** * Instantiate a TimeStampSecVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampSecVector(String name, FieldType fieldType, BufferAllocator allocator) { super(name, fieldType, allocator); } /** * Instantiate a TimeStampSecVector. This doesn't allocate any memory for the data in vector. * * @param field Field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampSecVector(Field field, BufferAllocator allocator) { super(field, allocator); } @Override protected FieldReader getReaderImpl() { return new TimeStampSecReaderImpl(TimeStampSecVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.TIMESTAMPSEC; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableTimeStampSecHolder holder) { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getLong((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public LocalDateTime getObject(int index) { if (isSet(index) == 0) { return null; } else { final long secs = valueBuffer.getLong((long) index * TYPE_WIDTH); final long millis = java.util.concurrent.TimeUnit.SECONDS.toMillis(secs); return DateUtility.getLocalDateTimeFromEpochMilli(millis); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableTimeStampSecHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, TimeStampSecHolder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, NullableTimeStampSecHolder)} except that it handles the case when * index is greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableTimeStampSecHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, TimeStampSecHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, TimeStampSecHolder holder) { handleSafe(index); set(index, holder); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { TimeStampSecVector to = new TimeStampSecVector(ref, field.getFieldType(), allocator); return new TransferImpl(to); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { TimeStampSecVector to = new TimeStampSecVector(field, allocator); return new TransferImpl(to); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((TimeStampSecVector) to); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/TimeStampVector.java000066400000000000000000000142641515451615200314400ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * TimeStampVector is an abstract interface for fixed width vector (8 bytes) of timestamp values * which could be null. A validity buffer (bit vector) is maintained to track which elements in the * vector are null. */ public abstract class TimeStampVector extends BaseFixedWidthVector { public static final byte TYPE_WIDTH = 8; /** * Instantiate a TimeStampVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a TimeStampVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public TimeStampVector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public long get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getLong((long) index * TYPE_WIDTH); } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ protected void setValue(int index, long value) { valueBuffer.setLong((long) index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, long value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Same as {@link #set(int, long)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, long value) { handleSafe(index); set(index, value); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void set(int index, int isSet, long value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, long)} except that it handles the case when index is greater than * or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setSafe(int index, int isSet, long value) { handleSafe(index); set(index, isSet, value); } /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static long get(final ArrowBuf buffer, final int index) { return buffer.getLong((long) index * TYPE_WIDTH); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** {@link TransferPair} for {@link TimeStampVector}. */ public class TransferImpl implements TransferPair { TimeStampVector to; public TransferImpl(TimeStampVector to) { this.to = to; } @Override public TimeStampVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, TimeStampVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/TinyIntVector.java000066400000000000000000000276141515451615200311360ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.TinyIntReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableTinyIntHolder; import org.apache.arrow.vector.holders.TinyIntHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * TinyIntVector implements a fixed width (1 bytes) vector of byte values which could be null. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public final class TinyIntVector extends BaseFixedWidthVector implements BaseIntVector, ValueIterableVector { public static final byte TYPE_WIDTH = 1; /** * Instantiate a TinyIntVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public TinyIntVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.TINYINT.getType()), allocator); } /** * Instantiate a TinyIntVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public TinyIntVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a TinyIntVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public TinyIntVector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new TinyIntReaderImpl(TinyIntVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.TINYINT; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public byte get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getByte(index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableTinyIntHolder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getByte(index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Byte getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getByte(index * TYPE_WIDTH); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, int value) { valueBuffer.setByte(index * TYPE_WIDTH, value); } private void setValue(int index, byte value) { valueBuffer.setByte(index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, int value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, byte value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableTinyIntHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, TinyIntHolder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, int)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, int value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, byte)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, byte value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableTinyIntHolder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableTinyIntHolder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, TinyIntHolder)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, TinyIntHolder holder) { handleSafe(index); set(index, holder); } /** * Store the given value at a particular position in the vector. isSet indicates whether the value * is NULL or not. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void set(int index, int isSet, byte value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, byte)} except that it handles the case when index is greater than * or equal to current value capacity of the vector. * * @param index position of the new value * @param isSet 0 for NULL value, 1 otherwise * @param value element value */ public void setSafe(int index, int isSet, byte value) { handleSafe(index); set(index, isSet, value); } /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method should not be used externally. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static byte get(final ArrowBuf buffer, final int index) { return buffer.getByte(index * TYPE_WIDTH); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((TinyIntVector) to); } @Override public void setWithPossibleTruncate(int index, long value) { this.setSafe(index, (int) value); } @Override public void setUnsafeWithPossibleTruncate(int index, long value) { this.set(index, (int) value); } @Override public long getValueAsLong(int index) { return this.get(index); } private class TransferImpl implements TransferPair { TinyIntVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new TinyIntVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new TinyIntVector(field, allocator); } public TransferImpl(TinyIntVector to) { this.to = to; } @Override public TinyIntVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, TinyIntVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/TypeLayout.java000066400000000000000000000436421515451615200304730ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.Arrays.asList; import java.util.ArrayList; import java.util.Collections; import java.util.List; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BufferLayout.BufferType; 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.Binary; import org.apache.arrow.vector.types.pojo.ArrowType.BinaryView; import org.apache.arrow.vector.types.pojo.ArrowType.Bool; import org.apache.arrow.vector.types.pojo.ArrowType.Date; import org.apache.arrow.vector.types.pojo.ArrowType.Decimal; import org.apache.arrow.vector.types.pojo.ArrowType.Duration; import org.apache.arrow.vector.types.pojo.ArrowType.FixedSizeBinary; import org.apache.arrow.vector.types.pojo.ArrowType.FixedSizeList; import org.apache.arrow.vector.types.pojo.ArrowType.FloatingPoint; import org.apache.arrow.vector.types.pojo.ArrowType.Int; import org.apache.arrow.vector.types.pojo.ArrowType.Interval; import org.apache.arrow.vector.types.pojo.ArrowType.LargeBinary; import org.apache.arrow.vector.types.pojo.ArrowType.LargeUtf8; import org.apache.arrow.vector.types.pojo.ArrowType.Map; import org.apache.arrow.vector.types.pojo.ArrowType.Null; import org.apache.arrow.vector.types.pojo.ArrowType.RunEndEncoded; import org.apache.arrow.vector.types.pojo.ArrowType.Struct; import org.apache.arrow.vector.types.pojo.ArrowType.Time; import org.apache.arrow.vector.types.pojo.ArrowType.Timestamp; import org.apache.arrow.vector.types.pojo.ArrowType.Union; import org.apache.arrow.vector.types.pojo.ArrowType.Utf8; import org.apache.arrow.vector.types.pojo.ArrowType.Utf8View; /** * The buffer layout of vectors for a given type. It defines its own buffers followed by the buffers * for the children if it is a nested type (Struct_, List, Union) */ public class TypeLayout { /** Constructs a new {@link TypeLayout} for the given arrowType. */ public static TypeLayout getTypeLayout(final ArrowType arrowType) { TypeLayout layout = arrowType.accept( new ArrowTypeVisitor() { @Override public TypeLayout visit(Int type) { return newFixedWidthTypeLayout(BufferLayout.dataBuffer(type.getBitWidth())); } @Override public TypeLayout visit(Union type) { List vectors; switch (type.getMode()) { case Dense: vectors = asList( BufferLayout.typeBuffer(), BufferLayout.offsetBuffer() // offset to find the vector ); break; case Sparse: vectors = asList( BufferLayout.typeBuffer() // type of the value at the index or 0 if null ); break; default: throw new UnsupportedOperationException( "Unsupported Union Mode: " + type.getMode()); } return new TypeLayout(vectors); } @Override public TypeLayout visit(Struct type) { List vectors = asList(BufferLayout.validityVector()); return new TypeLayout(vectors); } @Override public TypeLayout visit(Timestamp type) { return newFixedWidthTypeLayout(BufferLayout.dataBuffer(64)); } @Override public TypeLayout visit(ArrowType.List type) { List vectors = asList(BufferLayout.validityVector(), BufferLayout.offsetBuffer()); return new TypeLayout(vectors); } @Override public TypeLayout visit(ArrowType.ListView type) { List vectors = asList( BufferLayout.validityVector(), BufferLayout.offsetBuffer(), BufferLayout.sizeBuffer()); return new TypeLayout(vectors); } @Override public TypeLayout visit(ArrowType.LargeListView type) { List vectors = asList( BufferLayout.validityVector(), BufferLayout.largeOffsetBuffer(), BufferLayout.largeSizeBuffer()); return new TypeLayout(vectors); } @Override public TypeLayout visit(ArrowType.LargeList type) { List vectors = asList(BufferLayout.validityVector(), BufferLayout.largeOffsetBuffer()); return new TypeLayout(vectors); } @Override public TypeLayout visit(FixedSizeList type) { List vectors = asList(BufferLayout.validityVector()); return new TypeLayout(vectors); } @Override public TypeLayout visit(Map type) { List vectors = asList(BufferLayout.validityVector(), BufferLayout.offsetBuffer()); return new TypeLayout(vectors); } @Override public TypeLayout visit(FloatingPoint type) { int bitWidth; switch (type.getPrecision()) { case HALF: bitWidth = 16; break; case SINGLE: bitWidth = 32; break; case DOUBLE: bitWidth = 64; break; default: throw new UnsupportedOperationException( "Unsupported Precision: " + type.getPrecision()); } return newFixedWidthTypeLayout(BufferLayout.dataBuffer(bitWidth)); } @Override public TypeLayout visit(Decimal type) { return newFixedWidthTypeLayout(BufferLayout.dataBuffer(type.getBitWidth())); } @Override public TypeLayout visit(FixedSizeBinary type) { return newFixedWidthTypeLayout( new BufferLayout(BufferType.DATA, type.getByteWidth() * 8)); } @Override public TypeLayout visit(Bool type) { return newFixedWidthTypeLayout(BufferLayout.booleanVector()); } @Override public TypeLayout visit(Binary type) { return newVariableWidthTypeLayout(); } @Override public TypeLayout visit(ArrowType.BinaryView type) { return newVariableWidthViewTypeLayout(); } @Override public TypeLayout visit(Utf8 type) { return newVariableWidthTypeLayout(); } @Override public TypeLayout visit(Utf8View type) { return newVariableWidthViewTypeLayout(); } @Override public TypeLayout visit(LargeUtf8 type) { return newLargeVariableWidthTypeLayout(); } @Override public TypeLayout visit(LargeBinary type) { return newLargeVariableWidthTypeLayout(); } private TypeLayout newVariableWidthTypeLayout() { return newPrimitiveTypeLayout( BufferLayout.validityVector(), BufferLayout.offsetBuffer(), BufferLayout.byteVector()); } private TypeLayout newVariableWidthViewTypeLayout() { return new TypeLayout( false, BufferLayout.validityVector(), BufferLayout.viewVector()); } private TypeLayout newLargeVariableWidthTypeLayout() { // NOTE: only considers the non variadic buffers return newPrimitiveTypeLayout( BufferLayout.validityVector(), BufferLayout.largeOffsetBuffer(), BufferLayout.byteVector()); } private TypeLayout newPrimitiveTypeLayout(BufferLayout... vectors) { return new TypeLayout(asList(vectors)); } public TypeLayout newFixedWidthTypeLayout(BufferLayout dataVector) { return newPrimitiveTypeLayout(BufferLayout.validityVector(), dataVector); } @Override public TypeLayout visit(Null type) { return new TypeLayout(Collections.emptyList()); } @Override public TypeLayout visit(Date type) { switch (type.getUnit()) { case DAY: return newFixedWidthTypeLayout(BufferLayout.dataBuffer(32)); case MILLISECOND: return newFixedWidthTypeLayout(BufferLayout.dataBuffer(64)); default: throw new UnsupportedOperationException("Unknown unit " + type.getUnit()); } } @Override public TypeLayout visit(Time type) { return newFixedWidthTypeLayout(BufferLayout.dataBuffer(type.getBitWidth())); } @Override public TypeLayout visit(Interval type) { switch (type.getUnit()) { case DAY_TIME: return newFixedWidthTypeLayout(BufferLayout.dataBuffer(64)); case YEAR_MONTH: return newFixedWidthTypeLayout(BufferLayout.dataBuffer(32)); case MONTH_DAY_NANO: return newFixedWidthTypeLayout(BufferLayout.dataBuffer(128)); default: throw new UnsupportedOperationException("Unknown unit " + type.getUnit()); } } @Override public TypeLayout visit(Duration type) { return newFixedWidthTypeLayout(BufferLayout.dataBuffer(64)); } @Override public TypeLayout visit(RunEndEncoded type) { return new TypeLayout(Collections.emptyList()); } }); return layout; } /** Gets the number of {@link BufferLayout}s for the given arrowType. */ public static int getTypeBufferCount(final ArrowType arrowType) { return arrowType.accept( new ArrowTypeVisitor() { /** * All fixed width vectors have a common number of buffers 2: one validity buffer, plus a * data buffer. */ static final int FIXED_WIDTH_BUFFER_COUNT = 2; /** * All variable width vectors have a common number of buffers 3: a validity buffer, an * offset buffer, and a data buffer. */ static final int VARIABLE_WIDTH_BUFFER_COUNT = 3; @Override public Integer visit(Int type) { return FIXED_WIDTH_BUFFER_COUNT; } @Override public Integer visit(Union type) { switch (type.getMode()) { case Dense: // TODO: validate this return 2; case Sparse: // type buffer return 1; default: throw new UnsupportedOperationException( "Unsupported Union Mode: " + type.getMode()); } } @Override public Integer visit(Struct type) { // validity buffer return 1; } @Override public Integer visit(Timestamp type) { return FIXED_WIDTH_BUFFER_COUNT; } @Override public Integer visit(ArrowType.List type) { // validity buffer + offset buffer return 2; } @Override public Integer visit(ArrowType.ListView type) { // validity buffer + offset buffer + size buffer return 3; } @Override public Integer visit(ArrowType.LargeList type) { // validity buffer + offset buffer return 2; } @Override public Integer visit(ArrowType.LargeListView type) { // validity buffer + offset buffer + size buffer return 3; } @Override public Integer visit(FixedSizeList type) { // validity buffer return 1; } @Override public Integer visit(Map type) { // validity buffer + offset buffer return 2; } @Override public Integer visit(FloatingPoint type) { return FIXED_WIDTH_BUFFER_COUNT; } @Override public Integer visit(Decimal type) { return FIXED_WIDTH_BUFFER_COUNT; } @Override public Integer visit(FixedSizeBinary type) { return FIXED_WIDTH_BUFFER_COUNT; } @Override public Integer visit(Bool type) { return FIXED_WIDTH_BUFFER_COUNT; } @Override public Integer visit(Binary type) { return VARIABLE_WIDTH_BUFFER_COUNT; } @Override public Integer visit(BinaryView type) { // NOTE: only consider the validity and view buffers return 2; } @Override public Integer visit(Utf8 type) { return VARIABLE_WIDTH_BUFFER_COUNT; } @Override public Integer visit(Utf8View type) { // NOTE: only consider the validity and view buffers return 2; } @Override public Integer visit(LargeUtf8 type) { return VARIABLE_WIDTH_BUFFER_COUNT; } @Override public Integer visit(LargeBinary type) { return VARIABLE_WIDTH_BUFFER_COUNT; } @Override public Integer visit(Null type) { return 0; } @Override public Integer visit(Date type) { return FIXED_WIDTH_BUFFER_COUNT; } @Override public Integer visit(Time type) { return FIXED_WIDTH_BUFFER_COUNT; } @Override public Integer visit(Interval type) { return FIXED_WIDTH_BUFFER_COUNT; } @Override public Integer visit(Duration type) { return FIXED_WIDTH_BUFFER_COUNT; } @Override public Integer visit(RunEndEncoded type) { return 0; } }); } private final List bufferLayouts; private final boolean isFixedBufferCount; /** * Constructs a new {@link TypeLayout}. * * @param bufferLayouts the individual {@linkplain BufferLayout}s for the given type * @param isFixedBufferCount whether the number of buffers is fixed */ public TypeLayout(List bufferLayouts, boolean isFixedBufferCount) { super(); this.bufferLayouts = Preconditions.checkNotNull(bufferLayouts); this.isFixedBufferCount = isFixedBufferCount; } public TypeLayout(List bufferLayouts) { this(bufferLayouts, true); } public TypeLayout(BufferLayout... bufferLayouts) { this(asList(bufferLayouts), true); } public TypeLayout(boolean isFixedBufferCount, BufferLayout... bufferLayouts) { this(asList(bufferLayouts), isFixedBufferCount); } /** Returns the individual {@linkplain BufferLayout}s for the given type. */ public List getBufferLayouts() { return bufferLayouts; } /** * Returns the types of each buffer for this layout. A layout can consist of multiple buffers for * example a validity bitmap buffer, a value buffer or an offset buffer. */ public List getBufferTypes() { List types = new ArrayList<>(bufferLayouts.size()); for (BufferLayout vector : bufferLayouts) { types.add(vector.getType()); } return types; } /** * Determines whether the buffer count is fixed for the given type. * * @return true if the buffer count is fixed, false otherwise */ public boolean isFixedBufferCount() { return isFixedBufferCount; } @Override public String toString() { return bufferLayouts.toString(); } @Override public int hashCode() { return bufferLayouts.hashCode(); } @Override public boolean equals(Object obj) { if (this == obj) { return true; } if (obj == null) { return false; } if (!(obj instanceof TypeLayout)) { return false; } TypeLayout other = (TypeLayout) obj; return bufferLayouts.equals(other.bufferLayouts); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/UInt1Vector.java000066400000000000000000000265501515451615200304760ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.UInt1ReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableUInt1Holder; import org.apache.arrow.vector.holders.UInt1Holder; import org.apache.arrow.vector.types.Types.MinorType; 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.ValueVectorUtility; /** * UInt1Vector implements a fixed width (1 bytes) vector of integer values which could be null. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public final class UInt1Vector extends BaseFixedWidthVector implements BaseIntVector, ValueIterableVector { /** The mask to use when promoting the unsigned byte value to an integer. */ public static final int PROMOTION_MASK = 0xFF; /** The maximum 8-bit unsigned integer. */ public static final byte MAX_UINT1 = (byte) 0XFF; public static final byte TYPE_WIDTH = 1; public UInt1Vector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.UINT1.getType()), allocator); } public UInt1Vector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Constructor for UInt1Vector. * * @param field Field type * @param allocator Allocator type */ public UInt1Vector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new UInt1ReaderImpl(UInt1Vector.this); } @Override public MinorType getMinorType() { return MinorType.UINT1; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    To avoid overflow, the returned type is one step up from the signed type. * *

    This method is mainly meant for integration tests. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static short getNoOverflow(final ArrowBuf buffer, final int index) { byte b = buffer.getByte(index * TYPE_WIDTH); return (short) (PROMOTION_MASK & b); } /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public byte get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getByte(index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableUInt1Holder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getByte(index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Byte getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getByte(index * TYPE_WIDTH); } } /** * Returns the value stored at index without the potential for overflow. * * @param index position of element * @return element at given index */ public Short getObjectNoOverflow(int index) { if (isSet(index) == 0) { return null; } else { return getNoOverflow(valueBuffer, index); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, int value) { valueBuffer.setByte(index * TYPE_WIDTH, value); } private void setValue(int index, byte value) { valueBuffer.setByte(index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, int value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, byte value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableUInt1Holder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, UInt1Holder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, int)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, int value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, byte)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, byte value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableUInt1Holder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableUInt1Holder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, UInt1Holder)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, UInt1Holder holder) { handleSafe(index); set(index, holder); } /** * Sets the value at index to value isSet > 0, otherwise sets the index position to invalid/null. */ public void set(int index, int isSet, byte value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, byte)} but will reallocate the buffer if index is larger than * current capacity. */ public void setSafe(int index, int isSet, byte value) { handleSafe(index); set(index, isSet, value); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((UInt1Vector) to); } @Override public void setWithPossibleTruncate(int index, long value) { this.setSafe(index, (int) value); } @Override public void setUnsafeWithPossibleTruncate(int index, long value) { this.set(index, (int) value); } @Override public long getValueAsLong(int index) { return this.get(index) & PROMOTION_MASK; } @Override public String toString() { return ValueVectorUtility.getToString( this, 0, getValueCount(), (v, i) -> v.getObjectNoOverflow(i)); } private class TransferImpl implements TransferPair { UInt1Vector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new UInt1Vector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new UInt1Vector(field, allocator); } public TransferImpl(UInt1Vector to) { this.to = to; } @Override public UInt1Vector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, UInt1Vector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/UInt2Vector.java000066400000000000000000000255751515451615200305050ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.UInt2ReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableUInt2Holder; import org.apache.arrow.vector.holders.UInt2Holder; import org.apache.arrow.vector.types.Types.MinorType; 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.ValueVectorUtility; /** * UInt2Vector implements a fixed width (2 bytes) vector of integer values which could be null. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public final class UInt2Vector extends BaseFixedWidthVector implements BaseIntVector, ValueIterableVector { /** The maximum 16-bit unsigned integer. */ public static final char MAX_UINT2 = (char) 0XFFFF; public static final byte TYPE_WIDTH = 2; public UInt2Vector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.UINT2.getType()), allocator); } public UInt2Vector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Constructor for UInt2Vector type. * * @param field Field type * @param allocator Allocator type */ public UInt2Vector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new UInt2ReaderImpl(UInt2Vector.this); } @Override public MinorType getMinorType() { return MinorType.UINT2; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    This method is mainly meant for integration tests. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static char get(final ArrowBuf buffer, final int index) { return buffer.getChar((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public char get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getChar((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableUInt2Holder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getChar((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Character getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getChar((long) index * TYPE_WIDTH); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, int value) { valueBuffer.setChar((long) index * TYPE_WIDTH, value); } private void setValue(int index, char value) { valueBuffer.setChar((long) index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, int value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, char value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableUInt2Holder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, UInt2Holder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, int)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, int value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, char)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, char value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableUInt2Holder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableUInt2Holder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, UInt2Holder)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, UInt2Holder holder) { handleSafe(index); set(index, holder); } /** * Sets the given index to value is isSet is positive, otherwise sets the position as * invalid/null. */ public void set(int index, int isSet, char value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, char)} but will reallocate the buffer if index is larger than * current capacity. */ public void setSafe(int index, int isSet, char value) { handleSafe(index); set(index, isSet, value); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((UInt2Vector) to); } @Override public void setWithPossibleTruncate(int index, long value) { this.setSafe(index, (int) value); } @Override public void setUnsafeWithPossibleTruncate(int index, long value) { this.set(index, (int) value); } @Override public long getValueAsLong(int index) { return this.get(index); } @Override public String toString() { return ValueVectorUtility.getToString( this, 0, getValueCount(), (v, i) -> v.isNull(i) ? "null" : Integer.toString(v.get(i) & 0x0000ffff)); } private class TransferImpl implements TransferPair { UInt2Vector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new UInt2Vector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new UInt2Vector(field, allocator); } public TransferImpl(UInt2Vector to) { this.to = to; } @Override public UInt2Vector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, UInt2Vector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/UInt4Vector.java000066400000000000000000000252061515451615200304760ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.UInt4ReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableUInt4Holder; import org.apache.arrow.vector.holders.UInt4Holder; import org.apache.arrow.vector.types.Types.MinorType; 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.ValueVectorUtility; /** * UInt4Vector implements a fixed width (4 bytes) vector of integer values which could be null. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public final class UInt4Vector extends BaseFixedWidthVector implements BaseIntVector, ValueIterableVector { /** The mask to use when promoting the unsigned int value to a long int. */ public static final long PROMOTION_MASK = 0x00000000FFFFFFFFL; /** The maximum 32-bit unsigned integer. */ public static final int MAX_UINT4 = 0XFFFFFFFF; public static final byte TYPE_WIDTH = 4; public UInt4Vector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.UINT4.getType()), allocator); } public UInt4Vector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Constructor for UInt4Vector. * * @param field Field type * @param allocator Allocator type */ public UInt4Vector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new UInt4ReaderImpl(UInt4Vector.this); } @Override public MinorType getMinorType() { return MinorType.UINT4; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    To avoid overflow, the returned type is one step up from the signed type. * *

    This method is mainly meant for integration tests. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static long getNoOverflow(final ArrowBuf buffer, final int index) { long l = buffer.getInt((long) index * TYPE_WIDTH); return PROMOTION_MASK & l; } /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public int get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getInt((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableUInt4Holder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getInt((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Integer getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getInt((long) index * TYPE_WIDTH); } } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ public Long getObjectNoOverflow(int index) { if (isSet(index) == 0) { return null; } else { return getNoOverflow(valueBuffer, index); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, int value) { valueBuffer.setInt((long) index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, int value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableUInt4Holder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, UInt4Holder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, int)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, int value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableUInt4Holder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableUInt4Holder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, UInt4Holder)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, UInt4Holder holder) { handleSafe(index); set(index, holder); } /** * Sets the value at index to value isSet > 0, otherwise sets the index position to invalid/null. */ public void set(int index, int isSet, int value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, int)} but will reallocate if the buffer if index is larger than * the current capacity. */ public void setSafe(int index, int isSet, int value) { handleSafe(index); set(index, isSet, value); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((UInt4Vector) to); } @Override public void setWithPossibleTruncate(int index, long value) { this.setSafe(index, (int) value); } @Override public void setUnsafeWithPossibleTruncate(int index, long value) { this.set(index, (int) value); } @Override public long getValueAsLong(int index) { return this.get(index) & PROMOTION_MASK; } @Override public String toString() { return ValueVectorUtility.getToString( this, 0, getValueCount(), (v, i) -> v.getObjectNoOverflow(i)); } private class TransferImpl implements TransferPair { UInt4Vector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new UInt4Vector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new UInt4Vector(field, allocator); } public TransferImpl(UInt4Vector to) { this.to = to; } @Override public UInt4Vector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, UInt4Vector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/UInt8Vector.java000066400000000000000000000252441515451615200305040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import java.math.BigInteger; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.impl.UInt8ReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableUInt8Holder; import org.apache.arrow.vector.holders.UInt8Holder; import org.apache.arrow.vector.types.Types.MinorType; 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.ValueVectorUtility; /** * UInt8Vector implements a fixed width vector (8 bytes) of integer values which could be null. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public final class UInt8Vector extends BaseFixedWidthVector implements BaseIntVector, ValueIterableVector { /** The maximum 64-bit unsigned long integer. */ public static final long MAX_UINT8 = 0XFFFFFFFFFFFFFFFFL; public static final byte TYPE_WIDTH = 8; public UInt8Vector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.UINT8.getType()), allocator); } public UInt8Vector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Constructor for UInt8Vector. * * @param field Field type * @param allocator Allocator type. */ public UInt8Vector(Field field, BufferAllocator allocator) { super(field, allocator, TYPE_WIDTH); } @Override protected FieldReader getReaderImpl() { return new UInt8ReaderImpl(UInt8Vector.this); } @Override public MinorType getMinorType() { return MinorType.UINT8; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ private static final BigInteger SAFE_CONVERSION_MASK = new BigInteger("ffffffffffffffff", 16); /** * Given a data buffer, get the value stored at a particular position in the vector. * *

    To avoid overflow, the returned type is one step up from the signed type. * *

    This method is mainly meant for integration tests. * * @param buffer data buffer * @param index position of the element. * @return value stored at the index. */ public static BigInteger getNoOverflow(final ArrowBuf buffer, final int index) { BigInteger l = BigInteger.valueOf(buffer.getLong((long) index * TYPE_WIDTH)); return SAFE_CONVERSION_MASK.and(l); } /** * Get the element at the given index from the vector. * * @param index position of element * @return element at given index */ public long get(int index) throws IllegalStateException { if (NULL_CHECKING_ENABLED && isSet(index) == 0) { throw new IllegalStateException("Value at index is null"); } return valueBuffer.getLong((long) index * TYPE_WIDTH); } /** * Get the element at the given index from the vector and sets the state in holder. If element at * given index is null, holder.isSet will be zero. * * @param index position of element */ public void get(int index, NullableUInt8Holder holder) { if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.value = valueBuffer.getLong((long) index * TYPE_WIDTH); } /** * Same as {@link #get(int)}. * * @param index position of element * @return element at given index */ @Override public Long getObject(int index) { if (isSet(index) == 0) { return null; } else { return valueBuffer.getLong((long) index * TYPE_WIDTH); } } /** * Returns the value stored at index without the potential for overflow. * * @param index position of element * @return element at given index */ public BigInteger getObjectNoOverflow(int index) { if (isSet(index) == 0) { return null; } else { return getNoOverflow(valueBuffer, index); } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ private void setValue(int index, long value) { valueBuffer.setLong((long) index * TYPE_WIDTH, value); } /** * Set the element at the given index to the given value. * * @param index position of element * @param value value of element */ public void set(int index, long value) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, value); } /** * Set the element at the given index to the value set in data holder. If the value in holder is * not indicated as set, element in the at the given index will be null. * * @param index position of element * @param holder nullable data holder for value of element */ public void set(int index, NullableUInt8Holder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); } else if (holder.isSet > 0) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Set the element at the given index to the value set in data holder. * * @param index position of element * @param holder data holder for value of element */ public void set(int index, UInt8Holder holder) { BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } /** * Same as {@link #set(int, long)} except that it handles the case when index is greater than or * equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param value value of element */ public void setSafe(int index, long value) { handleSafe(index); set(index, value); } /** * Same as {@link #set(int, NullableUInt8Holder)} except that it handles the case when index is * greater than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder nullable data holder for value of element */ public void setSafe(int index, NullableUInt8Holder holder) throws IllegalArgumentException { handleSafe(index); set(index, holder); } /** * Same as {@link #set(int, UInt8Holder)} except that it handles the case when index is greater * than or equal to existing value capacity {@link #getValueCapacity()}. * * @param index position of element * @param holder data holder for value of element */ public void setSafe(int index, UInt8Holder holder) { handleSafe(index); set(index, holder); } /** Sets value at index is isSet is positive otherwise sets the index to invalid/null. */ public void set(int index, int isSet, long value) { if (isSet > 0) { set(index, value); } else { BitVectorHelper.unsetBit(validityBuffer, index); } } /** * Same as {@link #set(int, int, long)} but will reallocate if index is greater than current * capacity. */ public void setSafe(int index, int isSet, long value) { handleSafe(index); set(index, isSet, value); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param field Field object used by the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((UInt8Vector) to); } @Override public void setWithPossibleTruncate(int index, long value) { this.setSafe(index, value); } @Override public void setUnsafeWithPossibleTruncate(int index, long value) { this.set(index, value); } @Override public long getValueAsLong(int index) { return this.get(index); } @Override public String toString() { return ValueVectorUtility.getToString( this, 0, getValueCount(), (v, i) -> v.getObjectNoOverflow(i)); } private class TransferImpl implements TransferPair { UInt8Vector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new UInt8Vector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new UInt8Vector(field, allocator); } public TransferImpl(UInt8Vector to) { this.to = to; } @Override public UInt8Vector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, UInt8Vector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/UuidVector.java000066400000000000000000000341731515451615200304440ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ValueIterableVector.java000066400000000000000000000031341515451615200322530ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.Iterator; public interface ValueIterableVector extends ValueVector { /** * Get an Iterable that can be used to iterate over the values in the vector. * * @return an Iterable for the vector's values */ default Iterator getValueIterator() { return new Iterator() { private int index = 0; @Override public boolean hasNext() { return index < ValueIterableVector.this.getValueCount(); } @Override public T next() { return (T) ValueIterableVector.this.getObject(index++); } }; } /** * Get an Iterator for the values in the vector. * * @return an Iterator for the values in the vector */ default Iterable getValueIterable() { return this::getValueIterator; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ValueVector.java000066400000000000000000000260151515451615200306060ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.io.Closeable; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.OutOfMemoryException; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.util.CallBack; import org.apache.arrow.vector.util.TransferPair; import org.apache.arrow.vector.util.ValueVectorUtility; /** * An abstraction that is used to store a sequence of values in an individual column. * *

    A {@link ValueVector value vector} stores underlying data in-memory in a columnar fashion that * is compact and efficient. The column whose data is stored, is referred by {@link #getField()}. * *

    It is important that vector is allocated before attempting to read or write. * *

    There are a few "rules" around vectors: * *

      *
    • 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. *
    * *

    Please note that the current implementation doesn't enforce those rules, hence we may find few * places that deviate from these rules (e.g. offset vectors in Variable Length and Repeated vector) * *

    This interface "should" strive to guarantee this order of operation: * *

    * * allocate > mutate > setvaluecount > access > clear (or allocate to start the process * over). * *
    */ public interface ValueVector extends Closeable, Iterable { /** * Allocate new buffers. ValueVector implements logic to determine how much to allocate. * * @throws OutOfMemoryException Thrown if no memory can be allocated. */ void allocateNew() throws OutOfMemoryException; /** * Allocates new buffers. ValueVector implements logic to determine how much to allocate. * * @return Returns true if allocation was successful. */ boolean allocateNewSafe(); /** * Allocate new buffer with double capacity, and copy data into the new buffer. Replace vector's * buffer with new buffer, and release old one */ void reAlloc(); /** * Get the allocator associated with the vector. CAVEAT: Some ValueVector subclasses (e.g. * NullVector) do not require an allocator for data storage and may return null. * * @return Returns nullable allocator. */ BufferAllocator getAllocator(); /** * Set the initial record capacity. * * @param numRecords the initial record capacity. */ void setInitialCapacity(int numRecords); /** * Returns the maximum number of values that can be stored in this vector instance. * * @return the maximum number of values that can be stored in this vector instance. */ int getValueCapacity(); /** Alternative to clear(). Allows use as an AutoCloseable in try-with-resources. */ @Override void close(); /** * Release any owned ArrowBuf and reset the ValueVector to the initial state. If the vector has * any child vectors, they will also be cleared. */ void clear(); /** * Reset the ValueVector to the initial state without releasing any owned ArrowBuf. Buffer * capacities will remain unchanged and any previous data will be zeroed out. This includes * buffers for data, validity, offset, etc. If the vector has any child vectors, they will also be * reset. */ void reset(); /** * Get information about how this field is materialized. * * @return the field corresponding to this vector */ Field getField(); MinorType getMinorType(); /** * To transfer quota responsibility. * * @param allocator the target allocator * @return a {@link org.apache.arrow.vector.util.TransferPair transfer pair}, creating a new * target vector of the same type. */ TransferPair getTransferPair(BufferAllocator allocator); /** * To transfer quota responsibility. * * @param ref the name of the vector * @param allocator the target allocator * @return a {@link org.apache.arrow.vector.util.TransferPair transfer pair}, creating a new * target vector of the same type. */ TransferPair getTransferPair(String ref, BufferAllocator allocator); /** * To transfer quota responsibility. * * @param field the Field object used by the target vector * @param allocator the target allocator * @return a {@link org.apache.arrow.vector.util.TransferPair transfer pair}, creating a new * target vector of the same type. */ TransferPair getTransferPair(Field field, BufferAllocator allocator); /** * To transfer quota responsibility. * * @param ref the name of the vector * @param allocator the target allocator * @param callBack A schema change callback. * @return a {@link org.apache.arrow.vector.util.TransferPair transfer pair}, creating a new * target vector of the same type. */ TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack); /** * To transfer quota responsibility. * * @param field the Field object used by the target vector * @param allocator the target allocator * @param callBack A schema change callback. * @return a {@link org.apache.arrow.vector.util.TransferPair transfer pair}, creating a new * target vector of the same type. */ TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack); /** * Makes a new transfer pair used to transfer underlying buffers. * * @param target the target for the transfer * @return a new {@link org.apache.arrow.vector.util.TransferPair transfer pair} that is used to * transfer underlying buffers into the target vector. */ TransferPair makeTransferPair(ValueVector target); /** * Get a reader for this vector. * * @return a {@link org.apache.arrow.vector.complex.reader.FieldReader field reader} that supports * reading values from this vector. */ FieldReader getReader(); /** * Get the number of bytes used by this vector. * * @return the number of bytes that is used by this vector instance. */ int getBufferSize(); /** * Returns the number of bytes that is used by this vector if it holds the given number of values. * The result will be the same as if setValueCount() were called, followed by calling * getBufferSize(), but without any of the closing side-effects that setValueCount() implies wrt * finishing off the population of a vector. Some operations might wish to use this to determine * how much memory has been used by a vector so far, even though it is not finished being * populated. * * @param valueCount the number of values to assume this vector contains * @return the buffer size if this vector is holding valueCount values */ int getBufferSizeFor(int valueCount); /** * Return the underlying buffers associated with this vector. Note that this doesn't impact the * reference counts for this buffer so it only should be used for in-context access. Also note * that this buffer changes regularly thus external classes shouldn't hold a reference to it * (unless they change it). * * @param clear Whether to clear vector before returning; the buffers will still be refcounted; * but the returned array will be the only reference to them * @return The underlying {@link ArrowBuf buffers} that is used by this vector instance. */ ArrowBuf[] getBuffers(boolean clear); /** * Gets the underlying buffer associated with validity vector. * * @return buffer */ ArrowBuf getValidityBuffer(); /** * Gets the underlying buffer associated with data vector. * * @return buffer */ ArrowBuf getDataBuffer(); /** * Gets the underlying buffer associated with offset vector. * * @return buffer */ ArrowBuf getOffsetBuffer(); /** * Gets the number of values. * * @return number of values in the vector */ int getValueCount(); /** Set number of values in the vector. */ void setValueCount(int valueCount); /** * Get friendly type object from the vector. * * @param index index of object to get * @return friendly type object, null if value is unset */ Object getObject(int index); /** * Returns number of null elements in the vector. * * @return number of null elements */ int getNullCount(); /** * Check whether an element in the vector is null. * * @param index index to check for null * @return true if element is null */ boolean isNull(int index); /** Returns hashCode of element in index with the default hasher. */ int hashCode(int index); /** Returns hashCode of element in index with the given hasher. */ int hashCode(int index, ArrowBufHasher hasher); /** * Copy a cell value from a particular index in source vector to a particular position in this * vector. * * @param fromIndex position to copy from in source vector * @param thisIndex position to copy to in this vector * @param from source vector */ void copyFrom(int fromIndex, int thisIndex, ValueVector from); /** * Same as {@link #copyFrom(int, int, ValueVector)} except that it handles the case when the * capacity of the vector needs to be expanded before copy. * * @param fromIndex position to copy from in source vector * @param thisIndex position to copy to in this vector * @param from source vector */ void copyFromSafe(int fromIndex, int thisIndex, ValueVector from); /** * Accept a generic {@link VectorVisitor} and return the result. * * @param the output result type. * @param the input data together with visitor. */ OUT accept(VectorVisitor visitor, IN value); /** * Gets the name of the vector. * * @return the name of the vector. */ String getName(); default void validate() { ValueVectorUtility.validate(this); } default void validateFull() { ValueVectorUtility.validateFull(this); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/VarBinaryVector.java000066400000000000000000000257511515451615200314350ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReusableBuffer; import org.apache.arrow.vector.complex.impl.VarBinaryReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableVarBinaryHolder; import org.apache.arrow.vector.holders.VarBinaryHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * VarBinaryVector implements a variable width vector of binary values which could be NULL. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public final class VarBinaryVector extends BaseVariableWidthVector implements ValueIterableVector { /** * Instantiate a VarBinaryVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public VarBinaryVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.VARBINARY.getType()), allocator); } /** * Instantiate a VarBinaryVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public VarBinaryVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a VarBinaryVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public VarBinaryVector(Field field, BufferAllocator allocator) { super(field, allocator); } @Override protected FieldReader getReaderImpl() { return new VarBinaryReaderImpl(VarBinaryVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.VARBINARY; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the variable length element at specified index as byte array. * * @param index position of element to get * @return array of bytes for non-null element, null otherwise */ public byte[] get(int index) { assert index >= 0; if (NULL_CHECKING_ENABLED && 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; } /** * Read the value at the given position to the given output buffer. The caller is responsible for * checking for nullity first. * * @param index position of element. * @param buffer the buffer to write into. */ @Override public void read(int index, ReusableBuffer buffer) { final int startOffset = getStartOffset(index); final int dataLength = getEndOffset(index) - startOffset; buffer.set(valueBuffer, startOffset, dataLength); } /** * Get the variable length element at specified index as Text. * * @param index position of element to get * @return byte array for non-null element, null otherwise */ @Override public byte[] getObject(int 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; } /** * Get the variable length element at specified index and sets the state in provided holder. * * @param index position of element to get * @param holder data holder to be populated by this function */ public void get(int index, NullableVarBinaryHolder holder) { assert index >= 0; if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.start = getStartOffset(index); holder.end = getEndOffset(index); holder.buffer = valueBuffer; } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the variable length element at the specified index to the data buffer supplied in the * holder. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void set(int index, VarBinaryHolder holder) { assert index >= 0; fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); final int dataLength = holder.end - holder.start; final int startOffset = getStartOffset(index); offsetBuffer.setInt((index + 1) * ((long) OFFSET_WIDTH), startOffset + dataLength); valueBuffer.setBytes(startOffset, holder.buffer, holder.start, dataLength); lastSet = index; } /** * Same as {@link #set(int, VarBinaryHolder)} except that it handles the case where index and * length of new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void setSafe(int index, VarBinaryHolder holder) { assert index >= 0; final int dataLength = holder.end - holder.start; handleSafe(index, dataLength); fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); final int startOffset = getStartOffset(index); offsetBuffer.setInt((index + 1) * ((long) OFFSET_WIDTH), startOffset + dataLength); valueBuffer.setBytes(startOffset, holder.buffer, holder.start, dataLength); lastSet = index; } /** * Set the variable length element at the specified index to the data buffer supplied in the * holder. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void set(int index, NullableVarBinaryHolder holder) { assert index >= 0; fillHoles(index); BitVectorHelper.setValidityBit(validityBuffer, index, holder.isSet); final int startOffset = getStartOffset(index); if (holder.isSet != 0) { final int dataLength = holder.end - holder.start; offsetBuffer.setInt((index + 1) * ((long) OFFSET_WIDTH), startOffset + dataLength); valueBuffer.setBytes(startOffset, holder.buffer, holder.start, dataLength); } else { offsetBuffer.setInt((index + 1) * ((long) OFFSET_WIDTH), startOffset); } lastSet = index; } /** * Same as {@link #set(int, NullableVarBinaryHolder)} except that it handles the case where index * and length of new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void setSafe(int index, NullableVarBinaryHolder holder) { assert index >= 0; if (holder.isSet != 0) { final int dataLength = holder.end - holder.start; handleSafe(index, dataLength); fillHoles(index); final int startOffset = getStartOffset(index); offsetBuffer.setInt((index + 1) * ((long) OFFSET_WIDTH), startOffset + dataLength); valueBuffer.setBytes(startOffset, holder.buffer, holder.start, dataLength); } else { fillEmpties(index + 1); } BitVectorHelper.setValidityBit(validityBuffer, index, holder.isSet); lastSet = index; } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising of this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((VarBinaryVector) to); } private class TransferImpl implements TransferPair { VarBinaryVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new VarBinaryVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new VarBinaryVector(field, allocator); } public TransferImpl(VarBinaryVector to) { this.to = to; } @Override public VarBinaryVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, VarBinaryVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/VarCharVector.java000066400000000000000000000276051515451615200310660ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReusableBuffer; import org.apache.arrow.vector.complex.impl.VarCharReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableVarCharHolder; import org.apache.arrow.vector.holders.VarCharHolder; import org.apache.arrow.vector.types.Types.MinorType; 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; import org.apache.arrow.vector.validate.ValidateUtil; /** * VarCharVector implements a variable width vector of VARCHAR values which could be NULL. A * validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public final class VarCharVector extends BaseVariableWidthVector implements ValueIterableVector { /** * Instantiate a VarCharVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public VarCharVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.VARCHAR.getType()), allocator); } /** * Instantiate a VarCharVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public VarCharVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a VarCharVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public VarCharVector(Field field, BufferAllocator allocator) { super(field, allocator); } @Override protected FieldReader getReaderImpl() { return new VarCharReaderImpl(VarCharVector.this); } /** * Get minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.VARCHAR; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the variable length element at specified index as byte array. * * @param index position of element to get * @return array of bytes for non-null element, null otherwise */ public byte[] get(int index) { assert index >= 0; if (NULL_CHECKING_ENABLED && 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; } /** * Get the variable length element at specified index as Text. * * @param index position of element to get * @return Text object for non-null element, null otherwise */ @Override public Text getObject(int index) { assert index >= 0; if (isSet(index) == 0) { return null; } final Text result = new Text(); read(index, result); return result; } /** * Read the value at the given position to the given output buffer. The caller is responsible for * checking for nullity first. * * @param index position of element. * @param buffer the buffer to write into. */ @Override public void read(int index, ReusableBuffer buffer) { final int startOffset = getStartOffset(index); final int dataLength = getEndOffset(index) - startOffset; buffer.set(valueBuffer, startOffset, dataLength); } /** * Get the variable length element at specified index and sets the state in provided holder. * * @param index position of element to get * @param holder data holder to be populated by this function */ public void get(int index, NullableVarCharHolder holder) { assert index >= 0; if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; holder.start = getStartOffset(index); holder.end = getEndOffset(index); holder.buffer = valueBuffer; } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the variable length element at the specified index to the data buffer supplied in the * holder. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void set(int index, VarCharHolder holder) { assert index >= 0; fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); final int dataLength = holder.end - holder.start; final int startOffset = getStartOffset(index); offsetBuffer.setInt((index + 1) * ((long) OFFSET_WIDTH), startOffset + dataLength); valueBuffer.setBytes(startOffset, holder.buffer, holder.start, dataLength); lastSet = index; } /** * Same as {@link #set(int, VarCharHolder)} except that it handles the case where index and length * of new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void setSafe(int index, VarCharHolder holder) { assert index >= 0; final int dataLength = holder.end - holder.start; handleSafe(index, dataLength); fillHoles(index); BitVectorHelper.setBit(validityBuffer, index); final int startOffset = getStartOffset(index); offsetBuffer.setInt((index + 1) * ((long) OFFSET_WIDTH), startOffset + dataLength); valueBuffer.setBytes(startOffset, holder.buffer, holder.start, dataLength); lastSet = index; } /** * Set the variable length element at the specified index to the data buffer supplied in the * holder. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void set(int index, NullableVarCharHolder holder) { assert index >= 0; fillHoles(index); BitVectorHelper.setValidityBit(validityBuffer, index, holder.isSet); final int startOffset = getStartOffset(index); if (holder.isSet != 0) { final int dataLength = holder.end - holder.start; offsetBuffer.setInt((index + 1) * ((long) OFFSET_WIDTH), startOffset + dataLength); valueBuffer.setBytes(startOffset, holder.buffer, holder.start, dataLength); } else { offsetBuffer.setInt((index + 1) * ((long) OFFSET_WIDTH), startOffset); } lastSet = index; } /** * Same as {@link #set(int, NullableVarCharHolder)} except that it handles the case where index * and length of new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void setSafe(int index, NullableVarCharHolder holder) { assert index >= 0; if (holder.isSet != 0) { final int dataLength = holder.end - holder.start; handleSafe(index, dataLength); fillHoles(index); final int startOffset = getStartOffset(index); offsetBuffer.setInt((index + 1) * ((long) OFFSET_WIDTH), startOffset + dataLength); valueBuffer.setBytes(startOffset, holder.buffer, holder.start, dataLength); } else { fillEmpties(index + 1); } BitVectorHelper.setValidityBit(validityBuffer, index, holder.isSet); lastSet = index; } /** * Set the variable length element at the specified index to the content in supplied Text. * * @param index position of the element to set * @param text Text object with data */ public void set(int index, Text text) { set(index, text.getBytes(), 0, (int) text.getLength()); } /** * Same as {@link #set(int, NullableVarCharHolder)} except that it handles the case where index * and length of new element are beyond the existing capacity of the vector. * * @param index position of the element to set. * @param text Text object with data */ public void setSafe(int index, Text text) { setSafe(index, text.getBytes(), 0, (int) text.getLength()); } @Override public void validateScalars() { for (int i = 0; i < getValueCount(); ++i) { byte[] value = get(i); if (value != null) { ValidateUtil.validateOrThrow( Text.validateUTF8NoThrow(value), "Non-UTF-8 data in VarCharVector at position " + i + "."); } } } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising of this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((VarCharVector) to); } private class TransferImpl implements TransferPair { VarCharVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new VarCharVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new VarCharVector(field, allocator); } public TransferImpl(VarCharVector to) { this.to = to; } @Override public VarCharVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, VarCharVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/VariableWidthFieldVector.java000066400000000000000000000112161515451615200332200ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.arrow.vector; import java.nio.ByteBuffer; import org.apache.arrow.memory.ReusableBuffer; /** A base interface for common functionalities in variable width vectors. */ public interface VariableWidthFieldVector extends VariableWidthVector, FieldVector, VectorDefinitionSetter { /** * Set the variable length element at the specified index to the supplied byte array. * * @param index position of the element to set * @param value array of bytes with data */ void set(int index, byte[] value); /** * Set the variable length element at the specified index to the supplied byte array. * * @param index position of the element to set * @param value array of bytes with data * @param start start position in the array * @param length length of the data to write */ void set(int index, byte[] value, int start, int length); /** * Set the variable length element at the specified index to the supplied ByteBuffer. * * @param index position of the element to set * @param value ByteBuffer with data * @param start start position in the ByteBuffer * @param length length of the data to write */ void set(int index, ByteBuffer value, int start, int length); /** * Set the variable length element at the specified index to the supplied byte array, and it * handles the case where index and length of a new element are beyond the existing capacity of * the vector. * * @param index position of the element to set * @param value array of bytes to write */ void setSafe(int index, byte[] value); /** * Set the variable length element at the specified index to the supplied byte array, and it * handles the case where index and length of a new element are beyond the existing capacity. * * @param index position of the element to set * @param value array of bytes with data * @param start start position in the array * @param length length of the data to write */ void setSafe(int index, byte[] value, int start, int length); /** * Set the variable length element at the specified index to the supplied ByteBuffer, and it * handles the case where index and length of a new element are beyond the existing capacity. * * @param index position of the element to set * @param value ByteBuffer with data * @param start start position in the ByteBuffer * @param length length of the data to write */ void setSafe(int index, ByteBuffer value, int start, int length); /** * Get the variable length element at the specified index. * * @param index position of the element to get * @return byte array with the data */ byte[] get(int index); /** * Get the variable length element at the specified index using a ReusableBuffer. * * @param index position of the element to get * @param buffer ReusableBuffer to write the data to */ void read(int index, ReusableBuffer buffer); /** * Get the index of the last non-null element in the vector. * * @return index of the last non-null element */ int getLastSet(); /** * Set the index of the last non-null element in the vector. * * @param value desired index of last non-null element */ void setLastSet(int value); /** * Get the variable length element at specified index as Text. * * @param index position of an element to get * @return greater than length 0 for a non-null element, 0 otherwise */ int getValueLength(int index); /** * Create holes in the vector upto the given index (exclusive). Holes will be created from the * current last-set position in the vector. * * @param index target index */ void fillEmpties(int index); /** * Sets the value length for an element. * * @param index position of the element to set * @param length length of the element */ void setValueLengthSafe(int index, int length); } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/VariableWidthVector.java000066400000000000000000000036211515451615200322550ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** Interface vectors that contain variable width members (e.g. Strings, Lists, etc). */ public interface VariableWidthVector extends ElementAddressableVector, DensityAwareVector { /** * Allocate a new memory space for this vector. Must be called prior to using the ValueVector. * * @param totalBytes Desired size of the underlying data buffer. * @param valueCount Number of values in the vector. */ void allocateNew(long totalBytes, int valueCount); /** * Allocate a new memory space for this vector. Must be called prior to using the ValueVector. The * initial size in bytes is either default (or) reused from previous allocation * * @param valueCount Number of values in the vector. */ void allocateNew(int valueCount); /** * Provide the maximum amount of variable width bytes that can be stored in this vector. * * @return the byte capacity of this vector */ int getByteCapacity(); /** * Provide the number of bytes contained in the valueBuffer. * * @return the number of bytes in valueBuffer. */ int sizeOfValueBuffer(); } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/VectorDefinitionSetter.java000066400000000000000000000017461515451615200330150ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** Interface for setting a specific index values as defined/valid on a vector. */ public interface VectorDefinitionSetter { void setIndexDefined(int index); } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java000066400000000000000000000147741515451615200307510ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.Preconditions.checkArgument; import java.util.ArrayList; import java.util.Collections; import java.util.Iterator; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.util.Collections2; import org.apache.arrow.vector.compression.CompressionCodec; import org.apache.arrow.vector.compression.CompressionUtil; import org.apache.arrow.vector.compression.NoCompressionCodec; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.types.pojo.Field; /** Loads buffers into vectors. */ public class VectorLoader { private final VectorSchemaRoot root; private final CompressionCodec.Factory factory; /** * A flag indicating if decompression is needed. This will affect the behavior of releasing * buffers. */ private boolean decompressionNeeded; /** * Construct with a root to load and will create children in root based on schema. * * @param root the root to add vectors to based on schema */ public VectorLoader(VectorSchemaRoot root) { this(root, CompressionCodec.Factory.INSTANCE); } /** * Construct with a root to load and will create children in root based on schema. * * @param root the root to add vectors to based on schema. * @param factory the factory to create codec. */ public VectorLoader(VectorSchemaRoot root, CompressionCodec.Factory factory) { this.root = root; this.factory = factory; } /** * Loads the record batch in the vectors. will not close the record batch * * @param recordBatch the batch to load */ public void load(ArrowRecordBatch recordBatch) { Iterator buffers = recordBatch.getBuffers().iterator(); Iterator nodes = recordBatch.getNodes().iterator(); CompressionUtil.CodecType codecType = CompressionUtil.CodecType.fromCompressionType(recordBatch.getBodyCompression().getCodec()); decompressionNeeded = codecType != CompressionUtil.CodecType.NO_COMPRESSION; CompressionCodec codec = decompressionNeeded ? factory.createCodec(codecType) : NoCompressionCodec.INSTANCE; Iterator variadicBufferCounts = Collections.emptyIterator(); ; if (recordBatch.getVariadicBufferCounts() != null && !recordBatch.getVariadicBufferCounts().isEmpty()) { variadicBufferCounts = recordBatch.getVariadicBufferCounts().iterator(); } for (FieldVector fieldVector : root.getFieldVectors()) { loadBuffers(fieldVector, fieldVector.getField(), buffers, nodes, codec, variadicBufferCounts); } root.setRowCount(recordBatch.getLength()); if (nodes.hasNext() || buffers.hasNext() || variadicBufferCounts.hasNext()) { throw new IllegalArgumentException( "not all nodes, buffers and variadicBufferCounts were consumed. nodes: " + Collections2.toString(nodes) + " buffers: " + Collections2.toString(buffers) + " variadicBufferCounts: " + Collections2.toString(variadicBufferCounts)); } } private void loadBuffers( FieldVector vector, Field field, Iterator buffers, Iterator nodes, CompressionCodec codec, Iterator variadicBufferCounts) { checkArgument(nodes.hasNext(), "no more field nodes for field %s and vector %s", field, vector); ArrowFieldNode fieldNode = nodes.next(); // variadicBufferLayoutCount will be 0 for vectors of a type except BaseVariableWidthViewVector long variadicBufferLayoutCount = 0; if (vector instanceof BaseVariableWidthViewVector) { if (variadicBufferCounts.hasNext()) { variadicBufferLayoutCount = variadicBufferCounts.next(); } else { throw new IllegalStateException( "No variadicBufferCounts available for BaseVariableWidthViewVector"); } } int bufferLayoutCount = (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 = nextBuf.writerIndex() > 0 ? codec.decompress(vector.getAllocator(), nextBuf) : nextBuf; ownBuffers.add(bufferToAdd); if (decompressionNeeded) { // decompression performed nextBuf.getReferenceManager().retain(); } } try { vector.loadFieldBuffers(fieldNode, ownBuffers); if (decompressionNeeded) { for (ArrowBuf buf : ownBuffers) { buf.close(); } } } catch (RuntimeException e) { throw new IllegalArgumentException( "Could not load buffers for field " + field + ". error message: " + e.getMessage(), e); } List children = field.getChildren(); if (children.size() > 0) { List childrenFromFields = vector.getChildrenFromFields(); checkArgument( children.size() == childrenFromFields.size(), "should have as many children as in the schema: found %s expected %s", childrenFromFields.size(), children.size()); for (int i = 0; i < childrenFromFields.size(); i++) { Field child = children.get(i); FieldVector fieldVector = childrenFromFields.get(i); loadBuffers(fieldVector, child, buffers, nodes, codec, variadicBufferCounts); } } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java000066400000000000000000000352211515451615200315750ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.stream.Collectors; import java.util.stream.StreamSupport; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.compare.ApproxEqualsVisitor; import org.apache.arrow.vector.compare.Range; import org.apache.arrow.vector.compare.VectorEqualsVisitor; import org.apache.arrow.vector.compare.VectorValueEqualizer; import org.apache.arrow.vector.compare.util.ValueEpsilonEqualizers; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.TransferPair; /** * Holder for a set of vectors to be loaded/unloaded. A VectorSchemaRoot is a container that can * hold batches, batches flow through VectorSchemaRoot as part of a pipeline. Note this is different * from other implementations (i.e. in C++ and Python, a RecordBatch is a collection of equal-length * vector instances and was created each time for a new batch). * *

    The recommended usage for VectorSchemaRoot is creating a single VectorSchemaRoot based on the * known schema and populated data over and over into the same VectorSchemaRoot in a stream of * batches rather than create a new VectorSchemaRoot instance each time (see Flight or * ArrowFileWriter for better understanding). Thus at any one point a VectorSchemaRoot may have data * or may have no data (say it was transferred downstream or not yet populated). */ public class VectorSchemaRoot implements AutoCloseable { private Schema schema; private int rowCount; private final List fieldVectors; private final Map fieldVectorsMap = new LinkedHashMap<>(); /** Constructs new instance containing each of the vectors. */ public VectorSchemaRoot(Iterable vectors) { this( StreamSupport.stream(vectors.spliterator(), false) .map(t -> t.getField()) .collect(Collectors.toList()), StreamSupport.stream(vectors.spliterator(), false).collect(Collectors.toList())); } /** Constructs a new instance containing the children of parent but not the parent itself. */ public VectorSchemaRoot(FieldVector parent) { this(parent.getField().getChildren(), parent.getChildrenFromFields(), parent.getValueCount()); } /** * Constructs a new instance. * * @param fields The types of each vector. * @param fieldVectors The data vectors (must be equal in size to fields. */ public VectorSchemaRoot(List fields, List fieldVectors) { this( new Schema(fields), fieldVectors, fieldVectors.size() == 0 ? 0 : fieldVectors.get(0).getValueCount()); } /** * Constructs a new instance. * * @param fields The types of each vector. * @param fieldVectors The data vectors (must be equal in size to fields. * @param rowCount The number of rows contained. */ public VectorSchemaRoot(List fields, List fieldVectors, int rowCount) { this(new Schema(fields), fieldVectors, rowCount); } /** * Constructs a new instance. * * @param schema The schema for the vectors. * @param fieldVectors The data vectors. * @param rowCount The number of rows */ public VectorSchemaRoot(Schema schema, List fieldVectors, int rowCount) { if (schema.getFields().size() != fieldVectors.size()) { throw new IllegalArgumentException( "Fields must match field vectors. Found " + fieldVectors.size() + " vectors and " + schema.getFields().size() + " fields"); } this.schema = schema; this.rowCount = rowCount; this.fieldVectors = fieldVectors; for (int i = 0; i < schema.getFields().size(); ++i) { Field field = schema.getFields().get(i); FieldVector vector = fieldVectors.get(i); fieldVectorsMap.put(field, vector); } } /** Creates a new set of empty vectors corresponding to the given schema. */ public static VectorSchemaRoot create(Schema schema, BufferAllocator allocator) { List fieldVectors = new ArrayList<>(schema.getFields().size()); for (Field field : schema.getFields()) { FieldVector vector = field.createVector(allocator); fieldVectors.add(vector); } if (fieldVectors.size() != schema.getFields().size()) { throw new IllegalArgumentException( "The root vector did not create the right number of children. found " + fieldVectors.size() + " expected " + schema.getFields().size()); } return new VectorSchemaRoot(schema, fieldVectors, 0); } /** Constructs a new instance from vectors. */ public static VectorSchemaRoot of(FieldVector... vectors) { return new VectorSchemaRoot(Arrays.stream(vectors).collect(Collectors.toList())); } /** * Do an adaptive allocation of each vector for memory purposes. Sizes will be based on previously * defined initial allocation for each vector (and subsequent size learned). */ public void allocateNew() { for (FieldVector v : fieldVectors) { v.allocateNew(); } rowCount = 0; } /** * Release all the memory for each vector held in this root. This DOES NOT remove vectors from the * container. */ public void clear() { for (FieldVector v : fieldVectors) { v.clear(); } rowCount = 0; } public List getFieldVectors() { return Collections.unmodifiableList(fieldVectors); } /** * gets a vector by name. * *

    if name occurs multiple times this returns the first inserted entry for name */ public FieldVector getVector(String name) { for (Map.Entry entry : fieldVectorsMap.entrySet()) { if (entry.getKey().getName().equals(name)) { return entry.getValue(); } } return null; } public FieldVector getVector(Field field) { return fieldVectorsMap.get(field); } public FieldVector getVector(int index) { Preconditions.checkArgument(index >= 0 && index < fieldVectors.size()); return fieldVectors.get(index); } /** * Add vector to the record batch, producing a new VectorSchemaRoot. * * @param index field index * @param vector vector to be added. * @return out VectorSchemaRoot with vector added */ public VectorSchemaRoot addVector(int index, FieldVector vector) { Preconditions.checkNotNull(vector); Preconditions.checkArgument(index >= 0 && index <= fieldVectors.size()); List newVectors = new ArrayList<>(); 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)); } } return new VectorSchemaRoot(newVectors); } /** * Remove vector from the record batch, producing a new VectorSchemaRoot. * * @param index field index * @return out VectorSchemaRoot with vector removed */ public VectorSchemaRoot removeVector(int index) { Preconditions.checkArgument(index >= 0 && index < fieldVectors.size()); List newVectors = new ArrayList<>(); for (int i = 0; i < fieldVectors.size(); i++) { if (i != index) { newVectors.add(fieldVectors.get(i)); } } return new VectorSchemaRoot(newVectors); } public Schema getSchema() { return schema; } public int getRowCount() { return rowCount; } /** * Set the row count of all the vectors in this container. Also sets the value count for each root * level contained FieldVector. * * @param rowCount Number of records. */ public void setRowCount(int rowCount) { this.rowCount = rowCount; for (FieldVector v : fieldVectors) { v.setValueCount(rowCount); } } @Override public void close() { try { AutoCloseables.close(fieldVectors); } catch (RuntimeException ex) { throw ex; } catch (Exception ex) { // should never happen since FieldVector.close() doesn't throw IOException throw new RuntimeException(ex); } } private void printRow(StringBuilder sb, List row) { boolean first = true; for (Object v : row) { if (first) { first = false; } else { sb.append("\t"); } sb.append(v); } sb.append("\n"); } /** Returns a tab separated value of vectors (based on their java object representation). */ public String contentToTSVString() { StringBuilder sb = new StringBuilder(); List row = new ArrayList<>(schema.getFields().size()); for (Field field : schema.getFields()) { row.add(field.getName()); } printRow(sb, row); for (int i = 0; i < rowCount; i++) { row.clear(); for (FieldVector v : fieldVectors) { row.add(v.getObject(i)); } printRow(sb, row); } return sb.toString(); } /** * Synchronizes the schema from the current vectors. In some cases, the schema and the actual * vector structure may be different. This can be caused by a promoted writer (For details, please * see {@link org.apache.arrow.vector.complex.impl.PromotableWriter}). For example, when writing * different types of data to a {@link org.apache.arrow.vector.complex.ListVector} may lead to * such a case. When this happens, this method should be called to bring the schema and vector * structure in a synchronized state. * * @return true if the schema is updated, false otherwise. */ public boolean syncSchema() { List oldFields = this.schema.getFields(); List newFields = this.fieldVectors.stream().map(ValueVector::getField).collect(Collectors.toList()); if (!oldFields.equals(newFields)) { this.schema = new Schema(newFields); return true; } return false; } /** * Slice this root from desired index. * * @param index start position of the slice * @return the sliced root */ public VectorSchemaRoot slice(int index) { return slice(index, this.rowCount - index); } /** * Slice this root at desired index and length. * * @param index start position of the slice * @param length length of the slice * @return the sliced root */ public VectorSchemaRoot slice(int index, int length) { Preconditions.checkArgument(index >= 0, "expecting non-negative index"); Preconditions.checkArgument(length >= 0, "expecting non-negative length"); Preconditions.checkArgument(index + length <= rowCount, "index + length should <= rowCount"); List sliceVectors = fieldVectors.stream() .map( v -> { TransferPair transferPair = v.getTransferPair(v.getAllocator()); transferPair.splitAndTransfer(index, length); return (FieldVector) transferPair.getTo(); }) .collect(Collectors.toList()); return new VectorSchemaRoot(sliceVectors); } /** Determine if two VectorSchemaRoots are exactly equal. */ public boolean equals(VectorSchemaRoot other) { if (other == null) { return false; } if (!this.schema.equals(other.schema)) { return false; } if (this.rowCount != other.rowCount) { return false; } for (int i = 0; i < fieldVectors.size(); i++) { FieldVector vector = fieldVectors.get(i); FieldVector otherVector = other.fieldVectors.get(i); if (!VectorEqualsVisitor.vectorEquals(vector, otherVector)) { return false; } } return true; } /** * Determine if two VectorSchemaRoots are approximately equal using the given functions to * calculate difference between float/double values. Note that approx equals are in regards to * floating point values, other values are comparing to exactly equals. * * @param floatDiffFunction function to calculate difference between float values. * @param doubleDiffFunction function to calculate difference between double values. */ public boolean approxEquals( VectorSchemaRoot other, VectorValueEqualizer floatDiffFunction, VectorValueEqualizer doubleDiffFunction) { Preconditions.checkNotNull(floatDiffFunction); Preconditions.checkNotNull(doubleDiffFunction); if (other == null) { return false; } if (!this.schema.equals(other.schema)) { return false; } if (this.rowCount != other.rowCount) { return false; } Range range = new Range(0, 0, 0); for (int i = 0; i < fieldVectors.size(); i++) { FieldVector vector = fieldVectors.get(i); FieldVector otherVector = other.fieldVectors.get(i); if (vector.getValueCount() != otherVector.getValueCount()) { return false; } ApproxEqualsVisitor visitor = new ApproxEqualsVisitor(vector, otherVector, floatDiffFunction, doubleDiffFunction); range.setLength(vector.getValueCount()); if (!visitor.rangeEquals(range)) { return false; } } return true; } /** * Determine if two VectorSchemaRoots are approximately equal using default functions to calculate * difference between float/double values. */ public boolean approxEquals(VectorSchemaRoot other) { VectorValueEqualizer floatDiffFunction = new ValueEpsilonEqualizers.Float4EpsilonEqualizer( ApproxEqualsVisitor.DEFAULT_FLOAT_EPSILON); VectorValueEqualizer doubleDiffFunction = new ValueEpsilonEqualizers.Float8EpsilonEqualizer( ApproxEqualsVisitor.DEFAULT_DOUBLE_EPSILON); return approxEquals(other, floatDiffFunction, doubleDiffFunction); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java000066400000000000000000000127571515451615200313130ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ArrayList; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.compression.CompressionCodec; import org.apache.arrow.vector.compression.CompressionUtil; import org.apache.arrow.vector.compression.NoCompressionCodec; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; /** * Helper class that handles converting a {@link VectorSchemaRoot} to a {@link ArrowRecordBatch}. */ public class VectorUnloader { private final VectorSchemaRoot root; private final boolean includeNullCount; private final CompressionCodec codec; private final boolean alignBuffers; /** Constructs a new instance of the given set of vectors. */ public VectorUnloader(VectorSchemaRoot root) { this(root, true, NoCompressionCodec.INSTANCE, true); } /** * Constructs a new instance. * * @param root The set of vectors to serialize to an {@link ArrowRecordBatch}. * @param includeNullCount Controls whether null count is copied to the {@link ArrowRecordBatch} * @param alignBuffers Controls if buffers get aligned to 8-byte boundaries. */ public VectorUnloader(VectorSchemaRoot root, boolean includeNullCount, boolean alignBuffers) { this(root, includeNullCount, NoCompressionCodec.INSTANCE, alignBuffers); } /** * Constructs a new instance. * * @param root The set of vectors to serialize to an {@link ArrowRecordBatch}. * @param includeNullCount Controls whether null count is copied to the {@link ArrowRecordBatch} * @param codec the codec for compressing data. If it is null, then no compression is needed. * @param alignBuffers Controls if buffers get aligned to 8-byte boundaries. */ public VectorUnloader( VectorSchemaRoot root, boolean includeNullCount, CompressionCodec codec, boolean alignBuffers) { this.root = root; this.includeNullCount = includeNullCount; this.codec = codec == null ? NoCompressionCodec.INSTANCE : codec; this.alignBuffers = alignBuffers; } /** * Performs the depth first traversal of the Vectors to create an {@link ArrowRecordBatch} * suitable for serialization. */ public ArrowRecordBatch getRecordBatch() { List nodes = new ArrayList<>(); List buffers = new ArrayList<>(); List variadicBufferCounts = new ArrayList<>(); for (FieldVector vector : root.getFieldVectors()) { appendNodes(vector, nodes, buffers, variadicBufferCounts); } // Do NOT retain buffers in ArrowRecordBatch constructor since we have already retained them. return new ArrowRecordBatch( root.getRowCount(), nodes, buffers, CompressionUtil.createBodyCompression(codec), variadicBufferCounts, alignBuffers, /*retainBuffers*/ false); } private long getVariadicBufferCount(FieldVector vector) { if (vector instanceof BaseVariableWidthViewVector) { return ((BaseVariableWidthViewVector) vector).getDataBuffers().size(); } return 0L; } private void appendNodes( FieldVector vector, List nodes, List buffers, List variadicBufferCounts) { nodes.add( new ArrowFieldNode(vector.getValueCount(), includeNullCount ? vector.getNullCount() : -1)); List fieldBuffers = vector.getFieldBuffers(); long variadicBufferCount = getVariadicBufferCount(vector); int expectedBufferCount = (int) (TypeLayout.getTypeBufferCount(vector.getField().getType()) + variadicBufferCount); // only update variadicBufferCounts for vectors that have variadic buffers if (vector instanceof BaseVariableWidthViewVector) { variadicBufferCounts.add(variadicBufferCount); } if (fieldBuffers.size() != expectedBufferCount) { throw new IllegalArgumentException( String.format( "wrong number of buffers for field %s in vector %s. found: %s", vector.getField(), vector.getClass().getSimpleName(), fieldBuffers)); } for (ArrowBuf buf : fieldBuffers) { // If the codec is NoCompressionCodec, then it will return the input buffer unchanged. In that // case, // we need to retain it for ArrowRecordBatch. Otherwise, it will return a new buffer, and also // close // the input buffer. In that case, we need to retain the input buffer still to avoid modifying // the source VectorSchemaRoot. buf.getReferenceManager().retain(); buffers.add(codec.compress(vector.getAllocator(), buf)); } for (FieldVector child : vector.getChildrenFromFields()) { appendNodes(child, nodes, buffers, variadicBufferCounts); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ViewVarBinaryVector.java000066400000000000000000000237661515451615200322740ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReusableBuffer; import org.apache.arrow.vector.complex.impl.ViewVarBinaryReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableViewVarBinaryHolder; import org.apache.arrow.vector.holders.ViewVarBinaryHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** * ViewVarBinaryVector implements a variable width view vector of binary values which could be NULL. * A validity buffer (bit vector) is maintained to track which elements in the vector are null. */ public final class ViewVarBinaryVector extends BaseVariableWidthViewVector implements ValueIterableVector { /** * Instantiate a ViewVarBinaryVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public ViewVarBinaryVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.VIEWVARBINARY.getType()), allocator); } /** * Instantiate a ViewVarBinaryVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public ViewVarBinaryVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a ViewVarBinaryVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public ViewVarBinaryVector(Field field, BufferAllocator allocator) { super(field, allocator); } @Override protected FieldReader getReaderImpl() { return new ViewVarBinaryReaderImpl(ViewVarBinaryVector.this); } /** * Get a minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.VIEWVARBINARY; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the variable length element at specified index as a byte array. * * @param index position of an element to get * @return array of bytes for a non-null element, null otherwise */ public byte[] get(int index) { assert index >= 0; if (NULL_CHECKING_ENABLED && isSet(index) == 0) { return null; } return getData(index); } /** * Read the value at the given position to the given output buffer. The caller is responsible for * checking for nullity first. * * @param index position of an element. * @param buffer the buffer to write into. */ @Override public void read(int index, ReusableBuffer buffer) { getData(index, buffer); } /** * Get the variable length element at a specified index as a byte array. * * @param index position of an element to get * @return byte array for a non-null element, null otherwise */ @Override public byte[] getObject(int index) { if (isSet(index) == 0) { return null; } return getData(index); } /** * Get the variable length element at specified index and sets the state in provided holder. * * @param index position of an element to get * @param holder data holder to be populated by this function */ public void get(int index, NullableViewVarBinaryHolder holder) { final int dataLength = getValueLength(index); if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; if (dataLength > INLINE_SIZE) { // data is in the data buffer // get buffer index final int bufferIndex = viewBuffer.getInt(((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH); // get data offset final int dataOffset = viewBuffer.getInt( ((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH + BUF_INDEX_WIDTH); holder.buffer = dataBuffers.get(bufferIndex); holder.start = dataOffset; holder.end = dataOffset + dataLength; } else { final long dataOffset = ((long) index * ELEMENT_SIZE) + LENGTH_WIDTH; // data is in the value buffer holder.buffer = viewBuffer; holder.start = (int) dataOffset; holder.end = (int) dataOffset + dataLength; } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the variable length element at the specified index to the data buffer supplied in the * holder. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void set(int index, ViewVarBinaryHolder holder) { int start = holder.start; int length = holder.end - start; setBytes(index, holder.buffer, start, length); lastSet = index; } /** * Same as {@link #set(int, ViewVarBinaryHolder)} except that it handles the case where index and * length of a new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void setSafe(int index, ViewVarBinaryHolder holder) { int length = holder.end - holder.start; handleSafe(index, length); set(index, holder); } /** * Set the variable length element at the specified index to the data buffer supplied in the * holder. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void set(int index, NullableViewVarBinaryHolder holder) { if (holder.isSet == 0) { setNull(index); } else { BitVectorHelper.setBit(validityBuffer, index); int start = holder.start; int length = holder.end - start; setBytes(index, holder.buffer, start, length); } lastSet = index; } /** * Same as {@link #set(int, NullableViewVarBinaryHolder)} except that it handles the case where * index and length of a new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void setSafe(int index, NullableViewVarBinaryHolder holder) { int length = holder.end - holder.start; handleSafe(index, length); set(index, holder); } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to target vector * @return {@link TransferPair} */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((ViewVarBinaryVector) to); } private class TransferImpl implements TransferPair { ViewVarBinaryVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new ViewVarBinaryVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new ViewVarBinaryVector(field, allocator); } public TransferImpl(ViewVarBinaryVector to) { this.to = to; } @Override public ViewVarBinaryVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, ViewVarBinaryVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ViewVarCharVector.java000066400000000000000000000271051515451615200317140ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.NullCheckingForGet.NULL_CHECKING_ENABLED; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReusableBuffer; import org.apache.arrow.vector.complex.impl.ViewVarCharReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.NullableViewVarCharHolder; import org.apache.arrow.vector.holders.ViewVarCharHolder; import org.apache.arrow.vector.types.Types.MinorType; 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; import org.apache.arrow.vector.validate.ValidateUtil; /** * ViewVarCharVector implements a view of a variable width vector of VARCHAR values which could be * NULL. A validity buffer (bit vector) is maintained to track which elements in the vector are * null. A viewBuffer keeps track of all values in the vector, and an external data buffer is kept * to keep longer strings (>12). */ public final class ViewVarCharVector extends BaseVariableWidthViewVector implements ValueIterableVector { /** * Instantiate a ViewVarCharVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param allocator allocator for memory management. */ public ViewVarCharVector(String name, BufferAllocator allocator) { this(name, FieldType.nullable(MinorType.VIEWVARCHAR.getType()), allocator); } /** * Instantiate a ViewVarCharVector. This doesn't allocate any memory for the data in vector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector * @param allocator allocator for memory management. */ public ViewVarCharVector(String name, FieldType fieldType, BufferAllocator allocator) { this(new Field(name, fieldType, null), allocator); } /** * Instantiate a ViewVarCharVector. This doesn't allocate any memory for the data in vector. * * @param field field materialized by this vector * @param allocator allocator for memory management. */ public ViewVarCharVector(Field field, BufferAllocator allocator) { super(field, allocator); } @Override protected FieldReader getReaderImpl() { return new ViewVarCharReaderImpl(ViewVarCharVector.this); } /** * Get a minor type for this vector. The vector holds values belonging to a particular type. * * @return {@link org.apache.arrow.vector.types.Types.MinorType} */ @Override public MinorType getMinorType() { return MinorType.VIEWVARCHAR; } /*----------------------------------------------------------------* | | | vector value retrieval methods | | | *----------------------------------------------------------------*/ /** * Get the variable length element at specified index as a byte array. * * @param index position of an element to get * @return array of bytes for a non-null element, null otherwise */ public byte[] get(int index) { assert index >= 0; if (NULL_CHECKING_ENABLED && isSet(index) == 0) { return null; } return getData(index); } /** * Get the variable length element at specified index as Text. * * @param index position of an element to get * @return Text object for a non-null element, null otherwise */ @Override public Text getObject(int index) { assert index >= 0; if (isSet(index) == 0) { return null; } final Text result = new Text(); read(index, result); return result; } /** * Read the value at the given position to the given output buffer. The caller is responsible for * checking for nullity first. * * @param index position of an element. * @param buffer the buffer to write into. */ @Override public void read(int index, ReusableBuffer buffer) { getData(index, buffer); } /** * Get the variable length element at specified index and sets the state in provided holder. * * @param index position of an element to get * @param holder data holder to be populated by this function */ public void get(int index, NullableViewVarCharHolder holder) { final int dataLength = getValueLength(index); if (isSet(index) == 0) { holder.isSet = 0; return; } holder.isSet = 1; if (dataLength > INLINE_SIZE) { // data is in the data buffer // get buffer index final int bufferIndex = viewBuffer.getInt(((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH); // get data offset final int dataOffset = viewBuffer.getInt( ((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH + BUF_INDEX_WIDTH); holder.buffer = dataBuffers.get(bufferIndex); holder.start = dataOffset; holder.end = dataOffset + dataLength; } else { final long dataOffset = ((long) index * ELEMENT_SIZE) + LENGTH_WIDTH; // data is in the value buffer holder.buffer = viewBuffer; holder.start = (int) dataOffset; holder.end = (int) dataOffset + dataLength; } } /*----------------------------------------------------------------* | | | vector value setter methods | | | *----------------------------------------------------------------*/ /** * Set the variable length element at the specified index to the data buffer supplied in the * holder. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void set(int index, ViewVarCharHolder holder) { int start = holder.start; int length = holder.end - start; setBytes(index, holder.buffer, start, length); lastSet = index; } /** * Same as {@link #set(int, ViewVarCharHolder)} except that it handles the case where index and * length of a new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void setSafe(int index, ViewVarCharHolder holder) { int length = holder.end - holder.start; handleSafe(index, length); set(index, holder); } /** * Set the variable length element at the specified index to the data buffer supplied in the * holder. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void set(int index, NullableViewVarCharHolder holder) { if (holder.isSet == 0) { setNull(index); } else { BitVectorHelper.setBit(validityBuffer, index); int start = holder.start; int length = holder.end - start; setBytes(index, holder.buffer, start, length); } lastSet = index; } /** * Same as {@link #set(int, NullableViewVarCharHolder)} except that it handles the case where * index and length of a new element are beyond the existing capacity of the vector. * * @param index position of the element to set * @param holder holder that carries data buffer. */ public void setSafe(int index, NullableViewVarCharHolder holder) { int length = holder.end - holder.start; handleSafe(index, length); set(index, holder); } /** * Set the variable length element at the specified index to the content in supplied Text. * * @param index position of the element to set * @param text Text object with data */ public void set(int index, Text text) { set(index, text.getBytes(), 0, (int) text.getLength()); } /** * Same as {@link #set(int, NullableViewVarCharHolder)} except that it handles the case where * index and length of a new element are beyond the existing capacity of the vector. * * @param index position of the element to set. * @param text Text object with data */ public void setSafe(int index, Text text) { setSafe(index, text.getBytes(), 0, (int) text.getLength()); } @Override public void validateScalars() { for (int i = 0; i < getValueCount(); ++i) { byte[] value = get(i); if (value != null) { ValidateUtil.validateOrThrow( Text.validateUTF8NoThrow(value), "Non-UTF-8 data in VarCharVector at position " + i + "."); } } } /*----------------------------------------------------------------* | | | vector transfer | | | *----------------------------------------------------------------*/ /** * Construct a TransferPair comprising this and a target vector of the same type. * * @param ref name of the target vector * @param allocator allocator for the target vector * @return {@link TransferPair} (UnsupportedOperationException) */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new TransferImpl(ref, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param field The field materialized by this vector. * @param allocator allocator for the target vector * @return {@link TransferPair} (UnsupportedOperationException) */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator); } /** * Construct a TransferPair with a desired target vector of the same type. * * @param to the target for the transfer * @return {@link TransferPair} (UnsupportedOperationException) */ @Override public TransferPair makeTransferPair(ValueVector to) { return new TransferImpl((ViewVarCharVector) to); } private class TransferImpl implements TransferPair { ViewVarCharVector to; public TransferImpl(String ref, BufferAllocator allocator) { to = new ViewVarCharVector(ref, field.getFieldType(), allocator); } public TransferImpl(Field field, BufferAllocator allocator) { to = new ViewVarCharVector(field, allocator); } public TransferImpl(ViewVarCharVector to) { this.to = to; } @Override public ViewVarCharVector getTo() { return to; } @Override public void transfer() { transferTo(to); } @Override public void splitAndTransfer(int startIndex, int length) { splitAndTransferTo(startIndex, length, to); } @Override public void copyValueSafe(int fromIndex, int toIndex) { to.copyFromSafe(fromIndex, toIndex, ViewVarCharVector.this); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ZeroVector.java000066400000000000000000000063651515451615200304570ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.ArrowBufPointer; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.CallBack; import org.apache.arrow.vector.util.TransferPair; /** A zero length vector of any type. */ public final class ZeroVector extends NullVector { public static final ZeroVector INSTANCE = new ZeroVector(); /** * Instantiate a ZeroVector. * * @param name name of the vector */ public ZeroVector(String name) { super(name); } /** * Instantiate a ZeroVector. * * @param name name of the vector * @param fieldType type of Field materialized by this vector. */ public ZeroVector(String name, FieldType fieldType) { super(name, fieldType); } /** * Instantiate a ZeroVector. * * @param field field materialized by this vector. */ public ZeroVector(Field field) { super(field); } @Deprecated public ZeroVector() {} @Override public int getValueCount() { return 0; } @Override public void setValueCount(int valueCount) {} @Override public int getNullCount() { return 0; } @Override public boolean isNull(int index) { throw new IndexOutOfBoundsException(); } @Override public int hashCode(int index) { return 0; } @Override public int hashCode(int index, ArrowBufHasher hasher) { return ArrowBufPointer.NULL_HASH_CODE; } @Override public int getValueCapacity() { return 0; } @Override public TransferPair getTransferPair(BufferAllocator allocator) { return defaultPair; } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return defaultPair; } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return defaultPair; } @Override public TransferPair makeTransferPair(ValueVector target) { return defaultPair; } private final TransferPair defaultPair = new TransferPair() { @Override public void transfer() {} @Override public void splitAndTransfer(int startIndex, int length) {} @Override public ValueVector getTo() { return ZeroVector.this; } @Override public void copyValueSafe(int from, int to) {} }; } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/compare/000077500000000000000000000000001515451615200271265ustar00rootroot00000000000000ApproxEqualsVisitor.java000066400000000000000000000120771515451615200337250ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/compare/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compare; import java.util.function.BiFunction; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.compare.util.ValueEpsilonEqualizers; /** Visitor to compare floating point vectors approximately. */ public class ApproxEqualsVisitor extends RangeEqualsVisitor { /** Functions to calculate difference between float/double values. */ private final VectorValueEqualizer floatDiffFunction; private final VectorValueEqualizer doubleDiffFunction; /** Default epsilons for diff functions. */ public static final float DEFAULT_FLOAT_EPSILON = 1.0E-6f; public static final double DEFAULT_DOUBLE_EPSILON = 1.0E-6; /** * Constructs a new instance with default tolerances. * * @param left left vector * @param right right vector */ public ApproxEqualsVisitor(ValueVector left, ValueVector right) { this(left, right, DEFAULT_FLOAT_EPSILON, DEFAULT_DOUBLE_EPSILON); } /** * Constructs a new instance. * * @param left left vector * @param right right vector * @param floatEpsilon difference for float values * @param doubleEpsilon difference for double values */ public ApproxEqualsVisitor( ValueVector left, ValueVector right, float floatEpsilon, double doubleEpsilon) { this( left, right, new ValueEpsilonEqualizers.Float4EpsilonEqualizer(floatEpsilon), new ValueEpsilonEqualizers.Float8EpsilonEqualizer(doubleEpsilon)); } /** Constructs a new instance. */ public ApproxEqualsVisitor( ValueVector left, ValueVector right, VectorValueEqualizer floatDiffFunction, VectorValueEqualizer doubleDiffFunction) { this(left, right, floatDiffFunction, doubleDiffFunction, DEFAULT_TYPE_COMPARATOR); } /** * Constructs a new instance. * * @param left the left vector. * @param right the right vector. * @param floatDiffFunction the equalizer for float values. * @param doubleDiffFunction the equalizer for double values. * @param typeComparator type comparator to compare vector type. */ public ApproxEqualsVisitor( ValueVector left, ValueVector right, VectorValueEqualizer floatDiffFunction, VectorValueEqualizer doubleDiffFunction, BiFunction typeComparator) { super(left, right, typeComparator); this.floatDiffFunction = floatDiffFunction; this.doubleDiffFunction = doubleDiffFunction; } @Override public Boolean visit(BaseFixedWidthVector left, Range range) { if (left instanceof Float4Vector) { if (!validate(left)) { return false; } return float4ApproxEquals(range); } else if (left instanceof Float8Vector) { if (!validate(left)) { return false; } return float8ApproxEquals(range); } else { return super.visit(left, range); } } @Override protected ApproxEqualsVisitor createInnerVisitor( ValueVector left, ValueVector right, BiFunction typeComparator) { return new ApproxEqualsVisitor( left, right, floatDiffFunction.clone(), doubleDiffFunction.clone(), typeComparator); } private boolean float4ApproxEquals(Range range) { Float4Vector leftVector = (Float4Vector) getLeft(); Float4Vector rightVector = (Float4Vector) getRight(); for (int i = 0; i < range.getLength(); i++) { int leftIndex = range.getLeftStart() + i; int rightIndex = range.getRightStart() + i; if (!floatDiffFunction.valuesEqual(leftVector, leftIndex, rightVector, rightIndex)) { return false; } } return true; } private boolean float8ApproxEquals(Range range) { Float8Vector leftVector = (Float8Vector) getLeft(); Float8Vector rightVector = (Float8Vector) getRight(); for (int i = 0; i < range.getLength(); i++) { int leftIndex = range.getLeftStart() + i; int rightIndex = range.getRightStart() + i; if (!doubleDiffFunction.valuesEqual(leftVector, leftIndex, rightVector, rightIndex)) { return false; } } return true; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/compare/Range.java000066400000000000000000000037331515451615200310330ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compare; /** Wrapper for the parameters of comparing a range of values in two vectors. */ public class Range { /** Start position in the left vector. */ private int leftStart = -1; /** Start position in the right vector. */ private int rightStart = -1; /** Length of the range. */ private int length = -1; /** Constructs a new instance. */ public Range() {} /** * Constructs a new instance. * * @param leftStart start index in left vector * @param rightStart start index in right vector * @param length length of range */ public Range(int leftStart, int rightStart, int length) { this.leftStart = leftStart; this.rightStart = rightStart; this.length = length; } public int getLeftStart() { return leftStart; } public int getRightStart() { return rightStart; } public int getLength() { return length; } public Range setLeftStart(int leftStart) { this.leftStart = leftStart; return this; } public Range setRightStart(int rightStart) { this.rightStart = rightStart; return this; } public Range setLength(int length) { this.length = length; return this; } } RangeEqualsVisitor.java000066400000000000000000000751421515451615200335120ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/compare/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compare; import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import java.util.List; import java.util.function.BiFunction; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.util.ByteFunctionHelpers; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.BaseLargeVariableWidthVector; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.BaseVariableWidthViewVector; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.ExtensionTypeVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.NullVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.BaseLargeRepeatedValueViewVector; import org.apache.arrow.vector.complex.BaseRepeatedValueVector; import org.apache.arrow.vector.complex.BaseRepeatedValueViewVector; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.LargeListVector; 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.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; /** Visitor to compare a range of values for vectors. */ public class RangeEqualsVisitor implements VectorVisitor { private ValueVector left; private ValueVector right; private BiFunction typeComparator; private boolean typeCompareResult; /** Default type comparator. */ public static final BiFunction DEFAULT_TYPE_COMPARATOR = (v1, v2) -> new TypeEqualsVisitor(v2).equals(v1); /** * Constructs a new instance with default type comparator. * * @param left left vector * @param right right vector */ public RangeEqualsVisitor(ValueVector left, ValueVector right) { this(left, right, DEFAULT_TYPE_COMPARATOR); } /** * Constructs a new instance. * * @param left left vector * @param right right vector * @param typeComparator type comparator to compare vector type. */ public RangeEqualsVisitor( ValueVector left, ValueVector right, BiFunction typeComparator) { this.left = left; this.right = right; this.typeComparator = typeComparator; Preconditions.checkArgument(left != null, "left vector cannot be null"); Preconditions.checkArgument(right != null, "right vector cannot be null"); // type usually checks only once unless the left vector is changed. checkType(); } private void checkType() { if (typeComparator == null || left == right) { typeCompareResult = true; } else { typeCompareResult = typeComparator.apply(left, right); } } /** Validate the passed left vector, if it is changed, reset and check type. */ protected boolean validate(ValueVector left) { if (left != this.left) { this.left = left; checkType(); } return typeCompareResult; } /** Check range equals. */ public boolean rangeEquals(Range range) { if (!typeCompareResult) { return false; } Preconditions.checkArgument( range.getLeftStart() >= 0, "leftStart %s must be non negative.", range.getLeftStart()); Preconditions.checkArgument( range.getRightStart() >= 0, "rightStart %s must be non negative.", range.getRightStart()); Preconditions.checkArgument( range.getRightStart() + range.getLength() <= right.getValueCount(), "(rightStart + length) %s out of range[0, %s].", range.getRightStart() + range.getLength(), right.getValueCount()); Preconditions.checkArgument( range.getLeftStart() + range.getLength() <= left.getValueCount(), "(leftStart + length) %s out of range[0, %s].", range.getLeftStart() + range.getLength(), left.getValueCount()); return left.accept(this, range); } public ValueVector getLeft() { return left; } public ValueVector getRight() { return right; } @Override public Boolean visit(BaseFixedWidthVector left, Range range) { if (!validate(left)) { return false; } return compareBaseFixedWidthVectors(range); } @Override public Boolean visit(BaseVariableWidthVector left, Range range) { if (!validate(left)) { return false; } return compareBaseVariableWidthVectors(range); } @Override public Boolean visit(BaseLargeVariableWidthVector left, Range range) { if (!validate(left)) { return false; } return compareBaseLargeVariableWidthVectors(range); } @Override public Boolean visit(BaseVariableWidthViewVector left, Range range) { if (!validate(left)) { return false; } return compareBaseVariableWidthViewVectors(range); } @Override public Boolean visit(ListVector left, Range range) { if (!validate(left)) { return false; } return compareListVectors(range); } @Override public Boolean visit(FixedSizeListVector left, Range range) { if (!validate(left)) { return false; } return compareFixedSizeListVectors(range); } @Override public Boolean visit(LargeListVector left, Range range) { if (!validate(left)) { return false; } return compareLargeListVectors(range); } @Override public Boolean visit(NonNullableStructVector left, Range range) { if (!validate(left)) { return false; } return compareStructVectors(range); } @Override public Boolean visit(UnionVector left, Range range) { if (!validate(left)) { return false; } return compareUnionVectors(range); } @Override public Boolean visit(DenseUnionVector left, Range range) { if (!validate(left)) { return false; } return compareDenseUnionVectors(range); } @Override public Boolean visit(NullVector left, Range range) { if (!validate(left)) { return false; } return true; } @Override public Boolean visit(RunEndEncodedVector left, Range range) { if (!validate(left)) { return false; } return compareRunEndEncodedVectors(range); } @Override public Boolean visit(ExtensionTypeVector left, Range range) { if (!(right instanceof ExtensionTypeVector) || !validate(left)) { return false; } ValueVector rightUnderlying = ((ExtensionTypeVector) right).getUnderlyingVector(); TypeEqualsVisitor typeVisitor = new TypeEqualsVisitor(rightUnderlying); RangeEqualsVisitor underlyingVisitor = createInnerVisitor( left.getUnderlyingVector(), rightUnderlying, (l, r) -> typeVisitor.equals(l)); return underlyingVisitor.rangeEquals(range); } @Override public Boolean visit(ListViewVector left, Range range) { if (!validate(left)) { return false; } return compareListViewVectors(range); } @Override public Boolean visit(LargeListViewVector left, Range range) { if (!validate(left)) { return false; } return compareLargeListViewVectors(range); } protected boolean compareRunEndEncodedVectors(Range range) { RunEndEncodedVector leftVector = (RunEndEncodedVector) left; RunEndEncodedVector rightVector = (RunEndEncodedVector) right; 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); while (nextRun(leftIterator, rightIterator)) { int leftPhysicalIndex = leftIterator.getRunIndex(); int rightPhysicalIndex = rightIterator.getRunIndex(); if (leftIterator.getRunLength() != rightIterator.getRunLength() || !leftValuesVector.accept( innerVisitor, new Range(leftPhysicalIndex, rightPhysicalIndex, 1))) { return false; } } 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( ValueVector leftInner, ValueVector rightInner, BiFunction typeComparator) { return new RangeEqualsVisitor(leftInner, rightInner, typeComparator); } protected boolean compareUnionVectors(Range range) { UnionVector leftVector = (UnionVector) left; UnionVector rightVector = (UnionVector) right; Range subRange = new Range(0, 0, 1); for (int i = 0; i < range.getLength(); i++) { subRange.setLeftStart(range.getLeftStart() + i).setRightStart(range.getRightStart() + i); ValueVector leftSubVector = leftVector.getVector(range.getLeftStart() + i); ValueVector rightSubVector = rightVector.getVector(range.getRightStart() + i); if (leftSubVector == null || rightSubVector == null) { if (leftSubVector == rightSubVector) { continue; } else { return false; } } TypeEqualsVisitor typeVisitor = new TypeEqualsVisitor(rightSubVector); RangeEqualsVisitor visitor = createInnerVisitor( leftSubVector, rightSubVector, (left, right) -> typeVisitor.equals(left)); if (!visitor.rangeEquals(subRange)) { return false; } } return true; } protected boolean compareDenseUnionVectors(Range range) { DenseUnionVector leftVector = (DenseUnionVector) left; DenseUnionVector rightVector = (DenseUnionVector) right; Range subRange = new Range(0, 0, 1); for (int i = 0; i < range.getLength(); i++) { boolean isLeftNull = leftVector.isNull(range.getLeftStart() + i); boolean isRightNull = rightVector.isNull(range.getRightStart() + i); // compare nullabilities if (isLeftNull || isRightNull) { if (isLeftNull != isRightNull) { // exactly one slot is null, unequal return false; } else { // both slots are null, pass this iteration continue; } } // compare type ids byte leftTypeId = leftVector.getTypeId(range.getLeftStart() + i); byte rightTypeId = rightVector.getTypeId(range.getRightStart() + i); if (leftTypeId != rightTypeId) { return false; } ValueVector leftSubVector = leftVector.getVectorByType(leftTypeId); ValueVector rightSubVector = rightVector.getVectorByType(rightTypeId); if (leftSubVector == null || rightSubVector == null) { if (leftSubVector != rightSubVector) { // exactly one of the sub-vectors is null, unequal return false; } else { // both sub-vectors are null, pass this iteration continue; } } // compare values int leftOffset = leftVector.getOffset(range.getLeftStart() + i); int rightOffset = rightVector.getOffset(range.getRightStart() + i); subRange.setLeftStart(leftOffset).setRightStart(rightOffset); TypeEqualsVisitor typeVisitor = new TypeEqualsVisitor(rightSubVector); RangeEqualsVisitor visitor = createInnerVisitor( leftSubVector, rightSubVector, (left, right) -> typeVisitor.equals(left)); if (!visitor.rangeEquals(subRange)) { return false; } } return true; } private boolean compareStructVectorsInternal( NonNullableStructVector leftVector, NonNullableStructVector rightVector, Range range) { List leftChildNames = leftVector.getChildFieldNames(); for (String name : leftChildNames) { RangeEqualsVisitor visitor = createInnerVisitor( leftVector.getChild(name), rightVector.getChild(name), /*type comparator*/ null); if (!visitor.rangeEquals(range)) { return false; } } return true; } protected boolean compareStructVectors(Range range) { NonNullableStructVector leftVector = (NonNullableStructVector) left; NonNullableStructVector rightVector = (NonNullableStructVector) right; List leftChildNames = leftVector.getChildFieldNames(); if (!leftChildNames.equals(rightVector.getChildFieldNames())) { return false; } if (!(leftVector instanceof StructVector || rightVector instanceof StructVector)) { // neither struct vector is nullable return compareStructVectorsInternal(leftVector, rightVector, range); } Range subRange = new Range(0, 0, 0); boolean lastIsNull = true; int lastNullIndex = -1; for (int i = 0; i < range.getLength(); i++) { int leftIndex = range.getLeftStart() + i; int rightIndex = range.getRightStart() + i; boolean isLeftNull = leftVector.isNull(leftIndex); boolean isRightNull = rightVector.isNull(rightIndex); if (isLeftNull != isRightNull) { // exactly one slot is null, unequal return false; } if (isLeftNull) { // slots are null if (!lastIsNull) { subRange .setLeftStart(range.getLeftStart() + lastNullIndex + 1) .setRightStart(range.getRightStart() + lastNullIndex + 1) .setLength(i - (lastNullIndex + 1)); if (!compareStructVectorsInternal(leftVector, rightVector, subRange)) { return false; } } lastIsNull = true; lastNullIndex = i; } else { // slots are not null lastIsNull = false; } } if (!lastIsNull) { subRange .setLeftStart(range.getLeftStart() + lastNullIndex + 1) .setRightStart(range.getRightStart() + lastNullIndex + 1) .setLength(range.getLength() - (lastNullIndex + 1)); return compareStructVectorsInternal(leftVector, rightVector, subRange); } return true; } protected boolean compareBaseFixedWidthVectors(Range range) { BaseFixedWidthVector leftVector = (BaseFixedWidthVector) left; BaseFixedWidthVector rightVector = (BaseFixedWidthVector) right; for (int i = 0; i < range.getLength(); i++) { int leftIndex = range.getLeftStart() + i; int rightIndex = range.getRightStart() + i; boolean isNull = leftVector.isNull(leftIndex); if (isNull != rightVector.isNull(rightIndex)) { return false; } int typeWidth = leftVector.getTypeWidth(); if (!isNull) { if (!(leftVector instanceof BitVector)) { int startIndexLeft = typeWidth * leftIndex; int endIndexLeft = typeWidth * (leftIndex + 1); int startIndexRight = typeWidth * rightIndex; int endIndexRight = typeWidth * (rightIndex + 1); int ret = ByteFunctionHelpers.equal( leftVector.getDataBuffer(), startIndexLeft, endIndexLeft, rightVector.getDataBuffer(), startIndexRight, endIndexRight); if (ret == 0) { return false; } } else { boolean ret = ((BitVector) leftVector).get(leftIndex) == ((BitVector) rightVector).get(rightIndex); if (!ret) { return false; } } } } return true; } protected boolean compareBaseVariableWidthVectors(Range range) { BaseVariableWidthVector leftVector = (BaseVariableWidthVector) left; BaseVariableWidthVector rightVector = (BaseVariableWidthVector) right; for (int i = 0; i < range.getLength(); i++) { int leftIndex = range.getLeftStart() + i; int rightIndex = range.getRightStart() + i; boolean isNull = leftVector.isNull(leftIndex); if (isNull != rightVector.isNull(rightIndex)) { return false; } int offsetWidth = BaseVariableWidthVector.OFFSET_WIDTH; if (!isNull) { final int startIndexLeft = leftVector.getOffsetBuffer().getInt(leftIndex * offsetWidth); final int endIndexLeft = leftVector.getOffsetBuffer().getInt((leftIndex + 1) * offsetWidth); final int startIndexRight = rightVector.getOffsetBuffer().getInt(rightIndex * offsetWidth); final int endIndexRight = rightVector.getOffsetBuffer().getInt((rightIndex + 1) * offsetWidth); int ret = ByteFunctionHelpers.equal( leftVector.getDataBuffer(), startIndexLeft, endIndexLeft, rightVector.getDataBuffer(), startIndexRight, endIndexRight); if (ret == 0) { return false; } } } return true; } protected boolean compareBaseLargeVariableWidthVectors(Range range) { BaseLargeVariableWidthVector leftVector = (BaseLargeVariableWidthVector) left; BaseLargeVariableWidthVector rightVector = (BaseLargeVariableWidthVector) right; for (int i = 0; i < range.getLength(); i++) { int leftIndex = range.getLeftStart() + i; int rightIndex = range.getRightStart() + i; boolean isNull = leftVector.isNull(leftIndex); if (isNull != rightVector.isNull(rightIndex)) { return false; } int offsetWidth = BaseLargeVariableWidthVector.OFFSET_WIDTH; if (!isNull) { final long startIndexLeft = leftVector.getOffsetBuffer().getLong((long) leftIndex * offsetWidth); final long endIndexLeft = leftVector.getOffsetBuffer().getLong((long) (leftIndex + 1) * offsetWidth); final long startIndexRight = rightVector.getOffsetBuffer().getLong((long) rightIndex * offsetWidth); final long endIndexRight = rightVector.getOffsetBuffer().getLong((long) (rightIndex + 1) * offsetWidth); int ret = ByteFunctionHelpers.equal( leftVector.getDataBuffer(), startIndexLeft, endIndexLeft, rightVector.getDataBuffer(), startIndexRight, endIndexRight); if (ret == 0) { return false; } } } return true; } protected boolean compareBaseVariableWidthViewVectors(Range range) { BaseVariableWidthViewVector leftVector = (BaseVariableWidthViewVector) left; BaseVariableWidthViewVector rightVector = (BaseVariableWidthViewVector) right; final ArrowBuf leftViewBuffer = leftVector.getDataBuffer(); final ArrowBuf rightViewBuffer = rightVector.getDataBuffer(); final int elementSize = BaseVariableWidthViewVector.ELEMENT_SIZE; final int lengthWidth = BaseVariableWidthViewVector.LENGTH_WIDTH; final int prefixWidth = BaseVariableWidthViewVector.PREFIX_WIDTH; final int bufIndexWidth = BaseVariableWidthViewVector.BUF_INDEX_WIDTH; List leftDataBuffers = leftVector.getDataBuffers(); List rightDataBuffers = rightVector.getDataBuffers(); for (int i = 0; i < range.getLength(); i++) { int leftIndex = range.getLeftStart() + i; int rightIndex = range.getRightStart() + i; boolean isNull = leftVector.isNull(leftIndex); if (isNull != rightVector.isNull(rightIndex)) { return false; } if (isNull) { continue; } int startLeftByteOffset = leftIndex * elementSize; int startRightByteOffset = rightIndex * elementSize; int leftDataBufferValueLength = leftVector.getValueLength(leftIndex); int rightDataBufferValueLength = rightVector.getValueLength(rightIndex); if (leftDataBufferValueLength != rightDataBufferValueLength) { return false; } if (leftDataBufferValueLength > BaseVariableWidthViewVector.INLINE_SIZE) { // if the value is stored in the dataBuffers int leftDataBufferIndex = leftViewBuffer.getInt(startLeftByteOffset + lengthWidth + prefixWidth); int rightDataBufferIndex = rightViewBuffer.getInt(startRightByteOffset + lengthWidth + prefixWidth); final int leftDataOffset = leftViewBuffer.getInt(startLeftByteOffset + lengthWidth + prefixWidth + bufIndexWidth); final int rightDataOffset = rightViewBuffer.getInt( startRightByteOffset + lengthWidth + prefixWidth + bufIndexWidth); ArrowBuf leftDataBuffer = leftDataBuffers.get(leftDataBufferIndex); ArrowBuf rightDataBuffer = rightDataBuffers.get(rightDataBufferIndex); // check equality in the considered string stored in the dataBuffers int retDataBuf = ByteFunctionHelpers.equal( leftDataBuffer, leftDataOffset, leftDataOffset + leftDataBufferValueLength, rightDataBuffer, rightDataOffset, rightDataOffset + rightDataBufferValueLength); if (retDataBuf == 0) { return false; } } else { // if the value is stored in the view final int leftDataOffset = startLeftByteOffset + lengthWidth; final int rightDataOffset = startRightByteOffset + lengthWidth; // check equality in the considered string stored in the view int retDataBuf = ByteFunctionHelpers.equal( leftViewBuffer, leftDataOffset, leftDataOffset + leftDataBufferValueLength, rightViewBuffer, rightDataOffset, rightDataOffset + rightDataBufferValueLength); if (retDataBuf == 0) { return false; } } } return true; } protected boolean compareListVectors(Range range) { ListVector leftVector = (ListVector) left; ListVector rightVector = (ListVector) right; RangeEqualsVisitor innerVisitor = createInnerVisitor( leftVector.getDataVector(), rightVector.getDataVector(), /*type comparator*/ null); Range innerRange = new Range(); for (int i = 0; i < range.getLength(); i++) { int leftIndex = range.getLeftStart() + i; int rightIndex = range.getRightStart() + i; boolean isNull = leftVector.isNull(leftIndex); if (isNull != rightVector.isNull(rightIndex)) { return false; } int offsetWidth = BaseRepeatedValueVector.OFFSET_WIDTH; if (!isNull) { final int startIndexLeft = leftVector.getOffsetBuffer().getInt(leftIndex * offsetWidth); final int endIndexLeft = leftVector.getOffsetBuffer().getInt((leftIndex + 1) * offsetWidth); final int startIndexRight = rightVector.getOffsetBuffer().getInt(rightIndex * offsetWidth); final int endIndexRight = rightVector.getOffsetBuffer().getInt((rightIndex + 1) * offsetWidth); if ((endIndexLeft - startIndexLeft) != (endIndexRight - startIndexRight)) { return false; } innerRange = innerRange .setRightStart(startIndexRight) .setLeftStart(startIndexLeft) .setLength(endIndexLeft - startIndexLeft); if (!innerVisitor.rangeEquals(innerRange)) { return false; } } } return true; } protected boolean compareFixedSizeListVectors(Range range) { FixedSizeListVector leftVector = (FixedSizeListVector) left; FixedSizeListVector rightVector = (FixedSizeListVector) right; if (leftVector.getListSize() != rightVector.getListSize()) { return false; } int listSize = leftVector.getListSize(); RangeEqualsVisitor innerVisitor = createInnerVisitor( leftVector.getDataVector(), rightVector.getDataVector(), /*type comparator*/ null); Range innerRange = new Range(0, 0, listSize); for (int i = 0; i < range.getLength(); i++) { int leftIndex = range.getLeftStart() + i; int rightIndex = range.getRightStart() + i; boolean isNull = leftVector.isNull(leftIndex); if (isNull != rightVector.isNull(rightIndex)) { return false; } if (!isNull) { final int startIndexLeft = leftIndex * listSize; final int endIndexLeft = (leftIndex + 1) * listSize; final int startIndexRight = rightIndex * listSize; final int endIndexRight = (rightIndex + 1) * listSize; if ((endIndexLeft - startIndexLeft) != (endIndexRight - startIndexRight)) { return false; } innerRange = innerRange.setLeftStart(startIndexLeft).setRightStart(startIndexRight); if (!innerVisitor.rangeEquals(innerRange)) { return false; } } } return true; } protected boolean compareLargeListVectors(Range range) { LargeListVector leftVector = (LargeListVector) left; LargeListVector rightVector = (LargeListVector) right; RangeEqualsVisitor innerVisitor = createInnerVisitor( leftVector.getDataVector(), rightVector.getDataVector(), /*type comparator*/ null); Range innerRange = new Range(); for (int i = 0; i < range.getLength(); i++) { int leftIndex = range.getLeftStart() + i; int rightIndex = range.getRightStart() + i; boolean isNull = leftVector.isNull(leftIndex); if (isNull != rightVector.isNull(rightIndex)) { return false; } long offsetWidth = LargeListVector.OFFSET_WIDTH; if (!isNull) { final long startIndexLeft = leftVector.getOffsetBuffer().getLong((long) leftIndex * offsetWidth); final long endIndexLeft = leftVector.getOffsetBuffer().getLong((long) (leftIndex + 1) * offsetWidth); final long startIndexRight = rightVector.getOffsetBuffer().getLong((long) rightIndex * offsetWidth); final long endIndexRight = rightVector.getOffsetBuffer().getLong((long) (rightIndex + 1) * offsetWidth); if ((endIndexLeft - startIndexLeft) != (endIndexRight - startIndexRight)) { return false; } innerRange = innerRange // TODO revisit these casts when long indexing is finished .setRightStart(checkedCastToInt(startIndexRight)) .setLeftStart(checkedCastToInt(startIndexLeft)) .setLength(checkedCastToInt(endIndexLeft - startIndexLeft)); if (!innerVisitor.rangeEquals(innerRange)) { return false; } } } return true; } protected boolean compareListViewVectors(Range range) { ListViewVector leftVector = (ListViewVector) left; ListViewVector rightVector = (ListViewVector) right; RangeEqualsVisitor innerVisitor = createInnerVisitor( leftVector.getDataVector(), rightVector.getDataVector(), /*type comparator*/ null); Range innerRange = new Range(); for (int i = 0; i < range.getLength(); i++) { int leftIndex = range.getLeftStart() + i; int rightIndex = range.getRightStart() + i; boolean isNull = leftVector.isNull(leftIndex); if (isNull != rightVector.isNull(rightIndex)) { return false; } int offsetWidth = BaseRepeatedValueViewVector.OFFSET_WIDTH; int sizeWidth = BaseRepeatedValueViewVector.SIZE_WIDTH; if (!isNull) { final int startIndexLeft = leftVector.getOffsetBuffer().getInt((long) leftIndex * offsetWidth); final int leftSize = leftVector.getSizeBuffer().getInt((long) leftIndex * sizeWidth); final int startIndexRight = rightVector.getOffsetBuffer().getInt((long) rightIndex * offsetWidth); final int rightSize = rightVector.getSizeBuffer().getInt((long) rightIndex * sizeWidth); if (leftSize != rightSize) { return false; } innerRange = innerRange .setRightStart(startIndexRight) .setLeftStart(startIndexLeft) .setLength(leftSize); if (!innerVisitor.rangeEquals(innerRange)) { return false; } } } return true; } protected boolean compareLargeListViewVectors(Range range) { LargeListViewVector leftVector = (LargeListViewVector) left; LargeListViewVector rightVector = (LargeListViewVector) right; RangeEqualsVisitor innerVisitor = createInnerVisitor( leftVector.getDataVector(), rightVector.getDataVector(), /*type comparator*/ null); Range innerRange = new Range(); for (int i = 0; i < range.getLength(); i++) { int leftIndex = range.getLeftStart() + i; int rightIndex = range.getRightStart() + i; boolean isNull = leftVector.isNull(leftIndex); if (isNull != rightVector.isNull(rightIndex)) { return false; } int offsetWidth = BaseLargeRepeatedValueViewVector.OFFSET_WIDTH; int sizeWidth = BaseLargeRepeatedValueViewVector.SIZE_WIDTH; if (!isNull) { final int startIndexLeft = leftVector.getOffsetBuffer().getInt((long) leftIndex * offsetWidth); final int leftSize = leftVector.getSizeBuffer().getInt((long) leftIndex * sizeWidth); final int startIndexRight = rightVector.getOffsetBuffer().getInt((long) rightIndex * offsetWidth); final int rightSize = rightVector.getSizeBuffer().getInt((long) rightIndex * sizeWidth); if (leftSize != rightSize) { return false; } innerRange = innerRange .setRightStart(startIndexRight) .setLeftStart(startIndexLeft) .setLength(leftSize); if (!innerVisitor.rangeEquals(innerRange)) { return false; } } } return true; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/compare/TypeEqualsVisitor.java000066400000000000000000000131731515451615200334520ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compare; import java.util.List; import java.util.Objects; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.BaseLargeVariableWidthVector; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.BaseVariableWidthViewVector; import org.apache.arrow.vector.ExtensionTypeVector; import org.apache.arrow.vector.NullVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.LargeListVector; 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.NonNullableStructVector; import org.apache.arrow.vector.complex.RunEndEncodedVector; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.types.pojo.Field; /** Visitor to compare type equals for vectors. */ public class TypeEqualsVisitor implements VectorVisitor { private final ValueVector right; private final boolean checkName; private final boolean checkMetadata; /** Construct an instance. */ public TypeEqualsVisitor(ValueVector right) { this(right, true, true); } /** * Construct an instance. * * @param right right vector * @param checkName whether checks names * @param checkMetadata whether checks metadata */ public TypeEqualsVisitor(ValueVector right, boolean checkName, boolean checkMetadata) { this.right = right; this.checkName = checkName; this.checkMetadata = checkMetadata; } /** Check type equals without passing IN param in VectorVisitor. */ public boolean equals(ValueVector left) { return left.accept(this, null); } @Override public Boolean visit(BaseFixedWidthVector left, Void value) { return compareField(left.getField(), right.getField()); } @Override public Boolean visit(BaseVariableWidthVector left, Void value) { return compareField(left.getField(), right.getField()); } @Override public Boolean visit(BaseLargeVariableWidthVector left, Void value) { return compareField(left.getField(), right.getField()); } @Override public Boolean visit(BaseVariableWidthViewVector left, Void value) { return compareField(left.getField(), right.getField()); } @Override public Boolean visit(ListVector left, Void value) { return compareField(left.getField(), right.getField()); } @Override public Boolean visit(FixedSizeListVector left, Void value) { return compareField(left.getField(), right.getField()); } @Override public Boolean visit(LargeListVector left, Void value) { return compareField(left.getField(), right.getField()); } @Override public Boolean visit(NonNullableStructVector left, Void value) { return compareField(left.getField(), right.getField()); } @Override public Boolean visit(UnionVector left, Void value) { return compareField(left.getField(), right.getField()); } @Override public Boolean visit(DenseUnionVector left, Void value) { return compareField(left.getField(), right.getField()); } @Override public Boolean visit(NullVector left, Void value) { return compareField(left.getField(), right.getField()); } @Override public Boolean visit(ExtensionTypeVector left, Void value) { return compareField(left.getField(), right.getField()); } @Override public Boolean visit(ListViewVector left, Void value) { return compareField(left.getField(), right.getField()); } @Override public Boolean visit(LargeListViewVector left, Void value) { return compareField(left.getField(), right.getField()); } @Override public Boolean visit(RunEndEncodedVector left, Void value) { return compareField(left.getField(), right.getField()); } private boolean compareField(Field leftField, Field rightField) { if (leftField == rightField) { return true; } return (!checkName || Objects.equals(leftField.getName(), rightField.getName())) && Objects.equals(leftField.isNullable(), rightField.isNullable()) && Objects.equals(leftField.getType(), rightField.getType()) && Objects.equals(leftField.getDictionary(), rightField.getDictionary()) && (!checkMetadata || Objects.equals(leftField.getMetadata(), rightField.getMetadata())) && compareChildren(leftField.getChildren(), rightField.getChildren()); } private boolean compareChildren(List leftChildren, List rightChildren) { if (leftChildren.size() != rightChildren.size()) { return false; } for (int i = 0; i < leftChildren.size(); i++) { if (!compareField(leftChildren.get(i), rightChildren.get(i))) { return false; } } return true; } } VectorEqualsVisitor.java000066400000000000000000000041651515451615200337150ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/compare/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compare; import static org.apache.arrow.vector.compare.RangeEqualsVisitor.DEFAULT_TYPE_COMPARATOR; import java.util.function.BiFunction; import org.apache.arrow.vector.ValueVector; /** Visitor to compare vectors equal. */ public class VectorEqualsVisitor { /** * Checks if two vectors are equals with default type comparator. * * @param left the left vector to compare. * @param right the right vector to compare. * @return true if the vectors are equal, and false otherwise. */ public static boolean vectorEquals(ValueVector left, ValueVector right) { return vectorEquals(left, right, DEFAULT_TYPE_COMPARATOR); } /** * Checks if two vectors are equals. * * @param left the left vector to compare. * @param right the right vector to compare. * @param typeComparator type comparator to compare vector type. * @return true if the vectors are equal, and false otherwise. */ public static boolean vectorEquals( ValueVector left, ValueVector right, BiFunction typeComparator) { if (left.getValueCount() != right.getValueCount()) { return false; } RangeEqualsVisitor visitor = new RangeEqualsVisitor(left, right, typeComparator); return visitor.rangeEquals(new Range(0, 0, left.getValueCount())); } } VectorValueEqualizer.java000066400000000000000000000032311515451615200340320ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/compare/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compare; import org.apache.arrow.vector.ValueVector; /** * A function to determine if two vectors are equal at specified positions. * * @param the vector type. */ public interface VectorValueEqualizer extends Cloneable { /** * Checks if the vectors are equal at the given positions, given that the values at both positions * are non-null. * * @param vector1 the first vector. * @param index1 index in the first vector. * @param vector2 the second vector. * @param index2 index in the second vector. * @return true if the two values are considered to be equal, and false otherwise. */ boolean valuesEqual(V vector1, int index1, V vector2, int index2); /** * Creates a equalizer of the same type. * * @return the newly created equalizer. */ VectorValueEqualizer clone(); } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/compare/VectorVisitor.java000066400000000000000000000057101515451615200326160ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compare; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.BaseLargeVariableWidthVector; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.BaseVariableWidthViewVector; import org.apache.arrow.vector.ExtensionTypeVector; import org.apache.arrow.vector.NullVector; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.LargeListVector; 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.NonNullableStructVector; import org.apache.arrow.vector.complex.RunEndEncodedVector; import org.apache.arrow.vector.complex.UnionVector; /** * Generic visitor to visit a {@link org.apache.arrow.vector.ValueVector}. * * @param the output result type. * @param the input data together with visitor. */ public interface VectorVisitor { OUT visit(BaseFixedWidthVector left, IN value); OUT visit(BaseVariableWidthVector left, IN value); OUT visit(BaseLargeVariableWidthVector left, IN value); OUT visit(BaseVariableWidthViewVector left, IN value); OUT visit(ListVector left, IN value); OUT visit(FixedSizeListVector left, IN value); OUT visit(LargeListVector left, IN value); OUT visit(NonNullableStructVector left, IN value); OUT visit(UnionVector left, IN value); OUT visit(DenseUnionVector left, IN value); OUT visit(NullVector left, IN value); OUT visit(ExtensionTypeVector left, IN value); default OUT visit(ListViewVector left, IN value) { throw new UnsupportedOperationException("VectorVisitor for ListViewVector is not supported."); } default OUT visit(LargeListViewVector left, IN value) { throw new UnsupportedOperationException( "VectorVisitor for LargeListViewVector is not supported."); } default OUT visit(RunEndEncodedVector left, IN value) { throw new UnsupportedOperationException( "VectorVisitor for LargeListViewVector is not supported."); } ; } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/compare/util/000077500000000000000000000000001515451615200301035ustar00rootroot00000000000000ValueEpsilonEqualizers.java000066400000000000000000000104511515451615200353430ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/compare/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compare.util; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.FloatingPointVector; import org.apache.arrow.vector.compare.VectorValueEqualizer; /** * Vector value equalizers that regard values as equal if their difference is within a small * threshold (epsilon). */ public class ValueEpsilonEqualizers { private ValueEpsilonEqualizers() {} /** Difference function for floating point values. */ public static class FloatingPointEpsilonEqualizer implements VectorValueEqualizer { private final double epsilon; public FloatingPointEpsilonEqualizer(double epsilon) { this.epsilon = epsilon; } @Override public final boolean valuesEqual( FloatingPointVector vector1, int index1, FloatingPointVector vector2, int index2) { boolean isNull1 = vector1.isNull(index1); boolean isNull2 = vector2.isNull(index2); if (isNull1 || isNull2) { return isNull1 == isNull2; } double d1 = vector1.getValueAsDouble(index1); double d2 = vector2.getValueAsDouble(index2); if (Double.isNaN(d1)) { return Double.isNaN(d2); } if (Double.isInfinite(d1)) { return Double.isInfinite(d2) && Math.signum(d1) == Math.signum(d2); } return Math.abs(d1 - d2) <= epsilon; } @Override public VectorValueEqualizer clone() { return new FloatingPointEpsilonEqualizer(epsilon); } } /** Difference function for float values. */ public static class Float4EpsilonEqualizer implements VectorValueEqualizer { private final float epsilon; public Float4EpsilonEqualizer(float epsilon) { this.epsilon = epsilon; } @Override public final boolean valuesEqual( Float4Vector vector1, int index1, Float4Vector vector2, int index2) { boolean isNull1 = vector1.isNull(index1); boolean isNull2 = vector2.isNull(index2); if (isNull1 || isNull2) { return isNull1 == isNull2; } float f1 = vector1.get(index1); float f2 = vector2.get(index2); if (Float.isNaN(f1)) { return Float.isNaN(f2); } if (Float.isInfinite(f1)) { return Float.isInfinite(f2) && Math.signum(f1) == Math.signum(f2); } return Math.abs(f1 - f2) <= epsilon; } @Override public VectorValueEqualizer clone() { return new Float4EpsilonEqualizer(epsilon); } } /** Difference function for double values. */ public static class Float8EpsilonEqualizer implements VectorValueEqualizer { private final double epsilon; public Float8EpsilonEqualizer(double epsilon) { this.epsilon = epsilon; } @Override public final boolean valuesEqual( Float8Vector vector1, int index1, Float8Vector vector2, int index2) { boolean isNull1 = vector1.isNull(index1); boolean isNull2 = vector2.isNull(index2); if (isNull1 || isNull2) { return isNull1 == isNull2; } double d1 = vector1.get(index1); double d2 = vector2.get(index2); if (Double.isNaN(d1)) { return Double.isNaN(d2); } if (Double.isInfinite(d1)) { return Double.isInfinite(d2) && Math.signum(d1) == Math.signum(d2); } return Math.abs(d1 - d2) <= epsilon; } @Override public VectorValueEqualizer clone() { return new Float8EpsilonEqualizer(epsilon); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/000077500000000000000000000000001515451615200271475ustar00rootroot00000000000000AbstractContainerVector.java000066400000000000000000000125501515451615200345270ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.OutOfMemoryException; import org.apache.arrow.vector.DensityAwareVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.ArrowType.FixedSizeList; import org.apache.arrow.vector.types.pojo.ArrowType.List; import org.apache.arrow.vector.types.pojo.ArrowType.ListView; import org.apache.arrow.vector.types.pojo.ArrowType.Struct; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.CallBack; import org.apache.arrow.vector.util.ValueVectorUtility; /** * Base class for composite vectors. * *

    This class implements common functionality of composite vectors. */ public abstract class AbstractContainerVector implements ValueVector, DensityAwareVector { static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(AbstractContainerVector.class); protected final String name; protected final BufferAllocator allocator; protected final CallBack callBack; protected AbstractContainerVector(String name, BufferAllocator allocator, CallBack callBack) { this.name = name; this.allocator = allocator; this.callBack = callBack; } /** Representation of vector suitable for debugging. */ @Override public String toString() { return ValueVectorUtility.getToString(this, 0, getValueCount()); } @Override public void allocateNew() throws OutOfMemoryException { if (!allocateNewSafe()) { throw new OutOfMemoryException(); } } @Override public BufferAllocator getAllocator() { return allocator; } /** * Returns a {@link org.apache.arrow.vector.ValueVector} corresponding to the given field name if * exists or null. * * @param name the name of the child to return * @return the corresponding FieldVector */ public FieldVector getChild(String name) { return getChild(name, FieldVector.class); } /** Clears out all underlying child vectors. */ @Override public void close() { for (ValueVector vector : (Iterable) this) { vector.close(); } } protected T typeify(ValueVector v, Class clazz) { if (clazz.isAssignableFrom(v.getClass())) { return clazz.cast(v); } throw new IllegalStateException( String.format( "Vector requested [%s] was different than type stored [%s]. Arrow " + "doesn't yet support heterogeneous types.", clazz.getSimpleName(), v.getClass().getSimpleName())); } protected boolean supportsDirectRead() { return false; } // return the number of child vectors public abstract int size(); // add a new vector with the input FieldType or return the existing vector if we already added one // with the same name public abstract T addOrGet( String name, FieldType fieldType, Class clazz); // return the child vector with the input name public abstract T getChild(String name, Class clazz); // return the child vector's ordinal in the composite container public abstract VectorWithOrdinal getChildVectorWithOrdinal(String name); public StructVector addOrGetStruct(String name) { return addOrGet(name, FieldType.nullable(new Struct()), StructVector.class); } public ListVector addOrGetList(String name) { return addOrGet(name, FieldType.nullable(new List()), ListVector.class); } public ListViewVector addOrGetListView(String name) { return addOrGet(name, FieldType.nullable(new ListView()), ListViewVector.class); } public UnionVector addOrGetUnion(String name) { return addOrGet(name, FieldType.nullable(MinorType.UNION.getType()), UnionVector.class); } public FixedSizeListVector addOrGetFixedSizeList(String name, int listSize) { return addOrGet( name, FieldType.nullable(new FixedSizeList(listSize)), FixedSizeListVector.class); } public MapVector addOrGetMap(String name, boolean keysSorted) { return addOrGet(name, FieldType.nullable(new ArrowType.Map(keysSorted)), MapVector.class); } @Override public void copyFrom(int fromIndex, int thisIndex, ValueVector from) { throw new UnsupportedOperationException(); } @Override public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) { throw new UnsupportedOperationException(); } @Override public String getName() { return name; } } AbstractStructVector.java000066400000000000000000000347361515451615200341030ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.util.ArrayList; import java.util.Collections; import java.util.Iterator; import java.util.List; import java.util.Locale; import java.util.stream.Collectors; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.CallBack; import org.apache.arrow.vector.util.PromotableMultiMapWithOrdinal; import org.apache.arrow.vector.util.ValueVectorUtility; /** Base class for StructVectors. Currently used by NonNullableStructVector */ public abstract class AbstractStructVector extends AbstractContainerVector { private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(AbstractContainerVector.class); private static final String STRUCT_CONFLICT_POLICY_ENV = "ARROW_STRUCT_CONFLICT_POLICY"; private static final String STRUCT_CONFLICT_POLICY_JVM = "arrow.struct.conflict.policy"; private static final ConflictPolicy DEFAULT_CONFLICT_POLICY; // Maintains a map with key as field name and value is the vector itself private final PromotableMultiMapWithOrdinal vectors; protected final boolean allowConflictPolicyChanges; private ConflictPolicy conflictPolicy; static { 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)); } catch (Exception e) { conflictPolicy = ConflictPolicy.CONFLICT_REPLACE; } DEFAULT_CONFLICT_POLICY = conflictPolicy; } /** Policy to determine how to react when duplicate columns are encountered. */ public enum ConflictPolicy { // 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. This is the default behaviour CONFLICT_REPLACE, // Refuse the new field and error out. CONFLICT_ERROR } /** Base constructor that sets default conflict policy to APPEND. */ protected AbstractStructVector( String name, BufferAllocator allocator, CallBack callBack, ConflictPolicy conflictPolicy, boolean allowConflictPolicyChanges) { super(name, allocator, callBack); this.conflictPolicy = conflictPolicy == null ? DEFAULT_CONFLICT_POLICY : conflictPolicy; this.vectors = new PromotableMultiMapWithOrdinal<>(allowConflictPolicyChanges, this.conflictPolicy); this.allowConflictPolicyChanges = allowConflictPolicyChanges; } /** Set conflict policy and return last conflict policy state. */ public ConflictPolicy setConflictPolicy(ConflictPolicy conflictPolicy) { ConflictPolicy tmp = this.conflictPolicy; this.conflictPolicy = conflictPolicy; this.vectors.setConflictPolicy(conflictPolicy); return tmp; } public ConflictPolicy getConflictPolicy() { return conflictPolicy; } @Override public void close() { for (final ValueVector valueVector : vectors.values()) { valueVector.close(); } vectors.clear(); super.close(); } @Override public boolean allocateNewSafe() { /* boolean to keep track if all the memory allocation were successful * Used in the case of composite vectors when we need to allocate multiple * buffers for multiple vectors. If one of the allocations failed we need to * clear all the memory that we allocated */ boolean success = false; try { for (final ValueVector v : vectors.values()) { if (!v.allocateNewSafe()) { return false; } } success = true; } finally { if (!success) { clear(); } } return true; } @Override public void reAlloc() { for (final ValueVector v : vectors.values()) { v.reAlloc(); } } /** * Adds a new field with the given parameters or replaces the existing one and consequently * returns the resultant {@link org.apache.arrow.vector.ValueVector}. * *

    Execution takes place in the following order: * *

      *
    • if field is new, create and insert a new vector of desired type. *
    • if field exists and existing vector is of desired vector type, return the vector. *
    • if field exists and null filled, clear the existing vector; create and insert a new * vector of desired type. *
    • otherwise, throw an {@link java.lang.IllegalStateException} *
    * * @param childName the name of the field * @param fieldType the type for the vector * @param clazz class of expected vector type * @param class type of expected vector type * @return resultant {@link org.apache.arrow.vector.ValueVector} * @throws java.lang.IllegalStateException raised if there is a hard schema change */ @Override public T addOrGet(String childName, FieldType fieldType, Class clazz) { final ValueVector existing = getChild(childName); boolean create = false; if (existing == null) { create = true; } else if (clazz.isAssignableFrom(existing.getClass())) { return clazz.cast(existing); } else if (nullFilled(existing)) { existing.clear(); create = true; } if (create) { final T vector = clazz.cast(fieldType.createNewSingleVector(childName, allocator, callBack)); putChild(childName, vector); if (callBack != null) { callBack.doWork(); } return vector; } final String message = "Arrow does not support schema change yet. Existing[%s] and desired[%s] vector types " + "mismatch"; throw new IllegalStateException( String.format(message, existing.getClass().getSimpleName(), clazz.getSimpleName())); } private boolean nullFilled(ValueVector vector) { return BitVectorHelper.checkAllBitsEqualTo( vector.getValidityBuffer(), vector.getValueCount(), false); } /** * Returns a {@link org.apache.arrow.vector.ValueVector} corresponding to the given ordinal * identifier. * * @param id the ordinal of the child to return * @return the corresponding child */ public ValueVector getChildByOrdinal(int id) { return vectors.getByOrdinal(id); } /** * Returns a {@link org.apache.arrow.vector.ValueVector} instance of subtype of T corresponding to * the given field name if exists or null. * *

    If there is more than one element for name this will return the first inserted. * * @param name the name of the child to return * @param clazz the expected type of the child * @return the child corresponding to this name */ @Override public T getChild(String name, Class clazz) { final FieldVector f = vectors.get(name); if (f == null) { return null; } return typeify(f, clazz); } protected ValueVector add(String childName, FieldType fieldType) { FieldVector vector = fieldType.createNewSingleVector(childName, allocator, callBack); putChild(childName, vector); if (callBack != null) { callBack.doWork(); } return vector; } /** * Inserts the vector with the given name if it does not exist else replaces it with the new * value. * *

    Note that this method does not enforce any vector type check nor throws a schema change * exception. * * @param name the name of the child to add * @param vector the vector to add as a child */ protected void putChild(String name, FieldVector vector) { putVector(name, vector); } private void put(String name, FieldVector vector, boolean overwrite) { final boolean old = vectors.put( Preconditions.checkNotNull(name, "field name cannot be null"), Preconditions.checkNotNull(vector, "vector cannot be null"), overwrite); if (old) { logger.debug("Field [{}] mutated to [{}] ", name, vector.getClass().getSimpleName()); } } /** * Inserts the input vector into the map if it does not exist. * *

    If the field name already exists the conflict is handled according to the currently set * ConflictPolicy * * @param name field name * @param vector vector to be inserted */ protected void putVector(String name, FieldVector vector) { switch (conflictPolicy) { case CONFLICT_APPEND: put(name, vector, false); break; case CONFLICT_IGNORE: if (!vectors.containsKey(name)) { put(name, vector, false); } break; case CONFLICT_REPLACE: if (vectors.containsKey(name)) { vectors.removeAll(name); } put(name, vector, true); break; case CONFLICT_ERROR: if (vectors.containsKey(name)) { throw new IllegalStateException( String.format( "Vector already exists: Existing[%s], Requested[%s] ", vector.getClass().getSimpleName(), vector.getField().getFieldType())); } put(name, vector, false); break; default: throw new IllegalStateException( String.format("%s type not a valid conflict state", conflictPolicy)); } } /** * Get child vectors. * * @return a sequence of underlying child vectors. */ protected List getChildren() { int size = vectors.size(); List children = new ArrayList<>(); for (int i = 0; i < size; i++) { children.add(vectors.getByOrdinal(i)); } return children; } /** Get child field names. */ public List getChildFieldNames() { return getChildren().stream() .map(child -> child.getField().getName()) .collect(Collectors.toList()); } /** * Get the number of child vectors. * * @return the number of underlying child vectors. */ @Override public int size() { return vectors.size(); } @Override public Iterator iterator() { return Collections.unmodifiableCollection(vectors.values()).iterator(); } /** * Get primitive child vectors. * * @return a list of scalar child vectors recursing the entire vector hierarchy. */ public List getPrimitiveVectors() { final List primitiveVectors = new ArrayList<>(); for (final FieldVector v : vectors.values()) { primitiveVectors.addAll(getPrimitiveVectors(v)); } return primitiveVectors; } private List getPrimitiveVectors(FieldVector v) { final List primitives = new ArrayList<>(); if (v instanceof AbstractStructVector) { AbstractStructVector structVector = (AbstractStructVector) v; primitives.addAll(structVector.getPrimitiveVectors()); } else if (v instanceof ListVector) { ListVector listVector = (ListVector) v; primitives.addAll(getPrimitiveVectors(listVector.getDataVector())); } else if (v instanceof FixedSizeListVector) { FixedSizeListVector listVector = (FixedSizeListVector) v; primitives.addAll(getPrimitiveVectors(listVector.getDataVector())); } else if (v instanceof UnionVector) { UnionVector unionVector = (UnionVector) v; for (final FieldVector vector : unionVector.getChildrenFromFields()) { primitives.addAll(getPrimitiveVectors(vector)); } } else { primitives.add(v); } return primitives; } /** * Get a child vector by name. If duplicate names this returns the first inserted. * * @param name the name of the child to return * @return a vector with its corresponding ordinal mapping if field exists or null. */ @Override public VectorWithOrdinal getChildVectorWithOrdinal(String name) { final int ordinal = vectors.getOrdinal(name); if (ordinal < 0) { return null; } final ValueVector vector = vectors.getByOrdinal(ordinal); return new VectorWithOrdinal(vector, ordinal); } /** * Return the underlying buffers associated with this vector. Note that this doesn't impact the * reference counts for this buffer, so it only should be used for in-context access. Also note * that this buffer changes regularly, thus external classes shouldn't hold a reference to it * (unless they change it). * * @param clear Whether to clear vector before returning, the buffers will still be refcounted but * the returned array will be the only reference to them. Also, this won't clear the child * buffers. * @return The underlying {@link ArrowBuf buffers} that is used by this vector instance. */ @Override public ArrowBuf[] getBuffers(boolean clear) { final List buffers = new ArrayList<>(); for (final ValueVector vector : vectors.values()) { for (final ArrowBuf buf : vector.getBuffers(false)) { buffers.add(buf); if (clear) { buf.getReferenceManager().retain(1); } } if (clear) { vector.clear(); } } return buffers.toArray(new ArrowBuf[buffers.size()]); } @Override public int getBufferSize() { int actualBufSize = 0; for (final ValueVector v : vectors.values()) { for (final ArrowBuf buf : v.getBuffers(false)) { actualBufSize += (int) buf.writerIndex(); } } return actualBufSize; } @Override public String toString() { return ValueVectorUtility.getToString(this, 0, getValueCount()); } } BaseLargeRepeatedValueViewVector.java000066400000000000000000000337241515451615200362560ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.memory.util.LargeMemoryUtil.capAtMaxInt; import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import java.util.Collections; import java.util.Iterator; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.CommonUtil; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.AddOrGetResult; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.BaseValueVector; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.DensityAwareVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.NullVector; import org.apache.arrow.vector.UInt4Vector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.ZeroVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.CallBack; import org.apache.arrow.vector.util.OversizedAllocationException; import org.apache.arrow.vector.util.SchemaChangeRuntimeException; public abstract class BaseLargeRepeatedValueViewVector extends BaseValueVector implements RepeatedValueVector, FieldVector { public static final FieldVector DEFAULT_DATA_VECTOR = ZeroVector.INSTANCE; public static final String DATA_VECTOR_NAME = "$data$"; public static final byte OFFSET_WIDTH = 8; public static final byte SIZE_WIDTH = 8; protected ArrowBuf offsetBuffer; protected ArrowBuf sizeBuffer; protected FieldVector vector; protected final CallBack repeatedCallBack; protected long offsetAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH; protected long sizeAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * SIZE_WIDTH; private final String name; protected String defaultDataVectorName = DATA_VECTOR_NAME; protected BaseLargeRepeatedValueViewVector( String name, BufferAllocator allocator, CallBack callBack) { this(name, allocator, DEFAULT_DATA_VECTOR, callBack); } protected BaseLargeRepeatedValueViewVector( String name, BufferAllocator allocator, FieldVector vector, CallBack callBack) { super(allocator); this.name = name; this.offsetBuffer = allocator.getEmpty(); this.sizeBuffer = allocator.getEmpty(); this.vector = Preconditions.checkNotNull(vector, "data vector cannot be null"); this.repeatedCallBack = callBack; this.valueCount = 0; } @Override public String getName() { return name; } @Override public boolean allocateNewSafe() { boolean dataAlloc = false; try { allocateBuffers(); dataAlloc = vector.allocateNewSafe(); } catch (Exception e) { clear(); return false; } finally { if (!dataAlloc) { clear(); } } return dataAlloc; } private void allocateBuffers() { offsetBuffer = allocateBuffers(offsetAllocationSizeInBytes); sizeBuffer = allocateBuffers(sizeAllocationSizeInBytes); } protected ArrowBuf allocateBuffers(final long size) { final int curSize = (int) size; ArrowBuf buffer = allocator.buffer(curSize); buffer.readerIndex(0); buffer.setZero(0, buffer.capacity()); return buffer; } @Override public void reAlloc() { reallocateBuffers(); vector.reAlloc(); } protected void reallocateBuffers() { reallocOffsetBuffer(); reallocSizeBuffer(); } private void reallocOffsetBuffer() { final long currentBufferCapacity = offsetBuffer.capacity(); long newAllocationSize = currentBufferCapacity * 2; if (newAllocationSize == 0) { if (offsetAllocationSizeInBytes > 0) { newAllocationSize = offsetAllocationSizeInBytes; } else { newAllocationSize = INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH * 2; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); newAllocationSize = Math.min(newAllocationSize, (long) OFFSET_WIDTH * Integer.MAX_VALUE); assert newAllocationSize >= 1; if (newAllocationSize > MAX_ALLOCATION_SIZE || newAllocationSize <= offsetBuffer.capacity()) { throw new OversizedAllocationException("Unable to expand the buffer"); } final ArrowBuf newBuf = allocator.buffer(newAllocationSize); newBuf.setBytes(0, offsetBuffer, 0, currentBufferCapacity); newBuf.setZero(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); offsetBuffer.getReferenceManager().release(1); offsetBuffer = newBuf; offsetAllocationSizeInBytes = newAllocationSize; } private void reallocSizeBuffer() { final long currentBufferCapacity = sizeBuffer.capacity(); long newAllocationSize = currentBufferCapacity * 2; if (newAllocationSize == 0) { if (sizeAllocationSizeInBytes > 0) { newAllocationSize = sizeAllocationSizeInBytes; } else { newAllocationSize = INITIAL_VALUE_ALLOCATION * SIZE_WIDTH * 2; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); newAllocationSize = Math.min(newAllocationSize, (long) SIZE_WIDTH * Integer.MAX_VALUE); assert newAllocationSize >= 1; if (newAllocationSize > MAX_ALLOCATION_SIZE || newAllocationSize <= sizeBuffer.capacity()) { throw new OversizedAllocationException("Unable to expand the buffer"); } final ArrowBuf newBuf = allocator.buffer(newAllocationSize); newBuf.setBytes(0, sizeBuffer, 0, currentBufferCapacity); newBuf.setZero(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); sizeBuffer.getReferenceManager().release(1); sizeBuffer = newBuf; sizeAllocationSizeInBytes = newAllocationSize; } @Override public FieldVector getDataVector() { return vector; } @Override public void setInitialCapacity(int numRecords) { offsetAllocationSizeInBytes = (long) (numRecords) * OFFSET_WIDTH; sizeAllocationSizeInBytes = (long) (numRecords) * SIZE_WIDTH; if (vector instanceof BaseFixedWidthVector || vector instanceof BaseVariableWidthVector) { vector.setInitialCapacity(numRecords * RepeatedValueVector.DEFAULT_REPEAT_PER_RECORD); } else { vector.setInitialCapacity(numRecords); } } @Override public void setInitialCapacity(int numRecords, double density) { if ((numRecords * density) >= Integer.MAX_VALUE) { throw new OversizedAllocationException("Requested amount of memory is more than max allowed"); } offsetAllocationSizeInBytes = (long) numRecords * OFFSET_WIDTH; sizeAllocationSizeInBytes = (long) numRecords * SIZE_WIDTH; int innerValueCapacity = Math.max((int) (numRecords * density), 1); if (vector instanceof DensityAwareVector) { ((DensityAwareVector) vector).setInitialCapacity(innerValueCapacity, density); } else { vector.setInitialCapacity(innerValueCapacity); } } /** * Specialized version of setInitialTotalCapacity() for LargeListViewVector. This is used by some * callers when they want to explicitly control and be conservative about memory allocated for * inner data vector. This is very useful when we are working with memory constraints for a query * and have a fixed amount of memory reserved for the record batch. In such cases, we are likely * to face OOM or related problems when we reserve memory for a record batch with value count x * and do setInitialCapacity(x) such that each vector allocates only what is necessary and not the * default amount, but the multiplier forces the memory requirement to go beyond what was needed. * * @param numRecords value count * @param totalNumberOfElements the total number of elements to allow for in this vector across * all records. */ public void setInitialTotalCapacity(int numRecords, int totalNumberOfElements) { offsetAllocationSizeInBytes = (long) numRecords * OFFSET_WIDTH; sizeAllocationSizeInBytes = (long) numRecords * SIZE_WIDTH; vector.setInitialCapacity(totalNumberOfElements); } @Override public int getValueCapacity() { throw new UnsupportedOperationException( "Get value capacity is not supported in RepeatedValueVector"); } protected int getOffsetBufferValueCapacity() { return checkedCastToInt(offsetBuffer.capacity() / OFFSET_WIDTH); } protected int getSizeBufferValueCapacity() { return capAtMaxInt(sizeBuffer.capacity() / SIZE_WIDTH); } @Override public int getBufferSize() { if (valueCount == 0) { return 0; } return (valueCount * OFFSET_WIDTH) + (valueCount * SIZE_WIDTH) + vector.getBufferSize(); } @Override public int getBufferSizeFor(int valueCount) { if (valueCount == 0) { return 0; } int innerVectorValueCount = 0; for (int i = 0; i < valueCount; i++) { innerVectorValueCount += sizeBuffer.getInt(i * SIZE_WIDTH); } return (valueCount * OFFSET_WIDTH) + (valueCount * SIZE_WIDTH) + vector.getBufferSizeFor(checkedCastToInt(innerVectorValueCount)); } @Override public Iterator iterator() { return Collections.singleton(getDataVector()).iterator(); } @Override public void clear() { offsetBuffer = releaseBuffer(offsetBuffer); sizeBuffer = releaseBuffer(sizeBuffer); vector.clear(); valueCount = 0; super.clear(); } @Override public void reset() { offsetBuffer.setZero(0, offsetBuffer.capacity()); sizeBuffer.setZero(0, sizeBuffer.capacity()); vector.reset(); valueCount = 0; } @Override public ArrowBuf[] getBuffers(boolean clear) { return new ArrowBuf[0]; } @Override public int getValueCount() { return valueCount; } @Override public void setValueCount(int valueCount) { this.valueCount = valueCount; while (valueCount > getOffsetBufferValueCapacity()) { reallocateBuffers(); } final int childValueCount = valueCount == 0 ? 0 : getMaxViewEndChildVector(); vector.setValueCount(childValueCount); } /** * Get the end of the child vector via the maximum view length. This method deduces the length by * considering the condition i.e., argmax_i(offsets[i] + size[i]). * * @return the end of the child vector. */ protected int getMaxViewEndChildVector() { int maxOffsetSizeSum = offsetBuffer.getInt(0) + sizeBuffer.getInt(0); for (int i = 0; i < valueCount; i++) { int currentOffset = offsetBuffer.getInt((long) i * OFFSET_WIDTH); int currentSize = sizeBuffer.getInt((long) i * SIZE_WIDTH); int currentSum = currentOffset + currentSize; maxOffsetSizeSum = Math.max(maxOffsetSizeSum, currentSum); } return maxOffsetSizeSum; } /** * Get the end of the child vector via the maximum view length of the child vector by index. * * @return the end of the child vector by index */ protected int getMaxViewEndChildVectorByIndex(int index) { int maxOffsetSizeSum = offsetBuffer.getInt(0) + sizeBuffer.getInt(0); for (int i = 0; i < index; i++) { int currentOffset = offsetBuffer.getInt((long) i * OFFSET_WIDTH); int currentSize = sizeBuffer.getInt((long) i * SIZE_WIDTH); int currentSum = currentOffset + currentSize; maxOffsetSizeSum = Math.max(maxOffsetSizeSum, currentSum); } return maxOffsetSizeSum; } /** * Initialize the data vector (and execute callback) if it hasn't already been done, returns the * data vector. */ public AddOrGetResult addOrGetVector(FieldType fieldType) { boolean created = false; if (vector instanceof NullVector) { vector = fieldType.createNewSingleVector(defaultDataVectorName, allocator, repeatedCallBack); // returned vector must have the same field created = true; if (repeatedCallBack != null && // not a schema change if changing from ZeroVector to ZeroVector (fieldType.getType().getTypeID() != ArrowType.ArrowTypeID.Null)) { repeatedCallBack.doWork(); } } if (vector.getField().getType().getTypeID() != fieldType.getType().getTypeID()) { final String msg = String.format( "Inner vector type mismatch. Requested type: [%s], actual type: [%s]", fieldType.getType().getTypeID(), vector.getField().getType().getTypeID()); throw new SchemaChangeRuntimeException(msg); } return new AddOrGetResult<>((T) vector, created); } protected void replaceDataVector(FieldVector v) { vector.clear(); vector = v; } public abstract boolean isEmpty(int index); /** * Start a new value at the given index. * * @param index the index to start the new value at * @return the offset in the data vector where the new value starts */ public int startNewValue(int index) { while (index >= getOffsetBufferValueCapacity()) { reallocOffsetBuffer(); } while (index >= getSizeBufferValueCapacity()) { reallocSizeBuffer(); } if (index > 0) { final int prevOffset = getMaxViewEndChildVectorByIndex(index); offsetBuffer.setInt((long) index * OFFSET_WIDTH, prevOffset); } setValueCount(index + 1); return offsetBuffer.getInt((long) index * OFFSET_WIDTH); } @Override @Deprecated public UInt4Vector getOffsetVector() { throw new UnsupportedOperationException("There is no inner offset vector"); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/BaseListVector.java000066400000000000000000000022521515451615200327040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import org.apache.arrow.vector.FieldVector; /** Abstraction for all list type vectors. */ public interface BaseListVector extends FieldVector { /** Get data vector start index with the given list index. */ int getElementStartIndex(int index); /** Get data vector end index with the given list index. */ int getElementEndIndex(int index); } BaseRepeatedValueVector.java000066400000000000000000000334341515451615200344460ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.memory.util.LargeMemoryUtil.capAtMaxInt; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.Iterator; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.CommonUtil; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.AddOrGetResult; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.BaseValueVector; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.DensityAwareVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.NullVector; import org.apache.arrow.vector.UInt4Vector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.ZeroVector; import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeID; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.CallBack; import org.apache.arrow.vector.util.OversizedAllocationException; import org.apache.arrow.vector.util.SchemaChangeRuntimeException; /** Base class for Vectors that contain repeated values. */ public abstract class BaseRepeatedValueVector extends BaseValueVector implements RepeatedValueVector, BaseListVector { public static final FieldVector DEFAULT_DATA_VECTOR = ZeroVector.INSTANCE; public static final String DATA_VECTOR_NAME = "$data$"; public static final byte OFFSET_WIDTH = 4; protected ArrowBuf offsetBuffer; protected FieldVector vector; protected final CallBack repeatedCallBack; protected long offsetAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH; private final String name; protected String defaultDataVectorName = DATA_VECTOR_NAME; protected BaseRepeatedValueVector(String name, BufferAllocator allocator, CallBack callBack) { this(name, allocator, DEFAULT_DATA_VECTOR, callBack); } protected BaseRepeatedValueVector( String name, BufferAllocator allocator, FieldVector vector, CallBack callBack) { super(allocator); this.name = name; this.offsetBuffer = allocator.getEmpty(); this.vector = Preconditions.checkNotNull(vector, "data vector cannot be null"); this.repeatedCallBack = callBack; this.valueCount = 0; } @Override public String getName() { return name; } @Override public boolean allocateNewSafe() { boolean dataAlloc = false; try { offsetBuffer = allocateOffsetBuffer(offsetAllocationSizeInBytes); dataAlloc = vector.allocateNewSafe(); } catch (Exception e) { e.printStackTrace(); clear(); return false; } finally { if (!dataAlloc) { clear(); } } return dataAlloc; } protected ArrowBuf allocateOffsetBuffer(final long size) { final int curSize = (int) size; ArrowBuf offsetBuffer = allocator.buffer(curSize); offsetBuffer.readerIndex(0); offsetAllocationSizeInBytes = curSize; offsetBuffer.setZero(0, offsetBuffer.capacity()); return offsetBuffer; } @Override public void reAlloc() { reallocOffsetBuffer(); vector.reAlloc(); } protected void reallocOffsetBuffer() { final long currentBufferCapacity = offsetBuffer.capacity(); long newAllocationSize = currentBufferCapacity * 2; if (newAllocationSize == 0) { if (offsetAllocationSizeInBytes > 0) { newAllocationSize = offsetAllocationSizeInBytes; } else { newAllocationSize = INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH * 2; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); newAllocationSize = Math.min(newAllocationSize, (long) OFFSET_WIDTH * Integer.MAX_VALUE); assert newAllocationSize >= 1; if (newAllocationSize > MAX_ALLOCATION_SIZE || newAllocationSize <= offsetBuffer.capacity()) { throw new OversizedAllocationException("Unable to expand the buffer"); } final ArrowBuf newBuf = allocator.buffer(newAllocationSize); newBuf.setBytes(0, offsetBuffer, 0, currentBufferCapacity); newBuf.setZero(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); offsetBuffer.getReferenceManager().release(1); offsetBuffer = newBuf; offsetAllocationSizeInBytes = newAllocationSize; } /** * Get the offset vector. * * @return the underlying offset vector or null if none exists. * @deprecated This API will be removed, as the current implementations no longer hold inner * offset vectors. */ @Override @Deprecated public UInt4Vector getOffsetVector() { throw new UnsupportedOperationException("There is no inner offset vector"); } @Override public FieldVector getDataVector() { return vector; } @Override public void setInitialCapacity(int numRecords) { offsetAllocationSizeInBytes = (numRecords + 1L) * OFFSET_WIDTH; if (vector instanceof BaseFixedWidthVector || vector instanceof BaseVariableWidthVector) { vector.setInitialCapacity(numRecords * RepeatedValueVector.DEFAULT_REPEAT_PER_RECORD); } else { vector.setInitialCapacity(numRecords); } } /** * Specialized version of setInitialCapacity() for ListVector. This is used by some callers when * they want to explicitly control and be conservative about memory allocated for inner data * vector. This is very useful when we are working with memory constraints for a query and have a * fixed amount of memory reserved for the record batch. In such cases, we are likely to face OOM * or related problems when we reserve memory for a record batch with value count x and do * setInitialCapacity(x) such that each vector allocates only what is necessary and not the * default amount but the multiplier forces the memory requirement to go beyond what was needed. * * @param numRecords value count * @param density density of ListVector. Density is the average size of list per position in the * List vector. For example, a density value of 10 implies each position in the list vector * has a list of 10 values. A density value of 0.1 implies out of 10 positions in the list * vector, 1 position has a list of size 1 and remaining positions are null (no lists) or * empty lists. This helps in tightly controlling the memory we provision for inner data * vector. */ @Override public void setInitialCapacity(int numRecords, double density) { if ((numRecords * density) >= Integer.MAX_VALUE) { throw new OversizedAllocationException("Requested amount of memory is more than max allowed"); } offsetAllocationSizeInBytes = (numRecords + 1L) * OFFSET_WIDTH; int innerValueCapacity = Math.max((int) (numRecords * density), 1); if (vector instanceof DensityAwareVector) { ((DensityAwareVector) vector).setInitialCapacity(innerValueCapacity, density); } else { vector.setInitialCapacity(innerValueCapacity); } } /** * Specialized version of setInitialTotalCapacity() for ListVector. This is used by some callers * when they want to explicitly control and be conservative about memory allocated for inner data * vector. This is very useful when we are working with memory constraints for a query and have a * fixed amount of memory reserved for the record batch. In such cases, we are likely to face OOM * or related problems when we reserve memory for a record batch with value count x and do * setInitialCapacity(x) such that each vector allocates only what is necessary and not the * default amount but the multiplier forces the memory requirement to go beyond what was needed. * * @param numRecords value count * @param totalNumberOfElements the total number of elements to to allow for in this vector across * all records. */ public void setInitialTotalCapacity(int numRecords, int totalNumberOfElements) { offsetAllocationSizeInBytes = (numRecords + 1L) * OFFSET_WIDTH; vector.setInitialCapacity(totalNumberOfElements); } @Override public int getValueCapacity() { final int offsetValueCapacity = Math.max(getOffsetBufferValueCapacity() - 1, 0); if (vector == DEFAULT_DATA_VECTOR) { return offsetValueCapacity; } return Math.min(vector.getValueCapacity(), offsetValueCapacity); } protected int getOffsetBufferValueCapacity() { return capAtMaxInt(offsetBuffer.capacity() / OFFSET_WIDTH); } @Override public int getBufferSize() { if (valueCount == 0) { return 0; } return ((valueCount + 1) * OFFSET_WIDTH) + vector.getBufferSize(); } @Override public int getBufferSizeFor(int valueCount) { if (valueCount == 0) { return 0; } int innerVectorValueCount = offsetBuffer.getInt(valueCount * OFFSET_WIDTH); return ((valueCount + 1) * OFFSET_WIDTH) + vector.getBufferSizeFor(innerVectorValueCount); } @Override public Iterator iterator() { return Collections.singleton(getDataVector()).iterator(); } @Override public void clear() { offsetBuffer = releaseBuffer(offsetBuffer); vector.clear(); valueCount = 0; super.clear(); } @Override public void reset() { offsetBuffer.setZero(0, offsetBuffer.capacity()); vector.reset(); valueCount = 0; } /** * Return the underlying buffers associated with this vector. Note that this doesn't impact the * reference counts for this buffer, so it only should be used for in-context access. Also note * that this buffer changes regularly, thus external classes shouldn't hold a reference to it * (unless they change it). * * @param clear Whether to clear vector before returning, the buffers will still be refcounted but * the returned array will be the only reference to them. Also, this won't clear the child * buffers. * @return The underlying {@link ArrowBuf buffers} that is used by this vector instance. */ @Override public ArrowBuf[] getBuffers(boolean clear) { final ArrowBuf[] buffers; if (getBufferSize() == 0) { buffers = new ArrowBuf[0]; } else { List list = new ArrayList<>(); list.add(offsetBuffer); list.addAll(Arrays.asList(vector.getBuffers(false))); buffers = list.toArray(new ArrowBuf[list.size()]); } if (clear) { for (ArrowBuf buffer : buffers) { buffer.getReferenceManager().retain(); } clear(); } return buffers; } /** * Get value indicating if inner vector is set. * * @return 1 if inner vector is explicitly set via #addOrGetVector else 0 */ public int size() { return vector == DEFAULT_DATA_VECTOR ? 0 : 1; } /** * Initialize the data vector (and execute callback) if it hasn't already been done, returns the * data vector. */ public AddOrGetResult addOrGetVector(FieldType fieldType) { boolean created = false; if (vector instanceof NullVector) { vector = fieldType.createNewSingleVector(defaultDataVectorName, allocator, repeatedCallBack); // returned vector must have the same field created = true; if (repeatedCallBack != null && // not a schema change if changing from ZeroVector to ZeroVector (fieldType.getType().getTypeID() != ArrowTypeID.Null)) { repeatedCallBack.doWork(); } } if (vector.getField().getType().getTypeID() != fieldType.getType().getTypeID()) { final String msg = String.format( "Inner vector type mismatch. Requested type: [%s], actual type: [%s]", fieldType.getType().getTypeID(), vector.getField().getType().getTypeID()); throw new SchemaChangeRuntimeException(msg); } return new AddOrGetResult<>((T) vector, created); } protected void replaceDataVector(FieldVector v) { vector.clear(); vector = v; } @Override public int getValueCount() { return valueCount; } /* returns the value count for inner data vector for this list vector */ public int getInnerValueCount() { return vector.getValueCount(); } /** Returns the value count for inner data vector at a particular index. */ public int getInnerValueCountAt(int index) { return offsetBuffer.getInt((index + 1) * OFFSET_WIDTH) - offsetBuffer.getInt(index * OFFSET_WIDTH); } /** Return if value at index is empty. */ public abstract boolean isEmpty(int index); /** Starts a new repeated value. */ public int startNewValue(int index) { while (index >= getOffsetBufferValueCapacity()) { reallocOffsetBuffer(); } int offset = offsetBuffer.getInt(index * OFFSET_WIDTH); offsetBuffer.setInt((index + 1) * OFFSET_WIDTH, offset); setValueCount(index + 1); return offset; } /** Preallocates the number of repeated values. */ @Override public void setValueCount(int valueCount) { this.valueCount = valueCount; while (valueCount > getOffsetBufferValueCapacity()) { reallocOffsetBuffer(); } final int childValueCount = valueCount == 0 ? 0 : offsetBuffer.getInt(valueCount * OFFSET_WIDTH); vector.setValueCount(childValueCount); } } BaseRepeatedValueViewVector.java000066400000000000000000000334641515451615200353040ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.memory.util.LargeMemoryUtil.capAtMaxInt; import java.util.Collections; import java.util.Iterator; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.CommonUtil; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.AddOrGetResult; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.BaseValueVector; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.DensityAwareVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.NullVector; import org.apache.arrow.vector.UInt4Vector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.ZeroVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.CallBack; import org.apache.arrow.vector.util.OversizedAllocationException; import org.apache.arrow.vector.util.SchemaChangeRuntimeException; public abstract class BaseRepeatedValueViewVector extends BaseValueVector implements RepeatedValueVector, BaseListVector { public static final FieldVector DEFAULT_DATA_VECTOR = ZeroVector.INSTANCE; public static final String DATA_VECTOR_NAME = "$data$"; public static final byte OFFSET_WIDTH = 4; public static final byte SIZE_WIDTH = 4; protected ArrowBuf offsetBuffer; protected ArrowBuf sizeBuffer; protected FieldVector vector; protected final CallBack repeatedCallBack; protected long offsetAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH; protected long sizeAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * SIZE_WIDTH; private final String name; protected String defaultDataVectorName = DATA_VECTOR_NAME; protected BaseRepeatedValueViewVector(String name, BufferAllocator allocator, CallBack callBack) { this(name, allocator, DEFAULT_DATA_VECTOR, callBack); } protected BaseRepeatedValueViewVector( String name, BufferAllocator allocator, FieldVector vector, CallBack callBack) { super(allocator); this.name = name; this.offsetBuffer = allocator.getEmpty(); this.sizeBuffer = allocator.getEmpty(); this.vector = Preconditions.checkNotNull(vector, "data vector cannot be null"); this.repeatedCallBack = callBack; this.valueCount = 0; } @Override public String getName() { return name; } @Override public boolean allocateNewSafe() { boolean dataAlloc = false; try { allocateBuffers(); dataAlloc = vector.allocateNewSafe(); } catch (Exception e) { clear(); return false; } finally { if (!dataAlloc) { clear(); } } return dataAlloc; } private void allocateBuffers() { offsetBuffer = allocateBuffers(offsetAllocationSizeInBytes); sizeBuffer = allocateBuffers(sizeAllocationSizeInBytes); } protected ArrowBuf allocateBuffers(final long size) { final int curSize = (int) size; ArrowBuf buffer = allocator.buffer(curSize); buffer.readerIndex(0); buffer.setZero(0, buffer.capacity()); return buffer; } @Override public void reAlloc() { reallocateBuffers(); vector.reAlloc(); } protected void reallocateBuffers() { reallocOffsetBuffer(); reallocSizeBuffer(); } private void reallocOffsetBuffer() { final long currentBufferCapacity = offsetBuffer.capacity(); long newAllocationSize = currentBufferCapacity * 2; if (newAllocationSize == 0) { if (offsetAllocationSizeInBytes > 0) { newAllocationSize = offsetAllocationSizeInBytes; } else { newAllocationSize = INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH * 2; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); newAllocationSize = Math.min(newAllocationSize, (long) OFFSET_WIDTH * Integer.MAX_VALUE); assert newAllocationSize >= 1; if (newAllocationSize > MAX_ALLOCATION_SIZE || newAllocationSize <= offsetBuffer.capacity()) { throw new OversizedAllocationException("Unable to expand the buffer"); } final ArrowBuf newBuf = allocator.buffer(newAllocationSize); newBuf.setBytes(0, offsetBuffer, 0, currentBufferCapacity); newBuf.setZero(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); offsetBuffer.getReferenceManager().release(1); offsetBuffer = newBuf; offsetAllocationSizeInBytes = newAllocationSize; } private void reallocSizeBuffer() { final long currentBufferCapacity = sizeBuffer.capacity(); long newAllocationSize = currentBufferCapacity * 2; if (newAllocationSize == 0) { if (sizeAllocationSizeInBytes > 0) { newAllocationSize = sizeAllocationSizeInBytes; } else { newAllocationSize = INITIAL_VALUE_ALLOCATION * SIZE_WIDTH * 2; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); newAllocationSize = Math.min(newAllocationSize, (long) SIZE_WIDTH * Integer.MAX_VALUE); assert newAllocationSize >= 1; if (newAllocationSize > MAX_ALLOCATION_SIZE || newAllocationSize <= sizeBuffer.capacity()) { throw new OversizedAllocationException("Unable to expand the buffer"); } final ArrowBuf newBuf = allocator.buffer(newAllocationSize); newBuf.setBytes(0, sizeBuffer, 0, currentBufferCapacity); newBuf.setZero(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); sizeBuffer.getReferenceManager().release(1); sizeBuffer = newBuf; sizeAllocationSizeInBytes = newAllocationSize; } @Override public FieldVector getDataVector() { return vector; } @Override public void setInitialCapacity(int numRecords) { offsetAllocationSizeInBytes = (numRecords) * OFFSET_WIDTH; sizeAllocationSizeInBytes = (numRecords) * SIZE_WIDTH; if (vector instanceof BaseFixedWidthVector || vector instanceof BaseVariableWidthVector) { vector.setInitialCapacity(numRecords * RepeatedValueVector.DEFAULT_REPEAT_PER_RECORD); } else { vector.setInitialCapacity(numRecords); } } @Override public void setInitialCapacity(int numRecords, double density) { if ((numRecords * density) >= Integer.MAX_VALUE) { throw new OversizedAllocationException("Requested amount of memory is more than max allowed"); } offsetAllocationSizeInBytes = numRecords * OFFSET_WIDTH; sizeAllocationSizeInBytes = numRecords * SIZE_WIDTH; int innerValueCapacity = Math.max((int) (numRecords * density), 1); if (vector instanceof DensityAwareVector) { ((DensityAwareVector) vector).setInitialCapacity(innerValueCapacity, density); } else { vector.setInitialCapacity(innerValueCapacity); } } /** * Specialized version of setInitialTotalCapacity() for ListViewVector. This is used by some * callers when they want to explicitly control and be conservative about memory allocated for * inner data vector. This is very useful when we are working with memory constraints for a query * and have a fixed amount of memory reserved for the record batch. In such cases, we are likely * to face OOM or related problems when we reserve memory for a record batch with value count x * and do setInitialCapacity(x) such that each vector allocates only what is necessary and not the * default amount, but the multiplier forces the memory requirement to go beyond what was needed. * * @param numRecords value count * @param totalNumberOfElements the total number of elements to allow for in this vector across * all records. */ public void setInitialTotalCapacity(int numRecords, int totalNumberOfElements) { offsetAllocationSizeInBytes = numRecords * OFFSET_WIDTH; sizeAllocationSizeInBytes = numRecords * SIZE_WIDTH; vector.setInitialCapacity(totalNumberOfElements); } @Override public int getValueCapacity() { throw new UnsupportedOperationException( "Get value capacity is not supported in RepeatedValueVector"); } protected int getOffsetBufferValueCapacity() { return capAtMaxInt(offsetBuffer.capacity() / OFFSET_WIDTH); } protected int getSizeBufferValueCapacity() { return capAtMaxInt(sizeBuffer.capacity() / SIZE_WIDTH); } @Override public int getBufferSize() { if (valueCount == 0) { return 0; } return (valueCount * OFFSET_WIDTH) + (valueCount * SIZE_WIDTH) + vector.getBufferSize(); } @Override public int getBufferSizeFor(int valueCount) { if (valueCount == 0) { return 0; } int innerVectorValueCount = 0; for (int i = 0; i < valueCount; i++) { innerVectorValueCount += sizeBuffer.getInt(i * SIZE_WIDTH); } return (valueCount * OFFSET_WIDTH) + (valueCount * SIZE_WIDTH) + vector.getBufferSizeFor(innerVectorValueCount); } @Override public Iterator iterator() { return Collections.singleton(getDataVector()).iterator(); } @Override public void clear() { offsetBuffer = releaseBuffer(offsetBuffer); sizeBuffer = releaseBuffer(sizeBuffer); vector.clear(); valueCount = 0; super.clear(); } @Override public void reset() { offsetBuffer.setZero(0, offsetBuffer.capacity()); sizeBuffer.setZero(0, sizeBuffer.capacity()); vector.reset(); valueCount = 0; } @Override public ArrowBuf[] getBuffers(boolean clear) { return new ArrowBuf[0]; } @Override public int getValueCount() { return valueCount; } @Override public void setValueCount(int valueCount) { this.valueCount = valueCount; while (valueCount > getOffsetBufferValueCapacity()) { reallocateBuffers(); } final int childValueCount = valueCount == 0 ? 0 : getMaxViewEndChildVector(); vector.setValueCount(childValueCount); } /** * Get the end of the child vector via the maximum view length. This method deduces the length by * considering the condition i.e., argmax_i(offsets[i] + size[i]). * * @return the end of the child vector. */ protected int getMaxViewEndChildVector() { int maxOffsetSizeSum = offsetBuffer.getInt(0) + sizeBuffer.getInt(0); for (int i = 0; i < valueCount; i++) { int currentOffset = offsetBuffer.getInt(i * OFFSET_WIDTH); int currentSize = sizeBuffer.getInt(i * SIZE_WIDTH); int currentSum = currentOffset + currentSize; maxOffsetSizeSum = Math.max(maxOffsetSizeSum, currentSum); } return maxOffsetSizeSum; } /** * Get the end of the child vector via the maximum view length of the child vector by index. * * @return the end of the child vector by index */ protected int getMaxViewEndChildVectorByIndex(int index) { int maxOffsetSizeSum = offsetBuffer.getInt(0) + sizeBuffer.getInt(0); // int minOffset = offsetBuffer.getInt(0); for (int i = 0; i < index; i++) { int currentOffset = offsetBuffer.getInt(i * OFFSET_WIDTH); int currentSize = sizeBuffer.getInt(i * SIZE_WIDTH); int currentSum = currentOffset + currentSize; maxOffsetSizeSum = Math.max(maxOffsetSizeSum, currentSum); } return maxOffsetSizeSum; } /** * Initialize the data vector (and execute callback) if it hasn't already been done, returns the * data vector. */ public AddOrGetResult addOrGetVector(FieldType fieldType) { boolean created = false; if (vector instanceof NullVector) { vector = fieldType.createNewSingleVector(defaultDataVectorName, allocator, repeatedCallBack); // returned vector must have the same field created = true; if (repeatedCallBack != null && // not a schema change if changing from ZeroVector to ZeroVector (fieldType.getType().getTypeID() != ArrowType.ArrowTypeID.Null)) { repeatedCallBack.doWork(); } } if (vector.getField().getType().getTypeID() != fieldType.getType().getTypeID()) { final String msg = String.format( "Inner vector type mismatch. Requested type: [%s], actual type: [%s]", fieldType.getType().getTypeID(), vector.getField().getType().getTypeID()); throw new SchemaChangeRuntimeException(msg); } return new AddOrGetResult<>((T) vector, created); } protected void replaceDataVector(FieldVector v) { vector.clear(); vector = v; } public abstract boolean isEmpty(int index); /** * Start a new value at the given index. * * @param index the index to start the new value at * @return the offset in the data vector where the new value starts */ public int startNewValue(int index) { while (index >= getOffsetBufferValueCapacity()) { reallocOffsetBuffer(); } while (index >= getSizeBufferValueCapacity()) { reallocSizeBuffer(); } if (index > 0) { final int prevOffset = getMaxViewEndChildVectorByIndex(index); offsetBuffer.setInt(index * OFFSET_WIDTH, prevOffset); } setValueCount(index + 1); return offsetBuffer.getInt(index * OFFSET_WIDTH); } @Override @Deprecated public UInt4Vector getOffsetVector() { throw new UnsupportedOperationException("There is no inner offset vector"); } } EmptyValuePopulator.java000066400000000000000000000040041515451615200337320ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.UInt4Vector; /** A helper class that is used to track and populate empty values in repeated value vectors. */ public class EmptyValuePopulator { private final UInt4Vector offsets; public EmptyValuePopulator(UInt4Vector offsets) { this.offsets = Preconditions.checkNotNull(offsets, "offsets cannot be null"); } /** * Marks all values since the last set as empty. The last set value is obtained from underlying * offsets vector. * * @param lastIndex the last index (inclusive) in the offsets vector until which empty population * takes place * @throws java.lang.IndexOutOfBoundsException if lastIndex is negative or greater than offsets * capacity. */ public void populate(int lastIndex) { if (lastIndex < 0) { throw new IndexOutOfBoundsException("index cannot be negative"); } final int lastSet = Math.max(offsets.getValueCount() - 1, 0); final int previousEnd = offsets.get(lastSet); // 0 ? 0 : accessor.get(lastSet); for (int i = lastSet; i < lastIndex; i++) { offsets.setSafe(i + 1, previousEnd); } offsets.setValueCount(lastIndex + 1); } } FixedSizeListVector.java000066400000000000000000000505341515451615200336530ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static java.util.Collections.singletonList; import static org.apache.arrow.memory.util.LargeMemoryUtil.capAtMaxInt; import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import static org.apache.arrow.util.Preconditions.checkArgument; import static org.apache.arrow.vector.complex.BaseRepeatedValueVector.DATA_VECTOR_NAME; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.Iterator; import java.util.List; import java.util.Objects; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.OutOfMemoryException; 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.hash.ArrowBufHasher; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.AddOrGetResult; import org.apache.arrow.vector.BaseValueVector; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.BufferBacked; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueIterableVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.ZeroVector; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.impl.UnionFixedSizeListReader; import org.apache.arrow.vector.complex.impl.UnionFixedSizeListWriter; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; 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.CallBack; import org.apache.arrow.vector.util.JsonStringArrayList; import org.apache.arrow.vector.util.OversizedAllocationException; import org.apache.arrow.vector.util.SchemaChangeRuntimeException; import org.apache.arrow.vector.util.TransferPair; /** A ListVector where every list value is of the same size. */ public class FixedSizeListVector extends BaseValueVector implements BaseListVector, PromotableVector, ValueIterableVector> { public static FixedSizeListVector empty(String name, int size, BufferAllocator allocator) { FieldType fieldType = FieldType.nullable(new ArrowType.FixedSizeList(size)); return new FixedSizeListVector(name, allocator, fieldType, null); } private FieldVector vector; private final int listSize; private Field field; private UnionFixedSizeListReader reader; private int validityAllocationSizeInBytes; /** * Creates a new instance. * * @param name The name for the vector. * @param allocator The allocator to use for creating/reallocating buffers for the vector. * @param fieldType The underlying data type of the vector. * @param unusedSchemaChangeCallback Currently unused. */ public FixedSizeListVector( String name, BufferAllocator allocator, FieldType fieldType, CallBack unusedSchemaChangeCallback) { this(new Field(name, fieldType, null), allocator, unusedSchemaChangeCallback); } /** * Creates a new instance. * * @param field The field materialized by this vector. * @param allocator The allocator to use for creating/reallocating buffers for the vector. * @param unusedSchemaChangeCallback Currently unused. */ public FixedSizeListVector( Field field, BufferAllocator allocator, CallBack unusedSchemaChangeCallback) { super(allocator); this.field = field; this.validityBuffer = allocator.getEmpty(); this.vector = ZeroVector.INSTANCE; this.listSize = ((ArrowType.FixedSizeList) field.getFieldType().getType()).getListSize(); Preconditions.checkArgument(listSize >= 0, "list size must be non-negative"); this.valueCount = 0; this.validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION); } @Override public Field getField() { if (field.getChildren().contains(getDataVector().getField())) { return field; } field = new Field( field.getName(), field.getFieldType(), Collections.singletonList(getDataVector().getField())); return field; } @Override public MinorType getMinorType() { return MinorType.FIXED_SIZE_LIST; } @Override public String getName() { return field.getName(); } /** Get the fixed size for each list. */ public int getListSize() { return listSize; } @Override public void initializeChildrenFromFields(List children) { checkArgument( children.size() == 1, "Lists have one child Field. Found: %s", children.isEmpty() ? "none" : children); Field field = children.get(0); AddOrGetResult addOrGetVector = addOrGetVector(field.getFieldType()); checkArgument( addOrGetVector.isCreated(), "Child vector already existed: %s", addOrGetVector.getVector()); addOrGetVector.getVector().initializeChildrenFromFields(field.getChildren()); this.field = new Field(this.field.getName(), this.field.getFieldType(), children); } @Override public List getChildrenFromFields() { return singletonList(vector); } @Override public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers) { if (ownBuffers.size() != 1) { throw new IllegalArgumentException( "Illegal buffer count, expected " + 1 + ", got: " + ownBuffers.size()); } ArrowBuf bitBuffer = ownBuffers.get(0); validityBuffer.getReferenceManager().release(); validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); valueCount = fieldNode.getLength(); validityAllocationSizeInBytes = checkedCastToInt(validityBuffer.capacity()); } @Override public List getFieldBuffers() { List result = new ArrayList<>(1); setReaderAndWriterIndex(); result.add(validityBuffer); return result; } private void setReaderAndWriterIndex() { validityBuffer.readerIndex(0); validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); } /** * Get the inner vectors. * * @return the inner vectors for this field as defined by the TypeLayout */ @Deprecated @Override public List getFieldInnerVectors() { throw new UnsupportedOperationException("There are no inner vectors. Use getFieldBuffers"); } @Override protected FieldReader getReaderImpl() { return new UnionFixedSizeListReader(this); } @Override public UnionFixedSizeListReader getReader() { reader = (UnionFixedSizeListReader) super.getReader(); return reader; } private void invalidateReader() { reader = null; } @Override public void allocateNew() throws OutOfMemoryException { if (!allocateNewSafe()) { throw new OutOfMemoryException("Failure while allocating memory"); } } @Override public boolean allocateNewSafe() { /* boolean to keep track if all the memory allocation were successful * Used in the case of composite vectors when we need to allocate multiple * buffers for multiple vectors. If one of the allocations failed we need to * clear all the memory that we allocated */ boolean success = false; try { /* we are doing a new allocation -- release the current buffers */ clear(); /* allocate validity buffer */ allocateValidityBuffer(validityAllocationSizeInBytes); success = vector.allocateNewSafe(); } finally { if (!success) { clear(); } } return success; } @Override protected void allocateValidityBuffer(final long size) { super.allocateValidityBuffer(size); validityAllocationSizeInBytes = (int) size; } @Override public void reAlloc() { reallocValidityBuffer(); vector.reAlloc(); } private void reallocValidityBuffer() { final int currentBufferCapacity = checkedCastToInt(validityBuffer.capacity()); long newAllocationSize = currentBufferCapacity * 2L; if (newAllocationSize == 0) { if (validityAllocationSizeInBytes > 0) { newAllocationSize = validityAllocationSizeInBytes; } else { newAllocationSize = BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); assert newAllocationSize >= 1; if (newAllocationSize > MAX_ALLOCATION_SIZE) { throw new OversizedAllocationException("Unable to expand the buffer"); } final ArrowBuf newBuf = allocator.buffer(newAllocationSize); newBuf.setBytes(0, validityBuffer, 0, currentBufferCapacity); newBuf.setZero(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); validityBuffer.getReferenceManager().release(1); validityBuffer = newBuf; validityAllocationSizeInBytes = (int) newAllocationSize; } public FieldVector getDataVector() { return vector; } /** * Start a new value in the list vector. * * @param index index of the value to start */ public int startNewValue(int index) { while (index >= getValidityBufferValueCapacity()) { reallocValidityBuffer(); } BitVectorHelper.setBit(validityBuffer, index); return index * listSize; } public UnionFixedSizeListWriter getWriter() { return new UnionFixedSizeListWriter(this); } @Override public void setInitialCapacity(int numRecords) { validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); vector.setInitialCapacity(numRecords * listSize); } @Override public int getValueCapacity() { if (vector == ZeroVector.INSTANCE || listSize == 0) { return 0; } return Math.min(vector.getValueCapacity() / listSize, getValidityBufferValueCapacity()); } @Override public int getBufferSize() { if (getValueCount() == 0) { return 0; } return BitVectorHelper.getValidityBufferSizeFromCount(valueCount) + vector.getBufferSize(); } @Override public int getBufferSizeFor(int valueCount) { if (valueCount == 0) { return 0; } return BitVectorHelper.getValidityBufferSizeFromCount(valueCount) + vector.getBufferSizeFor(valueCount * listSize); } @Override public Iterator iterator() { return Collections.singleton(vector).iterator(); } @Override public void clear() { validityBuffer = releaseBuffer(validityBuffer); vector.clear(); valueCount = 0; super.clear(); } @Override public void reset() { validityBuffer.setZero(0, validityBuffer.capacity()); vector.reset(); valueCount = 0; } /** * Return the underlying buffers associated with this vector. Note that this doesn't impact the * reference counts for this buffer, so it only should be used for in-context access. Also note * that this buffer changes regularly, thus external classes shouldn't hold a reference to it * (unless they change it). * * @param clear Whether to clear vector before returning, the buffers will still be refcounted but * the returned array will be the only reference to them. Also, this won't clear the child * buffers. * @return The underlying {@link ArrowBuf buffers} that is used by this vector instance. */ @Override public ArrowBuf[] getBuffers(boolean clear) { setReaderAndWriterIndex(); final ArrowBuf[] buffers; if (getBufferSize() == 0) { buffers = new ArrowBuf[0]; } else { List list = new ArrayList<>(); list.add(validityBuffer); list.addAll(Arrays.asList(vector.getBuffers(false))); buffers = list.toArray(new ArrowBuf[list.size()]); } if (clear) { for (ArrowBuf buffer : buffers) { buffer.getReferenceManager().retain(); } clear(); } return buffers; } /** * Get value indicating if inner vector is set. * * @return 1 if inner vector is explicitly set via #addOrGetVector else 0 */ public int size() { return vector == ZeroVector.INSTANCE ? 0 : 1; } @Override @SuppressWarnings("unchecked") public AddOrGetResult addOrGetVector(FieldType type) { boolean created = false; if (vector == ZeroVector.INSTANCE) { vector = type.createNewSingleVector(DATA_VECTOR_NAME, allocator, null); invalidateReader(); created = true; } // returned vector must have the same field if (!Objects.equals(vector.getField().getType(), type.getType())) { final String msg = String.format( "Inner vector type mismatch. Requested type: [%s], actual type: [%s]", type.getType(), vector.getField().getType()); throw new SchemaChangeRuntimeException(msg); } return new AddOrGetResult<>((T) vector, created); } @Override public void copyFromSafe(int inIndex, int outIndex, ValueVector from) { copyFrom(inIndex, outIndex, from); } @Override public void copyFrom(int fromIndex, int thisIndex, ValueVector from) { Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); TransferPair pair = from.makeTransferPair(this); pair.copyValueSafe(fromIndex, thisIndex); } @Override public UnionVector promoteToUnion() { UnionVector vector = new UnionVector(getName(), allocator, /* field type */ null, /* call-back */ null); this.vector.clear(); this.vector = vector; invalidateReader(); return vector; } @Override public long getValidityBufferAddress() { return validityBuffer.memoryAddress(); } @Override public long getDataBufferAddress() { throw new UnsupportedOperationException(); } @Override public long getOffsetBufferAddress() { throw new UnsupportedOperationException(); } @Override public ArrowBuf getValidityBuffer() { return validityBuffer; } @Override public ArrowBuf getDataBuffer() { throw new UnsupportedOperationException(); } @Override public ArrowBuf getOffsetBuffer() { throw new UnsupportedOperationException(); } @Override public List getObject(int index) { if (isSet(index) == 0) { return null; } final List vals = new JsonStringArrayList<>(listSize); for (int i = 0; i < listSize; i++) { vals.add(vector.getObject(index * listSize + i)); } return vals; } /** Returns whether the value at index null. */ @Override public boolean isNull(int index) { return (isSet(index) == 0); } /** Returns non-zero when the value at index is non-null. */ 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; } @Override public int getNullCount() { return BitVectorHelper.getNullCount(validityBuffer, valueCount); } @Override public int getValueCount() { return valueCount; } /** Returns the number of elements the validity buffer can represent with its current capacity. */ private int getValidityBufferValueCapacity() { return capAtMaxInt(validityBuffer.capacity() * 8); } /** Sets the value at index to null. Reallocates if index is larger than capacity. */ @Override public void setNull(int index) { while (index >= getValidityBufferValueCapacity()) { reallocValidityBuffer(); } BitVectorHelper.unsetBit(validityBuffer, index); } /** Sets the value at index to not-null. Reallocates if index is larger than capacity. */ public void setNotNull(int index) { while (index >= getValidityBufferValueCapacity()) { reallocValidityBuffer(); } BitVectorHelper.setBit(validityBuffer, index); } @Override public void setValueCount(int valueCount) { this.valueCount = valueCount; while (valueCount > getValidityBufferValueCapacity()) { reallocValidityBuffer(); } vector.setValueCount(valueCount * listSize); } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return getTransferPair(ref, allocator, null); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return getTransferPair(field, allocator, null); } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return new TransferImpl(ref, allocator, callBack); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { return new TransferImpl(field, allocator, callBack); } @Override public TransferPair makeTransferPair(ValueVector target) { return new TransferImpl((FixedSizeListVector) target); } @Override public int hashCode(int index) { return hashCode(index, null); } @Override public int hashCode(int index, ArrowBufHasher hasher) { if (isSet(index) == 0) { return ArrowBufPointer.NULL_HASH_CODE; } int hash = 0; for (int i = 0; i < listSize; i++) { hash = ByteFunctionHelpers.combineHash(hash, vector.hashCode(index * listSize + i, hasher)); } return hash; } @Override public OUT accept(VectorVisitor visitor, IN value) { return visitor.visit(this, value); } @Override public int getElementStartIndex(int index) { return listSize * index; } @Override public int getElementEndIndex(int index) { return listSize * (index + 1); } private class TransferImpl implements TransferPair { FixedSizeListVector to; TransferPair dataPair; public TransferImpl(String name, BufferAllocator allocator, CallBack callBack) { this(new FixedSizeListVector(name, allocator, field.getFieldType(), callBack)); } public TransferImpl(Field field, BufferAllocator allocator, CallBack callBack) { this(new FixedSizeListVector(field, allocator, callBack)); } public TransferImpl(FixedSizeListVector to) { this.to = to; if (!(vector instanceof ZeroVector)) { to.addOrGetVector(vector.getField().getFieldType()); } dataPair = vector.makeTransferPair(to.vector); } @Override public void transfer() { to.clear(); dataPair.transfer(); to.validityBuffer = BaseValueVector.transferBuffer(validityBuffer, to.allocator); to.setValueCount(valueCount); clear(); } @Override public void splitAndTransfer(int startIndex, int length) { Preconditions.checkArgument( startIndex >= 0 && length >= 0 && startIndex + length <= valueCount, "Invalid parameters startIndex: %s, length: %s for valueCount: %s", startIndex, length, valueCount); final int startPoint = listSize * startIndex; final int sliceLength = listSize * length; to.clear(); /* splitAndTransfer validity buffer */ splitAndTransferValidityBuffer(startIndex, length, to); /* splitAndTransfer data buffer */ dataPair.splitAndTransfer(startPoint, sliceLength); to.setValueCount(length); } @Override public ValueVector getTo() { return to; } @Override public void copyValueSafe(int fromIndex, int toIndex) { while (toIndex >= to.getValueCapacity()) { to.reAlloc(); } BitVectorHelper.setValidityBit(to.validityBuffer, toIndex, isSet(fromIndex)); int fromOffset = fromIndex * listSize; int toOffset = toIndex * listSize; for (int i = 0; i < listSize; i++) { dataPair.copyValueSafe(fromOffset + i, toOffset + i); } } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java000066400000000000000000001101041515451615200330600ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static java.util.Collections.singletonList; import static org.apache.arrow.memory.util.LargeMemoryUtil.capAtMaxInt; import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import static org.apache.arrow.util.Preconditions.checkArgument; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.OutOfMemoryException; 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; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.BaseValueVector; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.BufferBacked; import org.apache.arrow.vector.DensityAwareVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.NullVector; import org.apache.arrow.vector.UInt4Vector; import org.apache.arrow.vector.ValueIterableVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.ZeroVector; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.impl.ComplexCopier; 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.ipc.message.ArrowFieldNode; 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.CallBack; import org.apache.arrow.vector.util.JsonStringArrayList; import org.apache.arrow.vector.util.OversizedAllocationException; import org.apache.arrow.vector.util.SchemaChangeRuntimeException; import org.apache.arrow.vector.util.TransferPair; /** * A list vector contains lists of a specific type of elements. Its structure contains 3 elements. * *
      *
    1. A validity buffer. *
    2. An offset buffer, that denotes lists boundaries. *
    3. A child data vector that contains the elements of lists. *
    * * This is the LargeList variant of list, it has a 64-bit wide offset * *

    WARNING: Currently Arrow in Java doesn't support 64-bit vectors. This class follows the * expected behaviour of a LargeList but doesn't actually support allocating a 64-bit vector. It has * little use until 64-bit vectors are supported and should be used with caution. todo review * checkedCastToInt usage in this class. Once int64 indexed vectors are supported these checks * aren't needed. */ public class LargeListVector extends BaseValueVector implements RepeatedValueVector, FieldVector, PromotableVector, ValueIterableVector> { public static LargeListVector empty(String name, BufferAllocator allocator) { return new LargeListVector( name, allocator, FieldType.nullable(ArrowType.LargeList.INSTANCE), null); } public static final FieldVector DEFAULT_DATA_VECTOR = ZeroVector.INSTANCE; public static final String DATA_VECTOR_NAME = "$data$"; public static final byte OFFSET_WIDTH = 8; protected ArrowBuf offsetBuffer; protected FieldVector vector; protected final CallBack callBack; protected long offsetAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH; protected String defaultDataVectorName = DATA_VECTOR_NAME; protected UnionLargeListReader reader; private Field field; private int validityAllocationSizeInBytes; /** The maximum index that is actually set. */ private int lastSet; /** * Constructs a new instance. * * @param name The name of the instance. * @param allocator The allocator to use for allocating/reallocating buffers. * @param fieldType The type of this list. * @param callBack A schema change callback. */ public LargeListVector( String name, BufferAllocator allocator, FieldType fieldType, CallBack callBack) { this(new Field(name, fieldType, null), allocator, callBack); } /** * Creates a new instance. * * @param field The field materialized by this vector. * @param allocator The allocator to use for creating/reallocating buffers for the vector. * @param callBack A schema change callback. */ public LargeListVector(Field field, BufferAllocator allocator, CallBack callBack) { super(allocator); this.field = field; this.validityBuffer = allocator.getEmpty(); this.callBack = callBack; this.validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION); this.lastSet = -1; this.offsetBuffer = allocator.getEmpty(); this.vector = vector == null ? DEFAULT_DATA_VECTOR : vector; this.valueCount = 0; } @Override public void initializeChildrenFromFields(List children) { checkArgument( children.size() == 1, "Lists have one child Field. Found: %s", children.isEmpty() ? "none" : children); Field field = children.get(0); AddOrGetResult addOrGetVector = addOrGetVector(field.getFieldType()); checkArgument( addOrGetVector.isCreated(), "Child vector already existed: %s", addOrGetVector.getVector()); addOrGetVector.getVector().initializeChildrenFromFields(field.getChildren()); this.field = new Field(this.field.getName(), this.field.getFieldType(), children); } @Override public void setInitialCapacity(int 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); } else { vector.setInitialCapacity(numRecords); } } /** * Specialized version of setInitialCapacity() for ListVector. This is used by some callers when * they want to explicitly control and be conservative about memory allocated for inner data * vector. This is very useful when we are working with memory constraints for a query and have a * fixed amount of memory reserved for the record batch. In such cases, we are likely to face OOM * or related problems when we reserve memory for a record batch with value count x and do * setInitialCapacity(x) such that each vector allocates only what is necessary and not the * default amount but the multiplier forces the memory requirement to go beyond what was needed. * * @param numRecords value count * @param density density of ListVector. Density is the average size of list per position in the * List vector. For example, a density value of 10 implies each position in the list vector * has a list of 10 values. A density value of 0.1 implies out of 10 positions in the list * vector, 1 position has a list of size 1 and remaining positions are null (no lists) or * empty lists. This helps in tightly controlling the memory we provision for inner data * vector. */ @Override public void setInitialCapacity(int numRecords, double density) { validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); if ((numRecords * density) >= Integer.MAX_VALUE) { throw new OversizedAllocationException("Requested amount of memory is more than max allowed"); } offsetAllocationSizeInBytes = (numRecords + 1L) * OFFSET_WIDTH; int innerValueCapacity = Math.max((int) (numRecords * density), 1); if (vector instanceof DensityAwareVector) { ((DensityAwareVector) vector).setInitialCapacity(innerValueCapacity, density); } else { vector.setInitialCapacity(innerValueCapacity); } } /** * Specialized version of setInitialTotalCapacity() for ListVector. This is used by some callers * when they want to explicitly control and be conservative about memory allocated for inner data * vector. This is very useful when we are working with memory constraints for a query and have a * fixed amount of memory reserved for the record batch. In such cases, we are likely to face OOM * or related problems when we reserve memory for a record batch with value count x and do * setInitialCapacity(x) such that each vector allocates only what is necessary and not the * default amount but the multiplier forces the memory requirement to go beyond what was needed. * * @param numRecords value count * @param totalNumberOfElements the total number of elements to to allow for in this vector across * all records. */ public void setInitialTotalCapacity(int numRecords, int totalNumberOfElements) { offsetAllocationSizeInBytes = (numRecords + 1L) * OFFSET_WIDTH; vector.setInitialCapacity(totalNumberOfElements); } /** * Get the density of this ListVector. * * @return density */ public double getDensity() { if (valueCount == 0) { return 0.0D; } final long startOffset = offsetBuffer.getLong(0L); final long endOffset = offsetBuffer.getLong((long) valueCount * OFFSET_WIDTH); final double totalListSize = endOffset - startOffset; return totalListSize / valueCount; } @Override public List getChildrenFromFields() { return singletonList(getDataVector()); } /** * Load the buffers of this vector with provided source buffers. The caller manages the source * buffers and populates them before invoking this method. * * @param fieldNode the fieldNode indicating the value count * @param ownBuffers the buffers for this Field (own buffers only, children not included) */ @Override public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers) { if (ownBuffers.size() != 2) { throw new IllegalArgumentException( "Illegal buffer count, expected " + 2 + ", got: " + ownBuffers.size()); } ArrowBuf bitBuffer = ownBuffers.get(0); ArrowBuf offBuffer = ownBuffers.get(1); validityBuffer.getReferenceManager().release(); validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); offsetBuffer.getReferenceManager().release(); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); validityAllocationSizeInBytes = checkedCastToInt(validityBuffer.capacity()); offsetAllocationSizeInBytes = offsetBuffer.capacity(); lastSet = fieldNode.getLength() - 1; valueCount = fieldNode.getLength(); } /** * Get the buffers belonging to this vector. * * @return the inner buffers. */ @Override public List getFieldBuffers() { List result = new ArrayList<>(2); setReaderAndWriterIndex(); result.add(validityBuffer); result.add(offsetBuffer); return result; } /** * Export the buffers of the fields for C Data Interface. This method traverse the buffers and * export buffer and buffer's memory address into a list of buffers and a pointer to the list of * buffers. */ @Override public void exportCDataBuffers(List buffers, ArrowBuf buffersPtr, long nullValue) { exportBuffer(validityBuffer, buffers, buffersPtr, nullValue, true); if (offsetBuffer.capacity() == 0) { // Empty offset buffer is allowed for historical reason. // To export it through C Data interface, we need to allocate a buffer with one offset. // We set `retain = false` to explicitly not increase the ref count for the exported buffer. // The ref count of the newly created buffer (i.e., 1) already represents the usage // at imported side. exportBuffer(allocateOffsetBuffer(OFFSET_WIDTH), buffers, buffersPtr, nullValue, false); } else { exportBuffer(offsetBuffer, buffers, buffersPtr, nullValue, true); } } /** Set the reader and writer indexes for the inner buffers. */ private void setReaderAndWriterIndex() { validityBuffer.readerIndex(0); offsetBuffer.readerIndex(0); if (valueCount == 0) { validityBuffer.writerIndex(0); } else { 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); } /** * Get the inner vectors. * * @return the inner vectors for this field as defined by the TypeLayout * @deprecated This API will be removed as the current implementations no longer support inner * vectors. */ @Deprecated @Override public List getFieldInnerVectors() { throw new UnsupportedOperationException("There are no inner vectors. Use getFieldBuffers"); } /** Same as {@link #allocateNewSafe()}. */ @Override public void allocateNew() throws OutOfMemoryException { if (!allocateNewSafe()) { throw new OutOfMemoryException("Failure while allocating memory"); } } /** * Allocate memory for the vector. We internally use a default value count of 4096 to allocate * memory for at least these many elements in the vector. * * @return false if memory allocation fails, true otherwise. */ @Override public boolean allocateNewSafe() { boolean success = false; try { /* we are doing a new allocation -- release the current buffers */ clear(); /* allocate validity buffer */ allocateValidityBuffer(validityAllocationSizeInBytes); /* allocate offset and data buffer */ boolean dataAlloc = false; try { offsetBuffer = allocateOffsetBuffer(offsetAllocationSizeInBytes); dataAlloc = vector.allocateNewSafe(); } catch (Exception e) { e.printStackTrace(); clear(); success = false; } finally { if (!dataAlloc) { clear(); } } success = dataAlloc; } finally { if (!success) { clear(); } } return success; } @Override protected void allocateValidityBuffer(final long size) { super.allocateValidityBuffer(size); validityAllocationSizeInBytes = (int) size; } protected ArrowBuf allocateOffsetBuffer(final long size) { ArrowBuf offsetBuffer = allocator.buffer(size); offsetBuffer.readerIndex(0); offsetAllocationSizeInBytes = size; offsetBuffer.setZero(0, offsetBuffer.capacity()); return offsetBuffer; } /** * Resize the vector to increase the capacity. The internal behavior is to double the current * value capacity. */ @Override public void reAlloc() { /* reallocate the validity buffer */ reallocValidityBuffer(); /* reallocate the offset and data */ reallocOffsetBuffer(); vector.reAlloc(); } private void reallocValidityAndOffsetBuffers() { reallocOffsetBuffer(); reallocValidityBuffer(); } protected void reallocOffsetBuffer() { final long currentBufferCapacity = offsetBuffer.capacity(); long newAllocationSize = currentBufferCapacity * 2; if (newAllocationSize == 0) { if (offsetAllocationSizeInBytes > 0) { newAllocationSize = offsetAllocationSizeInBytes; } else { newAllocationSize = INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH * 2; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); newAllocationSize = Math.min(newAllocationSize, (long) OFFSET_WIDTH * Integer.MAX_VALUE); assert newAllocationSize >= 1; if (newAllocationSize > MAX_ALLOCATION_SIZE || newAllocationSize <= offsetBuffer.capacity()) { throw new OversizedAllocationException("Unable to expand the buffer"); } final ArrowBuf newBuf = allocator.buffer(newAllocationSize); newBuf.setBytes(0, offsetBuffer, 0, currentBufferCapacity); newBuf.setZero(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); offsetBuffer.getReferenceManager().release(1); offsetBuffer = newBuf; offsetAllocationSizeInBytes = newAllocationSize; } private void reallocValidityBuffer() { final int currentBufferCapacity = checkedCastToInt(validityBuffer.capacity()); long newAllocationSize = currentBufferCapacity * 2L; if (newAllocationSize == 0) { if (validityAllocationSizeInBytes > 0) { newAllocationSize = validityAllocationSizeInBytes; } else { newAllocationSize = BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); assert newAllocationSize >= 1; if (newAllocationSize > MAX_ALLOCATION_SIZE) { throw new OversizedAllocationException("Unable to expand the buffer"); } final ArrowBuf newBuf = allocator.buffer(newAllocationSize); newBuf.setBytes(0, validityBuffer, 0, currentBufferCapacity); newBuf.setZero(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); validityBuffer.getReferenceManager().release(1); validityBuffer = newBuf; validityAllocationSizeInBytes = (int) newAllocationSize; } /** * Same as {@link #copyFrom(int, int, ValueVector)} except that it handles the case when the * capacity of the vector needs to be expanded before copy. * * @param inIndex position to copy from in source vector * @param outIndex position to copy to in this vector * @param from source vector */ @Override public void copyFromSafe(int inIndex, int outIndex, ValueVector from) { copyFrom(inIndex, outIndex, from); } /** * Copy a cell value from a particular index in source vector to a particular position in this * vector. * * @param inIndex position to copy from in source vector * @param outIndex position to copy to in this vector * @param from source vector */ @Override public void copyFrom(int inIndex, int outIndex, ValueVector from) { Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); FieldReader in = from.getReader(); in.setPosition(inIndex); UnionLargeListWriter out = getWriter(); out.setPosition(outIndex); ComplexCopier.copy(in, out); } /** * Get the offset vector. * * @return the underlying offset vector or null if none exists. * @deprecated This API will be removed, as the current implementations no longer hold inner * offset vectors. */ @Override @Deprecated public UInt4Vector getOffsetVector() { throw new UnsupportedOperationException("There is no inner offset vector"); } /** * Get the inner data vector for this list vector. * * @return data vector */ @Override public FieldVector getDataVector() { return vector; } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return getTransferPair(ref, allocator, null); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return getTransferPair(field, allocator, null); } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return new TransferImpl(ref, allocator, callBack); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { return new TransferImpl(field, allocator, callBack); } @Override public TransferPair makeTransferPair(ValueVector target) { return new TransferImpl((LargeListVector) target); } @Override public long getValidityBufferAddress() { return validityBuffer.memoryAddress(); } @Override public long getDataBufferAddress() { throw new UnsupportedOperationException(); } @Override public long getOffsetBufferAddress() { return offsetBuffer.memoryAddress(); } @Override public ArrowBuf getValidityBuffer() { return validityBuffer; } @Override public ArrowBuf getDataBuffer() { throw new UnsupportedOperationException(); } @Override public ArrowBuf getOffsetBuffer() { return offsetBuffer; } @Override public int getValueCount() { return valueCount; } @Override public int hashCode(int index) { return hashCode(index, null); } @Override public int hashCode(int index, ArrowBufHasher hasher) { if (isSet(index) == 0) { return ArrowBufPointer.NULL_HASH_CODE; } int hash = 0; final long start = offsetBuffer.getLong((long) index * OFFSET_WIDTH); final long end = offsetBuffer.getLong(((long) index + 1L) * OFFSET_WIDTH); for (long i = start; i < end; i++) { hash = ByteFunctionHelpers.combineHash(hash, vector.hashCode(checkedCastToInt(i), hasher)); } return hash; } @Override public OUT accept(VectorVisitor visitor, IN value) { return visitor.visit(this, value); } public UnionLargeListWriter getWriter() { return new UnionLargeListWriter(this); } protected void replaceDataVector(FieldVector v) { vector.clear(); vector = v; } @Override public UnionVector promoteToUnion() { UnionVector vector = new UnionVector("$data$", allocator, /* field type */ null, callBack); replaceDataVector(vector); invalidateReader(); if (callBack != null) { callBack.doWork(); } return vector; } private class TransferImpl implements TransferPair { LargeListVector to; TransferPair dataTransferPair; public TransferImpl(String name, BufferAllocator allocator, CallBack callBack) { this(new LargeListVector(name, allocator, field.getFieldType(), callBack)); } public TransferImpl(Field field, BufferAllocator allocator, CallBack callBack) { this(new LargeListVector(field, allocator, callBack)); } public TransferImpl(LargeListVector to) { this.to = to; to.addOrGetVector(vector.getField().getFieldType()); if (to.getDataVector() instanceof ZeroVector) { to.addOrGetVector(vector.getField().getFieldType()); } dataTransferPair = getDataVector().makeTransferPair(to.getDataVector()); } /** * Transfer this vector'data to another vector. The memory associated with this vector is * transferred to the allocator of target vector for accounting and management purposes. */ @Override public void transfer() { to.clear(); dataTransferPair.transfer(); to.validityBuffer = transferBuffer(validityBuffer, to.allocator); to.offsetBuffer = transferBuffer(offsetBuffer, to.allocator); to.lastSet = lastSet; if (valueCount > 0) { to.setValueCount(valueCount); } clear(); } /** * Slice this vector at desired index and length and transfer the corresponding data to the * target vector. * * @param startIndex start position of the split in source vector. * @param length length of the split. */ @Override public void splitAndTransfer(int startIndex, int length) { Preconditions.checkArgument( startIndex >= 0 && length >= 0 && startIndex + length <= valueCount, "Invalid parameters startIndex: %s, length: %s for valueCount: %s", startIndex, length, valueCount); final long startPoint = offsetBuffer.getLong((long) startIndex * OFFSET_WIDTH); final long sliceLength = offsetBuffer.getLong((long) (startIndex + length) * OFFSET_WIDTH) - startPoint; to.clear(); to.offsetBuffer = to.allocateOffsetBuffer((length + 1) * OFFSET_WIDTH); /* splitAndTransfer offset buffer */ for (int i = 0; i < length + 1; i++) { final long relativeOffset = offsetBuffer.getLong((long) (startIndex + i) * OFFSET_WIDTH) - startPoint; to.offsetBuffer.setLong((long) i * OFFSET_WIDTH, relativeOffset); } /* splitAndTransfer validity buffer */ splitAndTransferValidityBuffer(startIndex, length, to); /* splitAndTransfer data buffer */ dataTransferPair.splitAndTransfer( checkedCastToInt(startPoint), checkedCastToInt(sliceLength)); to.lastSet = length - 1; to.setValueCount(length); } @Override public ValueVector getTo() { return to; } @Override public void copyValueSafe(int from, int to) { this.to.copyFrom(from, to, LargeListVector.this); } } @Override protected FieldReader getReaderImpl() { return new UnionLargeListReader(this); } @Override public UnionLargeListReader getReader() { reader = (UnionLargeListReader) super.getReader(); return reader; } /** * Initialize the data vector (and execute callback) if it hasn't already been done, returns the * data vector. */ @Override public AddOrGetResult addOrGetVector(FieldType fieldType) { boolean created = false; if (vector instanceof NullVector) { vector = fieldType.createNewSingleVector(defaultDataVectorName, allocator, callBack); // returned vector must have the same field created = true; if (callBack != null && // not a schema change if changing from ZeroVector to ZeroVector (fieldType.getType().getTypeID() != ArrowType.ArrowTypeID.Null)) { callBack.doWork(); } } if (vector.getField().getType().getTypeID() != fieldType.getType().getTypeID()) { final String msg = String.format( "Inner vector type mismatch. Requested type: [%s], actual type: [%s]", fieldType.getType().getTypeID(), vector.getField().getType().getTypeID()); throw new SchemaChangeRuntimeException(msg); } invalidateReader(); return new AddOrGetResult<>((T) vector, created); } /** * Get the size (number of bytes) of underlying buffers used by this vector. * * @return size of underlying buffers. */ @Override public int getBufferSize() { if (valueCount == 0) { return 0; } final int offsetBufferSize = (valueCount + 1) * OFFSET_WIDTH; final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); return offsetBufferSize + validityBufferSize + vector.getBufferSize(); } @Override public int getBufferSizeFor(int valueCount) { if (valueCount == 0) { return 0; } final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); long innerVectorValueCount = offsetBuffer.getLong((long) valueCount * OFFSET_WIDTH); return ((valueCount + 1) * OFFSET_WIDTH) + vector.getBufferSizeFor(checkedCastToInt(innerVectorValueCount)) + validityBufferSize; } @Override public Field getField() { if (field.getChildren().contains(getDataVector().getField())) { return field; } field = new Field( field.getName(), field.getFieldType(), Collections.singletonList(getDataVector().getField())); return field; } @Override public MinorType getMinorType() { return MinorType.LARGELIST; } @Override public String getName() { return field.getName(); } @Override public void clear() { offsetBuffer = releaseBuffer(offsetBuffer); vector.clear(); valueCount = 0; super.clear(); validityBuffer = releaseBuffer(validityBuffer); lastSet = -1; } @Override public void reset() { offsetBuffer.setZero(0, offsetBuffer.capacity()); vector.reset(); valueCount = 0; validityBuffer.setZero(0, validityBuffer.capacity()); lastSet = -1; } /** * Return the underlying buffers associated with this vector. Note that this doesn't impact the * reference counts for this buffer, so it only should be used for in-context access. Also note * that this buffer changes regularly, thus external classes shouldn't hold a reference to it * (unless they change it). * * @param clear Whether to clear vector before returning, the buffers will still be refcounted but * the returned array will be the only reference to them. Also, this won't clear the child * buffers. * @return The underlying {@link ArrowBuf buffers} that is used by this vector instance. */ @Override public ArrowBuf[] getBuffers(boolean clear) { setReaderAndWriterIndex(); final ArrowBuf[] buffers; if (getBufferSize() == 0) { buffers = new ArrowBuf[0]; } else { List list = new ArrayList<>(); list.add(offsetBuffer); list.add(validityBuffer); list.addAll(Arrays.asList(vector.getBuffers(false))); buffers = list.toArray(new ArrowBuf[list.size()]); } if (clear) { for (ArrowBuf buffer : buffers) { buffer.getReferenceManager().retain(); } clear(); } return buffers; } protected void invalidateReader() { reader = null; } /** * Get the element in the list vector at a particular index. * * @param index position of the element * @return Object at given position */ @Override public List getObject(int index) { if (isSet(index) == 0) { return null; } 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))); } return vals; } /** * Check if element at given index is null. * * @param index position of element * @return true if element at given index is null, false otherwise */ @Override public boolean isNull(int index) { return (isSet(index) == 0); } /** * Check if element at given index is empty list. * * @param index position of element * @return true if element at given index is empty list or NULL, false otherwise */ public boolean isEmpty(int index) { if (isNull(index)) { return true; } else { final long start = offsetBuffer.getLong((long) index * OFFSET_WIDTH); final long end = offsetBuffer.getLong(((long) index + 1L) * OFFSET_WIDTH); return start == end; } } /** * Same as {@link #isNull(int)}. * * @param index position of element * @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; } /** * Get the number of elements that are null in the vector. * * @return the number of null elements. */ @Override public int getNullCount() { return BitVectorHelper.getNullCount(validityBuffer, valueCount); } /** * Get the current value capacity for the vector. * * @return number of elements that vector can hold. */ @Override public int getValueCapacity() { return getValidityAndOffsetValueCapacity(); } protected int getOffsetBufferValueCapacity() { return checkedCastToInt(offsetBuffer.capacity() / OFFSET_WIDTH); } private int getValidityAndOffsetValueCapacity() { final int offsetValueCapacity = Math.max(getOffsetBufferValueCapacity() - 1, 0); return Math.min(offsetValueCapacity, getValidityBufferValueCapacity()); } private int getValidityBufferValueCapacity() { return capAtMaxInt(validityBuffer.capacity() * 8); } /** * Sets the list at index to be not-null. Reallocates validity buffer if index is larger than * current capacity. */ public void setNotNull(int index) { while (index >= getValidityAndOffsetValueCapacity()) { reallocValidityAndOffsetBuffers(); } BitVectorHelper.setBit(validityBuffer, index); lastSet = index; } /** * Sets list at index to be null. * * @param index position in vector */ @Override public void setNull(int index) { while (index >= getValidityAndOffsetValueCapacity()) { reallocValidityAndOffsetBuffers(); } if (lastSet >= index) { lastSet = index - 1; } for (int i = lastSet + 1; i <= index; i++) { final int currentOffset = offsetBuffer.getInt(i * OFFSET_WIDTH); offsetBuffer.setInt((i + 1) * OFFSET_WIDTH, currentOffset); } BitVectorHelper.unsetBit(validityBuffer, index); } /** * Start a new value in the list vector. * * @param index index of the value to start */ public long startNewValue(long index) { while (index >= getValidityAndOffsetValueCapacity()) { reallocValidityAndOffsetBuffers(); } for (int i = lastSet + 1; i <= index; i++) { final long currentOffset = offsetBuffer.getLong((long) i * OFFSET_WIDTH); offsetBuffer.setLong(((long) i + 1L) * OFFSET_WIDTH, currentOffset); } BitVectorHelper.setBit(validityBuffer, index); lastSet = checkedCastToInt(index); return offsetBuffer.getLong(((long) lastSet + 1L) * OFFSET_WIDTH); } /** * End the current value. * * @param index index of the value to end * @param size number of elements in the list that was written */ public void endValue(int index, long size) { final long currentOffset = offsetBuffer.getLong(((long) index + 1L) * OFFSET_WIDTH); offsetBuffer.setLong(((long) index + 1L) * OFFSET_WIDTH, currentOffset + size); } /** * Sets the value count for the vector. * *

    Important note: The underlying vector does not support 64-bit allocations yet. This may * throw if attempting to hold larger than what a 32-bit vector can store. * * @param valueCount value count */ @Override public void setValueCount(int valueCount) { this.valueCount = valueCount; if (valueCount > 0) { while (valueCount > getValidityAndOffsetValueCapacity()) { /* check if validity and offset buffers need to be re-allocated */ reallocValidityAndOffsetBuffers(); } for (int i = lastSet + 1; i < valueCount; i++) { /* fill the holes with offsets */ final long currentOffset = offsetBuffer.getLong((long) i * OFFSET_WIDTH); offsetBuffer.setLong(((long) i + 1L) * OFFSET_WIDTH, currentOffset); } } /* valueCount for the data vector is the current end offset */ final long childValueCount = (valueCount == 0) ? 0 : offsetBuffer.getLong(((long) lastSet + 1L) * OFFSET_WIDTH); /* set the value count of data vector and this will take care of * checking whether data buffer needs to be reallocated. * TODO: revisit when 64-bit vectors are supported */ Preconditions.checkArgument( childValueCount <= Integer.MAX_VALUE || childValueCount >= Integer.MIN_VALUE, "LargeListVector doesn't yet support 64-bit allocations: %s", childValueCount); vector.setValueCount((int) childValueCount); } public void setLastSet(int value) { lastSet = value; } public int getLastSet() { return lastSet; } public long getElementStartIndex(int index) { return offsetBuffer.getLong((long) index * OFFSET_WIDTH); } public long getElementEndIndex(int index) { return offsetBuffer.getLong(((long) index + 1L) * OFFSET_WIDTH); } } LargeListViewVector.java000066400000000000000000001000761515451615200336430ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static java.util.Collections.singletonList; import static org.apache.arrow.memory.util.LargeMemoryUtil.capAtMaxInt; import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import static org.apache.arrow.util.Preconditions.checkArgument; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.OutOfMemoryException; 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.hash.ArrowBufHasher; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.AddOrGetResult; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.BufferBacked; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueIterableVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.ZeroVector; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.impl.UnionLargeListViewReader; import org.apache.arrow.vector.complex.impl.UnionLargeListViewWriter; import org.apache.arrow.vector.complex.impl.UnionListReader; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; 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.CallBack; import org.apache.arrow.vector.util.JsonStringArrayList; import org.apache.arrow.vector.util.OversizedAllocationException; import org.apache.arrow.vector.util.TransferPair; /** * A large list view vector contains lists of a specific type of elements. Its structure contains 3 * elements. * *

      *
    1. A validity buffer. *
    2. An offset buffer, that denotes lists starting positions. *
    3. A size buffer, that denotes sizes of the lists. *
    4. A child data vector that contains the elements of lists. *
    * * This is the LargeListView variant of listview, it has a 64-bit wide offset * *

    WARNING: Currently Arrow in Java doesn't support 64-bit vectors. This class follows the * expected behaviour of a LargeList but doesn't actually support allocating a 64-bit vector. It has * little use until 64-bit vectors are supported and should be used with caution. todo review * checkedCastToInt usage in this class. Once int64 indexed vectors are supported these checks * aren't needed. */ public class LargeListViewVector extends BaseLargeRepeatedValueViewVector implements PromotableVector, ValueIterableVector> { protected UnionLargeListViewReader reader; private CallBack callBack; protected Field field; protected int validityAllocationSizeInBytes; public static LargeListViewVector empty(String name, BufferAllocator allocator) { return new LargeListViewVector( name, allocator, FieldType.nullable(ArrowType.LargeListView.INSTANCE), null); } /** * Constructs a new instance. * * @param name The name of the instance. * @param allocator The allocator to use for allocating/reallocating buffers. * @param fieldType The type of this list. * @param callBack A schema change callback. */ public LargeListViewVector( String name, BufferAllocator allocator, FieldType fieldType, CallBack callBack) { this(new Field(name, fieldType, null), allocator, callBack); } /** * Constructs a new instance. * * @param field The field materialized by this vector. * @param allocator The allocator to use for allocating/reallocating buffers. * @param callBack A schema change callback. */ public LargeListViewVector(Field field, BufferAllocator allocator, CallBack callBack) { super(field.getName(), allocator, callBack); this.validityBuffer = allocator.getEmpty(); this.field = field; this.callBack = callBack; this.validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION); } @Override public void initializeChildrenFromFields(List children) { checkArgument( children.size() == 1, "ListViews have one child Field. Found: %s", children.isEmpty() ? "none" : children); Field field = children.get(0); AddOrGetResult addOrGetVector = addOrGetVector(field.getFieldType()); checkArgument( addOrGetVector.isCreated(), "Child vector already existed: %s", addOrGetVector.getVector()); addOrGetVector.getVector().initializeChildrenFromFields(field.getChildren()); this.field = new Field(this.field.getName(), this.field.getFieldType(), children); } @Override public void setInitialCapacity(int numRecords) { validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialCapacity(numRecords); } /** * Specialized version of setInitialCapacity() for LargeListViewVector. This is used by some * callers when they want to explicitly control and be conservative about memory allocated for * inner data vector. This is very useful when we are working with memory constraints for a query * and have a fixed amount of memory reserved for the record batch. In such cases, we are likely * to face OOM or related problems when we reserve memory for a record batch with value count x * and do setInitialCapacity(x) such that each vector allocates only what is necessary and not the * default amount, but the multiplier forces the memory requirement to go beyond what was needed. * * @param numRecords value count * @param density density of LargeListViewVector. Density is the average size of a list per * position in the LargeListViewVector. For example, a density value of 10 implies each * position in the list vector has a list of 10 values. A density value of 0.1 implies out of * 10 positions in the list vector, 1 position has a list of size 1, and the remaining * positions are null (no lists) or empty lists. This helps in tightly controlling the memory * we provision for inner data vector. */ @Override public void setInitialCapacity(int numRecords, double density) { validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialCapacity(numRecords, density); } /** * Specialized version of setInitialTotalCapacity() for LargeListViewVector. This is used by some * callers when they want to explicitly control and be conservative about memory allocated for * inner data vector. This is very useful when we are working with memory constraints for a query * and have a fixed amount of memory reserved for the record batch. In such cases, we are likely * to face OOM or related problems when we reserve memory for a record batch with value count x * and do setInitialCapacity(x) such that each vector allocates only what is necessary and not the * default amount, but the multiplier forces the memory requirement to go beyond what was needed. * * @param numRecords value count * @param totalNumberOfElements the total number of elements to allow for in this vector across * all records. */ @Override public void setInitialTotalCapacity(int numRecords, int totalNumberOfElements) { validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialTotalCapacity(numRecords, totalNumberOfElements); } @Override public List getChildrenFromFields() { return singletonList(getDataVector()); } /** * Load the buffers associated with this Field. * * @param fieldNode the fieldNode * @param ownBuffers the buffers for this Field (own buffers only, children not included) */ @Override public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers) { if (ownBuffers.size() != 3) { throw new IllegalArgumentException( "Illegal buffer count, expected " + 3 + ", got: " + ownBuffers.size()); } ArrowBuf bitBuffer = ownBuffers.get(0); ArrowBuf offBuffer = ownBuffers.get(1); ArrowBuf szBuffer = ownBuffers.get(2); validityBuffer.getReferenceManager().release(); validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); offsetBuffer.getReferenceManager().release(); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); sizeBuffer.getReferenceManager().release(); sizeBuffer = szBuffer.getReferenceManager().retain(szBuffer, allocator); validityAllocationSizeInBytes = checkedCastToInt(validityBuffer.capacity()); offsetAllocationSizeInBytes = offsetBuffer.capacity(); sizeAllocationSizeInBytes = sizeBuffer.capacity(); valueCount = fieldNode.getLength(); } /** Set the reader and writer indexes for the inner buffers. */ private void setReaderAndWriterIndex() { validityBuffer.readerIndex(0); offsetBuffer.readerIndex(0); sizeBuffer.readerIndex(0); if (valueCount == 0) { validityBuffer.writerIndex(0); offsetBuffer.writerIndex(0); sizeBuffer.writerIndex(0); } else { validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); offsetBuffer.writerIndex((long) valueCount * OFFSET_WIDTH); sizeBuffer.writerIndex((long) valueCount * SIZE_WIDTH); } } @Override public List getFieldBuffers() { List result = new ArrayList<>(2); setReaderAndWriterIndex(); result.add(validityBuffer); result.add(offsetBuffer); result.add(sizeBuffer); return result; } /** * Export the buffers of the fields for C Data Interface. This method traverses the buffers and * export buffer and buffer's memory address into a list of buffers and a pointer to the list of * buffers. */ @Override public void exportCDataBuffers(List buffers, ArrowBuf buffersPtr, long nullValue) { exportBuffer(validityBuffer, buffers, buffersPtr, nullValue, true); exportBuffer(offsetBuffer, buffers, buffersPtr, nullValue, true); exportBuffer(sizeBuffer, buffers, buffersPtr, nullValue, true); } @Override public void allocateNew() throws OutOfMemoryException { if (!allocateNewSafe()) { throw new OutOfMemoryException("Failure while allocating memory"); } } @Override public boolean allocateNewSafe() { boolean success = false; try { /* release the current buffers, hence this is a new allocation * Note that, the `clear` method call below is releasing validityBuffer * calling the superclass clear method which is releasing the associated buffers * (sizeBuffer and offsetBuffer). */ clear(); /* allocate validity buffer */ allocateValidityBuffer(validityAllocationSizeInBytes); /* allocate offset, data and sizes buffer */ success = super.allocateNewSafe(); } finally { if (!success) { clear(); } } return success; } @Override protected void allocateValidityBuffer(final long size) { super.allocateValidityBuffer(size); validityAllocationSizeInBytes = (int) size; } @Override public void reAlloc() { /* reallocate the validity buffer */ reallocValidityBuffer(); /* reallocate the offset, size, and data */ super.reAlloc(); } protected void reallocValidityAndSizeAndOffsetBuffers() { reallocateBuffers(); reallocValidityBuffer(); } private void reallocValidityBuffer() { final int currentBufferCapacity = checkedCastToInt(validityBuffer.capacity()); long newAllocationSize = getNewAllocationSize(currentBufferCapacity); final ArrowBuf newBuf = allocator.buffer(newAllocationSize); newBuf.setBytes(0, validityBuffer, 0, currentBufferCapacity); newBuf.setZero(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); validityBuffer.getReferenceManager().release(1); validityBuffer = newBuf; validityAllocationSizeInBytes = (int) newAllocationSize; } private long getNewAllocationSize(int currentBufferCapacity) { long newAllocationSize = currentBufferCapacity * 2L; if (newAllocationSize == 0) { if (validityAllocationSizeInBytes > 0) { newAllocationSize = validityAllocationSizeInBytes; } else { newAllocationSize = BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); assert newAllocationSize >= 1; if (newAllocationSize > MAX_ALLOCATION_SIZE) { throw new OversizedAllocationException("Unable to expand the buffer"); } return newAllocationSize; } @Override public void copyFromSafe(int inIndex, int outIndex, ValueVector from) { throw new UnsupportedOperationException( "LargeListViewVector does not support copyFromSafe operation yet."); } @Override public void copyFrom(int inIndex, int outIndex, ValueVector from) { throw new UnsupportedOperationException( "LargeListViewVector does not support copyFrom operation yet."); } @Override public FieldVector getDataVector() { return vector; } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return getTransferPair(ref, allocator, null); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return getTransferPair(field, allocator, null); } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return new TransferImpl(ref, allocator, callBack); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { return new TransferImpl(field, allocator, callBack); } @Override public TransferPair makeTransferPair(ValueVector target) { return new TransferImpl((LargeListViewVector) target); } @Override public long getValidityBufferAddress() { return validityBuffer.memoryAddress(); } @Override public long getDataBufferAddress() { throw new UnsupportedOperationException(); } @Override public long getOffsetBufferAddress() { return offsetBuffer.memoryAddress(); } @Override public ArrowBuf getValidityBuffer() { return validityBuffer; } @Override public ArrowBuf getDataBuffer() { throw new UnsupportedOperationException(); } @Override public ArrowBuf getOffsetBuffer() { return offsetBuffer; } public ArrowBuf getSizeBuffer() { return sizeBuffer; } public long getSizeBufferAddress() { return sizeBuffer.memoryAddress(); } /** * Get the hash code for the element at the given index. * * @param index position of the element * @return hash code for the element at the given index */ @Override public int hashCode(int index) { return hashCode(index, null); } /** * Get the hash code for the element at the given index. * * @param index position of the element * @param hasher hasher to use * @return hash code for the element at the given index */ @Override public int hashCode(int index, ArrowBufHasher hasher) { if (isSet(index) == 0) { return ArrowBufPointer.NULL_HASH_CODE; } int hash = 0; final int start = offsetBuffer.getInt((long) index * OFFSET_WIDTH); final int end = sizeBuffer.getInt((long) index * OFFSET_WIDTH); for (int i = start; i < end; i++) { hash = ByteFunctionHelpers.combineHash(hash, vector.hashCode(checkedCastToInt(i), hasher)); } return hash; } @Override public OUT accept(VectorVisitor visitor, IN value) { return visitor.visit(this, value); } private class TransferImpl implements TransferPair { LargeListViewVector to; TransferPair dataTransferPair; public TransferImpl(String name, BufferAllocator allocator, CallBack callBack) { this(new LargeListViewVector(name, allocator, field.getFieldType(), callBack)); } public TransferImpl(Field field, BufferAllocator allocator, CallBack callBack) { this(new LargeListViewVector(field, allocator, callBack)); } public TransferImpl(LargeListViewVector to) { this.to = to; to.addOrGetVector(vector.getField().getFieldType()); if (to.getDataVector() instanceof ZeroVector) { to.addOrGetVector(vector.getField().getFieldType()); } dataTransferPair = getDataVector().makeTransferPair(to.getDataVector()); } @Override public void transfer() { to.clear(); dataTransferPair.transfer(); to.validityBuffer = transferBuffer(validityBuffer, to.allocator); to.offsetBuffer = transferBuffer(offsetBuffer, to.allocator); to.sizeBuffer = transferBuffer(sizeBuffer, to.allocator); if (valueCount > 0) { to.setValueCount(valueCount); } clear(); } @Override public void splitAndTransfer(int startIndex, int length) { Preconditions.checkArgument( startIndex >= 0 && length >= 0 && startIndex + length <= valueCount, "Invalid parameters startIndex: %s, length: %s for valueCount: %s", startIndex, length, valueCount); to.clear(); if (length > 0) { // 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); /* splitAndTransfer the size buffer */ int maxOffsetAndSizeSum = Integer.MIN_VALUE; int minOffsetValue = Integer.MAX_VALUE; 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); maxOffsetAndSizeSum = Math.max(maxOffsetAndSizeSum, offsetValue + sizeValue); minOffsetValue = Math.min(minOffsetValue, offsetValue); } /* splitAndTransfer the offset buffer */ for (int i = 0; i < length; i++) { final int offsetValue = offsetBuffer.getInt((long) (startIndex + i) * OFFSET_WIDTH); final int relativeOffset = offsetValue - minOffsetValue; to.offsetBuffer.setInt((long) i * OFFSET_WIDTH, relativeOffset); } /* splitAndTransfer the validity buffer */ splitAndTransferValidityBuffer(startIndex, length, to); /* splitAndTransfer the data buffer */ final int childSliceLength = maxOffsetAndSizeSum - minOffsetValue; dataTransferPair.splitAndTransfer(minOffsetValue, childSliceLength); to.setValueCount(length); } } @Override public ValueVector getTo() { return to; } @Override public void copyValueSafe(int from, int to) { this.to.copyFrom(from, to, LargeListViewVector.this); } } @Override protected FieldReader getReaderImpl() { throw new UnsupportedOperationException( "LargeListViewVector does not support getReaderImpl operation yet."); } @Override public UnionListReader getReader() { throw new UnsupportedOperationException( "LargeListViewVector does not support getReader operation yet."); } /** * Get the size (number of bytes) of underlying buffers used by this vector. * * @return size of underlying buffers. */ @Override public int getBufferSize() { if (valueCount == 0) { return 0; } final int offsetBufferSize = valueCount * OFFSET_WIDTH; final int sizeBufferSize = valueCount * SIZE_WIDTH; final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); return offsetBufferSize + sizeBufferSize + validityBufferSize + vector.getBufferSize(); } /** * Get the size (number of bytes) of underlying buffers used by this. * * @param valueCount the number of values to assume this vector contains * @return size of underlying buffers. */ @Override public int getBufferSizeFor(int valueCount) { if (valueCount == 0) { return 0; } final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); return super.getBufferSizeFor(valueCount) + validityBufferSize; } /** * Get the field associated with the list view vector. * * @return the field */ @Override public Field getField() { if (field.getChildren().contains(getDataVector().getField())) { return field; } field = new Field( field.getName(), field.getFieldType(), Collections.singletonList(getDataVector().getField())); return field; } /** * Get the minor type for the vector. * * @return the minor type */ @Override public MinorType getMinorType() { return MinorType.LARGELISTVIEW; } /** Clear the vector data. */ @Override public void clear() { // calling superclass clear method which is releasing the sizeBufer and offsetBuffer super.clear(); validityBuffer = releaseBuffer(validityBuffer); } /** Release the buffers associated with this vector. */ @Override public void reset() { super.reset(); validityBuffer.setZero(0, validityBuffer.capacity()); } /** * Return the underlying buffers associated with this vector. Note that this doesn't impact the * reference counts for this buffer, so it only should be used for in-context access. Also note * that this buffer changes regularly, thus external classes shouldn't hold a reference to it * (unless they change it). * * @param clear Whether to clear vector before returning, the buffers will still be refcounted but * the returned array will be the only reference to them. Also, this won't clear the child * buffers. * @return The underlying {@link ArrowBuf buffers} that is used by this vector instance. */ @Override public ArrowBuf[] getBuffers(boolean clear) { setReaderAndWriterIndex(); final ArrowBuf[] buffers; if (getBufferSize() == 0) { buffers = new ArrowBuf[0]; } else { List list = new ArrayList<>(); // the order must be validity, offset and size buffers list.add(validityBuffer); list.add(offsetBuffer); list.add(sizeBuffer); list.addAll(Arrays.asList(vector.getBuffers(false))); buffers = list.toArray(new ArrowBuf[list.size()]); } if (clear) { for (ArrowBuf buffer : buffers) { buffer.getReferenceManager().retain(); } clear(); } return buffers; } /** * Get the element in the list view vector at a particular index. * * @param index position of the element * @return Object at given position */ @Override public List getObject(int index) { if (isSet(index) == 0) { return null; } 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))); } return vals; } /** * Check if an element at given index is null. * * @param index position of an element * @return true if an element at given index is null, false otherwise */ @Override public boolean isNull(int index) { return (isSet(index) == 0); } /** * Check if an element at given index is an empty list. * * @param index position of an element * @return true if an element at given index is an empty list or NULL, false otherwise */ @Override public boolean isEmpty(int index) { if (isNull(index)) { return true; } else { return sizeBuffer.getInt(index * SIZE_WIDTH) == 0; } } /** * Same as {@link #isNull(int)}. * * @param index position of the element * @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; } /** * Get the number of elements that are null in the vector. * * @return the number of null elements. */ @Override public int getNullCount() { return BitVectorHelper.getNullCount(validityBuffer, valueCount); } /** * Get the value capacity by considering validity and offset capacity. Note that the size buffer * capacity is not considered here since it has the same capacity as the offset buffer. * * @return the value capacity */ @Override public int getValueCapacity() { return getValidityAndOffsetValueCapacity(); } private int getValidityAndSizeValueCapacity() { final int offsetValueCapacity = Math.max(getOffsetBufferValueCapacity(), 0); final int sizeValueCapacity = Math.max(getSizeBufferValueCapacity(), 0); return Math.min(offsetValueCapacity, sizeValueCapacity); } private int getValidityAndOffsetValueCapacity() { final int offsetValueCapacity = Math.max(getOffsetBufferValueCapacity(), 0); return Math.min(offsetValueCapacity, getValidityBufferValueCapacity()); } private int getValidityBufferValueCapacity() { return capAtMaxInt(validityBuffer.capacity() * 8); } /** * Set the element at the given index to null. * * @param index the value to change */ @Override public void setNull(int index) { while (index >= getValidityAndSizeValueCapacity()) { reallocValidityAndSizeAndOffsetBuffers(); } offsetBuffer.setInt(index * OFFSET_WIDTH, 0); sizeBuffer.setInt(index * SIZE_WIDTH, 0); BitVectorHelper.unsetBit(validityBuffer, index); } /** * Start new value in the ListView vector. * * @param index index of the value to start * @return offset of the new value */ @Override public int startNewValue(int index) { while (index >= getValidityAndSizeValueCapacity()) { reallocValidityAndSizeAndOffsetBuffers(); } if (index > 0) { final int prevOffset = getMaxViewEndChildVectorByIndex(index); offsetBuffer.setInt(index * OFFSET_WIDTH, prevOffset); } BitVectorHelper.setBit(validityBuffer, index); return offsetBuffer.getInt(index * OFFSET_WIDTH); } /** * Validate the invariants of the offset and size buffers. 0 <= offsets[i] <= length of the child * array 0 <= offsets[i] + size[i] <= length of the child array * * @param offset the offset at a given index * @param size the size at a given index */ private void validateInvariants(int offset, int size) { if (offset < 0) { throw new IllegalArgumentException("Offset cannot be negative"); } if (size < 0) { throw new IllegalArgumentException("Size cannot be negative"); } // 0 <= offsets[i] <= length of the child array if (offset > this.vector.getValueCount()) { throw new IllegalArgumentException("Offset is out of bounds."); } // 0 <= offsets[i] + size[i] <= length of the child array if (offset + size > this.vector.getValueCount()) { throw new IllegalArgumentException("Offset + size <= length of the child array."); } } /** * Set the offset at the given index. Make sure to use this function after updating `field` vector * and using `setValidity` * * @param index index of the value to set * @param value value to set */ public void setOffset(int index, int value) { validateInvariants(value, sizeBuffer.getInt(index * SIZE_WIDTH)); offsetBuffer.setInt(index * OFFSET_WIDTH, value); } /** * Set the size at the given index. Make sure to use this function after using `setOffset`. * * @param index index of the value to set * @param value value to set */ public void setSize(int index, int value) { validateInvariants(offsetBuffer.getInt(index * SIZE_WIDTH), value); sizeBuffer.setInt(index * SIZE_WIDTH, value); } /** * Set the validity at the given index. * * @param index index of the value to set * @param value value to set (0 for unset and 1 for a set) */ public void setValidity(int index, int value) { if (value == 0) { BitVectorHelper.unsetBit(validityBuffer, index); } else { BitVectorHelper.setBit(validityBuffer, index); } } /** * Sets the value count for the vector. * *

    Important note: The underlying vector does not support 64-bit allocations yet. This may * throw if attempting to hold larger than what a 32-bit vector can store. * * @param valueCount value count */ @Override public void setValueCount(int valueCount) { this.valueCount = valueCount; if (valueCount > 0) { while (valueCount > getValidityAndSizeValueCapacity()) { /* check if validity and offset buffers need to be re-allocated */ reallocValidityAndSizeAndOffsetBuffers(); } } /* valueCount for the data vector is the current end offset */ final long childValueCount = (valueCount == 0) ? 0 : getMaxViewEndChildVector(); /* set the value count of data vector and this will take care of * checking whether data buffer needs to be reallocated. * TODO: revisit when 64-bit vectors are supported */ Preconditions.checkArgument( childValueCount <= Integer.MAX_VALUE && childValueCount >= 0, "LargeListViewVector doesn't yet support 64-bit allocations: %s", childValueCount); vector.setValueCount((int) childValueCount); } @Override public AddOrGetResult addOrGetVector(FieldType fieldType) { AddOrGetResult result = super.addOrGetVector(fieldType); invalidateReader(); return result; } @Override public UnionVector promoteToUnion() { UnionVector vector = new UnionVector("$data$", allocator, /* field type*/ null, callBack); replaceDataVector(vector); invalidateReader(); if (callBack != null) { callBack.doWork(); } return vector; } private void invalidateReader() { reader = null; } @Deprecated @Override public List getFieldInnerVectors() { throw new UnsupportedOperationException("There are no inner vectors. Use getFieldBuffers"); } public UnionLargeListViewWriter getWriter() { return new UnionLargeListViewWriter(this); } @Override public int getValueCount() { return valueCount; } /** * Get the density of this LargeListViewVector. * * @return density */ public double getDensity() { if (valueCount == 0) { return 0.0D; } final double totalListSize = getMaxViewEndChildVector(); return totalListSize / valueCount; } /** Validating LargeListViewVector creation based on the specification guideline. */ @Override public void validate() { for (int i = 0; i < valueCount; i++) { final int offset = offsetBuffer.getInt((long) i * OFFSET_WIDTH); final int size = sizeBuffer.getInt((long) i * SIZE_WIDTH); validateInvariants(offset, size); } } /** * End the current value. * * @param index index of the value to end * @param size number of elements in the list that was written */ public void endValue(int index, int size) { sizeBuffer.setInt((long) index * SIZE_WIDTH, size); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java000066400000000000000000000737161515451615200321260ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static java.util.Collections.singletonList; import static org.apache.arrow.memory.util.LargeMemoryUtil.capAtMaxInt; import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import static org.apache.arrow.util.Preconditions.checkArgument; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.OutOfMemoryException; 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.hash.ArrowBufHasher; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.AddOrGetResult; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.BufferBacked; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueIterableVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.ZeroVector; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.impl.ComplexCopier; 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.FieldWriter; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; 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.CallBack; import org.apache.arrow.vector.util.JsonStringArrayList; import org.apache.arrow.vector.util.OversizedAllocationException; import org.apache.arrow.vector.util.TransferPair; /** * A list vector contains lists of a specific type of elements. Its structure contains 3 elements. * *

      *
    1. A validity buffer. *
    2. An offset buffer, that denotes lists boundaries. *
    3. A child data vector that contains the elements of lists. *
    * * The latter two are managed by its superclass. */ public class ListVector extends BaseRepeatedValueVector implements PromotableVector, ValueIterableVector> { public static ListVector empty(String name, BufferAllocator allocator) { return new ListVector(name, allocator, FieldType.nullable(ArrowType.List.INSTANCE), null); } protected UnionListReader reader; private CallBack callBack; protected Field field; protected int validityAllocationSizeInBytes; /** The maximum index that is actually set. */ protected int lastSet; /** * Constructs a new instance. * * @param name The name of the instance. * @param allocator The allocator to use for allocating/reallocating buffers. * @param fieldType The type of this list. * @param callBack A schema change callback. */ public ListVector( String name, BufferAllocator allocator, FieldType fieldType, CallBack callBack) { this(new Field(name, fieldType, null), allocator, callBack); } /** * Constructs a new instance. * * @param field The field materialized by this vector. * @param allocator The allocator to use for allocating/reallocating buffers. * @param callBack A schema change callback. */ public ListVector(Field field, BufferAllocator allocator, CallBack callBack) { super(field.getName(), allocator, callBack); this.validityBuffer = allocator.getEmpty(); this.field = field; this.callBack = callBack; this.validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION); this.lastSet = -1; } @Override public void initializeChildrenFromFields(List children) { checkArgument( children.size() == 1, "Lists have one child Field. Found: %s", children.isEmpty() ? "none" : children); Field field = children.get(0); AddOrGetResult addOrGetVector = addOrGetVector(field.getFieldType()); checkArgument( addOrGetVector.isCreated(), "Child vector already existed: %s", addOrGetVector.getVector()); addOrGetVector.getVector().initializeChildrenFromFields(field.getChildren()); this.field = new Field(this.field.getName(), this.field.getFieldType(), children); } @Override public void setInitialCapacity(int numRecords) { validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialCapacity(numRecords); } /** * Specialized version of setInitialCapacity() for ListVector. This is used by some callers when * they want to explicitly control and be conservative about memory allocated for inner data * vector. This is very useful when we are working with memory constraints for a query and have a * fixed amount of memory reserved for the record batch. In such cases, we are likely to face OOM * or related problems when we reserve memory for a record batch with value count x and do * setInitialCapacity(x) such that each vector allocates only what is necessary and not the * default amount but the multiplier forces the memory requirement to go beyond what was needed. * * @param numRecords value count * @param density density of ListVector. Density is the average size of list per position in the * List vector. For example, a density value of 10 implies each position in the list vector * has a list of 10 values. A density value of 0.1 implies out of 10 positions in the list * vector, 1 position has a list of size 1 and remaining positions are null (no lists) or * empty lists. This helps in tightly controlling the memory we provision for inner data * vector. */ @Override public void setInitialCapacity(int numRecords, double density) { validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialCapacity(numRecords, density); } /** * Specialized version of setInitialTotalCapacity() for ListVector. This is used by some callers * when they want to explicitly control and be conservative about memory allocated for inner data * vector. This is very useful when we are working with memory constraints for a query and have a * fixed amount of memory reserved for the record batch. In such cases, we are likely to face OOM * or related problems when we reserve memory for a record batch with value count x and do * setInitialCapacity(x) such that each vector allocates only what is necessary and not the * default amount but the multiplier forces the memory requirement to go beyond what was needed. * * @param numRecords value count * @param totalNumberOfElements the total number of elements to to allow for in this vector across * all records. */ @Override public void setInitialTotalCapacity(int numRecords, int totalNumberOfElements) { validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialTotalCapacity(numRecords, totalNumberOfElements); } /** * Get the density of this ListVector. * * @return density */ public double getDensity() { if (valueCount == 0) { return 0.0D; } final int startOffset = offsetBuffer.getInt(0); final int endOffset = offsetBuffer.getInt(valueCount * OFFSET_WIDTH); final double totalListSize = endOffset - startOffset; return totalListSize / valueCount; } @Override public List getChildrenFromFields() { return singletonList(getDataVector()); } /** * Load the buffers of this vector with provided source buffers. The caller manages the source * buffers and populates them before invoking this method. * * @param fieldNode the fieldNode indicating the value count * @param ownBuffers the buffers for this Field (own buffers only, children not included) */ @Override public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers) { if (ownBuffers.size() != 2) { throw new IllegalArgumentException( "Illegal buffer count, expected " + 2 + ", got: " + ownBuffers.size()); } ArrowBuf bitBuffer = ownBuffers.get(0); ArrowBuf offBuffer = ownBuffers.get(1); validityBuffer.getReferenceManager().release(); validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); offsetBuffer.getReferenceManager().release(); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); validityAllocationSizeInBytes = checkedCastToInt(validityBuffer.capacity()); offsetAllocationSizeInBytes = offsetBuffer.capacity(); lastSet = fieldNode.getLength() - 1; valueCount = fieldNode.getLength(); } /** * Get the buffers belonging to this vector. * * @return the inner buffers. */ @Override public List getFieldBuffers() { List result = new ArrayList<>(2); setReaderAndWriterIndex(); result.add(validityBuffer); result.add(offsetBuffer); return result; } /** * Export the buffers of the fields for C Data Interface. This method traverse the buffers and * export buffer and buffer's memory address into a list of buffers and a pointer to the list of * buffers. */ @Override public void exportCDataBuffers(List buffers, ArrowBuf buffersPtr, long nullValue) { exportBuffer(validityBuffer, buffers, buffersPtr, nullValue, true); if (offsetBuffer.capacity() == 0) { // Empty offset buffer is allowed for historical reason. // To export it through C Data interface, we need to allocate a buffer with one offset. // We set `retain = false` to explicitly not increase the ref count for the exported buffer. // The ref count of the newly created buffer (i.e., 1) already represents the usage // at imported side. exportBuffer(allocateOffsetBuffer(OFFSET_WIDTH), buffers, buffersPtr, nullValue, false); } else { exportBuffer(offsetBuffer, buffers, buffersPtr, nullValue, true); } } /** Set the reader and writer indexes for the inner buffers. */ private void setReaderAndWriterIndex() { validityBuffer.readerIndex(0); offsetBuffer.readerIndex(0); if (valueCount == 0) { validityBuffer.writerIndex(0); } else { 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); } /** * Get the inner vectors. * * @return the inner vectors for this field as defined by the TypeLayout * @deprecated This API will be removed as the current implementations no longer support inner * vectors. */ @Deprecated @Override public List getFieldInnerVectors() { throw new UnsupportedOperationException("There are no inner vectors. Use getFieldBuffers"); } public UnionListWriter getWriter() { return new UnionListWriter(this); } /** Same as {@link #allocateNewSafe()}. */ @Override public void allocateNew() throws OutOfMemoryException { if (!allocateNewSafe()) { throw new OutOfMemoryException("Failure while allocating memory"); } } /** * Allocate memory for the vector. We internally use a default value count of 4096 to allocate * memory for at least these many elements in the vector. * * @return false if memory allocation fails, true otherwise. */ @Override public boolean allocateNewSafe() { boolean success = false; try { /* we are doing a new allocation -- release the current buffers */ clear(); /* allocate validity buffer */ allocateValidityBuffer(validityAllocationSizeInBytes); /* allocate offset and data buffer */ success = super.allocateNewSafe(); } finally { if (!success) { clear(); } } return success; } @Override protected void allocateValidityBuffer(final long size) { super.allocateValidityBuffer(size); validityAllocationSizeInBytes = (int) size; } /** * Resize the vector to increase the capacity. The internal behavior is to double the current * value capacity. */ @Override public void reAlloc() { /* reallocate the validity buffer */ reallocValidityBuffer(); /* reallocate the offset and data */ super.reAlloc(); } protected void reallocValidityAndOffsetBuffers() { reallocOffsetBuffer(); reallocValidityBuffer(); } private void reallocValidityBuffer() { final int currentBufferCapacity = checkedCastToInt(validityBuffer.capacity()); long newAllocationSize = getNewAllocationSize(currentBufferCapacity); final ArrowBuf newBuf = allocator.buffer(newAllocationSize); newBuf.setBytes(0, validityBuffer, 0, currentBufferCapacity); newBuf.setZero(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); validityBuffer.getReferenceManager().release(1); validityBuffer = newBuf; validityAllocationSizeInBytes = (int) newAllocationSize; } private long getNewAllocationSize(int currentBufferCapacity) { long newAllocationSize = currentBufferCapacity * 2L; if (newAllocationSize == 0) { if (validityAllocationSizeInBytes > 0) { newAllocationSize = validityAllocationSizeInBytes; } else { newAllocationSize = BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); assert newAllocationSize >= 1; if (newAllocationSize > MAX_ALLOCATION_SIZE) { throw new OversizedAllocationException("Unable to expand the buffer"); } return newAllocationSize; } /** * Same as {@link #copyFrom(int, int, ValueVector)} except that it handles the case when the * capacity of the vector needs to be expanded before copy. * * @param inIndex position to copy from in source vector * @param outIndex position to copy to in this vector * @param from source vector */ @Override public void copyFromSafe(int inIndex, int outIndex, ValueVector from) { copyFrom(inIndex, outIndex, from); } /** * Copy a cell value from a particular index in source vector to a particular position in this * vector. * * @param inIndex position to copy from in source vector * @param outIndex position to copy to in this vector * @param from source vector */ @Override public void copyFrom(int inIndex, int outIndex, ValueVector from) { Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); FieldReader in = from.getReader(); in.setPosition(inIndex); FieldWriter out = getWriter(); out.setPosition(outIndex); ComplexCopier.copy(in, out); } /** * Get the inner data vector for this list vector. * * @return data vector */ @Override public FieldVector getDataVector() { return vector; } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return getTransferPair(ref, allocator, null); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return getTransferPair(field, allocator, null); } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return new TransferImpl(ref, allocator, callBack); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { return new TransferImpl(field, allocator, callBack); } @Override public TransferPair makeTransferPair(ValueVector target) { return new TransferImpl((ListVector) target); } @Override public long getValidityBufferAddress() { return validityBuffer.memoryAddress(); } @Override public long getDataBufferAddress() { throw new UnsupportedOperationException(); } @Override public long getOffsetBufferAddress() { return offsetBuffer.memoryAddress(); } @Override public ArrowBuf getValidityBuffer() { return validityBuffer; } @Override public ArrowBuf getDataBuffer() { throw new UnsupportedOperationException(); } @Override public ArrowBuf getOffsetBuffer() { return offsetBuffer; } @Override public int hashCode(int index) { return hashCode(index, null); } @Override public int hashCode(int index, ArrowBufHasher hasher) { if (isSet(index) == 0) { return ArrowBufPointer.NULL_HASH_CODE; } int hash = 0; final int start = offsetBuffer.getInt(index * OFFSET_WIDTH); final int end = offsetBuffer.getInt((index + 1) * OFFSET_WIDTH); for (int i = start; i < end; i++) { hash = ByteFunctionHelpers.combineHash(hash, vector.hashCode(i, hasher)); } return hash; } @Override public OUT accept(VectorVisitor visitor, IN value) { return visitor.visit(this, value); } private class TransferImpl implements TransferPair { ListVector to; TransferPair dataTransferPair; public TransferImpl(String name, BufferAllocator allocator, CallBack callBack) { this(new ListVector(name, allocator, field.getFieldType(), callBack)); } public TransferImpl(Field field, BufferAllocator allocator, CallBack callBack) { this(new ListVector(field, allocator, callBack)); } public TransferImpl(ListVector to) { this.to = to; to.addOrGetVector(vector.getField().getFieldType()); if (to.getDataVector() instanceof ZeroVector) { to.addOrGetVector(vector.getField().getFieldType()); } dataTransferPair = getDataVector().makeTransferPair(to.getDataVector()); } /** * Transfer this vector'data to another vector. The memory associated with this vector is * transferred to the allocator of target vector for accounting and management purposes. */ @Override public void transfer() { to.clear(); dataTransferPair.transfer(); to.validityBuffer = transferBuffer(validityBuffer, to.allocator); to.offsetBuffer = transferBuffer(offsetBuffer, to.allocator); to.lastSet = lastSet; if (valueCount > 0) { to.setValueCount(valueCount); } clear(); } /** * Slice this vector at desired index and length and transfer the corresponding data to the * target vector. * * @param startIndex start position of the split in source vector. * @param length length of the split. */ @Override public void splitAndTransfer(int startIndex, int length) { Preconditions.checkArgument( startIndex >= 0 && length >= 0 && startIndex + length <= valueCount, "Invalid parameters startIndex: %s, length: %s for valueCount: %s", startIndex, length, valueCount); to.clear(); if (length > 0) { final int startPoint = offsetBuffer.getInt(startIndex * OFFSET_WIDTH); final int sliceLength = offsetBuffer.getInt((startIndex + length) * OFFSET_WIDTH) - startPoint; to.offsetBuffer = to.allocateOffsetBuffer((length + 1) * OFFSET_WIDTH); /* splitAndTransfer offset buffer */ for (int i = 0; i < length + 1; i++) { final int relativeOffset = offsetBuffer.getInt((startIndex + i) * OFFSET_WIDTH) - startPoint; to.offsetBuffer.setInt(i * OFFSET_WIDTH, relativeOffset); } /* splitAndTransfer validity buffer */ splitAndTransferValidityBuffer(startIndex, length, to); /* splitAndTransfer data buffer */ dataTransferPair.splitAndTransfer(startPoint, sliceLength); to.lastSet = length - 1; to.setValueCount(length); } } @Override public ValueVector getTo() { return to; } @Override public void copyValueSafe(int from, int to) { this.to.copyFrom(from, to, ListVector.this); } } @Override protected FieldReader getReaderImpl() { return new UnionListReader(this); } @Override public UnionListReader getReader() { reader = (UnionListReader) super.getReader(); return reader; } /** Initialize the child data vector to field type. */ @Override public AddOrGetResult addOrGetVector(FieldType fieldType) { AddOrGetResult result = super.addOrGetVector(fieldType); invalidateReader(); return result; } /** * Get the size (number of bytes) of underlying buffers used by this vector. * * @return size of underlying buffers. */ @Override public int getBufferSize() { if (valueCount == 0) { return 0; } final int offsetBufferSize = (valueCount + 1) * OFFSET_WIDTH; final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); return offsetBufferSize + validityBufferSize + vector.getBufferSize(); } @Override public int getBufferSizeFor(int valueCount) { if (valueCount == 0) { return 0; } final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); return super.getBufferSizeFor(valueCount) + validityBufferSize; } @Override public Field getField() { if (field.getChildren().contains(getDataVector().getField())) { return field; } field = new Field( field.getName(), field.getFieldType(), Collections.singletonList(getDataVector().getField())); return field; } @Override public MinorType getMinorType() { return MinorType.LIST; } @Override public void clear() { super.clear(); validityBuffer = releaseBuffer(validityBuffer); lastSet = -1; } @Override public void reset() { super.reset(); validityBuffer.setZero(0, validityBuffer.capacity()); lastSet = -1; } /** * Return the underlying buffers associated with this vector. Note that this doesn't impact the * reference counts for this buffer, so it only should be used for in-context access. Also note * that this buffer changes regularly, thus external classes shouldn't hold a reference to it * (unless they change it). * * @param clear Whether to clear vector before returning, the buffers will still be refcounted but * the returned array will be the only reference to them. Also, this won't clear the child * buffers. * @return The underlying {@link ArrowBuf buffers} that is used by this vector instance. */ @Override public ArrowBuf[] getBuffers(boolean clear) { setReaderAndWriterIndex(); final ArrowBuf[] buffers; if (getBufferSize() == 0) { buffers = new ArrowBuf[0]; } else { List list = new ArrayList<>(); list.add(offsetBuffer); list.add(validityBuffer); list.addAll(Arrays.asList(vector.getBuffers(false))); buffers = list.toArray(new ArrowBuf[list.size()]); } if (clear) { for (ArrowBuf buffer : buffers) { buffer.getReferenceManager().retain(); } clear(); } return buffers; } @Override public UnionVector promoteToUnion() { UnionVector vector = new UnionVector("$data$", allocator, /* field type*/ null, callBack); replaceDataVector(vector); invalidateReader(); if (callBack != null) { callBack.doWork(); } return vector; } protected void invalidateReader() { fieldReader = null; reader = null; } /** * Get the element in the list vector at a particular index. * * @param index position of the element * @return Object at given position */ @Override public List getObject(int index) { if (isSet(index) == 0) { return null; } 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)); } return vals; } /** * Check if element at given index is null. * * @param index position of element * @return true if element at given index is null, false otherwise */ @Override public boolean isNull(int index) { return (isSet(index) == 0); } /** * Check if element at given index is empty list. * * @param index position of element * @return true if element at given index is empty list or NULL, false otherwise */ @Override public boolean isEmpty(int index) { if (isNull(index)) { return true; } else { final int start = offsetBuffer.getInt(index * OFFSET_WIDTH); final int end = offsetBuffer.getInt((index + 1) * OFFSET_WIDTH); return start == end; } } /** * Same as {@link #isNull(int)}. * * @param index position of element * @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; } /** * Get the number of elements that are null in the vector. * * @return the number of null elements. */ @Override public int getNullCount() { return BitVectorHelper.getNullCount(validityBuffer, valueCount); } /** * Get the current value capacity for the vector. * * @return number of elements that vector can hold. */ @Override public int getValueCapacity() { return getValidityAndOffsetValueCapacity(); } private int getValidityAndOffsetValueCapacity() { final int offsetValueCapacity = Math.max(getOffsetBufferValueCapacity() - 1, 0); return Math.min(offsetValueCapacity, getValidityBufferValueCapacity()); } private int getValidityBufferValueCapacity() { return capAtMaxInt(validityBuffer.capacity() * 8); } /** * Sets the list at index to be not-null. Reallocates validity buffer if index is larger than * current capacity. */ public void setNotNull(int index) { while (index >= getValidityAndOffsetValueCapacity()) { reallocValidityAndOffsetBuffers(); } BitVectorHelper.setBit(validityBuffer, index); lastSet = index; } /** * Sets list at index to be null. * * @param index position in vector */ @Override public void setNull(int index) { while (index >= getValidityAndOffsetValueCapacity()) { reallocValidityAndOffsetBuffers(); } if (lastSet >= index) { lastSet = index - 1; } for (int i = lastSet + 1; i <= index; i++) { final int currentOffset = offsetBuffer.getInt(i * OFFSET_WIDTH); offsetBuffer.setInt((i + 1) * OFFSET_WIDTH, currentOffset); } BitVectorHelper.unsetBit(validityBuffer, index); lastSet = index; } /** * Start a new value in the list vector. * * @param index index of the value to start */ @Override public int startNewValue(int index) { while (index >= getValidityAndOffsetValueCapacity()) { reallocValidityAndOffsetBuffers(); } if (lastSet >= index) { lastSet = index - 1; } for (int i = lastSet + 1; i <= index; i++) { final int currentOffset = offsetBuffer.getInt(i * OFFSET_WIDTH); offsetBuffer.setInt((i + 1) * OFFSET_WIDTH, currentOffset); } BitVectorHelper.setBit(validityBuffer, index); lastSet = index; return offsetBuffer.getInt((lastSet + 1) * OFFSET_WIDTH); } /** * End the current value. * * @param index index of the value to end * @param size number of elements in the list that was written */ public void endValue(int index, int size) { final int currentOffset = offsetBuffer.getInt((index + 1) * OFFSET_WIDTH); offsetBuffer.setInt((index + 1) * OFFSET_WIDTH, currentOffset + size); } /** * Sets the value count for the vector. * * @param valueCount value count */ @Override public void setValueCount(int valueCount) { this.valueCount = valueCount; if (valueCount > 0) { while (valueCount > getValidityAndOffsetValueCapacity()) { /* check if validity and offset buffers need to be re-allocated */ reallocValidityAndOffsetBuffers(); } for (int i = lastSet + 1; i < valueCount; i++) { /* fill the holes with offsets */ final int currentOffset = offsetBuffer.getInt(i * OFFSET_WIDTH); offsetBuffer.setInt((i + 1) * OFFSET_WIDTH, currentOffset); } } /* valueCount for the data vector is the current end offset */ final int childValueCount = (valueCount == 0) ? 0 : offsetBuffer.getInt((lastSet + 1) * OFFSET_WIDTH); /* set the value count of data vector and this will take care of * checking whether data buffer needs to be reallocated. */ vector.setValueCount(childValueCount); } public void setLastSet(int value) { lastSet = value; } public int getLastSet() { return lastSet; } @Override public int getElementStartIndex(int index) { return offsetBuffer.getInt(index * OFFSET_WIDTH); } @Override public int getElementEndIndex(int index) { return offsetBuffer.getInt((index + 1) * OFFSET_WIDTH); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java000066400000000000000000000763271515451615200327620ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static java.util.Collections.singletonList; import static org.apache.arrow.memory.util.LargeMemoryUtil.capAtMaxInt; import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import static org.apache.arrow.util.Preconditions.checkArgument; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.OutOfMemoryException; 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.hash.ArrowBufHasher; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.AddOrGetResult; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.BufferBacked; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueIterableVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.ZeroVector; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.impl.ComplexCopier; import org.apache.arrow.vector.complex.impl.UnionListViewReader; import org.apache.arrow.vector.complex.impl.UnionListViewWriter; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.complex.writer.FieldWriter; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; 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.CallBack; import org.apache.arrow.vector.util.JsonStringArrayList; import org.apache.arrow.vector.util.OversizedAllocationException; import org.apache.arrow.vector.util.TransferPair; /** * A list view vector contains lists of a specific type of elements. Its structure contains four * elements. * *
      *
    1. A validity buffer. *
    2. An offset buffer, that denotes lists starts. *
    3. A size buffer, that denotes lists ends. *
    4. A child data vector that contains the elements of lists. *
    * * The latter three are managed by its superclass. */ /* * TODO: consider merging the functionality in `BaseRepeatedValueVector` into this class. */ public class ListViewVector extends BaseRepeatedValueViewVector implements PromotableVector, ValueIterableVector> { protected UnionListViewReader reader; private CallBack callBack; protected Field field; protected int validityAllocationSizeInBytes; public static ListViewVector empty(String name, BufferAllocator allocator) { return new ListViewVector( name, allocator, FieldType.nullable(ArrowType.ListView.INSTANCE), null); } /** * Constructs a new instance. * * @param name The name of the instance. * @param allocator The allocator to use for allocating/reallocating buffers. * @param fieldType The type of this list. * @param callBack A schema change callback. */ public ListViewVector( String name, BufferAllocator allocator, FieldType fieldType, CallBack callBack) { this(new Field(name, fieldType, null), allocator, callBack); } /** * Constructs a new instance. * * @param field The field materialized by this vector. * @param allocator The allocator to use for allocating/reallocating buffers. * @param callBack A schema change callback. */ public ListViewVector(Field field, BufferAllocator allocator, CallBack callBack) { super(field.getName(), allocator, callBack); this.validityBuffer = allocator.getEmpty(); this.field = field; this.callBack = callBack; this.validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION); } @Override public void initializeChildrenFromFields(List children) { checkArgument( children.size() == 1, "ListViews have one child Field. Found: %s", children.isEmpty() ? "none" : children); Field field = children.get(0); AddOrGetResult addOrGetVector = addOrGetVector(field.getFieldType()); checkArgument( addOrGetVector.isCreated(), "Child vector already existed: %s", addOrGetVector.getVector()); addOrGetVector.getVector().initializeChildrenFromFields(field.getChildren()); this.field = new Field(this.field.getName(), this.field.getFieldType(), children); } @Override public void setInitialCapacity(int numRecords) { validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialCapacity(numRecords); } /** * Specialized version of setInitialCapacity() for ListViewVector. This is used by some callers * when they want to explicitly control and be conservative about memory allocated for inner data * vector. This is very useful when we are working with memory constraints for a query and have a * fixed amount of memory reserved for the record batch. In such cases, we are likely to face OOM * or related problems when we reserve memory for a record batch with value count x and do * setInitialCapacity(x) such that each vector allocates only what is necessary and not the * default amount, but the multiplier forces the memory requirement to go beyond what was needed. * * @param numRecords value count * @param density density of ListViewVector. Density is the average size of a list per position in * the ListViewVector. For example, a density value of 10 implies each position in the list * vector has a list of 10 values. A density value of 0.1 implies out of 10 positions in the * list vector, 1 position has a list of size 1, and the remaining positions are null (no * lists) or empty lists. This helps in tightly controlling the memory we provision for inner * data vector. */ @Override public void setInitialCapacity(int numRecords, double density) { validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialCapacity(numRecords, density); } /** * Specialized version of setInitialTotalCapacity() for ListViewVector. This is used by some * callers when they want to explicitly control and be conservative about memory allocated for * inner data vector. This is very useful when we are working with memory constraints for a query * and have a fixed amount of memory reserved for the record batch. In such cases, we are likely * to face OOM or related problems when we reserve memory for a record batch with value count x * and do setInitialCapacity(x) such that each vector allocates only what is necessary and not the * default amount, but the multiplier forces the memory requirement to go beyond what was needed. * * @param numRecords value count * @param totalNumberOfElements the total number of elements to allow for in this vector across * all records. */ @Override public void setInitialTotalCapacity(int numRecords, int totalNumberOfElements) { validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialTotalCapacity(numRecords, totalNumberOfElements); } @Override public List getChildrenFromFields() { return singletonList(getDataVector()); } /** * Load the buffers associated with this Field. * * @param fieldNode the fieldNode * @param ownBuffers the buffers for this Field (own buffers only, children not included) */ @Override public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers) { if (ownBuffers.size() != 3) { throw new IllegalArgumentException( "Illegal buffer count, expected " + 3 + ", got: " + ownBuffers.size()); } ArrowBuf bitBuffer = ownBuffers.get(0); ArrowBuf offBuffer = ownBuffers.get(1); ArrowBuf szBuffer = ownBuffers.get(2); validityBuffer.getReferenceManager().release(); validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); offsetBuffer.getReferenceManager().release(); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); sizeBuffer.getReferenceManager().release(); sizeBuffer = szBuffer.getReferenceManager().retain(szBuffer, allocator); validityAllocationSizeInBytes = checkedCastToInt(validityBuffer.capacity()); offsetAllocationSizeInBytes = offsetBuffer.capacity(); sizeAllocationSizeInBytes = sizeBuffer.capacity(); valueCount = fieldNode.getLength(); } /** Set the reader and writer indexes for the inner buffers. */ private void setReaderAndWriterIndex() { validityBuffer.readerIndex(0); offsetBuffer.readerIndex(0); sizeBuffer.readerIndex(0); if (valueCount == 0) { validityBuffer.writerIndex(0); offsetBuffer.writerIndex(0); sizeBuffer.writerIndex(0); } else { validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); offsetBuffer.writerIndex((long) valueCount * OFFSET_WIDTH); sizeBuffer.writerIndex((long) valueCount * SIZE_WIDTH); } } @Override public List getFieldBuffers() { List result = new ArrayList<>(2); setReaderAndWriterIndex(); result.add(validityBuffer); result.add(offsetBuffer); result.add(sizeBuffer); return result; } /** * Export the buffers of the fields for C Data Interface. This method traverses the buffers and * export buffer and buffer's memory address into a list of buffers and a pointer to the list of * buffers. */ @Override public void exportCDataBuffers(List buffers, ArrowBuf buffersPtr, long nullValue) { exportBuffer(validityBuffer, buffers, buffersPtr, nullValue, true); exportBuffer(offsetBuffer, buffers, buffersPtr, nullValue, true); exportBuffer(sizeBuffer, buffers, buffersPtr, nullValue, true); } @Override public void allocateNew() throws OutOfMemoryException { if (!allocateNewSafe()) { throw new OutOfMemoryException("Failure while allocating memory"); } } @Override public boolean allocateNewSafe() { boolean success = false; try { /* release the current buffers, hence this is a new allocation * Note that, the `clear` method call below is releasing validityBuffer * calling the superclass clear method which is releasing the associated buffers * (sizeBuffer and offsetBuffer). */ clear(); /* allocate validity buffer */ allocateValidityBuffer(validityAllocationSizeInBytes); /* allocate offset, data and sizes buffer */ success = super.allocateNewSafe(); } finally { if (!success) { clear(); } } return success; } @Override protected void allocateValidityBuffer(final long size) { super.allocateValidityBuffer(size); validityAllocationSizeInBytes = (int) size; } @Override public void reAlloc() { /* reallocate the validity buffer */ reallocValidityBuffer(); /* reallocate the offset, size, and data */ super.reAlloc(); } protected void reallocValidityAndSizeAndOffsetBuffers() { reallocateBuffers(); reallocValidityBuffer(); } private void reallocValidityBuffer() { final int currentBufferCapacity = checkedCastToInt(validityBuffer.capacity()); long newAllocationSize = getNewAllocationSize(currentBufferCapacity); final ArrowBuf newBuf = allocator.buffer(newAllocationSize); newBuf.setBytes(0, validityBuffer, 0, currentBufferCapacity); newBuf.setZero(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); validityBuffer.getReferenceManager().release(1); validityBuffer = newBuf; validityAllocationSizeInBytes = (int) newAllocationSize; } private long getNewAllocationSize(int currentBufferCapacity) { long newAllocationSize = currentBufferCapacity * 2L; if (newAllocationSize == 0) { if (validityAllocationSizeInBytes > 0) { newAllocationSize = validityAllocationSizeInBytes; } else { newAllocationSize = BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); assert newAllocationSize >= 1; if (newAllocationSize > MAX_ALLOCATION_SIZE) { throw new OversizedAllocationException("Unable to expand the buffer"); } return newAllocationSize; } @Override public void copyFromSafe(int inIndex, int outIndex, ValueVector from) { copyFrom(inIndex, outIndex, from); } @Override public OUT accept(VectorVisitor visitor, IN value) { return visitor.visit(this, value); } @Override public void copyFrom(int inIndex, int outIndex, ValueVector from) { Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); FieldReader in = from.getReader(); in.setPosition(inIndex); FieldWriter out = getWriter(); out.setPosition(outIndex); ComplexCopier.copy(in, out); } @Override public FieldVector getDataVector() { return vector; } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return getTransferPair(ref, allocator, null); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return getTransferPair(field, allocator, null); } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return new TransferImpl(ref, allocator, callBack); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { return new TransferImpl(field, allocator, callBack); } @Override public TransferPair makeTransferPair(ValueVector target) { return new TransferImpl((ListViewVector) target); } @Override public long getValidityBufferAddress() { return validityBuffer.memoryAddress(); } @Override public long getDataBufferAddress() { throw new UnsupportedOperationException(); } @Override public long getOffsetBufferAddress() { return offsetBuffer.memoryAddress(); } @Override public ArrowBuf getValidityBuffer() { return validityBuffer; } @Override public ArrowBuf getDataBuffer() { throw new UnsupportedOperationException(); } @Override public ArrowBuf getOffsetBuffer() { return offsetBuffer; } public ArrowBuf getSizeBuffer() { return sizeBuffer; } public long getSizeBufferAddress() { return sizeBuffer.memoryAddress(); } /** * Get the hash code for the element at the given index. * * @param index position of the element * @return hash code for the element at the given index */ @Override public int hashCode(int index) { return hashCode(index, null); } /** * Get the hash code for the element at the given index. * * @param index position of the element * @param hasher hasher to use * @return hash code for the element at the given index */ @Override public int hashCode(int index, ArrowBufHasher hasher) { if (isSet(index) == 0) { return ArrowBufPointer.NULL_HASH_CODE; } int hash = 0; 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; TransferPair dataTransferPair; public TransferImpl(String name, BufferAllocator allocator, CallBack callBack) { this(new ListViewVector(name, allocator, field.getFieldType(), callBack)); } public TransferImpl(Field field, BufferAllocator allocator, CallBack callBack) { this(new ListViewVector(field, allocator, callBack)); } public TransferImpl(ListViewVector to) { this.to = to; to.addOrGetVector(vector.getField().getFieldType()); if (to.getDataVector() instanceof ZeroVector) { to.addOrGetVector(vector.getField().getFieldType()); } dataTransferPair = getDataVector().makeTransferPair(to.getDataVector()); } @Override public void transfer() { to.clear(); dataTransferPair.transfer(); to.validityBuffer = transferBuffer(validityBuffer, to.allocator); to.offsetBuffer = transferBuffer(offsetBuffer, to.allocator); to.sizeBuffer = transferBuffer(sizeBuffer, to.allocator); if (valueCount > 0) { to.setValueCount(valueCount); } clear(); } @Override public void splitAndTransfer(int startIndex, int length) { Preconditions.checkArgument( startIndex >= 0 && length >= 0 && startIndex + length <= valueCount, "Invalid parameters startIndex: %s, length: %s for valueCount: %s", startIndex, length, valueCount); to.clear(); if (length > 0) { // 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); /* splitAndTransfer the size buffer */ int maxOffsetAndSizeSum = -1; int minOffsetValue = -1; for (int i = 0; i < length; i++) { final int offsetValue = getElementStartIndex(startIndex + i); final int sizeValue = getElementSize(startIndex + i); to.setElementSizeBuffer(i, sizeValue); if (maxOffsetAndSizeSum < offsetValue + sizeValue) { maxOffsetAndSizeSum = offsetValue + sizeValue; } if (minOffsetValue == -1 || minOffsetValue > offsetValue) { minOffsetValue = offsetValue; } } /* splitAndTransfer the offset buffer */ for (int i = 0; i < length; i++) { final int offsetValue = getElementStartIndex(startIndex + i); final int relativeOffset = offsetValue - minOffsetValue; to.setElementOffsetBuffer(i, relativeOffset); } /* splitAndTransfer the validity buffer */ splitAndTransferValidityBuffer(startIndex, length, to); /* splitAndTransfer the data buffer */ final int childSliceLength = maxOffsetAndSizeSum - minOffsetValue; dataTransferPair.splitAndTransfer(minOffsetValue, childSliceLength); to.setValueCount(length); } } @Override public ValueVector getTo() { return to; } @Override public void copyValueSafe(int from, int to) { this.to.copyFrom(from, to, ListViewVector.this); } } @Override protected FieldReader getReaderImpl() { return new UnionListViewReader(this); } @Override public UnionListViewReader getReader() { reader = (UnionListViewReader) super.getReader(); return reader; } /** * Get the size (number of bytes) of underlying buffers used by this vector. * * @return size of underlying buffers. */ @Override public int getBufferSize() { if (valueCount == 0) { return 0; } final int offsetBufferSize = valueCount * OFFSET_WIDTH; final int sizeBufferSize = valueCount * SIZE_WIDTH; final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); return offsetBufferSize + sizeBufferSize + validityBufferSize + vector.getBufferSize(); } /** * Get the size (number of bytes) of underlying buffers used by this. * * @param valueCount the number of values to assume this vector contains * @return size of underlying buffers. */ @Override public int getBufferSizeFor(int valueCount) { if (valueCount == 0) { return 0; } final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); return super.getBufferSizeFor(valueCount) + validityBufferSize; } /** * Get the field associated with the list view vector. * * @return the field */ @Override public Field getField() { if (field.getChildren().contains(getDataVector().getField())) { return field; } field = new Field( field.getName(), field.getFieldType(), Collections.singletonList(getDataVector().getField())); return field; } /** * Get the minor type for the vector. * * @return the minor type */ @Override public MinorType getMinorType() { return MinorType.LISTVIEW; } /** Clear the vector data. */ @Override public void clear() { // calling superclass clear method which is releasing the sizeBufer and offsetBuffer super.clear(); validityBuffer = releaseBuffer(validityBuffer); } /** Release the buffers associated with this vector. */ @Override public void reset() { super.reset(); validityBuffer.setZero(0, validityBuffer.capacity()); } /** * Return the underlying buffers associated with this vector. Note that this doesn't impact the * reference counts for this buffer, so it only should be used for in-context access. Also note * that this buffer changes regularly, thus external classes shouldn't hold a reference to it * (unless they change it). * * @param clear Whether to clear vector before returning, the buffers will still be refcounted but * the returned array will be the only reference to them. Also, this won't clear the child * buffers. * @return The underlying {@link ArrowBuf buffers} that is used by this vector instance. */ @Override public ArrowBuf[] getBuffers(boolean clear) { setReaderAndWriterIndex(); final ArrowBuf[] buffers; if (getBufferSize() == 0) { buffers = new ArrowBuf[0]; } else { List list = new ArrayList<>(); // the order must be validity, offset and size buffers list.add(validityBuffer); list.add(offsetBuffer); list.add(sizeBuffer); list.addAll(Arrays.asList(vector.getBuffers(false))); buffers = list.toArray(new ArrowBuf[list.size()]); } if (clear) { for (ArrowBuf buffer : buffers) { buffer.getReferenceManager().retain(); } clear(); } return buffers; } /** * Get the element in the list view vector at a particular index. * * @param index position of the element * @return Object at given position */ @Override public List getObject(int index) { if (isSet(index) == 0) { return null; } 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)); } return vals; } /** * Check if an element at given index is null. * * @param index position of an element * @return true if an element at given index is null, false otherwise */ @Override public boolean isNull(int index) { return (isSet(index) == 0); } /** * Check if an element at given index is an empty list. * * @param index position of an element * @return true if an element at given index is an empty list or NULL, false otherwise */ @Override public boolean isEmpty(int index) { if (isNull(index)) { return true; } else { return getElementSize(index) == 0; } } /** * Same as {@link #isNull(int)}. * * @param index position of the element * @return 1 if element at given index is not null, 0 otherwise */ public int isSet(int index) { return BitVectorHelper.get(validityBuffer, index); } /** * Get the number of elements that are null in the vector. * * @return the number of null elements. */ @Override public int getNullCount() { return BitVectorHelper.getNullCount(validityBuffer, valueCount); } /** * Get the value capacity by considering validity and offset capacity. Note that the size buffer * capacity is not considered here since it has the same capacity as the offset buffer. * * @return the value capacity */ @Override public int getValueCapacity() { return getValidityAndOffsetValueCapacity(); } private int getValidityAndSizeValueCapacity() { final int offsetValueCapacity = Math.max(getOffsetBufferValueCapacity(), 0); final int sizeValueCapacity = Math.max(getSizeBufferValueCapacity(), 0); return Math.min(offsetValueCapacity, sizeValueCapacity); } private int getValidityAndOffsetValueCapacity() { final int offsetValueCapacity = Math.max(getOffsetBufferValueCapacity(), 0); return Math.min(offsetValueCapacity, getValidityBufferValueCapacity()); } private int getValidityBufferValueCapacity() { return capAtMaxInt(validityBuffer.capacity() * 8); } /** * Set the element at the given index to null. * * @param index the value to change */ @Override public void setNull(int index) { while (index >= getValidityAndSizeValueCapacity()) { reallocValidityAndSizeAndOffsetBuffers(); } setElementOffsetBuffer(index, 0); setElementSizeBuffer(index, 0); BitVectorHelper.unsetBit(validityBuffer, index); } /** * Start new value in the ListView vector. * * @param index index of the value to start * @return offset of the new value */ @Override public int startNewValue(int index) { while (index >= getValidityAndSizeValueCapacity()) { reallocValidityAndSizeAndOffsetBuffers(); } if (index > 0) { final int prevOffset = getMaxViewEndChildVectorByIndex(index); setElementOffsetBuffer(index, prevOffset); } BitVectorHelper.setBit(validityBuffer, index); return getElementStartIndex(index); } /** * Validate the invariants of the offset and size buffers. 0 <= offsets[i] <= length of the child * array 0 <= offsets[i] + size[i] <= length of the child array * * @param offset the offset at a given index * @param size the size at a given index */ private void validateInvariants(int offset, int size) { if (offset < 0) { throw new IllegalArgumentException("Offset cannot be negative"); } if (size < 0) { throw new IllegalArgumentException("Size cannot be negative"); } // 0 <= offsets[i] <= length of the child array if (offset > this.vector.getValueCount()) { throw new IllegalArgumentException("Offset is out of bounds."); } // 0 <= offsets[i] + size[i] <= length of the child array if (offset + size > this.vector.getValueCount()) { throw new IllegalArgumentException("Offset + size <= length of the child array."); } } /** * Set the offset at the given index. Make sure to use this function after updating `field` vector * and using `setValidity` * * @param index index of the value to set * @param value value to set */ public void setOffset(int index, int value) { validateInvariants(value, getElementSize(index)); setElementOffsetBuffer(index, value); } /** * Set the size at the given index. Make sure to use this function after using `setOffset`. * * @param index index of the value to set * @param value value to set */ public void setSize(int index, int value) { validateInvariants(getElementStartIndex(index), value); setElementSizeBuffer(index, value); } /** * Set the validity at the given index. * * @param index index of the value to set * @param value value to set (0 for unset and 1 for a set) */ public void setValidity(int index, int value) { if (value == 0) { BitVectorHelper.unsetBit(validityBuffer, index); } else { BitVectorHelper.setBit(validityBuffer, index); } } @Override public void setValueCount(int valueCount) { this.valueCount = valueCount; if (valueCount > 0) { while (valueCount > getValidityAndSizeValueCapacity()) { /* check if validity and offset buffers need to be re-allocated */ reallocValidityAndSizeAndOffsetBuffers(); } } /* valueCount for the data vector is the current end offset */ final int childValueCount = (valueCount == 0) ? 0 : getMaxViewEndChildVector(); /* set the value count of data vector and this will take care of * checking whether data buffer needs to be reallocated. */ vector.setValueCount(childValueCount); } @Override public int getElementStartIndex(int index) { 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 getElementStartIndex(index) + getElementSize(index); } @Override public AddOrGetResult addOrGetVector(FieldType fieldType) { AddOrGetResult result = super.addOrGetVector(fieldType); invalidateReader(); return result; } @Override public UnionVector promoteToUnion() { UnionVector vector = new UnionVector("$data$", allocator, /* field type*/ null, callBack); replaceDataVector(vector); invalidateReader(); if (callBack != null) { callBack.doWork(); } return vector; } private void invalidateReader() { reader = null; } @Deprecated @Override public List getFieldInnerVectors() { throw new UnsupportedOperationException("There are no inner vectors. Use getFieldBuffers"); } public UnionListViewWriter getWriter() { return new UnionListViewWriter(this); } @Override public int getValueCount() { return valueCount; } /** * Get the density of this ListVector. * * @return density */ public double getDensity() { if (valueCount == 0) { return 0.0D; } final double totalListSize = getMaxViewEndChildVector(); return totalListSize / valueCount; } /** Validating ListViewVector creation based on the specification guideline. */ @Override public void validate() { for (int i = 0; i < valueCount; i++) { final int offset = getElementStartIndex(i); final int size = getElementSize(i); validateInvariants(offset, size); } } /** * End the current value. * * @param index index of the value to end * @param size number of elements in the list that was written */ public void endValue(int index, int size) { setElementSizeBuffer(index, size); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/MapVector.java000066400000000000000000000215131515451615200317140ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.util.Preconditions.checkArgument; import java.util.List; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.AddOrGetResult; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.ZeroVector; import org.apache.arrow.vector.complex.impl.UnionMapReader; import org.apache.arrow.vector.complex.impl.UnionMapWriter; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType.Map; 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; /** * A MapVector is used to store entries of key/value pairs. It is a container vector that is * composed of a list of struct values with "key" and "value" fields. The MapVector is nullable, but * if a map is set at a given index, there must be an entry. In other words, the StructVector data * is non-nullable. Also for a given entry, the "key" is non-nullable, however the "value" can be * null. */ public class MapVector extends ListVector { public static final String KEY_NAME = "key"; public static final String VALUE_NAME = "value"; public static final String DATA_VECTOR_NAME = "entries"; /** * Construct an empty MapVector with no data. Child vectors must be added subsequently. * * @param name The name of the vector. * @param allocator The allocator used for allocating/reallocating buffers. * @param keysSorted True if the map keys have been pre-sorted. * @return a new instance of MapVector. */ public static MapVector empty(String name, BufferAllocator allocator, boolean keysSorted) { return new MapVector(name, allocator, FieldType.nullable(new Map(keysSorted)), null); } /** * Construct a MapVector instance. * * @param name The name of the vector. * @param allocator The allocator used for allocating/reallocating buffers. * @param fieldType The type definition of the MapVector. * @param callBack A schema change callback. */ public MapVector(String name, BufferAllocator allocator, FieldType fieldType, CallBack callBack) { super(name, allocator, fieldType, callBack); defaultDataVectorName = DATA_VECTOR_NAME; } public MapVector(Field field, BufferAllocator allocator, CallBack callBack) { super(field, allocator, callBack); defaultDataVectorName = DATA_VECTOR_NAME; } /** * Initialize child vectors of the map from the given list of fields. * * @param children List of fields that will be children of this MapVector. */ @Override public void initializeChildrenFromFields(List children) { checkArgument( children.size() == 1, "Maps have one List child. Found: %s", children.isEmpty() ? "none" : children); Field structField = children.get(0); MinorType minorType = Types.getMinorTypeForArrowType(structField.getType()); checkArgument( minorType == MinorType.STRUCT && !structField.isNullable(), "Map data should be a non-nullable struct type"); checkArgument( structField.getChildren().size() == 2, "Map data should be a struct with 2 children. Found: %s", children); Field keyField = structField.getChildren().get(0); checkArgument(!keyField.isNullable(), "Map data key type should be a non-nullable"); AddOrGetResult addOrGetVector = addOrGetVector(structField.getFieldType()); checkArgument( addOrGetVector.isCreated(), "Child vector already existed: %s", addOrGetVector.getVector()); addOrGetVector.getVector().initializeChildrenFromFields(structField.getChildren()); this.field = new Field(this.field.getName(), this.field.getFieldType(), children); } /** Get the writer for this MapVector instance. */ @Override public UnionMapWriter getWriter() { return new UnionMapWriter(this); } /** Get the reader for this MapVector instance. */ @Override public UnionMapReader getReader() { if (reader == null) { reader = new UnionMapReader(this); } return (UnionMapReader) reader; } @Override public MinorType getMinorType() { return MinorType.MAP; } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return getTransferPair(ref, allocator, null); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new TransferImpl(field, allocator, null); } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return new TransferImpl(ref, allocator, callBack); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { return new TransferImpl(field, allocator, callBack); } @Override public TransferPair makeTransferPair(ValueVector target) { return new MapVector.TransferImpl((MapVector) target); } private class TransferImpl implements TransferPair { MapVector to; TransferPair dataTransferPair; public TransferImpl(String name, BufferAllocator allocator, CallBack callBack) { this(new MapVector(name, allocator, field.getFieldType(), callBack)); } public TransferImpl(Field field, BufferAllocator allocator, CallBack callBack) { this(new MapVector(field, allocator, callBack)); } public TransferImpl(MapVector to) { this.to = to; to.addOrGetVector(vector.getField().getFieldType()); if (to.getDataVector() instanceof ZeroVector) { to.addOrGetVector(vector.getField().getFieldType()); } dataTransferPair = getDataVector().makeTransferPair(to.getDataVector()); } /** * Transfer this vector'data to another vector. The memory associated with this vector is * transferred to the allocator of target vector for accounting and management purposes. */ @Override public void transfer() { to.clear(); dataTransferPair.transfer(); to.validityBuffer = transferBuffer(validityBuffer, to.allocator); to.offsetBuffer = transferBuffer(offsetBuffer, to.allocator); to.lastSet = lastSet; if (valueCount > 0) { to.setValueCount(valueCount); } clear(); } /** * Slice this vector at desired index and length and transfer the corresponding data to the * target vector. * * @param startIndex start position of the split in source vector. * @param length length of the split. */ @Override public void splitAndTransfer(int startIndex, int length) { Preconditions.checkArgument( startIndex >= 0 && length >= 0 && startIndex + length <= valueCount, "Invalid parameters startIndex: %s, length: %s for valueCount: %s", startIndex, length, valueCount); to.clear(); if (length > 0) { final int startPoint = offsetBuffer.getInt(startIndex * OFFSET_WIDTH); final int sliceLength = offsetBuffer.getInt((startIndex + length) * OFFSET_WIDTH) - startPoint; to.offsetBuffer = to.allocateOffsetBuffer((length + 1) * OFFSET_WIDTH); /* splitAndTransfer offset buffer */ for (int i = 0; i < length + 1; i++) { final int relativeOffset = offsetBuffer.getInt((startIndex + i) * OFFSET_WIDTH) - startPoint; to.offsetBuffer.setInt(i * OFFSET_WIDTH, relativeOffset); } /* splitAndTransfer validity buffer */ splitAndTransferValidityBuffer(startIndex, length, to); /* splitAndTransfer data buffer */ dataTransferPair.splitAndTransfer(startPoint, sliceLength); to.lastSet = length - 1; to.setValueCount(length); } } @Override public ValueVector getTo() { return to; } @Override public void copyValueSafe(int from, int to) { this.to.copyFrom(from, to, MapVector.this); } } } NonNullableStructVector.java000066400000000000000000000363511515451615200345440ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.util.ArrayList; import java.util.Collection; import java.util.Iterator; import java.util.List; import java.util.Map; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; 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.DensityAwareVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueIterableVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.impl.SingleStructReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.ComplexHolder; 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.CallBack; import org.apache.arrow.vector.util.JsonStringHashMap; import org.apache.arrow.vector.util.TransferPair; /** * A struct vector that has no null values (and no validity buffer). Child Vectors are handled in * {@link AbstractStructVector}. */ public class NonNullableStructVector extends AbstractStructVector implements ValueIterableVector> { /** * Construct a new empty instance which replaces an existing field with the new one in case of * name conflict. */ public static NonNullableStructVector empty(String name, BufferAllocator allocator) { FieldType fieldType = new FieldType(false, ArrowType.Struct.INSTANCE, null, null); return new NonNullableStructVector( name, allocator, fieldType, null, ConflictPolicy.CONFLICT_REPLACE, false); } /** Construct a new empty instance which preserve fields with identical names. */ public static NonNullableStructVector emptyWithDuplicates( String name, BufferAllocator allocator) { FieldType fieldType = new FieldType(false, ArrowType.Struct.INSTANCE, null, null); return new NonNullableStructVector( name, allocator, fieldType, null, ConflictPolicy.CONFLICT_APPEND, true); } private final SingleStructReaderImpl reader = new SingleStructReaderImpl(this); protected Field field; public int valueCount; /** * Constructs a new instance. * * @param name The name of the instance. * @param allocator The allocator to use to allocating/reallocating buffers. * @param fieldType The type of this list. */ public NonNullableStructVector( String name, BufferAllocator allocator, FieldType fieldType, CallBack callBack) { this(new Field(name, fieldType, null), allocator, callBack); } /** * Constructs a new instance. * * @param field The field materialized by this vector. * @param allocator The allocator to use to allocating/reallocating buffers. * @param callBack A schema change callback. */ public NonNullableStructVector(Field field, BufferAllocator allocator, CallBack callBack) { this(field, allocator, callBack, null, true); } /** * Constructs a new instance. * * @param name The name of the instance. * @param allocator The allocator to use to allocating/reallocating buffers. * @param fieldType The type of this list. * @param callBack A schema change callback. * @param conflictPolicy How to handle duplicate field names in the struct. */ public NonNullableStructVector( String name, BufferAllocator allocator, FieldType fieldType, CallBack callBack, ConflictPolicy conflictPolicy, boolean allowConflictPolicyChanges) { this( new Field(name, fieldType, null), allocator, callBack, conflictPolicy, allowConflictPolicyChanges); } /** * Constructs a new instance. * * @param field The field materialized by this vector. * @param allocator The allocator to use to allocating/reallocating buffers. * @param callBack A schema change callback. * @param conflictPolicy How to handle duplicate field names in the struct. */ public NonNullableStructVector( Field field, BufferAllocator allocator, CallBack callBack, ConflictPolicy conflictPolicy, boolean allowConflictPolicyChanges) { super(field.getName(), allocator, callBack, conflictPolicy, allowConflictPolicyChanges); this.field = field; this.valueCount = 0; } @Override public FieldReader getReader() { return reader; } private transient StructTransferPair ephPair; /** * Copies the element at fromIndex in the provided vector to thisIndex. Reallocates buffers if * thisIndex is larger then current capacity. */ @Override public void copyFrom(int fromIndex, int thisIndex, ValueVector from) { Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); if (ephPair == null || ephPair.from != from) { ephPair = (StructTransferPair) from.makeTransferPair(this); } ephPair.copyValueSafe(fromIndex, thisIndex); } @Override public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) { copyFrom(fromIndex, thisIndex, from); } @Override protected boolean supportsDirectRead() { return true; } public Iterator fieldNameIterator() { return getChildFieldNames().iterator(); } @Override public void setInitialCapacity(int numRecords) { for (final ValueVector v : this) { v.setInitialCapacity(numRecords); } } @Override public void setInitialCapacity(int valueCount, double density) { for (final ValueVector vector : this) { if (vector instanceof DensityAwareVector) { ((DensityAwareVector) vector).setInitialCapacity(valueCount, density); } else { vector.setInitialCapacity(valueCount); } } } @Override public int getBufferSize() { if (valueCount == 0 || size() == 0) { return 0; } long buffer = 0; for (final ValueVector v : this) { buffer += v.getBufferSize(); } return (int) buffer; } @Override public int getBufferSizeFor(final int valueCount) { if (valueCount == 0) { return 0; } long bufferSize = 0; for (final ValueVector v : this) { bufferSize += v.getBufferSizeFor(valueCount); } return (int) bufferSize; } @Override public ArrowBuf getValidityBuffer() { throw new UnsupportedOperationException(); } @Override public ArrowBuf getDataBuffer() { throw new UnsupportedOperationException(); } @Override public ArrowBuf getOffsetBuffer() { throw new UnsupportedOperationException(); } @Override public TransferPair getTransferPair(BufferAllocator allocator) { return getTransferPair(name, allocator, null); } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return new StructTransferPair( this, new NonNullableStructVector( name, allocator, field.getFieldType(), callBack, getConflictPolicy(), allowConflictPolicyChanges), false); } @Override public TransferPair makeTransferPair(ValueVector to) { return new StructTransferPair(this, (NonNullableStructVector) to); } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new StructTransferPair( this, new NonNullableStructVector( ref, allocator, field.getFieldType(), callBack, getConflictPolicy(), allowConflictPolicyChanges), false); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new StructTransferPair( this, new NonNullableStructVector( field, allocator, callBack, getConflictPolicy(), allowConflictPolicyChanges), false); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { return new StructTransferPair( this, new NonNullableStructVector( field, allocator, callBack, getConflictPolicy(), allowConflictPolicyChanges), false); } /** {@link TransferPair} for this this class. */ protected static class StructTransferPair implements TransferPair { private final TransferPair[] pairs; private final NonNullableStructVector from; private final NonNullableStructVector to; public StructTransferPair(NonNullableStructVector from, NonNullableStructVector to) { this(from, to, true); } protected StructTransferPair( NonNullableStructVector from, NonNullableStructVector to, boolean allocate) { this.from = from; this.to = to; this.pairs = new TransferPair[from.size()]; this.to.ephPair = null; int i = 0; FieldVector vector; for (String child : from.getChildFieldNames()) { int preSize = to.size(); vector = from.getChild(child); if (vector == null) { continue; } // DRILL-1872: we add the child fields for the vector, looking up the field by name. For a // map vector, // the child fields may be nested fields of the top level child. For example if the // structure // of a child field is oa.oab.oabc then we add oa, then add oab to oa then oabc to oab. // But the children member of a Materialized field is a HashSet. If the fields are added in // the // children HashSet, and the hashCode of the Materialized field includes the hash code of // the // children, the hashCode value of oa changes *after* the field has been added to the // HashSet. // (This is similar to what happens in ScanBatch where the children cannot be added till // they are // read). To take care of this, we ensure that the hashCode of the MaterializedField does // not // include the hashCode of the children but is based only on MaterializedField$key. final FieldVector newVector = to.addOrGet(child, vector.getField().getFieldType(), vector.getClass()); if (allocate && to.size() != preSize) { newVector.allocateNew(); } pairs[i++] = vector.makeTransferPair(newVector); } } @Override public void transfer() { for (final TransferPair p : pairs) { p.transfer(); } to.valueCount = from.valueCount; from.clear(); } @Override public ValueVector getTo() { return to; } @Override public void copyValueSafe(int from, int to) { for (TransferPair p : pairs) { p.copyValueSafe(from, to); } } @Override public void splitAndTransfer(int startIndex, int length) { for (TransferPair p : pairs) { p.splitAndTransfer(startIndex, length); } to.setValueCount(length); } } @Override public int getValueCapacity() { if (size() == 0) { return 0; } return getChildren().stream().mapToInt(child -> child.getValueCapacity()).min().getAsInt(); } @Override public Map getObject(int index) { Map vv = new JsonStringHashMap<>(); for (String child : getChildFieldNames()) { ValueVector v = getChild(child); if (v != null && index < v.getValueCount()) { Object value = v.getObject(index); if (value != null) { vv.put(child, value); } } } return vv; } @Override public int hashCode(int index) { return hashCode(index, null); } @Override public int hashCode(int index, ArrowBufHasher hasher) { int hash = 0; for (FieldVector v : getChildren()) { if (index < v.getValueCount()) { hash = ByteFunctionHelpers.combineHash(hash, v.hashCode(index, hasher)); } } return hash; } @Override public OUT accept(VectorVisitor visitor, IN value) { return visitor.visit(this, value); } @Override public boolean isNull(int index) { return false; } @Override public int getNullCount() { return 0; } public void get(int index, ComplexHolder holder) { reader.setPosition(index); holder.reader = reader; } @Override public int getValueCount() { return valueCount; } public ValueVector getVectorById(int id) { return getChildByOrdinal(id); } /** Gets a child vector by ordinal position and casts to the specified class. */ public V getVectorById(int id, Class clazz) { ValueVector untyped = getVectorById(id); if (clazz.isInstance(untyped)) { return clazz.cast(untyped); } throw new ClassCastException( "Id " + id + " had the wrong type. Expected " + clazz.getCanonicalName() + " but was " + untyped.getClass().getCanonicalName()); } @Override public void setValueCount(int valueCount) { for (final ValueVector v : getChildren()) { v.setValueCount(valueCount); } NonNullableStructVector.this.valueCount = valueCount; } @Override public void clear() { for (final ValueVector v : getChildren()) { v.clear(); } valueCount = 0; } @Override public void reset() { for (final ValueVector v : getChildren()) { v.reset(); } valueCount = 0; } @Override public Field getField() { List children = new ArrayList<>(); for (ValueVector child : getChildren()) { children.add(child.getField()); } if (children.isEmpty() || field.getChildren().equals(children)) { return field; } field = new Field(field.getName(), field.getFieldType(), children); return field; } @Override public MinorType getMinorType() { return MinorType.STRUCT; } @Override public void close() { final Collection vectors = getChildren(); for (final FieldVector v : vectors) { v.close(); } vectors.clear(); valueCount = 0; super.close(); } /** Initializes the struct's members from the given Fields. */ public void initializeChildrenFromFields(List children) { for (Field field : children) { FieldVector vector = (FieldVector) this.add(field.getName(), field.getFieldType()); vector.initializeChildrenFromFields(field.getChildren()); } } public List getChildrenFromFields() { return getChildren(); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/Positionable.java000066400000000000000000000020561515451615200324450ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** Get and set position in a particular data structure. */ @SuppressWarnings("unused") // Used in when instantiating freemarker templates. public interface Positionable { int getPosition(); void setPosition(int index); } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/PromotableVector.java000066400000000000000000000023021515451615200332760ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import org.apache.arrow.vector.AddOrGetResult; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.types.pojo.FieldType; /** Vector that can store multiple {@linkplain FieldType} vectors as children. */ public interface PromotableVector { AddOrGetResult addOrGetVector(FieldType type); UnionVector promoteToUnion(); } RepeatedFixedWidthVectorLike.java000066400000000000000000000024731515451615200354420ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * A {@link org.apache.arrow.vector.ValueVector} mix-in that can be used in conjunction with {@link * RepeatedValueVector} subtypes. */ public interface RepeatedFixedWidthVectorLike { /** * Allocate a new memory space for this vector. Must be called prior to using the ValueVector. * * @param valueCount Number of separate repeating groupings. * @param innerValueCount Number of supported values in the vector. */ void allocateNew(int valueCount, int innerValueCount); } RepeatedValueVector.java000066400000000000000000000036041515451615200336470ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import org.apache.arrow.vector.DensityAwareVector; import org.apache.arrow.vector.UInt4Vector; import org.apache.arrow.vector.ValueVector; /** * An abstraction representing repeated value vectors. * *

    A repeated vector contains values that may either be flat or nested. A value consists of zero * or more cells(inner values). Current design maintains data and offsets vectors. Each cell is * stored in the data vector. Repeated vector uses the offset vector to determine the sequence of * cells pertaining to an individual value. */ public interface RepeatedValueVector extends ValueVector, DensityAwareVector { int DEFAULT_REPEAT_PER_RECORD = 5; /** * Get the offset vector. * * @return the underlying offset vector or null if none exists. * @deprecated This API will be removed, as the current implementations no longer hold inner * offset vectors. */ @Deprecated UInt4Vector getOffsetVector(); /** * Get the data vector. * * @return the underlying data vector or null if none exists. */ ValueVector getDataVector(); } RepeatedVariableWidthVectorLike.java000066400000000000000000000031461515451615200361260ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * A {@link org.apache.arrow.vector.ValueVector} mix-in that can be used in conjunction with * variable {@link RepeatedValueVector} subtypes (e.g. Strings, Lists, etc). */ public interface RepeatedVariableWidthVectorLike { /** * Allocate a new memory space for this vector. Must be called prior to using the ValueVector. * * @param totalBytes Desired size of the underlying data buffer. * @param parentValueCount Number of separate repeating groupings. * @param childValueCount Number of supported values in the vector. */ void allocateNew(int totalBytes, int parentValueCount, int childValueCount); /** * Provide the maximum amount of variable width bytes that can be stored int his vector. * * @return the byte capacity */ int getByteCapacity(); } RunEndEncodedVector.java000066400000000000000000000722201515451615200335760ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.apache.arrow.util.Preconditions.checkArgument; import java.util.Arrays; import java.util.Collections; import java.util.Iterator; import java.util.List; import java.util.NoSuchElementException; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; 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; import org.apache.arrow.vector.BufferBacked; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.SmallIntVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.ZeroVector; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.complex.writer.FieldWriter; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; 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.CallBack; import org.apache.arrow.vector.util.TransferPair; /** * A run-end encoded vector contains only two child vectors: a run_end vector of type int and a * values vector of any type. There are no buffers associated with the parent vector. */ public class RunEndEncodedVector extends BaseValueVector implements FieldVector { public static final FieldVector DEFAULT_VALUE_VECTOR = ZeroVector.INSTANCE; public static final FieldVector DEFAULT_RUN_END_VECTOR = ZeroVector.INSTANCE; public static RunEndEncodedVector empty(String name, BufferAllocator allocator) { return new RunEndEncodedVector( name, allocator, FieldType.notNullable(ArrowType.RunEndEncoded.INSTANCE), null); } protected final CallBack callBack; protected Field field; protected FieldVector runEndsVector; protected FieldVector valuesVector; protected int valueCount; /** * Constructs a new instance. * * @param name The name of the instance. * @param allocator The allocator to use for allocating/reallocating buffers. * @param fieldType The type of the array that is run-end encoded. * @param callBack A schema change callback. */ public RunEndEncodedVector( String name, BufferAllocator allocator, FieldType fieldType, CallBack callBack) { this(new Field(name, fieldType, null), allocator, callBack); } /** * Constructs a new instance. * * @param field The field materialized by this vector. * @param allocator The allocator to use for allocating/reallocating buffers. * @param callBack A schema change callback. */ public RunEndEncodedVector(Field field, BufferAllocator allocator, CallBack callBack) { this(field, allocator, DEFAULT_RUN_END_VECTOR, DEFAULT_VALUE_VECTOR, callBack); } /** * Constructs a new instance. * * @param field The field materialized by this vector. * @param allocator The allocator to use for allocating/reallocating buffers. * @param runEndsVector The vector represents run ends. Only Zero vector or type int vector with * size 16, 32 is allowed * @param valuesVector The vector represents values * @param callBack A schema change callback. */ public RunEndEncodedVector( Field field, BufferAllocator allocator, FieldVector runEndsVector, FieldVector valuesVector, CallBack callBack) { super(allocator); this.field = field; this.callBack = callBack; this.valueCount = 0; this.runEndsVector = runEndsVector; this.valuesVector = valuesVector; } /** ValueVector interface */ /** * Allocate new buffers. ValueVector implements logic to determine how much to allocate. * * @throws OutOfMemoryException Thrown if no memory can be allocated. */ @Override public void allocateNew() throws OutOfMemoryException { if (!allocateNewSafe()) { throw new OutOfMemoryException("Failure while allocating memory"); } } /** * Allocates new buffers. ValueVector implements logic to determine how much to allocate. * * @return Returns true if allocation was successful. */ @Override public boolean allocateNewSafe() { initializeChildrenFromFields(field.getChildren()); for (FieldVector v : getChildrenFromFields()) { boolean isAllocated = v.allocateNewSafe(); if (!isAllocated) { v.clear(); return false; } } return true; } /** * Allocate new buffer with double capacity, and copy data into the new buffer. Replace vector's * buffer with new buffer, and release old one */ @Override public void reAlloc() { for (FieldVector v : getChildrenFromFields()) { v.reAlloc(); } } @Override public BufferAllocator getAllocator() { return allocator; } @Override protected FieldReader getReaderImpl() { throw new UnsupportedOperationException("Not yet implemented."); } /** * Set the initial record capacity. * * @param numRecords the initial record capacity. */ @Override public void setInitialCapacity(int numRecords) {} /** * Returns the maximum number of values that can be stored in this vector instance. * * @return the maximum number of values that can be stored in this vector instance. */ @Override public int getValueCapacity() { return getChildrenFromFields().stream() .mapToInt(item -> item != null ? item.getValueCapacity() : 0) .min() .orElseThrow(NoSuchElementException::new); } /** Alternative to clear(). Allows use as an AutoCloseable in try-with-resources. */ @Override public void close() { for (FieldVector v : getChildrenFromFields()) { v.close(); } } /** * Release any owned ArrowBuf and reset the ValueVector to the initial state. If the vector has * any child vectors, they will also be cleared. */ @Override public void clear() { for (FieldVector v : getChildrenFromFields()) { v.clear(); } this.valueCount = 0; } /** * Reset the ValueVector to the initial state without releasing any owned ArrowBuf. Buffer * capacities will remain unchanged and any previous data will be zeroed out. This includes * buffers for data, validity, offset, etc. If the vector has any child vectors, they will also be * reset. */ @Override public void reset() { for (FieldVector v : getChildrenFromFields()) { v.reset(); } valueCount = 0; } /** * Get information about how this field is materialized. * * @return the field corresponding to this vector */ @Override public Field getField() { return field; } @Override public MinorType getMinorType() { return MinorType.RUNENDENCODED; } /** * To transfer quota responsibility. * * @param ref the name of the vector * @param allocator the target allocator * @return a {@link org.apache.arrow.vector.util.TransferPair transfer pair}, creating a new * target vector of the same type. */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return getTransferPair(ref, allocator, null); } /** * To transfer quota responsibility. * * @param field the Field object used by the target vector * @param allocator the target allocator * @return a {@link org.apache.arrow.vector.util.TransferPair transfer pair}, creating a new * target vector of the same type. */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return getTransferPair(field, allocator, null); } /** * To transfer quota responsibility. * * @param ref the name of the vector * @param allocator the target allocator * @param callBack A schema change callback. * @return a {@link org.apache.arrow.vector.util.TransferPair transfer pair}, creating a new * target vector of the same type. */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return new TransferImpl(ref, allocator, callBack); } /** * To transfer quota responsibility. * * @param field the Field object used by the target vector * @param allocator the target allocator * @param callBack A schema change callback. * @return a {@link org.apache.arrow.vector.util.TransferPair transfer pair}, creating a new * target vector of the same type. */ @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { return new TransferImpl(field, allocator, callBack); } /** * Makes a new transfer pair used to transfer underlying buffers. * * @param target the target for the transfer * @return a new {@link org.apache.arrow.vector.util.TransferPair transfer pair} that is used to * transfer underlying buffers into the target vector. */ @Override public TransferPair makeTransferPair(ValueVector target) { return new TransferImpl((RunEndEncodedVector) target); } private class TransferImpl implements TransferPair { RunEndEncodedVector to; TransferPair dataTransferPair; TransferPair reeTransferPair; public TransferImpl(String name, BufferAllocator allocator, CallBack callBack) { this(new RunEndEncodedVector(name, allocator, field.getFieldType(), callBack)); } public TransferImpl(Field field, BufferAllocator allocator, CallBack callBack) { this(new RunEndEncodedVector(field, allocator, callBack)); } public TransferImpl(RunEndEncodedVector to) { this.to = to; if (to.getRunEndsVector() instanceof ZeroVector) { to.initializeChildrenFromFields(field.getChildren()); } reeTransferPair = getRunEndsVector().makeTransferPair(to.getRunEndsVector()); dataTransferPair = getValuesVector().makeTransferPair(to.getValuesVector()); } /** * Transfer the vector data to another vector. The memory associated with this vector is * transferred to the allocator of target vector for accounting and management purposes. */ @Override public void transfer() { to.clear(); dataTransferPair.transfer(); reeTransferPair.transfer(); if (valueCount > 0) { to.setValueCount(valueCount); } clear(); } /** * Slice this vector at the desired index and length, then transfer the corresponding data to * the target vector. * * @param startIndex start position of the split in source vector. * @param length length of the split. */ @Override public void splitAndTransfer(int startIndex, int length) { to.clear(); if (length <= 0) { return; } int physicalStartIndex = getPhysicalIndex(startIndex); int physicalEndIndex = getPhysicalIndex(startIndex + length - 1); int physicalLength = physicalEndIndex - physicalStartIndex + 1; dataTransferPair.splitAndTransfer(physicalStartIndex, physicalLength); FieldVector toRunEndsVector = to.runEndsVector; if (startIndex == 0) { if (((BaseIntVector) runEndsVector).getValueAsLong(physicalEndIndex) == length) { reeTransferPair.splitAndTransfer(physicalStartIndex, physicalLength); } else { reeTransferPair.splitAndTransfer(physicalStartIndex, physicalLength - 1); toRunEndsVector.setValueCount(physicalLength); if (toRunEndsVector instanceof SmallIntVector) { ((SmallIntVector) toRunEndsVector).set(physicalEndIndex, length); } else if (toRunEndsVector instanceof IntVector) { ((IntVector) toRunEndsVector).set(physicalEndIndex, length); } else if (toRunEndsVector instanceof BigIntVector) { ((BigIntVector) toRunEndsVector).set(physicalEndIndex, length); } else { throw new IllegalArgumentException( "Run-end vector and must be of type int with size 16, 32, or 64 bits."); } } } else { shiftRunEndsVector( toRunEndsVector, startIndex, length, physicalStartIndex, physicalEndIndex, physicalLength); } getTo().setValueCount(length); } private void shiftRunEndsVector( ValueVector toRunEndVector, int startIndex, int length, int physicalStartIndex, int physicalEndIndex, int physicalLength) { toRunEndVector.setValueCount(physicalLength); toRunEndVector.getValidityBuffer().setOne(0, toRunEndVector.getValidityBuffer().capacity()); ArrowBuf fromRunEndBuffer = runEndsVector.getDataBuffer(); ArrowBuf toRunEndBuffer = toRunEndVector.getDataBuffer(); int physicalLastIndex = physicalLength - 1; if (toRunEndVector instanceof SmallIntVector) { byte typeWidth = SmallIntVector.TYPE_WIDTH; for (int i = 0; i < physicalLastIndex; i++) { toRunEndBuffer.setShort( (long) i * typeWidth, fromRunEndBuffer.getShort((long) (i + physicalStartIndex) * typeWidth) - startIndex); } int lastEnd = Math.min( fromRunEndBuffer.getShort((long) physicalEndIndex * typeWidth) - startIndex, length); toRunEndBuffer.setShort((long) physicalLastIndex * typeWidth, lastEnd); } else if (toRunEndVector instanceof IntVector) { byte typeWidth = IntVector.TYPE_WIDTH; for (int i = 0; i < physicalLastIndex; i++) { toRunEndBuffer.setInt( (long) i * typeWidth, fromRunEndBuffer.getInt((long) (i + physicalStartIndex) * typeWidth) - startIndex); } int lastEnd = Math.min( fromRunEndBuffer.getInt((long) physicalEndIndex * typeWidth) - startIndex, length); toRunEndBuffer.setInt((long) physicalLastIndex * typeWidth, lastEnd); } else if (toRunEndVector instanceof BigIntVector) { byte typeWidth = BigIntVector.TYPE_WIDTH; for (int i = 0; i < physicalLastIndex; i++) { toRunEndBuffer.setLong( (long) i * typeWidth, fromRunEndBuffer.getLong((long) (i + physicalStartIndex) * typeWidth) - startIndex); } long lastEnd = Math.min( fromRunEndBuffer.getLong((long) physicalEndIndex * typeWidth) - startIndex, length); toRunEndBuffer.setLong((long) physicalLastIndex * typeWidth, lastEnd); } else { throw new IllegalArgumentException( "Run-end vector and must be of type int with size 16, 32, or 64 bits."); } } @Override public ValueVector getTo() { return to; } @Override public void copyValueSafe(int from, int to) { this.to.copyFrom(from, to, RunEndEncodedVector.this); } } /** * Get a reader for this vector. * * @return a {@link org.apache.arrow.vector.complex.reader.FieldReader field reader} that supports * reading values from this vector. */ @Override public FieldReader getReader() { throw new UnsupportedOperationException("Not yet implemented."); } /** * Get a writer for this vector. * * @return a {@link org.apache.arrow.vector.complex.writer.FieldWriter field writer} that supports * writing values to this vector. */ public FieldWriter getWriter() { throw new UnsupportedOperationException("Not yet implemented."); } /** * Get the number of bytes used by this vector. * * @return the number of bytes that is used by this vector instance. */ @Override public int getBufferSize() { int bufferSize = 0; for (FieldVector v : getChildrenFromFields()) { bufferSize += v.getBufferSize(); } return bufferSize; } /** * Returns the number of bytes that is used by this vector if it holds the given number of values. * The result will be the same as if setValueCount() were called, followed by calling * getBufferSize(), but without any of the closing side-effects that setValueCount() implies wrt * finishing off the population of a vector. Some operations might wish to use this to determine * how much memory has been used by a vector so far, even though it is not finished being * populated. * * @param valueCount the number of values to assume this vector contains * @return the buffer size if this vector is holding valueCount values */ @Override public int getBufferSizeFor(int valueCount) { return 0; } /** * Return the underlying buffers associated with this vector. Note that this doesn't impact the * reference counts for this buffer so it only should be used for in-context access. Also note * that this buffer changes regularly thus external classes shouldn't hold a reference to it * (unless they change it). * * @param clear Whether to clear vector before returning; the buffers will still be refcounted; * but the returned array will be the only reference to them * @return The underlying {@link ArrowBuf buffers} that is used by this vector instance. */ @Override public ArrowBuf[] getBuffers(boolean clear) { return new ArrowBuf[0]; } /** * Gets the underlying buffer associated with validity vector. * * @return buffer */ @Override public ArrowBuf getValidityBuffer() { throw new UnsupportedOperationException( "Run-end encoded vectors do not have a validity buffer."); } /** * Gets the underlying buffer associated with data vector. * * @return buffer */ @Override public ArrowBuf getDataBuffer() { throw new UnsupportedOperationException("Run-end encoded vectors do not have a data buffer."); } /** * Gets the underlying buffer associated with offset vector. * * @return buffer */ @Override public ArrowBuf getOffsetBuffer() { throw new UnsupportedOperationException("Run-end encoded vectors do not have a offset buffer."); } /** * Gets the number of values. * * @return number of values in the vector */ @Override public int getValueCount() { return valueCount; } /** Set number of values in the vector. */ @Override public void setValueCount(int valueCount) { this.valueCount = valueCount; } /** * Get friendly type object from the vector. * * @param index index of object to get * @return friendly type object */ @Override public Object getObject(int index) { checkIndex(index); int physicalIndex = getPhysicalIndex(index); return valuesVector.getObject(physicalIndex); } /** * Get the run end of giving index. * * @param index index of the run end to get * @return the run end of giving index */ public int getRunEnd(int index) { checkIndex(index); int physicalIndex = getPhysicalIndex(index); return (int) ((BaseIntVector) runEndsVector).getValueAsLong(physicalIndex); } /** * Returns number of null elements in the vector. * * @return number of null elements */ @Override public int getNullCount() { // Null count is always 0 for run-end encoded array return 0; } /** * Check whether an element in the vector is null. * * @param index index to check for null * @return true if element is null */ @Override public boolean isNull(int index) { int physicalIndex = getPhysicalIndex(runEndsVector, index); return valuesVector.isNull(physicalIndex); } /** Returns hashCode of element in index with the default hasher. */ @Override public int hashCode(int index) { return hashCode(index, null); } /** Returns hashCode of element in index with the given hasher. */ @Override public int hashCode(int index, ArrowBufHasher hasher) { int hash = 0; for (FieldVector v : getChildrenFromFields()) { if (index < v.getValueCount()) { hash = ByteFunctionHelpers.combineHash(hash, v.hashCode(index, hasher)); } } return hash; } /** * Accept a generic {@link VectorVisitor} and return the result. * * @param the output result type. * @param the input data together with visitor. */ @Override public OUT accept(VectorVisitor visitor, IN value) { return visitor.visit(this, value); } /** * Gets the name of the vector. * * @return the name of the vector. */ @Override public String getName() { return this.field.getName(); } @Override public Iterator iterator() { return Collections.unmodifiableCollection(getChildrenFromFields()).iterator(); } /** FieldVector interface */ /** * Initializes the child vectors to be later loaded with loadBuffers. * * @param children the schema containing the run_ends column first and the values column second */ @Override public void initializeChildrenFromFields(List children) { checkArgument( children.size() == 2, "Run-end encoded vectors must have two child Fields. Found: %s", children.isEmpty() ? "none" : children); checkArgument( Arrays.asList( MinorType.SMALLINT.getType(), MinorType.INT.getType(), MinorType.BIGINT.getType()) .contains(children.get(0).getType()), "The first field represents the run-end vector and must be of type int " + "with size 16, 32, or 64 bits. Found: %s", children.get(0).getType()); runEndsVector = (BaseIntVector) children.get(0).createVector(allocator); valuesVector = children.get(1).createVector(allocator); field = new Field(field.getName(), field.getFieldType(), children); } /** * The returned list is the same size as the list passed to initializeChildrenFromFields. * * @return the children according to schema (empty for primitive types) */ @Override public List getChildrenFromFields() { return Arrays.asList(runEndsVector, valuesVector); } /** * Loads data in the vectors. (ownBuffers must be the same size as getFieldVectors()) * * @param fieldNode the fieldNode * @param ownBuffers the buffers for this Field (own buffers only, children not included) */ @Override public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers) { if (!ownBuffers.isEmpty()) { throw new UnsupportedOperationException( "Run-end encoded vectors do not have any associated buffers."); } this.valueCount = fieldNode.getLength(); } /** * Get the buffers of the fields, (same size as getFieldVectors() since it is their content). * * @return the buffers containing the data for this vector (ready for reading) */ @Override public List getFieldBuffers() { return List.of(); } /** * Get the inner vectors. * * @return the inner vectors for this field as defined by the TypeLayout * @deprecated This API will be removed as the current implementations no longer support inner * vectors. */ @Deprecated @Override public List getFieldInnerVectors() { throw new UnsupportedOperationException("There are no inner vectors. Use getFieldBuffers()."); } /** * Gets the starting address of the underlying buffer associated with validity vector. * * @return buffer address */ @Override public long getValidityBufferAddress() { throw new UnsupportedOperationException( "Run-end encoded vectors do not have a validity buffer."); } /** * Gets the starting address of the underlying buffer associated with data vector. * * @return buffer address */ @Override public long getDataBufferAddress() { throw new UnsupportedOperationException("Run-end encoded vectors do not have a data buffer."); } /** * Gets the starting address of the underlying buffer associated with offset vector. * * @return buffer address */ @Override public long getOffsetBufferAddress() { throw new UnsupportedOperationException( "Run-end encoded vectors do not have an offset buffer."); } /** * Set the element at the given index to null. * * @param index the value to change */ @Override public void setNull(int index) { throw new UnsupportedOperationException( "Run-end encoded vectors do not have a validity buffer."); } public FieldVector getRunEndsVector() { return runEndsVector; } public FieldVector getValuesVector() { return valuesVector; } private void checkIndex(int logicalIndex) { if (logicalIndex < 0 || logicalIndex >= valueCount) { throw new IndexOutOfBoundsException( String.format("index: %s, expected range (0, %s)", logicalIndex, valueCount)); } } /** * The physical index is the index of the first value that is larger than logical index. e.g. if * run_ends is [1,2,3], the physical index of logical index from 0 to 5 is [0, 1, 1, 2, 2, 2] */ public int getPhysicalIndex(int logicalIndex) { return getPhysicalIndex(runEndsVector, logicalIndex); } static int getPhysicalIndex(FieldVector runEndVector, int logicalIndex) { if (runEndVector == null || runEndVector.getValueCount() == 0) { return -1; } int low = 0; int high = runEndVector.getValueCount() - 1; int result = -1; while (low <= high) { int mid = low + (high - low) / 2; long valueAsLong = ((BaseIntVector) runEndVector).getValueAsLong(mid); if (valueAsLong > logicalIndex) { result = mid; high = mid - 1; } else { low = mid + 1; } } 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; } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/StateTool.java000066400000000000000000000031231515451615200317270ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import java.util.Arrays; /** Utility methods for state machines based on enums. */ public class StateTool { private StateTool() {} static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(StateTool.class); /** * Verifies currentState is in one of expectedStates, throws an * IllegalArgumentException if it isn't. */ public static > void check(T currentState, T... expectedStates) { for (T s : expectedStates) { if (s == currentState) { return; } } throw new IllegalArgumentException( String.format( "Expected to be in one of these states %s but was actually in " + "state %s", Arrays.toString(expectedStates), currentState)); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java000066400000000000000000000540261515451615200324700ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; 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; import java.util.List; import java.util.Map; 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.CommonUtil; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BaseValueVector; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.BufferBacked; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueIterableVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.impl.NullableStructReaderImpl; import org.apache.arrow.vector.complex.impl.NullableStructWriter; import org.apache.arrow.vector.holders.ComplexHolder; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; 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; import org.apache.arrow.vector.util.CallBack; import org.apache.arrow.vector.util.OversizedAllocationException; import org.apache.arrow.vector.util.TransferPair; /** * A Struct vector consists of nullability/validity buffer and children vectors that make up the * struct's fields. The children vectors are handled by the parent class. */ public class StructVector extends NonNullableStructVector implements FieldVector, ValueIterableVector> { /** * Construct a new empty instance which replaces an existing field with the new one in case of * name conflict. */ public static StructVector empty(String name, BufferAllocator allocator) { FieldType fieldType = FieldType.nullable(Struct.INSTANCE); return new StructVector( name, allocator, fieldType, null, ConflictPolicy.CONFLICT_REPLACE, false); } /** Construct a new empty instance which preserve fields with identical names. */ public static StructVector emptyWithDuplicates(String name, BufferAllocator allocator) { FieldType fieldType = new FieldType(false, ArrowType.Struct.INSTANCE, null, null); return new StructVector(name, allocator, fieldType, null, ConflictPolicy.CONFLICT_APPEND, true); } private final NullableStructReaderImpl reader = new NullableStructReaderImpl(this); private final NullableStructWriter writer = new NullableStructWriter(this); protected ArrowBuf validityBuffer; private int validityAllocationSizeInBytes; /** * Constructs a new instance. * * @param name The name of the instance. * @param allocator The allocator to use to allocating/reallocating buffers. * @param fieldType The type of this list. * @param callBack A schema change callback. */ public StructVector( String name, BufferAllocator allocator, FieldType fieldType, CallBack callBack) { super(name, checkNotNull(allocator), fieldType, callBack); this.validityBuffer = allocator.getEmpty(); this.validityAllocationSizeInBytes = getValidityBufferSizeFromCount(BaseValueVector.INITIAL_VALUE_ALLOCATION); } /** * Constructs a new instance. * * @param name The name of the instance. * @param allocator The allocator to use to allocating/reallocating buffers. * @param fieldType The type of this list. * @param callBack A schema change callback. * @param conflictPolicy policy to determine how duplicate names are handled. * @param allowConflictPolicyChanges whether duplicate names are allowed at all. */ public StructVector( String name, BufferAllocator allocator, FieldType fieldType, CallBack callBack, ConflictPolicy conflictPolicy, boolean allowConflictPolicyChanges) { super( name, checkNotNull(allocator), fieldType, callBack, conflictPolicy, allowConflictPolicyChanges); this.validityBuffer = allocator.getEmpty(); this.validityAllocationSizeInBytes = getValidityBufferSizeFromCount(BaseValueVector.INITIAL_VALUE_ALLOCATION); } /** * Constructs a new instance. * * @param field The field materialized by this vector. * @param allocator The allocator to use to allocating/reallocating buffers. * @param callBack A schema change callback. */ public StructVector(Field field, BufferAllocator allocator, CallBack callBack) { super(field, checkNotNull(allocator), callBack); this.validityBuffer = allocator.getEmpty(); this.validityAllocationSizeInBytes = getValidityBufferSizeFromCount(BaseValueVector.INITIAL_VALUE_ALLOCATION); } /** * Constructs a new instance. * * @param field The field materialized by this vector. * @param allocator The allocator to use to allocating/reallocating buffers. * @param callBack A schema change callback. * @param conflictPolicy policy to determine how duplicate names are handled. * @param allowConflictPolicyChanges whether duplicate names are allowed at all. */ public StructVector( Field field, BufferAllocator allocator, CallBack callBack, ConflictPolicy conflictPolicy, boolean allowConflictPolicyChanges) { super(field, checkNotNull(allocator), callBack, conflictPolicy, allowConflictPolicyChanges); this.validityBuffer = allocator.getEmpty(); this.validityAllocationSizeInBytes = getValidityBufferSizeFromCount(BaseValueVector.INITIAL_VALUE_ALLOCATION); } @Override public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers) { if (ownBuffers.size() != 1) { throw new IllegalArgumentException( "Illegal buffer count, expected " + 1 + ", got: " + ownBuffers.size()); } ArrowBuf bitBuffer = ownBuffers.get(0); validityBuffer.getReferenceManager().release(); validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); valueCount = fieldNode.getLength(); validityAllocationSizeInBytes = checkedCastToInt(validityBuffer.capacity()); } @Override public List getFieldBuffers() { List result = new ArrayList<>(1); setReaderAndWriterIndex(); result.add(validityBuffer); return result; } private void setReaderAndWriterIndex() { validityBuffer.readerIndex(0); validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount)); } /** * Get the inner vectors. * * @return the inner vectors for this field as defined by the TypeLayout * @deprecated This API will be removed as the current implementations no longer support inner * vectors. */ @Deprecated @Override public List getFieldInnerVectors() { throw new UnsupportedOperationException("There are no inner vectors. Use getFieldBuffers"); } @Override public NullableStructReaderImpl getReader() { return reader; } public NullableStructWriter getWriter() { return writer; } @Override public TransferPair getTransferPair(BufferAllocator allocator) { return new NullableStructTransferPair( this, new StructVector( name, allocator, field.getFieldType(), null, getConflictPolicy(), allowConflictPolicyChanges), false); } @Override public TransferPair makeTransferPair(ValueVector to) { return new NullableStructTransferPair(this, (StructVector) to, false); } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { return new NullableStructTransferPair( this, new StructVector( ref, allocator, field.getFieldType(), null, getConflictPolicy(), allowConflictPolicyChanges), false); } @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { return new NullableStructTransferPair( this, new StructVector( ref, allocator, field.getFieldType(), callBack, getConflictPolicy(), allowConflictPolicyChanges), false); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { return new NullableStructTransferPair( this, new StructVector(field, allocator, null, getConflictPolicy(), allowConflictPolicyChanges), false); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { return new NullableStructTransferPair( this, new StructVector( field, allocator, callBack, getConflictPolicy(), allowConflictPolicyChanges), false); } /** {@link TransferPair} for this (nullable) {@link StructVector}. */ protected class NullableStructTransferPair extends StructTransferPair { private StructVector target; protected NullableStructTransferPair(StructVector from, StructVector to, boolean allocate) { super(from, to, allocate); this.target = to; } @Override public void transfer() { target.clear(); target.validityBuffer = BaseValueVector.transferBuffer(validityBuffer, target.allocator); super.transfer(); clear(); } @Override public void copyValueSafe(int fromIndex, int toIndex) { while (toIndex >= target.getValidityBufferValueCapacity()) { target.reallocValidityBuffer(); } BitVectorHelper.setValidityBit(target.validityBuffer, toIndex, isSet(fromIndex)); super.copyValueSafe(fromIndex, toIndex); } @Override public void splitAndTransfer(int startIndex, int length) { Preconditions.checkArgument( startIndex >= 0 && length >= 0 && startIndex + length <= valueCount, "Invalid parameters startIndex: %s, length: %s for valueCount: %s", startIndex, length, valueCount); target.clear(); splitAndTransferValidityBuffer(startIndex, length, target); super.splitAndTransfer(startIndex, length); } } /* * transfer the validity. */ private void splitAndTransferValidityBuffer(int startIndex, int length, StructVector 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); } } } } /** * Get the value capacity of the internal validity buffer. * * @return number of elements that validity buffer can hold */ private int getValidityBufferValueCapacity() { return checkedCastToInt(validityBuffer.capacity() * 8); } /** * Get the current value capacity for the vector. * * @return number of elements that vector can hold. */ @Override public int getValueCapacity() { return Math.min(getValidityBufferValueCapacity(), super.getValueCapacity()); } /** * Return the underlying buffers associated with this vector. Note that this doesn't impact the * reference counts for this buffer, so it only should be used for in-context access. Also note * that this buffer changes regularly, thus external classes shouldn't hold a reference to it * (unless they change it). * * @param clear Whether to clear vector before returning, the buffers will still be refcounted but * the returned array will be the only reference to them. Also, this won't clear the child * buffers. * @return The underlying {@link ArrowBuf buffers} that is used by this vector instance. */ @Override public ArrowBuf[] getBuffers(boolean clear) { setReaderAndWriterIndex(); final ArrowBuf[] buffers; if (getBufferSize() == 0) { buffers = new ArrowBuf[0]; } else { List list = new ArrayList<>(); list.add(validityBuffer); list.addAll(Arrays.asList(super.getBuffers(false))); buffers = list.toArray(new ArrowBuf[list.size()]); } if (clear) { for (ArrowBuf buffer : buffers) { buffer.getReferenceManager().retain(); } clear(); } return buffers; } /** Close the vector and release the associated buffers. */ @Override public void close() { clearValidityBuffer(); super.close(); } /** Same as {@link #close()}. */ @Override public void clear() { clearValidityBuffer(); super.clear(); } /** Reset this vector to empty, does not release buffers. */ @Override public void reset() { super.reset(); validityBuffer.setZero(0, validityBuffer.capacity()); } /** Release the validity buffer. */ private void clearValidityBuffer() { validityBuffer.getReferenceManager().release(); validityBuffer = allocator.getEmpty(); } /** * Get the size (number of bytes) of underlying buffers used by this vector. * * @return size of underlying buffers. */ @Override public int getBufferSize() { if (valueCount == 0) { return 0; } return super.getBufferSize() + getValidityBufferSizeFromCount(valueCount); } /** * Get the potential buffer size for a particular number of records. * * @param valueCount desired number of elements in the vector * @return estimated size of underlying buffers if the vector holds a given number of elements */ @Override public int getBufferSizeFor(final int valueCount) { if (valueCount == 0) { return 0; } return super.getBufferSizeFor(valueCount) + getValidityBufferSizeFromCount(valueCount); } @Override public void setInitialCapacity(int numRecords) { validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords); super.setInitialCapacity(numRecords); } @Override public void setInitialCapacity(int numRecords, double density) { validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords); super.setInitialCapacity(numRecords, density); } @Override public boolean allocateNewSafe() { /* Boolean to keep track if all the memory allocations were successful * Used in the case of composite vectors when we need to allocate multiple * buffers for multiple vectors. If one of the allocations failed we need to * clear all the memory that we allocated */ boolean success = false; try { clear(); allocateValidityBuffer(validityAllocationSizeInBytes); success = super.allocateNewSafe(); } finally { if (!success) { clear(); } } 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 public void reAlloc() { /* reallocate the validity buffer */ reallocValidityBuffer(); super.reAlloc(); } private void reallocValidityBuffer() { final int currentBufferCapacity = checkedCastToInt(validityBuffer.capacity()); long newAllocationSize = getNewAllocationSize(currentBufferCapacity); final ArrowBuf newBuf = allocator.buffer(newAllocationSize); newBuf.setBytes(0, validityBuffer, 0, currentBufferCapacity); newBuf.setZero(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); validityBuffer.getReferenceManager().release(1); validityBuffer = newBuf; validityAllocationSizeInBytes = (int) newAllocationSize; } private long getNewAllocationSize(int currentBufferCapacity) { long newAllocationSize = currentBufferCapacity * 2L; if (newAllocationSize == 0) { if (validityAllocationSizeInBytes > 0) { newAllocationSize = validityAllocationSizeInBytes; } else { newAllocationSize = getValidityBufferSizeFromCount(BaseValueVector.INITIAL_VALUE_ALLOCATION) * 2L; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); assert newAllocationSize >= 1; if (newAllocationSize > BaseValueVector.MAX_ALLOCATION_SIZE) { throw new OversizedAllocationException("Unable to expand the buffer"); } return newAllocationSize; } @Override public long getValidityBufferAddress() { return validityBuffer.memoryAddress(); } @Override public long getDataBufferAddress() { throw new UnsupportedOperationException(); } @Override public long getOffsetBufferAddress() { throw new UnsupportedOperationException(); } @Override public ArrowBuf getValidityBuffer() { return validityBuffer; } @Override public ArrowBuf getDataBuffer() { throw new UnsupportedOperationException(); } @Override public ArrowBuf getOffsetBuffer() { throw new UnsupportedOperationException(); } @Override public Map getObject(int index) { if (isSet(index) == 0) { return null; } else { return super.getObject(index); } } @Override public int hashCode(int index) { return hashCode(index, null); } @Override public int hashCode(int index, ArrowBufHasher hasher) { if (isSet(index) == 0) { return ArrowBufPointer.NULL_HASH_CODE; } else { return super.hashCode(index, hasher); } } @Override public void get(int index, ComplexHolder holder) { holder.isSet = isSet(index); if (holder.isSet == 0) { holder.reader = null; return; } super.get(index, holder); } /** Return the number of null values in the vector. */ @Override public int getNullCount() { return BitVectorHelper.getNullCount(validityBuffer, valueCount); } /** Returns true if the value at the provided index is null. */ @Override public boolean isNull(int index) { return isSet(index) == 0; } /** Returns true the value at the given index is set (i.e. not null). */ 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; } /** * Marks the value at index as being set. Reallocates the validity buffer if index is larger than * current capacity. */ public void setIndexDefined(int index) { while (index >= getValidityBufferValueCapacity()) { /* realloc the inner buffers if needed */ reallocValidityBuffer(); } BitVectorHelper.setBit(validityBuffer, index); } /** Marks the value at index as null/not set. */ @Override public void setNull(int index) { while (index >= getValidityBufferValueCapacity()) { /* realloc the inner buffers if needed */ reallocValidityBuffer(); } BitVectorHelper.unsetBit(validityBuffer, index); } @Override public void setValueCount(int valueCount) { Preconditions.checkArgument(valueCount >= 0); while (valueCount > getValidityBufferValueCapacity()) { /* realloc the inner buffers if needed */ reallocValidityBuffer(); } super.setValueCount(valueCount); this.valueCount = valueCount; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/VectorWithOrdinal.java000066400000000000000000000023731515451615200334260ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import org.apache.arrow.vector.ValueVector; /** * Tuple of a {@link ValueVector} and an index into a data structure containing the {@link * ValueVector}. Useful for composite types to determine the index of a child. */ public class VectorWithOrdinal { public final ValueVector vector; public final int ordinal; public VectorWithOrdinal(ValueVector v, int ordinal) { this.vector = v; this.ordinal = ordinal; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/000077500000000000000000000000001515451615200301105ustar00rootroot00000000000000AbstractBaseReader.java000066400000000000000000000066641515451615200343710ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.Iterator; import org.apache.arrow.vector.complex.reader.FieldReader; 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.holders.DenseUnionHolder; import org.apache.arrow.vector.holders.UnionHolder; /** * Base class providing common functionality for {@link FieldReader} implementations. * *

    This includes tracking the current index and throwing implementations of optional methods. */ abstract class AbstractBaseReader implements FieldReader { static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(AbstractBaseReader.class); private int index; public AbstractBaseReader() { super(); } @Override public int getPosition() { return index; } @Override public void setPosition(int index) { this.index = index; } protected int idx() { return index; } @Override public void reset() { index = 0; } @Override public Iterator iterator() { throw new IllegalStateException("The current reader doesn't support reading as a map."); } @Override public boolean next() { throw new IllegalStateException("The current reader doesn't support getting next information."); } @Override public int size() { throw new IllegalStateException("The current reader doesn't support getting size information."); } @Override public void read(UnionHolder holder) { holder.reader = this; holder.isSet = this.isSet() ? 1 : 0; } @Override public void read(int index, UnionHolder holder) { throw new IllegalStateException("The current reader doesn't support reading union type"); } @Override public void copyAsValue(UnionWriter writer) { throw new IllegalStateException("The current reader doesn't support reading union type"); } @Override public void read(DenseUnionHolder holder) { holder.reader = this; holder.isSet = this.isSet() ? 1 : 0; } @Override public void read(int index, DenseUnionHolder holder) { throw new IllegalStateException("The current reader doesn't support reading dense union type"); } @Override public void copyAsValue(DenseUnionWriter writer) { throw new IllegalStateException("The current reader doesn't support reading dense union type"); } @Override public void copyAsValue(ListWriter writer) { ComplexCopier.copy(this, (FieldWriter) writer); } @Override public void copyAsValue(MapWriter writer) { ComplexCopier.copy(this, (FieldWriter) writer); } } AbstractBaseWriter.java000066400000000000000000000030511515451615200344260ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex.writer.FieldWriter; /** * Base class providing common functionality for {@link FieldWriter} implementations. * *

    Currently this only includes index tracking. */ abstract class AbstractBaseWriter implements FieldWriter { // private static final org.slf4j.Logger logger = // org.slf4j.LoggerFactory.getLogger(AbstractBaseWriter.class); private int index; @Override public String toString() { return super.toString() + "[index = " + index + "]"; } int idx() { return index; } @Override public int getPosition() { return index; } @Override public void setPosition(int index) { this.index = index; } @Override public void end() {} } AbstractExtensionTypeWriter.java000066400000000000000000000034741515451615200364030ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } ComplexWriterImpl.java000066400000000000000000000170771515451615200343360ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.Preconditions; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.ListViewVector; import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.NonNullableStructVector; import org.apache.arrow.vector.complex.StateTool; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.writer.BaseWriter.ComplexWriter; import org.apache.arrow.vector.types.pojo.Field; /** Concrete implementation of {@link ComplexWriter}. */ public class ComplexWriterImpl extends AbstractFieldWriter implements ComplexWriter { private NullableStructWriter structRoot; private UnionListWriter listRoot; private UnionListViewWriter listViewRoot; private UnionMapWriter mapRoot; private final NonNullableStructVector container; Mode mode = Mode.INIT; private final String name; private final boolean unionEnabled; private final NullableStructWriterFactory nullableStructWriterFactory; private enum Mode { INIT, STRUCT, LIST, LISTVIEW, MAP } /** * Constructs a new instance. * * @param name The name of the writer (for tracking). * @param container A container for the data field to be written. * @param unionEnabled Unused. * @param caseSensitive Whether field names are case-sensitive (if false field names will be * lowercase. */ public ComplexWriterImpl( String name, NonNullableStructVector container, boolean unionEnabled, boolean caseSensitive) { this.name = name; this.container = container; this.unionEnabled = unionEnabled; nullableStructWriterFactory = caseSensitive ? NullableStructWriterFactory.getNullableCaseSensitiveStructWriterFactoryInstance() : NullableStructWriterFactory.getNullableStructWriterFactoryInstance(); } public ComplexWriterImpl(String name, NonNullableStructVector container, boolean unionEnabled) { this(name, container, unionEnabled, false); } public ComplexWriterImpl(String name, NonNullableStructVector container) { this(name, container, false); } @Override public Field getField() { return container.getField(); } @Override public int getValueCapacity() { return container.getValueCapacity(); } private void check(Mode... modes) { StateTool.check(mode, modes); } @Override public void reset() { setPosition(0); } @Override public void close() throws Exception { clear(); structRoot.close(); if (listRoot != null) { listRoot.close(); } if (listViewRoot != null) { listViewRoot.close(); } } @Override public void clear() { switch (mode) { case STRUCT: structRoot.clear(); break; case LIST: listRoot.clear(); break; case LISTVIEW: listViewRoot.clear(); break; case MAP: mapRoot.clear(); break; default: break; } } @Override public void setValueCount(int count) { switch (mode) { case STRUCT: structRoot.setValueCount(count); break; case LIST: listRoot.setValueCount(count); break; case LISTVIEW: listViewRoot.setValueCount(count); break; case MAP: mapRoot.setValueCount(count); break; default: break; } } @Override public void setPosition(int index) { super.setPosition(index); switch (mode) { case STRUCT: structRoot.setPosition(index); break; case LIST: listRoot.setPosition(index); break; case LISTVIEW: listViewRoot.setPosition(index); break; case MAP: mapRoot.setPosition(index); break; default: break; } } /** * Returns a StructWriter, initializing it necessary from the constructor this instance was * constructed with. */ public StructWriter directStruct() { Preconditions.checkArgument(name == null); switch (mode) { case INIT: structRoot = nullableStructWriterFactory.build((StructVector) container); structRoot.setPosition(idx()); mode = Mode.STRUCT; break; case STRUCT: break; default: check(Mode.INIT, Mode.STRUCT); } return structRoot; } @Override public StructWriter rootAsStruct() { switch (mode) { case INIT: // TODO allow dictionaries in complex types StructVector struct = container.addOrGetStruct(name); structRoot = nullableStructWriterFactory.build(struct); structRoot.setPosition(idx()); mode = Mode.STRUCT; break; case STRUCT: break; default: check(Mode.INIT, Mode.STRUCT); } return structRoot; } @Override public void allocate() { if (structRoot != null) { structRoot.allocate(); } else if (listRoot != null) { listRoot.allocate(); } } @Override public ListWriter rootAsList() { switch (mode) { case INIT: int vectorCount = container.size(); // TODO allow dictionaries in complex types ListVector listVector = container.addOrGetList(name); if (container.size() > vectorCount) { listVector.allocateNew(); } listRoot = new UnionListWriter(listVector, nullableStructWriterFactory); listRoot.setPosition(idx()); mode = Mode.LIST; break; case LIST: break; default: check(Mode.INIT, Mode.STRUCT); } return listRoot; } @Override public ListWriter rootAsListView() { switch (mode) { case INIT: int vectorCount = container.size(); // TODO allow dictionaries in complex types ListViewVector listVector = container.addOrGetListView(name); if (container.size() > vectorCount) { listVector.allocateNew(); } listViewRoot = new UnionListViewWriter(listVector, nullableStructWriterFactory); listViewRoot.setPosition(idx()); mode = Mode.LISTVIEW; break; case LISTVIEW: break; default: check(Mode.INIT, Mode.STRUCT); } return listViewRoot; } @Override public MapWriter rootAsMap(boolean keysSorted) { switch (mode) { case INIT: int vectorCount = container.size(); // TODO allow dictionaries in complex types MapVector mapVector = container.addOrGetMap(name, keysSorted); if (container.size() > vectorCount) { mapVector.allocateNew(); } mapRoot = new UnionMapWriter(mapVector); mapRoot.setPosition(idx()); mode = Mode.MAP; break; case MAP: break; default: check(Mode.INIT, Mode.STRUCT); } return mapRoot; } } NullableStructReaderImpl.java000066400000000000000000000040361515451615200356070ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex.NonNullableStructVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter; import org.apache.arrow.vector.types.pojo.Field; /** * An {@link org.apache.arrow.vector.complex.reader.FieldReader} for reading nullable struct * vectors. */ public class NullableStructReaderImpl extends SingleStructReaderImpl { private StructVector nullableStructVector; public NullableStructReaderImpl(NonNullableStructVector vector) { super(vector); this.nullableStructVector = (StructVector) vector; } @Override public Field getField() { return nullableStructVector.getField(); } @Override public void copyAsValue(StructWriter writer) { NullableStructWriter impl = (NullableStructWriter) writer; impl.container.copyFromSafe(idx(), impl.idx(), nullableStructVector); } @Override public void copyAsField(String name, StructWriter writer) { NullableStructWriter impl = (NullableStructWriter) writer.struct(name); impl.container.copyFromSafe(idx(), impl.idx(), nullableStructVector); } @Override public boolean isSet() { return !nullableStructVector.isNull(idx()); } } NullableStructWriterFactory.java000066400000000000000000000037541515451615200363750ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex.StructVector; /** * A factory for {@link NullableStructWriter} instances. The factory allows for configuring if field * names should be considered case-sensitive. */ public class NullableStructWriterFactory { private final boolean caseSensitive; private static final NullableStructWriterFactory nullableStructWriterFactory = new NullableStructWriterFactory(false); private static final NullableStructWriterFactory nullableCaseSensitiveWriterFactory = new NullableStructWriterFactory(true); public NullableStructWriterFactory(boolean caseSensitive) { this.caseSensitive = caseSensitive; } /** Creates a new instance. */ public NullableStructWriter build(StructVector container) { return this.caseSensitive ? new NullableCaseSensitiveStructWriter(container) : new NullableStructWriter(container); } public static NullableStructWriterFactory getNullableStructWriterFactoryInstance() { return nullableStructWriterFactory; } public static NullableStructWriterFactory getNullableCaseSensitiveStructWriterFactoryInstance() { return nullableCaseSensitiveWriterFactory; } } NullableUuidHolderReaderImpl.java000066400000000000000000000106131515451615200363650ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } } SingleListReaderImpl.java000066400000000000000000000052231515451615200347200ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex.AbstractContainerVector; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.complex.writer.BaseWriter.ListWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter; import org.apache.arrow.vector.types.Types.MinorType; /** An implementation of {@link AbstractFieldReader} for lists vectors. */ @SuppressWarnings("unused") public class SingleListReaderImpl extends AbstractFieldReader { private final String name; private final AbstractContainerVector container; private FieldReader reader; /** * Constructs a new instance. * * @param name The name of field to read in container. * @param container The container holding a list. */ public SingleListReaderImpl(String name, AbstractContainerVector container) { super(); this.name = name; this.container = container; } @Override public void setPosition(int index) { super.setPosition(index); if (reader != null) { reader.setPosition(index); } } @Override public Object readObject() { return reader.readObject(); } @Override public FieldReader reader() { if (reader == null) { reader = container.getChild(name).getReader(); setPosition(idx()); } return reader; } @Override public MinorType getMinorType() { return MinorType.LIST; } @Override public boolean isSet() { return false; } @Override public void copyAsValue(ListWriter writer) { throw new UnsupportedOperationException( "Generic list copying not yet supported. Please resolve to typed list."); } @Override public void copyAsField(String name, StructWriter writer) { throw new UnsupportedOperationException( "Generic list copying not yet supported. Please resolve to typed list."); } } SingleStructReaderImpl.java000066400000000000000000000061121515451615200352670ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.HashMap; import java.util.Map; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.NonNullableStructVector; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; /** * {@link FieldReader} for a single {@link org.apache.arrow.vector.complex.NonNullableStructVector}. */ @SuppressWarnings("unused") public class SingleStructReaderImpl extends AbstractFieldReader { private final NonNullableStructVector vector; private final Map fields = new HashMap<>(); public SingleStructReaderImpl(NonNullableStructVector vector) { this.vector = vector; } private void setChildrenPosition(int index) { for (FieldReader r : fields.values()) { r.setPosition(index); } } @Override public Field getField() { return vector.getField(); } @Override public FieldReader reader(String name) { FieldReader reader = fields.get(name); if (reader == null) { ValueVector child = vector.getChild(name); if (child == null) { reader = NullReader.INSTANCE; } else { reader = child.getReader(); } fields.put(name, reader); reader.setPosition(idx()); } return reader; } @Override public void setPosition(int index) { super.setPosition(index); for (FieldReader r : fields.values()) { r.setPosition(index); } } @Override public Object readObject() { return vector.getObject(idx()); } @Override public MinorType getMinorType() { return MinorType.STRUCT; } @Override public boolean isSet() { return true; } @Override public java.util.Iterator iterator() { return vector.fieldNameIterator(); } @Override public void copyAsValue(StructWriter writer) { SingleStructWriter impl = (SingleStructWriter) writer; impl.container.copyFromSafe(idx(), impl.idx(), vector); } @Override public void copyAsField(String name, StructWriter writer) { SingleStructWriter impl = (SingleStructWriter) writer.struct(name); impl.container.copyFromSafe(idx(), impl.idx(), vector); } } StructOrListWriterImpl.java000066400000000000000000000103441515451615200353360ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex.writer.BaseWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.StructOrListWriter; import org.apache.arrow.vector.complex.writer.BigIntWriter; import org.apache.arrow.vector.complex.writer.BitWriter; import org.apache.arrow.vector.complex.writer.Float4Writer; import org.apache.arrow.vector.complex.writer.Float8Writer; import org.apache.arrow.vector.complex.writer.IntWriter; import org.apache.arrow.vector.complex.writer.VarBinaryWriter; import org.apache.arrow.vector.complex.writer.VarCharWriter; /** Concrete implementation of {@link StructOrListWriter}. */ public class StructOrListWriterImpl implements StructOrListWriter { public final BaseWriter.StructWriter struct; public final BaseWriter.ListWriter list; /** * Constructs a new instance using a {@link BaseWriter.StructWriter} (instead of an {@link * BaseWriter.ListWriter}). */ public StructOrListWriterImpl(final BaseWriter.StructWriter writer) { this.struct = writer; this.list = null; } /** * Constructs a new instance using a {@link BaseWriter.ListWriter} (instead of a {@link * BaseWriter.StructWriter}). */ public StructOrListWriterImpl(final BaseWriter.ListWriter writer) { this.struct = null; this.list = writer; } /** Start writing to either the list or the struct. */ public void start() { if (struct != null) { struct.start(); } else { list.startList(); } } /** Finish writing to the list or struct. */ public void end() { if (struct != null) { struct.end(); } else { list.endList(); } } /** Creates a new writer for a struct with the given name. */ public StructOrListWriter struct(final String name) { assert struct != null; return new StructOrListWriterImpl(struct.struct(name)); } /** * Creates a new writer for a list of structs. * * @param name Unused. * @deprecated use {@link #listOfStruct(String)} instead. */ @Deprecated public StructOrListWriter listoftstruct(final String name) { return listOfStruct(name); } /** * Creates a new writer for a list of structs. * * @param name Unused. */ public StructOrListWriter listOfStruct(final String name) { assert list != null; return new StructOrListWriterImpl(list.struct()); } public StructOrListWriter list(final String name) { assert struct != null; return new StructOrListWriterImpl(struct.list(name)); } public boolean isStructWriter() { return struct != null; } public boolean isListWriter() { return list != null; } public VarCharWriter varChar(final String name) { return (struct != null) ? struct.varChar(name) : list.varChar(); } public IntWriter integer(final String name) { return (struct != null) ? struct.integer(name) : list.integer(); } public BigIntWriter bigInt(final String name) { return (struct != null) ? struct.bigInt(name) : list.bigInt(); } public Float4Writer float4(final String name) { return (struct != null) ? struct.float4(name) : list.float4(); } public Float8Writer float8(final String name) { return (struct != null) ? struct.float8(name) : list.float8(); } public BitWriter bit(final String name) { return (struct != null) ? struct.bit(name) : list.bit(); } public VarBinaryWriter binary(final String name) { return (struct != null) ? struct.varBinary(name) : list.varBinary(); } } UnionExtensionWriter.java000066400000000000000000000041541515451615200350620ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } UnionFixedSizeListReader.java000066400000000000000000000056271515451615200355700ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ValueVector; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.complex.writer.BaseWriter.ListWriter; import org.apache.arrow.vector.complex.writer.FieldWriter; import org.apache.arrow.vector.holders.UnionHolder; import org.apache.arrow.vector.types.Types.MinorType; /** Reader for fixed size list vectors. */ public class UnionFixedSizeListReader extends AbstractFieldReader { private final FixedSizeListVector vector; private final ValueVector data; private final int listSize; private int currentOffset; /** Constructs a new instance that reads data in vector. */ public UnionFixedSizeListReader(FixedSizeListVector vector) { this.vector = vector; this.data = vector.getDataVector(); this.listSize = vector.getListSize(); } @Override public boolean isSet() { return !vector.isNull(idx()); } @Override public FieldReader reader() { return data.getReader(); } @Override public Object readObject() { return vector.getObject(idx()); } @Override public MinorType getMinorType() { return vector.getMinorType(); } @Override public void setPosition(int index) { super.setPosition(index); data.getReader().setPosition(index * listSize); currentOffset = 0; } @Override public void read(int index, UnionHolder holder) { setPosition(idx()); for (int i = -1; i < index; i++) { if (!next()) { throw new IndexOutOfBoundsException("Requested " + index + ", size " + listSize); } } holder.reader = data.getReader(); holder.isSet = vector.isNull(idx()) ? 0 : 1; } @Override public int size() { return listSize; } @Override public boolean next() { if (currentOffset < listSize) { data.getReader().setPosition(idx() * listSize + currentOffset++); return true; } else { return false; } } @Override public void copyAsValue(ListWriter writer) { ComplexCopier.copy(this, (FieldWriter) writer); } } UnionLargeListReader.java000066400000000000000000000061761515451615200347300ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.LargeListVector; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.complex.writer.FieldWriter; import org.apache.arrow.vector.holders.UnionHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; /** {@link FieldReader} for list of union types. */ public class UnionLargeListReader extends AbstractFieldReader { private LargeListVector vector; private ValueVector data; private static final long OFFSET_WIDTH = 8L; public UnionLargeListReader(LargeListVector vector) { this.vector = vector; this.data = vector.getDataVector(); } @Override public Field getField() { return vector.getField(); } @Override public boolean isSet() { return !vector.isNull(idx()); } private long currentOffset; private long maxOffset; @Override public void setPosition(int index) { super.setPosition(index); currentOffset = vector.getOffsetBuffer().getLong((long) index * OFFSET_WIDTH) - 1; maxOffset = vector.getOffsetBuffer().getLong(((long) index + 1L) * OFFSET_WIDTH); } @Override public FieldReader reader() { return data.getReader(); } @Override public Object readObject() { return vector.getObject(idx()); } @Override public MinorType getMinorType() { return MinorType.LARGELIST; } @Override public void read(int index, UnionHolder holder) { setPosition(index); for (int i = -1; i < index; i++) { next(); } holder.reader = data.getReader(); holder.isSet = data.getReader().isSet() ? 1 : 0; } @Override public int size() { int size = checkedCastToInt(maxOffset - currentOffset - 1); // todo revisit when int64 vectors are done return size < 0 ? 0 : size; } @Override public boolean next() { if (currentOffset + 1 < maxOffset) { data.getReader() .setPosition( checkedCastToInt(++currentOffset)); // todo revisit when int64 vectors are done return true; } else { return false; } } public void copyAsValue(UnionLargeListWriter writer) { ComplexCopier.copy(this, (FieldWriter) writer); } } UnionLargeListViewReader.java000066400000000000000000000067121515451615200355570ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.BaseLargeRepeatedValueViewVector; import org.apache.arrow.vector.complex.LargeListViewVector; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.UnionHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; /** {@link FieldReader} for largeListView of union types. */ public class UnionLargeListViewReader extends AbstractFieldReader { private final LargeListViewVector vector; private final ValueVector data; private int currentOffset; private int size; /** * Constructor for UnionLargeListViewReader. * * @param vector the vector to read from */ public UnionLargeListViewReader(LargeListViewVector vector) { this.vector = vector; this.data = vector.getDataVector(); } @Override public Field getField() { return vector.getField(); } @Override public boolean isSet() { return !vector.isNull(idx()); } @Override public void setPosition(int index) { super.setPosition(index); if (vector.getOffsetBuffer().capacity() == 0) { currentOffset = 0; size = 0; } else { currentOffset = vector .getOffsetBuffer() .getInt(index * (long) BaseLargeRepeatedValueViewVector.OFFSET_WIDTH); size = vector.getSizeBuffer().getInt(index * (long) BaseLargeRepeatedValueViewVector.SIZE_WIDTH); } } @Override public FieldReader reader() { return data.getReader(); } @Override public Object readObject() { return vector.getObject(idx()); } @Override public MinorType getMinorType() { return MinorType.LISTVIEW; } @Override public void read(int index, UnionHolder holder) { setPosition(idx()); for (int i = -1; i < index; i++) { next(); } holder.reader = data.getReader(); holder.isSet = data.getReader().isSet() ? 1 : 0; } @Override public int size() { return Math.max(size, 0); } @Override public boolean next() { // Here, the currentOffSet keeps track of the current position in the vector inside the list at // set position. // And, size keeps track of the elements count in the list, so to make sure we traverse // the full list, we need to check if the currentOffset is less than the currentOffset + size if (currentOffset < currentOffset + size) { data.getReader().setPosition(checkedCastToInt(currentOffset++)); return true; } else { return false; } } } UnionListReader.java000066400000000000000000000060551515451615200337510ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ValueVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.complex.writer.BaseWriter.ListWriter; import org.apache.arrow.vector.complex.writer.FieldWriter; import org.apache.arrow.vector.holders.UnionHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; /** {@link FieldReader} for list of union types. */ public class UnionListReader extends AbstractFieldReader { private ListVector vector; private ValueVector data; private static final int OFFSET_WIDTH = 4; public UnionListReader(ListVector vector) { this.vector = vector; this.data = vector.getDataVector(); } @Override public Field getField() { return vector.getField(); } @Override public boolean isSet() { return !vector.isNull(idx()); } private int currentOffset; private int maxOffset; @Override public void setPosition(int index) { super.setPosition(index); if (vector.getOffsetBuffer().capacity() == 0) { currentOffset = 0; maxOffset = 0; } else { currentOffset = vector.getOffsetBuffer().getInt(index * (long) OFFSET_WIDTH) - 1; maxOffset = vector.getOffsetBuffer().getInt((index + 1) * (long) OFFSET_WIDTH); } } @Override public FieldReader reader() { return data.getReader(); } @Override public Object readObject() { return vector.getObject(idx()); } @Override public MinorType getMinorType() { return MinorType.LIST; } @Override public void read(int index, UnionHolder holder) { setPosition(idx()); for (int i = -1; i < index; i++) { next(); } holder.reader = data.getReader(); holder.isSet = data.getReader().isSet() ? 1 : 0; } @Override public int size() { int size = maxOffset - currentOffset - 1; return size < 0 ? 0 : size; } @Override public boolean next() { if (currentOffset + 1 < maxOffset) { data.getReader().setPosition(++currentOffset); return true; } else { return false; } } @Override public void copyAsValue(ListWriter writer) { ComplexCopier.copy(this, (FieldWriter) writer); } } UnionListViewReader.java000066400000000000000000000064201515451615200346000ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ValueVector; import org.apache.arrow.vector.complex.BaseRepeatedValueViewVector; import org.apache.arrow.vector.complex.ListViewVector; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.UnionHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; /** {@link FieldReader} for listview of union types. */ public class UnionListViewReader extends AbstractFieldReader { private final ListViewVector vector; private final ValueVector data; private int currentOffset; private int size; /** * Constructor for UnionListViewReader. * * @param vector the vector to read from */ public UnionListViewReader(ListViewVector vector) { this.vector = vector; this.data = vector.getDataVector(); } @Override public Field getField() { return vector.getField(); } @Override public boolean isSet() { return !vector.isNull(idx()); } @Override public void setPosition(int index) { super.setPosition(index); if (vector.getOffsetBuffer().capacity() == 0) { currentOffset = 0; size = 0; } else { currentOffset = vector.getOffsetBuffer().getInt(index * (long) BaseRepeatedValueViewVector.OFFSET_WIDTH); size = vector.getSizeBuffer().getInt(index * (long) BaseRepeatedValueViewVector.SIZE_WIDTH); } } @Override public FieldReader reader() { return data.getReader(); } @Override public Object readObject() { return vector.getObject(idx()); } @Override public MinorType getMinorType() { return MinorType.LISTVIEW; } @Override public void read(int index, UnionHolder holder) { setPosition(idx()); for (int i = -1; i < index; i++) { next(); } holder.reader = data.getReader(); holder.isSet = data.getReader().isSet() ? 1 : 0; } @Override public int size() { return Math.max(size, 0); } @Override public boolean next() { // Here, the currentOffSet keeps track of the current position in the vector inside the list at // set position. // And, size keeps track of the elements count in the list, so to make sure we traverse // the full list, we need to check if the currentOffset is less than the currentOffset + size if (currentOffset < currentOffset + size) { data.getReader().setPosition(currentOffset++); return true; } else { return false; } } } UnionMapReader.java000066400000000000000000000041321515451615200335450ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex.MapVector; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.types.Types.MinorType; /** Reader for a MapVector. */ public class UnionMapReader extends UnionListReader { private String keyName = MapVector.KEY_NAME; private String valueName = MapVector.VALUE_NAME; /** * Construct a new reader for the given vector. * * @param vector Vector to read from. */ public UnionMapReader(MapVector vector) { super(vector); } /** * Set the key, value field names to read. * * @param key Field name for key. * @param value Field name for value. */ public void setKeyValueNames(String key, String value) { keyName = key; valueName = value; } /** * Start reading a key from the map entry. * * @return reader that can be used to read the key. */ public FieldReader key() { return reader().reader(keyName); } /** * Start reading a value element from the map entry. * * @return reader that can be used to read the value. */ public FieldReader value() { return reader().reader(valueName); } /** Return the MinorType of the reader as MAP. */ @Override public MinorType getMinorType() { return MinorType.MAP; } } UuidReaderImpl.java000066400000000000000000000055261515451615200335570ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } UuidWriterImpl.java000066400000000000000000000057031515451615200336260ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/reader/000077500000000000000000000000001515451615200304115ustar00rootroot00000000000000ExtensionReader.java000066400000000000000000000027541515451615200343040ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/reader/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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(); } FieldReader.java000066400000000000000000000033471515451615200333520ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/reader/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex.reader.BaseReader.ListReader; import org.apache.arrow.vector.complex.reader.BaseReader.MapReader; import org.apache.arrow.vector.complex.reader.BaseReader.RepeatedListReader; import org.apache.arrow.vector.complex.reader.BaseReader.RepeatedMapReader; import org.apache.arrow.vector.complex.reader.BaseReader.RepeatedStructReader; import org.apache.arrow.vector.complex.reader.BaseReader.ScalarReader; import org.apache.arrow.vector.complex.reader.BaseReader.StructReader; /** * Composite of all Reader types (e.g. {@link StructReader}, {@link ScalarReader}, etc). Each reader * type is in essence a way of iterating over a {@link org.apache.arrow.vector.ValueVector}. */ public interface FieldReader extends StructReader, ListReader, MapReader, ScalarReader, RepeatedStructReader, RepeatedListReader, RepeatedMapReader {} apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/writer/000077500000000000000000000000001515451615200304635ustar00rootroot00000000000000FieldWriter.java000066400000000000000000000027531515451615200334760ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/complex/writer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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; import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter; /** * 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, ExtensionWriter { void allocate(); void clear(); } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/compression/000077500000000000000000000000001515451615200300415ustar00rootroot00000000000000AbstractCompressionCodec.java000066400000000000000000000110641515451615200355520ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/compression/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compression; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.MemoryUtil; import org.apache.arrow.util.Preconditions; /** * The base class for concrete compression codecs, providing common logic for all compression * codecs. */ public abstract class AbstractCompressionCodec implements CompressionCodec { @Override public ArrowBuf compress(BufferAllocator allocator, ArrowBuf uncompressedBuffer) { if (uncompressedBuffer.writerIndex() == 0L) { // shortcut for empty buffer ArrowBuf compressedBuffer = allocator.buffer(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH); compressedBuffer.setLong(0, 0); compressedBuffer.writerIndex(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH); uncompressedBuffer.close(); return compressedBuffer; } 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 compressedBuffer.close(); // XXX: this makes a copy of uncompressedBuffer compressedBuffer = CompressionUtil.packageRawBuffer(allocator, uncompressedBuffer); } else { writeUncompressedLength(compressedBuffer, uncompressedLength); } uncompressedBuffer.close(); return compressedBuffer; } @Override public ArrowBuf decompress(BufferAllocator allocator, ArrowBuf compressedBuffer) { Preconditions.checkArgument( compressedBuffer.writerIndex() >= CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH, "Not enough data to decompress."); long decompressedLength = readUncompressedLength(compressedBuffer); if (decompressedLength == 0L) { // shortcut for empty buffer compressedBuffer.close(); return allocator.getEmpty(); } if (decompressedLength == CompressionUtil.NO_COMPRESSION_LENGTH) { // no compression return CompressionUtil.extractUncompressedBuffer(compressedBuffer); } ArrowBuf decompressedBuffer = doDecompress(allocator, compressedBuffer); compressedBuffer.close(); return decompressedBuffer; } protected void writeUncompressedLength(ArrowBuf compressedBuffer, long uncompressedLength) { if (!MemoryUtil.LITTLE_ENDIAN) { uncompressedLength = Long.reverseBytes(uncompressedLength); } // first 8 bytes reserved for uncompressed length, according to the specification compressedBuffer.setLong(0, uncompressedLength); } protected long readUncompressedLength(ArrowBuf compressedBuffer) { long decompressedLength = compressedBuffer.getLong(0); if (!MemoryUtil.LITTLE_ENDIAN) { decompressedLength = Long.reverseBytes(decompressedLength); } return decompressedLength; } /** * The method that actually performs the data compression. The layout of the returned compressed * buffer is the compressed data, plus 8 bytes reserved at the beginning of the buffer for the * uncompressed data size. * *

    Please note that this method is not responsible for releasing the uncompressed buffer. */ protected abstract ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf uncompressedBuffer); /** * The method that actually performs the data decompression. The layout of the compressed buffer * is the compressed data, plus 8 bytes at the beginning of the buffer storing the uncompressed * data size. * *

    Please note that this method is not responsible for releasing the compressed buffer. */ protected abstract ArrowBuf doDecompress(BufferAllocator allocator, ArrowBuf compressedBuffer); } CompressionCodec.java000066400000000000000000000102111515451615200340570ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/compression/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compression; import java.util.EnumMap; import java.util.Map; import java.util.ServiceLoader; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; /** The codec for compression/decompression. */ public interface CompressionCodec { /** * Compress a buffer. * * @param allocator the allocator for allocating memory for compressed buffer. * @param uncompressedBuffer the buffer to compress. Implementation of this method should take * care of releasing this buffer. * @return the compressed buffer */ ArrowBuf compress(BufferAllocator allocator, ArrowBuf uncompressedBuffer); /** * Decompress a buffer. * * @param allocator the allocator for allocating memory for decompressed buffer. * @param compressedBuffer the buffer to be decompressed. Implementation of this method should * take care of releasing this buffer. * @return the decompressed buffer. */ ArrowBuf decompress(BufferAllocator allocator, ArrowBuf compressedBuffer); /** * Gets the type of the codec. * * @return the type of the codec. */ CompressionUtil.CodecType getCodecType(); /** Factory to create compression codec. */ interface Factory { /** * This combines all the available factories registered as service providers in the module path. * For each {@link CompressionUtil.CodecType compression codec type}, it will use whatever * factory supports it, i.e. doesn't throw on `createCodec(type)`. If multiple factories * registered as service providers support the same codec type, the first one encountered while * iterating over the {@link ServiceLoader} will be selected. A codec type that is not supported * by any registered service provider will fall back to {@link * NoCompressionCodec.Factory#INSTANCE} for backwards compatibility. */ Factory INSTANCE = bestEffort(); /** Creates the codec based on the codec type. */ CompressionCodec createCodec(CompressionUtil.CodecType codecType); /** Creates the codec based on the codec type and compression level. */ CompressionCodec createCodec(CompressionUtil.CodecType codecType, int compressionLevel); private static Factory bestEffort() { final ServiceLoader serviceLoader = ServiceLoader.load(Factory.class); final Map factories = new EnumMap<>(CompressionUtil.CodecType.class); for (Factory factory : serviceLoader) { for (CompressionUtil.CodecType codecType : CompressionUtil.CodecType.values()) { try { factory.createCodec(codecType); // will throw if not supported factories.putIfAbsent(codecType, factory); } catch (Throwable ignored) { } } } final Factory fallback = NoCompressionCodec.Factory.INSTANCE; return new Factory() { @Override public CompressionCodec createCodec(CompressionUtil.CodecType codecType) { return factories.getOrDefault(codecType, fallback).createCodec(codecType); } @Override public CompressionCodec createCodec( CompressionUtil.CodecType codecType, int compressionLevel) { return factories .getOrDefault(codecType, fallback) .createCodec(codecType, compressionLevel); } }; } } } CompressionUtil.java000066400000000000000000000067031515451615200337720ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/compression/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compression; import org.apache.arrow.flatbuf.BodyCompressionMethod; import org.apache.arrow.flatbuf.CompressionType; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.ipc.message.ArrowBodyCompression; /** Utilities for data compression/decompression. */ public class CompressionUtil { /** * Compression codec types corresponding to flat buffer implementation in {@link CompressionType}. */ public enum CodecType { NO_COMPRESSION(NoCompressionCodec.COMPRESSION_TYPE), LZ4_FRAME(org.apache.arrow.flatbuf.CompressionType.LZ4_FRAME), ZSTD(org.apache.arrow.flatbuf.CompressionType.ZSTD); private final byte type; CodecType(byte type) { this.type = type; } public byte getType() { return type; } /** Gets the codec type from the compression type defined in {@link CompressionType}. */ public static CodecType fromCompressionType(byte type) { for (CodecType codecType : values()) { if (codecType.type == type) { return codecType; } } return NO_COMPRESSION; } } public static final long SIZE_OF_UNCOMPRESSED_LENGTH = 8L; /** * Special flag to indicate no compression. (e.g. when the compressed buffer has a larger size.) */ public static final long NO_COMPRESSION_LENGTH = -1L; private CompressionUtil() {} /** * Creates the {@link ArrowBodyCompression} object, given the {@link CompressionCodec}. The * implementation of this method should depend on the values of {@link * org.apache.arrow.flatbuf.CompressionType#names}. */ public static ArrowBodyCompression createBodyCompression(CompressionCodec codec) { return new ArrowBodyCompression(codec.getCodecType().getType(), BodyCompressionMethod.BUFFER); } /** Process compression by compressing the buffer as is. */ public static ArrowBuf packageRawBuffer(BufferAllocator allocator, ArrowBuf inputBuffer) { ArrowBuf compressedBuffer = allocator.buffer(SIZE_OF_UNCOMPRESSED_LENGTH + inputBuffer.writerIndex()); compressedBuffer.setLong(0, NO_COMPRESSION_LENGTH); compressedBuffer.setBytes( SIZE_OF_UNCOMPRESSED_LENGTH, inputBuffer, 0, inputBuffer.writerIndex()); compressedBuffer.writerIndex(SIZE_OF_UNCOMPRESSED_LENGTH + inputBuffer.writerIndex()); return compressedBuffer; } /** Process decompression by slicing the buffer that contains the uncompressed bytes. */ public static ArrowBuf extractUncompressedBuffer(ArrowBuf inputBuffer) { return inputBuffer.slice( SIZE_OF_UNCOMPRESSED_LENGTH, inputBuffer.writerIndex() - SIZE_OF_UNCOMPRESSED_LENGTH); } } NoCompressionCodec.java000066400000000000000000000053731515451615200343710ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/compression/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compression; import org.apache.arrow.flatbuf.BodyCompressionMethod; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.ipc.message.ArrowBodyCompression; /** The default compression codec that does no compression. */ public class NoCompressionCodec implements CompressionCodec { public static final NoCompressionCodec INSTANCE = new NoCompressionCodec(); public static final byte COMPRESSION_TYPE = -1; public static final ArrowBodyCompression DEFAULT_BODY_COMPRESSION = new ArrowBodyCompression(COMPRESSION_TYPE, BodyCompressionMethod.BUFFER); private NoCompressionCodec() {} @Override public ArrowBuf compress(BufferAllocator allocator, ArrowBuf uncompressedBuffer) { return uncompressedBuffer; } @Override public ArrowBuf decompress(BufferAllocator allocator, ArrowBuf compressedBuffer) { return compressedBuffer; } @Override public CompressionUtil.CodecType getCodecType() { return CompressionUtil.CodecType.NO_COMPRESSION; } /** The default factory that creates a {@link NoCompressionCodec}. */ public static class Factory implements CompressionCodec.Factory { public static final NoCompressionCodec.Factory INSTANCE = new NoCompressionCodec.Factory(); @Override public CompressionCodec createCodec(CompressionUtil.CodecType codecType) { switch (codecType) { case NO_COMPRESSION: return NoCompressionCodec.INSTANCE; case LZ4_FRAME: case ZSTD: throw new IllegalArgumentException( "Please add arrow-compression module to use CommonsCompressionFactory for " + codecType); default: throw new IllegalArgumentException("Unsupported codec type: " + codecType); } } @Override public CompressionCodec createCodec(CompressionUtil.CodecType codecType, int compressionLevel) { return createCodec(codecType); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/dictionary/000077500000000000000000000000001515451615200276455ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/dictionary/Dictionary.java000066400000000000000000000042371515451615200326230ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dictionary; import java.util.Objects; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.compare.VectorEqualsVisitor; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.DictionaryEncoding; /** * A dictionary (integer to Value mapping) that is used to facilitate dictionary encoding * compression. */ public class Dictionary { private final DictionaryEncoding encoding; private final FieldVector dictionary; public Dictionary(FieldVector dictionary, DictionaryEncoding encoding) { this.dictionary = dictionary; this.encoding = encoding; } public FieldVector getVector() { return dictionary; } public DictionaryEncoding getEncoding() { return encoding; } public ArrowType getVectorType() { return dictionary.getField().getType(); } @Override public String toString() { return "Dictionary " + encoding + " " + dictionary; } @Override public boolean equals(Object o) { if (this == o) { return true; } if (!(o instanceof Dictionary)) { return false; } Dictionary that = (Dictionary) o; return Objects.equals(encoding, that.encoding) && new VectorEqualsVisitor().vectorEquals(that.dictionary, dictionary); } @Override public int hashCode() { return Objects.hash(encoding, dictionary); } } DictionaryEncoder.java000066400000000000000000000201521515451615200340360ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dictionary; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.memory.util.hash.SimpleHasher; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BaseIntVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueVector; 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; /** * Encoder/decoder for Dictionary encoded {@link ValueVector}. Dictionary encoding produces an * integer {@link ValueVector}. Each entry in the Vector is index into the dictionary which can hold * values of any type. */ public class DictionaryEncoder { private final DictionaryHashTable hashTable; private final Dictionary dictionary; private final BufferAllocator allocator; /** Construct an instance. */ public DictionaryEncoder(Dictionary dictionary, BufferAllocator allocator) { this(dictionary, allocator, SimpleHasher.INSTANCE); } /** Construct an instance. */ public DictionaryEncoder( Dictionary dictionary, BufferAllocator allocator, ArrowBufHasher hasher) { this.dictionary = dictionary; this.allocator = allocator; hashTable = new DictionaryHashTable(dictionary.getVector(), hasher); } /** * Dictionary encodes a vector with a provided dictionary. The dictionary must contain all values * in the vector. * * @param vector vector to encode * @param dictionary dictionary used for encoding * @return dictionary encoded vector */ public static ValueVector encode(ValueVector vector, Dictionary dictionary) { DictionaryEncoder encoder = new DictionaryEncoder(dictionary, vector.getAllocator()); return encoder.encode(vector); } /** * Decodes a dictionary encoded array using the provided dictionary. * * @param indices dictionary encoded values, must be int type * @param dictionary dictionary used to decode the values * @return vector with values restored from dictionary */ public static ValueVector decode(ValueVector indices, Dictionary dictionary) { return decode(indices, dictionary, indices.getAllocator()); } /** * Decodes a dictionary encoded array using the provided dictionary. * * @param indices dictionary encoded values, must be int type * @param dictionary dictionary used to decode the values * @param allocator allocator the decoded values use * @return vector with values restored from dictionary */ public static ValueVector decode( ValueVector indices, Dictionary dictionary, BufferAllocator allocator) { int count = indices.getValueCount(); ValueVector dictionaryVector = dictionary.getVector(); int dictionaryCount = dictionaryVector.getValueCount(); // copy the dictionary values into the decoded vector TransferPair transfer = dictionaryVector.getTransferPair(allocator); transfer.getTo().allocateNewSafe(); try { BaseIntVector baseIntVector = (BaseIntVector) indices; retrieveIndexVector(baseIntVector, transfer, dictionaryCount, 0, count); ValueVector decoded = transfer.getTo(); decoded.setValueCount(count); return decoded; } catch (Exception e) { AutoCloseables.close(e, transfer.getTo()); throw e; } } /** * Get the indexType according to the dictionary vector valueCount. * * @param valueCount dictionary vector valueCount. * @return index type. */ @SuppressWarnings("ComparisonOutOfRange") public static ArrowType.Int getIndexType(int valueCount) { Preconditions.checkArgument(valueCount >= 0); if (valueCount <= Byte.MAX_VALUE) { return new ArrowType.Int(8, true); } else if (valueCount <= Character.MAX_VALUE) { return new ArrowType.Int(16, true); } else if (valueCount <= Integer.MAX_VALUE) { // this comparison will always evaluate to true return new ArrowType.Int(32, true); } else { return new ArrowType.Int(64, true); } } /** * Populates indices between start and end with the encoded values of vector. * * @param vector the vector to encode * @param indices the index vector * @param encoding the hash table for encoding * @param start the start index * @param end the end index */ static void buildIndexVector( ValueVector vector, BaseIntVector indices, DictionaryHashTable encoding, int start, int end) { for (int i = start; i < end; i++) { if (!vector.isNull(i)) { // if it's null leave it null // note: this may fail if value was not included in the dictionary int encoded = encoding.getIndex(i, vector); if (encoded == -1) { throw new IllegalArgumentException( "Dictionary encoding not defined for value:" + vector.getObject(i)); } indices.setWithPossibleTruncate(i, encoded); } } } /** * Retrieve values to target vector from index vector. * * @param indices the index vector * @param transfer the {@link TransferPair} to copy dictionary data into target vector. * @param dictionaryCount the value count of dictionary vector. * @param start the start index * @param end the end index */ static void retrieveIndexVector( BaseIntVector indices, TransferPair transfer, int dictionaryCount, int start, int end) { for (int i = start; i < end; i++) { if (!indices.isNull(i)) { int indexAsInt = (int) indices.getValueAsLong(i); if (indexAsInt > dictionaryCount) { throw new IllegalArgumentException( "Provided dictionary does not contain value for index " + indexAsInt); } transfer.copyValueSafe(indexAsInt, i); } } } /** Encodes a vector with the built hash table in this encoder. */ public ValueVector encode(ValueVector vector) { Field valueField = vector.getField(); FieldType indexFieldType = new FieldType( valueField.isNullable(), dictionary.getEncoding().getIndexType(), dictionary.getEncoding(), valueField.getMetadata()); Field indexField = new Field(valueField.getName(), indexFieldType, null); // vector to hold our indices (dictionary encoded values) FieldVector createdVector = indexField.createVector(allocator); if (!(createdVector instanceof BaseIntVector)) { throw new IllegalArgumentException( "Dictionary encoding does not have a valid int type:" + createdVector.getClass()); } BaseIntVector indices = (BaseIntVector) createdVector; indices.allocateNew(); try { buildIndexVector(vector, indices, hashTable, 0, vector.getValueCount()); indices.setValueCount(vector.getValueCount()); return indices; } catch (Exception e) { AutoCloseables.close(e, indices); throw e; } } /** * Decodes a vector with the dictionary in this encoder. * *

    {@link DictionaryEncoder#decode(ValueVector, Dictionary, BufferAllocator)} should be used * instead if only decoding is required as it can avoid building the {@link DictionaryHashTable} * which only makes sense when encoding. */ public ValueVector decode(ValueVector indices) { return decode(indices, dictionary, allocator); } } DictionaryHashTable.java000066400000000000000000000176341515451615200343250ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dictionary; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.memory.util.hash.SimpleHasher; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.compare.Range; import org.apache.arrow.vector.compare.RangeEqualsVisitor; /** * HashTable used for Dictionary encoding. It holds two vectors (the vector to encode and dictionary * vector) It stores the index in dictionary vector and for a given index in encode vector, it could * return dictionary index. */ public class DictionaryHashTable { /** Represents a null value in map. */ static final int NULL_VALUE = -1; /** The default initial capacity - MUST be a power of two. */ static final int DEFAULT_INITIAL_CAPACITY = 1 << 4; /** * The maximum capacity, used if a higher value is implicitly specified by either of the * constructors with arguments. */ static final int MAXIMUM_CAPACITY = 1 << 30; /** The load factor used when none specified in constructor. */ static final float DEFAULT_LOAD_FACTOR = 0.75f; static final DictionaryHashTable.Entry[] EMPTY_TABLE = {}; /** * The table, initialized on first use, and resized as necessary. When allocated, length is always * a power of two. */ transient DictionaryHashTable.Entry[] table = EMPTY_TABLE; /** The number of key-value mappings contained in this map. */ transient int size; /** The next size value at which to resize (capacity * load factor). */ int threshold; /** The load factor for the hash table. */ final float loadFactor; private final ValueVector dictionary; private final ArrowBufHasher hasher; /** Constructs an empty map with the specified initial capacity and load factor. */ public DictionaryHashTable(int initialCapacity, ValueVector dictionary, ArrowBufHasher hasher) { if (initialCapacity < 0) { throw new IllegalArgumentException("Illegal initial capacity: " + initialCapacity); } if (initialCapacity > MAXIMUM_CAPACITY) { initialCapacity = MAXIMUM_CAPACITY; } this.loadFactor = DEFAULT_LOAD_FACTOR; this.threshold = initialCapacity; this.dictionary = dictionary; this.hasher = hasher; // build hash table for (int i = 0; i < this.dictionary.getValueCount(); i++) { put(i); } } public DictionaryHashTable(ValueVector dictionary, ArrowBufHasher hasher) { this(DEFAULT_INITIAL_CAPACITY, dictionary, hasher); } public DictionaryHashTable(ValueVector dictionary) { this(dictionary, SimpleHasher.INSTANCE); } /** Compute the capacity with given threshold and create init table. */ private void inflateTable(int threshold) { int capacity = roundUpToPowerOf2(threshold); this.threshold = (int) Math.min(capacity * loadFactor, MAXIMUM_CAPACITY + 1); table = new DictionaryHashTable.Entry[capacity]; } /** Computes the storage location in an array for the given hashCode. */ static int indexFor(int h, int length) { return h & (length - 1); } /** Returns a power of two size for the given size. */ static final int roundUpToPowerOf2(int size) { int n = size - 1; n |= n >>> 1; n |= n >>> 2; n |= n >>> 4; n |= n >>> 8; n |= n >>> 16; return (n < 0) ? 1 : (n >= MAXIMUM_CAPACITY) ? MAXIMUM_CAPACITY : n + 1; } /** * get the corresponding dictionary index with the given index in vector which to encode. * * @param indexInArray index in vector. * @return dictionary vector index or -1 if no value equals. */ public int getIndex(int indexInArray, ValueVector toEncode) { int hash = toEncode.hashCode(indexInArray, this.hasher); int index = indexFor(hash, table.length); RangeEqualsVisitor equalVisitor = new RangeEqualsVisitor(dictionary, toEncode, null); Range range = new Range(0, 0, 1); for (DictionaryHashTable.Entry e = table[index]; e != null; e = e.next) { if (e.hash == hash) { int dictIndex = e.index; range = range.setRightStart(indexInArray).setLeftStart(dictIndex); if (equalVisitor.rangeEquals(range)) { return dictIndex; } } } return NULL_VALUE; } /** put the index of dictionary vector to build hash table. */ private void put(int indexInDictionary) { if (table == EMPTY_TABLE) { inflateTable(threshold); } int hash = dictionary.hashCode(indexInDictionary, this.hasher); int i = indexFor(hash, table.length); for (DictionaryHashTable.Entry e = table[i]; e != null; e = e.next) { if (e.hash == hash && e.index == indexInDictionary) { // already has this index, return return; } } addEntry(hash, indexInDictionary, i); } /** Create a new Entry at the specific position of table. */ void createEntry(int hash, int index, int bucketIndex) { DictionaryHashTable.Entry e = table[bucketIndex]; table[bucketIndex] = new DictionaryHashTable.Entry(hash, index, e); size++; } /** Add Entry at the specified location of the table. */ void addEntry(int hash, int index, int bucketIndex) { if ((size >= threshold) && (null != table[bucketIndex])) { resize(2 * table.length); bucketIndex = indexFor(hash, table.length); } createEntry(hash, index, bucketIndex); } /** Resize table with given new capacity. */ void resize(int newCapacity) { DictionaryHashTable.Entry[] oldTable = table; int oldCapacity = oldTable.length; if (oldCapacity == MAXIMUM_CAPACITY) { threshold = Integer.MAX_VALUE; return; } DictionaryHashTable.Entry[] newTable = new DictionaryHashTable.Entry[newCapacity]; transfer(newTable); table = newTable; threshold = (int) Math.min(newCapacity * loadFactor, MAXIMUM_CAPACITY + 1); } /** * Transfer entries into new table from old table. * * @param newTable new table */ void transfer(DictionaryHashTable.Entry[] newTable) { int newCapacity = newTable.length; for (DictionaryHashTable.Entry e : table) { while (null != e) { DictionaryHashTable.Entry next = e.next; int i = indexFor(e.hash, newCapacity); e.next = newTable[i]; newTable[i] = e; e = next; } } } /** Returns the number of mappings in this Map. */ public int size() { return size; } /** Removes all elements from this map, leaving it empty. */ public void clear() { size = 0; for (int i = 0; i < table.length; i++) { table[i] = null; } } /** Class to keep dictionary index data within hash table. */ static class Entry { // dictionary index int index; DictionaryHashTable.Entry next; int hash; Entry(int hash, int index, DictionaryHashTable.Entry next) { this.index = index; this.hash = hash; this.next = next; } public final int getIndex() { return this.index; } @Override public int hashCode() { return hash; } public final boolean equals(Object o) { if (!(o instanceof DictionaryHashTable.Entry)) { return false; } DictionaryHashTable.Entry e = (DictionaryHashTable.Entry) o; if (index == e.getIndex()) { return true; } return false; } } } DictionaryProvider.java000066400000000000000000000055561515451615200342640ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dictionary; import java.util.HashMap; import java.util.Map; import java.util.Set; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.VisibleForTesting; /** A manager for association of dictionary IDs to their corresponding {@link Dictionary}. */ public interface DictionaryProvider { /** Return the dictionary for the given ID. */ Dictionary lookup(long id); /** Get all dictionary IDs. */ Set getDictionaryIds(); /** Implementation of {@link DictionaryProvider} that is backed by a hash-map. */ class MapDictionaryProvider implements AutoCloseable, DictionaryProvider { private final Map map; /** Constructs a new instance from the given dictionaries. */ public MapDictionaryProvider(Dictionary... dictionaries) { this.map = new HashMap<>(); for (Dictionary dictionary : dictionaries) { put(dictionary); } } /** * Initialize the map structure from another provider, but with empty vectors. * * @param other the {@link DictionaryProvider} to copy the ids and fields from * @param allocator allocator to create the empty vectors */ // This is currently called using JPype by the integration tests. @VisibleForTesting public void copyStructureFrom(DictionaryProvider other, BufferAllocator allocator) { for (Long id : other.getDictionaryIds()) { Dictionary otherDict = other.lookup(id); Dictionary newDict = new Dictionary( otherDict.getVector().getField().createVector(allocator), otherDict.getEncoding()); put(newDict); } } public void put(Dictionary dictionary) { map.put(dictionary.getEncoding().getId(), dictionary); } @Override public final Set getDictionaryIds() { return map.keySet(); } @Override public Dictionary lookup(long id) { return map.get(id); } @Override public void close() { for (Dictionary dictionary : map.values()) { dictionary.getVector().close(); } } } } ListSubfieldEncoder.java000066400000000000000000000147261515451615200343340ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dictionary; import java.util.Collections; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.memory.util.hash.SimpleHasher; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.BaseIntVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.BaseListVector; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; /** Sub fields encoder/decoder for Dictionary encoded {@link BaseListVector}. */ public class ListSubfieldEncoder { private final DictionaryHashTable hashTable; private final Dictionary dictionary; private final BufferAllocator allocator; public ListSubfieldEncoder(Dictionary dictionary, BufferAllocator allocator) { this(dictionary, allocator, SimpleHasher.INSTANCE); } /** Construct an instance. */ public ListSubfieldEncoder( Dictionary dictionary, BufferAllocator allocator, ArrowBufHasher hasher) { this.dictionary = dictionary; this.allocator = allocator; BaseListVector dictVector = (BaseListVector) dictionary.getVector(); hashTable = new DictionaryHashTable(getDataVector(dictVector), hasher); } private static FieldVector getDataVector(BaseListVector vector) { return vector.getChildrenFromFields().get(0); } private static BaseListVector cloneVector(BaseListVector vector, BufferAllocator allocator) { final FieldType fieldType = vector.getField().getFieldType(); BaseListVector cloned = (BaseListVector) fieldType.createNewSingleVector( vector.getField().getName(), allocator, /* schemaCallBack= */ null); final ArrowFieldNode fieldNode = new ArrowFieldNode(vector.getValueCount(), vector.getNullCount()); cloned.loadFieldBuffers(fieldNode, vector.getFieldBuffers()); return cloned; } /** * Dictionary encodes subfields for complex vector with a provided dictionary. The dictionary must * contain all values in the sub fields vector. * * @param vector vector to encode * @return dictionary encoded vector */ public BaseListVector encodeListSubField(BaseListVector vector) { final int valueCount = vector.getValueCount(); FieldType indexFieldType = new FieldType( vector.getField().isNullable(), dictionary.getEncoding().getIndexType(), dictionary.getEncoding(), vector.getField().getMetadata()); Field valueField = new Field(vector.getField().getName(), indexFieldType, null); // clone list vector and initialize data vector BaseListVector encoded = cloneVector(vector, allocator); try { encoded.initializeChildrenFromFields(Collections.singletonList(valueField)); BaseIntVector indices = (BaseIntVector) getDataVector(encoded); ValueVector dataVector = getDataVector(vector); for (int i = 0; i < valueCount; i++) { if (!vector.isNull(i)) { int start = vector.getElementStartIndex(i); int end = vector.getElementEndIndex(i); DictionaryEncoder.buildIndexVector(dataVector, indices, hashTable, start, end); } } return encoded; } catch (Exception e) { AutoCloseables.close(e, encoded); throw e; } } /** * Decodes a dictionary subfields encoded vector using the provided dictionary. * *

    {@link ListSubfieldEncoder#decodeListSubField(BaseListVector, Dictionary, BufferAllocator)} * should be used instead if only decoding is required as it can avoid building the {@link * DictionaryHashTable} which only makes sense when encoding. * * @param vector dictionary encoded vector, its data vector must be int type * @return vector with values restored from dictionary */ public BaseListVector decodeListSubField(BaseListVector vector) { return decodeListSubField(vector, dictionary, allocator); } /** * Decodes a dictionary subfields encoded vector using the provided dictionary. * * @param vector dictionary encoded vector, its data vector must be int type * @param dictionary dictionary used to decode the values * @param allocator allocator the decoded values use * @return vector with values restored from dictionary */ public static BaseListVector decodeListSubField( BaseListVector vector, Dictionary dictionary, BufferAllocator allocator) { int valueCount = vector.getValueCount(); BaseListVector dictionaryVector = (BaseListVector) dictionary.getVector(); int dictionaryValueCount = getDataVector(dictionaryVector).getValueCount(); // clone list vector and initialize data vector BaseListVector decoded = cloneVector(vector, allocator); try { Field dataVectorField = getDataVector(dictionaryVector).getField(); decoded.initializeChildrenFromFields(Collections.singletonList(dataVectorField)); // get data vector ValueVector dataVector = getDataVector(decoded); TransferPair transfer = getDataVector(dictionaryVector).makeTransferPair(dataVector); BaseIntVector indices = (BaseIntVector) getDataVector(vector); for (int i = 0; i < valueCount; i++) { if (!vector.isNull(i)) { int start = vector.getElementStartIndex(i); int end = vector.getElementEndIndex(i); DictionaryEncoder.retrieveIndexVector( indices, transfer, dictionaryValueCount, start, end); } } return decoded; } catch (Exception e) { AutoCloseables.close(e, decoded); throw e; } } } StructSubfieldEncoder.java000066400000000000000000000220711515451615200346750ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/dictionary/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.dictionary; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.memory.util.hash.SimpleHasher; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BaseIntVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; 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.TransferPair; /** * Sub fields encoder/decoder for Dictionary encoded {@link StructVector}. Notes that child vectors * within struct vector can either be dictionary encodable or not. */ public class StructSubfieldEncoder { private final BufferAllocator allocator; private final DictionaryProvider.MapDictionaryProvider provider; private final Map dictionaryIdToHashTable; /** Construct an instance. */ public StructSubfieldEncoder( BufferAllocator allocator, DictionaryProvider.MapDictionaryProvider provider) { this(allocator, provider, SimpleHasher.INSTANCE); } /** Construct an instance. */ public StructSubfieldEncoder( BufferAllocator allocator, DictionaryProvider.MapDictionaryProvider provider, ArrowBufHasher hasher) { this.allocator = allocator; this.provider = provider; this.dictionaryIdToHashTable = new HashMap<>(); provider .getDictionaryIds() .forEach( id -> dictionaryIdToHashTable.put( id, new DictionaryHashTable(provider.lookup(id).getVector(), hasher))); } private static FieldVector getChildVector(StructVector vector, int index) { return vector.getChildrenFromFields().get(index); } private static StructVector cloneVector(StructVector vector, BufferAllocator allocator) { final FieldType fieldType = vector.getField().getFieldType(); StructVector cloned = (StructVector) fieldType.createNewSingleVector( vector.getField().getName(), allocator, /* schemaCallback= */ null); final ArrowFieldNode fieldNode = new ArrowFieldNode(vector.getValueCount(), vector.getNullCount()); cloned.loadFieldBuffers(fieldNode, vector.getFieldBuffers()); return cloned; } /** * Dictionary encodes subfields for complex vector with a provided dictionary. The dictionary must * contain all values in the sub fields vector. * * @param vector vector to encode * @param columnToDictionaryId the mappings between child vector index and dictionary id. A null * dictionary id indicates the child vector is not encodable. * @return dictionary encoded vector */ public StructVector encode(StructVector vector, Map columnToDictionaryId) { final int valueCount = vector.getValueCount(); final int childCount = vector.getChildrenFromFields().size(); List childrenFields = new ArrayList<>(); // initialize child fields. for (int i = 0; i < childCount; i++) { FieldVector childVector = getChildVector(vector, i); Long dictionaryId = columnToDictionaryId.get(i); // A null dictionaryId indicates the child vector shouldn't be encoded. if (dictionaryId == null) { childrenFields.add(childVector.getField()); } else { Dictionary dictionary = provider.lookup(dictionaryId); Preconditions.checkNotNull(dictionary, "Dictionary not found with id:" + dictionaryId); FieldType indexFieldType = new FieldType( childVector.getField().isNullable(), dictionary.getEncoding().getIndexType(), dictionary.getEncoding()); childrenFields.add( new Field(childVector.getField().getName(), indexFieldType, /* children= */ null)); } } // clone list vector and initialize data vector StructVector encoded = cloneVector(vector, allocator); try { encoded.initializeChildrenFromFields(childrenFields); encoded.setValueCount(valueCount); for (int index = 0; index < childCount; index++) { FieldVector childVector = getChildVector(vector, index); FieldVector encodedChildVector = getChildVector(encoded, index); Long dictionaryId = columnToDictionaryId.get(index); if (dictionaryId != null) { BaseIntVector indices = (BaseIntVector) encodedChildVector; DictionaryEncoder.buildIndexVector( childVector, indices, dictionaryIdToHashTable.get(dictionaryId), 0, valueCount); } else { childVector.makeTransferPair(encodedChildVector).splitAndTransfer(0, valueCount); } } return encoded; } catch (Exception e) { AutoCloseables.close(e, encoded); throw e; } } /** * Decodes a dictionary subfields encoded vector using the provided dictionary. * *

    {@link StructSubfieldEncoder#decode(StructVector, DictionaryProvider.MapDictionaryProvider, * BufferAllocator)} should be used instead if only decoding is required as it can avoid building * the {@link DictionaryHashTable} which only makes sense when encoding. * * @param vector dictionary encoded vector, its child vector must be int type * @return vector with values restored from dictionary */ public StructVector decode(StructVector vector) { return decode(vector, provider, allocator); } /** * Decodes a dictionary subfields encoded vector using the provided dictionary. * * @param vector dictionary encoded vector, its data vector must be int type * @param provider dictionary provider used to decode the values * @param allocator allocator the decoded values use * @return vector with values restored from dictionary */ public static StructVector decode( StructVector vector, DictionaryProvider.MapDictionaryProvider provider, BufferAllocator allocator) { final int valueCount = vector.getValueCount(); final int childCount = vector.getChildrenFromFields().size(); // clone list vector and initialize child vectors StructVector decoded = cloneVector(vector, allocator); try { List childFields = new ArrayList<>(); for (int i = 0; i < childCount; i++) { FieldVector childVector = getChildVector(vector, i); Dictionary dictionary = getChildVectorDictionary(childVector, provider); // childVector is not encoded. if (dictionary == null) { childFields.add(childVector.getField()); } else { childFields.add(dictionary.getVector().getField()); } } decoded.initializeChildrenFromFields(childFields); decoded.setValueCount(valueCount); for (int index = 0; index < childCount; index++) { // get child vector FieldVector childVector = getChildVector(vector, index); FieldVector decodedChildVector = getChildVector(decoded, index); Dictionary dictionary = getChildVectorDictionary(childVector, provider); if (dictionary == null) { childVector.makeTransferPair(decodedChildVector).splitAndTransfer(0, valueCount); } else { TransferPair transfer = dictionary.getVector().makeTransferPair(decodedChildVector); BaseIntVector indices = (BaseIntVector) childVector; DictionaryEncoder.retrieveIndexVector(indices, transfer, valueCount, 0, valueCount); } } return decoded; } catch (Exception e) { AutoCloseables.close(e, decoded); throw e; } } /** Get the child vector dictionary, return null if not dictionary encoded. */ private static Dictionary getChildVectorDictionary( FieldVector childVector, DictionaryProvider.MapDictionaryProvider provider) { DictionaryEncoding dictionaryEncoding = childVector.getField().getDictionary(); if (dictionaryEncoding != null) { Dictionary dictionary = provider.lookup(dictionaryEncoding.getId()); Preconditions.checkNotNull(dictionary, "Dictionary not found with id:" + dictionary); return dictionary; } return null; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/extension/000077500000000000000000000000001515451615200275145ustar00rootroot00000000000000InvalidExtensionMetadataException.java000066400000000000000000000022111515451615200370770ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/extension/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** The extension metadata was malformed. */ public class InvalidExtensionMetadataException extends RuntimeException { public InvalidExtensionMetadataException(String message) { super(message); } public InvalidExtensionMetadataException(String message, Throwable cause) { super(message, cause); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/extension/OpaqueType.java000066400000000000000000000321161515451615200324560ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.node.ObjectNode; import java.util.Collections; import java.util.Objects; import java.util.concurrent.atomic.AtomicBoolean; import org.apache.arrow.memory.BufferAllocator; 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.DurationVector; 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.IntervalDayVector; import org.apache.arrow.vector.IntervalMonthDayNanoVector; import org.apache.arrow.vector.IntervalYearVector; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.LargeVarCharVector; import org.apache.arrow.vector.NullVector; 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.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; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; /** * Opaque is a placeholder for a type from an external (usually non-Arrow) system that could not be * interpreted. */ public class OpaqueType extends ArrowType.ExtensionType { private static final AtomicBoolean registered = new AtomicBoolean(false); public static final String EXTENSION_NAME = "arrow.opaque"; private final ArrowType storageType; private final String typeName; private final String vendorName; /** Register the extension type so it can be used globally. */ public static void ensureRegistered() { if (!registered.getAndSet(true)) { // The values don't matter, we just need an instance ExtensionTypeRegistry.register(new OpaqueType(Types.MinorType.NULL.getType(), "", "")); } } /** * Create a new type instance. * * @param storageType The underlying Arrow type. * @param typeName The name of the unknown type. * @param vendorName The name of the originating system of the unknown type. */ public OpaqueType(ArrowType storageType, String typeName, String vendorName) { this.storageType = Objects.requireNonNull(storageType, "storageType"); this.typeName = Objects.requireNonNull(typeName, "typeName"); this.vendorName = Objects.requireNonNull(vendorName, "vendorName"); } @Override public ArrowType storageType() { return storageType; } public String typeName() { return typeName; } public String vendorName() { return vendorName; } @Override public String extensionName() { return EXTENSION_NAME; } @Override public boolean extensionEquals(ExtensionType other) { return other != null && EXTENSION_NAME.equals(other.extensionName()) && other instanceof OpaqueType && storageType.equals(other.storageType()) && typeName.equals(((OpaqueType) other).typeName()) && vendorName.equals(((OpaqueType) other).vendorName()); } @Override public String serialize() { ObjectMapper mapper = new ObjectMapper(); ObjectNode object = mapper.createObjectNode(); object.put("type_name", typeName); object.put("vendor_name", vendorName); try { return mapper.writeValueAsString(object); } catch (JsonProcessingException e) { throw new RuntimeException("Could not serialize " + this, e); } } @Override public ArrowType deserialize(ArrowType storageType, String serializedData) { ObjectMapper mapper = new ObjectMapper(); JsonNode object; try { object = mapper.readTree(serializedData); } catch (JsonProcessingException e) { throw new InvalidExtensionMetadataException("Extension metadata is invalid", e); } JsonNode typeName = object.get("type_name"); JsonNode vendorName = object.get("vendor_name"); if (typeName == null) { throw new InvalidExtensionMetadataException("typeName is missing"); } if (vendorName == null) { throw new InvalidExtensionMetadataException("vendorName is missing"); } if (!typeName.isTextual()) { throw new InvalidExtensionMetadataException("typeName should be string, was " + typeName); } if (!vendorName.isTextual()) { throw new InvalidExtensionMetadataException("vendorName should be string, was " + vendorName); } return new OpaqueType(storageType, typeName.asText(), vendorName.asText()); } @Override public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) { // XXX: fieldType is supposed to be the extension type final Field field = new Field(name, fieldType, Collections.emptyList()); final FieldVector underlyingVector = storageType.accept(new UnderlyingVectorTypeVisitor(name, allocator)); return new OpaqueVector(field, allocator, underlyingVector); } @Override 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(" + storageType + ", typeName='" + typeName + '\'' + ", vendorName='" + vendorName + '\'' + ')'; } private static class UnderlyingVectorTypeVisitor implements ArrowTypeVisitor { private final String name; private final BufferAllocator allocator; UnderlyingVectorTypeVisitor(String name, BufferAllocator allocator) { this.name = name; this.allocator = allocator; } @Override public FieldVector visit(Null type) { return new NullVector(name); } private RuntimeException unsupported(ArrowType type) { throw new UnsupportedOperationException( "OpaqueType#getUnderlyingVector is not supported for storage type: " + type); } @Override public FieldVector visit(Struct type) { throw unsupported(type); } @Override public FieldVector visit(List type) { throw unsupported(type); } @Override public FieldVector visit(LargeList type) { throw unsupported(type); } @Override public FieldVector visit(FixedSizeList type) { throw unsupported(type); } @Override public FieldVector visit(Union type) { throw unsupported(type); } @Override public FieldVector visit(Map type) { throw unsupported(type); } @Override public FieldVector visit(Int type) { return new IntVector(name, allocator); } @Override public FieldVector visit(FloatingPoint type) { switch (type.getPrecision()) { case HALF: return new Float2Vector(name, allocator); case SINGLE: return new Float4Vector(name, allocator); case DOUBLE: return new Float8Vector(name, allocator); default: throw unsupported(type); } } @Override public FieldVector visit(Utf8 type) { return new VarCharVector(name, allocator); } @Override public FieldVector visit(Utf8View type) { return new ViewVarCharVector(name, allocator); } @Override public FieldVector visit(LargeUtf8 type) { return new LargeVarCharVector(name, allocator); } @Override public FieldVector visit(Binary type) { return new VarBinaryVector(name, allocator); } @Override public FieldVector visit(BinaryView type) { return new ViewVarBinaryVector(name, allocator); } @Override public FieldVector visit(LargeBinary type) { return new LargeVarBinaryVector(name, allocator); } @Override public FieldVector visit(FixedSizeBinary type) { return new FixedSizeBinaryVector(Field.nullable(name, type), allocator); } @Override public FieldVector visit(Bool type) { return new BitVector(name, allocator); } @Override public FieldVector visit(Decimal type) { if (type.getBitWidth() == 128) { return new DecimalVector(Field.nullable(name, type), allocator); } else if (type.getBitWidth() == 256) { return new Decimal256Vector(Field.nullable(name, type), allocator); } throw unsupported(type); } @Override public FieldVector visit(Date type) { switch (type.getUnit()) { case DAY: return new DateDayVector(name, allocator); case MILLISECOND: return new DateMilliVector(name, allocator); default: throw unsupported(type); } } @Override public FieldVector visit(Time type) { switch (type.getUnit()) { case SECOND: return new TimeSecVector(name, allocator); case MILLISECOND: return new TimeMilliVector(name, allocator); case MICROSECOND: return new TimeMicroVector(name, allocator); case NANOSECOND: return new TimeNanoVector(name, allocator); default: throw unsupported(type); } } @Override public FieldVector visit(Timestamp type) { if (type.getTimezone() == null || type.getTimezone().isEmpty()) { switch (type.getUnit()) { case SECOND: return new TimeStampSecVector(Field.nullable(name, type), allocator); case MILLISECOND: return new TimeStampMilliVector(Field.nullable(name, type), allocator); case MICROSECOND: return new TimeStampMicroVector(Field.nullable(name, type), allocator); case NANOSECOND: return new TimeStampNanoVector(Field.nullable(name, type), allocator); default: throw unsupported(type); } } switch (type.getUnit()) { case SECOND: return new TimeStampSecTZVector(Field.nullable(name, type), allocator); case MILLISECOND: return new TimeStampMilliTZVector(Field.nullable(name, type), allocator); case MICROSECOND: return new TimeStampMicroTZVector(Field.nullable(name, type), allocator); case NANOSECOND: return new TimeStampNanoTZVector(Field.nullable(name, type), allocator); default: throw unsupported(type); } } @Override public FieldVector visit(Interval type) { switch (type.getUnit()) { case YEAR_MONTH: return new IntervalYearVector(name, allocator); case DAY_TIME: return new IntervalDayVector(name, allocator); case MONTH_DAY_NANO: return new IntervalMonthDayNanoVector(name, allocator); default: throw unsupported(type); } } @Override public FieldVector visit(Duration type) { return new DurationVector(Field.nullable(name, type), allocator); } @Override public FieldVector visit(ListView type) { throw unsupported(type); } @Override public FieldVector visit(LargeListView type) { throw unsupported(type); } @Override public FieldVector visit(RunEndEncoded type) { throw unsupported(type); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/extension/OpaqueVector.java000066400000000000000000000036771515451615200330110ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.memory.util.hash.ArrowBufHasher; import org.apache.arrow.vector.ExtensionTypeVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueIterableVector; import org.apache.arrow.vector.types.pojo.Field; /** * Opaque is a wrapper for (usually binary) data from an external (often non-Arrow) system that * could not be interpreted. */ public class OpaqueVector extends ExtensionTypeVector implements ValueIterableVector { private final Field field; public OpaqueVector(Field field, BufferAllocator allocator, FieldVector underlyingVector) { super(field, allocator, underlyingVector); this.field = field; } @Override public Field getField() { return field; } @Override public Object getObject(int index) { return getUnderlyingVector().getObject(index); } @Override public int hashCode(int index) { return hashCode(index, null); } @Override public int hashCode(int index, ArrowBufHasher hasher) { return getUnderlyingVector().hashCode(index, hasher); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/extension/UuidType.java000066400000000000000000000076341515451615200321410ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/holders/000077500000000000000000000000001515451615200271405ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/holders/ComplexHolder.java000066400000000000000000000020671515451615200325550ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex.reader.FieldReader; /** Represents a single value of a complex type (e.g. Union, Struct). */ public class ComplexHolder implements ValueHolder { public FieldReader reader; public int isSet; } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/holders/DenseUnionHolder.java000066400000000000000000000023571515451615200332170ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex.reader.FieldReader; import org.apache.arrow.vector.types.Types; /** {@link ValueHolder} for Dense Union Vectors. */ public class DenseUnionHolder implements ValueHolder { public FieldReader reader; public int isSet; public byte typeId; public Types.MinorType getMinorType() { return reader.getMinorType(); } public boolean isSet() { return isSet == 1; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/holders/ExtensionHolder.java000066400000000000000000000021341515451615200331150ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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(); } NullableUuidHolder.java000066400000000000000000000032241515451615200334500ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/holders/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; } } RepeatedListHolder.java000066400000000000000000000020161515451615200334460ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/holders/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** {@link ValueHolder} for a nested {@link org.apache.arrow.vector.complex.ListVector}. */ public final class RepeatedListHolder implements ValueHolder { public int start; public int end; } RepeatedStructHolder.java000066400000000000000000000017461515451615200340300ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/holders/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** {@link ValueHolder} for a list of structs. */ public final class RepeatedStructHolder implements ValueHolder { public int start; public int end; } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/holders/UnionHolder.java000066400000000000000000000023221515451615200322300ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.complex.reader.FieldReader; import org.apache.arrow.vector.types.Types.MinorType; /** {@link ValueHolder} for Union Vectors. */ public class UnionHolder implements ValueHolder { public FieldReader reader; public int isSet; public MinorType getMinorType() { return reader.getMinorType(); } public boolean isSet() { return isSet == 1; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/holders/UuidHolder.java000066400000000000000000000031331515451615200320470ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/holders/ValueHolder.java000066400000000000000000000024711515451615200322210ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * Wrapper object for an individual value in Arrow. * *

    ValueHolders are designed to be mutable wrapper objects for defining clean APIs that access * data in Arrow. For performance, object creation is avoided at all costs throughout execution. For * this reason, ValueHolders are disallowed from implementing any methods, this allows for them to * be replaced by their java primitive inner members during optimization of run-time generated code. */ public interface ValueHolder {} apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/000077500000000000000000000000001515451615200262535ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/ArrowFileReader.java000066400000000000000000000201071515451615200321330ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.channels.SeekableByteChannel; import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; import org.apache.arrow.flatbuf.Footer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.VisibleForTesting; import org.apache.arrow.vector.compression.CompressionCodec; import org.apache.arrow.vector.ipc.message.ArrowBlock; import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch; import org.apache.arrow.vector.ipc.message.ArrowFooter; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.validate.MetadataV4UnionChecker; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** An implementation of {@link ArrowReader} that reads the standard arrow binary file format. */ public class ArrowFileReader extends ArrowReader { private static final Logger LOGGER = LoggerFactory.getLogger(ArrowFileReader.class); private SeekableReadChannel in; private ArrowFooter footer; private int currentDictionaryBatch = 0; private int currentRecordBatch = 0; public ArrowFileReader( SeekableReadChannel in, BufferAllocator allocator, CompressionCodec.Factory compressionFactory) { super(allocator, compressionFactory); this.in = in; } public ArrowFileReader( SeekableByteChannel in, BufferAllocator allocator, CompressionCodec.Factory compressionFactory) { this(new SeekableReadChannel(in), allocator, compressionFactory); } public ArrowFileReader(SeekableReadChannel in, BufferAllocator allocator) { this(in, allocator, CompressionCodec.Factory.INSTANCE); } public ArrowFileReader(SeekableByteChannel in, BufferAllocator allocator) { this(new SeekableReadChannel(in), allocator); } @Override public long bytesRead() { return in.bytesRead(); } @Override protected void closeReadSource() throws IOException { in.close(); } @Override protected Schema readSchema() throws IOException { if (footer == null) { if (in.size() <= (ArrowMagic.MAGIC_LENGTH * 2 + 4)) { throw new InvalidArrowFileException("file too small: " + in.size()); } ByteBuffer buffer = ByteBuffer.allocate(4 + ArrowMagic.MAGIC_LENGTH); long footerLengthOffset = in.size() - buffer.remaining(); in.setPosition(footerLengthOffset); in.readFully(buffer); buffer.flip(); byte[] array = buffer.array(); if (!ArrowMagic.validateMagic(Arrays.copyOfRange(array, 4, array.length))) { throw new InvalidArrowFileException( "missing Magic number " + Arrays.toString(buffer.array())); } int footerLength = MessageSerializer.bytesToInt(array); if (footerLength <= 0 || footerLength + ArrowMagic.MAGIC_LENGTH * 2 + 4 > in.size() || footerLength > footerLengthOffset) { throw new InvalidArrowFileException("invalid footer length: " + footerLength); } long footerOffset = footerLengthOffset - footerLength; LOGGER.debug("Footer starts at {}, length: {}", footerOffset, footerLength); ByteBuffer footerBuffer = ByteBuffer.allocate(footerLength); in.setPosition(footerOffset); in.readFully(footerBuffer); footerBuffer.flip(); Footer footerFB = Footer.getRootAsFooter(footerBuffer); this.footer = new ArrowFooter(footerFB); } MetadataV4UnionChecker.checkRead(footer.getSchema(), footer.getMetadataVersion()); return footer.getSchema(); } @Override public void initialize() throws IOException { super.initialize(); // empty stream, has no dictionaries in IPC. if (footer.getRecordBatches().size() == 0) { return; } // Read and load all dictionaries from schema for (int i = 0; i < dictionaries.size(); i++) { ArrowDictionaryBatch dictionaryBatch = readDictionary(); loadDictionary(dictionaryBatch); } } /** Get custom metadata. */ public Map getMetaData() { if (footer != null) { return footer.getMetaData(); } return new HashMap<>(); } /** * Read a dictionary batch from the source, will be invoked after the schema has been read and * called N times, where N is the number of dictionaries indicated by the schema Fields. * * @return the read ArrowDictionaryBatch * @throws IOException on error */ public ArrowDictionaryBatch readDictionary() throws IOException { if (currentDictionaryBatch >= footer.getDictionaries().size()) { throw new IOException( "Requested more dictionaries than defined in footer: " + currentDictionaryBatch); } ArrowBlock block = footer.getDictionaries().get(currentDictionaryBatch++); return readDictionaryBatch(in, block, allocator); } /** Returns true if a batch was read, false if no more batches. */ @Override public boolean loadNextBatch() throws IOException { prepareLoadNextBatch(); if (currentRecordBatch < footer.getRecordBatches().size()) { ArrowBlock block = footer.getRecordBatches().get(currentRecordBatch++); ArrowRecordBatch batch = readRecordBatch(in, block, allocator); loadRecordBatch(batch); return true; } else { return false; } } public List getDictionaryBlocks() throws IOException { ensureInitialized(); return footer.getDictionaries(); } /** Returns the {@link ArrowBlock} metadata from the file. */ public List getRecordBlocks() throws IOException { ensureInitialized(); return footer.getRecordBatches(); } /** Loads record batch for the given block. */ public boolean loadRecordBatch(ArrowBlock block) throws IOException { ensureInitialized(); int blockIndex = footer.getRecordBatches().indexOf(block); if (blockIndex == -1) { throw new IllegalArgumentException("Arrow block does not exist in record batches: " + block); } currentRecordBatch = blockIndex; return loadNextBatch(); } @VisibleForTesting ArrowFooter getFooter() { return footer; } private ArrowDictionaryBatch readDictionaryBatch( SeekableReadChannel in, ArrowBlock block, BufferAllocator allocator) throws IOException { LOGGER.debug( "DictionaryRecordBatch at {}, metadata: {}, body: {}", block.getOffset(), block.getMetadataLength(), block.getBodyLength()); in.setPosition(block.getOffset()); ArrowDictionaryBatch batch = MessageSerializer.deserializeDictionaryBatch(in, block, allocator); if (batch == null) { throw new IOException("Invalid file. No batch at offset: " + block.getOffset()); } return batch; } private ArrowRecordBatch readRecordBatch( SeekableReadChannel in, ArrowBlock block, BufferAllocator allocator) throws IOException { LOGGER.debug( "RecordBatch at {}, metadata: {}, body: {}", block.getOffset(), block.getMetadataLength(), block.getBodyLength()); in.setPosition(block.getOffset()); ArrowRecordBatch batch = MessageSerializer.deserializeRecordBatch(in, block, allocator); if (batch == null) { throw new IOException("Invalid file. No batch at offset: " + block.getOffset()); } return batch; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/ArrowFileWriter.java000066400000000000000000000136711515451615200322150ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import java.io.IOException; import java.nio.channels.WritableByteChannel; import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; import org.apache.arrow.util.VisibleForTesting; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.compression.CompressionCodec; import org.apache.arrow.vector.compression.CompressionUtil; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.ipc.message.ArrowBlock; import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch; import org.apache.arrow.vector.ipc.message.ArrowFooter; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.ipc.message.IpcOption; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * {@link ArrowWriter} that writes out a Arrow files * (https://arrow.apache.org/docs/format/IPC.html#file-format). */ public class ArrowFileWriter extends ArrowWriter { private static final Logger LOGGER = LoggerFactory.getLogger(ArrowFileWriter.class); // All ArrowBlocks written are saved in these lists to be passed to ArrowFooter in endInternal. private final List dictionaryBlocks = new ArrayList<>(); private final List recordBlocks = new ArrayList<>(); private Map metaData; private boolean dictionariesWritten = false; public ArrowFileWriter( VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out) { super(root, provider, out); } public ArrowFileWriter( VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out, Map metaData) { super(root, provider, out); this.metaData = metaData; } public ArrowFileWriter( VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out, IpcOption option) { super(root, provider, out, option); } public ArrowFileWriter( VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out, Map metaData, IpcOption option) { super(root, provider, out, option); this.metaData = metaData; } public ArrowFileWriter( VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out, Map metaData, IpcOption option, CompressionCodec.Factory compressionFactory, CompressionUtil.CodecType codecType) { this(root, provider, out, metaData, option, compressionFactory, codecType, Optional.empty()); } public ArrowFileWriter( VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out, Map metaData, IpcOption option, CompressionCodec.Factory compressionFactory, CompressionUtil.CodecType codecType, Optional compressionLevel) { super(root, provider, out, option, compressionFactory, codecType, compressionLevel); this.metaData = metaData; } @Override protected void startInternal(WriteChannel out) throws IOException { ArrowMagic.writeMagic(out, true); } @Override protected ArrowBlock writeDictionaryBatch(ArrowDictionaryBatch batch) throws IOException { ArrowBlock block = super.writeDictionaryBatch(batch); dictionaryBlocks.add(block); return block; } @Override protected ArrowBlock writeRecordBatch(ArrowRecordBatch batch) throws IOException { ArrowBlock block = super.writeRecordBatch(batch); recordBlocks.add(block); return block; } @Override protected void endInternal(WriteChannel out) throws IOException { if (!option.write_legacy_ipc_format) { out.writeIntLittleEndian(MessageSerializer.IPC_CONTINUATION_TOKEN); } out.writeIntLittleEndian(0); long footerStart = out.getCurrentPosition(); out.write( new ArrowFooter(schema, dictionaryBlocks, recordBlocks, metaData, option.metadataVersion), false); int footerLength = (int) (out.getCurrentPosition() - footerStart); if (footerLength <= 0) { throw new InvalidArrowFileException("invalid footer"); } out.writeIntLittleEndian(footerLength); LOGGER.debug("Footer starts at {}, length: {}", footerStart, footerLength); ArrowMagic.writeMagic(out, false); LOGGER.debug("magic written, now at {}", out.getCurrentPosition()); } @Override protected void ensureDictionariesWritten(DictionaryProvider provider, Set dictionaryIdsUsed) throws IOException { if (dictionariesWritten) { return; } dictionariesWritten = true; // Write out all dictionaries required. // Replacement dictionaries are not supported in the IPC file format. for (long id : dictionaryIdsUsed) { Dictionary dictionary = provider.lookup(id); writeDictionaryBatch(dictionary); } } @VisibleForTesting public List getRecordBlocks() { return recordBlocks; } @VisibleForTesting public List getDictionaryBlocks() { return dictionaryBlocks; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/ArrowMagic.java000066400000000000000000000027211515451615200311530ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.util.Arrays; /** * Magic header/footer helpers for {@link ArrowFileWriter} and {@link ArrowFileReader} formatted * files. */ class ArrowMagic { private ArrowMagic() {} private static final byte[] MAGIC = "ARROW1".getBytes(StandardCharsets.UTF_8); public static final int MAGIC_LENGTH = MAGIC.length; public static void writeMagic(WriteChannel out, boolean align) throws IOException { out.write(MAGIC); if (align) { out.align(); } } public static boolean validateMagic(byte[] array) { return Arrays.equals(MAGIC, array); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/ArrowReader.java000066400000000000000000000200051515451615200313300ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import java.io.IOException; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Set; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.compression.CompressionCodec; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch; 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.apache.arrow.vector.util.DictionaryUtility; import org.apache.arrow.vector.util.VectorBatchAppender; /** Abstract class to read Schema and ArrowRecordBatches. */ public abstract class ArrowReader implements DictionaryProvider, AutoCloseable { protected final BufferAllocator allocator; private VectorLoader loader; private VectorSchemaRoot root; protected Map dictionaries; private boolean initialized = false; private final CompressionCodec.Factory compressionFactory; protected ArrowReader(BufferAllocator allocator) { this(allocator, CompressionCodec.Factory.INSTANCE); } protected ArrowReader(BufferAllocator allocator, CompressionCodec.Factory compressionFactory) { this.allocator = allocator; this.compressionFactory = compressionFactory; } /** * Returns the vector schema root. This will be loaded with new values on every call to * loadNextBatch. * * @return the vector schema root * @throws IOException if reading of schema fails */ public VectorSchemaRoot getVectorSchemaRoot() throws IOException { ensureInitialized(); return root; } /** * Returns any dictionaries that were loaded along with ArrowRecordBatches. * * @return Map of dictionaries to dictionary id, empty if no dictionaries loaded * @throws IOException if reading of schema fails */ public Map getDictionaryVectors() throws IOException { ensureInitialized(); return dictionaries; } /** * Lookup a dictionary that has been loaded using the dictionary id. * * @param id Unique identifier for a dictionary * @return the requested dictionary or null if not found */ @Override public Dictionary lookup(long id) { if (!initialized) { throw new IllegalStateException("Unable to lookup until reader has been initialized"); } return dictionaries.get(id); } @Override public Set getDictionaryIds() { return dictionaries.keySet(); } /** * Load the next ArrowRecordBatch to the vector schema root if available. * * @return true if a batch was read, false on EOS * @throws IOException on error */ public abstract boolean loadNextBatch() throws IOException; /** * Return the number of bytes read from the ReadChannel. * * @return number of bytes read */ public abstract long bytesRead(); /** * Close resources, including vector schema root and dictionary vectors, and the underlying read * source. * * @throws IOException on error */ @Override public void close() throws IOException { close(true); } /** * Close resources, including vector schema root and dictionary vectors. If the flag * closeReadChannel is true then close the underlying read source, otherwise leave it open. * * @param closeReadSource Flag to control if closing the underlying read source * @throws IOException on error */ public void close(boolean closeReadSource) throws IOException { if (initialized) { root.close(); for (Dictionary dictionary : dictionaries.values()) { dictionary.getVector().close(); } } if (closeReadSource) { closeReadSource(); } } /** * Close the underlying read source. * * @throws IOException on error */ protected abstract void closeReadSource() throws IOException; /** * Read the Schema from the source, will be invoked at the beginning the initialization. * * @return the read Schema * @throws IOException on error */ protected abstract Schema readSchema() throws IOException; /** * Initialize if not done previously. * * @throws IOException on error */ protected void ensureInitialized() throws IOException { if (!initialized) { initialize(); initialized = true; } } /** Reads the schema and initializes the vectors. */ protected void initialize() throws IOException { Schema originalSchema = readSchema(); List fields = new ArrayList<>(originalSchema.getFields().size()); List vectors = new ArrayList<>(originalSchema.getFields().size()); Map dictionaries = new HashMap<>(); // Convert fields with dictionaries to have the index type for (Field field : originalSchema.getFields()) { Field updated = DictionaryUtility.toMemoryFormat(field, allocator, dictionaries); fields.add(updated); vectors.add(updated.createVector(allocator)); } Schema schema = new Schema(fields, originalSchema.getCustomMetadata()); this.root = new VectorSchemaRoot(schema, vectors, 0); this.loader = new VectorLoader(root, compressionFactory); this.dictionaries = Collections.unmodifiableMap(dictionaries); } /** * Ensure the reader has been initialized and reset the VectorSchemaRoot row count to 0. * * @throws IOException on error */ protected void prepareLoadNextBatch() throws IOException { ensureInitialized(); root.setRowCount(0); } /** * Load an ArrowRecordBatch to the readers VectorSchemaRoot. * * @param batch the record batch to load */ protected void loadRecordBatch(ArrowRecordBatch batch) { try { loader.load(batch); } finally { batch.close(); } } /** * Load an ArrowDictionaryBatch to the readers dictionary vectors. * * @param dictionaryBatch dictionary batch to load */ protected void loadDictionary(ArrowDictionaryBatch dictionaryBatch) { long id = dictionaryBatch.getDictionaryId(); Dictionary dictionary = dictionaries.get(id); if (dictionary == null) { throw new IllegalArgumentException("Dictionary ID " + id + " not defined in schema"); } FieldVector vector = dictionary.getVector(); // if is deltaVector, concat it with non-delta vector with the same ID. if (dictionaryBatch.isDelta()) { try (FieldVector deltaVector = vector.getField().createVector(allocator)) { load(dictionaryBatch, deltaVector); VectorBatchAppender.batchAppend(vector, deltaVector); } return; } load(dictionaryBatch, vector); } private void load(ArrowDictionaryBatch dictionaryBatch, FieldVector vector) { VectorSchemaRoot root = new VectorSchemaRoot( Collections.singletonList(vector.getField()), Collections.singletonList(vector), 0); VectorLoader loader = new VectorLoader(root, this.compressionFactory); try { loader.load(dictionaryBatch.getDictionary()); } finally { dictionaryBatch.close(); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/ArrowStreamReader.java000066400000000000000000000201621515451615200325100ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import java.io.IOException; import java.io.InputStream; import java.nio.channels.Channels; import java.nio.channels.ReadableByteChannel; import org.apache.arrow.flatbuf.MessageHeader; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.compression.CompressionCodec; import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.ipc.message.MessageChannelReader; import org.apache.arrow.vector.ipc.message.MessageResult; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.MetadataVersion; import org.apache.arrow.vector.types.pojo.DictionaryEncoding; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.validate.MetadataV4UnionChecker; /** This class reads from an input stream and produces ArrowRecordBatches. */ public class ArrowStreamReader extends ArrowReader { private MessageChannelReader messageReader; private int loadedDictionaryCount; /** * Constructs a streaming reader using a MessageChannelReader. Non-blocking. * * @param messageReader reader used to get messages from a ReadChannel * @param allocator to allocate new buffers * @param compressionFactory the factory to create compression codec. */ public ArrowStreamReader( MessageChannelReader messageReader, BufferAllocator allocator, CompressionCodec.Factory compressionFactory) { super(allocator, compressionFactory); this.messageReader = messageReader; } /** * Constructs a streaming reader using a MessageChannelReader. Non-blocking. * * @param messageReader reader used to get messages from a ReadChannel * @param allocator to allocate new buffers */ public ArrowStreamReader(MessageChannelReader messageReader, BufferAllocator allocator) { this(messageReader, allocator, CompressionCodec.Factory.INSTANCE); } /** * Constructs a streaming reader from a ReadableByteChannel input. Non-blocking. * * @param in ReadableByteChannel to read messages from * @param allocator to allocate new buffers * @param compressionFactory the factory to create compression codec. */ public ArrowStreamReader( ReadableByteChannel in, BufferAllocator allocator, CompressionCodec.Factory compressionFactory) { this(new MessageChannelReader(new ReadChannel(in), allocator), allocator, compressionFactory); } /** * Constructs a streaming reader from a ReadableByteChannel input. Non-blocking. * * @param in ReadableByteChannel to read messages from * @param allocator to allocate new buffers */ public ArrowStreamReader(ReadableByteChannel in, BufferAllocator allocator) { this(new MessageChannelReader(new ReadChannel(in), allocator), allocator); } /** * Constructs a streaming reader from a ReadableByteChannel input. Non-blocking. * * @param in InputStream to read messages from * @param allocator to allocate new buffers * @param compressionFactory the factory to create compression codec. */ public ArrowStreamReader( InputStream in, BufferAllocator allocator, CompressionCodec.Factory compressionFactory) { this(Channels.newChannel(in), allocator, compressionFactory); } /** * Constructs a streaming reader from an InputStream. Non-blocking. * * @param in InputStream to read messages from * @param allocator to allocate new buffers */ public ArrowStreamReader(InputStream in, BufferAllocator allocator) { this(Channels.newChannel(in), allocator); } /** * Get the number of bytes read from the stream since constructing the reader. * * @return number of bytes */ @Override public long bytesRead() { return messageReader.bytesRead(); } /** * Closes the underlying read source. * * @throws IOException on error */ @Override protected void closeReadSource() throws IOException { messageReader.close(); } /** * Load the next ArrowRecordBatch to the vector schema root if available. * * @return true if a batch was read, false on EOS * @throws IOException on error */ public boolean loadNextBatch() throws IOException { prepareLoadNextBatch(); MessageResult result = messageReader.readNext(); // Reached EOS if (result == null) { return false; } if (result.getMessage().headerType() == MessageHeader.RecordBatch) { ArrowBuf bodyBuffer = result.getBodyBuffer(); // For zero-length batches, need an empty buffer to deserialize the batch if (bodyBuffer == null) { bodyBuffer = allocator.getEmpty(); } ArrowRecordBatch batch = MessageSerializer.deserializeRecordBatch(result.getMessage(), bodyBuffer); loadRecordBatch(batch); checkDictionaries(); return true; } else if (result.getMessage().headerType() == MessageHeader.DictionaryBatch) { // if it's dictionary message, read dictionary message out and continue to read unless get a // batch or eos. ArrowDictionaryBatch dictionaryBatch = readDictionary(result); loadDictionary(dictionaryBatch); loadedDictionaryCount++; return loadNextBatch(); } else { throw new IOException( "Expected RecordBatch or DictionaryBatch but header was " + result.getMessage().headerType()); } } /** When read a record batch, check whether its dictionaries are available. */ private void checkDictionaries() throws IOException { // if all dictionaries are loaded, return. if (loadedDictionaryCount == dictionaries.size()) { return; } for (FieldVector vector : getVectorSchemaRoot().getFieldVectors()) { DictionaryEncoding encoding = vector.getField().getDictionary(); if (encoding != null) { // if the dictionaries it needs is not available and the vector is not all null, something // was wrong. if (!dictionaries.containsKey(encoding.getId()) && vector.getNullCount() < vector.getValueCount()) { throw new IOException("The dictionary was not available, id was:" + encoding.getId()); } } } } /** * Reads the schema message from the beginning of the stream. * * @return the deserialized arrow schema */ @Override protected Schema readSchema() throws IOException { MessageResult result = messageReader.readNext(); if (result == null) { throw new IOException("Unexpected end of input. Missing schema."); } if (result.getMessage().headerType() != MessageHeader.Schema) { throw new IOException("Expected schema but header was " + result.getMessage().headerType()); } final Schema schema = MessageSerializer.deserializeSchema(result.getMessage()); MetadataV4UnionChecker.checkRead( schema, MetadataVersion.fromFlatbufID(result.getMessage().version())); return schema; } private ArrowDictionaryBatch readDictionary(MessageResult result) throws IOException { ArrowBuf bodyBuffer = result.getBodyBuffer(); // For zero-length batches, need an empty buffer to deserialize the batch if (bodyBuffer == null) { bodyBuffer = allocator.getEmpty(); } return MessageSerializer.deserializeDictionaryBatch(result.getMessage(), bodyBuffer); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/ArrowStreamWriter.java000066400000000000000000000152311515451615200325630ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import java.io.IOException; import java.io.OutputStream; import java.nio.channels.Channels; import java.nio.channels.WritableByteChannel; import java.util.HashMap; import java.util.Map; import java.util.Optional; import java.util.Set; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.compare.VectorEqualsVisitor; import org.apache.arrow.vector.compression.CompressionCodec; import org.apache.arrow.vector.compression.CompressionUtil; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.ipc.message.IpcOption; import org.apache.arrow.vector.ipc.message.MessageSerializer; /** Writer for the Arrow stream format to send ArrowRecordBatches over a WriteChannel. */ public class ArrowStreamWriter extends ArrowWriter { private final Map previousDictionaries = new HashMap<>(); /** * Construct an ArrowStreamWriter with an optional DictionaryProvider for the OutputStream. * * @param root Existing VectorSchemaRoot with vectors to be written. * @param provider DictionaryProvider for any vectors that are dictionary encoded. (Optional, can * be null) * @param out OutputStream for writing. */ public ArrowStreamWriter(VectorSchemaRoot root, DictionaryProvider provider, OutputStream out) { this(root, provider, Channels.newChannel(out)); } /** * Construct an ArrowStreamWriter with an optional DictionaryProvider for the WritableByteChannel. */ public ArrowStreamWriter( VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out) { this(root, provider, out, IpcOption.DEFAULT); } /** * Construct an ArrowStreamWriter with an optional DictionaryProvider for the WritableByteChannel. * * @param root Existing VectorSchemaRoot with vectors to be written. * @param provider DictionaryProvider for any vectors that are dictionary encoded. (Optional, can * be null) * @param option IPC write options * @param out WritableByteChannel for writing. */ public ArrowStreamWriter( VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out, IpcOption option) { super(root, provider, out, option); } /** * Construct an ArrowStreamWriter with compression enabled. * * @param root Existing VectorSchemaRoot with vectors to be written. * @param provider DictionaryProvider for any vectors that are dictionary encoded. (Optional, can * be null) * @param option IPC write options * @param compressionFactory Compression codec factory * @param codecType Codec type * @param out WritableByteChannel for writing. */ public ArrowStreamWriter( VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out, IpcOption option, CompressionCodec.Factory compressionFactory, CompressionUtil.CodecType codecType) { this(root, provider, out, option, compressionFactory, codecType, Optional.empty()); } /** * Construct an ArrowStreamWriter with compression enabled. * * @param root Existing VectorSchemaRoot with vectors to be written. * @param provider DictionaryProvider for any vectors that are dictionary encoded. (Optional, can * be null) * @param option IPC write options * @param compressionFactory Compression codec factory * @param codecType Codec type * @param compressionLevel Compression level * @param out WritableByteChannel for writing. */ public ArrowStreamWriter( VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out, IpcOption option, CompressionCodec.Factory compressionFactory, CompressionUtil.CodecType codecType, Optional compressionLevel) { super(root, provider, out, option, compressionFactory, codecType, compressionLevel); } /** * Write an EOS identifier to the WriteChannel. * * @param out Open WriteChannel with an active Arrow stream. * @param option IPC write option * @throws IOException on error */ public static void writeEndOfStream(WriteChannel out, IpcOption option) throws IOException { if (!option.write_legacy_ipc_format) { out.writeIntLittleEndian(MessageSerializer.IPC_CONTINUATION_TOKEN); } out.writeIntLittleEndian(0); } @Override protected void endInternal(WriteChannel out) throws IOException { writeEndOfStream(out, option); } @Override protected void ensureDictionariesWritten(DictionaryProvider provider, Set dictionaryIdsUsed) throws IOException { // write out any dictionaries that have changes for (long id : dictionaryIdsUsed) { Dictionary dictionary = provider.lookup(id); FieldVector vector = dictionary.getVector(); if (previousDictionaries.containsKey(id) && VectorEqualsVisitor.vectorEquals(vector, previousDictionaries.get(id))) { // Dictionary was previously written and hasn't changed continue; } writeDictionaryBatch(dictionary); // Store a copy of the vector in case it is later mutated if (previousDictionaries.containsKey(id)) { previousDictionaries.get(id).close(); } previousDictionaries.put(id, copyVector(vector)); } } @Override public void close() { super.close(); try { AutoCloseables.close(previousDictionaries.values()); } catch (Exception e) { throw new RuntimeException(e); } } private static FieldVector copyVector(FieldVector source) { FieldVector copy = source.getField().createVector(source.getAllocator()); copy.allocateNew(); for (int i = 0; i < source.getValueCount(); i++) { copy.copyFromSafe(i, i, source); } copy.setValueCount(source.getValueCount()); return copy; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/ArrowWriter.java000066400000000000000000000204171515451615200314110ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import java.io.IOException; import java.nio.channels.WritableByteChannel; import java.util.ArrayList; import java.util.Collections; import java.util.HashSet; import java.util.List; import java.util.Optional; import java.util.Set; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; import org.apache.arrow.vector.compression.CompressionCodec; import org.apache.arrow.vector.compression.CompressionUtil; import org.apache.arrow.vector.compression.NoCompressionCodec; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.ipc.message.ArrowBlock; import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.ipc.message.IpcOption; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.DictionaryUtility; import org.apache.arrow.vector.validate.MetadataV4UnionChecker; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** Abstract base class for implementing Arrow writers for IPC over a WriteChannel. */ public abstract class ArrowWriter implements AutoCloseable { protected static final Logger LOGGER = LoggerFactory.getLogger(ArrowWriter.class); // schema with fields in message format, not memory format protected final Schema schema; protected final WriteChannel out; private final VectorUnloader unloader; private final DictionaryProvider dictionaryProvider; private final Set dictionaryIdsUsed = new HashSet<>(); private final CompressionCodec.Factory compressionFactory; private final CompressionUtil.CodecType codecType; private final Optional compressionLevel; private boolean started = false; private boolean ended = false; private final CompressionCodec codec; protected IpcOption option; protected ArrowWriter( VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out) { this(root, provider, out, IpcOption.DEFAULT); } protected ArrowWriter( VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out, IpcOption option) { this( root, provider, out, option, NoCompressionCodec.Factory.INSTANCE, CompressionUtil.CodecType.NO_COMPRESSION, Optional.empty()); } /** * Note: fields are not closed when the writer is closed. * * @param root the vectors to write to the output * @param provider where to find the dictionaries * @param out the output where to write * @param option IPC write options * @param compressionFactory Compression codec factory * @param codecType Compression codec * @param compressionLevel Compression level */ protected ArrowWriter( VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out, IpcOption option, CompressionCodec.Factory compressionFactory, CompressionUtil.CodecType codecType, Optional compressionLevel) { this.out = new WriteChannel(out); this.option = option; this.dictionaryProvider = provider; this.compressionFactory = compressionFactory; this.codecType = codecType; this.compressionLevel = compressionLevel; this.codec = this.compressionLevel.isPresent() ? this.compressionFactory.createCodec(this.codecType, this.compressionLevel.get()) : this.compressionFactory.createCodec(this.codecType); this.unloader = new VectorUnloader(root, /*includeNullCount*/ true, codec, /*alignBuffers*/ true); List fields = new ArrayList<>(root.getSchema().getFields().size()); MetadataV4UnionChecker.checkForUnion( root.getSchema().getFields().iterator(), option.metadataVersion); // Convert fields with dictionaries to have dictionary type for (Field field : root.getSchema().getFields()) { fields.add(DictionaryUtility.toMessageFormat(field, provider, dictionaryIdsUsed)); } this.schema = new Schema(fields, root.getSchema().getCustomMetadata()); } public void start() throws IOException { ensureStarted(); } /** Writes the record batch currently loaded in this instance's VectorSchemaRoot. */ public void writeBatch() throws IOException { ensureStarted(); ensureDictionariesWritten(dictionaryProvider, dictionaryIdsUsed); try (ArrowRecordBatch batch = unloader.getRecordBatch()) { writeRecordBatch(batch); } } protected void writeDictionaryBatch(Dictionary dictionary) throws IOException { FieldVector vector = dictionary.getVector(); long id = dictionary.getEncoding().getId(); int count = vector.getValueCount(); VectorSchemaRoot dictRoot = new VectorSchemaRoot( Collections.singletonList(vector.getField()), Collections.singletonList(vector), count); VectorUnloader unloader = new VectorUnloader(dictRoot, /*includeNullCount*/ true, this.codec, /*alignBuffers*/ true); ArrowRecordBatch batch = unloader.getRecordBatch(); ArrowDictionaryBatch dictionaryBatch = new ArrowDictionaryBatch(id, batch, false); try { writeDictionaryBatch(dictionaryBatch); } finally { try { dictionaryBatch.close(); } catch (Exception e) { throw new RuntimeException("Error occurred while closing dictionary.", e); } } } protected ArrowBlock writeDictionaryBatch(ArrowDictionaryBatch batch) throws IOException { ArrowBlock block = MessageSerializer.serialize(out, batch, option); if (LOGGER.isDebugEnabled()) { LOGGER.debug( "DictionaryRecordBatch at {}, metadata: {}, body: {}", block.getOffset(), block.getMetadataLength(), block.getBodyLength()); } return block; } protected ArrowBlock writeRecordBatch(ArrowRecordBatch batch) throws IOException { ArrowBlock block = MessageSerializer.serialize(out, batch, option); if (LOGGER.isDebugEnabled()) { LOGGER.debug( "RecordBatch at {}, metadata: {}, body: {}", block.getOffset(), block.getMetadataLength(), block.getBodyLength()); } return block; } public void end() throws IOException { ensureStarted(); ensureEnded(); } public long bytesWritten() { return out.getCurrentPosition(); } private void ensureStarted() throws IOException { if (!started) { started = true; startInternal(out); // write the schema - for file formats this is duplicated in the footer, but matches // the streaming format MessageSerializer.serialize(out, schema, option); } } /** * Write dictionaries after schema and before recordBatches, dictionaries won't be written if * empty stream (only has schema data in IPC). */ protected abstract void ensureDictionariesWritten( DictionaryProvider provider, Set dictionaryIdsUsed) throws IOException; private void ensureEnded() throws IOException { if (!ended) { ended = true; endInternal(out); } } protected void startInternal(WriteChannel out) throws IOException {} protected void endInternal(WriteChannel out) throws IOException {} @Override public void close() { try { end(); out.close(); } catch (Exception e) { throw new RuntimeException(e); } } } InvalidArrowFileException.java000066400000000000000000000022111515451615200341130ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; /** * Exception indicating a problem with an Arrow File * (https://arrow.apache.org/docs/format/IPC.html#file-format). */ public class InvalidArrowFileException extends RuntimeException { private static final long serialVersionUID = 1L; public InvalidArrowFileException(String message) { super(message); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java000066400000000000000000001014701515451615200317550ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import static com.fasterxml.jackson.core.JsonToken.END_ARRAY; 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; import static org.apache.arrow.vector.BufferLayout.BufferType.TYPE; import static org.apache.arrow.vector.BufferLayout.BufferType.VALIDITY; import static org.apache.arrow.vector.BufferLayout.BufferType.VARIADIC_DATA_BUFFERS; import static org.apache.arrow.vector.BufferLayout.BufferType.VIEWS; import com.fasterxml.jackson.core.JsonParseException; import com.fasterxml.jackson.core.JsonParser; import com.fasterxml.jackson.core.JsonParser.Feature; import com.fasterxml.jackson.core.JsonToken; import com.fasterxml.jackson.databind.MapperFeature; import com.fasterxml.jackson.databind.MappingJsonFactory; import com.fasterxml.jackson.databind.ObjectMapper; import java.io.File; import java.io.IOException; import java.math.BigDecimal; import java.math.BigInteger; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Set; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BaseVariableWidthViewVector; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.BufferLayout.BufferType; import org.apache.arrow.vector.Decimal256Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.IntervalDayVector; import org.apache.arrow.vector.IntervalMonthDayNanoVector; import org.apache.arrow.vector.SmallIntVector; import org.apache.arrow.vector.TinyIntVector; import org.apache.arrow.vector.TypeLayout; 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.message.ArrowFieldNode; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.ArrowType.LargeListView; import org.apache.arrow.vector.types.pojo.ArrowType.ListView; import org.apache.arrow.vector.types.pojo.ArrowType.Union; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.DecimalUtility; import org.apache.arrow.vector.util.DictionaryUtility; import org.apache.commons.codec.DecoderException; import org.apache.commons.codec.binary.Hex; /** * A reader for JSON files that translates them into vectors. This reader is used for integration * tests. * *

    This class uses a streaming parser API, method naming tends to reflect this implementation * detail. */ public class JsonFileReader implements AutoCloseable, DictionaryProvider { private final JsonParser parser; private final BufferAllocator allocator; private Schema schema; private Map dictionaries; private Boolean started = false; /** * Constructs a new instance. * * @param inputFile The file to read. * @param allocator The allocator to use for allocating buffers. */ public JsonFileReader(File inputFile, BufferAllocator allocator) throws JsonParseException, IOException { super(); this.allocator = allocator; MappingJsonFactory jsonFactory = new MappingJsonFactory( new ObjectMapper() // ignore case for enums .configure(MapperFeature.ACCEPT_CASE_INSENSITIVE_ENUMS, true)); this.parser = jsonFactory.createParser(inputFile); // Allow reading NaN for floating point values this.parser.configure(Feature.ALLOW_NON_NUMERIC_NUMBERS, true); } @Override public Dictionary lookup(long id) { if (!started) { throw new IllegalStateException("Unable to lookup until after read() has started"); } return dictionaries.get(id); } @Override public Set getDictionaryIds() { return dictionaries.keySet(); } /** Reads the beginning (schema section) of the json file and returns it. */ public Schema start() throws JsonParseException, IOException { readToken(START_OBJECT); { Schema originalSchema = readNextField("schema", Schema.class); List fields = new ArrayList<>(); dictionaries = new HashMap<>(); // Convert fields with dictionaries to have the index type for (Field field : originalSchema.getFields()) { fields.add(DictionaryUtility.toMemoryFormat(field, allocator, dictionaries)); } this.schema = new Schema(fields, originalSchema.getCustomMetadata()); if (!dictionaries.isEmpty()) { nextFieldIs("dictionaries"); readDictionaryBatches(); } nextFieldIs("batches"); readToken(START_ARRAY); started = true; return this.schema; } } private void readDictionaryBatches() throws JsonParseException, IOException { readToken(START_ARRAY); JsonToken token = parser.nextToken(); boolean haveDictionaryBatch = token == START_OBJECT; while (haveDictionaryBatch) { // Lookup what dictionary for the batch about to be read long id = readNextField("id", Long.class); Dictionary dict = dictionaries.get(id); if (dict == null) { throw new IllegalArgumentException( "Dictionary with id: " + id + " missing encoding from schema Field"); } // Read the dictionary record batch nextFieldIs("data"); FieldVector vector = dict.getVector(); List fields = Collections.singletonList(vector.getField()); List vectors = Collections.singletonList(vector); VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, vector.getValueCount()); read(root); readToken(END_OBJECT); token = parser.nextToken(); haveDictionaryBatch = token == START_OBJECT; } if (token != END_ARRAY) { throw new IllegalArgumentException( "Invalid token: " + token + " expected end of array at " + parser.getTokenLocation()); } } /** Reads the next record batch from the file into root. */ public boolean read(VectorSchemaRoot root) throws IOException { JsonToken t = parser.nextToken(); if (t == START_OBJECT) { { int count = readNextField("count", Integer.class); nextFieldIs("columns"); readToken(START_ARRAY); { for (Field field : root.getSchema().getFields()) { FieldVector vector = root.getVector(field); readFromJsonIntoVector(field, vector); } } readToken(END_ARRAY); root.setRowCount(count); } readToken(END_OBJECT); return true; } else if (t == END_ARRAY) { root.setRowCount(0); return false; } else { throw new IllegalArgumentException("Invalid token: " + t); } } /** Returns the next record batch from the file. */ public VectorSchemaRoot read() throws IOException { JsonToken t = parser.nextToken(); if (t == START_OBJECT) { VectorSchemaRoot recordBatch = VectorSchemaRoot.create(schema, allocator); { int count = readNextField("count", Integer.class); recordBatch.setRowCount(count); nextFieldIs("columns"); readToken(START_ARRAY); { for (Field field : schema.getFields()) { FieldVector vector = recordBatch.getVector(field); readFromJsonIntoVector(field, vector); } } readToken(END_ARRAY); } readToken(END_OBJECT); return recordBatch; } else if (t == END_ARRAY) { return null; } else { throw new IllegalArgumentException("Invalid token: " + t); } } /** * Skips a number of record batches in the file. * * @param numBatches the number of batches to skip * @return the actual number of skipped batches. */ // This is currently called using JPype by the integration tests. public int skip(int numBatches) throws IOException { for (int i = 0; i < numBatches; ++i) { JsonToken t = parser.nextToken(); if (t == START_OBJECT) { parser.skipChildren(); assert parser.getCurrentToken() == END_OBJECT; } else if (t == END_ARRAY) { return i; } else { throw new IllegalArgumentException("Invalid token: " + t); } } return numBatches; } private abstract class BufferReader { protected abstract ArrowBuf read(BufferAllocator allocator, int count) throws IOException; ArrowBuf readBuffer(BufferAllocator allocator, int count) throws IOException { readToken(START_ARRAY); ArrowBuf buf = read(allocator, count); readToken(END_ARRAY); return buf; } } /** * Read all the variadic data buffers from the parser. * * @param allocator BufferAllocator * @param variadicBuffersCount Number of variadic buffers * @return List of ArrowBuf * @throws IOException throws IOException in a failure */ List readVariadicBuffers(BufferAllocator allocator, int variadicBuffersCount) throws IOException { readToken(START_ARRAY); ArrayList dataBuffers = new ArrayList<>(variadicBuffersCount); for (int i = 0; i < variadicBuffersCount; i++) { parser.nextToken(); final byte[] value; String variadicStr = parser.readValueAs(String.class); if (variadicStr == null) { value = new byte[0]; } else { value = decodeHexSafe(variadicStr); } ArrowBuf buf = allocator.buffer(value.length); buf.writeBytes(value); dataBuffers.add(buf); } readToken(END_ARRAY); return dataBuffers; } private ArrowBuf readViewBuffers( BufferAllocator allocator, int count, List variadicBufferIndices, MinorType type) throws IOException { readToken(START_ARRAY); ArrayList values = new ArrayList<>(count); long bufferSize = 0L; for (int i = 0; i < count; i++) { readToken(START_OBJECT); final int length = readNextField("SIZE", Integer.class); byte[] value; if (length > BaseVariableWidthViewVector.INLINE_SIZE) { // PREFIX_HEX final byte[] prefix = decodeHexSafe(readNextField("PREFIX_HEX", String.class)); // BUFFER_INDEX final int bufferIndex = readNextField("BUFFER_INDEX", Integer.class); if (variadicBufferIndices.isEmpty()) { variadicBufferIndices.add(bufferIndex); } else { int lastBufferIndex = variadicBufferIndices.get(variadicBufferIndices.size() - 1); if (lastBufferIndex != bufferIndex) { variadicBufferIndices.add(bufferIndex); } } // OFFSET final int offset = readNextField("OFFSET", Integer.class); ByteBuffer buffer = ByteBuffer.allocate(BaseVariableWidthViewVector.ELEMENT_SIZE) .order(ByteOrder.LITTLE_ENDIAN); // Allocate a ByteBuffer of size 16 bytes buffer.putInt(length); // Write 'length' to bytes 0-3 buffer.put(prefix); // Write 'prefix' to bytes 4-7 buffer.putInt(bufferIndex); // Write 'bufferIndex' to bytes 8-11 buffer.putInt(offset); // Write 'offset' to bytes 12-15 value = buffer.array(); // Convert the ByteBuffer to a byte array } else { // in-line ByteBuffer buffer = ByteBuffer.allocate(BaseVariableWidthViewVector.ELEMENT_SIZE) .order(ByteOrder.LITTLE_ENDIAN); // Allocate a ByteBuffer of size 16 bytes buffer.putInt(length); // Write 'length' to bytes 0-3 // INLINE if (type == MinorType.VIEWVARCHAR) { buffer.put(readNextField("INLINED", String.class).getBytes(StandardCharsets.UTF_8)); } else { String inlined = readNextField("INLINED", String.class); if (inlined == null) { buffer.put(new byte[length]); } else { buffer.put(decodeHexSafe(inlined)); } } value = buffer.array(); // Convert the ByteBuffer to a byte array } values.add(value); bufferSize += value.length; readToken(END_OBJECT); } ArrowBuf buf = allocator.buffer(bufferSize); for (byte[] value : values) { buf.writeBytes(value); } readToken(END_ARRAY); return buf; } private class BufferHelper { BufferReader BIT = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final int bufferSize = getValidityBufferSizeFromCount(count); ArrowBuf buf = allocator.buffer(bufferSize); // C++ integration test fails without this. buf.setZero(0, bufferSize); for (int i = 0; i < count; i++) { parser.nextToken(); BitVectorHelper.setValidityBit(buf, i, parser.readValueAs(Boolean.class) ? 1 : 0); } buf.writerIndex(bufferSize); return buf; } }; BufferReader DAY_MILLIS = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * IntervalDayVector.TYPE_WIDTH; ArrowBuf buf = allocator.buffer(size); for (int i = 0; i < count; i++) { readToken(START_OBJECT); buf.writeInt(readNextField("days", Integer.class)); buf.writeInt(readNextField("milliseconds", Integer.class)); readToken(END_OBJECT); } return buf; } }; BufferReader MONTH_DAY_NANOS = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * IntervalMonthDayNanoVector.TYPE_WIDTH; ArrowBuf buf = allocator.buffer(size); for (int i = 0; i < count; i++) { readToken(START_OBJECT); buf.writeInt(readNextField("months", Integer.class)); buf.writeInt(readNextField("days", Integer.class)); buf.writeLong(readNextField("nanoseconds", Long.class)); readToken(END_OBJECT); } return buf; } }; BufferReader INT1 = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * TinyIntVector.TYPE_WIDTH; ArrowBuf buf = allocator.buffer(size); for (int i = 0; i < count; i++) { parser.nextToken(); buf.writeByte(parser.getByteValue()); } return buf; } }; BufferReader INT2 = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * SmallIntVector.TYPE_WIDTH; ArrowBuf buf = allocator.buffer(size); for (int i = 0; i < count; i++) { parser.nextToken(); buf.writeShort(parser.getShortValue()); } return buf; } }; BufferReader INT4 = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * IntVector.TYPE_WIDTH; ArrowBuf buf = allocator.buffer(size); for (int i = 0; i < count; i++) { parser.nextToken(); buf.writeInt(parser.getIntValue()); } return buf; } }; BufferReader INT8 = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * BigIntVector.TYPE_WIDTH; ArrowBuf buf = allocator.buffer(size); for (int i = 0; i < count; i++) { parser.nextToken(); String value = parser.getValueAsString(); buf.writeLong(Long.valueOf(value)); } return buf; } }; BufferReader UINT1 = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * TinyIntVector.TYPE_WIDTH; ArrowBuf buf = allocator.buffer(size); for (int i = 0; i < count; i++) { parser.nextToken(); buf.writeByte(parser.getShortValue() & 0xFF); } return buf; } }; BufferReader UINT2 = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * SmallIntVector.TYPE_WIDTH; ArrowBuf buf = allocator.buffer(size); for (int i = 0; i < count; i++) { parser.nextToken(); buf.writeShort(parser.getIntValue() & 0xFFFF); } return buf; } }; BufferReader UINT4 = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * IntVector.TYPE_WIDTH; ArrowBuf buf = allocator.buffer(size); for (int i = 0; i < count; i++) { parser.nextToken(); buf.writeInt((int) parser.getLongValue()); } return buf; } }; BufferReader UINT8 = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * BigIntVector.TYPE_WIDTH; ArrowBuf buf = allocator.buffer(size); for (int i = 0; i < count; i++) { parser.nextToken(); BigInteger value = new BigInteger(parser.getValueAsString()); buf.writeLong(value.longValue()); } return buf; } }; BufferReader FLOAT4 = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * Float4Vector.TYPE_WIDTH; ArrowBuf buf = allocator.buffer(size); for (int i = 0; i < count; i++) { parser.nextToken(); buf.writeFloat(parser.getFloatValue()); } return buf; } }; BufferReader FLOAT8 = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * Float8Vector.TYPE_WIDTH; ArrowBuf buf = allocator.buffer(size); for (int i = 0; i < count; i++) { parser.nextToken(); buf.writeDouble(parser.getDoubleValue()); } return buf; } }; BufferReader DECIMAL = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * DecimalVector.TYPE_WIDTH; ArrowBuf buf = allocator.buffer(size); for (int i = 0; i < count; i++) { parser.nextToken(); BigDecimal decimalValue = new BigDecimal(parser.readValueAs(String.class)); DecimalUtility.writeBigDecimalToArrowBuf( decimalValue, buf, i, DecimalVector.TYPE_WIDTH); } buf.writerIndex(size); return buf; } }; BufferReader DECIMAL256 = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * Decimal256Vector.TYPE_WIDTH; ArrowBuf buf = allocator.buffer(size); for (int i = 0; i < count; i++) { parser.nextToken(); BigDecimal decimalValue = new BigDecimal(parser.readValueAs(String.class)); DecimalUtility.writeBigDecimalToArrowBuf( decimalValue, buf, i, Decimal256Vector.TYPE_WIDTH); } buf.writerIndex(size); return buf; } }; ArrowBuf readBinaryValues(BufferAllocator allocator, int count) throws IOException { ArrayList values = new ArrayList<>(count); long bufferSize = 0L; for (int i = 0; i < count; i++) { parser.nextToken(); final byte[] value = decodeHexSafe(parser.readValueAs(String.class)); values.add(value); bufferSize += value.length; } ArrowBuf buf = allocator.buffer(bufferSize); for (byte[] value : values) { buf.writeBytes(value); } return buf; } ArrowBuf readStringValues(BufferAllocator allocator, int count) throws IOException { ArrayList values = new ArrayList<>(count); long bufferSize = 0L; for (int i = 0; i < count; i++) { parser.nextToken(); final byte[] value = parser.getValueAsString().getBytes(StandardCharsets.UTF_8); values.add(value); bufferSize += value.length; } ArrowBuf buf = allocator.buffer(bufferSize); for (byte[] value : values) { buf.writeBytes(value); } return buf; } BufferReader FIXEDSIZEBINARY = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { return readBinaryValues(allocator, count); } }; BufferReader VARCHAR = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { return readStringValues(allocator, count); } }; BufferReader LARGEVARCHAR = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { return readStringValues(allocator, count); } }; BufferReader VARBINARY = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { return readBinaryValues(allocator, count); } }; BufferReader LARGEVARBINARY = new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { return readBinaryValues(allocator, count); } }; } private List readIntoBuffer( BufferAllocator allocator, BufferType bufferType, MinorType type, int count, List variadicBufferIndices) throws IOException { ArrowBuf buf; BufferHelper helper = new BufferHelper(); BufferReader reader; if (bufferType.equals(VALIDITY)) { reader = helper.BIT; } else if (bufferType.equals(OFFSET) || bufferType.equals(SIZE)) { if (type == MinorType.LARGELIST || type == MinorType.LARGEVARCHAR || type == MinorType.LARGEVARBINARY || type == MinorType.LARGELISTVIEW) { reader = helper.INT8; } else { reader = helper.INT4; } } else if (bufferType.equals(TYPE)) { reader = helper.INT1; } else if (bufferType.equals(DATA)) { switch (type) { case BIT: reader = helper.BIT; break; case TINYINT: reader = helper.INT1; break; case SMALLINT: reader = helper.INT2; break; case INT: reader = helper.INT4; break; case BIGINT: reader = helper.INT8; break; case UINT1: reader = helper.UINT1; break; case UINT2: reader = helper.UINT2; break; case UINT4: reader = helper.UINT4; break; case UINT8: reader = helper.UINT8; break; case FLOAT4: reader = helper.FLOAT4; break; case FLOAT8: reader = helper.FLOAT8; break; case DECIMAL: reader = helper.DECIMAL; break; case DECIMAL256: reader = helper.DECIMAL256; break; case FIXEDSIZEBINARY: reader = helper.FIXEDSIZEBINARY; break; case VARCHAR: reader = helper.VARCHAR; break; case LARGEVARCHAR: reader = helper.LARGEVARCHAR; break; case VARBINARY: reader = helper.VARBINARY; break; case LARGEVARBINARY: reader = helper.LARGEVARBINARY; break; case DATEDAY: reader = helper.INT4; break; case DATEMILLI: reader = helper.INT8; break; case TIMESEC: case TIMEMILLI: reader = helper.INT4; break; case TIMEMICRO: case TIMENANO: reader = helper.INT8; break; case TIMESTAMPNANO: case TIMESTAMPMICRO: case TIMESTAMPMILLI: case TIMESTAMPSEC: case TIMESTAMPNANOTZ: case TIMESTAMPMICROTZ: case TIMESTAMPMILLITZ: case TIMESTAMPSECTZ: reader = helper.INT8; break; case INTERVALYEAR: reader = helper.INT4; break; case INTERVALDAY: reader = helper.DAY_MILLIS; break; case INTERVALMONTHDAYNANO: reader = helper.MONTH_DAY_NANOS; break; case DURATION: reader = helper.INT8; break; default: throw new UnsupportedOperationException("Cannot read array of type " + type); } } else if (bufferType.equals(VIEWS)) { return Collections.singletonList( readViewBuffers(allocator, count, variadicBufferIndices, type)); } else if (bufferType.equals(VARIADIC_DATA_BUFFERS)) { return readVariadicBuffers(allocator, variadicBufferIndices.size()); } else { throw new InvalidArrowFileException("Unrecognized buffer type " + bufferType); } buf = reader.readBuffer(allocator, count); Preconditions.checkNotNull(buf); return Collections.singletonList(buf); } private void readFromJsonIntoVector(Field field, FieldVector vector) throws IOException { ArrowType type = field.getType(); TypeLayout typeLayout = TypeLayout.getTypeLayout(type); List vectorTypes = typeLayout.getBufferTypes(); List vectorBuffers = new ArrayList<>(vectorTypes.size()); List variadicBufferIndices = new ArrayList<>(); if (!typeLayout.isFixedBufferCount()) { vectorTypes.add(VARIADIC_DATA_BUFFERS); } /* * The order of inner buffers is: * Fixed width vector: * -- validity buffer * -- data buffer * Variable width vector: * -- validity buffer * -- offset buffer * -- data buffer * * This is similar to what getFieldInnerVectors() used to give but now that we don't have * inner vectors anymore, we will work directly at the buffer level -- populate buffers * locally as we read from Json parser and do loadFieldBuffers on the vector followed by * releasing the local buffers. */ readToken(START_OBJECT); { // If currently reading dictionaries, field name is not important so don't check String name = readNextField("name", String.class); if (started && !Objects.equals(field.getName(), name)) { throw new IllegalArgumentException( "Expected field " + field.getName() + " but got " + name); } /* Initialize the vector with required capacity but don't allocateNew since we would * be doing loadFieldBuffers. */ int valueCount = readNextField("count", Integer.class); vector.setInitialCapacity(valueCount); for (int v = 0; v < vectorTypes.size(); v++) { BufferType bufferType = vectorTypes.get(v); nextFieldIs(bufferType.getName()); int innerBufferValueCount = valueCount; if (bufferType.equals(OFFSET) && !(type instanceof Union) && !(type instanceof ListView) && !(type instanceof LargeListView)) { /* offset buffer has 1 additional value capacity except for dense unions and ListView */ innerBufferValueCount = valueCount + 1; } vectorBuffers.addAll( readIntoBuffer( allocator, bufferType, vector.getMinorType(), innerBufferValueCount, variadicBufferIndices)); } int nullCount; if (type instanceof ArrowType.RunEndEncoded || type instanceof Union) { nullCount = 0; } else if (type instanceof ArrowType.Null) { nullCount = valueCount; } else { nullCount = BitVectorHelper.getNullCount(vectorBuffers.get(0), valueCount); } final ArrowFieldNode fieldNode = new ArrowFieldNode(valueCount, nullCount); vector.loadFieldBuffers(fieldNode, vectorBuffers); /* read child vectors (if any) */ List fields = field.getChildren(); if (!fields.isEmpty()) { List vectorChildren = vector.getChildrenFromFields(); if (fields.size() != vectorChildren.size()) { throw new IllegalArgumentException( "fields and children are not the same size: " + fields.size() + " != " + vectorChildren.size()); } nextFieldIs("children"); readToken(START_ARRAY); for (int i = 0; i < fields.size(); i++) { Field childField = fields.get(i); FieldVector childVector = vectorChildren.get(i); readFromJsonIntoVector(childField, childVector); } readToken(END_ARRAY); } } readToken(END_OBJECT); for (ArrowBuf buffer : vectorBuffers) { buffer.getReferenceManager().release(); } } private byte[] decodeHexSafe(String hexString) throws IOException { try { return Hex.decodeHex(hexString.toCharArray()); } catch (DecoderException e) { throw new IOException("Unable to decode hex string: " + hexString, e); } } @Override public void close() throws IOException { parser.close(); if (dictionaries != null) { for (Dictionary dictionary : dictionaries.values()) { dictionary.getVector().close(); } } } private T readNextField(String expectedFieldName, Class c) throws IOException, JsonParseException { nextFieldIs(expectedFieldName); parser.nextToken(); return parser.readValueAs(c); } private void nextFieldIs(String expectedFieldName) throws IOException, JsonParseException { String name = parser.nextFieldName(); if (name == null || !name.equals(expectedFieldName)) { throw new IllegalStateException("Expected " + expectedFieldName + " but got " + name); } } private void readToken(JsonToken expected) throws JsonParseException, IOException { JsonToken t = parser.nextToken(); if (t != expected) { throw new IllegalStateException("Expected " + expected + " but got " + t); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileWriter.java000066400000000000000000000610401515451615200320250ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import static org.apache.arrow.vector.BufferLayout.BufferType.*; import com.fasterxml.jackson.core.JsonEncoding; import com.fasterxml.jackson.core.JsonGenerator; import com.fasterxml.jackson.core.util.DefaultPrettyPrinter; import com.fasterxml.jackson.core.util.DefaultPrettyPrinter.NopIndenter; import com.fasterxml.jackson.databind.MappingJsonFactory; import java.io.File; import java.io.IOException; import java.math.BigDecimal; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.HashSet; import java.util.List; import java.util.Set; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.util.Preconditions; 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.BitVectorHelper; import org.apache.arrow.vector.BufferLayout.BufferType; 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.DurationVector; 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.IntervalDayVector; import org.apache.arrow.vector.IntervalMonthDayNanoVector; import org.apache.arrow.vector.IntervalYearVector; 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.TypeLayout; 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.VectorSchemaRoot; import org.apache.arrow.vector.complex.BaseLargeRepeatedValueViewVector; import org.apache.arrow.vector.complex.BaseRepeatedValueViewVector; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.DecimalUtility; import org.apache.arrow.vector.util.DictionaryUtility; import org.apache.commons.codec.binary.Hex; /** * A writer that converts binary Vectors into an internal, unstable JSON format suitable * for integration testing. * *

    This writer does NOT implement a JSON dataset format like JSONL. */ public class JsonFileWriter implements AutoCloseable { /** Configuration POJO for writing JSON files. */ public static final class JSONWriteConfig { private final boolean pretty; private JSONWriteConfig(boolean pretty) { this.pretty = pretty; } private JSONWriteConfig() { this.pretty = false; } public JSONWriteConfig pretty(boolean pretty) { return new JSONWriteConfig(pretty); } } public static JSONWriteConfig config() { return new JSONWriteConfig(); } private final JsonGenerator generator; private Schema schema; /** Constructs a new writer that will output to outputFile. */ public JsonFileWriter(File outputFile) throws IOException { this(outputFile, config()); } /** Constructs a new writer that will output to outputFile with the given options. */ public JsonFileWriter(File outputFile, JSONWriteConfig config) throws IOException { MappingJsonFactory jsonFactory = new MappingJsonFactory(); this.generator = jsonFactory.createGenerator(outputFile, JsonEncoding.UTF8); if (config.pretty) { DefaultPrettyPrinter prettyPrinter = new DefaultPrettyPrinter(); prettyPrinter.indentArraysWith(NopIndenter.instance); this.generator.setPrettyPrinter(prettyPrinter); } // Allow writing of floating point NaN values not as strings this.generator.configure(JsonGenerator.Feature.QUOTE_NON_NUMERIC_NUMBERS, false); } /** Writes out the "header" of the file including the schema and any dictionaries required. */ public void start(Schema schema, DictionaryProvider provider) throws IOException { List fields = new ArrayList<>(schema.getFields().size()); Set dictionaryIdsUsed = new HashSet<>(); this.schema = schema; // Store original Schema to ensure batches written match // Convert fields with dictionaries to have dictionary type for (Field field : schema.getFields()) { fields.add(DictionaryUtility.toMessageFormat(field, provider, dictionaryIdsUsed)); } Schema updatedSchema = new Schema(fields, schema.getCustomMetadata()); generator.writeStartObject(); generator.writeObjectField("schema", updatedSchema); // Write all dictionaries that were used if (!dictionaryIdsUsed.isEmpty()) { writeDictionaryBatches(generator, dictionaryIdsUsed, provider); } // Start writing of record batches generator.writeArrayFieldStart("batches"); } private void writeDictionaryBatches( JsonGenerator generator, Set dictionaryIdsUsed, DictionaryProvider provider) throws IOException { generator.writeArrayFieldStart("dictionaries"); for (Long id : dictionaryIdsUsed) { generator.writeStartObject(); generator.writeObjectField("id", id); generator.writeFieldName("data"); Dictionary dictionary = provider.lookup(id); FieldVector vector = dictionary.getVector(); List fields = Collections.singletonList(vector.getField()); List vectors = Collections.singletonList(vector); VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, vector.getValueCount()); writeBatch(root); generator.writeEndObject(); } generator.writeEndArray(); } /** Writes the record batch to the JSON file. */ public void write(VectorSchemaRoot recordBatch) throws IOException { if (!recordBatch.getSchema().equals(schema)) { throw new IllegalArgumentException("record batches must have the same schema: " + schema); } writeBatch(recordBatch); } private void writeBatch(VectorSchemaRoot recordBatch) throws IOException { generator.writeStartObject(); { generator.writeObjectField("count", recordBatch.getRowCount()); generator.writeArrayFieldStart("columns"); for (Field field : recordBatch.getSchema().getFields()) { FieldVector vector = recordBatch.getVector(field); writeFromVectorIntoJson(field, vector); } generator.writeEndArray(); } generator.writeEndObject(); } private void writeFromVectorIntoJson(Field field, FieldVector vector) throws IOException { TypeLayout typeLayout = TypeLayout.getTypeLayout(field.getType()); List vectorTypes = typeLayout.getBufferTypes(); List vectorBuffers = vector.getFieldBuffers(); if (typeLayout.isFixedBufferCount()) { if (vectorTypes.size() != vectorBuffers.size()) { throw new IllegalArgumentException( "vector types and inner vector buffers are not the same size: " + vectorTypes.size() + " != " + vectorBuffers.size()); } } else { vectorTypes.add(VARIADIC_DATA_BUFFERS); } generator.writeStartObject(); { generator.writeObjectField("name", field.getName()); int valueCount = vector.getValueCount(); generator.writeObjectField("count", valueCount); for (int v = 0; v < vectorTypes.size(); v++) { BufferType bufferType = vectorTypes.get(v); ArrowBuf vectorBuffer = vectorBuffers.get(v); // Note that in JSON format we cannot have VARIADIC_DATA_BUFFERS repeated, // thus the values are only written to a single entity. generator.writeArrayFieldStart(bufferType.getName()); final int bufferValueCount = (bufferType.equals(OFFSET) && vector.getMinorType() != MinorType.DENSEUNION && vector.getMinorType() != MinorType.LISTVIEW && vector.getMinorType() != MinorType.LARGELISTVIEW) ? valueCount + 1 : valueCount; for (int i = 0; i < bufferValueCount; i++) { if (bufferType.equals(DATA) && (vector.getMinorType() == MinorType.VARCHAR || vector.getMinorType() == MinorType.VARBINARY)) { writeValueToGenerator(bufferType, vectorBuffer, vectorBuffers.get(v - 1), vector, i); } else if (bufferType.equals(VIEWS) && (vector.getMinorType() == MinorType.VIEWVARCHAR || vector.getMinorType() == MinorType.VIEWVARBINARY)) { // writing views ArrowBuf viewBuffer = vectorBuffers.get(1); List dataBuffers = vectorBuffers.subList(v + 1, vectorBuffers.size()); writeValueToViewGenerator(bufferType, viewBuffer, dataBuffers, vector, i); } else if (bufferType.equals(VARIADIC_DATA_BUFFERS) && (vector.getMinorType() == MinorType.VIEWVARCHAR || vector.getMinorType() == MinorType.VIEWVARBINARY)) { ArrowBuf viewBuffer = vectorBuffers.get(1); // check if this is v-1 List dataBuffers = vectorBuffers.subList(v, vectorBuffers.size()); if (!dataBuffers.isEmpty()) { writeValueToDataBufferGenerator(bufferType, viewBuffer, dataBuffers, vector); // The variadic buffers are written at once and doesn't require iterating for // each index. // So, break the loop. break; } } else if (bufferType.equals(OFFSET) && vector.getValueCount() == 0 && (vector.getMinorType() == MinorType.LIST || vector.getMinorType() == MinorType.LISTVIEW || vector.getMinorType() == MinorType.MAP || vector.getMinorType() == MinorType.VARBINARY || vector.getMinorType() == MinorType.VARCHAR)) { // Empty vectors may not have allocated an offsets buffer try (ArrowBuf vectorBufferTmp = vector.getAllocator().buffer(4)) { vectorBufferTmp.setInt(0, 0); writeValueToGenerator(bufferType, vectorBufferTmp, null, vector, i); } } else if (bufferType.equals(OFFSET) && vector.getValueCount() == 0 && (vector.getMinorType() == MinorType.LARGELIST || vector.getMinorType() == MinorType.LARGELISTVIEW || vector.getMinorType() == MinorType.LARGEVARBINARY || vector.getMinorType() == MinorType.LARGEVARCHAR)) { // Empty vectors may not have allocated an offsets buffer try (ArrowBuf vectorBufferTmp = vector.getAllocator().buffer(8)) { vectorBufferTmp.setLong(0, 0); writeValueToGenerator(bufferType, vectorBufferTmp, null, vector, i); } } else { writeValueToGenerator(bufferType, vectorBuffer, null, vector, i); } } generator.writeEndArray(); } List fields = field.getChildren(); List children = vector.getChildrenFromFields(); if (fields.size() != children.size()) { throw new IllegalArgumentException( "fields and children are not the same size: " + fields.size() + " != " + children.size()); } if (fields.size() > 0) { generator.writeArrayFieldStart("children"); for (int i = 0; i < fields.size(); i++) { Field childField = fields.get(i); FieldVector childVector = children.get(i); writeFromVectorIntoJson(childField, childVector); } generator.writeEndArray(); } } generator.writeEndObject(); } /** * Get data of a view by index. * * @param viewBuffer view buffer * @param dataBuffers data buffers * @param index index of the view * @return byte array of the view */ private byte[] getView(final ArrowBuf viewBuffer, final List dataBuffers, int index) { final int dataLength = viewBuffer.getInt((long) index * BaseVariableWidthViewVector.ELEMENT_SIZE); byte[] result = new byte[dataLength]; final int inlineSize = BaseVariableWidthViewVector.INLINE_SIZE; final int elementSize = BaseVariableWidthViewVector.ELEMENT_SIZE; final int lengthWidth = BaseVariableWidthViewVector.LENGTH_WIDTH; final int prefixWidth = BaseVariableWidthViewVector.PREFIX_WIDTH; final int bufIndexWidth = BaseVariableWidthViewVector.BUF_INDEX_WIDTH; if (dataLength > inlineSize) { // data is in the data buffer // get buffer index final int bufferIndex = viewBuffer.getInt(((long) index * elementSize) + lengthWidth + prefixWidth); // get data offset final int dataOffset = viewBuffer.getInt( ((long) index * elementSize) + lengthWidth + prefixWidth + bufIndexWidth); dataBuffers.get(bufferIndex).getBytes(dataOffset, result, 0, dataLength); } else { // data is in the view buffer viewBuffer.getBytes((long) index * elementSize + lengthWidth, result, 0, dataLength); } return result; } private void writeValueToViewGenerator( BufferType bufferType, ArrowBuf viewBuffer, List dataBuffers, FieldVector vector, final int index) throws IOException { Preconditions.checkNotNull(viewBuffer); byte[] b = getView(viewBuffer, dataBuffers, index); final int elementSize = BaseVariableWidthViewVector.ELEMENT_SIZE; final int lengthWidth = BaseVariableWidthViewVector.LENGTH_WIDTH; final int prefixWidth = BaseVariableWidthViewVector.PREFIX_WIDTH; final int bufIndexWidth = BaseVariableWidthViewVector.BUF_INDEX_WIDTH; final int length = viewBuffer.getInt((long) index * elementSize); generator.writeStartObject(); generator.writeFieldName("SIZE"); generator.writeObject(length); if (length > 12) { byte[] prefix = Arrays.copyOfRange(b, 0, prefixWidth); final int bufferIndex = viewBuffer.getInt(((long) index * elementSize) + lengthWidth + prefixWidth); // get data offset final int dataOffset = viewBuffer.getInt( ((long) index * elementSize) + lengthWidth + prefixWidth + bufIndexWidth); generator.writeFieldName("PREFIX_HEX"); generator.writeString(Hex.encodeHexString(prefix)); generator.writeFieldName("BUFFER_INDEX"); generator.writeObject(bufferIndex); generator.writeFieldName("OFFSET"); generator.writeObject(dataOffset); } else { generator.writeFieldName("INLINED"); if (vector.getMinorType() == MinorType.VIEWVARCHAR) { generator.writeString(new String(b, "UTF-8")); } else { generator.writeString(Hex.encodeHexString(b)); } } generator.writeEndObject(); } private void writeValueToDataBufferGenerator( BufferType bufferType, ArrowBuf viewBuffer, List dataBuffers, FieldVector vector) throws IOException { if (bufferType.equals(VARIADIC_DATA_BUFFERS)) { Preconditions.checkNotNull(viewBuffer); Preconditions.checkArgument(!dataBuffers.isEmpty()); for (int i = 0; i < dataBuffers.size(); i++) { ArrowBuf dataBuf = dataBuffers.get(i); byte[] result = new byte[(int) dataBuf.writerIndex()]; dataBuf.getBytes(0, result); if (result != null) { generator.writeString(Hex.encodeHexString(result)); } } } } private void writeValueToGenerator( BufferType bufferType, ArrowBuf buffer, ArrowBuf offsetBuffer, FieldVector vector, final int index) throws IOException { if (bufferType.equals(TYPE)) { generator.writeNumber(buffer.getByte(index * TinyIntVector.TYPE_WIDTH)); } else if (bufferType.equals(OFFSET)) { switch (vector.getMinorType()) { case VARCHAR: case VARBINARY: case LIST: case MAP: generator.writeNumber(buffer.getInt((long) index * BaseVariableWidthVector.OFFSET_WIDTH)); break; case LISTVIEW: generator.writeNumber( buffer.getInt((long) index * BaseRepeatedValueViewVector.OFFSET_WIDTH)); break; case LARGELISTVIEW: generator.writeNumber( buffer.getInt((long) index * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); break; case LARGELIST: case LARGEVARBINARY: case LARGEVARCHAR: generator.writeNumber( buffer.getLong((long) index * BaseLargeVariableWidthVector.OFFSET_WIDTH)); break; default: throw new IllegalArgumentException("Type has no offset buffer: " + vector.getField()); } } else if (bufferType.equals(VALIDITY)) { generator.writeNumber(vector.isNull(index) ? 0 : 1); } else if (bufferType.equals(DATA)) { switch (vector.getMinorType()) { case TINYINT: generator.writeNumber(TinyIntVector.get(buffer, index)); break; case SMALLINT: generator.writeNumber(SmallIntVector.get(buffer, index)); break; case INT: generator.writeNumber(IntVector.get(buffer, index)); break; case BIGINT: generator.writeString(String.valueOf(BigIntVector.get(buffer, index))); break; case UINT1: generator.writeNumber(UInt1Vector.getNoOverflow(buffer, index)); break; case UINT2: generator.writeNumber(UInt2Vector.get(buffer, index)); break; case UINT4: generator.writeNumber(UInt4Vector.getNoOverflow(buffer, index)); break; case UINT8: generator.writeString(UInt8Vector.getNoOverflow(buffer, index).toString()); break; case FLOAT4: generator.writeNumber(Float4Vector.get(buffer, index)); break; case FLOAT8: generator.writeNumber(Float8Vector.get(buffer, index)); break; case DATEDAY: generator.writeNumber(DateDayVector.get(buffer, index)); break; case DATEMILLI: generator.writeNumber(DateMilliVector.get(buffer, index)); break; case TIMESEC: generator.writeNumber(TimeSecVector.get(buffer, index)); break; case TIMEMILLI: generator.writeNumber(TimeMilliVector.get(buffer, index)); break; case TIMEMICRO: generator.writeNumber(TimeMicroVector.get(buffer, index)); break; case TIMENANO: generator.writeNumber(TimeNanoVector.get(buffer, index)); break; case TIMESTAMPSEC: generator.writeNumber(TimeStampSecVector.get(buffer, index)); break; case TIMESTAMPMILLI: generator.writeNumber(TimeStampMilliVector.get(buffer, index)); break; case TIMESTAMPMICRO: generator.writeNumber(TimeStampMicroVector.get(buffer, index)); break; case TIMESTAMPNANO: generator.writeNumber(TimeStampNanoVector.get(buffer, index)); break; case TIMESTAMPSECTZ: generator.writeNumber(TimeStampSecTZVector.get(buffer, index)); break; case TIMESTAMPMILLITZ: generator.writeNumber(TimeStampMilliTZVector.get(buffer, index)); break; case TIMESTAMPMICROTZ: generator.writeNumber(TimeStampMicroTZVector.get(buffer, index)); break; case TIMESTAMPNANOTZ: generator.writeNumber(TimeStampNanoTZVector.get(buffer, index)); break; case DURATION: generator.writeNumber(DurationVector.get(buffer, index)); break; case INTERVALYEAR: generator.writeNumber(IntervalYearVector.getTotalMonths(buffer, index)); break; case INTERVALDAY: generator.writeStartObject(); generator.writeObjectField("days", IntervalDayVector.getDays(buffer, index)); generator.writeObjectField( "milliseconds", IntervalDayVector.getMilliseconds(buffer, index)); generator.writeEndObject(); break; case INTERVALMONTHDAYNANO: generator.writeStartObject(); generator.writeObjectField("months", IntervalMonthDayNanoVector.getMonths(buffer, index)); generator.writeObjectField("days", IntervalMonthDayNanoVector.getDays(buffer, index)); generator.writeObjectField( "nanoseconds", IntervalMonthDayNanoVector.getNanoseconds(buffer, index)); generator.writeEndObject(); break; case BIT: generator.writeNumber(BitVectorHelper.get(buffer, index)); break; case VARBINARY: { Preconditions.checkNotNull(offsetBuffer); String hexString = Hex.encodeHexString(BaseVariableWidthVector.get(buffer, offsetBuffer, index)); generator.writeObject(hexString); break; } case FIXEDSIZEBINARY: int byteWidth = ((FixedSizeBinaryVector) vector).getByteWidth(); String fixedSizeHexString = Hex.encodeHexString(FixedSizeBinaryVector.get(buffer, index, byteWidth)); generator.writeObject(fixedSizeHexString); break; case VARCHAR: { Preconditions.checkNotNull(offsetBuffer); byte[] b = (BaseVariableWidthVector.get(buffer, offsetBuffer, index)); generator.writeString(new String(b, "UTF-8")); break; } case DECIMAL: { int scale = ((DecimalVector) vector).getScale(); BigDecimal decimalValue = DecimalUtility.getBigDecimalFromArrowBuf( buffer, index, scale, DecimalVector.TYPE_WIDTH); // We write the unscaled value, because the scale is stored in the type metadata. generator.writeString(decimalValue.unscaledValue().toString()); break; } case DECIMAL256: { int scale = ((Decimal256Vector) vector).getScale(); BigDecimal decimalValue = DecimalUtility.getBigDecimalFromArrowBuf( buffer, index, scale, Decimal256Vector.TYPE_WIDTH); // We write the unscaled value, because the scale is stored in the type metadata. generator.writeString(decimalValue.unscaledValue().toString()); break; } default: throw new UnsupportedOperationException("minor type: " + vector.getMinorType()); } } else if (bufferType.equals(SIZE)) { if (vector.getMinorType() == MinorType.LISTVIEW) { generator.writeNumber(buffer.getInt((long) index * BaseRepeatedValueViewVector.SIZE_WIDTH)); } else { generator.writeNumber( buffer.getInt((long) index * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); } } } @Override public void close() throws IOException { generator.writeEndArray(); generator.writeEndObject(); generator.close(); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/ReadChannel.java000066400000000000000000000060721515451615200312670ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.channels.ReadableByteChannel; import org.apache.arrow.memory.ArrowBuf; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** Adapter around {@link ReadableByteChannel} that reads into {@linkplain ArrowBuf}s. */ public class ReadChannel implements AutoCloseable { private static final Logger LOGGER = LoggerFactory.getLogger(ReadChannel.class); private ReadableByteChannel in; private long bytesRead = 0; public ReadChannel(ReadableByteChannel in) { this.in = in; } public long bytesRead() { return bytesRead; } /** * Reads bytes into buffer until it is full (buffer.remaining() == 0). Returns the number of bytes * read which can be less than full if there are no more. * * @param buffer The buffer to read to * @return the number of byte read * @throws IOException if nit enough bytes left to read */ public int readFully(ByteBuffer buffer) throws IOException { if (LOGGER.isDebugEnabled()) { LOGGER.debug("Reading buffer with size: {}", buffer.remaining()); } int totalRead = 0; while (buffer.remaining() != 0) { int read = in.read(buffer); if (read == -1) { this.bytesRead += totalRead; return totalRead; } totalRead += read; if (read == 0) { break; } } this.bytesRead += totalRead; return totalRead; } /** * Reads up to len into buffer. Returns bytes read. * * @param buffer the buffer to read to * @param length the amount of bytes to read * @return the number of bytes read * @throws IOException if nit enough bytes left to read */ public long readFully(ArrowBuf buffer, long length) throws IOException { boolean fullRead = true; long bytesLeft = length; while (fullRead && bytesLeft > 0) { int bytesToRead = (int) Math.min(bytesLeft, Integer.MAX_VALUE); int n = readFully(buffer.nioBuffer(buffer.writerIndex(), bytesToRead)); buffer.writerIndex(buffer.writerIndex() + n); fullRead = n == bytesToRead; bytesLeft -= n; } return length - bytesLeft; } @Override public void close() throws IOException { if (this.in != null) { in.close(); in = null; } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/SeekableReadChannel.java000066400000000000000000000024761515451615200327270ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import java.io.IOException; import java.nio.channels.SeekableByteChannel; /** An {@link ReadChannel} that supports seeking to a random position. */ public class SeekableReadChannel extends ReadChannel { private final SeekableByteChannel in; public SeekableReadChannel(SeekableByteChannel in) { super(in); this.in = in; } public void setPosition(long position) throws IOException { in.position(position); } public long size() throws IOException { return in.size(); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/WriteChannel.java000066400000000000000000000114131515451615200315010ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import com.google.flatbuffers.FlatBufferBuilder; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.channels.WritableByteChannel; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.ipc.message.FBSerializable; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * Wrapper around a WritableByteChannel that maintains the position as well adding some common * serialization utilities. * *

    All write methods in this class follow full write semantics, i.e., write calls only return * after requested data has been fully written. Note this is different from java WritableByteChannel * interface where partial write is allowed * *

    Please note that objects of this class are not thread-safe. */ public class WriteChannel implements AutoCloseable { private static final Logger LOGGER = LoggerFactory.getLogger(WriteChannel.class); private static final byte[] ZERO_BYTES = new byte[8]; private final byte[] intBuf = new byte[4]; private long currentPosition = 0; private final WritableByteChannel out; public WriteChannel(WritableByteChannel out) { this.out = out; } @Override public void close() throws IOException { out.close(); } public long getCurrentPosition() { return currentPosition; } public long write(byte[] buffer) throws IOException { return write(ByteBuffer.wrap(buffer)); } long write(byte[] buffer, int offset, int length) throws IOException { return write(ByteBuffer.wrap(buffer, offset, length)); } /** Writes zeroCount zeros the underlying channel. */ public long writeZeros(long zeroCount) throws IOException { long bytesWritten = 0; long wholeWordsEnd = zeroCount - 8; while (bytesWritten <= wholeWordsEnd) { bytesWritten += write(ZERO_BYTES); } if (bytesWritten < zeroCount) { bytesWritten += write(ZERO_BYTES, 0, (int) (zeroCount - bytesWritten)); } return bytesWritten; } /** Writes enough bytes to align the channel to an 8-byte boundary. */ public long align() throws IOException { int trailingByteSize = (int) (currentPosition % 8); if (trailingByteSize != 0) { // align on 8 byte boundaries return writeZeros(8 - trailingByteSize); } return 0; } /** Writes all data from buffer to the underlying channel. */ public long write(ByteBuffer buffer) throws IOException { long length = buffer.remaining(); while (buffer.hasRemaining()) { out.write(buffer); } currentPosition += length; return length; } /** Writes v in little-endian format to the underlying channel. */ public long writeIntLittleEndian(int v) throws IOException { MessageSerializer.intToBytes(v, intBuf); return write(intBuf); } /** Writes the buffer to the underlying channel. */ public void write(ArrowBuf buffer) throws IOException { long bytesWritten = 0; while (bytesWritten < buffer.readableBytes()) { int bytesToWrite = (int) Math.min(Integer.MAX_VALUE, buffer.readableBytes() - bytesWritten); ByteBuffer nioBuffer = buffer.nioBuffer(buffer.readerIndex() + bytesWritten, bytesToWrite); write(nioBuffer); bytesWritten += bytesToWrite; } } /** * Writes the serialized flatbuffer to the underlying channel. If withSizePrefix is true then the * length in bytes of the buffer will first be written in little endian format. */ public long write(FBSerializable writer, boolean withSizePrefix) throws IOException { ByteBuffer buffer = serialize(writer); if (withSizePrefix) { writeIntLittleEndian(buffer.remaining()); } return write(buffer); } /** Serializes writer to a ByteBuffer. */ public static ByteBuffer serialize(FBSerializable writer) { FlatBufferBuilder builder = new FlatBufferBuilder(); int root = writer.writeTo(builder); builder.finish(root); return builder.dataBuffer(); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/message/000077500000000000000000000000001515451615200276775ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/message/ArrowBlock.java000066400000000000000000000052001515451615200326040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message; import com.google.flatbuffers.FlatBufferBuilder; import org.apache.arrow.flatbuf.Block; /** Metadata for an arrow message in a channel. */ public class ArrowBlock implements FBSerializable { private final long offset; private final int metadataLength; private final long bodyLength; /** * Constructs a new instance. * * @param offset The offset into the channel file where the block was written. * @param metadataLength The length of the flatbuffer metadata in the block. * @param bodyLength The length of data in the block. */ public ArrowBlock(long offset, int metadataLength, long bodyLength) { super(); this.offset = offset; this.metadataLength = metadataLength; this.bodyLength = bodyLength; } public long getOffset() { return offset; } public int getMetadataLength() { return metadataLength; } public long getBodyLength() { return bodyLength; } @Override public int writeTo(FlatBufferBuilder builder) { return Block.createBlock(builder, offset, metadataLength, bodyLength); } @Override public int hashCode() { final int prime = 31; int result = 1; result = prime * result + (int) (bodyLength ^ (bodyLength >>> 32)); result = prime * result + metadataLength; result = prime * result + (int) (offset ^ (offset >>> 32)); return result; } @Override public boolean equals(Object obj) { if (this == obj) { return true; } if (obj == null) { return false; } if (getClass() != obj.getClass()) { return false; } ArrowBlock other = (ArrowBlock) obj; if (bodyLength != other.bodyLength) { return false; } if (metadataLength != other.metadataLength) { return false; } if (offset != other.offset) { return false; } return true; } } ArrowBodyCompression.java000066400000000000000000000031161515451615200346160ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/message/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message; import com.google.flatbuffers.FlatBufferBuilder; import org.apache.arrow.flatbuf.BodyCompression; /** Compression information about data written to a channel. */ public class ArrowBodyCompression implements FBSerializable { private final byte codec; private final byte method; public ArrowBodyCompression(byte codec, byte method) { this.codec = codec; this.method = method; } @Override public int writeTo(FlatBufferBuilder builder) { return BodyCompression.createBodyCompression(builder, codec, method); } public byte getCodec() { return codec; } public byte getMethod() { return method; } @Override public String toString() { return "ArrowBodyCompression [codec=" + codec + ", method=" + method + "]"; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/message/ArrowBuffer.java000066400000000000000000000044271515451615200327750ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message; import com.google.flatbuffers.FlatBufferBuilder; import org.apache.arrow.flatbuf.Buffer; /** Metadata for a buffer written to a channel. */ public class ArrowBuffer implements FBSerializable { private long offset; private long size; /** * Constructs a new instance. * * @param offset The offset to the start of the buffer in the channel. * @param size The size of the buffer. */ public ArrowBuffer(long offset, long size) { super(); this.offset = offset; this.size = size; } public long getOffset() { return offset; } public long getSize() { return size; } @Override public int hashCode() { final int prime = 31; int result = 1; result = prime * result + (int) (offset ^ (offset >>> 32)); result = prime * result + (int) (size ^ (size >>> 32)); return result; } @Override public boolean equals(Object obj) { if (this == obj) { return true; } if (obj == null) { return false; } if (getClass() != obj.getClass()) { return false; } ArrowBuffer other = (ArrowBuffer) obj; if (offset != other.offset) { return false; } if (size != other.size) { return false; } return true; } @Override public int writeTo(FlatBufferBuilder builder) { return Buffer.createBuffer(builder, offset, size); } @Override public String toString() { return "ArrowBuffer [offset=" + offset + ", size=" + size + "]"; } } ArrowDictionaryBatch.java000066400000000000000000000053241515451615200345510ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/message/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message; import com.google.flatbuffers.FlatBufferBuilder; import org.apache.arrow.flatbuf.DictionaryBatch; import org.apache.arrow.flatbuf.MessageHeader; /** * POJO wrapper around a Dictionary Batch IPC messages. * (https://arrow.apache.org/docs/format/IPC.html#dictionary-batches) */ public class ArrowDictionaryBatch implements ArrowMessage { private final long dictionaryId; private final ArrowRecordBatch dictionary; private final boolean isDelta; @Deprecated public ArrowDictionaryBatch(long dictionaryId, ArrowRecordBatch dictionary) { this(dictionaryId, dictionary, false); } /** Constructs new instance. */ public ArrowDictionaryBatch(long dictionaryId, ArrowRecordBatch dictionary, boolean isDelta) { this.dictionaryId = dictionaryId; this.dictionary = dictionary; this.isDelta = isDelta; } public boolean isDelta() { return isDelta; } public byte getMessageType() { return MessageHeader.DictionaryBatch; } public long getDictionaryId() { return dictionaryId; } public ArrowRecordBatch getDictionary() { return dictionary; } @Override public int writeTo(FlatBufferBuilder builder) { int dataOffset = dictionary.writeTo(builder); DictionaryBatch.startDictionaryBatch(builder); DictionaryBatch.addId(builder, dictionaryId); DictionaryBatch.addData(builder, dataOffset); DictionaryBatch.addIsDelta(builder, isDelta); return DictionaryBatch.endDictionaryBatch(builder); } @Override public long computeBodyLength() { return dictionary.computeBodyLength(); } @Override public T accepts(ArrowMessageVisitor visitor) { return visitor.visit(this); } @Override public String toString() { return "ArrowDictionaryBatch [dictionaryId=" + dictionaryId + ", dictionary=" + dictionary + "]"; } @Override public void close() { dictionary.close(); } } ArrowFieldNode.java000066400000000000000000000035361515451615200333360ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/message/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message; import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import com.google.flatbuffers.FlatBufferBuilder; import org.apache.arrow.flatbuf.FieldNode; /** Metadata about Vectors/Arrays that is written to a channel. */ public class ArrowFieldNode implements FBSerializable { private final int length; private final int nullCount; /** * Constructs a new instance. * * @param length The number of values written. * @param nullCount The number of null values. */ public ArrowFieldNode(long length, long nullCount) { super(); this.length = checkedCastToInt(length); this.nullCount = checkedCastToInt(nullCount); } @Override public int writeTo(FlatBufferBuilder builder) { return FieldNode.createFieldNode(builder, length, nullCount); } public int getNullCount() { return nullCount; } public int getLength() { return length; } @Override public String toString() { return "ArrowFieldNode [length=" + length + ", nullCount=" + nullCount + "]"; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/message/ArrowFooter.java000066400000000000000000000156511515451615200330230ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message; import static org.apache.arrow.vector.ipc.message.FBSerializables.writeAllStructsToVector; import static org.apache.arrow.vector.ipc.message.FBSerializables.writeKeyValues; import com.google.flatbuffers.FlatBufferBuilder; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import org.apache.arrow.flatbuf.Block; import org.apache.arrow.flatbuf.Footer; import org.apache.arrow.flatbuf.KeyValue; import org.apache.arrow.vector.types.MetadataVersion; import org.apache.arrow.vector.types.pojo.Schema; /** Footer metadata for the arrow file format. */ public class ArrowFooter implements FBSerializable { private final Schema schema; private final List dictionaries; private final List recordBatches; private final Map metaData; private final MetadataVersion metadataVersion; public ArrowFooter(Schema schema, List dictionaries, List recordBatches) { this(schema, dictionaries, recordBatches, null); } /** * Constructs a new instance. * * @param schema The schema for record batches in the file. * @param dictionaries The dictionaries relevant to the file. * @param recordBatches The recordBatches written to the file. * @param metaData user-defined k-v meta data. */ public ArrowFooter( Schema schema, List dictionaries, List recordBatches, Map metaData) { this(schema, dictionaries, recordBatches, metaData, MetadataVersion.DEFAULT); } /** * Constructs a new instance. * * @param schema The schema for record batches in the file. * @param dictionaries The dictionaries relevant to the file. * @param recordBatches The recordBatches written to the file. * @param metaData user-defined k-v meta data. * @param metadataVersion The Arrow metadata version. */ public ArrowFooter( Schema schema, List dictionaries, List recordBatches, Map metaData, MetadataVersion metadataVersion) { this.schema = schema; this.dictionaries = dictionaries; this.recordBatches = recordBatches; this.metaData = metaData; this.metadataVersion = metadataVersion; } /** Constructs from the corresponding Flatbuffer message. */ public ArrowFooter(Footer footer) { this( Schema.convertSchema(footer.schema()), dictionaries(footer), recordBatches(footer), metaData(footer), MetadataVersion.fromFlatbufID(footer.version())); } private static List recordBatches(Footer footer) { List recordBatches = new ArrayList<>(); Block tempBlock = new Block(); int recordBatchesLength = footer.recordBatchesLength(); for (int i = 0; i < recordBatchesLength; i++) { Block block = footer.recordBatches(tempBlock, i); recordBatches.add(new ArrowBlock(block.offset(), block.metaDataLength(), block.bodyLength())); } return recordBatches; } private static List dictionaries(Footer footer) { List dictionaries = new ArrayList<>(); Block tempBlock = new Block(); int dictionariesLength = footer.dictionariesLength(); for (int i = 0; i < dictionariesLength; i++) { Block block = footer.dictionaries(tempBlock, i); dictionaries.add(new ArrowBlock(block.offset(), block.metaDataLength(), block.bodyLength())); } return dictionaries; } private static Map metaData(Footer footer) { Map metaData = new HashMap<>(); int metaDataLength = footer.customMetadataLength(); for (int i = 0; i < metaDataLength; i++) { KeyValue kv = footer.customMetadata(i); metaData.put(kv.key(), kv.value()); } return metaData; } public Schema getSchema() { return schema; } public List getDictionaries() { return dictionaries; } public List getRecordBatches() { return recordBatches; } public Map getMetaData() { return metaData; } public MetadataVersion getMetadataVersion() { return metadataVersion; } @Override public int writeTo(FlatBufferBuilder builder) { int schemaIndex = schema.getSchema(builder); Footer.startDictionariesVector(builder, dictionaries.size()); int dicsOffset = writeAllStructsToVector(builder, dictionaries); Footer.startRecordBatchesVector(builder, recordBatches.size()); int rbsOffset = writeAllStructsToVector(builder, recordBatches); int metaDataOffset = 0; if (metaData != null) { metaDataOffset = writeKeyValues(builder, metaData); } Footer.startFooter(builder); Footer.addSchema(builder, schemaIndex); Footer.addDictionaries(builder, dicsOffset); Footer.addRecordBatches(builder, rbsOffset); Footer.addCustomMetadata(builder, metaDataOffset); Footer.addVersion(builder, metadataVersion.toFlatbufID()); return Footer.endFooter(builder); } @Override public int hashCode() { final int prime = 31; int result = 1; result = prime * result + ((dictionaries == null) ? 0 : dictionaries.hashCode()); result = prime * result + ((recordBatches == null) ? 0 : recordBatches.hashCode()); result = prime * result + ((schema == null) ? 0 : schema.hashCode()); return result; } @Override public boolean equals(Object obj) { if (this == obj) { return true; } if (obj == null) { return false; } if (getClass() != obj.getClass()) { return false; } ArrowFooter other = (ArrowFooter) obj; if (dictionaries == null) { if (other.dictionaries != null) { return false; } } else if (!dictionaries.equals(other.dictionaries)) { return false; } if (recordBatches == null) { if (other.recordBatches != null) { return false; } } else if (!recordBatches.equals(other.recordBatches)) { return false; } if (schema == null) { if (other.schema != null) { return false; } } else if (!schema.equals(other.schema)) { return false; } return true; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/message/ArrowMessage.java000066400000000000000000000026661515451615200331530ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message; /** Interface for Arrow IPC messages (https://arrow.apache.org/docs/format/IPC.html). */ public interface ArrowMessage extends FBSerializable, AutoCloseable { long computeBodyLength(); T accepts(ArrowMessageVisitor visitor); /** Returns the flatbuffer enum value indicating the type of the message. */ byte getMessageType(); /** * Visitor interface for implementations of {@link ArrowMessage}. * * @param The type of value to return after visiting. */ interface ArrowMessageVisitor { T visit(ArrowDictionaryBatch message); T visit(ArrowRecordBatch message); } } ArrowRecordBatch.java000066400000000000000000000312421515451615200336600ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/message/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message; import com.google.flatbuffers.FlatBufferBuilder; import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.stream.Collectors; import org.apache.arrow.flatbuf.RecordBatch; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.compression.NoCompressionCodec; import org.apache.arrow.vector.util.DataSizeRoundingUtil; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * POJO representation of a RecordBatch IPC message * (https://arrow.apache.org/docs/format/Columnar.html#recordbatch-message). */ public class ArrowRecordBatch implements ArrowMessage { private static final Logger LOGGER = LoggerFactory.getLogger(ArrowRecordBatch.class); /** Number of records. */ private final int length; /** Nodes correspond to the pre-ordered flattened logical schema. */ private final List nodes; private final List buffers; private final ArrowBodyCompression bodyCompression; private final List buffersLayout; private final List variadicBufferCounts; private boolean closed = false; public ArrowRecordBatch(int length, List nodes, List buffers) { this(length, nodes, buffers, NoCompressionCodec.DEFAULT_BODY_COMPRESSION, null, true); } public ArrowRecordBatch( int length, List nodes, List buffers, ArrowBodyCompression bodyCompression) { this(length, nodes, buffers, bodyCompression, null, true); } /** * Construct a record batch from nodes. * * @param length how many rows in this batch * @param nodes field level info * @param buffers will be retained until this recordBatch is closed * @param bodyCompression compression info. * @param alignBuffers Whether to align buffers to an 8 byte boundary. */ public ArrowRecordBatch( int length, List nodes, List buffers, ArrowBodyCompression bodyCompression, boolean alignBuffers) { this(length, nodes, buffers, bodyCompression, null, alignBuffers, /*retainBuffers*/ true); } /** * Construct a record batch from nodes. * * @param length how many rows in this batch * @param nodes field level info * @param buffers will be retained until this recordBatch is closed * @param bodyCompression compression info. * @param alignBuffers Whether to align buffers to an 8 byte boundary. * @param retainBuffers Whether to retain() each source buffer in the constructor. If false, the * caller is responsible for retaining the buffers beforehand. */ public ArrowRecordBatch( int length, List nodes, List buffers, ArrowBodyCompression bodyCompression, boolean alignBuffers, boolean retainBuffers) { this(length, nodes, buffers, bodyCompression, null, alignBuffers, retainBuffers); } /** * Construct a record batch from nodes. * * @param length how many rows in this batch * @param nodes field level info * @param buffers will be retained until this recordBatch is closed * @param bodyCompression compression info. * @param variadicBufferCounts the number of buffers in each variadic section. * @param alignBuffers Whether to align buffers to an 8 byte boundary. */ public ArrowRecordBatch( int length, List nodes, List buffers, ArrowBodyCompression bodyCompression, List variadicBufferCounts, boolean alignBuffers) { this( length, nodes, buffers, bodyCompression, variadicBufferCounts, alignBuffers, /*retainBuffers*/ true); } /** * Construct a record batch from nodes. * * @param length how many rows in this batch * @param nodes field level info * @param buffers will be retained until this recordBatch is closed * @param bodyCompression compression info. * @param variadicBufferCounts the number of buffers in each variadic section. * @param alignBuffers Whether to align buffers to an 8 byte boundary. * @param retainBuffers Whether to retain() each source buffer in the constructor. If false, the * caller is responsible for retaining the buffers beforehand. */ public ArrowRecordBatch( int length, List nodes, List buffers, ArrowBodyCompression bodyCompression, List variadicBufferCounts, boolean alignBuffers, boolean retainBuffers) { super(); this.length = length; this.nodes = nodes; this.buffers = buffers; Preconditions.checkArgument(bodyCompression != null, "body compression cannot be null"); this.bodyCompression = bodyCompression; this.variadicBufferCounts = variadicBufferCounts; List arrowBuffers = new ArrayList<>(buffers.size()); long offset = 0; for (ArrowBuf arrowBuf : buffers) { if (retainBuffers) { arrowBuf.getReferenceManager().retain(); } long size = arrowBuf.readableBytes(); arrowBuffers.add(new ArrowBuffer(offset, size)); if (LOGGER.isTraceEnabled()) { LOGGER.trace("Buffer in RecordBatch at {}, length: {}", offset, size); } offset += size; if (alignBuffers) { // align on 8 byte boundaries offset = DataSizeRoundingUtil.roundUpTo8Multiple(offset); } } this.buffersLayout = Collections.unmodifiableList(arrowBuffers); } // clone constructor // this constructor is different from the public ones in that the reference manager's // retain method is not called, so the first dummy parameter is used // to distinguish this from the public constructor. private ArrowRecordBatch( boolean dummy, int length, List nodes, List buffers, ArrowBodyCompression bodyCompression, List variadicBufferCounts) { this.length = length; this.nodes = nodes; this.buffers = buffers; Preconditions.checkArgument(bodyCompression != null, "body compression cannot be null"); this.bodyCompression = bodyCompression; this.variadicBufferCounts = variadicBufferCounts; this.closed = false; List arrowBuffers = new ArrayList<>(); long offset = 0; for (ArrowBuf arrowBuf : buffers) { long size = arrowBuf.readableBytes(); arrowBuffers.add(new ArrowBuffer(offset, size)); offset += size; } this.buffersLayout = Collections.unmodifiableList(arrowBuffers); } public byte getMessageType() { return org.apache.arrow.flatbuf.MessageHeader.RecordBatch; } public int getLength() { return length; } public ArrowBodyCompression getBodyCompression() { return bodyCompression; } /** * Get the nodes in this record batch. * * @return the FieldNodes corresponding to the schema */ public List getNodes() { return nodes; } /** * Get the record batch buffers. * * @return the buffers containing the data */ public List getBuffers() { if (closed) { throw new IllegalStateException("already closed"); } return buffers; } /** * Get the record batch variadic buffer counts. * * @return the variadic buffer counts */ public List getVariadicBufferCounts() { return variadicBufferCounts; } /** * Create a new ArrowRecordBatch which has the same information as this batch but whose buffers * are owned by that Allocator. * *

    This will also close this record batch and make it no longer useful. * * @return A cloned ArrowRecordBatch */ public ArrowRecordBatch cloneWithTransfer(final BufferAllocator allocator) { final List newBufs = buffers.stream() .map( buf -> (buf.getReferenceManager() .transferOwnership(buf, allocator) .getTransferredBuffer()) .writerIndex(buf.writerIndex())) .collect(Collectors.toList()); close(); return new ArrowRecordBatch( false, length, nodes, newBufs, bodyCompression, variadicBufferCounts); } /** * Get the serialized layout. * * @return the serialized layout if we send the buffers on the wire */ public List getBuffersLayout() { return buffersLayout; } @Override public int writeTo(FlatBufferBuilder builder) { RecordBatch.startNodesVector(builder, nodes.size()); int nodesOffset = FBSerializables.writeAllStructsToVector(builder, nodes); RecordBatch.startBuffersVector(builder, buffers.size()); int buffersOffset = FBSerializables.writeAllStructsToVector(builder, buffersLayout); int compressOffset = 0; if (bodyCompression.getCodec() != NoCompressionCodec.COMPRESSION_TYPE) { compressOffset = bodyCompression.writeTo(builder); } // Start the variadicBufferCounts vector. int variadicBufferCountsOffset = 0; if (variadicBufferCounts != null && !variadicBufferCounts.isEmpty()) { variadicBufferCountsOffset = variadicBufferCounts.size(); int elementSizeInBytes = 8; // Size of long in bytes builder.startVector(elementSizeInBytes, variadicBufferCountsOffset, elementSizeInBytes); // Add each long to the builder. Note that elements should be added in reverse order. for (int i = variadicBufferCounts.size() - 1; i >= 0; i--) { long value = variadicBufferCounts.get(i); builder.addLong(value); } // End the vector. This returns an offset that you can use to refer to the vector. variadicBufferCountsOffset = builder.endVector(); } RecordBatch.startRecordBatch(builder); RecordBatch.addLength(builder, length); RecordBatch.addNodes(builder, nodesOffset); RecordBatch.addBuffers(builder, buffersOffset); if (bodyCompression.getCodec() != NoCompressionCodec.COMPRESSION_TYPE) { RecordBatch.addCompression(builder, compressOffset); } // Add the variadicBufferCounts to the RecordBatch if (variadicBufferCounts != null && !variadicBufferCounts.isEmpty()) { RecordBatch.addVariadicBufferCounts(builder, variadicBufferCountsOffset); } return RecordBatch.endRecordBatch(builder); } @Override public T accepts(ArrowMessageVisitor visitor) { return visitor.visit(this); } /** Releases the buffers. */ @Override public void close() { if (!closed) { closed = true; for (ArrowBuf arrowBuf : buffers) { arrowBuf.getReferenceManager().release(); } } } @Override public String toString() { int variadicBufCount = 0; if (variadicBufferCounts != null && !variadicBufferCounts.isEmpty()) { variadicBufCount = variadicBufferCounts.size(); } return "ArrowRecordBatch [length=" + length + ", nodes=" + nodes + ", #buffers=" + buffers.size() + ", #variadicBufferCounts=" + variadicBufCount + ", buffersLayout=" + buffersLayout + ", closed=" + closed + "]"; } /** Computes the size of the serialized body for this recordBatch. */ @Override public long computeBodyLength() { long size = 0; List buffers = getBuffers(); List buffersLayout = getBuffersLayout(); if (buffers.size() != buffersLayout.size()) { throw new IllegalStateException( "the layout does not match: " + buffers.size() + " != " + buffersLayout.size()); } for (int i = 0; i < buffers.size(); i++) { ArrowBuf buffer = buffers.get(i); ArrowBuffer layout = buffersLayout.get(i); size = layout.getOffset() + buffer.readableBytes(); // round up size to the next multiple of 8 size = DataSizeRoundingUtil.roundUpTo8Multiple(size); } return size; } } FBSerializable.java000066400000000000000000000021411515451615200332770ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/message/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message; import com.google.flatbuffers.FlatBufferBuilder; /** Interface for serializing to FlatBuffers. */ public interface FBSerializable { /** Returns the number of bytes taken to serialize the data in builder after writing to it. */ int writeTo(FlatBufferBuilder builder); } FBSerializables.java000066400000000000000000000050571515451615200334730ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/message/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message; import com.google.flatbuffers.FlatBufferBuilder; import java.util.ArrayList; import java.util.Collections; import java.util.Iterator; import java.util.List; import java.util.Map; import org.apache.arrow.flatbuf.KeyValue; /** Utility methods for {@linkplain org.apache.arrow.vector.ipc.message.FBSerializable}s. */ public class FBSerializables { private FBSerializables() {} /** * Writes every element of all to builder and calls {@link FlatBufferBuilder#endVector()} * afterwards. Returns the number of result of calling endVector. */ public static int writeAllStructsToVector( FlatBufferBuilder builder, List all) { // struct vectors have to be created in reverse order List reversed = new ArrayList<>(all); Collections.reverse(reversed); for (FBSerializable element : reversed) { element.writeTo(builder); } return builder.endVector(); } /** Writes map data with string type. */ public static int writeKeyValues(FlatBufferBuilder builder, Map metaData) { int[] metadataOffsets = new int[metaData.size()]; Iterator> metadataIterator = metaData.entrySet().iterator(); for (int i = 0; i < metadataOffsets.length; i++) { Map.Entry kv = metadataIterator.next(); int keyOffset = builder.createString(kv.getKey()); int valueOffset = builder.createString(kv.getValue()); KeyValue.startKeyValue(builder); KeyValue.addKey(builder, keyOffset); KeyValue.addValue(builder, valueOffset); metadataOffsets[i] = KeyValue.endKeyValue(builder); } return org.apache.arrow.flatbuf.Field.createCustomMetadataVector(builder, metadataOffsets); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/message/IpcOption.java000066400000000000000000000030111515451615200324410ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message; import org.apache.arrow.vector.types.MetadataVersion; /** IPC options, now only use for write. */ public class IpcOption { // Write the pre-0.15.0 encapsulated IPC message format // consisting of a 4-byte prefix instead of 8 byte public final boolean write_legacy_ipc_format; // The metadata version. Defaults to V5. public final MetadataVersion metadataVersion; public IpcOption() { this(false, MetadataVersion.DEFAULT); } public IpcOption(boolean writeLegacyIpcFormat, MetadataVersion metadataVersion) { this.write_legacy_ipc_format = writeLegacyIpcFormat; this.metadataVersion = metadataVersion; } public static final IpcOption DEFAULT = new IpcOption(); } MessageChannelReader.java000066400000000000000000000055511515451615200344710ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/message/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message; import java.io.IOException; import org.apache.arrow.flatbuf.Message; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.ipc.ReadChannel; /** Reads a sequence of messages using a ReadChannel. */ public class MessageChannelReader implements AutoCloseable { protected ReadChannel in; protected BufferAllocator allocator; /** * Construct a MessageReader to read streaming messages from an existing ReadChannel. * * @param in Channel to read messages from * @param allocator BufferAllocator used to read Message body into an ArrowBuf. */ public MessageChannelReader(ReadChannel in, BufferAllocator allocator) { this.in = in; this.allocator = allocator; } /** * Read a message from the ReadChannel and return a MessageResult containing the Message metadata * and optional message body data. Once the end-of-stream has been reached, a null value will be * returned. If the message has no body, then MessageResult.getBodyBuffer() returns null. * * @return MessageResult or null if reached end-of-stream * @throws IOException on error */ public MessageResult readNext() throws IOException { // Read the flatbuf message and check for end-of-stream MessageMetadataResult result = MessageSerializer.readMessage(in); if (result == null) { return null; } Message message = result.getMessage(); ArrowBuf bodyBuffer = null; // Read message body data if defined in message if (result.messageHasBody()) { long bodyLength = result.getMessageBodyLength(); bodyBuffer = MessageSerializer.readMessageBody(in, bodyLength, allocator); } return new MessageResult(message, bodyBuffer); } /** * Get the number of bytes read from the ReadChannel. * * @return number of bytes */ public long bytesRead() { return in.bytesRead(); } /** * Close the ReadChannel. * * @throws IOException on error */ @Override public void close() throws IOException { in.close(); } } MessageMetadataResult.java000066400000000000000000000072721515451615200347170ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/message/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message; import java.nio.ByteBuffer; import org.apache.arrow.flatbuf.Message; /** * Class to hold resulting Message metadata and buffer containing the serialized Flatbuffer message * when reading messages from a ReadChannel. This handles Message metadata only and does not include * the message body data, which should be subsequently read into an ArrowBuf. */ public class MessageMetadataResult { /** * Construct a container to hold a deserialized Message metadata, and buffer with the serialized * Message as read from a ReadChannel. * * @param messageLength the length of the serialized Flatbuffer message in bytes * @param messageBuffer contains the serialized Flatbuffer Message metadata * @param message the deserialized Flatbuffer Message metadata description */ MessageMetadataResult(int messageLength, ByteBuffer messageBuffer, Message message) { this.messageLength = messageLength; this.messageBuffer = messageBuffer; this.message = message; } /** * Creates a new {@link MessageMetadataResult} by parsing it from the beginning of the buffer. * * @param messageLength The length of the serialized flatbuffer message in bytes (might not be * equal to the buffer size). */ public static MessageMetadataResult create(ByteBuffer buffer, int messageLength) { return new MessageMetadataResult(messageLength, buffer, Message.getRootAsMessage(buffer)); } /** * Get the length of the message metadata in bytes, not including the body length. * * @return number of bytes in the message metadata buffer. */ public int getMessageLength() { return messageLength; } /** * Get the buffer containing the raw message metadata bytes, not including the message body data. * * @return buffer containing the message metadata. */ public ByteBuffer getMessageBuffer() { return messageBuffer; } /** Returns the bytes remaining in the buffer after parsing the message from it. */ public int bytesAfterMessage() { return message.getByteBuffer().remaining(); } public byte headerType() { return message.headerType(); } /** * Check if the message is followed by a body. This will be true if the message has a body length * > 0, which indicates that a message body needs to be read from the input source. * * @return true if message has a defined body */ public boolean messageHasBody() { return message.bodyLength() > 0; } /** * Get the length of the message body. * * @return number of bytes of the message body */ public long getMessageBodyLength() { return message.bodyLength(); } /** * Get the realized flatbuf Message metadata description. * * @return Message metadata */ public Message getMessage() { return message; } private final int messageLength; private final ByteBuffer messageBuffer; private final Message message; } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/message/MessageResult.java000066400000000000000000000036071515451615200333330ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message; import org.apache.arrow.flatbuf.Message; import org.apache.arrow.memory.ArrowBuf; /** * Class to hold the Message metadata and body data when reading messages through a * MessageChannelReader. */ public class MessageResult { /** * Construct with a valid Message metadata and optional ArrowBuf containing message body data, if * any. * * @param message Deserialized Flatbuffer Message metadata description * @param bodyBuffer Optional ArrowBuf containing message body data, null if message has no body */ MessageResult(Message message, ArrowBuf bodyBuffer) { this.message = message; this.bodyBuffer = bodyBuffer; } /** * Get the Message metadata. * * @return the Flatbuffer Message metadata */ public Message getMessage() { return message; } /** * Get the message body data. * * @return an ArrowBuf containing the message body data or null if the message has no body */ public ArrowBuf getBodyBuffer() { return bodyBuffer; } private final Message message; private final ArrowBuf bodyBuffer; } MessageSerializer.java000066400000000000000000000731361515451615200341130ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/ipc/message/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message; import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import com.google.flatbuffers.FlatBufferBuilder; import java.io.IOException; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.List; import org.apache.arrow.flatbuf.Buffer; import org.apache.arrow.flatbuf.DictionaryBatch; import org.apache.arrow.flatbuf.FieldNode; import org.apache.arrow.flatbuf.Message; import org.apache.arrow.flatbuf.MessageHeader; import org.apache.arrow.flatbuf.MetadataVersion; import org.apache.arrow.flatbuf.RecordBatch; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.compression.NoCompressionCodec; import org.apache.arrow.vector.ipc.ReadChannel; import org.apache.arrow.vector.ipc.WriteChannel; import org.apache.arrow.vector.types.pojo.Schema; /** * Utility class for serializing Messages. Messages are all serialized a similar way. 1. 4 byte * little endian message header prefix 2. FB serialized Message: This includes it the body length, * which is the serialized body and the type of the message. 3. Serialized message. * *

    For schema messages, the serialization is simply the FB serialized Schema. * *

    For RecordBatch messages the serialization is: 1. 4 byte little endian batch metadata header * 2. FB serialized RowBatch 3. Padding to align to 8 byte boundary. 4. serialized RowBatch buffers. */ public class MessageSerializer { // This 0xFFFFFFFF value is the first 4 bytes of a valid IPC message public static final int IPC_CONTINUATION_TOKEN = -1; /** * Convert an array of 4 bytes in little-endian to an native-endian i32 value. * * @param bytes byte array with minimum length of 4 in little-endian * @return converted an native-endian 32-bit integer */ public static int bytesToInt(byte[] bytes) { return ((bytes[3] & 255) << 24) + ((bytes[2] & 255) << 16) + ((bytes[1] & 255) << 8) + ((bytes[0] & 255)); } /** * Convert an integer to a little endian 4 byte array. * * @param value integer value input * @param bytes existing byte array with minimum length of 4 to contain the conversion output */ public static void intToBytes(int value, byte[] bytes) { bytes[3] = (byte) (value >>> 24); bytes[2] = (byte) (value >>> 16); bytes[1] = (byte) (value >>> 8); bytes[0] = (byte) (value); } /** * Convert a long to a little-endian 8 byte array. * * @param value long value input * @param bytes existing byte array with minimum length of 8 to contain the conversion output */ public static void longToBytes(long value, byte[] bytes) { bytes[7] = (byte) (value >>> 56); bytes[6] = (byte) (value >>> 48); bytes[5] = (byte) (value >>> 40); bytes[4] = (byte) (value >>> 32); bytes[3] = (byte) (value >>> 24); bytes[2] = (byte) (value >>> 16); bytes[1] = (byte) (value >>> 8); bytes[0] = (byte) (value); } public static int writeMessageBuffer( WriteChannel out, int messageLength, ByteBuffer messageBuffer) throws IOException { return writeMessageBuffer(out, messageLength, messageBuffer, IpcOption.DEFAULT); } /** * Write the serialized Message metadata, prefixed by the length, to the output Channel. This * ensures that it aligns to an 8 byte boundary and will adjust the message length to include any * padding used for alignment. * * @param out Output Channel * @param messageLength Number of bytes in the message buffer, written as little Endian prefix * @param messageBuffer Message metadata buffer to be written, this does not include any message * body data which should be subsequently written to the Channel * @param option IPC write options * @return Number of bytes written * @throws IOException on error */ public static int writeMessageBuffer( WriteChannel out, int messageLength, ByteBuffer messageBuffer, IpcOption option) throws IOException { // if write the pre-0.15.0 encapsulated IPC message format consisting of a 4-byte prefix instead // of 8 byte int prefixSize = option.write_legacy_ipc_format ? 4 : 8; // ensure that message aligns to 8 byte padding - prefix_size bytes, then message body if ((messageLength + prefixSize) % 8 != 0) { messageLength += 8 - (messageLength + prefixSize) % 8; } if (!option.write_legacy_ipc_format) { out.writeIntLittleEndian(IPC_CONTINUATION_TOKEN); } out.writeIntLittleEndian(messageLength); out.write(messageBuffer); out.align(); // any bytes written are already captured by our size modification above return messageLength + prefixSize; } /** Serialize a schema object. */ public static long serialize(WriteChannel out, Schema schema) throws IOException { return serialize(out, schema, IpcOption.DEFAULT); } /** * Serialize a schema object. * * @param out where to write the schema * @param schema the object to serialize to out * @return the number of bytes written * @throws IOException if something went wrong */ public static long serialize(WriteChannel out, Schema schema, IpcOption option) throws IOException { long start = out.getCurrentPosition(); Preconditions.checkArgument(start % 8 == 0, "out is not aligned"); ByteBuffer serializedMessage = serializeMetadata(schema, option); int messageLength = serializedMessage.remaining(); int bytesWritten = writeMessageBuffer(out, messageLength, serializedMessage, option); Preconditions.checkArgument(bytesWritten % 8 == 0, "out is not aligned"); return bytesWritten; } /** Returns the serialized flatbuffer bytes of the schema wrapped in a message table. */ @Deprecated public static ByteBuffer serializeMetadata(Schema schema) { return serializeMetadata(schema, IpcOption.DEFAULT); } /** Returns the serialized flatbuffer bytes of the schema wrapped in a message table. */ public static ByteBuffer serializeMetadata(Schema schema, IpcOption writeOption) { FlatBufferBuilder builder = new FlatBufferBuilder(); int schemaOffset = schema.getSchema(builder); return MessageSerializer.serializeMessage( builder, org.apache.arrow.flatbuf.MessageHeader.Schema, schemaOffset, 0, writeOption); } /** * Deserializes an Arrow Schema object from a schema message. Format is from serialize(). * * @param schemaMessage a Message of type MessageHeader.Schema * @return the deserialized Arrow Schema */ public static Schema deserializeSchema(Message schemaMessage) { Preconditions.checkArgument( schemaMessage.headerType() == MessageHeader.Schema, "Expected schema but result was: %s", schemaMessage.headerType()); return Schema.convertSchema( (org.apache.arrow.flatbuf.Schema) schemaMessage.header(new org.apache.arrow.flatbuf.Schema())); } /** * Deserializes an Arrow Schema read from the input channel. Format is from serialize(). * * @param in the channel to deserialize from * @return the deserialized Arrow Schema * @throws IOException if something went wrong */ public static Schema deserializeSchema(ReadChannel in) throws IOException { MessageMetadataResult result = readMessage(in); if (result == null) { throw new IOException("Unexpected end of input when reading Schema"); } if (result.getMessage().headerType() != MessageHeader.Schema) { throw new IOException("Expected schema but header was " + result.getMessage().headerType()); } return deserializeSchema(result); } /** * Deserializes an Arrow Schema object from a {@link MessageMetadataResult}. Format is from * serialize(). * * @param message a Message of type MessageHeader.Schema * @return the deserialized Arrow Schema */ public static Schema deserializeSchema(MessageMetadataResult message) { return deserializeSchema(message.getMessage()); } /** Serializes an ArrowRecordBatch. Returns the offset and length of the written batch. */ public static ArrowBlock serialize(WriteChannel out, ArrowRecordBatch batch) throws IOException { return serialize(out, batch, IpcOption.DEFAULT); } /** * Serializes an ArrowRecordBatch. Returns the offset and length of the written batch. * * @param out where to write the batch * @param batch the object to serialize to out * @return the serialized block metadata * @throws IOException if something went wrong */ public static ArrowBlock serialize(WriteChannel out, ArrowRecordBatch batch, IpcOption option) throws IOException { long start = out.getCurrentPosition(); long bodyLength = batch.computeBodyLength(); Preconditions.checkArgument(bodyLength % 8 == 0, "batch is not aligned"); ByteBuffer serializedMessage = serializeMetadata(batch, option); int metadataLength = serializedMessage.remaining(); int prefixSize = 4; if (!option.write_legacy_ipc_format) { out.writeIntLittleEndian(IPC_CONTINUATION_TOKEN); prefixSize = 8; } // calculate alignment bytes so that metadata length points to the correct location after // alignment int padding = (int) ((start + metadataLength + prefixSize) % 8); if (padding != 0) { metadataLength += (8 - padding); } out.writeIntLittleEndian(metadataLength); out.write(serializedMessage); // Align the output to 8 byte boundary. out.align(); long bufferLength = writeBatchBuffers(out, batch); Preconditions.checkArgument(bufferLength % 8 == 0, "out is not aligned"); // Metadata size in the Block account for the size prefix return new ArrowBlock(start, metadataLength + prefixSize, bufferLength); } /** * Write the Arrow buffers of the record batch to the output channel. * * @param out the output channel to write the buffers to * @param batch an ArrowRecordBatch containing buffers to be written * @return the number of bytes written * @throws IOException on error */ public static long writeBatchBuffers(WriteChannel out, ArrowRecordBatch batch) throws IOException { long bufferStart = out.getCurrentPosition(); List buffers = batch.getBuffers(); List buffersLayout = batch.getBuffersLayout(); for (int i = 0; i < buffers.size(); i++) { ArrowBuf buffer = buffers.get(i); ArrowBuffer layout = buffersLayout.get(i); long startPosition = bufferStart + layout.getOffset(); if (startPosition != out.getCurrentPosition()) { out.writeZeros(startPosition - out.getCurrentPosition()); } out.write(buffer); if (out.getCurrentPosition() != startPosition + layout.getSize()) { throw new IllegalStateException( "wrong buffer size: " + out.getCurrentPosition() + " != " + startPosition + layout.getSize()); } } out.align(); return out.getCurrentPosition() - bufferStart; } /** * Returns the serialized form of {@link RecordBatch} wrapped in a {@link * org.apache.arrow.flatbuf.Message}. */ @Deprecated public static ByteBuffer serializeMetadata(ArrowMessage message) { return serializeMetadata(message, IpcOption.DEFAULT); } /** * Returns the serialized form of {@link RecordBatch} wrapped in a {@link * org.apache.arrow.flatbuf.Message}. */ public static ByteBuffer serializeMetadata(ArrowMessage message, IpcOption writeOption) { FlatBufferBuilder builder = new FlatBufferBuilder(); int batchOffset = message.writeTo(builder); return serializeMessage( builder, message.getMessageType(), batchOffset, message.computeBodyLength(), writeOption); } /** * Deserializes an ArrowRecordBatch from a record batch message and data in an ArrowBuf. * * @param recordBatchMessage a Message of type MessageHeader.RecordBatch * @param bodyBuffer Arrow buffer containing the RecordBatch data * @return the deserialized ArrowRecordBatch * @throws IOException if something went wrong */ public static ArrowRecordBatch deserializeRecordBatch( Message recordBatchMessage, ArrowBuf bodyBuffer) throws IOException { RecordBatch recordBatchFB = (RecordBatch) recordBatchMessage.header(new RecordBatch()); return deserializeRecordBatch(recordBatchFB, bodyBuffer); } /** * Deserializes an ArrowRecordBatch read from the input channel. This uses the given allocator to * create an ArrowBuf for the batch body data. * * @param in Channel to read a RecordBatch message and data from * @param allocator BufferAllocator to allocate an Arrow buffer to read message body data * @return the deserialized ArrowRecordBatch * @throws IOException on error */ public static ArrowRecordBatch deserializeRecordBatch(ReadChannel in, BufferAllocator allocator) throws IOException { MessageMetadataResult result = readMessage(in); if (result == null) { throw new IOException("Unexpected end of input when reading a RecordBatch"); } if (result.getMessage().headerType() != MessageHeader.RecordBatch) { throw new IOException( "Expected RecordBatch but header was " + result.getMessage().headerType()); } long bodyLength = result.getMessageBodyLength(); ArrowBuf bodyBuffer = readMessageBody(in, bodyLength, allocator); return deserializeRecordBatch(result.getMessage(), bodyBuffer); } /** * Deserializes an ArrowRecordBatch knowing the size of the entire message up front. This * minimizes the number of reads to the underlying stream. * * @param in the channel to deserialize from * @param block the object to deserialize to * @param alloc to allocate buffers * @return the deserialized ArrowRecordBatch * @throws IOException if something went wrong */ public static ArrowRecordBatch deserializeRecordBatch( ReadChannel in, ArrowBlock block, BufferAllocator alloc) throws IOException { // Metadata length contains prefix_size bytes plus byte padding long totalLen = block.getMetadataLength() + block.getBodyLength(); ArrowBuf buffer = alloc.buffer(totalLen); if (in.readFully(buffer, totalLen) != totalLen) { throw new IOException("Unexpected end of input trying to read batch."); } int prefixSize = buffer.getInt(0) == IPC_CONTINUATION_TOKEN ? 8 : 4; ArrowBuf metadataBuffer = buffer.slice(prefixSize, block.getMetadataLength() - prefixSize); Message messageFB = Message.getRootAsMessage(metadataBuffer.nioBuffer().asReadOnlyBuffer()); RecordBatch recordBatchFB = (RecordBatch) messageFB.header(new RecordBatch()); // Now read the body final ArrowBuf body = buffer.slice(block.getMetadataLength(), totalLen - block.getMetadataLength()); return deserializeRecordBatch(recordBatchFB, body); } /** * Deserializes an ArrowRecordBatch given the Flatbuffer metadata and in-memory body. * * @param recordBatchFB Deserialized FlatBuffer record batch * @param body Read body of the record batch * @return ArrowRecordBatch from metadata and in-memory body * @throws IOException on error */ public static ArrowRecordBatch deserializeRecordBatch(RecordBatch recordBatchFB, ArrowBuf body) throws IOException { // Now read the body int nodesLength = recordBatchFB.nodesLength(); List nodes = new ArrayList<>(); for (int i = 0; i < nodesLength; ++i) { FieldNode node = recordBatchFB.nodes(i); if ((int) node.length() != node.length() || (int) node.nullCount() != node.nullCount()) { throw new IOException( "Cannot currently deserialize record batches with " + "node length larger than INT_MAX records."); } nodes.add(new ArrowFieldNode(node.length(), node.nullCount())); } List buffers = new ArrayList<>(); for (int i = 0; i < recordBatchFB.buffersLength(); ++i) { Buffer bufferFB = recordBatchFB.buffers(i); ArrowBuf vectorBuffer = body.slice(bufferFB.offset(), bufferFB.length()); buffers.add(vectorBuffer); } ArrowBodyCompression bodyCompression = recordBatchFB.compression() == null ? NoCompressionCodec.DEFAULT_BODY_COMPRESSION : new ArrowBodyCompression( recordBatchFB.compression().codec(), recordBatchFB.compression().method()); List variadicBufferCounts = new ArrayList<>(); for (int i = 0; i < recordBatchFB.variadicBufferCountsLength(); i++) { variadicBufferCounts.add(recordBatchFB.variadicBufferCounts(i)); } if ((int) recordBatchFB.length() != recordBatchFB.length()) { throw new IOException( "Cannot currently deserialize record batches with more than INT_MAX records."); } ArrowRecordBatch arrowRecordBatch = new ArrowRecordBatch( checkedCastToInt(recordBatchFB.length()), nodes, buffers, bodyCompression, variadicBufferCounts, /*alignBuffers*/ true); body.getReferenceManager().release(); return arrowRecordBatch; } /** * Reads a record batch based on the metadata in serializedMessage and the underlying data buffer. */ public static ArrowRecordBatch deserializeRecordBatch( MessageMetadataResult serializedMessage, ArrowBuf underlying) throws IOException { return deserializeRecordBatch(serializedMessage.getMessage(), underlying); } public static ArrowBlock serialize(WriteChannel out, ArrowDictionaryBatch batch) throws IOException { return serialize(out, batch, IpcOption.DEFAULT); } /** * Serializes a dictionary ArrowRecordBatch. Returns the offset and length of the written batch. * * @param out where to serialize * @param batch the batch to serialize * @param option options for IPC * @return the metadata of the serialized block * @throws IOException if something went wrong */ public static ArrowBlock serialize(WriteChannel out, ArrowDictionaryBatch batch, IpcOption option) throws IOException { long start = out.getCurrentPosition(); long bodyLength = batch.computeBodyLength(); Preconditions.checkArgument(bodyLength % 8 == 0, "batch is not aligned"); ByteBuffer serializedMessage = serializeMetadata(batch, option); int metadataLength = serializedMessage.remaining(); int prefixSize = 4; if (!option.write_legacy_ipc_format) { out.writeIntLittleEndian(IPC_CONTINUATION_TOKEN); prefixSize = 8; } // calculate alignment bytes so that metadata length points to the correct location after // alignment int padding = (int) ((start + metadataLength + prefixSize) % 8); if (padding != 0) { metadataLength += (8 - padding); } out.writeIntLittleEndian(metadataLength); out.write(serializedMessage); // Align the output to 8 byte boundary. out.align(); // write the embedded record batch long bufferLength = writeBatchBuffers(out, batch.getDictionary()); Preconditions.checkArgument(bufferLength % 8 == 0, "out is not aligned"); // Metadata size in the Block account for the size prefix return new ArrowBlock(start, metadataLength + prefixSize, bufferLength); } /** * Deserializes an ArrowDictionaryBatch from a dictionary batch Message and data in an ArrowBuf. * * @param message a message of type MessageHeader.DictionaryBatch * @param bodyBuffer Arrow buffer containing the DictionaryBatch data of type * MessageHeader.DictionaryBatch * @return the deserialized ArrowDictionaryBatch * @throws IOException if something went wrong */ public static ArrowDictionaryBatch deserializeDictionaryBatch( Message message, ArrowBuf bodyBuffer) throws IOException { DictionaryBatch dictionaryBatchFB = (DictionaryBatch) message.header(new DictionaryBatch()); ArrowRecordBatch recordBatch = deserializeRecordBatch(dictionaryBatchFB.data(), bodyBuffer); return new ArrowDictionaryBatch( dictionaryBatchFB.id(), recordBatch, dictionaryBatchFB.isDelta()); } /** * Deserializes an ArrowDictionaryBatch from a dictionary batch Message and data in an ArrowBuf. * * @param message a message of type MessageHeader.DictionaryBatch * @param bodyBuffer Arrow buffer containing the DictionaryBatch data of type * MessageHeader.DictionaryBatch * @return the deserialized ArrowDictionaryBatch * @throws IOException if something went wrong */ public static ArrowDictionaryBatch deserializeDictionaryBatch( MessageMetadataResult message, ArrowBuf bodyBuffer) throws IOException { return deserializeDictionaryBatch(message.getMessage(), bodyBuffer); } /** * Deserializes an ArrowDictionaryBatch read from the input channel. This uses the given allocator * to create an ArrowBuf for the batch body data. * * @param in Channel to read a DictionaryBatch message and data from * @param allocator BufferAllocator to allocate an Arrow buffer to read message body data * @return the deserialized ArrowDictionaryBatch * @throws IOException on error */ public static ArrowDictionaryBatch deserializeDictionaryBatch( ReadChannel in, BufferAllocator allocator) throws IOException { MessageMetadataResult result = readMessage(in); if (result == null) { throw new IOException("Unexpected end of input when reading a DictionaryBatch"); } if (result.getMessage().headerType() != MessageHeader.DictionaryBatch) { throw new IOException( "Expected DictionaryBatch but header was " + result.getMessage().headerType()); } long bodyLength = result.getMessageBodyLength(); ArrowBuf bodyBuffer = readMessageBody(in, bodyLength, allocator); return deserializeDictionaryBatch(result.getMessage(), bodyBuffer); } /** * Deserializes a DictionaryBatch knowing the size of the entire message up front. This minimizes * the number of reads to the underlying stream. * * @param in where to read from * @param block block metadata for deserializing * @param alloc to allocate new buffers * @return the deserialized ArrowDictionaryBatch * @throws IOException if something went wrong */ public static ArrowDictionaryBatch deserializeDictionaryBatch( ReadChannel in, ArrowBlock block, BufferAllocator alloc) throws IOException { // Metadata length contains integer prefix plus byte padding long totalLen = block.getMetadataLength() + block.getBodyLength(); ArrowBuf buffer = alloc.buffer(totalLen); if (in.readFully(buffer, totalLen) != totalLen) { throw new IOException("Unexpected end of input trying to read batch."); } int prefixSize = buffer.getInt(0) == IPC_CONTINUATION_TOKEN ? 8 : 4; ArrowBuf metadataBuffer = buffer.slice(prefixSize, block.getMetadataLength() - prefixSize); Message messageFB = Message.getRootAsMessage(metadataBuffer.nioBuffer().asReadOnlyBuffer()); DictionaryBatch dictionaryBatchFB = (DictionaryBatch) messageFB.header(new DictionaryBatch()); // Now read the body final ArrowBuf body = buffer.slice(block.getMetadataLength(), totalLen - block.getMetadataLength()); ArrowRecordBatch recordBatch = deserializeRecordBatch(dictionaryBatchFB.data(), body); return new ArrowDictionaryBatch( dictionaryBatchFB.id(), recordBatch, dictionaryBatchFB.isDelta()); } /** * Deserialize a message that is either an ArrowDictionaryBatch or ArrowRecordBatch. * * @param reader MessageChannelReader to read a sequence of messages from a ReadChannel * @return The deserialized record batch * @throws IOException if the message is not an ArrowDictionaryBatch or ArrowRecordBatch */ public static ArrowMessage deserializeMessageBatch(MessageChannelReader reader) throws IOException { MessageResult result = reader.readNext(); if (result == null) { return null; } else if (result.getMessage().bodyLength() > Integer.MAX_VALUE) { throw new IOException("Cannot currently deserialize record batches over 2GB"); } if (result.getMessage().version() != MetadataVersion.V4 && result.getMessage().version() != MetadataVersion.V5) { throw new IOException( "Received metadata with an incompatible version number: " + result.getMessage().version()); } switch (result.getMessage().headerType()) { case MessageHeader.RecordBatch: return deserializeRecordBatch(result.getMessage(), result.getBodyBuffer()); case MessageHeader.DictionaryBatch: return deserializeDictionaryBatch(result.getMessage(), result.getBodyBuffer()); default: throw new IOException("Unexpected message header type " + result.getMessage().headerType()); } } /** * Deserialize a message that is either an ArrowDictionaryBatch or ArrowRecordBatch. * * @param in ReadChannel to read messages from * @param alloc Allocator for message data * @return The deserialized record batch * @throws IOException if the message is not an ArrowDictionaryBatch or ArrowRecordBatch */ public static ArrowMessage deserializeMessageBatch(ReadChannel in, BufferAllocator alloc) throws IOException { return deserializeMessageBatch(new MessageChannelReader(in, alloc)); } @Deprecated public static ByteBuffer serializeMessage( FlatBufferBuilder builder, byte headerType, int headerOffset, long bodyLength) { return serializeMessage(builder, headerType, headerOffset, bodyLength, IpcOption.DEFAULT); } /** * Serializes a message header. * * @param builder to write the flatbuf to * @param headerType headerType field * @param headerOffset header offset field * @param bodyLength body length field * @param writeOption IPC write options * @return the corresponding ByteBuffer */ public static ByteBuffer serializeMessage( FlatBufferBuilder builder, byte headerType, int headerOffset, long bodyLength, IpcOption writeOption) { Message.startMessage(builder); Message.addHeaderType(builder, headerType); Message.addHeader(builder, headerOffset); Message.addVersion(builder, writeOption.metadataVersion.toFlatbufID()); Message.addBodyLength(builder, bodyLength); builder.finish(Message.endMessage(builder)); return builder.dataBuffer(); } /** * Read a Message from the input channel and return a MessageMetadataResult that contains the * Message metadata, buffer containing the serialized Message metadata as read, and length of the * Message in bytes. Returns null if the end-of-stream has been reached. * * @param in ReadChannel to read messages from * @return MessageMetadataResult with deserialized Message metadata and message information if a * valid Message was read, or null if end-of-stream * @throws IOException on error */ public static MessageMetadataResult readMessage(ReadChannel in) throws IOException { // Read the message size. There is an i32 little endian prefix. ByteBuffer buffer = ByteBuffer.allocate(4); if (in.readFully(buffer) == 4) { int messageLength = MessageSerializer.bytesToInt(buffer.array()); if (messageLength == IPC_CONTINUATION_TOKEN) { // Avoid breaking change in signature of ByteBuffer.clear() in JDK9+ ((java.nio.Buffer) buffer).clear(); // ARROW-6313, if the first 4 bytes are continuation message, read the next 4 for the length if (in.readFully(buffer) == 4) { messageLength = MessageSerializer.bytesToInt(buffer.array()); } } // Length of 0 indicates end of stream if (messageLength != 0) { // Read the message into the buffer. ByteBuffer messageBuffer = ByteBuffer.allocate(messageLength); if (in.readFully(messageBuffer) != messageLength) { throw new IOException("Unexpected end of stream trying to read message."); } // see https://github.com/apache/arrow/issues/41717 for reason why we cast to // java.nio.Buffer ByteBuffer rewindBuffer = (ByteBuffer) ((java.nio.Buffer) messageBuffer).rewind(); // Load the message. Message message = Message.getRootAsMessage(messageBuffer); return new MessageMetadataResult(messageLength, messageBuffer, message); } } return null; } /** * Read a Message body from the in channel into an ArrowBuf. * * @param in ReadChannel to read message body from * @param bodyLength Length in bytes of the message body to read * @param allocator Allocate the ArrowBuf to contain message body data * @return an ArrowBuf containing the message body data * @throws IOException on error */ public static ArrowBuf readMessageBody(ReadChannel in, long bodyLength, BufferAllocator allocator) throws IOException { ArrowBuf bodyBuffer = allocator.buffer(bodyLength); try { if (in.readFully(bodyBuffer, bodyLength) != bodyLength) { throw new IOException("Unexpected end of input trying to read batch."); } } catch (RuntimeException | IOException e) { bodyBuffer.close(); throw e; } return bodyBuffer; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/table/000077500000000000000000000000001515451615200265675ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/table/BaseTable.java000066400000000000000000000342771515451615200312710ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.table; import java.util.ArrayList; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.stream.Collectors; import org.apache.arrow.util.AutoCloseables; 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.complex.reader.FieldReader; 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.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.TransferPair; /** * Abstract base class for Table. * *

    This API is EXPERIMENTAL. */ public abstract class BaseTable implements AutoCloseable { /** The field vectors holding the data in this table. */ protected final List fieldVectors; /** * An optional DictionaryProvider. One must be present if any vector in the table is dictionary * encoded. */ protected DictionaryProvider dictionaryProvider; /** A map of Fields to FieldVectors used to select Fields. */ protected final Map fieldVectorsMap = new LinkedHashMap<>(); /** The schema for the table. */ protected Schema schema; /** * The number of rows of data in the table; not necessarily the same as the table row capacity. */ protected int rowCount; /** * Constructs new instance with the given rowCount, and containing the schema and each of the * given vectors. * * @param fieldVectors the FieldVectors containing the table's data * @param rowCount the number of rows in the table * @param provider a dictionary provider, may be null if none of the vectors in the table are * encoded */ public BaseTable(List fieldVectors, int rowCount, DictionaryProvider provider) { this.dictionaryProvider = provider; this.rowCount = rowCount; this.fieldVectors = new ArrayList<>(); List fields = new ArrayList<>(); for (FieldVector fv : fieldVectors) { TransferPair transferPair = fv.getTransferPair(fv.getAllocator()); transferPair.transfer(); FieldVector newVector = (FieldVector) transferPair.getTo(); newVector.setValueCount(rowCount); Field newField = newVector.getField(); this.fieldVectors.add(newVector); fields.add(newField); fieldVectorsMap.put(newField, newVector); } this.schema = new Schema(fields); } BaseTable() { this.fieldVectors = new ArrayList<>(); } /** * Returns a FieldReader for the vector with the given name. * * @param name The name of a vector in this Table (case-sensitive) * @return A FieldReader for the named FieldVector */ public FieldReader getReader(String name) { for (Map.Entry entry : fieldVectorsMap.entrySet()) { if (entry.getKey().getName().equals(name)) { return entry.getValue().getReader(); } } return null; } /** * Returns a FieldReader for the given field. * * @param field The field to be read * @return A FieldReader for the given field */ public FieldReader getReader(Field field) { return fieldVectorsMap.get(field).getReader(); } /** * Returns a FieldReader for the field at the given vector index. * * @param index The 0-based index of the field desired. * @return A FieldReader for the requested field */ public FieldReader getReader(int index) { Preconditions.checkArgument(index >= 0 && index < fieldVectors.size()); return fieldVectors.get(index).getReader(); } /** Returns the schema for this Table. */ public Schema getSchema() { return schema; } /** * Returns the Field with the given name if one exists in this table. * * @param fieldName the name of the field to return * @return a field with the given name if one is present * @throws IllegalArgumentException – if the field was not found */ public Field getField(String fieldName) { return getSchema().findField(fieldName); } /** * Returns a list of Field created by adding the given vector to the vectors in this Table. * * @param index field index * @param vector vector to be added. * @return out List of FieldVectors with vector added */ List insertVector(int index, FieldVector vector) { Preconditions.checkNotNull(vector); Preconditions.checkArgument(index >= 0 && index <= fieldVectors.size()); List newVectors = new ArrayList<>(); 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)); } } return newVectors; } /** * Returns a new List of FieldVectors created by removing the selected Vector from the list in * this Table. * * @param index field index * @return out List of FieldVectors like the list in this table, but with the argument removed */ List extractVector(int index) { Preconditions.checkArgument(index >= 0 && index < fieldVectors.size()); List newVectors = new ArrayList<>(); for (int i = 0; i < fieldVectors.size(); i++) { if (i != index) { newVectors.add(fieldVectors.get(i)); } } return newVectors; } /** Returns the number of vectors (columns) in this table. */ public int getVectorCount() { return fieldVectors.size(); } /** * Closes all the vectors holding data for this table and sets the rowcount to 0, preventing * enumeration. */ void clear() { close(); rowCount = 0; } /** Closes all the vectors holding data for this table. */ @Override public void close() { try { AutoCloseables.close(fieldVectors); } catch (RuntimeException ex) { throw ex; } catch (Exception ex) { // should never happen since FieldVector.close() doesn't throw IOException throw new RuntimeException(ex); } } /** Returns the number of rows in this table. */ public long getRowCount() { return rowCount; } /** * Returns a new VectorSchemaRoot with the data and schema from this table. Data is transferred to * the new VectorSchemaRoot, so this table is cleared and the rowCount is set to 0; * * @return a new VectorSchemaRoot */ public VectorSchemaRoot toVectorSchemaRoot() { VectorSchemaRoot vsr = new VectorSchemaRoot( fieldVectors.stream() .map( v -> { TransferPair transferPair = v.getTransferPair(v.getAllocator()); transferPair.transfer(); return (FieldVector) transferPair.getTo(); }) .collect(Collectors.toList())); clear(); return vsr; } /** * Returns the vector with the given name, or throws IllegalArgumentException if the name is not * found. Names are case-sensitive. * * @param columnName The name of the vector * @return the Vector with the given name, or null * @throws IllegalArgumentException if the name is not the name of a vector in the table. */ FieldVector getVector(String columnName) { for (Map.Entry entry : fieldVectorsMap.entrySet()) { if (entry.getKey().getName().equals(columnName)) { return entry.getValue(); } } throw new IllegalArgumentException( String.format("No vector named '%s' is present in the table", columnName)); } /** * Returns the vector at the given position. * * @param columnIndex The 0-based position of the vector */ FieldVector getVector(int columnIndex) { return fieldVectors.get(columnIndex); } /** * Returns a copy of the vector with the given name, or throws IllegalArgumentException if the * name is not found. Names are case-sensitive. * * @param columnName The name of the vector to copy * @return A copy of the Vector with the given name * @throws IllegalArgumentException if the name is not the name of a vector in the table. */ public FieldVector getVectorCopy(String columnName) { FieldVector source; for (Map.Entry entry : fieldVectorsMap.entrySet()) { if (entry.getKey().getName().equals(columnName)) { source = entry.getValue(); FieldVector copy = source.getField().createVector(source.getAllocator()); copy.allocateNew(); for (int i = 0; i < source.getValueCount(); i++) { copy.copyFromSafe(i, i, source); } copy.setValueCount(source.getValueCount()); return copy; } } throw new IllegalStateException( String.format("No vector named '%s' is present in the table", columnName)); } /** * Returns a copy of the vector at the given position. * * @param columnIndex The 0-based position of the vector to be copied */ public FieldVector getVectorCopy(int columnIndex) { FieldVector source = fieldVectors.get(columnIndex); FieldVector copy = source.getField().createVector(source.getAllocator()); copy.allocateNew(); for (int i = 0; i < source.getValueCount(); i++) { copy.copyFromSafe(i, i, source); } copy.setValueCount(source.getValueCount()); return copy; } /** * Returns an immutable Row object holding a reference to this table. The default character * encoding used by the cursor to decode Strings will be StandardCharsets.UTF_8 as this is the * only charset supported in Arrow format. */ public Row immutableRow() { return new Row(this); } /** Returns a tab separated value of vectors (based on their java object representation). */ public String contentToTSVString() { StringBuilder sb = new StringBuilder(); List row = new ArrayList<>(schema.getFields().size()); for (Field field : schema.getFields()) { row.add(field.getName()); } printRow(sb, row); for (int i = 0; i < rowCount; i++) { row.clear(); for (FieldVector v : fieldVectors) { row.add(v.getObject(i)); } printRow(sb, row); } return sb.toString(); } /** * Prints a single row without a header to the given StringBuilder. * * @param sb the StringBuilder to write to * @param row the row to write */ private void printRow(StringBuilder sb, List row) { boolean first = true; for (Object v : row) { if (first) { first = false; } else { sb.append("\t"); } sb.append(v); } sb.append("\n"); } /** * Returns true if the row at the given index has been deleted and false otherwise. * *

    If the index is larger than the number of rows, the method returns true. * * @param rowNumber The 0-based index of the possibly deleted row * @return true if the row at the index was deleted; false otherwise */ public boolean isRowDeleted(int rowNumber) { return false; } /** Returns the DictionaryProvider for this table. It can be used to decode an encoded values */ public DictionaryProvider getDictionaryProvider() { return dictionaryProvider; } /** * Returns a ValueVector containing the decoded version of the vector with the given name. * * @param vectorName The name of the vector to decode * @param dictionaryId The identifier for the dictionary to use when decoding. Must match the id * returned by the dictionary's getId() method. * @return A ValueVector */ public ValueVector decode(String vectorName, long dictionaryId) { Dictionary dictionary = getDictionary(dictionaryId); FieldVector vector = getVector(vectorName); if (vector == null) { throw new IllegalArgumentException( String.format("No vector with name '%s' is present in table", vectorName)); } DictionaryEncoder decoder = new DictionaryEncoder(dictionary, vector.getAllocator()); return decoder.decode(vector); } /** * Returns a ValueVector containing the encoded version of the vector with the given name. * * @param vectorName The name of the vector to encode * @param dictionaryId The identifier for the dictionary to use when encoding. Must match the id * returned by the dictionary's getId() method. * @return A ValueVector */ public ValueVector encode(String vectorName, long dictionaryId) { Dictionary dictionary = getDictionary(dictionaryId); FieldVector vector = getVector(vectorName); if (vector == null) { throw new IllegalArgumentException( String.format("No vector with name '%s' is present in table", vectorName)); } DictionaryEncoder decoder = new DictionaryEncoder(dictionary, vector.getAllocator()); return decoder.encode(vector); } /** * Returns the dictionary with given id. * * @param dictionaryId A long integer that is the id returned by the dictionary's getId() method */ private Dictionary getDictionary(long dictionaryId) { if (dictionaryProvider == null) { throw new IllegalStateException("No dictionary provider is present in table."); } Dictionary dictionary = dictionaryProvider.lookup(dictionaryId); if (dictionary == null) { throw new IllegalArgumentException("No dictionary with id '%n' exists in the table"); } return dictionary; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/table/Row.java000066400000000000000000002376211515451615200302140ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.table; import java.math.BigDecimal; import java.nio.charset.Charset; import java.nio.charset.StandardCharsets; import java.time.Duration; import java.time.LocalDateTime; import java.time.Period; import java.util.Iterator; import java.util.List; import java.util.NoSuchElementException; import org.apache.arrow.memory.ArrowBuf; 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.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.IntervalDayVector; import org.apache.arrow.vector.IntervalMonthDayNanoVector; import org.apache.arrow.vector.IntervalYearVector; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.LargeVarCharVector; import org.apache.arrow.vector.PeriodDuration; 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.ValueVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.ListVector; 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.NullableBitHolder; import org.apache.arrow.vector.holders.NullableDateDayHolder; import org.apache.arrow.vector.holders.NullableDateMilliHolder; import org.apache.arrow.vector.holders.NullableDecimalHolder; import org.apache.arrow.vector.holders.NullableDurationHolder; import org.apache.arrow.vector.holders.NullableFloat4Holder; import org.apache.arrow.vector.holders.NullableFloat8Holder; import org.apache.arrow.vector.holders.NullableIntHolder; import org.apache.arrow.vector.holders.NullableIntervalDayHolder; import org.apache.arrow.vector.holders.NullableIntervalMonthDayNanoHolder; import org.apache.arrow.vector.holders.NullableIntervalYearHolder; import org.apache.arrow.vector.holders.NullableSmallIntHolder; import org.apache.arrow.vector.holders.NullableTimeMicroHolder; import org.apache.arrow.vector.holders.NullableTimeMilliHolder; import org.apache.arrow.vector.holders.NullableTimeNanoHolder; import org.apache.arrow.vector.holders.NullableTimeSecHolder; import org.apache.arrow.vector.holders.NullableTimeStampMicroHolder; import org.apache.arrow.vector.holders.NullableTimeStampMicroTZHolder; import org.apache.arrow.vector.holders.NullableTimeStampMilliHolder; import org.apache.arrow.vector.holders.NullableTimeStampMilliTZHolder; import org.apache.arrow.vector.holders.NullableTimeStampNanoHolder; import org.apache.arrow.vector.holders.NullableTimeStampNanoTZHolder; import org.apache.arrow.vector.holders.NullableTimeStampSecHolder; import org.apache.arrow.vector.holders.NullableTimeStampSecTZHolder; import org.apache.arrow.vector.holders.NullableTinyIntHolder; import org.apache.arrow.vector.holders.NullableUInt1Holder; import org.apache.arrow.vector.holders.NullableUInt2Holder; import org.apache.arrow.vector.holders.NullableUInt4Holder; import org.apache.arrow.vector.holders.NullableUInt8Holder; /** * Row is a positionable, immutable cursor backed by a {@link Table}. * *

    Getters are provided for most vector types. The exceptions being {@link * org.apache.arrow.vector.NullVector}, which only contains null values and has no getter, and * {@link org.apache.arrow.vector.ZeroVector}, which is a zero-length vector of any type * *

    This API is EXPERIMENTAL. */ public class Row implements Iterator { /** * Returns the standard character set to use for decoding strings. The Arrow format only supports * UTF-8. */ private static final Charset DEFAULT_CHARACTER_SET = StandardCharsets.UTF_8; /** 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; /** * An iterator that returns every row in the table, deleted or not. The implemented next() and * hasNext() methods in Row wrap it with a filter to get only the non-deleted ones. */ private final Iterator iterator = intIterator(); /** * Constructs a new Row backed by the given table. * * @param table the table that this Row object represents */ public Row(BaseTable table) { this.table = table; } /** Resets the current row to -1 and returns this object. */ public Row resetPosition() { rowNumber = -1; return this; } /** * Moves this Row to the given 0-based row index. * * @return this Row for chaining */ public Row setPosition(int rowNumber) { this.rowNumber = rowNumber; this.nextRowSet = false; return this; } /** * For vectors other than Union and DenseUnion, returns true if the value at columnName is null, * and false otherwise. * *

    UnionVector#isNull always returns false, but the underlying vector may hold null values. */ public boolean isNull(String columnName) { ValueVector vector = table.getVector(columnName); return vector.isNull(rowNumber); } /** * For vectors other than Union and DenseUnion, returns true if the value at columnIndex is null, * and false otherwise. * *

    UnionVector#isNull always returns false, but the underlying vector may hold null values. */ public boolean isNull(int columnIndex) { ValueVector vector = table.getVector(columnIndex); return vector.isNull(rowNumber); } /** * Returns an object representing the value in the ExtensionTypeVector at the currentRow and * vectorIndex. An IllegalArgumentException is thrown if the column is not present in the Row and * a ClassCastException is thrown if the type is incorrect. */ public Object getExtensionType(int vectorIndex) { FieldVector vector = table.getVector(vectorIndex); return vector.getObject(rowNumber); } /** * Returns an object representing the value in the named ExtensionTypeVector at the currentRow. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type. * * @param columnName The name of the vector providing the result * @return The object in the named column at the current row */ public Object getExtensionType(String columnName) { FieldVector vector = table.getVector(columnName); return vector.getObject(rowNumber); } /** * Returns a Map from the column of the given vectorIndex at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type. */ public List getMap(int vectorIndex) { ListVector vector = (ListVector) table.getVector(vectorIndex); return vector.getObject(rowNumber); } /** * Returns a Map from the column of the given name at the current row. An IllegalArgumentException * is thrown if the column is not present in the Row and a ClassCastException is thrown if it has * a different type */ public List getMap(String columnName) { ListVector vector = (ListVector) table.getVector(columnName); return vector.getObject(rowNumber); } /** * Returns an Object from the column at vectorIndex at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type */ public Object getStruct(int vectorIndex) { StructVector vector = (StructVector) table.getVector(vectorIndex); return vector.getObject(rowNumber); } /** * Returns an Object from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type */ public Object getStruct(String columnName) { StructVector vector = (StructVector) table.getVector(columnName); return vector.getObject(rowNumber); } /** * Returns an Object from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type */ public Object getUnion(int vectorIndex) { UnionVector vector = (UnionVector) table.getVector(vectorIndex); return vector.getObject(rowNumber); } /** * Returns an Object from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type */ public Object getUnion(String columnName) { UnionVector vector = (UnionVector) table.getVector(columnName); return vector.getObject(rowNumber); } /** * Returns an Object from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type */ public Object getDenseUnion(String columnName) { DenseUnionVector vector = (DenseUnionVector) table.getVector(columnName); return vector.getObject(rowNumber); } /** * Returns an Object from the column with the given vectorIndex at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type */ public Object getDenseUnion(int vectorIndex) { DenseUnionVector vector = (DenseUnionVector) table.getVector(vectorIndex); return vector.getObject(rowNumber); } /** * Returns a List from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type */ public List getList(String columnName) { ListVector vector = (ListVector) table.getVector(columnName); return vector.getObject(rowNumber); } /** * Returns a List from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present and a ClassCastException is * thrown if it has a different type */ public List getList(int columnIndex) { ListVector vector = (ListVector) table.getVector(columnIndex); return vector.getObject(rowNumber); } /** * Returns an int from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type */ public int getInt(String columnName) { IntVector vector = (IntVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns an int from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present and a ClassCastException is * thrown if it has a different type */ public int getInt(int columnIndex) { IntVector vector = (IntVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value in the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type */ public void getInt(String columnName, NullableIntHolder holder) { IntVector vector = (IntVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value in the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present and a ClassCastException is * thrown if it has a different type */ public void getInt(int columnIndex, NullableIntHolder holder) { IntVector vector = (IntVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns an int from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type */ public int getUInt4(String columnName) { UInt4Vector vector = (UInt4Vector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns an int from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present and a ClassCastException is * thrown if it has a different type */ public int getUInt4(int columnIndex) { UInt4Vector vector = (UInt4Vector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value at the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type */ public void getUInt4(String columnName, NullableUInt4Holder holder) { UInt4Vector vector = (UInt4Vector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value at the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present and a ClassCastException is * thrown if it has a different type */ public void getUInt4(int columnIndex, NullableUInt4Holder holder) { UInt4Vector vector = (UInt4Vector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a short from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public short getSmallInt(String columnName) { SmallIntVector vector = (SmallIntVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a short from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public short getSmallInt(int columnIndex) { SmallIntVector vector = (SmallIntVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value in the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getSmallInt(String columnName, NullableSmallIntHolder holder) { SmallIntVector vector = (SmallIntVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value in the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getSmallInt(int columnIndex, NullableSmallIntHolder holder) { SmallIntVector vector = (SmallIntVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a char from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public char getUInt2(String columnName) { UInt2Vector vector = (UInt2Vector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a char from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public char getUInt2(int columnIndex) { UInt2Vector vector = (UInt2Vector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value in the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getUInt2(String columnName, NullableUInt2Holder holder) { UInt2Vector vector = (UInt2Vector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value in the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getUInt2(int columnIndex, NullableUInt2Holder holder) { UInt2Vector vector = (UInt2Vector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a byte from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public byte getTinyInt(String columnName) { TinyIntVector vector = (TinyIntVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a byte from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public byte getTinyInt(int columnIndex) { TinyIntVector vector = (TinyIntVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value in the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTinyInt(String columnName, NullableTinyIntHolder holder) { TinyIntVector vector = (TinyIntVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value in the column at the given index and current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTinyInt(int columnIndex, NullableTinyIntHolder holder) { TinyIntVector vector = (TinyIntVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a byte from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public byte getUInt1(String columnName) { UInt1Vector vector = (UInt1Vector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a byte from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public byte getUInt1(int columnIndex) { UInt1Vector vector = (UInt1Vector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value in the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getUInt1(String columnName, NullableUInt1Holder holder) { UInt1Vector vector = (UInt1Vector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value in the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getUInt1(int columnIndex, NullableUInt1Holder holder) { UInt1Vector vector = (UInt1Vector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a long from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getBigInt(String columnName) { BigIntVector vector = (BigIntVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a long from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getBigInt(int columnIndex) { BigIntVector vector = (BigIntVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value in the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getBigInt(String columnName, NullableBigIntHolder holder) { BigIntVector vector = (BigIntVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value in the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getBigInt(int columnIndex, NullableBigIntHolder holder) { BigIntVector vector = (BigIntVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a long from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getUInt8(String columnName) { UInt8Vector vector = (UInt8Vector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a long from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getUInt8(int columnIndex) { UInt8Vector vector = (UInt8Vector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value in the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getUInt8(String columnName, NullableUInt8Holder holder) { UInt8Vector vector = (UInt8Vector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value in the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getUInt8(int columnIndex, NullableUInt8Holder holder) { UInt8Vector vector = (UInt8Vector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a float from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public float getFloat4(String columnName) { Float4Vector vector = (Float4Vector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a float from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public float getFloat4(int columnIndex) { Float4Vector vector = (Float4Vector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value in the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getFloat4(String columnName, NullableFloat4Holder holder) { Float4Vector vector = (Float4Vector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value in the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getFloat4(int columnIndex, NullableFloat4Holder holder) { Float4Vector vector = (Float4Vector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a double from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public double getFloat8(String columnName) { Float8Vector vector = (Float8Vector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a double from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public double getFloat8(int columnIndex) { Float8Vector vector = (Float8Vector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value in the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getFloat8(String columnName, NullableFloat8Holder holder) { Float8Vector vector = (Float8Vector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value in the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getFloat8(int columnIndex, NullableFloat8Holder holder) { Float8Vector vector = (Float8Vector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns an int from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public int getBit(String columnName) { BitVector vector = (BitVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns an int from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public int getBit(int columnIndex) { BitVector vector = (BitVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value in the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getBit(String columnName, NullableBitHolder holder) { BitVector vector = (BitVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value in the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getBit(int columnIndex, NullableBitHolder holder) { BitVector vector = (BitVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a long from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public long getDateMilli(String columnName) { DateMilliVector vector = (DateMilliVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a long from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public long getDateMilli(int columnIndex) { DateMilliVector vector = (DateMilliVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value in the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getDateMilli(String columnName, NullableDateMilliHolder holder) { DateMilliVector vector = (DateMilliVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value in the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getDateMilli(int columnIndex, NullableDateMilliHolder holder) { DateMilliVector vector = (DateMilliVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns an int from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public int getDateDay(String columnName) { DateDayVector vector = (DateDayVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns an int from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public int getDateDay(int columnIndex) { DateDayVector vector = (DateDayVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value in the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getDateDay(String columnName, NullableDateDayHolder holder) { DateDayVector vector = (DateDayVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value in the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getDateDay(int columnIndex, NullableDateDayHolder holder) { DateDayVector vector = (DateDayVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a long from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getTimeNano(String columnName) { TimeNanoVector vector = (TimeNanoVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a long from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getTimeNano(int columnIndex) { TimeNanoVector vector = (TimeNanoVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value in the column with the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTimeNano(String columnName, NullableTimeNanoHolder holder) { TimeNanoVector vector = (TimeNanoVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value in the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public void getTimeNano(int columnIndex, NullableTimeNanoHolder holder) { TimeNanoVector vector = (TimeNanoVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a long from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public long getTimeMicro(String columnName) { TimeMicroVector vector = (TimeMicroVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a long from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public long getTimeMicro(int columnIndex) { TimeMicroVector vector = (TimeMicroVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public void getTimeMicro(String columnName, NullableTimeMicroHolder holder) { TimeMicroVector vector = (TimeMicroVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public void getTimeMicro(int columnIndex, NullableTimeMicroHolder holder) { TimeMicroVector vector = (TimeMicroVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns an int from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public int getTimeMilli(String columnName) { TimeMilliVector vector = (TimeMilliVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns an int from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public int getTimeMilli(int columnIndex) { TimeMilliVector vector = (TimeMilliVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public void getTimeMilli(String columnName, NullableTimeMilliHolder holder) { TimeMilliVector vector = (TimeMilliVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public void getTimeMilli(int columnIndex, NullableTimeMilliHolder holder) { TimeMilliVector vector = (TimeMilliVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a LocalDateTime from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public LocalDateTime getTimeMilliObj(String columnName) { TimeMilliVector vector = (TimeMilliVector) table.getVector(columnName); return vector.getObject(rowNumber); } /** * Returns a LocalDateTime from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public LocalDateTime getTimeMilliObj(int columnIndex) { TimeMilliVector vector = (TimeMilliVector) table.getVector(columnIndex); return vector.getObject(rowNumber); } /** * Returns an int from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public int getTimeSec(String columnName) { TimeSecVector vector = (TimeSecVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns an int from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public int getTimeSec(int columnIndex) { TimeSecVector vector = (TimeSecVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public void getTimeSec(String columnName, NullableTimeSecHolder holder) { TimeSecVector vector = (TimeSecVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public void getTimeSec(int columnIndex, NullableTimeSecHolder holder) { TimeSecVector vector = (TimeSecVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a long from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type. */ public long getTimeStampSec(String columnName) { TimeStampSecVector vector = (TimeStampSecVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a long from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getTimeStampSec(int columnIndex) { TimeStampSecVector vector = (TimeStampSecVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTimeStampSec(String columnName, NullableTimeStampSecHolder holder) { TimeStampSecVector vector = (TimeStampSecVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTimeStampSec(int columnIndex, NullableTimeStampSecHolder holder) { TimeStampSecVector vector = (TimeStampSecVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a LocalDateTime from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public LocalDateTime getTimeStampSecObj(String columnName) { TimeStampSecVector vector = (TimeStampSecVector) table.getVector(columnName); return vector.getObject(rowNumber); } /** * Returns a LocalDateTime from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public LocalDateTime getTimeStampSecObj(int columnIndex) { TimeStampSecVector vector = (TimeStampSecVector) table.getVector(columnIndex); return vector.getObject(rowNumber); } /** * Updates the holder with the value from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getTimeStampSecTZ(String columnName) { TimeStampSecTZVector vector = (TimeStampSecTZVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a long from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getTimeStampSecTZ(int columnIndex) { TimeStampSecTZVector vector = (TimeStampSecTZVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTimeStampSecTZ(String columnName, NullableTimeStampSecTZHolder holder) { TimeStampSecTZVector vector = (TimeStampSecTZVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTimeStampSecTZ(int columnIndex, NullableTimeStampSecTZHolder holder) { TimeStampSecTZVector vector = (TimeStampSecTZVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a long from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getTimeStampNano(String columnName) { TimeStampNanoVector vector = (TimeStampNanoVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a long from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getTimeStampNano(int columnIndex) { TimeStampNanoVector vector = (TimeStampNanoVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTimeStampNano(String columnName, NullableTimeStampNanoHolder holder) { TimeStampNanoVector vector = (TimeStampNanoVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTimeStampNano(int columnIndex, NullableTimeStampNanoHolder holder) { TimeStampNanoVector vector = (TimeStampNanoVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a LocalDateTime from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public LocalDateTime getTimeStampNanoObj(String columnName) { TimeStampNanoVector vector = (TimeStampNanoVector) table.getVector(columnName); return vector.getObject(rowNumber); } /** * Returns a LocalDateTime from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public LocalDateTime getTimeStampNanoObj(int columnIndex) { TimeStampNanoVector vector = (TimeStampNanoVector) table.getVector(columnIndex); return vector.getObject(rowNumber); } /** * Returns a long from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getTimeStampNanoTZ(String columnName) { TimeStampNanoTZVector vector = (TimeStampNanoTZVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a long from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getTimeStampNanoTZ(int columnIndex) { TimeStampNanoTZVector vector = (TimeStampNanoTZVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTimeStampNanoTZ(String columnName, NullableTimeStampNanoTZHolder holder) { TimeStampNanoTZVector vector = (TimeStampNanoTZVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTimeStampNanoTZ(int columnIndex, NullableTimeStampNanoTZHolder holder) { TimeStampNanoTZVector vector = (TimeStampNanoTZVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a long from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getTimeStampMilli(String columnName) { TimeStampMilliVector vector = (TimeStampMilliVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a long from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getTimeStampMilli(int columnIndex) { TimeStampMilliVector vector = (TimeStampMilliVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTimeStampMilli(String columnName, NullableTimeStampMilliHolder holder) { TimeStampMilliVector vector = (TimeStampMilliVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTimeStampMilli(int columnIndex, NullableTimeStampMilliHolder holder) { TimeStampMilliVector vector = (TimeStampMilliVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a LocalDateTime from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public LocalDateTime getTimeStampMilliObj(String columnName) { TimeStampMilliVector vector = (TimeStampMilliVector) table.getVector(columnName); return vector.getObject(rowNumber); } /** * Returns a LocalDateTime from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public LocalDateTime getTimeStampMilliObj(int columnIndex) { TimeStampMilliVector vector = (TimeStampMilliVector) table.getVector(columnIndex); return vector.getObject(rowNumber); } /** * Returns a long from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getTimeStampMilliTZ(String columnName) { TimeStampMilliTZVector vector = (TimeStampMilliTZVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a long from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getTimeStampMilliTZ(int columnIndex) { TimeStampMilliTZVector vector = (TimeStampMilliTZVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different types */ public void getTimeStampMilliTZ(String columnName, NullableTimeStampMilliTZHolder holder) { TimeStampMilliTZVector vector = (TimeStampMilliTZVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTimeStampMilliTZ(int columnIndex, NullableTimeStampMilliTZHolder holder) { TimeStampMilliTZVector vector = (TimeStampMilliTZVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a long from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getTimeStampMicro(String columnName) { TimeStampMicroVector vector = (TimeStampMicroVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a long from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getTimeStampMicro(int columnIndex) { TimeStampMicroVector vector = (TimeStampMicroVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTimeStampMicro(String columnName, NullableTimeStampMicroHolder holder) { TimeStampMicroVector vector = (TimeStampMicroVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTimeStampMicro(int columnIndex, NullableTimeStampMicroHolder holder) { TimeStampMicroVector vector = (TimeStampMicroVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a LocalDateTime from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public LocalDateTime getTimeStampMicroObj(String columnName) { TimeStampMicroVector vector = (TimeStampMicroVector) table.getVector(columnName); return vector.getObject(rowNumber); } /** * Returns a LocalDateTime from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public LocalDateTime getTimeStampMicroObj(int columnIndex) { TimeStampMicroVector vector = (TimeStampMicroVector) table.getVector(columnIndex); return vector.getObject(rowNumber); } /** * Returns a long from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getTimeStampMicroTZ(String columnName) { TimeStampMicroTZVector vector = (TimeStampMicroTZVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a long from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public long getTimeStampMicroTZ(int columnIndex) { TimeStampMicroTZVector vector = (TimeStampMicroTZVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTimeStampMicroTZ(String columnName, NullableTimeStampMicroTZHolder holder) { TimeStampMicroTZVector vector = (TimeStampMicroTZVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getTimeStampMicroTZ(int columnIndex, NullableTimeStampMicroTZHolder holder) { TimeStampMicroTZVector vector = (TimeStampMicroTZVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a Duration from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public Duration getDurationObj(String columnName) { DurationVector vector = (DurationVector) table.getVector(columnName); return vector.getObject(rowNumber); } /** * Returns a Duration from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public Duration getDurationObj(int columnIndex) { DurationVector vector = (DurationVector) table.getVector(columnIndex); return vector.getObject(rowNumber); } /** * Returns an ArrowBuf from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public ArrowBuf getDuration(String columnName) { DurationVector vector = (DurationVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns an ArrowBuf from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public ArrowBuf getDuration(int columnIndex) { DurationVector vector = (DurationVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getDuration(String columnName, NullableDurationHolder holder) { DurationVector vector = (DurationVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getDuration(int columnIndex, NullableDurationHolder holder) { DurationVector vector = (DurationVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a PeriodDuration from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public PeriodDuration getIntervalMonthDayNanoObj(String columnName) { IntervalMonthDayNanoVector vector = (IntervalMonthDayNanoVector) table.getVector(columnName); return vector.getObject(rowNumber); } /** * Returns a PeriodDuration from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public PeriodDuration getIntervalMonthDayNanoObj(int columnIndex) { IntervalMonthDayNanoVector vector = (IntervalMonthDayNanoVector) table.getVector(columnIndex); return vector.getObject(rowNumber); } /** * Returns an ArrowBuf from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public ArrowBuf getIntervalMonthDayNano(String columnName) { IntervalMonthDayNanoVector vector = (IntervalMonthDayNanoVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns an ArrowBuf from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public ArrowBuf getIntervalMonthDayNano(int columnIndex) { IntervalMonthDayNanoVector vector = (IntervalMonthDayNanoVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getIntervalMonthDayNano( String columnName, NullableIntervalMonthDayNanoHolder holder) { IntervalMonthDayNanoVector vector = (IntervalMonthDayNanoVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getIntervalMonthDayNano(int columnIndex, NullableIntervalMonthDayNanoHolder holder) { IntervalMonthDayNanoVector vector = (IntervalMonthDayNanoVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns an ArrowBuf from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public ArrowBuf getIntervalDay(String columnName) { IntervalDayVector vector = (IntervalDayVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns an ArrowBuf from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public ArrowBuf getIntervalDay(int columnIndex) { IntervalDayVector vector = (IntervalDayVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getIntervalDay(String columnName, NullableIntervalDayHolder holder) { IntervalDayVector vector = (IntervalDayVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getIntervalDay(int columnIndex, NullableIntervalDayHolder holder) { IntervalDayVector vector = (IntervalDayVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Returns a Duration from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public Duration getIntervalDayObj(int columnIndex) { IntervalDayVector vector = (IntervalDayVector) table.getVector(columnIndex); return vector.getObject(rowNumber); } /** * Returns a Duration from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public Duration getIntervalDayObj(String columnName) { IntervalDayVector vector = (IntervalDayVector) table.getVector(columnName); return vector.getObject(rowNumber); } /** * Returns a Period from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type * * @return a Period of n MONTHS, not YEARS */ public Period getIntervalYearObj(String columnName) { IntervalYearVector vector = (IntervalYearVector) table.getVector(columnName); return vector.getObject(rowNumber); } /** * Returns a Period from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type * * @return a Period of n MONTHS, not YEARS */ public Period getIntervalYearObj(int columnIndex) { IntervalYearVector vector = (IntervalYearVector) table.getVector(columnIndex); return vector.getObject(rowNumber); } /** * Returns an int from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type * * @return the number of MONTHS in the interval (not YEARS) */ public int getIntervalYear(String columnName) { IntervalYearVector vector = (IntervalYearVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns an int from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type * * @return the number of MONTHS in the interval (not YEARS) */ public int getIntervalYear(int columnIndex) { IntervalYearVector vector = (IntervalYearVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Updates the holder with the value from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type * * @param holder a holder to store the interval. Note that the value of the holder represents * MONTHS not years */ public void getIntervalYear(String columnName, NullableIntervalYearHolder holder) { IntervalYearVector vector = (IntervalYearVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Updates the holder with the value from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type * * @param holder a holder to store the interval. Note that the value of the holder represents * MONTHS not years */ public void getIntervalYear(int columnIndex, NullableIntervalYearHolder holder) { IntervalYearVector vector = (IntervalYearVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Updates the value of the holder with data from vector at the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getDecimal(int columnIndex, NullableDecimalHolder holder) { DecimalVector vector = (DecimalVector) table.getVector(columnIndex); vector.get(rowNumber, holder); } /** * Updates the value of the holder with data from the vector with given name at the current row. * An IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public void getDecimal(String columnName, NullableDecimalHolder holder) { DecimalVector vector = (DecimalVector) table.getVector(columnName); vector.get(rowNumber, holder); } /** * Returns a BigDecimal from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public BigDecimal getDecimalObj(String columnName) { DecimalVector vector = (DecimalVector) table.getVector(columnName); return vector.getObject(rowNumber); } /** * Returns a BigDecimal from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public BigDecimal getDecimalObj(int columnIndex) { DecimalVector vector = (DecimalVector) table.getVector(columnIndex); return vector.getObject(rowNumber); } /** * Returns an ArrowBuf from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public ArrowBuf getDecimal(String columnName) { DecimalVector vector = (DecimalVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns an ArrowBuf from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public ArrowBuf getDecimal(int columnIndex) { DecimalVector vector = (DecimalVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Returns a byte[] from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public byte[] getVarBinary(String columnName) { VarBinaryVector vector = (VarBinaryVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a byte[] from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public byte[] getVarBinary(int columnIndex) { VarBinaryVector vector = (VarBinaryVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Returns a byte[] from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public byte[] getFixedSizeBinary(String columnName) { FixedSizeBinaryVector vector = (FixedSizeBinaryVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a byte[] from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public byte[] getFixedSizeBinary(int columnIndex) { FixedSizeBinaryVector vector = (FixedSizeBinaryVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Returns a byte[] from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public byte[] getLargeVarBinary(String columnName) { LargeVarBinaryVector vector = (LargeVarBinaryVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a byte[] from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present, and a ClassCastException is * thrown if it is present but has a different type */ public byte[] getLargeVarBinary(int columnIndex) { LargeVarBinaryVector vector = (LargeVarBinaryVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Returns a String from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type * *

    StandardCharsets.UTF_8 is used as the charset */ public String getVarCharObj(String columnName) { VarCharVector vector = (VarCharVector) table.getVector(columnName); return new String(vector.get(rowNumber), getDefaultCharacterSet()); } /** * Returns a String from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type * * @param columnIndex the index of the FieldVector holding the value */ public String getVarCharObj(int columnIndex) { VarCharVector vector = (VarCharVector) table.getVector(columnIndex); return new String(vector.get(rowNumber), getDefaultCharacterSet()); } /** * Returns a byte[] from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type * *

    StandardCharsets.UTF_8 is used as the charset */ public byte[] getVarChar(String columnName) { VarCharVector vector = (VarCharVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a byte[] from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type * * @param columnIndex the index of the FieldVector holding the value */ public byte[] getVarChar(int columnIndex) { VarCharVector vector = (VarCharVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Returns a String from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type * *

    StandardCharsets.UTF_8 is used as the charset, unless this cursor was created with a default * Charset */ public String getLargeVarCharObj(String columnName) { LargeVarCharVector vector = (LargeVarCharVector) table.getVector(columnName); return new String(vector.get(rowNumber), getDefaultCharacterSet()); } /** * Returns a String from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type */ public String getLargeVarCharObj(int columnIndex) { LargeVarCharVector vector = (LargeVarCharVector) table.getVector(columnIndex); return new String(vector.get(rowNumber), getDefaultCharacterSet()); } /** * Returns a byte[] from the column of the given name at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type * *

    StandardCharsets.UTF_8 is used as the charset, unless this cursor was created with a default * Charset */ public byte[] getLargeVarChar(String columnName) { LargeVarCharVector vector = (LargeVarCharVector) table.getVector(columnName); return vector.get(rowNumber); } /** * Returns a byte[] from the column with the given index at the current row. An * IllegalArgumentException is thrown if the column is not present in the Row and a * ClassCastException is thrown if it has a different type */ public byte[] getLargeVarChar(int columnIndex) { LargeVarCharVector vector = (LargeVarCharVector) table.getVector(columnIndex); return vector.get(rowNumber); } /** * Returns true if there is at least one more non-deleted row in the table that has yet to be * processed. */ @Override public boolean hasNext() { return nextRowSet || setNextObject(); } /** * Returns the next non-deleted row in the table. * * @throws NoSuchElementException if there are no more rows */ @Override public Row next() { if (!nextRowSet && !setNextObject()) { throw new NoSuchElementException(); } nextRowSet = false; return this; } /** * Set rowNumber to the next non-deleted row. If there are no more rows return false. Otherwise, * return true. */ private boolean setNextObject() { while (iterator.hasNext()) { final int row = iterator.next(); if (!rowIsDeleted(row)) { rowNumber = row; nextRowSet = true; return true; } } return false; } /** * Returns new internal iterator that processes every row, deleted or not. Use the wrapping next() * and hasNext() methods rather than using this iterator directly, unless you want to see any * deleted rows. */ private Iterator intIterator() { return new Iterator() { @Override public boolean hasNext() { return rowNumber < table.getRowCount() - 1; } @Override public Integer next() { rowNumber++; return rowNumber; } }; } public int getRowNumber() { return rowNumber; } private boolean rowIsDeleted(int rowNumber) { return table.isRowDeleted(rowNumber); } /** Returns the default character set for use with character vectors. */ public Charset getDefaultCharacterSet() { return DEFAULT_CHARACTER_SET; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/table/Table.java000066400000000000000000000162041515451615200304640ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.table; import java.util.ArrayList; import java.util.Arrays; import java.util.Iterator; import java.util.List; import java.util.Set; import java.util.stream.Collectors; import java.util.stream.StreamSupport; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.FieldVector; 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.types.pojo.DictionaryEncoding; import org.apache.arrow.vector.util.TransferPair; /** * Table is an immutable tabular data structure. * *

    See {@link VectorSchemaRoot} for batch processing use cases * *

    This API is EXPERIMENTAL. */ public class Table extends BaseTable implements Iterable { /** Constructs new instance containing each of the given vectors. */ public Table(Iterable vectors) { this(StreamSupport.stream(vectors.spliterator(), false).collect(Collectors.toList())); } /** Constructs a new instance from vectors. */ public static Table of(FieldVector... vectors) { return new Table(Arrays.stream(vectors).collect(Collectors.toList())); } /** * Constructs a new instance with the number of rows set to the value count of the first * FieldVector. * *

    All vectors must have the same value count. Although this is not checked, inconsistent * counts may lead to exceptions or other undefined behavior later. * * @param fieldVectors The data vectors (must be equal in size to fields. */ public Table(List fieldVectors) { this(fieldVectors, fieldVectors.size() == 0 ? 0 : fieldVectors.get(0).getValueCount()); } /** * Constructs a new instance. * * @param fieldVectors The data vectors. * @param rowCount The number of rows */ public Table(List fieldVectors, int rowCount) { super(fieldVectors, rowCount, null); } /** * Constructs a new instance. * * @param fieldVectors The data vectors. * @param rowCount The number of rows * @param provider A dictionary provider. May be null if none of the vectors is dictionary encoded */ public Table(List fieldVectors, int rowCount, DictionaryProvider provider) { super(fieldVectors, rowCount, provider); } /** * Constructs a new instance containing the data from the argument. Vectors are shared between the * Table and VectorSchemaRoot. Direct modification of those vectors is unsafe and should be * avoided. * * @param vsr The VectorSchemaRoot providing data for this Table */ public Table(VectorSchemaRoot vsr) { this(vsr.getFieldVectors(), vsr.getRowCount()); vsr.clear(); } /** Returns a deep copy of this table. */ public Table copy() { List vectorCopies = new ArrayList<>(); for (int i = 0; i < getVectorCount(); i++) { vectorCopies.add(getVectorCopy(i)); } DictionaryProvider providerCopy = null; if (dictionaryProvider != null) { Set ids = dictionaryProvider.getDictionaryIds(); Dictionary[] dictionaryCopies = new Dictionary[ids.size()]; int i = 0; for (Long id : ids) { Dictionary src = dictionaryProvider.lookup(id); FieldVector srcVector = src.getVector(); FieldVector destVector = srcVector.getField().createVector(srcVector.getAllocator()); destVector.copyFromSafe( 0, srcVector.getValueCount(), srcVector); // TODO: Remove safe copy for perf DictionaryEncoding srcEncoding = src.getEncoding(); Dictionary dest = new Dictionary( destVector, new DictionaryEncoding( srcEncoding.getId(), srcEncoding.isOrdered(), srcEncoding.getIndexType())); dictionaryCopies[i] = dest; i++; } providerCopy = new DictionaryProvider.MapDictionaryProvider(dictionaryCopies); } return new Table(vectorCopies, (int) getRowCount(), providerCopy); } /** * Returns a new Table created by adding the given vector to the vectors in this Table. * * @param index field index * @param vector vector to be added. * @return out a new Table with vector added */ public Table addVector(int index, FieldVector vector) { return new Table(insertVector(index, vector)); } /** * Returns a new Table created by removing the selected Vector from this Table. * * @param index field index * @return out a new Table with vector removed */ public Table removeVector(int index) { return new Table(extractVector(index)); } /** * Slice this table from desired index. Memory is NOT transferred from the vectors in this table * to new vectors in the target table. This table is unchanged. * * @param index start position of the slice * @return the sliced table */ public Table slice(int index) { return slice(index, this.rowCount - index); } /** * Slice this table at desired index and length. Memory is NOT transferred from the vectors in * this table to new vectors in the target table. This table is unchanged. * * @param index start position of the slice * @param length length of the slice * @return the sliced table */ public Table slice(int index, int length) { Preconditions.checkArgument(index >= 0, "expecting non-negative index"); Preconditions.checkArgument(length >= 0, "expecting non-negative length"); Preconditions.checkArgument(index + length <= rowCount, "index + length should <= rowCount"); if (index == 0 && length == rowCount) { return this; } List sliceVectors = fieldVectors.stream() .map( v -> { TransferPair transferPair = v.getTransferPair(v.getAllocator()); transferPair.splitAndTransfer(index, length); return (FieldVector) transferPair.getTo(); }) .collect(Collectors.toList()); return new Table(sliceVectors); } /** Returns a Row iterator for this Table. */ @Override public Iterator iterator() { return new Iterator() { private final Row row = new Row(Table.this); @Override public Row next() { row.next(); return row; } @Override public boolean hasNext() { return row.hasNext(); } }; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/table/package-info.java000066400000000000000000000020421515451615200317540ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.table; /* * Support for Table, an immutable, columnar, tabular data structure based on FieldVectors. * See the Arrow Java documentation for details: Table */ apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/types/000077500000000000000000000000001515451615200266445ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/types/DateUnit.java000066400000000000000000000030001515451615200312150ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.arrow.vector.types; /** Resolutions that dates can be stored at. */ public enum DateUnit { /** Days since epoch. */ DAY(org.apache.arrow.flatbuf.DateUnit.DAY), /** Milliseconds since epoch. */ MILLISECOND(org.apache.arrow.flatbuf.DateUnit.MILLISECOND); private static final DateUnit[] valuesByFlatbufId = new DateUnit[DateUnit.values().length]; static { for (DateUnit v : DateUnit.values()) { valuesByFlatbufId[v.flatbufID] = v; } } private final short flatbufID; DateUnit(short flatbufID) { this.flatbufID = flatbufID; } public short getFlatbufID() { return flatbufID; } public static DateUnit fromFlatbufID(short id) { return valuesByFlatbufId[id]; } } FloatingPointPrecision.java000066400000000000000000000033201515451615200340570ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/types/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.arrow.vector.types; import org.apache.arrow.flatbuf.Precision; /** Precisions of primitive floating point numbers. */ public enum FloatingPointPrecision { /** 16-bit (not a standard java type). */ HALF(Precision.HALF), /** 32-bit (i.e. float in java). */ SINGLE(Precision.SINGLE), /** 64-bit (i.e. double in java). */ DOUBLE(Precision.DOUBLE); private static final FloatingPointPrecision[] valuesByFlatbufId = new FloatingPointPrecision[FloatingPointPrecision.values().length]; static { for (FloatingPointPrecision v : FloatingPointPrecision.values()) { valuesByFlatbufId[v.flatbufID] = v; } } private final short flatbufID; private FloatingPointPrecision(short flatbufID) { this.flatbufID = flatbufID; } public short getFlatbufID() { return flatbufID; } public static FloatingPointPrecision fromFlatbufID(short id) { return valuesByFlatbufId[id]; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/types/IntervalUnit.java000066400000000000000000000035431515451615200321400ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.arrow.vector.types; /** Resolutions for Interval Vectors. */ public enum IntervalUnit { /** * Values are stored as number of months (which can be converted into years and months via * division). */ YEAR_MONTH(org.apache.arrow.flatbuf.IntervalUnit.YEAR_MONTH), /** Values are stored as some number of days and some number of milliseconds within that day. */ DAY_TIME(org.apache.arrow.flatbuf.IntervalUnit.DAY_TIME), /** Values are stored as number of months, days and nanoseconds. */ MONTH_DAY_NANO(org.apache.arrow.flatbuf.IntervalUnit.MONTH_DAY_NANO); private static final IntervalUnit[] valuesByFlatbufId = new IntervalUnit[IntervalUnit.values().length]; static { for (IntervalUnit v : IntervalUnit.values()) { valuesByFlatbufId[v.flatbufID] = v; } } private final short flatbufID; private IntervalUnit(short flatbufID) { this.flatbufID = flatbufID; } public short getFlatbufID() { return flatbufID; } public static IntervalUnit fromFlatbufID(short id) { return valuesByFlatbufId[id]; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/types/MetadataVersion.java000066400000000000000000000034351515451615200326020ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.arrow.vector.types; /** Metadata version for Arrow metadata. */ public enum MetadataVersion { /// 0.1.0 V1(org.apache.arrow.flatbuf.MetadataVersion.V1), /// 0.2.0 V2(org.apache.arrow.flatbuf.MetadataVersion.V2), /// 0.3.0 to 0.7.1 V3(org.apache.arrow.flatbuf.MetadataVersion.V3), /// 0.8.0 to 0.17.1 V4(org.apache.arrow.flatbuf.MetadataVersion.V4), /// >= 1.0.0 V5(org.apache.arrow.flatbuf.MetadataVersion.V5), ; public static final MetadataVersion DEFAULT = V5; private static final MetadataVersion[] valuesByFlatbufId = new MetadataVersion[MetadataVersion.values().length]; static { for (MetadataVersion v : MetadataVersion.values()) { valuesByFlatbufId[v.flatbufID] = v; } } private final short flatbufID; MetadataVersion(short flatbufID) { this.flatbufID = flatbufID; } public short toFlatbufID() { return flatbufID; } public static MetadataVersion fromFlatbufID(short id) { return valuesByFlatbufId[id]; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/types/TimeUnit.java000066400000000000000000000031041515451615200312430ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.arrow.vector.types; /** Resolutions that times can be stored with. */ public enum TimeUnit { SECOND(org.apache.arrow.flatbuf.TimeUnit.SECOND), MILLISECOND(org.apache.arrow.flatbuf.TimeUnit.MILLISECOND), MICROSECOND(org.apache.arrow.flatbuf.TimeUnit.MICROSECOND), NANOSECOND(org.apache.arrow.flatbuf.TimeUnit.NANOSECOND); private static final TimeUnit[] valuesByFlatbufId = new TimeUnit[TimeUnit.values().length]; static { for (TimeUnit v : TimeUnit.values()) { valuesByFlatbufId[v.flatbufID] = v; } } private final short flatbufID; TimeUnit(short flatbufID) { this.flatbufID = flatbufID; } public short getFlatbufID() { return flatbufID; } public static TimeUnit fromFlatbufID(short id) { return valuesByFlatbufId[id]; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/types/Types.java000066400000000000000000001134101515451615200306130ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.arrow.vector.types; import static org.apache.arrow.vector.types.FloatingPointPrecision.DOUBLE; import static org.apache.arrow.vector.types.FloatingPointPrecision.HALF; import static org.apache.arrow.vector.types.FloatingPointPrecision.SINGLE; import static org.apache.arrow.vector.types.UnionMode.Dense; import static org.apache.arrow.vector.types.UnionMode.Sparse; import org.apache.arrow.memory.BufferAllocator; 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.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; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.IntervalDayVector; import org.apache.arrow.vector.IntervalMonthDayNanoVector; import org.apache.arrow.vector.IntervalYearVector; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.LargeVarCharVector; 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.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; 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.MapVector; 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.BigIntWriterImpl; import org.apache.arrow.vector.complex.impl.BitWriterImpl; import org.apache.arrow.vector.complex.impl.DateDayWriterImpl; import org.apache.arrow.vector.complex.impl.DateMilliWriterImpl; import org.apache.arrow.vector.complex.impl.Decimal256WriterImpl; import org.apache.arrow.vector.complex.impl.DecimalWriterImpl; import org.apache.arrow.vector.complex.impl.DenseUnionWriter; import org.apache.arrow.vector.complex.impl.DurationWriterImpl; import org.apache.arrow.vector.complex.impl.FixedSizeBinaryWriterImpl; import org.apache.arrow.vector.complex.impl.Float2WriterImpl; import org.apache.arrow.vector.complex.impl.Float4WriterImpl; import org.apache.arrow.vector.complex.impl.Float8WriterImpl; import org.apache.arrow.vector.complex.impl.IntWriterImpl; import org.apache.arrow.vector.complex.impl.IntervalDayWriterImpl; import org.apache.arrow.vector.complex.impl.IntervalMonthDayNanoWriterImpl; import org.apache.arrow.vector.complex.impl.IntervalYearWriterImpl; import org.apache.arrow.vector.complex.impl.LargeVarBinaryWriterImpl; import org.apache.arrow.vector.complex.impl.LargeVarCharWriterImpl; import org.apache.arrow.vector.complex.impl.NullableStructWriter; import org.apache.arrow.vector.complex.impl.SmallIntWriterImpl; import org.apache.arrow.vector.complex.impl.TimeMicroWriterImpl; import org.apache.arrow.vector.complex.impl.TimeMilliWriterImpl; import org.apache.arrow.vector.complex.impl.TimeNanoWriterImpl; import org.apache.arrow.vector.complex.impl.TimeSecWriterImpl; import org.apache.arrow.vector.complex.impl.TimeStampMicroTZWriterImpl; import org.apache.arrow.vector.complex.impl.TimeStampMicroWriterImpl; import org.apache.arrow.vector.complex.impl.TimeStampMilliTZWriterImpl; import org.apache.arrow.vector.complex.impl.TimeStampMilliWriterImpl; import org.apache.arrow.vector.complex.impl.TimeStampNanoTZWriterImpl; import org.apache.arrow.vector.complex.impl.TimeStampNanoWriterImpl; import org.apache.arrow.vector.complex.impl.TimeStampSecTZWriterImpl; import org.apache.arrow.vector.complex.impl.TimeStampSecWriterImpl; import org.apache.arrow.vector.complex.impl.TinyIntWriterImpl; import org.apache.arrow.vector.complex.impl.UInt1WriterImpl; import org.apache.arrow.vector.complex.impl.UInt2WriterImpl; import org.apache.arrow.vector.complex.impl.UInt4WriterImpl; import org.apache.arrow.vector.complex.impl.UInt8WriterImpl; 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; import org.apache.arrow.vector.complex.impl.ViewVarBinaryWriterImpl; import org.apache.arrow.vector.complex.impl.ViewVarCharWriterImpl; import org.apache.arrow.vector.complex.writer.FieldWriter; 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.Binary; import org.apache.arrow.vector.types.pojo.ArrowType.BinaryView; import org.apache.arrow.vector.types.pojo.ArrowType.Bool; import org.apache.arrow.vector.types.pojo.ArrowType.Date; import org.apache.arrow.vector.types.pojo.ArrowType.Decimal; import org.apache.arrow.vector.types.pojo.ArrowType.Duration; import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; import org.apache.arrow.vector.types.pojo.ArrowType.FixedSizeBinary; import org.apache.arrow.vector.types.pojo.ArrowType.FixedSizeList; import org.apache.arrow.vector.types.pojo.ArrowType.FloatingPoint; import org.apache.arrow.vector.types.pojo.ArrowType.Int; import org.apache.arrow.vector.types.pojo.ArrowType.Interval; import org.apache.arrow.vector.types.pojo.ArrowType.LargeBinary; import org.apache.arrow.vector.types.pojo.ArrowType.LargeListView; import org.apache.arrow.vector.types.pojo.ArrowType.LargeUtf8; import org.apache.arrow.vector.types.pojo.ArrowType.List; import org.apache.arrow.vector.types.pojo.ArrowType.ListView; import org.apache.arrow.vector.types.pojo.ArrowType.Map; import org.apache.arrow.vector.types.pojo.ArrowType.Null; import org.apache.arrow.vector.types.pojo.ArrowType.RunEndEncoded; import org.apache.arrow.vector.types.pojo.ArrowType.Struct; import org.apache.arrow.vector.types.pojo.ArrowType.Time; import org.apache.arrow.vector.types.pojo.ArrowType.Timestamp; import org.apache.arrow.vector.types.pojo.ArrowType.Union; import org.apache.arrow.vector.types.pojo.ArrowType.Utf8; import org.apache.arrow.vector.types.pojo.ArrowType.Utf8View; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.CallBack; /** An enumeration of all logical types supported by this library. */ public class Types { /** The actual enumeration of types. */ public enum MinorType { NULL(Null.INSTANCE) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new NullVector(field.getName()); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return null; } }, STRUCT(Struct.INSTANCE) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new StructVector( field.getName(), allocator, field.getFieldType(), schemaChangeCallback); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new NullableStructWriter((StructVector) vector); } }, TINYINT(new Int(8, true)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new TinyIntVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new TinyIntWriterImpl((TinyIntVector) vector); } }, SMALLINT(new Int(16, true)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new SmallIntVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new SmallIntWriterImpl((SmallIntVector) vector); } }, INT(new Int(32, true)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new IntVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new IntWriterImpl((IntVector) vector); } }, BIGINT(new Int(64, true)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new BigIntVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new BigIntWriterImpl((BigIntVector) vector); } }, DATEDAY(new Date(DateUnit.DAY)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new DateDayVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new DateDayWriterImpl((DateDayVector) vector); } }, DATEMILLI(new Date(DateUnit.MILLISECOND)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new DateMilliVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new DateMilliWriterImpl((DateMilliVector) vector); } }, TIMESEC(new Time(TimeUnit.SECOND, 32)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new TimeSecVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new TimeSecWriterImpl((TimeSecVector) vector); } }, TIMEMILLI(new Time(TimeUnit.MILLISECOND, 32)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new TimeMilliVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new TimeMilliWriterImpl((TimeMilliVector) vector); } }, TIMEMICRO(new Time(TimeUnit.MICROSECOND, 64)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new TimeMicroVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new TimeMicroWriterImpl((TimeMicroVector) vector); } }, TIMENANO(new Time(TimeUnit.NANOSECOND, 64)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new TimeNanoVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new TimeNanoWriterImpl((TimeNanoVector) vector); } }, // time in second from the Unix epoch, 00:00:00 on 1 January 1970, UTC. TIMESTAMPSEC(new Timestamp(org.apache.arrow.vector.types.TimeUnit.SECOND, null)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new TimeStampSecVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new TimeStampSecWriterImpl((TimeStampSecVector) vector); } }, // time in millis from the Unix epoch, 00:00:00.000 on 1 January 1970, UTC. TIMESTAMPMILLI(new Timestamp(org.apache.arrow.vector.types.TimeUnit.MILLISECOND, null)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new TimeStampMilliVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new TimeStampMilliWriterImpl((TimeStampMilliVector) vector); } }, // time in microsecond from the Unix epoch, 00:00:00.000000 on 1 January 1970, UTC. TIMESTAMPMICRO(new Timestamp(org.apache.arrow.vector.types.TimeUnit.MICROSECOND, null)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new TimeStampMicroVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new TimeStampMicroWriterImpl((TimeStampMicroVector) vector); } }, // time in nanosecond from the Unix epoch, 00:00:00.000000000 on 1 January 1970, UTC. TIMESTAMPNANO(new Timestamp(org.apache.arrow.vector.types.TimeUnit.NANOSECOND, null)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new TimeStampNanoVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new TimeStampNanoWriterImpl((TimeStampNanoVector) vector); } }, INTERVALDAY(new Interval(IntervalUnit.DAY_TIME)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new IntervalDayVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new IntervalDayWriterImpl((IntervalDayVector) vector); } }, INTERVALMONTHDAYNANO(new Interval(IntervalUnit.MONTH_DAY_NANO)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new IntervalMonthDayNanoVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new IntervalMonthDayNanoWriterImpl((IntervalMonthDayNanoVector) vector); } }, DURATION(null) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new DurationVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new DurationWriterImpl((DurationVector) vector); } }, INTERVALYEAR(new Interval(IntervalUnit.YEAR_MONTH)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new IntervalYearVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new IntervalYearWriterImpl((IntervalYearVector) vector); } }, FLOAT2(new FloatingPoint(HALF)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new Float2Vector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new Float2WriterImpl((Float2Vector) vector); } }, // 4 byte ieee 754 FLOAT4(new FloatingPoint(SINGLE)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new Float4Vector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new Float4WriterImpl((Float4Vector) vector); } }, // 8 byte ieee 754 FLOAT8(new FloatingPoint(DOUBLE)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new Float8Vector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new Float8WriterImpl((Float8Vector) vector); } }, BIT(Bool.INSTANCE) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new BitVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new BitWriterImpl((BitVector) vector); } }, VARCHAR(Utf8.INSTANCE) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new VarCharVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new VarCharWriterImpl((VarCharVector) vector); } }, VIEWVARCHAR(Utf8View.INSTANCE) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new ViewVarCharVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new ViewVarCharWriterImpl((ViewVarCharVector) vector); } }, LARGEVARCHAR(LargeUtf8.INSTANCE) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new LargeVarCharVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new LargeVarCharWriterImpl((LargeVarCharVector) vector); } }, LARGEVARBINARY(LargeBinary.INSTANCE) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new LargeVarBinaryVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new LargeVarBinaryWriterImpl((LargeVarBinaryVector) vector); } }, VARBINARY(Binary.INSTANCE) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new VarBinaryVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new VarBinaryWriterImpl((VarBinaryVector) vector); } }, VIEWVARBINARY(BinaryView.INSTANCE) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new ViewVarBinaryVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new ViewVarBinaryWriterImpl((ViewVarBinaryVector) vector); } }, DECIMAL(null) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new DecimalVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new DecimalWriterImpl((DecimalVector) vector); } }, DECIMAL256(null) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new Decimal256Vector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new Decimal256WriterImpl((Decimal256Vector) vector); } }, FIXEDSIZEBINARY(null) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new FixedSizeBinaryVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new FixedSizeBinaryWriterImpl((FixedSizeBinaryVector) vector); } }, UINT1(new Int(8, false)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new UInt1Vector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new UInt1WriterImpl((UInt1Vector) vector); } }, UINT2(new Int(16, false)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new UInt2Vector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new UInt2WriterImpl((UInt2Vector) vector); } }, UINT4(new Int(32, false)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new UInt4Vector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new UInt4WriterImpl((UInt4Vector) vector); } }, UINT8(new Int(64, false)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new UInt8Vector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new UInt8WriterImpl((UInt8Vector) vector); } }, LIST(List.INSTANCE) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new ListVector( field.getName(), allocator, field.getFieldType(), schemaChangeCallback); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new UnionListWriter((ListVector) vector); } }, LISTVIEW(ListView.INSTANCE) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new ListViewVector( field.getName(), allocator, field.getFieldType(), schemaChangeCallback); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new UnionListWriter((ListVector) vector); } }, LARGELIST(ArrowType.LargeList.INSTANCE) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new LargeListVector( field.getName(), allocator, field.getFieldType(), schemaChangeCallback); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new UnionLargeListWriter((LargeListVector) vector); } }, LARGELISTVIEW(ArrowType.LargeListView.INSTANCE) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new LargeListViewVector( field.getName(), allocator, field.getFieldType(), schemaChangeCallback); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new UnionLargeListViewWriter((LargeListViewVector) vector); } }, FIXED_SIZE_LIST(null) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new FixedSizeListVector( field.getName(), allocator, field.getFieldType(), schemaChangeCallback); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { throw new UnsupportedOperationException( "FieldWriter not implemented for FixedSizeList " + "type"); } }, UNION(new Union(Sparse, null)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { if (field.getFieldType().getDictionary() != null) { throw new UnsupportedOperationException( "Dictionary encoding not supported for complex " + "types"); } return new UnionVector( field.getName(), allocator, field.getFieldType(), schemaChangeCallback); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new UnionWriter((UnionVector) vector); } }, DENSEUNION(new Union(Dense, null)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { if (field.getFieldType().getDictionary() != null) { throw new UnsupportedOperationException( "Dictionary encoding not supported for complex " + "types"); } return new DenseUnionVector( field.getName(), allocator, field.getFieldType(), schemaChangeCallback); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new DenseUnionWriter((DenseUnionVector) vector); } }, MAP(null) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new MapVector( field.getName(), allocator, field.getFieldType(), schemaChangeCallback); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new UnionMapWriter((MapVector) vector); } }, TIMESTAMPSECTZ(null) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new TimeStampSecTZVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new TimeStampSecTZWriterImpl((TimeStampSecTZVector) vector); } }, TIMESTAMPMILLITZ(null) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new TimeStampMilliTZVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new TimeStampMilliTZWriterImpl((TimeStampMilliTZVector) vector); } }, TIMESTAMPMICROTZ(null) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new TimeStampMicroTZVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new TimeStampMicroTZWriterImpl((TimeStampMicroTZVector) vector); } }, TIMESTAMPNANOTZ(null) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new TimeStampNanoTZVector(field, allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return new TimeStampNanoTZWriterImpl((TimeStampNanoTZVector) vector); } }, EXTENSIONTYPE(null) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return ((ExtensionType) field.getFieldType().getType()) .getNewVector(field.getName(), field.getFieldType(), allocator); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { return ((ExtensionTypeVector) vector) .getUnderlyingVector() .getMinorType() .getNewFieldWriter(vector); } }, RUNENDENCODED(RunEndEncoded.INSTANCE) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { return new RunEndEncodedVector(field, allocator, schemaChangeCallback); } @Override public FieldWriter getNewFieldWriter(ValueVector vector) { throw new UnsupportedOperationException( "FieldWriter for run-end encoded vector is not implemented yet."); } }, ; private final ArrowType type; MinorType(ArrowType type) { this.type = type; } /** Returns the {@link ArrowType} equivalent of this type. */ public final ArrowType getType() { if (type == null) { throw new UnsupportedOperationException("Cannot get simple type for type " + name()); } return type; } /** Constructs a new vector for the given type. */ public final FieldVector getNewVector( String name, FieldType fieldType, BufferAllocator allocator, CallBack schemaChangeCallback) { return getNewVector(new Field(name, fieldType, null), allocator, schemaChangeCallback); } /** Constructs a new vector for the given type. */ public abstract FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback); public abstract FieldWriter getNewFieldWriter(ValueVector vector); } /** Maps the ArrowType to the java implementations MinorType. */ public static MinorType getMinorTypeForArrowType(ArrowType arrowType) { return arrowType.accept( new ArrowTypeVisitor() { @Override public MinorType visit(Null type) { return MinorType.NULL; } @Override public MinorType visit(Struct type) { return MinorType.STRUCT; } @Override public MinorType visit(List type) { return MinorType.LIST; } @Override public MinorType visit(FixedSizeList type) { return MinorType.FIXED_SIZE_LIST; } @Override public MinorType visit(Union type) { switch (type.getMode()) { case Sparse: return MinorType.UNION; case Dense: return MinorType.DENSEUNION; default: throw new IllegalArgumentException( "only Dense or Sparse unions supported: " + type); } } @Override public MinorType visit(Map type) { return MinorType.MAP; } @Override public MinorType visit(ArrowType.LargeList type) { return MinorType.LARGELIST; } @Override public MinorType visit(Int type) { switch (type.getBitWidth()) { case 8: return type.getIsSigned() ? MinorType.TINYINT : MinorType.UINT1; case 16: return type.getIsSigned() ? MinorType.SMALLINT : MinorType.UINT2; case 32: return type.getIsSigned() ? MinorType.INT : MinorType.UINT4; case 64: return type.getIsSigned() ? MinorType.BIGINT : MinorType.UINT8; default: throw new IllegalArgumentException("only 8, 16, 32, 64 supported: " + type); } } @Override public MinorType visit(FloatingPoint type) { switch (type.getPrecision()) { case HALF: return MinorType.FLOAT2; case SINGLE: return MinorType.FLOAT4; case DOUBLE: return MinorType.FLOAT8; default: throw new IllegalArgumentException("unknown precision: " + type); } } @Override public MinorType visit(Utf8 type) { return MinorType.VARCHAR; } @Override public MinorType visit(Utf8View type) { return MinorType.VIEWVARCHAR; } @Override public Types.MinorType visit(LargeUtf8 type) { return MinorType.LARGEVARCHAR; } @Override public MinorType visit(Binary type) { return MinorType.VARBINARY; } @Override public MinorType visit(BinaryView type) { return MinorType.VIEWVARBINARY; } @Override public MinorType visit(LargeBinary type) { return MinorType.LARGEVARBINARY; } @Override public MinorType visit(Bool type) { return MinorType.BIT; } @Override public MinorType visit(Decimal type) { if (type.getBitWidth() == 256) { return MinorType.DECIMAL256; } return MinorType.DECIMAL; } @Override public MinorType visit(FixedSizeBinary type) { return MinorType.FIXEDSIZEBINARY; } @Override public MinorType visit(Date type) { switch (type.getUnit()) { case DAY: return MinorType.DATEDAY; case MILLISECOND: return MinorType.DATEMILLI; default: throw new IllegalArgumentException("unknown unit: " + type); } } @Override public MinorType visit(Time type) { switch (type.getUnit()) { case SECOND: return MinorType.TIMESEC; case MILLISECOND: return MinorType.TIMEMILLI; case MICROSECOND: return MinorType.TIMEMICRO; case NANOSECOND: return MinorType.TIMENANO; default: throw new IllegalArgumentException("unknown unit: " + type); } } @Override public MinorType visit(Timestamp type) { String tz = type.getTimezone(); switch (type.getUnit()) { case SECOND: return tz == null ? MinorType.TIMESTAMPSEC : MinorType.TIMESTAMPSECTZ; case MILLISECOND: return tz == null ? MinorType.TIMESTAMPMILLI : MinorType.TIMESTAMPMILLITZ; case MICROSECOND: return tz == null ? MinorType.TIMESTAMPMICRO : MinorType.TIMESTAMPMICROTZ; case NANOSECOND: return tz == null ? MinorType.TIMESTAMPNANO : MinorType.TIMESTAMPNANOTZ; default: throw new IllegalArgumentException("unknown unit: " + type); } } @Override public MinorType visit(Interval type) { switch (type.getUnit()) { case DAY_TIME: return MinorType.INTERVALDAY; case YEAR_MONTH: return MinorType.INTERVALYEAR; case MONTH_DAY_NANO: return MinorType.INTERVALMONTHDAYNANO; default: throw new IllegalArgumentException("unknown unit: " + type); } } @Override public MinorType visit(Duration type) { return MinorType.DURATION; } @Override public MinorType visit(ListView type) { return MinorType.LISTVIEW; } @Override public MinorType visit(LargeListView type) { return MinorType.LARGELISTVIEW; } @Override public MinorType visit(ExtensionType type) { return MinorType.EXTENSIONTYPE; } @Override public MinorType visit(RunEndEncoded type) { return MinorType.RUNENDENCODED; } }); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/types/UnionMode.java000066400000000000000000000036221515451615200314070ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.arrow.vector.types; /** Different memory layouts for Union Vectors. */ public enum UnionMode { /** * Each child vector is the same length as the overall vector, and there is one 8-bit integer * buffer to indicate the index of a child vector to use at any given position. */ Sparse(org.apache.arrow.flatbuf.UnionMode.Sparse), /** * Each child vector is of variable width. The parent vector contains both a child index vector * (like in {@link #Sparse}) and in addition a slot index buffer to determine the offset into the * child vector indicated by the index vector. */ Dense(org.apache.arrow.flatbuf.UnionMode.Dense); private static final UnionMode[] valuesByFlatbufId = new UnionMode[UnionMode.values().length]; static { for (UnionMode v : UnionMode.values()) { valuesByFlatbufId[v.flatbufID] = v; } } private final short flatbufID; private UnionMode(short flatbufID) { this.flatbufID = flatbufID; } public short getFlatbufID() { return flatbufID; } public static UnionMode fromFlatbufID(short id) { return valuesByFlatbufId[id]; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/types/pojo/000077500000000000000000000000001515451615200276135ustar00rootroot00000000000000DictionaryEncoding.java000066400000000000000000000051261515451615200341570ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/types/pojo/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.arrow.vector.types.pojo; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonGetter; import com.fasterxml.jackson.annotation.JsonProperty; import java.util.Objects; import org.apache.arrow.vector.types.pojo.ArrowType.Int; /** A POJO representation of Arrow Dictionary metadata. */ public class DictionaryEncoding { private final long id; private final boolean ordered; private final Int indexType; /** * Constructs a new instance. * * @param id The ID of the dictionary to use for encoding. * @param ordered Whether the keys in values in the dictionary are ordered. * @param indexType (nullable). The integer type to use for indexing in the dictionary. Defaults * to a signed 32 bit integer. */ @JsonCreator public DictionaryEncoding( @JsonProperty("id") long id, @JsonProperty("isOrdered") boolean ordered, @JsonProperty("indexType") Int indexType) { this.id = id; this.ordered = ordered; this.indexType = indexType == null ? new Int(32, true) : indexType; } public long getId() { return id; } @JsonGetter("isOrdered") public boolean isOrdered() { return ordered; } public Int getIndexType() { return indexType; } @Override public String toString() { return "DictionaryEncoding[id=" + id + ",ordered=" + ordered + ",indexType=" + indexType + "]"; } @Override public boolean equals(Object o) { if (this == o) { return true; } else if (!(o instanceof DictionaryEncoding)) { return false; } DictionaryEncoding that = (DictionaryEncoding) o; return id == that.id && ordered == that.ordered && Objects.equals(indexType, that.indexType); } @Override public int hashCode() { return Objects.hash(id, ordered, indexType); } } ExtensionTypeRegistry.java000066400000000000000000000027211515451615200347500ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/types/pojo/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.arrow.vector.types.pojo; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; /** A registry of recognized extension types. */ public final class ExtensionTypeRegistry { private static final ConcurrentMap registry = new ConcurrentHashMap<>(); public static void register(ExtensionType type) { registry.put(type.extensionName(), type); } public static void unregister(ExtensionType type) { registry.remove(type.extensionName()); } public static ExtensionType lookup(String name) { return registry.get(name); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/types/pojo/Field.java000066400000000000000000000266641515451615200315170ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.arrow.vector.types.pojo; import static org.apache.arrow.util.Preconditions.checkNotNull; import static org.apache.arrow.vector.complex.BaseRepeatedValueVector.DATA_VECTOR_NAME; import static org.apache.arrow.vector.types.pojo.ArrowType.getTypeForField; import static org.apache.arrow.vector.types.pojo.Schema.convertMetadata; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonInclude.Include; import com.fasterxml.jackson.annotation.JsonProperty; import com.google.flatbuffers.FlatBufferBuilder; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Map.Entry; import java.util.Objects; import java.util.stream.Collectors; import org.apache.arrow.flatbuf.KeyValue; import org.apache.arrow.flatbuf.Type; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Collections2; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** A POJO abstraction for the Flatbuffer description of Vector Type. */ public class Field { private static final Logger logger = LoggerFactory.getLogger(Field.class); public static Field nullablePrimitive(String name, ArrowType.PrimitiveType type) { return nullable(name, type); } public static Field nullable(String name, ArrowType type) { return new Field(name, FieldType.nullable(type), null); } public static Field notNullable(String name, ArrowType type) { return new Field(name, FieldType.notNullable(type), null); } private final String name; private final FieldType fieldType; private final List children; private Field( String name, boolean nullable, ArrowType type, DictionaryEncoding dictionary, List children, Map metadata) { this(name, new FieldType(nullable, type, dictionary, metadata), children); } @JsonCreator private Field( @JsonProperty("name") String name, @JsonProperty("nullable") boolean nullable, @JsonProperty("type") ArrowType type, @JsonProperty("dictionary") DictionaryEncoding dictionary, @JsonProperty("children") List children, @JsonProperty("metadata") List> metadata) { this(name, new FieldType(nullable, type, dictionary, convertMetadata(metadata)), children); } /** * Constructs a new Field object. * * @param name name of the field * @param fieldType type of the field * @param children child fields, if any */ public Field(String name, FieldType fieldType, List children) { this.name = name; this.fieldType = checkNotNull(fieldType); this.children = children == null ? Collections.emptyList() : Collections2.toImmutableList(children); } /** Construct a new vector of this type using the given allocator. */ public FieldVector createVector(BufferAllocator allocator) { FieldVector vector = fieldType.createNewSingleVector(this, allocator, null); vector.initializeChildrenFromFields(children); return vector; } /** Constructs a new instance from a flatbuffer representation of the field. */ public static Field convertField(org.apache.arrow.flatbuf.Field field) { Map metadata = new HashMap<>(); for (int i = 0; i < field.customMetadataLength(); i++) { KeyValue kv = field.customMetadata(i); String key = kv.key(); String value = kv.value(); metadata.put(key == null ? "" : key, value == null ? "" : value); } metadata = Collections.unmodifiableMap(metadata); String name = field.name(); boolean nullable = field.nullable(); ArrowType type = getTypeForField(field); if (metadata.containsKey(ExtensionType.EXTENSION_METADATA_KEY_NAME)) { final String extensionName = metadata.get(ExtensionType.EXTENSION_METADATA_KEY_NAME); final String extensionMetadata = metadata.getOrDefault(ExtensionType.EXTENSION_METADATA_KEY_METADATA, ""); ExtensionType extensionType = ExtensionTypeRegistry.lookup(extensionName); if (extensionType != null) { type = extensionType.deserialize(type, extensionMetadata); } else { // Otherwise, we haven't registered the type logger.info("Unrecognized extension type: {}", extensionName); } } DictionaryEncoding dictionary = null; org.apache.arrow.flatbuf.DictionaryEncoding dictionaryFB = field.dictionary(); if (dictionaryFB != null) { ArrowType.Int indexType = null; org.apache.arrow.flatbuf.Int indexTypeFB = dictionaryFB.indexType(); if (indexTypeFB != null) { indexType = new ArrowType.Int(indexTypeFB.bitWidth(), indexTypeFB.isSigned()); } dictionary = new DictionaryEncoding(dictionaryFB.id(), dictionaryFB.isOrdered(), indexType); } List children = new ArrayList<>(); for (int i = 0; i < field.childrenLength(); i++) { Field childField = convertField(field.children(i)); childField = mutateOriginalNameIfNeeded(field, childField); children.add(childField); } children = Collections.unmodifiableList(children); return new Field(name, nullable, type, dictionary, children, metadata); } /** * Helper method to ensure backward compatibility with schemas generated prior to ARROW-1347, * ARROW-1663. * * @param field the field to check * @param originalChildField original field which name might be mutated * @return original or mutated field */ private static Field mutateOriginalNameIfNeeded( org.apache.arrow.flatbuf.Field field, Field originalChildField) { if ((field.typeType() == Type.List || field.typeType() == Type.FixedSizeList) && originalChildField.getName().equals("[DEFAULT]")) { return new Field( DATA_VECTOR_NAME, originalChildField.isNullable(), originalChildField.getType(), originalChildField.getDictionary(), originalChildField.getChildren(), originalChildField.getMetadata()); } return originalChildField; } /** * Puts this object into builder and returns the length of the serialized flatbuffer. */ public int getField(FlatBufferBuilder builder) { int nameOffset = name == null ? -1 : builder.createString(name); int typeOffset = getType().getType(builder); int dictionaryOffset = -1; DictionaryEncoding dictionary = getDictionary(); if (dictionary != null) { int dictionaryType = dictionary.getIndexType().getType(builder); org.apache.arrow.flatbuf.DictionaryEncoding.startDictionaryEncoding(builder); org.apache.arrow.flatbuf.DictionaryEncoding.addId(builder, dictionary.getId()); org.apache.arrow.flatbuf.DictionaryEncoding.addIsOrdered(builder, dictionary.isOrdered()); org.apache.arrow.flatbuf.DictionaryEncoding.addIndexType(builder, dictionaryType); dictionaryOffset = org.apache.arrow.flatbuf.DictionaryEncoding.endDictionaryEncoding(builder); } int[] childrenData = new int[children.size()]; for (int i = 0; i < children.size(); i++) { childrenData[i] = children.get(i).getField(builder); } int childrenOffset = org.apache.arrow.flatbuf.Field.createChildrenVector(builder, childrenData); int[] metadataOffsets = new int[getMetadata().size()]; Iterator> metadataIterator = getMetadata().entrySet().iterator(); for (int i = 0; i < metadataOffsets.length; i++) { Entry kv = metadataIterator.next(); int keyOffset = builder.createString(kv.getKey()); int valueOffset = builder.createString(kv.getValue()); KeyValue.startKeyValue(builder); KeyValue.addKey(builder, keyOffset); KeyValue.addValue(builder, valueOffset); metadataOffsets[i] = KeyValue.endKeyValue(builder); } int metadataOffset = org.apache.arrow.flatbuf.Field.createCustomMetadataVector(builder, metadataOffsets); org.apache.arrow.flatbuf.Field.startField(builder); if (name != null) { org.apache.arrow.flatbuf.Field.addName(builder, nameOffset); } org.apache.arrow.flatbuf.Field.addNullable(builder, isNullable()); org.apache.arrow.flatbuf.Field.addTypeType(builder, getType().getTypeID().getFlatbufID()); org.apache.arrow.flatbuf.Field.addType(builder, typeOffset); org.apache.arrow.flatbuf.Field.addChildren(builder, childrenOffset); org.apache.arrow.flatbuf.Field.addCustomMetadata(builder, metadataOffset); if (dictionary != null) { org.apache.arrow.flatbuf.Field.addDictionary(builder, dictionaryOffset); } return org.apache.arrow.flatbuf.Field.endField(builder); } public String getName() { return name; } public boolean isNullable() { return fieldType.isNullable(); } public ArrowType getType() { return fieldType.getType(); } @JsonIgnore public FieldType getFieldType() { return fieldType; } @JsonInclude(Include.NON_NULL) public DictionaryEncoding getDictionary() { return fieldType.getDictionary(); } public List getChildren() { return children; } @JsonIgnore public Map getMetadata() { return fieldType.getMetadata(); } @JsonProperty("metadata") @JsonInclude(Include.NON_EMPTY) List> getMetadataForJson() { return convertMetadata(getMetadata()); } @Override public int hashCode() { return Objects.hash(name, isNullable(), getType(), getDictionary(), getMetadata(), children); } @Override public boolean equals(Object obj) { if (!(obj instanceof Field)) { return false; } Field that = (Field) obj; return Objects.equals(this.name, that.name) && this.isNullable() == that.isNullable() && Objects.equals(this.getType(), that.getType()) && Objects.equals(this.getDictionary(), that.getDictionary()) && Objects.equals(this.getMetadata(), that.getMetadata()) && Objects.equals(this.children, that.children); } @Override public String toString() { StringBuilder sb = new StringBuilder(); if (name != null) { sb.append(name).append(": "); } sb.append(getType()); if (getDictionary() != null) { sb.append("[dictionary: ").append(getDictionary().getId()).append("]"); } if (!children.isEmpty()) { sb.append("<") .append(children.stream().map(t -> t.toString()).collect(Collectors.joining(", "))) .append(">"); } if (!isNullable()) { sb.append(" not null"); } return sb.toString(); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/types/pojo/FieldType.java000066400000000000000000000107541515451615200323520ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.arrow.vector.types.pojo; import java.util.Collections; import java.util.HashMap; import java.util.Map; import java.util.Objects; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Collections2; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; import org.apache.arrow.vector.util.CallBack; /** * POJO representation of an Arrow field type. It consists of a logical type, nullability and * whether the field (column) is dictionary encoded. */ public class FieldType { public static FieldType nullable(ArrowType type) { return new FieldType(true, type, null, null); } public static FieldType notNullable(ArrowType type) { return new FieldType(false, type, null, null); } private final boolean nullable; private final ArrowType type; private final DictionaryEncoding dictionary; private final Map metadata; public FieldType(boolean nullable, ArrowType type, DictionaryEncoding dictionary) { this(nullable, type, dictionary, null); } /** * Constructs a new instance. * * @param nullable Whether the Vector is nullable * @param type The logical arrow type of the field. * @param dictionary The dictionary encoding of the field. * @param metadata Custom metadata for the field. */ public FieldType( boolean nullable, ArrowType type, DictionaryEncoding dictionary, Map metadata) { super(); this.nullable = nullable; this.type = Preconditions.checkNotNull(type); this.dictionary = dictionary; if (type instanceof ExtensionType) { // Save the extension type name/metadata final Map extensionMetadata = new HashMap<>(); extensionMetadata.put( ExtensionType.EXTENSION_METADATA_KEY_NAME, ((ExtensionType) type).extensionName()); extensionMetadata.put( ExtensionType.EXTENSION_METADATA_KEY_METADATA, ((ExtensionType) type).serialize()); if (metadata != null) { extensionMetadata.putAll(metadata); } this.metadata = Collections.unmodifiableMap(extensionMetadata); } else { this.metadata = metadata == null ? java.util.Collections.emptyMap() : Collections2.immutableMapCopy(metadata); } } public boolean isNullable() { return nullable; } public ArrowType getType() { return type; } public DictionaryEncoding getDictionary() { return dictionary; } public Map getMetadata() { return metadata; } public FieldVector createNewSingleVector( String name, BufferAllocator allocator, CallBack schemaCallBack) { MinorType minorType = Types.getMinorTypeForArrowType(type); return minorType.getNewVector(name, this, allocator, schemaCallBack); } public FieldVector createNewSingleVector( Field field, BufferAllocator allocator, CallBack schemaCallBack) { MinorType minorType = Types.getMinorTypeForArrowType(type); return minorType.getNewVector(field, allocator, schemaCallBack); } @Override public int hashCode() { return Objects.hash(nullable, type, dictionary, metadata); } @Override public boolean equals(Object obj) { if (!(obj instanceof FieldType)) { return false; } FieldType that = (FieldType) obj; return this.isNullable() == that.isNullable() && Objects.equals(this.getType(), that.getType()) && Objects.equals(this.getDictionary(), that.getDictionary()) && Objects.equals(this.getMetadata(), that.getMetadata()); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/types/pojo/Schema.java000066400000000000000000000245341515451615200316660ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.arrow.vector.types.pojo; import static org.apache.arrow.vector.types.pojo.Field.convertField; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonInclude.Include; import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.ObjectReader; import com.fasterxml.jackson.databind.ObjectWriter; import com.fasterxml.jackson.databind.util.ByteBufferBackedInputStream; import com.google.flatbuffers.FlatBufferBuilder; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.nio.channels.Channels; import java.util.AbstractMap; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.stream.Collectors; import org.apache.arrow.flatbuf.Endianness; import org.apache.arrow.flatbuf.KeyValue; import org.apache.arrow.util.Collections2; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.ipc.ReadChannel; import org.apache.arrow.vector.ipc.WriteChannel; import org.apache.arrow.vector.ipc.message.FBSerializables; import org.apache.arrow.vector.ipc.message.MessageSerializer; /** An Arrow Schema. */ public class Schema { /** * Search for a field by name in given the list of fields. * * @param fields the list of the fields * @param name the name of the field to return * @return the corresponding field * @throws IllegalArgumentException if the field was not found */ public static Field findField(List fields, String name) { for (Field field : fields) { if (field.getName().equals(name)) { return field; } } throw new IllegalArgumentException(String.format("field %s not found in %s", name, fields)); } static final String METADATA_KEY = "key"; static final String METADATA_VALUE = "value"; private static final ObjectMapper mapper = new ObjectMapper(); private static final ObjectWriter writer = mapper.writerWithDefaultPrettyPrinter(); private static final ObjectReader reader = mapper.readerFor(Schema.class); private static final boolean LITTLE_ENDIAN = ByteOrder.nativeOrder() == ByteOrder.LITTLE_ENDIAN; public static Schema fromJSON(String json) throws IOException { return reader.readValue(Preconditions.checkNotNull(json)); } /** * Deserialize a schema that has been serialized using {@link #toByteArray()}. * * @param buffer the bytes to deserialize. * @return The deserialized schema. */ @Deprecated public static Schema deserialize(ByteBuffer buffer) { return convertSchema(org.apache.arrow.flatbuf.Schema.getRootAsSchema(buffer)); } /** * Deserialize a schema that has been serialized as a message using {@link #serializeAsMessage()}. * * @param buffer the bytes to deserialize. * @return The deserialized schema. */ public static Schema deserializeMessage(ByteBuffer buffer) { ByteBufferBackedInputStream stream = new ByteBufferBackedInputStream(buffer); try (ReadChannel channel = new ReadChannel(Channels.newChannel(stream))) { return MessageSerializer.deserializeSchema(channel); } catch (IOException ex) { throw new RuntimeException(ex); } } /** Converts a flatbuffer schema to its POJO representation. */ public static Schema convertSchema(org.apache.arrow.flatbuf.Schema schema) { List fields = new ArrayList<>(); for (int i = 0; i < schema.fieldsLength(); i++) { fields.add(convertField(schema.fields(i))); } Map metadata = new HashMap<>(); for (int i = 0; i < schema.customMetadataLength(); i++) { KeyValue kv = schema.customMetadata(i); String key = kv.key(); String value = kv.value(); metadata.put(key == null ? "" : key, value == null ? "" : value); } return new Schema( true, Collections.unmodifiableList(fields), Collections.unmodifiableMap(metadata)); } private final List fields; private final Map metadata; public Schema(Iterable fields) { this(fields, (Map) null); } /** Constructor with metadata. */ public Schema(Iterable fields, Map metadata) { this( true, Collections2.toImmutableList(fields), metadata == null ? Collections.emptyMap() : Collections2.immutableMapCopy(metadata)); } /** Constructor used for JSON deserialization. */ @JsonCreator private Schema( @JsonProperty("fields") Iterable fields, @JsonProperty("metadata") List> metadata) { this(fields, convertMetadata(metadata)); } /** * Private constructor to bypass automatic collection copy. * * @param ignored an ignored argument. Its only purpose is to prevent using the constructor by * accident because of type collisions (List vs Iterable). */ private Schema(boolean ignored, List fields, Map metadata) { this.fields = fields; this.metadata = metadata; } static Map convertMetadata(List> metadata) { return (metadata == null) ? null : metadata.stream() .map( e -> new AbstractMap.SimpleImmutableEntry<>( e.get(METADATA_KEY), e.get(METADATA_VALUE))) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); } static List> convertMetadata(Map metadata) { return (metadata == null) ? null : metadata.entrySet().stream() .map(Schema::convertEntryToKeyValueMap) .collect(Collectors.toList()); } private static Map convertEntryToKeyValueMap(Map.Entry entry) { Map map = new HashMap<>(2); map.put(METADATA_KEY, entry.getKey()); map.put(METADATA_VALUE, entry.getValue()); return Collections.unmodifiableMap(map); } public List getFields() { return fields; } @JsonIgnore public Map getCustomMetadata() { return metadata; } @JsonProperty("metadata") @JsonInclude(Include.NON_EMPTY) List> getCustomMetadataForJson() { return convertMetadata(getCustomMetadata()); } /** * Search for a field by name in this Schema. * * @param name the name of the field to return * @return the corresponding field * @throws IllegalArgumentException if the field was not found */ public Field findField(String name) { return findField(getFields(), name); } /** Returns the JSON string representation of this schema. */ public String toJson() { try { return writer.writeValueAsString(this); } catch (JsonProcessingException e) { // this should not happen throw new RuntimeException(e); } } /** Adds this schema to the builder returning the size of the builder after adding. */ public int getSchema(FlatBufferBuilder builder) { int[] fieldOffsets = new int[fields.size()]; for (int i = 0; i < fields.size(); i++) { fieldOffsets[i] = fields.get(i).getField(builder); } int fieldsOffset = org.apache.arrow.flatbuf.Schema.createFieldsVector(builder, fieldOffsets); int metadataOffset = FBSerializables.writeKeyValues(builder, metadata); org.apache.arrow.flatbuf.Schema.startSchema(builder); org.apache.arrow.flatbuf.Schema.addEndianness( builder, (LITTLE_ENDIAN ? Endianness.Little : Endianness.Big)); org.apache.arrow.flatbuf.Schema.addFields(builder, fieldsOffset); org.apache.arrow.flatbuf.Schema.addCustomMetadata(builder, metadataOffset); return org.apache.arrow.flatbuf.Schema.endSchema(builder); } /** * Returns the serialized flatbuffer bytes of the schema wrapped in a message table. Use {@link * #deserializeMessage(ByteBuffer)} to rebuild the Schema. */ public byte[] serializeAsMessage() { ByteArrayOutputStream out = new ByteArrayOutputStream(); try (WriteChannel channel = new WriteChannel(Channels.newChannel(out))) { MessageSerializer.serialize(channel, this); return out.toByteArray(); } catch (IOException ex) { throw new RuntimeException(ex); } } /** * Returns the serialized flatbuffer representation of this schema. * * @deprecated This method does not encapsulate the schema in a Message payload which is * incompatible with other languages. Use {@link #serializeAsMessage()} instead. */ @Deprecated public byte[] toByteArray() { FlatBufferBuilder builder = new FlatBufferBuilder(); int schemaOffset = this.getSchema(builder); builder.finish(schemaOffset); ByteBuffer bb = builder.dataBuffer(); byte[] bytes = new byte[bb.remaining()]; bb.get(bytes); return bytes; } @Override public int hashCode() { return Objects.hash(fields, metadata); } @Override public boolean equals(Object obj) { if (!(obj instanceof Schema)) { return false; } return Objects.equals(this.fields, ((Schema) obj).fields) && Objects.equals(this.metadata, ((Schema) obj).metadata); } @Override public String toString() { String meta = metadata.isEmpty() ? "" : "(metadata: " + metadata.toString() + ")"; return "Schema<" + fields.stream().map(t -> t.toString()).collect(Collectors.joining(", ")) + ">" + meta; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/000077500000000000000000000000001515451615200264555ustar00rootroot00000000000000ByteArrayReadableSeekableByteChannel.java000066400000000000000000000046611515451615200363430ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.io.IOException; import java.nio.ByteBuffer; import java.nio.channels.SeekableByteChannel; /** A {@link SeekableByteChannel} backed by a byte array. */ public class ByteArrayReadableSeekableByteChannel implements SeekableByteChannel { private byte[] byteArray; private int position = 0; /** Construct a new object using the given byteArray as a backing store. */ public ByteArrayReadableSeekableByteChannel(byte[] byteArray) { if (byteArray == null) { throw new NullPointerException(); } this.byteArray = byteArray; } @Override public boolean isOpen() { return byteArray != null; } @Override public void close() throws IOException { byteArray = null; } @Override public int read(final ByteBuffer dst) throws IOException { int remainingInBuf = byteArray.length - this.position; int length = Math.min(dst.remaining(), remainingInBuf); dst.put(this.byteArray, this.position, length); this.position += length; return length; } @Override public long position() throws IOException { return this.position; } @Override public SeekableByteChannel position(final long newPosition) throws IOException { this.position = (int) newPosition; return this; } @Override public long size() throws IOException { return this.byteArray.length; } @Override public int write(final ByteBuffer src) throws IOException { throw new UnsupportedOperationException("Read only"); } @Override public SeekableByteChannel truncate(final long size) throws IOException { throw new UnsupportedOperationException("Read only"); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/CallBack.java000066400000000000000000000017041515451615200307560ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** Generic callback interface to be notified of events on value vectors. */ public interface CallBack { void doWork(); } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/DataSizeRoundingUtil.java000066400000000000000000000056011515451615200333720ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** Utilities for rounding data size. */ public final class DataSizeRoundingUtil { /** The mask for rounding an integer to a multiple of 8. (i.e. clear the lowest 3 bits) */ public static int ROUND_8_MASK_INT = 0xFFFFFFF8; /** The mask for rounding a long integer to a multiple of 8. (i.e. clear the lowest 3 bits) */ public static long ROUND_8_MASK_LONG = 0xFFFFFFFFFFFFFFF8L; /** The number of bits to shift for dividing by 8. */ public static int DIVIDE_BY_8_SHIFT_BITS = 3; /** * Round up the number to the nearest multiple of 8. * * @param input the number to round. * @return the rounded number. */ public static int roundUpTo8Multiple(int input) { return (input + 7) & ROUND_8_MASK_INT; } /** * Round up the number to the nearest multiple of 8. * * @param input the number to round. * @return the rounded number */ public static long roundUpTo8Multiple(long input) { return (input + 7L) & ROUND_8_MASK_LONG; } /** * Round down the number to the nearest multiple of 8. * * @param input the number to round. * @return the rounded number. */ public static int roundDownTo8Multiple(int input) { return input & ROUND_8_MASK_INT; } /** * Round down the number to the nearest multiple of 8. * * @param input the number to round. * @return the rounded number */ public static long roundDownTo8Multiple(long input) { return input & ROUND_8_MASK_LONG; } /** * A fast way to compute Math.ceil(input / 8.0). * * @param input the input number. * @return the computed number. */ public static int divideBy8Ceil(int input) { return (input + 7) >>> DIVIDE_BY_8_SHIFT_BITS; } /** * A fast way to compute Math.ceil(input / 8.0). * * @param input the input number. * @return the computed number. */ public static long divideBy8Ceil(long input) { return (input + 7) >>> (long) DIVIDE_BY_8_SHIFT_BITS; } public static long roundUpToMultipleOf16(long num) { return (num + 15) & 0xFFFFFFFFFFFFFFF0L; } private DataSizeRoundingUtil() {} } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/DateUtility.java000066400000000000000000000133731515451615200315700ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.time.Instant; import java.time.LocalDateTime; import java.time.format.DateTimeFormatter; import java.time.format.DateTimeFormatterBuilder; import java.time.temporal.ChronoUnit; import java.util.TimeZone; /** Utility class for Date, DateTime, TimeStamp, Interval data types. */ public class DateUtility { private DateUtility() {} private static final String UTC = "UTC"; public static final DateTimeFormatter formatDate = DateTimeFormatter.ofPattern("yyyy-MM-dd"); public static final DateTimeFormatter formatTimeStampMilli = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSS"); public static final DateTimeFormatter formatTimeStampTZ = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSS ZZZ"); public static final DateTimeFormatter formatTime = DateTimeFormatter.ofPattern("HH:mm:ss.SSS"); public static DateTimeFormatter dateTimeTZFormat = null; public static DateTimeFormatter timeFormat = null; public static final int yearsToMonths = 12; public static final int hoursToMillis = 60 * 60 * 1000; public static final int minutesToMillis = 60 * 1000; public static final int secondsToMillis = 1000; public static final int monthToStandardDays = 30; public static final long monthsToMillis = 2592000000L; // 30 * 24 * 60 * 60 * 1000 public static final int daysToStandardMillis = 24 * 60 * 60 * 1000; /** Returns the date time formatter used to parse date strings. */ public static DateTimeFormatter getDateTimeFormatter() { if (dateTimeTZFormat == null) { DateTimeFormatter dateFormatter = DateTimeFormatter.ofPattern("yyyy-MM-dd"); DateTimeFormatter optionalTime = DateTimeFormatter.ofPattern(" HH:mm:ss"); DateTimeFormatter optionalSec = DateTimeFormatter.ofPattern(".SSS"); DateTimeFormatter optionalZone = DateTimeFormatter.ofPattern(" ZZZ"); dateTimeTZFormat = new DateTimeFormatterBuilder() .append(dateFormatter) .appendOptional(optionalTime) .appendOptional(optionalSec) .appendOptional(optionalZone) .toFormatter(); } return dateTimeTZFormat; } /** Returns time formatter used to parse time strings. */ public static DateTimeFormatter getTimeFormatter() { if (timeFormat == null) { DateTimeFormatter timeFormatter = DateTimeFormatter.ofPattern("HH:mm:ss"); DateTimeFormatter optionalSec = DateTimeFormatter.ofPattern(".SSS"); timeFormat = new DateTimeFormatterBuilder() .append(timeFormatter) .appendOptional(optionalSec) .toFormatter(); } return timeFormat; } /** * Convert milliseconds from epoch to a LocalDateTime with timeZone offset. * * @param epochMillis milliseconds from epoch * @param timeZone current timeZone * @return LocalDateTime object with timeZone offset */ public static LocalDateTime getLocalDateTimeFromEpochMilli(long epochMillis, String timeZone) { final LocalDateTime localDateTime = LocalDateTime.ofInstant( Instant.ofEpochMilli(epochMillis), TimeZone.getTimeZone(timeZone).toZoneId()); return localDateTime; } /** Convert milliseconds from epoch to a LocalDateTime with UTC offset. */ public static LocalDateTime getLocalDateTimeFromEpochMilli(long epochMillis) { return getLocalDateTimeFromEpochMilli(epochMillis, UTC); } /** * Convert microseconds from epoch to a LocalDateTime with timeZone offset. * * @param epochMicros microseconds from epoch * @param timeZone current timeZone * @return LocalDateTime object with timeZone offset */ public static LocalDateTime getLocalDateTimeFromEpochMicro(long epochMicros, String timeZone) { final long millis = java.util.concurrent.TimeUnit.MICROSECONDS.toMillis(epochMicros); final long addl_micros = epochMicros - (millis * 1000); return DateUtility.getLocalDateTimeFromEpochMilli(millis, timeZone) .plus(addl_micros, ChronoUnit.MICROS); } /** Convert microseconds from epoch to a LocalDateTime with UTC offset. */ public static LocalDateTime getLocalDateTimeFromEpochMicro(long epochMicros) { return getLocalDateTimeFromEpochMicro(epochMicros, UTC); } /** * Convert nanoseconds from epoch to a LocalDateTime with timeZone offset. * * @param epochNanos nanoseconds from epoch * @param timeZone current timeZone * @return LocalDateTime object with timeZone offset */ public static LocalDateTime getLocalDateTimeFromEpochNano(long epochNanos, String timeZone) { final long millis = java.util.concurrent.TimeUnit.NANOSECONDS.toMillis(epochNanos); final long addl_nanos = epochNanos - (millis * 1000 * 1000); return DateUtility.getLocalDateTimeFromEpochMilli(millis, timeZone).plusNanos(addl_nanos); } /** Convert nanoseconds from epoch to a LocalDateTime with UTC offset. */ public static LocalDateTime getLocalDateTimeFromEpochNano(long epochNanos) { return getLocalDateTimeFromEpochNano(epochNanos, UTC); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/DecimalUtility.java000066400000000000000000000206111515451615200322420ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.math.BigDecimal; import java.math.BigInteger; import java.nio.ByteBuffer; import java.nio.ByteOrder; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.util.MemoryUtil; /** Utility methods for configurable precision Decimal values (e.g. {@link BigDecimal}). */ public class DecimalUtility { private DecimalUtility() {} public static final byte[] zeroes = new byte[] { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 }; public static final byte[] minus_one = new byte[] { -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1 }; private static final boolean LITTLE_ENDIAN = ByteOrder.nativeOrder() == ByteOrder.LITTLE_ENDIAN; /** * Read an ArrowType.Decimal at the given value index in the ArrowBuf and convert to a BigDecimal * with the given scale. */ public static BigDecimal getBigDecimalFromArrowBuf( ArrowBuf bytebuf, int index, int scale, int byteWidth) { byte[] value = new byte[byteWidth]; byte temp; final long startIndex = (long) index * byteWidth; bytebuf.getBytes(startIndex, value, 0, byteWidth); if (LITTLE_ENDIAN) { // Decimal stored as native endian, need to swap bytes to make BigDecimal if native endian is // LE int stop = byteWidth / 2; for (int i = 0, j; i < stop; i++) { temp = value[i]; j = (byteWidth - 1) - i; value[i] = value[j]; value[j] = temp; } } BigInteger unscaledValue = new BigInteger(value); return new BigDecimal(unscaledValue, scale); } /** * Read an ArrowType.Decimal from the ByteBuffer and convert to a BigDecimal with the given scale. */ public static BigDecimal getBigDecimalFromByteBuffer( ByteBuffer bytebuf, int scale, int byteWidth) { byte[] value = new byte[byteWidth]; bytebuf.get(value); BigInteger unscaledValue = new BigInteger(value); return new BigDecimal(unscaledValue, scale); } /** * Read an ArrowType.Decimal from the ArrowBuf at the given value index and return it as a byte * array. */ public static byte[] getByteArrayFromArrowBuf(ArrowBuf bytebuf, int index, int byteWidth) { final byte[] value = new byte[byteWidth]; final long startIndex = (long) index * byteWidth; bytebuf.getBytes(startIndex, value, 0, byteWidth); return value; } /** * Check that the BigDecimal scale equals the vectorScale and that the BigDecimal precision is * less than or equal to the vectorPrecision. If not, then an UnsupportedOperationException is * thrown, otherwise returns true. */ public static boolean checkPrecisionAndScale( BigDecimal value, int vectorPrecision, int vectorScale) { if (value.scale() != vectorScale) { throw new UnsupportedOperationException( "BigDecimal scale must equal that in the Arrow vector: " + value.scale() + " != " + vectorScale); } if (value.precision() > vectorPrecision) { throw new UnsupportedOperationException( "BigDecimal precision cannot be greater than that in the Arrow " + "vector: " + value.precision() + " > " + vectorPrecision); } return true; } /** * Check that the BigDecimal scale equals the vectorScale and that the BigDecimal precision is * less than or equal to the vectorPrecision. Return true if so, otherwise return false. */ public static boolean checkPrecisionAndScaleNoThrow( BigDecimal value, int vectorPrecision, int vectorScale) { return value.scale() == vectorScale && value.precision() < vectorPrecision; } /** * Check that the decimal scale equals the vectorScale and that the decimal precision is less than * or equal to the vectorPrecision. If not, then an UnsupportedOperationException is thrown, * otherwise returns true. */ public static boolean checkPrecisionAndScale( int decimalPrecision, int decimalScale, int vectorPrecision, int vectorScale) { if (decimalScale != vectorScale) { throw new UnsupportedOperationException( "BigDecimal scale must equal that in the Arrow vector: " + decimalScale + " != " + vectorScale); } if (decimalPrecision > vectorPrecision) { throw new UnsupportedOperationException( "BigDecimal precision cannot be greater than that in the Arrow " + "vector: " + decimalPrecision + " > " + vectorPrecision); } return true; } /** * Write the given BigDecimal to the ArrowBuf at the given value index. Will throw an * UnsupportedOperationException if the decimal size is greater than the Decimal vector byte * width. */ public static void writeBigDecimalToArrowBuf( BigDecimal value, ArrowBuf bytebuf, int index, int byteWidth) { final byte[] bytes = value.unscaledValue().toByteArray(); writeByteArrayToArrowBufHelper(bytes, bytebuf, index, byteWidth); } /** * Write the given long to the ArrowBuf at the given value index. This routine extends the * original sign bit to a new upper area in 128-bit or 256-bit. */ public static void writeLongToArrowBuf(long value, ArrowBuf bytebuf, int index, int byteWidth) { if (byteWidth != 16 && byteWidth != 32) { throw new UnsupportedOperationException( "DecimalUtility.writeLongToArrowBuf() currently supports " + "128-bit or 256-bit width data"); } final long addressOfValue = bytebuf.memoryAddress() + (long) index * byteWidth; final long padValue = Long.signum(value) == -1 ? -1L : 0L; if (LITTLE_ENDIAN) { MemoryUtil.putLong(addressOfValue, value); for (int i = 1; i <= (byteWidth - 8) / 8; i++) { MemoryUtil.putLong(addressOfValue + Long.BYTES * i, padValue); } } else { for (int i = 0; i < (byteWidth - 8) / 8; i++) { MemoryUtil.putLong(addressOfValue + Long.BYTES * i, padValue); } MemoryUtil.putLong(addressOfValue + Long.BYTES * (byteWidth - 8) / 8, value); } } /** * Write the given byte array to the ArrowBuf at the given value index. Will throw an * UnsupportedOperationException if the decimal size is greater than the Decimal vector byte * width. */ public static void writeByteArrayToArrowBuf( byte[] bytes, ArrowBuf bytebuf, int index, int byteWidth) { writeByteArrayToArrowBufHelper(bytes, bytebuf, index, byteWidth); } private static void writeByteArrayToArrowBufHelper( byte[] bytes, ArrowBuf bytebuf, int index, int byteWidth) { final long startIndex = (long) index * byteWidth; if (bytes.length > byteWidth) { throw new UnsupportedOperationException( "Decimal size greater than " + byteWidth + " bytes: " + bytes.length); } byte[] padBytes = bytes[0] < 0 ? minus_one : zeroes; if (LITTLE_ENDIAN) { // Decimal stored as native-endian, need to swap data bytes before writing to ArrowBuf if LE byte[] bytesLE = new byte[bytes.length]; for (int i = 0; i < bytes.length; i++) { bytesLE[i] = bytes[bytes.length - 1 - i]; } // Write LE data bytebuf.setBytes(startIndex, bytesLE, 0, bytes.length); bytebuf.setBytes(startIndex + bytes.length, padBytes, 0, byteWidth - bytes.length); } else { // Write BE data bytebuf.setBytes(startIndex + byteWidth - bytes.length, bytes, 0, bytes.length); bytebuf.setBytes(startIndex, padBytes, 0, byteWidth - bytes.length); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/DictionaryUtility.java000066400000000000000000000122501515451615200330110ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.Set; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; 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; /** Utility methods for working with Dictionaries used in Dictionary encodings. */ public class DictionaryUtility { private DictionaryUtility() {} /** * Convert field and child fields that have a dictionary encoding to message format, so fields * have the dictionary type. * *

    NOTE: in the message format, fields have the dictionary type in the memory format, they have * the index type */ public static Field toMessageFormat( Field field, DictionaryProvider provider, Set dictionaryIdsUsed) { if (!needConvertToMessageFormat(field)) { return field; } DictionaryEncoding encoding = field.getDictionary(); List children; ArrowType type; if (encoding == null) { type = field.getType(); children = field.getChildren(); } else { long id = encoding.getId(); Dictionary dictionary = provider.lookup(id); if (dictionary == null) { throw new IllegalArgumentException("Could not find dictionary with ID " + id); } type = dictionary.getVectorType(); children = dictionary.getVector().getField().getChildren(); dictionaryIdsUsed.add(id); } final List updatedChildren = new ArrayList<>(children.size()); for (Field child : children) { updatedChildren.add(toMessageFormat(child, provider, dictionaryIdsUsed)); } return new Field( field.getName(), new FieldType(field.isNullable(), type, encoding, field.getMetadata()), updatedChildren); } /** * Checks if it is required to convert the field to message format. * * @param field the field to check. * @return true if a conversion is required, and false otherwise. */ public static boolean needConvertToMessageFormat(Field field) { DictionaryEncoding encoding = field.getDictionary(); if (encoding != null) { // when encoding is not null, the type must be determined from the // dictionary, so conversion must be performed. return true; } List children = field.getChildren(); for (Field child : children) { if (needConvertToMessageFormat(child)) { return true; } } return false; } /** * Convert field and child fields that have a dictionary encoding to memory format, so fields have * the index type. */ public static Field toMemoryFormat( Field field, BufferAllocator allocator, Map dictionaries) { DictionaryEncoding encoding = field.getDictionary(); List children = field.getChildren(); if (encoding == null && children.isEmpty()) { return field; } List updatedChildren = new ArrayList<>(children.size()); for (Field child : children) { updatedChildren.add(toMemoryFormat(child, allocator, dictionaries)); } ArrowType type; List fieldChildren = null; if (encoding == null) { type = field.getType(); fieldChildren = updatedChildren; } else { // re-type the field for in-memory format type = encoding.getIndexType(); if (type == null) { type = new ArrowType.Int(32, true); } // get existing or create dictionary vector if (!dictionaries.containsKey(encoding.getId())) { // create a new dictionary vector for the values String dictName = "DICT" + encoding.getId(); Field dictionaryField = new Field( dictName, new FieldType(field.isNullable(), field.getType(), null, null), updatedChildren); FieldVector dictionaryVector = dictionaryField.createVector(allocator); dictionaries.put(encoding.getId(), new Dictionary(dictionaryVector, encoding)); } } return new Field( field.getName(), new FieldType(field.isNullable(), type, encoding, field.getMetadata()), fieldChildren); } } ElementAddressableVectorIterator.java000066400000000000000000000052711515451615200356660ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.Iterator; import org.apache.arrow.memory.util.ArrowBufPointer; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.memory.util.hash.SimpleHasher; import org.apache.arrow.vector.ElementAddressableVector; /** * Iterator for traversing elements of a {@link ElementAddressableVector}. * * @param vector type. */ public class ElementAddressableVectorIterator implements Iterator { private final T vector; /** Index of the next element to access. */ private int index = 0; private final ArrowBufPointer reusablePointer; /** * Constructs an iterator for the {@link ElementAddressableVector}. * * @param vector the vector to iterate. */ public ElementAddressableVectorIterator(T vector) { this(vector, SimpleHasher.INSTANCE); } /** * Constructs an iterator for the {@link ElementAddressableVector}. * * @param vector the vector to iterate. * @param hasher the hasher to calculate the hash code. */ public ElementAddressableVectorIterator(T vector, ArrowBufHasher hasher) { this.vector = vector; reusablePointer = new ArrowBufPointer(hasher); } @Override public boolean hasNext() { return index < vector.getValueCount(); } /** * Retrieves the next pointer from the vector. * * @return the pointer pointing to the next element in the vector. Note that the returned pointer * is only valid before the next call to this method. */ @Override public ArrowBufPointer next() { vector.getDataPointer(index, reusablePointer); index += 1; return reusablePointer; } /** * Retrieves the next pointer from the vector. * * @param outPointer the pointer to populate. */ public void next(ArrowBufPointer outPointer) { vector.getDataPointer(index, outPointer); index += 1; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/IntObjectHashMap.java000066400000000000000000000470361515451615200324550ustar00rootroot00000000000000/* * 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. */ package org.apache.arrow.vector.util; import java.util.AbstractCollection; import java.util.AbstractSet; import java.util.Arrays; import java.util.Collection; import java.util.Iterator; import java.util.Map; import java.util.NoSuchElementException; import java.util.Set; /** * A vendored specialized copy of Netty's IntObjectHashMap for use within Arrow. Avoids requiring * Netty in the Arrow core just for this one class. * * @param The value type stored in the map. */ class IntObjectHashMap implements IntObjectMap { /** Default initial capacity. Used if not specified in the constructor */ public static final int DEFAULT_CAPACITY = 8; /** Default load factor. Used if not specified in the constructor */ public static final float DEFAULT_LOAD_FACTOR = 0.5f; /** * Placeholder for null values, so we can use the actual null to mean available. (Better than * using a placeholder for available: less references for GC processing.) */ private static final Object NULL_VALUE = new Object(); /** The maximum number of elements allowed without allocating more space. */ private int maxSize; /** The load factor for the map. Used to calculate {@link #maxSize}. */ private final float loadFactor; private int[] keys; private V[] values; private int size; private int mask; private final Set keySet = new KeySet(); private final Set> entrySet = new EntrySet(); private final Iterable> entries = new Iterable>() { @Override public Iterator> iterator() { return new PrimitiveIterator(); } }; public IntObjectHashMap() { this(DEFAULT_CAPACITY, DEFAULT_LOAD_FACTOR); } public IntObjectHashMap(int initialCapacity) { this(initialCapacity, DEFAULT_LOAD_FACTOR); } public IntObjectHashMap(int initialCapacity, float loadFactor) { if (loadFactor <= 0.0f || loadFactor > 1.0f) { // Cannot exceed 1 because we can never store more than capacity elements; // using a bigger loadFactor would trigger rehashing before the desired load is reached. throw new IllegalArgumentException("loadFactor must be > 0 and <= 1"); } this.loadFactor = loadFactor; // Adjust the initial capacity if necessary. int capacity = safeFindNextPositivePowerOfTwo(initialCapacity); mask = capacity - 1; // Allocate the arrays. keys = new int[capacity]; @SuppressWarnings({"unchecked", "SuspiciousArrayCast"}) V[] temp = (V[]) new Object[capacity]; values = temp; // Initialize the maximum size value. maxSize = calcMaxSize(capacity); } private static T toExternal(T value) { assert value != null : "null is not a legitimate internal value. Concurrent Modification?"; return value == NULL_VALUE ? null : value; } @SuppressWarnings("unchecked") private static T toInternal(T value) { return value == null ? (T) NULL_VALUE : value; } @Override public V get(int key) { int index = indexOf(key); return index == -1 ? null : toExternal(values[index]); } @Override public V put(int key, V value) { int startIndex = hashIndex(key); int index = startIndex; for (; ; ) { if (values[index] == null) { // Found empty slot, use it. keys[index] = key; values[index] = toInternal(value); growSize(); return null; } if (keys[index] == key) { // Found existing entry with this key, just replace the value. V previousValue = values[index]; values[index] = toInternal(value); return toExternal(previousValue); } // Conflict, keep probing ... if ((index = probeNext(index)) == startIndex) { // Can only happen if the map was full at MAX_ARRAY_SIZE and couldn't grow. throw new IllegalStateException("Unable to insert"); } } } @Override public void putAll(Map sourceMap) { if (sourceMap instanceof IntObjectHashMap) { // Optimization - iterate through the arrays. @SuppressWarnings("unchecked") IntObjectHashMap source = (IntObjectHashMap) sourceMap; for (int i = 0; i < source.values.length; ++i) { V sourceValue = source.values[i]; if (sourceValue != null) { put(source.keys[i], sourceValue); } } return; } // Otherwise, just add each entry. for (Entry entry : sourceMap.entrySet()) { put(entry.getKey(), entry.getValue()); } } @Override public V remove(int key) { int index = indexOf(key); if (index == -1) { return null; } V prev = values[index]; removeAt(index); return toExternal(prev); } @Override public int size() { return size; } @Override public boolean isEmpty() { return size == 0; } @Override public void clear() { Arrays.fill(keys, (int) 0); Arrays.fill(values, null); size = 0; } @Override public boolean containsKey(int key) { return indexOf(key) >= 0; } @Override public boolean containsValue(Object value) { @SuppressWarnings("unchecked") V v1 = toInternal((V) value); for (V v2 : values) { // The map supports null values; this will be matched as NULL_VALUE.equals(NULL_VALUE). if (v2 != null && v2.equals(v1)) { return true; } } return false; } @Override public Iterable> entries() { return entries; } @Override public Collection values() { return new AbstractCollection() { @Override public Iterator iterator() { return new Iterator() { final PrimitiveIterator iter = new PrimitiveIterator(); @Override public boolean hasNext() { return iter.hasNext(); } @Override public V next() { return iter.next().value(); } @Override public void remove() { iter.remove(); } }; } @Override public int size() { return size; } }; } @Override public int hashCode() { // Hashcode is based on all non-zero, valid keys. We have to scan the whole keys // array, which may have different lengths for two maps of same size(), so the // capacity cannot be used as input for hashing but the size can. int hash = size; for (int key : keys) { // 0 can be a valid key or unused slot, but won't impact the hashcode in either case. // This way we can use a cheap loop without conditionals, or hard-to-unroll operations, // or the devastatingly bad memory locality of visiting value objects. // Also, it's important to use a hash function that does not depend on the ordering // of terms, only their values; since the map is an unordered collection and // entries can end up in different positions in different maps that have the same // elements, but with different history of puts/removes, due to conflicts. hash ^= hashCode(key); } return hash; } @Override public boolean equals(Object obj) { if (this == obj) { return true; } if (!(obj instanceof IntObjectMap)) { return false; } @SuppressWarnings("rawtypes") IntObjectMap other = (IntObjectMap) obj; if (size != other.size()) { return false; } for (int i = 0; i < values.length; ++i) { V value = values[i]; if (value != null) { int key = keys[i]; Object otherValue = other.get(key); if (value == NULL_VALUE) { if (otherValue != null) { return false; } } else if (!value.equals(otherValue)) { return false; } } } return true; } @Override public boolean containsKey(Object key) { return containsKey(objectToKey(key)); } @Override public V get(Object key) { return get(objectToKey(key)); } @Override public V put(Integer key, V value) { return put(objectToKey(key), value); } @Override public V remove(Object key) { return remove(objectToKey(key)); } @Override public Set keySet() { return keySet; } @Override public Set> entrySet() { return entrySet; } private int objectToKey(Object key) { return (int) (Integer) key; } /** * Locates the index for the given key. This method probes using double hashing. * * @param key the key for an entry in the map. * @return the index where the key was found, or {@code -1} if no entry is found for that key. */ private int indexOf(int key) { int startIndex = hashIndex(key); int index = startIndex; for (; ; ) { if (values[index] == null) { // It's available, so no chance that this value exists anywhere in the map. return -1; } if (key == keys[index]) { return index; } // Conflict, keep probing ... if ((index = probeNext(index)) == startIndex) { return -1; } } } /** Returns the hashed index for the given key. */ private int hashIndex(int key) { // The array lengths are always a power of two, so we can use a bitmask to stay inside the array // bounds. return hashCode(key) & mask; } /** Returns the hash code for the key. */ private static int hashCode(int key) { return key; } /** Get the next sequential index after {@code index} and wraps if necessary. */ private int probeNext(int index) { // The array lengths are always a power of two, so we can use a bitmask to stay inside the array // bounds. return (index + 1) & mask; } /** Grows the map size after an insertion. If necessary, performs a rehash of the map. */ private void growSize() { size++; if (size > maxSize) { if (keys.length == Integer.MAX_VALUE) { throw new IllegalStateException("Max capacity reached at size=" + size); } // Double the capacity. rehash(keys.length << 1); } } /** * Removes entry at the given index position. Also performs opportunistic, incremental rehashing * if necessary to not break conflict chains. * * @param index the index position of the element to remove. * @return {@code true} if the next item was moved back. {@code false} otherwise. */ private boolean removeAt(final int index) { --size; // Clearing the key is not strictly necessary (for GC like in a regular collection), // but recommended for security. The memory location is still fresh in the cache anyway. keys[index] = 0; values[index] = null; // In the interval from index to the next available entry, the arrays may have entries // that are displaced from their base position due to prior conflicts. Iterate these // entries and move them back if possible, optimizing future lookups. // Knuth Section 6.4 Algorithm R, also used by the JDK's IdentityHashMap. int nextFree = index; int i = probeNext(index); for (V value = values[i]; value != null; value = values[i = probeNext(i)]) { int key = keys[i]; int bucket = hashIndex(key); if (i < bucket && (bucket <= nextFree || nextFree <= i) || bucket <= nextFree && nextFree <= i) { // Move the displaced entry "back" to the first available position. keys[nextFree] = key; values[nextFree] = value; // Put the first entry after the displaced entry keys[i] = 0; values[i] = null; nextFree = i; } } return nextFree != index; } /** Calculates the maximum size allowed before rehashing. */ private int calcMaxSize(int capacity) { // Clip the upper bound so that there will always be at least one available slot. int upperBound = capacity - 1; return Math.min(upperBound, (int) (capacity * loadFactor)); } /** * Rehashes the map for the given capacity. * * @param newCapacity the new capacity for the map. */ private void rehash(int newCapacity) { int[] oldKeys = keys; V[] oldVals = values; keys = new int[newCapacity]; @SuppressWarnings({"unchecked", "SuspiciousArrayCast"}) V[] temp = (V[]) new Object[newCapacity]; values = temp; maxSize = calcMaxSize(newCapacity); mask = newCapacity - 1; // Insert to the new arrays. for (int i = 0; i < oldVals.length; ++i) { V oldVal = oldVals[i]; if (oldVal != null) { // Inlined put(), but much simpler: we don't need to worry about // duplicated keys, growing/rehashing, or failing to insert. int oldKey = oldKeys[i]; int index = hashIndex(oldKey); for (; ; ) { if (values[index] == null) { keys[index] = oldKey; values[index] = oldVal; break; } // Conflict, keep probing. Can wrap around, but never reaches startIndex again. index = probeNext(index); } } } } @Override public String toString() { if (isEmpty()) { return "{}"; } StringBuilder sb = new StringBuilder(4 * size); sb.append('{'); boolean first = true; for (int i = 0; i < values.length; ++i) { V value = values[i]; if (value != null) { if (!first) { sb.append(", "); } sb.append(keyToString(keys[i])) .append('=') .append(value == this ? "(this Map)" : toExternal(value)); first = false; } } return sb.append('}').toString(); } /** * Helper method called by {@link #toString()} in order to convert a single map key into a string. * This is protected to allow subclasses to override the appearance of a given key. */ protected String keyToString(int key) { return Integer.toString(key); } /** Set implementation for iterating over the entries of the map. */ private final class EntrySet extends AbstractSet> { @Override public Iterator> iterator() { return new MapIterator(); } @Override public int size() { return IntObjectHashMap.this.size(); } } /** Set implementation for iterating over the keys. */ private final class KeySet extends AbstractSet { @Override public int size() { return IntObjectHashMap.this.size(); } @Override public boolean contains(Object o) { return IntObjectHashMap.this.containsKey(o); } @Override public boolean remove(Object o) { return IntObjectHashMap.this.remove(o) != null; } @Override public boolean retainAll(Collection retainedKeys) { boolean changed = false; for (Iterator> iter = entries().iterator(); iter.hasNext(); ) { PrimitiveEntry entry = iter.next(); if (!retainedKeys.contains(entry.key())) { changed = true; iter.remove(); } } return changed; } @Override public void clear() { IntObjectHashMap.this.clear(); } @Override public Iterator iterator() { return new Iterator() { private final Iterator> iter = entrySet.iterator(); @Override public boolean hasNext() { return iter.hasNext(); } @Override public Integer next() { return iter.next().getKey(); } @Override public void remove() { iter.remove(); } }; } } /** * Iterator over primitive entries. Entry key/values are overwritten by each call to {@link * #next()}. */ private final class PrimitiveIterator implements Iterator>, PrimitiveEntry { private int prevIndex = -1; private int nextIndex = -1; private int entryIndex = -1; private void scanNext() { while (++nextIndex != values.length && values[nextIndex] == null) {} } @Override public boolean hasNext() { if (nextIndex == -1) { scanNext(); } return nextIndex != values.length; } @Override public PrimitiveEntry next() { if (!hasNext()) { throw new NoSuchElementException(); } prevIndex = nextIndex; scanNext(); // Always return the same Entry object, just change its index each time. entryIndex = prevIndex; return this; } @Override public void remove() { if (prevIndex == -1) { throw new IllegalStateException("next must be called before each remove."); } if (removeAt(prevIndex)) { // removeAt may move elements "back" in the array if they have been displaced because their // spot in the // array was occupied when they were inserted. If this occurs then the nextIndex is now // invalid and // should instead point to the prevIndex which now holds an element which was "moved back". nextIndex = prevIndex; } prevIndex = -1; } // Entry implementation. Since this implementation uses a single Entry, we coalesce that // into the Iterator object (potentially making loop optimization much easier). @Override public int key() { return keys[entryIndex]; } @Override public V value() { return toExternal(values[entryIndex]); } @Override public void setValue(V value) { values[entryIndex] = toInternal(value); } } /** Iterator used by the {@link Map} interface. */ private final class MapIterator implements Iterator> { private final PrimitiveIterator iter = new PrimitiveIterator(); @Override public boolean hasNext() { return iter.hasNext(); } @Override public Entry next() { if (!hasNext()) { throw new NoSuchElementException(); } iter.next(); return new MapEntry(iter.entryIndex); } @Override public void remove() { iter.remove(); } } /** A single entry in the map. */ final class MapEntry implements Entry { private final int entryIndex; MapEntry(int entryIndex) { this.entryIndex = entryIndex; } @Override public Integer getKey() { verifyExists(); return keys[entryIndex]; } @Override public V getValue() { verifyExists(); return toExternal(values[entryIndex]); } @Override public V setValue(V value) { verifyExists(); V prevValue = toExternal(values[entryIndex]); values[entryIndex] = toInternal(value); return prevValue; } private void verifyExists() { if (values[entryIndex] == null) { throw new IllegalStateException("The map entry has been removed"); } } } static int safeFindNextPositivePowerOfTwo(final int value) { return value <= 0 ? 1 : value >= 0x40000000 ? 0x40000000 : findNextPositivePowerOfTwo(value); } static int findNextPositivePowerOfTwo(final int value) { assert value > Integer.MIN_VALUE && value < 0x40000000; return 1 << (32 - Integer.numberOfLeadingZeros(value - 1)); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/IntObjectMap.java000066400000000000000000000047731515451615200316520ustar00rootroot00000000000000/* * 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. */ package org.apache.arrow.vector.util; import java.util.Iterator; import java.util.Map; /** * A vendored specialized copy of Netty's IntObjectMap for use within Arrow. Avoids requiring Netty * in the Arrow core just for this one class. * * @param the value type stored in the map. */ interface IntObjectMap extends Map { /** * A primitive entry in the map, provided by the iterator from {@link #entries()}. * * @param the value type stored in the map. */ interface PrimitiveEntry { /** Gets the key for this entry. */ int key(); /** Gets the value for this entry. */ V value(); /** Sets the value for this entry. */ void setValue(V value); } /** * Gets the value in the map with the specified key. * * @param key the key whose associated value is to be returned. * @return the value or {@code null} if the key was not found in the map. */ V get(int key); /** * Puts the given entry into the map. * * @param key the key of the entry. * @param value the value of the entry. * @return the previous value for this key or {@code null} if there was no previous mapping. */ V put(int key, V value); /** * Removes the entry with the specified key. * * @param key the key for the entry to be removed from this map. * @return the previous value for the key, or {@code null} if there was no mapping. */ V remove(int key); /** * Gets an iterable to traverse over the primitive entries contained in this map. As an * optimization, the {@link PrimitiveEntry}s returned by the {@link Iterator} may change as the * {@link Iterator} progresses. The caller should not rely on {@link PrimitiveEntry} key/value * stability. */ Iterable> entries(); /** Indicates whether or not this map contains a value for the specified key. */ boolean containsKey(int key); } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/JsonStringArrayList.java000066400000000000000000000032721515451615200332570ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import java.util.ArrayList; /** * Extension of {@link ArrayList} that {@link #toString()} method returns the serialized JSON * version of its members (or throws an exception if they can't be converted to JSON). * * @param Type of value held in the list. */ public class JsonStringArrayList extends ArrayList { private static final ObjectMapper MAPPER = ObjectMapperFactory.newObjectMapper(); public JsonStringArrayList() { super(); } public JsonStringArrayList(int size) { super(size); } @Override public final String toString() { try { return MAPPER.writeValueAsString(this); } catch (JsonProcessingException e) { throw new IllegalStateException("Cannot serialize array list to JSON string", e); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/JsonStringHashMap.java000066400000000000000000000031471515451615200326670ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import java.util.LinkedHashMap; /** * Simple class that extends the regular java.util.HashMap but overrides the toString() method of * the HashMap class to produce a JSON string instead * * @param The type of the key for the map. * @param The type of the value for the map. */ public class JsonStringHashMap extends LinkedHashMap { private static final ObjectMapper MAPPER = ObjectMapperFactory.newObjectMapper(); @Override public final String toString() { try { return MAPPER.writeValueAsString(this); } catch (JsonProcessingException e) { throw new IllegalStateException("Cannot serialize hash map to JSON string", e); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/MapWithOrdinal.java000066400000000000000000000043021515451615200322010ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.Collection; import java.util.Map; import java.util.Set; /** * An implementation of a map that supports constant time look-up by a generic key or an ordinal. * *

    This class extends the functionality a regular {@link Map} with ordinal lookup support. Upon * insertion an unused ordinal is assigned to the inserted (key, value) tuple. Upon update the same * ordinal id is re-used while value is replaced. Upon deletion of an existing item, its * corresponding ordinal is recycled and could be used by another item. * *

    For any instance with N items, this implementation guarantees that ordinals are in the range * of [0, N). However, the ordinal assignment is dynamic and may change after an insertion or * deletion. Consumers of this class are responsible for explicitly checking the ordinal * corresponding to a key via {@link MultiMapWithOrdinal#getOrdinal(Object)} before attempting to * execute a lookup with an ordinal. * * @param key type * @param value type */ public interface MapWithOrdinal { V getByOrdinal(int id); int getOrdinal(K key); int size(); boolean isEmpty(); V get(K key); Collection getAll(K key); boolean put(K key, V value, boolean overwrite); Collection values(); boolean containsKey(K key); boolean remove(K key, V value); boolean removeAll(K key); void clear(); Set keys(); } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/MapWithOrdinalImpl.java000066400000000000000000000160541515451615200330320ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.AbstractMap; import java.util.ArrayList; import java.util.Collection; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; import java.util.stream.Collectors; /** * An implementation of map that supports constant time look-up by a generic key or an ordinal. * *

    This class extends the functionality a regular {@link Map} with ordinal lookup support. Upon * insertion an unused ordinal is assigned to the inserted (key, value) tuple. Upon update the same * ordinal id is re-used while value is replaced. Upon deletion of an existing item, its * corresponding ordinal is recycled and could be used by another item. * *

    For any instance with N items, this implementation guarantees that ordinals are in the range * of [0, N). However, the ordinal assignment is dynamic and may change after an insertion or * deletion. Consumers of this class are responsible for explicitly checking the ordinal * corresponding to a key via {@link MapWithOrdinalImpl#getOrdinal(Object)} before attempting to * execute a lookup with an ordinal. * * @param key type * @param value type */ public class MapWithOrdinalImpl implements MapWithOrdinal { private final Map> primary = new LinkedHashMap<>(); private final IntObjectHashMap secondary = new IntObjectHashMap<>(); private final Map delegate = new Map() { @Override public boolean isEmpty() { return size() == 0; } @Override public int size() { return primary.size(); } @Override public boolean containsKey(Object key) { return primary.containsKey(key); } @Override public boolean containsValue(Object value) { return primary.containsValue(value); } @Override public V get(Object key) { Entry pair = primary.get(key); if (pair != null) { return pair.getValue(); } return null; } @Override public V put(K key, V value) { final Entry oldPair = primary.get(key); // if key exists try replacing otherwise, assign a new ordinal identifier final int ordinal = oldPair == null ? primary.size() : oldPair.getKey(); primary.put(key, new AbstractMap.SimpleImmutableEntry<>(ordinal, value)); secondary.put(ordinal, value); return oldPair == null ? null : oldPair.getValue(); } @Override public V remove(Object key) { final Entry oldPair = primary.remove(key); if (oldPair != null) { final int lastOrdinal = secondary.size(); final V last = secondary.get(lastOrdinal); // normalize mappings so that all numbers until primary.size() is assigned // swap the last element with the deleted one secondary.put(oldPair.getKey(), last); primary.put((K) key, new AbstractMap.SimpleImmutableEntry<>(oldPair.getKey(), last)); } return oldPair == null ? null : oldPair.getValue(); } @Override public void putAll(Map m) { throw new UnsupportedOperationException(); } @Override public void clear() { primary.clear(); secondary.clear(); } @Override public Set keySet() { return primary.keySet(); } @Override public Collection values() { return secondary.values(); } @Override public Set> entrySet() { return primary.entrySet().stream() .map( entry -> new AbstractMap.SimpleImmutableEntry<>( entry.getKey(), entry.getValue().getValue())) .collect(Collectors.toSet()); } }; /** * Returns the value corresponding to the given ordinal. * * @param id ordinal value for lookup * @return an instance of V */ @Override public V getByOrdinal(int id) { return secondary.get(id); } /** * Returns the ordinal corresponding to the given key. * * @param key key for ordinal lookup * @return ordinal value corresponding to key if it exists or -1 */ @Override public int getOrdinal(K key) { Map.Entry pair = primary.get(key); if (pair != null) { return pair.getKey(); } return -1; } @Override public int size() { return delegate.size(); } @Override public boolean isEmpty() { return delegate.isEmpty(); } @Override public Collection getAll(K key) { if (delegate.containsKey(key)) { List list = new ArrayList<>(1); list.add(get(key)); return list; } return null; } @Override public V get(K key) { return delegate.get(key); } /** * Inserts the tuple (key, value) into the map extending the semantics of {@link Map#put} with * automatic ordinal assignment. A new ordinal is assigned if key does not exists. Otherwise the * same ordinal is re-used but the value is replaced. * * @see java.util.Map#put */ @Override public boolean put(K key, V value, boolean overwrite) { return delegate.put(key, value) != null; } @Override public Collection values() { return delegate.values(); } @Override public boolean remove(K key, V value) { return false; } @Override public boolean containsKey(Object key) { return delegate.containsKey(key); } /** * Removes the element corresponding to the key if exists extending the semantics of {@link * java.util.Map#remove} with ordinal re-cycling. The ordinal corresponding to the given key may * be re-assigned to another tuple. It is important that consumer checks the ordinal value via * {@link MapWithOrdinalImpl#getOrdinal(Object)} before attempting to look-up by ordinal. * * @see java.util.Map#remove */ @Override public boolean removeAll(K key) { return delegate.remove(key) != null; } @Override public void clear() { delegate.clear(); } @Override public Set keys() { return delegate.keySet(); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/MultiMapWithOrdinal.java000066400000000000000000000160451515451615200332230ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.Collection; import java.util.HashSet; import java.util.LinkedHashMap; import java.util.Map; import java.util.Optional; import java.util.Set; import java.util.stream.Collectors; /** * An implementation of a multimap that supports constant time look-up by a generic key or an * ordinal. * *

    This class extends the functionality a regular {@link Map} with ordinal lookup support. Upon * insertion an unused ordinal is assigned to the inserted (key, value) tuple. Upon update the same * ordinal id is re-used while value is replaced. Upon deletion of an existing item, its * corresponding ordinal is recycled and could be used by another item. * *

    For any instance with N items, this implementation guarantees that ordinals are in the range * of [0, N). However, the ordinal assignment is dynamic and may change after an insertion or * deletion. Consumers of this class are responsible for explicitly checking the ordinal * corresponding to a key via {@link MultiMapWithOrdinal#getOrdinal(Object)} before attempting to * execute a lookup with an ordinal. * * @param key type * @param value type */ public class MultiMapWithOrdinal implements MapWithOrdinal { private final Map> keyToOrdinal = new LinkedHashMap<>(); private final IntObjectHashMap ordinalToValue = new IntObjectHashMap<>(); /** * Returns the value corresponding to the given ordinal. * * @param id ordinal value for lookup * @return an instance of V */ @Override public V getByOrdinal(int id) { return ordinalToValue.get(id); } /** * Returns the ordinal corresponding to the given key. * * @param key key for ordinal lookup * @return ordinal value corresponding to key if it exists or -1 */ @Override public int getOrdinal(K key) { Set pair = getOrdinals(key); if (!pair.isEmpty()) { return pair.iterator().next(); } return -1; } private Set getOrdinals(K key) { return keyToOrdinal.getOrDefault(key, new HashSet<>()); } @Override public int size() { return ordinalToValue.size(); } @Override public boolean isEmpty() { return ordinalToValue.isEmpty(); } /** get set of values for key. */ @Override public V get(K key) { Set ordinals = keyToOrdinal.get(key); if (ordinals == null) { return null; } return ordinals.stream().map(ordinalToValue::get).collect(Collectors.toList()).get(0); } /** get set of values for key. */ @Override public Collection getAll(K key) { Set ordinals = keyToOrdinal.get(key); if (ordinals == null) { return null; } return ordinals.stream().map(ordinalToValue::get).collect(Collectors.toList()); } /** * Inserts the tuple (key, value) into the multimap with automatic ordinal assignment. * *

    A new ordinal is assigned if key/value pair does not exists. * *

    If overwrite is true the existing key will be overwritten with value else value will be * appended to the multimap. */ @Override public boolean put(K key, V value, boolean overwrite) { if (overwrite) { removeAll(key); } Set ordinalSet = getOrdinals(key); int nextOrdinal = ordinalToValue.size(); ordinalToValue.put(nextOrdinal, value); boolean changed = ordinalSet.add(nextOrdinal); keyToOrdinal.put(key, ordinalSet); return changed; } @Override public Collection values() { return ordinalToValue.values(); } @Override public boolean containsKey(K key) { return keyToOrdinal.containsKey(key); } /** * Removes the element corresponding to the key/value if exists with ordinal re-cycling. * *

    The ordinal corresponding to the given key may be re-assigned to another tuple. It is * important that consumer checks the ordinal value via {@link * MultiMapWithOrdinal#getOrdinal(Object)} before attempting to look-up by ordinal. * *

    If the multimap is changed return true. */ @Override public synchronized boolean remove(K key, V value) { Set removalSet = getOrdinals(key); if (removalSet.isEmpty()) { return false; } Optional removeValue = removalSet.stream().map(ordinalToValue::get).filter(value::equals).findFirst(); if (!removeValue.isPresent()) { return false; } int removalOrdinal = removeKv(removalSet, key, value); int lastOrdinal = ordinalToValue.size(); if (lastOrdinal != removalOrdinal) { // we didn't remove the last ordinal swapOrdinal(lastOrdinal, removalOrdinal); } return true; } private void swapOrdinal(int lastOrdinal, int removalOrdinal) { V swapOrdinalValue = ordinalToValue.remove(lastOrdinal); ordinalToValue.put(removalOrdinal, swapOrdinalValue); K swapOrdinalKey = keyToOrdinal.entrySet().stream() .filter(kv -> kv.getValue().stream().anyMatch(o -> o == lastOrdinal)) .map(Map.Entry::getKey) .findFirst() .orElseThrow(() -> new IllegalStateException("MultimapWithOrdinal in bad state")); ordinalToValue.put(removalOrdinal, swapOrdinalValue); Set swapSet = getOrdinals(swapOrdinalKey); swapSet.remove(lastOrdinal); swapSet.add(removalOrdinal); keyToOrdinal.put(swapOrdinalKey, swapSet); } private int removeKv(Set removalSet, K key, V value) { Integer removalOrdinal = removalSet.stream() .filter(i -> ordinalToValue.get(i).equals(value)) .findFirst() .orElseThrow(() -> new IllegalStateException("MultimapWithOrdinal in bad state")); ordinalToValue.remove(removalOrdinal); removalSet.remove(removalOrdinal); if (removalSet.isEmpty()) { keyToOrdinal.remove(key); } else { keyToOrdinal.put(key, removalSet); } return removalOrdinal; } /** remove all entries of key. */ @Override public synchronized boolean removeAll(K key) { Collection values = this.getAll(key); if (values == null) { return false; } for (V v : values) { this.remove(key, v); } return true; } @Override public void clear() { ordinalToValue.clear(); keyToOrdinal.clear(); } @Override public Set keys() { return keyToOrdinal.keySet(); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/ObjectMapperFactory.java000066400000000000000000000024571515451615200332330ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.json.JsonMapper; import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; /** A {@link ObjectMapper} factory to read/write JSON. */ public final class ObjectMapperFactory { private ObjectMapperFactory() {} /** Creates a new {@link ObjectMapper} instance. */ public static ObjectMapper newObjectMapper() { return JsonMapper.builder().addModule(new JavaTimeModule()).build(); } } OversizedAllocationException.java000066400000000000000000000033371515451615200351060ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** * An exception that is used to signal that allocation request in bytes is greater than the maximum * allowed by {@link org.apache.arrow.memory.BufferAllocator#buffer(int) allocator}. * *

    Operators should handle this exception to split the batch and later resume the execution on * the next iteration. */ public class OversizedAllocationException extends RuntimeException { public OversizedAllocationException() { super(); } public OversizedAllocationException( String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) { super(message, cause, enableSuppression, writableStackTrace); } public OversizedAllocationException(String message, Throwable cause) { super(message, cause); } public OversizedAllocationException(String message) { super(message); } public OversizedAllocationException(Throwable cause) { super(cause); } } PromotableMultiMapWithOrdinal.java000066400000000000000000000070701515451615200351670ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.Collection; import java.util.Set; import org.apache.arrow.vector.complex.AbstractStructVector; /** * Implementation of MapWithOrdinal that allows for promotion to multimap when duplicate fields * exist. * * @param key type * @param value type */ public class PromotableMultiMapWithOrdinal implements MapWithOrdinal { private final MapWithOrdinalImpl mapWithOrdinal = new MapWithOrdinalImpl<>(); private final MultiMapWithOrdinal multiMapWithOrdinal = new MultiMapWithOrdinal<>(); private final boolean promotable; private AbstractStructVector.ConflictPolicy conflictPolicy; private MapWithOrdinal delegate; /** * Create promotable map. * * @param promotable if promotion is allowed, otherwise delegate to MapWithOrdinal. * @param conflictPolicy how to handle name conflicts. */ public PromotableMultiMapWithOrdinal( boolean promotable, AbstractStructVector.ConflictPolicy conflictPolicy) { this.promotable = promotable; this.conflictPolicy = conflictPolicy; delegate = mapWithOrdinal; } private void promote() { if (delegate == multiMapWithOrdinal || !promotable || conflictPolicy.equals(AbstractStructVector.ConflictPolicy.CONFLICT_REPLACE)) { return; } for (K key : mapWithOrdinal.keys()) { V value = mapWithOrdinal.get(key); multiMapWithOrdinal.put(key, value, false); } mapWithOrdinal.clear(); delegate = multiMapWithOrdinal; } @Override public V getByOrdinal(int id) { return delegate.getByOrdinal(id); } @Override public int getOrdinal(K key) { return delegate.getOrdinal(key); } @Override public int size() { return delegate.size(); } @Override public boolean isEmpty() { return delegate.isEmpty(); } @Override public V get(K key) { return delegate.get(key); } @Override public Collection getAll(K key) { return delegate.getAll(key); } @Override public boolean put(K key, V value, boolean overwrite) { if (delegate.containsKey(key)) { promote(); } return delegate.put(key, value, overwrite); } @Override public Collection values() { return delegate.values(); } @Override public boolean containsKey(K key) { return delegate.containsKey(key); } @Override public boolean remove(K key, V value) { return delegate.remove(key, value); } @Override public boolean removeAll(K key) { return delegate.removeAll(key); } @Override public void clear() { delegate.clear(); } @Override public Set keys() { return delegate.keys(); } public void setConflictPolicy(AbstractStructVector.ConflictPolicy conflictPolicy) { this.conflictPolicy = conflictPolicy; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/ReusableByteArray.java000066400000000000000000000067511515451615200327160ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.Arrays; import java.util.Base64; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.ReusableBuffer; /** A wrapper around byte arrays for repeated writing. */ public class ReusableByteArray implements ReusableBuffer { protected static final byte[] EMPTY_BYTES = new byte[0]; protected byte[] bytes; protected int length; public ReusableByteArray() { bytes = EMPTY_BYTES; } public ReusableByteArray(byte[] data) { bytes = Arrays.copyOfRange(data, 0, data.length); length = data.length; } /** * Get the number of bytes in the byte array. * * @return the number of bytes in the byte array */ @Override public long getLength() { return length; } @Override public byte[] getBuffer() { return bytes; } @Override public void set(ArrowBuf srcBytes, long start, long len) { setCapacity((int) len, false); srcBytes.getBytes(start, bytes, 0, (int) len); length = (int) len; } @Override public void set(byte[] srcBytes, long start, long len) { setCapacity((int) len, false); System.arraycopy(srcBytes, (int) start, bytes, 0, (int) len); length = (int) len; } @Override public boolean equals(Object o) { if (o == this) { return true; } else if (o == null) { return false; } if (!(o instanceof ReusableByteArray)) { return false; } final ReusableByteArray that = (ReusableByteArray) o; if (this.getLength() != that.getLength()) { return false; } for (int i = 0; i < length; i++) { if (bytes[i] != that.bytes[i]) { return false; } } return true; } @Override public int hashCode() { if (bytes == null) { return 0; } int result = 1; for (int i = 0; i < length; i++) { result = 31 * result + bytes[i]; } return result; } @Override public String toString() { return Base64.getEncoder().encodeToString(Arrays.copyOfRange(bytes, 0, length)); } /** * Sets the capacity of this object to at least len bytes. If the current * buffer is longer, then the capacity and existing content of the buffer are unchanged. If * len is larger than the current capacity, the Text object's capacity is increased to * match. * * @param len the number of bytes we need * @param keepData should the old data be kept */ protected void setCapacity(int len, boolean keepData) { if (bytes == null || bytes.length < len) { if (bytes != null && keepData) { bytes = Arrays.copyOf(bytes, Math.max(len, length << 1)); } else { bytes = new byte[len]; } } } } SchemaChangeRuntimeException.java000066400000000000000000000027621515451615200350010ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; /** Thrown when child vectors (e.g. in lists) don't match the expected type. */ public class SchemaChangeRuntimeException extends RuntimeException { public SchemaChangeRuntimeException() { super(); } public SchemaChangeRuntimeException( String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) { super(message, cause, enableSuppression, writableStackTrace); } public SchemaChangeRuntimeException(String message, Throwable cause) { super(message, cause); } public SchemaChangeRuntimeException(String message) { super(message); } public SchemaChangeRuntimeException(Throwable cause) { super(cause); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/SchemaUtility.java000066400000000000000000000044271515451615200321130ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.channels.Channels; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.ipc.ReadChannel; import org.apache.arrow.vector.ipc.WriteChannel; import org.apache.arrow.vector.ipc.message.MessageChannelReader; import org.apache.arrow.vector.ipc.message.MessageResult; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.pojo.Schema; /** Schema utility class including serialization and deserialization. */ public class SchemaUtility { private SchemaUtility() {} /** Deserialize Arrow schema from byte array. */ public static Schema deserialize(byte[] bytes, BufferAllocator allocator) throws IOException { try (MessageChannelReader schemaReader = new MessageChannelReader( new ReadChannel(new ByteArrayReadableSeekableByteChannel(bytes)), allocator)) { MessageResult result = schemaReader.readNext(); if (result == null) { throw new IOException("Unexpected end of input. Missing schema."); } return MessageSerializer.deserializeSchema(result.getMessage()); } } /** Serialize Arrow schema into byte array. */ public static byte[] serialize(Schema schema) throws IOException { final ByteArrayOutputStream out = new ByteArrayOutputStream(); MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), schema); return out.toByteArray(); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/Text.java000066400000000000000000000512301515451615200302450ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 com.fasterxml.jackson.core.JsonGenerationException; import com.fasterxml.jackson.core.JsonGenerator; import com.fasterxml.jackson.databind.SerializerProvider; import com.fasterxml.jackson.databind.annotation.JsonSerialize; import com.fasterxml.jackson.databind.ser.std.StdSerializer; import java.io.DataInput; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.CharBuffer; import java.nio.charset.CharacterCodingException; import java.nio.charset.Charset; import java.nio.charset.CharsetDecoder; import java.nio.charset.CharsetEncoder; import java.nio.charset.CodingErrorAction; import java.nio.charset.MalformedInputException; import java.text.CharacterIterator; import java.text.StringCharacterIterator; import java.util.Optional; /** * A simplified byte wrapper similar to Hadoop's Text class without all the dependencies. Lifted * from Hadoop 2.7.1 */ @JsonSerialize(using = Text.TextSerializer.class) public class Text extends ReusableByteArray { private static ThreadLocal ENCODER_FACTORY = new ThreadLocal() { @Override protected CharsetEncoder initialValue() { return Charset.forName("UTF-8") .newEncoder() .onMalformedInput(CodingErrorAction.REPORT) .onUnmappableCharacter(CodingErrorAction.REPORT); } }; private static ThreadLocal DECODER_FACTORY = new ThreadLocal() { @Override protected CharsetDecoder initialValue() { return Charset.forName("UTF-8") .newDecoder() .onMalformedInput(CodingErrorAction.REPORT) .onUnmappableCharacter(CodingErrorAction.REPORT); } }; public Text() { super(); } /** * Construct from a string. * * @param string initialize from that string */ public Text(String string) { set(string); } /** * Construct from another text. * * @param utf8 initialize from that Text */ public Text(Text utf8) { set(utf8); } /** * Construct from a byte array. * * @param utf8 initialize from that byte array */ public Text(byte[] utf8) { set(utf8); } /** * Get a copy of the bytes that is exactly the length of the data. See {@link #getBytes()} for * faster access to the underlying array. * * @return a copy of the underlying array */ public byte[] copyBytes() { byte[] result = new byte[length]; System.arraycopy(bytes, 0, result, 0, length); return result; } /** * Returns the raw bytes; however, only data up to {@link #getLength()} is valid. Please use * {@link #copyBytes()} if you need the returned array to be precisely the length of the data. * * @return the underlying array */ public byte[] getBytes() { return bytes; } /** * Returns the Unicode Scalar Value (32-bit integer value) for the character at position * . Note that this method avoids using the converter or doing String instantiation. * * @param position the index of the char we want to retrieve * @return the Unicode scalar value at position or -1 if the position is invalid or points to a * trailing byte */ public int charAt(int position) { if (position > this.length) { return -1; // too long } if (position < 0) { return -1; // duh. } ByteBuffer bb = (ByteBuffer) ByteBuffer.wrap(bytes).position(position); return bytesToCodePoint(bb.slice()); } public int find(String what) { return find(what, 0); } /** * Finds any occurrence of what in the backing buffer, starting as position * start. The starting position is measured in bytes and the return value is in terms of * byte position in the buffer. The backing buffer is not converted to a string for this * operation. * * @param what the string to search for * @param start where to start from * @return byte position of the first occurrence of the search string in the UTF-8 buffer or -1 if * not found */ public int find(String what, int start) { try { ByteBuffer src = ByteBuffer.wrap(this.bytes, 0, this.length); ByteBuffer tgt = encode(what); byte b = tgt.get(); src.position(start); while (src.hasRemaining()) { if (b == src.get()) { // matching first byte src.mark(); // save position in loop tgt.mark(); // save position in target boolean found = true; int pos = src.position() - 1; while (tgt.hasRemaining()) { if (!src.hasRemaining()) { // src expired first tgt.reset(); src.reset(); found = false; break; } if (!(tgt.get() == src.get())) { tgt.reset(); src.reset(); found = false; break; // no match } } if (found) { return pos; } } } return -1; // not found } catch (CharacterCodingException e) { // can't get here e.printStackTrace(); return -1; } } /** * Set to contain the contents of a string. * * @param string the string to initialize from */ public void set(String string) { try { ByteBuffer bb = encode(string, true); bytes = bb.array(); length = bb.limit(); } catch (CharacterCodingException e) { throw new RuntimeException("Should not have happened ", e); } } /** * Set to an utf8 byte array. * * @param utf8 the byte array to initialize from */ public void set(byte[] utf8) { set(utf8, 0, utf8.length); } /** * copy a text. * * @param other the text to initialize from */ public void set(Text other) { set(other.getBytes(), 0, (int) other.getLength()); } /** * Set the Text to range of bytes. * * @param utf8 the data to copy from * @param start the first position of the new string * @param len the number of bytes of the new string */ public void set(byte[] utf8, int start, int len) { super.set(utf8, start, len); } /** * Append a range of bytes to the end of the given text. * * @param utf8 the data to copy from * @param start the first position to append from utf8 * @param len the number of bytes to append */ public void append(byte[] utf8, int start, int len) { setCapacity(length + len, true); System.arraycopy(utf8, start, bytes, length, len); length += len; } /** * Clear the string to empty. * *

    Note: For performance reasons, this call does not clear the underlying byte array * that is retrievable via {@link #getBytes()}. In order to free the byte-array memory, call * {@link #set(byte[])} with an empty byte array (For example, new byte[0]). */ public void clear() { length = 0; } @Override public String toString() { try { return decode(bytes, 0, length); } catch (CharacterCodingException e) { throw new RuntimeException("Should not have happened ", e); } } /** * Read a Text object whose length is already known. This allows creating Text from a stream which * uses a different serialization format. * * @param in the input to initialize from * @param len how many bytes to read from in * @throws IOException if something bad happens */ public void readWithKnownLength(DataInput in, int len) throws IOException { setCapacity(len, false); in.readFully(bytes, 0, len); length = len; } @Override public boolean equals(Object o) { if (!(o instanceof Text)) { return false; } return super.equals(o); } // / STATIC UTILITIES FROM HERE DOWN /** * Converts the provided byte array to a String using the UTF-8 encoding. If the input is * malformed, replace by a default value. * * @param utf8 bytes to decode * @return the decoded string * @throws CharacterCodingException if this is not valid UTF-8 */ public static String decode(byte[] utf8) throws CharacterCodingException { return decode(ByteBuffer.wrap(utf8), true); } public static String decode(byte[] utf8, int start, int length) throws CharacterCodingException { return decode(ByteBuffer.wrap(utf8, start, length), true); } /** * Converts the provided byte array to a String using the UTF-8 encoding. If replace * is true, then malformed input is replaced with the substitution character, which is U+FFFD. * Otherwise the method throws a MalformedInputException. * * @param utf8 the bytes to decode * @param start where to start from * @param length length of the bytes to decode * @param replace whether to replace malformed characters with U+FFFD * @return the decoded string * @throws CharacterCodingException if the input could not be decoded */ public static String decode(byte[] utf8, int start, int length, boolean replace) throws CharacterCodingException { return decode(ByteBuffer.wrap(utf8, start, length), replace); } private static String decode(ByteBuffer utf8, boolean replace) throws CharacterCodingException { CharsetDecoder decoder = DECODER_FACTORY.get(); if (replace) { decoder.onMalformedInput(java.nio.charset.CodingErrorAction.REPLACE); decoder.onUnmappableCharacter(CodingErrorAction.REPLACE); } String str = decoder.decode(utf8).toString(); // set decoder back to its default value: REPORT if (replace) { decoder.onMalformedInput(CodingErrorAction.REPORT); decoder.onUnmappableCharacter(CodingErrorAction.REPORT); } return str; } /** * Converts the provided String to bytes using the UTF-8 encoding. If the input is malformed, * invalid chars are replaced by a default value. * * @param string the string to encode * @return ByteBuffer: bytes stores at ByteBuffer.array() and length is ByteBuffer.limit() * @throws CharacterCodingException if the string could not be encoded */ public static ByteBuffer encode(String string) throws CharacterCodingException { return encode(string, true); } /** * Converts the provided String to bytes using the UTF-8 encoding. If replace is * true, then malformed input is replaced with the substitution character, which is U+FFFD. * Otherwise the method throws a MalformedInputException. * * @param string the string to encode * @param replace whether to replace malformed characters with U+FFFD * @return ByteBuffer: bytes stores at ByteBuffer.array() and length is ByteBuffer.limit() * @throws CharacterCodingException if the string could not be encoded */ public static ByteBuffer encode(String string, boolean replace) throws CharacterCodingException { CharsetEncoder encoder = ENCODER_FACTORY.get(); if (replace) { encoder.onMalformedInput(CodingErrorAction.REPLACE); encoder.onUnmappableCharacter(CodingErrorAction.REPLACE); } ByteBuffer bytes = encoder.encode(CharBuffer.wrap(string.toCharArray())); if (replace) { encoder.onMalformedInput(CodingErrorAction.REPORT); encoder.onUnmappableCharacter(CodingErrorAction.REPORT); } return bytes; } public static final int DEFAULT_MAX_LEN = 1024 * 1024; // //// states for validateUTF8 private static final int LEAD_BYTE = 0; private static final int TRAIL_BYTE_1 = 1; private static final int TRAIL_BYTE = 2; /** * Check if a byte array contains valid utf-8. * * @param utf8 byte array * @return true if the input is valid UTF-8. False otherwise. */ public static boolean validateUTF8NoThrow(byte[] utf8) { return !validateUTF8Internal(utf8, 0, utf8.length).isPresent(); } /** * Check if a byte array contains valid utf-8. * * @param utf8 byte array * @throws MalformedInputException if the byte array contains invalid utf-8 */ public static void validateUTF8(byte[] utf8) throws MalformedInputException { validateUTF8(utf8, 0, utf8.length); } /** * Check to see if a byte array is valid utf-8. * * @param utf8 the array of bytes * @param start the offset of the first byte in the array * @param len the length of the byte sequence * @throws MalformedInputException if the byte array contains invalid bytes */ public static void validateUTF8(byte[] utf8, int start, int len) throws MalformedInputException { Optional result = validateUTF8Internal(utf8, start, len); if (result.isPresent()) { throw new MalformedInputException(result.get()); } } /** * Check to see if a byte array is valid utf-8. * * @param utf8 the array of bytes * @param start the offset of the first byte in the array * @param len the length of the byte sequence * @return the position where a malformed byte occurred or Optional.empty() if the byte array was * valid UTF-8. */ private static Optional validateUTF8Internal(byte[] utf8, int start, int len) { int count = start; int leadByte = 0; int length = 0; int state = LEAD_BYTE; while (count < start + len) { int aByte = utf8[count] & 0xFF; switch (state) { case LEAD_BYTE: leadByte = aByte; length = bytesFromUTF8[aByte]; switch (length) { case 0: // check for ASCII if (leadByte > 0x7F) { return Optional.of(count); } break; case 1: if (leadByte < 0xC2 || leadByte > 0xDF) { return Optional.of(count); } state = TRAIL_BYTE_1; break; case 2: if (leadByte < 0xE0 || leadByte > 0xEF) { return Optional.of(count); } state = TRAIL_BYTE_1; break; case 3: if (leadByte < 0xF0 || leadByte > 0xF4) { return Optional.of(count); } state = TRAIL_BYTE_1; break; default: // too long! Longest valid UTF-8 is 4 bytes (lead + three) // or if < 0 we got a trail byte in the lead byte position return Optional.of(count); } // switch (length) break; case TRAIL_BYTE_1: if (leadByte == 0xF0 && aByte < 0x90) { return Optional.of(count); } if (leadByte == 0xF4 && aByte > 0x8F) { return Optional.of(count); } if (leadByte == 0xE0 && aByte < 0xA0) { return Optional.of(count); } if (leadByte == 0xED && aByte > 0x9F) { return Optional.of(count); } // falls through to regular trail-byte test!! case TRAIL_BYTE: if (aByte < 0x80 || aByte > 0xBF) { return Optional.of(count); } if (--length == 0) { state = LEAD_BYTE; } else { state = TRAIL_BYTE; } break; default: break; } // switch (state) count++; } return Optional.empty(); } /** * Magic numbers for UTF-8. These are the number of bytes that follow a given lead byte. * Trailing bytes have the value -1. The values 4 and 5 are presented in this table, even though * valid UTF-8 cannot include the five and six byte sequences. */ static final int[] bytesFromUTF8 = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // trail bytes -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5 }; /** * Returns the next code point at the current position in the buffer. The buffer's position will * be incremented. Any mark set on this buffer will be changed by this method! * * @param bytes the incoming bytes * @return the corresponding unicode codepoint */ public static int bytesToCodePoint(ByteBuffer bytes) { bytes.mark(); byte b = bytes.get(); bytes.reset(); int extraBytesToRead = bytesFromUTF8[(b & 0xFF)]; if (extraBytesToRead < 0) { return -1; // trailing byte! } int ch = 0; switch (extraBytesToRead) { case 5: ch += (bytes.get() & 0xFF); ch <<= 6; /* remember, illegal UTF-8 */ // fall through case 4: ch += (bytes.get() & 0xFF); ch <<= 6; /* remember, illegal UTF-8 */ // fall through case 3: ch += (bytes.get() & 0xFF); ch <<= 6; // fall through case 2: ch += (bytes.get() & 0xFF); ch <<= 6; // fall through case 1: ch += (bytes.get() & 0xFF); ch <<= 6; // fall through case 0: ch += (bytes.get() & 0xFF); break; default: // do nothing } ch -= offsetsFromUTF8[extraBytesToRead]; return ch; } static final int[] offsetsFromUTF8 = { 0x00000000, 0x00003080, 0x000E2080, 0x03C82080, 0xFA082080, 0x82082080 }; /** * For the given string, returns the number of UTF-8 bytes required to encode the string. * * @param string text to encode * @return number of UTF-8 bytes required to encode */ public static int utf8Length(String string) { CharacterIterator iter = new StringCharacterIterator(string); char ch = iter.first(); int size = 0; while (ch != CharacterIterator.DONE) { if ((ch >= 0xD800) && (ch < 0xDC00)) { // surrogate pair? char trail = iter.next(); if ((trail > 0xDBFF) && (trail < 0xE000)) { // valid pair size += 4; } else { // invalid pair size += 3; iter.previous(); // rewind one } } else if (ch < 0x80) { size++; } else if (ch < 0x800) { size += 2; } else { // ch < 0x10000, that is, the largest char value size += 3; } ch = iter.next(); } return size; } /** JSON serializer for {@link Text}. */ public static class TextSerializer extends StdSerializer { public TextSerializer() { super(Text.class); } @Override public void serialize( Text text, JsonGenerator jsonGenerator, SerializerProvider serializerProvider) throws IOException, JsonGenerationException { jsonGenerator.writeString(text.toString()); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/TransferPair.java000066400000000000000000000021651515451615200317240ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.ValueVector; /** Interface for copying values between a pair of two vectors of the same type. */ public interface TransferPair { void transfer(); void splitAndTransfer(int startIndex, int length); ValueVector getTo(); void copyValueSafe(int from, int to); } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/UuidUtility.java000066400000000000000000000052111515451615200316110ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/Validator.java000066400000000000000000000202411515451615200312440ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.Objects; import org.apache.arrow.vector.FieldVector; 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.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.Schema; /** Utility class for validating arrow data structures. */ public class Validator { /** * Validate two arrow schemas are equal. * * @param schema1 the 1st schema to compare * @param schema2 the 2nd schema to compare * @throws IllegalArgumentException if they are different. */ public static void compareSchemas(Schema schema1, Schema schema2) { if (!schema2.equals(schema1)) { throw new IllegalArgumentException("Different schemas:\n" + schema2 + "\n" + schema1); } } /** Validate two Dictionary encodings and dictionaries with id's from the encodings. */ public static void compareDictionaries( List encodings1, List encodings2, DictionaryProvider provider1, DictionaryProvider provider2) { if (encodings1.size() != encodings2.size()) { throw new IllegalArgumentException( "Different dictionary encoding count:\n" + encodings1.size() + "\n" + encodings2.size()); } for (int i = 0; i < encodings1.size(); i++) { if (!encodings1.get(i).equals(encodings2.get(i))) { throw new IllegalArgumentException( "Different dictionary encodings:\n" + encodings1.get(i) + "\n" + encodings2.get(i)); } long id = encodings1.get(i).getId(); Dictionary dict1 = provider1.lookup(id); Dictionary dict2 = provider2.lookup(id); if (dict1 == null || dict2 == null) { throw new IllegalArgumentException( "The DictionaryProvider did not contain the required " + "dictionary with id: " + id + "\n" + dict1 + "\n" + dict2); } try { compareFieldVectors(dict1.getVector(), dict2.getVector()); } catch (IllegalArgumentException e) { throw new IllegalArgumentException("Different dictionaries:\n" + dict1 + "\n" + dict2, e); } } } /** Validate two dictionary providers are equal in structure and contents. */ public static void compareDictionaryProviders( DictionaryProvider provider1, DictionaryProvider provider2) { List ids1 = new ArrayList(provider1.getDictionaryIds()); List ids2 = new ArrayList(provider2.getDictionaryIds()); java.util.Collections.sort(ids1); java.util.Collections.sort(ids2); if (!ids1.equals(ids2)) { throw new IllegalArgumentException( "Different ids in dictionary providers:\n" + ids1 + "\n" + ids2); } for (long id : ids1) { Dictionary dict1 = provider1.lookup(id); Dictionary dict2 = provider2.lookup(id); try { compareFieldVectors(dict1.getVector(), dict2.getVector()); } catch (IllegalArgumentException e) { throw new IllegalArgumentException("Different dictionaries:\n" + dict1 + "\n" + dict2, e); } } } /** * Validate two arrow vectorSchemaRoot are equal. * * @param root1 the 1st schema to compare * @param root2 the 2nd schema to compare * @throws IllegalArgumentException if they are different. */ public static void compareVectorSchemaRoot(VectorSchemaRoot root1, VectorSchemaRoot root2) { compareSchemas(root2.getSchema(), root1.getSchema()); if (root1.getRowCount() != root2.getRowCount()) { throw new IllegalArgumentException( "Different row count:\n" + root1.getRowCount() + " != " + root2.getRowCount()); } List vectors1 = root1.getFieldVectors(); List vectors2 = root2.getFieldVectors(); if (vectors1.size() != vectors2.size()) { throw new IllegalArgumentException( "Different column count:\n" + vectors1.toString() + "\n!=\n" + vectors2.toString()); } for (int i = 0; i < vectors1.size(); i++) { compareFieldVectors(vectors1.get(i), vectors2.get(i)); } } /** * Validate two arrow FieldVectors are equal. * * @param vector1 the 1st VectorField to compare * @param vector2 the 2nd VectorField to compare * @throws IllegalArgumentException if they are different */ public static void compareFieldVectors(FieldVector vector1, FieldVector vector2) { Field field1 = vector1.getField(); if (!field1.equals(vector2.getField())) { throw new IllegalArgumentException( "Different Fields:\n" + field1 + "\n!=\n" + vector2.getField()); } int valueCount = vector1.getValueCount(); if (valueCount != vector2.getValueCount()) { throw new IllegalArgumentException( "Different value count for field " + field1 + " : " + valueCount + " != " + vector2.getValueCount()); } for (int j = 0; j < valueCount; j++) { Object obj1 = vector1.getObject(j); Object obj2 = vector2.getObject(j); if (!equals(field1.getType(), obj1, obj2)) { throw new IllegalArgumentException( "Different values in column:\n" + field1 + " at index " + j + ": " + obj1 + " != " + obj2); } } } static boolean equals(ArrowType type, final Object o1, final Object o2) { if (type instanceof ArrowType.FloatingPoint) { ArrowType.FloatingPoint fpType = (ArrowType.FloatingPoint) type; switch (fpType.getPrecision()) { case DOUBLE: return equalEnough((Double) o1, (Double) o2); case SINGLE: return equalEnough((Float) o1, (Float) o2); case HALF: default: throw new UnsupportedOperationException("unsupported precision: " + fpType); } } else if (type instanceof ArrowType.Binary || type instanceof ArrowType.LargeBinary || type instanceof ArrowType.FixedSizeBinary) { return Arrays.equals((byte[]) o1, (byte[]) o2); } else if (o1 instanceof byte[] && o2 instanceof byte[]) { return Arrays.equals((byte[]) o1, (byte[]) o2); } return Objects.equals(o1, o2); } static boolean equalEnough(Float f1, Float f2) { if (f1 == null || f2 == null) { return f1 == null && f2 == null; } if (f1.isNaN()) { return f2.isNaN(); } if (f1.isInfinite()) { return f2.isInfinite() && Math.signum(f1) == Math.signum(f2); } float average = Math.abs((f1 + f2) / 2); float differenceScaled = Math.abs(f1 - f2) / (average == 0.0f ? 1f : average); return differenceScaled < 1.0E-6f; } static boolean equalEnough(Double f1, Double f2) { if (f1 == null || f2 == null) { return f1 == null && f2 == null; } if (f1.isNaN()) { return f2.isNaN(); } if (f1.isInfinite()) { return f2.isInfinite() && Math.signum(f1) == Math.signum(f2); } double average = Math.abs((f1 + f2) / 2); double differenceScaled = Math.abs(f1 - f2) / (average == 0.0d ? 1d : average); return differenceScaled < 1.0E-12d; } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/ValueVectorUtility.java000066400000000000000000000143441515451615200331510ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.validate.ValidateUtil.validateOrThrow; import java.util.function.BiFunction; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.validate.ValidateVectorBufferVisitor; import org.apache.arrow.vector.validate.ValidateVectorDataVisitor; import org.apache.arrow.vector.validate.ValidateVectorTypeVisitor; /** Utility methods for {@link ValueVector}. */ public class ValueVectorUtility { private ValueVectorUtility() {} /** * Get the toString() representation of vector suitable for debugging. Note since vectors may have * millions of values, this method only shows max 20 values. Examples as below (v represents * value): *

  • vector with 0 value: [] *
  • vector with 5 values (no more than 20 values): [v0, v1, v2, v3, v4] *
  • vector with 100 values (more than 20 values): [v0, v1, v2, v3, v4, v5, v6, v7, v8, v9, ..., * v90, v91, v92, v93, v94, v95, v96, v97, v98, v99] */ public static String getToString(V vector, int start, int end) { return getToString(vector, start, end, (v, i) -> v.getObject(i)); } /** * Get the toString() representation of vector suitable for debugging. Note since vectors may have * millions of values, this method only shows at most 20 values. * * @param vector the vector for which to get toString representation. * @param start the starting index, inclusive. * @param end the end index, exclusive. * @param valueToString the function to transform individual elements to strings. */ public static String getToString( V vector, int start, int end, BiFunction valueToString) { Preconditions.checkNotNull(vector); final int length = end - start; Preconditions.checkArgument(length >= 0); Preconditions.checkArgument(start >= 0); Preconditions.checkArgument(end <= vector.getValueCount()); if (length == 0) { return "[]"; } final int window = 10; boolean skipComma = false; StringBuilder sb = new StringBuilder(); sb.append('['); for (int i = start; i < end; i++) { if (skipComma) { skipComma = false; } if (i - start >= window && i < end - window) { sb.append("..."); i = end - window - 1; skipComma = true; } else { sb.append(valueToString.apply(vector, i)); } if (i == end - 1) { sb.append(']'); } else { if (!skipComma) { sb.append(','); } sb.append(' '); } } return sb.toString(); } /** Utility to validate vector in O(1) time. */ public static void validate(ValueVector vector) { Preconditions.checkNotNull(vector); ValidateVectorTypeVisitor typeVisitor = new ValidateVectorTypeVisitor(); vector.accept(typeVisitor, null); ValidateVectorBufferVisitor bufferVisitor = new ValidateVectorBufferVisitor(); vector.accept(bufferVisitor, null); } /** Utility to validate vector in O(n) time, where n is the value count. */ public static void validateFull(ValueVector vector) { validate(vector); ValidateVectorDataVisitor dataVisitor = new ValidateVectorDataVisitor(); vector.accept(dataVisitor, null); } /** Utility to validate vector schema root in O(1) time. */ public static void validate(VectorSchemaRoot root) { Preconditions.checkNotNull(root); int valueCount = root.getRowCount(); validateOrThrow( valueCount >= 0, "The row count of vector schema root %s is negative.", valueCount); for (ValueVector childVec : root.getFieldVectors()) { validateOrThrow( valueCount == childVec.getValueCount(), "Child vector and vector schema root have different value counts. " + "Child vector value count %s, vector schema root value count %s", childVec.getValueCount(), valueCount); validate(childVec); } } /** Utility to validate vector in O(n) time, where n is the value count. */ public static void validateFull(VectorSchemaRoot root) { Preconditions.checkNotNull(root); int valueCount = root.getRowCount(); validateOrThrow( valueCount >= 0, "The row count of vector schema root %s is negative.", valueCount); for (ValueVector childVec : root.getFieldVectors()) { validateOrThrow( valueCount == childVec.getValueCount(), "Child vector and vector schema root have different value counts. " + "Child vector value count %s, vector schema root value count %s", childVec.getValueCount(), valueCount); validateFull(childVec); } } /** Pre allocate memory for BaseFixedWidthVector. */ public static void preAllocate(VectorSchemaRoot root, int targetSize) { for (ValueVector vector : root.getFieldVectors()) { if (vector instanceof BaseFixedWidthVector) { ((BaseFixedWidthVector) vector).allocateNew(targetSize); } } } /** Ensure capacity for BaseFixedWidthVector. */ public static void ensureCapacity(VectorSchemaRoot root, int targetCapacity) { for (ValueVector vector : root.getFieldVectors()) { if (vector instanceof BaseFixedWidthVector) { while (vector.getValueCapacity() < targetCapacity) { vector.reAlloc(); } } } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java000066400000000000000000000755241515451615200322560ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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; import org.apache.arrow.vector.complex.DenseUnionVector; 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.NonNullableStructVector; import org.apache.arrow.vector.complex.RunEndEncodedVector; import org.apache.arrow.vector.complex.UnionVector; /** Utility to append two vectors together. */ public class VectorAppender implements VectorVisitor { /** The targetVector to be appended. */ private final ValueVector targetVector; private final TypeEqualsVisitor typeVisitor; /** * Constructs a new targetVector appender, with the given targetVector. * * @param targetVector the targetVector to be appended. */ public VectorAppender(ValueVector targetVector) { this.targetVector = targetVector; typeVisitor = new TypeEqualsVisitor(targetVector, false, true); } @Override public ValueVector visit(BaseFixedWidthVector deltaVector, Void value) { Preconditions.checkArgument( targetVector.getField().getType().equals(deltaVector.getField().getType()), "The targetVector to append must have the same type as the targetVector being appended"); if (deltaVector.getValueCount() == 0) { return targetVector; // optimization, nothing to append, return } int newValueCount = targetVector.getValueCount() + deltaVector.getValueCount(); // make sure there is enough capacity while (targetVector.getValueCapacity() < newValueCount) { targetVector.reAlloc(); } // append validity buffer BitVectorHelper.concatBits( targetVector.getValidityBuffer(), targetVector.getValueCount(), deltaVector.getValidityBuffer(), deltaVector.getValueCount(), targetVector.getValidityBuffer()); // append data buffer if (targetVector instanceof BitVector) { // special processing for bit vector, as its type width is 0 BitVectorHelper.concatBits( targetVector.getDataBuffer(), targetVector.getValueCount(), deltaVector.getDataBuffer(), deltaVector.getValueCount(), targetVector.getDataBuffer()); } else { MemoryUtil.copyMemory( deltaVector.getDataBuffer().memoryAddress(), targetVector.getDataBuffer().memoryAddress() + deltaVector.getTypeWidth() * targetVector.getValueCount(), deltaVector.getTypeWidth() * deltaVector.getValueCount()); } targetVector.setValueCount(newValueCount); return targetVector; } @Override public ValueVector visit(BaseVariableWidthVector deltaVector, Void value) { Preconditions.checkArgument( targetVector.getField().getType().equals(deltaVector.getField().getType()), "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 newValueCount = targetVector.getValueCount() + deltaVector.getValueCount(); int targetDataSize = targetVector .getOffsetBuffer() .getInt((long) targetVector.getValueCount() * BaseVariableWidthVector.OFFSET_WIDTH); int deltaDataSize = deltaVector .getOffsetBuffer() .getInt((long) deltaVector.getValueCount() * BaseVariableWidthVector.OFFSET_WIDTH); int newValueCapacity = targetDataSize + deltaDataSize; // make sure there is enough capacity while (targetVector.getValueCapacity() < newValueCount) { ((BaseVariableWidthVector) targetVector).reallocValidityAndOffsetBuffers(); } while (targetVector.getDataBuffer().capacity() < newValueCapacity) { ((BaseVariableWidthVector) targetVector).reallocDataBuffer(); } // append validity buffer BitVectorHelper.concatBits( targetVector.getValidityBuffer(), targetVector.getValueCount(), deltaVector.getValidityBuffer(), deltaVector.getValueCount(), targetVector.getValidityBuffer()); // append data buffer MemoryUtil.copyMemory( deltaVector.getDataBuffer().memoryAddress(), targetVector.getDataBuffer().memoryAddress() + targetDataSize, deltaDataSize); // copy offset buffer MemoryUtil.copyMemory( deltaVector.getOffsetBuffer().memoryAddress() + BaseVariableWidthVector.OFFSET_WIDTH, targetVector.getOffsetBuffer().memoryAddress() + (targetVector.getValueCount() + 1) * BaseVariableWidthVector.OFFSET_WIDTH, deltaVector.getValueCount() * BaseVariableWidthVector.OFFSET_WIDTH); // increase each offset from the second buffer for (int i = 0; i < deltaVector.getValueCount(); i++) { int oldOffset = targetVector .getOffsetBuffer() .getInt( (long) (targetVector.getValueCount() + 1 + i) * BaseVariableWidthVector.OFFSET_WIDTH); targetVector .getOffsetBuffer() .setInt( (long) (targetVector.getValueCount() + 1 + i) * BaseVariableWidthVector.OFFSET_WIDTH, oldOffset + targetDataSize); } ((BaseVariableWidthVector) targetVector).setLastSet(newValueCount - 1); targetVector.setValueCount(newValueCount); return targetVector; } @Override public ValueVector visit(BaseLargeVariableWidthVector deltaVector, Void value) { Preconditions.checkArgument( targetVector.getField().getType().equals(deltaVector.getField().getType()), "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 newValueCount = targetVector.getValueCount() + deltaVector.getValueCount(); long targetDataSize = targetVector .getOffsetBuffer() .getLong( (long) targetVector.getValueCount() * BaseLargeVariableWidthVector.OFFSET_WIDTH); long deltaDataSize = deltaVector .getOffsetBuffer() .getLong( (long) deltaVector.getValueCount() * BaseLargeVariableWidthVector.OFFSET_WIDTH); long newValueCapacity = targetDataSize + deltaDataSize; // make sure there is enough capacity while (targetVector.getValueCapacity() < newValueCount) { ((BaseLargeVariableWidthVector) targetVector).reallocValidityAndOffsetBuffers(); } while (targetVector.getDataBuffer().capacity() < newValueCapacity) { ((BaseLargeVariableWidthVector) targetVector).reallocDataBuffer(); } // append validity buffer BitVectorHelper.concatBits( targetVector.getValidityBuffer(), targetVector.getValueCount(), deltaVector.getValidityBuffer(), deltaVector.getValueCount(), targetVector.getValidityBuffer()); // append data buffer MemoryUtil.copyMemory( deltaVector.getDataBuffer().memoryAddress(), targetVector.getDataBuffer().memoryAddress() + targetDataSize, deltaDataSize); // copy offset buffer MemoryUtil.copyMemory( deltaVector.getOffsetBuffer().memoryAddress() + BaseLargeVariableWidthVector.OFFSET_WIDTH, targetVector.getOffsetBuffer().memoryAddress() + (targetVector.getValueCount() + 1) * BaseLargeVariableWidthVector.OFFSET_WIDTH, deltaVector.getValueCount() * BaseLargeVariableWidthVector.OFFSET_WIDTH); // increase each offset from the second buffer for (int i = 0; i < deltaVector.getValueCount(); i++) { long oldOffset = targetVector .getOffsetBuffer() .getLong( (long) (targetVector.getValueCount() + 1 + i) * BaseLargeVariableWidthVector.OFFSET_WIDTH); targetVector .getOffsetBuffer() .setLong( (long) (targetVector.getValueCount() + 1 + i) * BaseLargeVariableWidthVector.OFFSET_WIDTH, oldOffset + targetDataSize); } ((BaseLargeVariableWidthVector) targetVector).setLastSet(newValueCount - 1); targetVector.setValueCount(newValueCount); return targetVector; } @Override 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 public ValueVector visit(ListVector 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 newValueCount = targetVector.getValueCount() + deltaVector.getValueCount(); int targetListSize = targetVector .getOffsetBuffer() .getInt((long) targetVector.getValueCount() * ListVector.OFFSET_WIDTH); int deltaListSize = deltaVector .getOffsetBuffer() .getInt((long) deltaVector.getValueCount() * ListVector.OFFSET_WIDTH); ListVector targetListVector = (ListVector) targetVector; // make sure the underlying vector has value count set targetListVector.getDataVector().setValueCount(targetListSize); deltaVector.getDataVector().setValueCount(deltaListSize); // make sure there is enough capacity while (targetVector.getValueCapacity() < newValueCount) { targetVector.reAlloc(); } // append validity buffer BitVectorHelper.concatBits( targetVector.getValidityBuffer(), targetVector.getValueCount(), deltaVector.getValidityBuffer(), deltaVector.getValueCount(), targetVector.getValidityBuffer()); // append offset buffer MemoryUtil.copyMemory( deltaVector.getOffsetBuffer().memoryAddress() + ListVector.OFFSET_WIDTH, targetVector.getOffsetBuffer().memoryAddress() + (targetVector.getValueCount() + 1) * ListVector.OFFSET_WIDTH, (long) deltaVector.getValueCount() * ListVector.OFFSET_WIDTH); // increase each offset from the second buffer for (int i = 0; i < deltaVector.getValueCount(); i++) { int oldOffset = targetVector .getOffsetBuffer() .getInt((long) (targetVector.getValueCount() + 1 + i) * ListVector.OFFSET_WIDTH); targetVector .getOffsetBuffer() .setInt( (long) (targetVector.getValueCount() + 1 + i) * ListVector.OFFSET_WIDTH, oldOffset + targetListSize); } targetListVector.setLastSet(newValueCount - 1); // append underlying vectors VectorAppender innerAppender = new VectorAppender(targetListVector.getDataVector()); deltaVector.getDataVector().accept(innerAppender, null); targetVector.setValueCount(newValueCount); return targetVector; } @Override public ValueVector visit(LargeListVector 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 newValueCount = targetVector.getValueCount() + deltaVector.getValueCount(); long targetListSize = targetVector .getOffsetBuffer() .getLong((long) targetVector.getValueCount() * LargeListVector.OFFSET_WIDTH); long deltaListSize = deltaVector .getOffsetBuffer() .getLong((long) deltaVector.getValueCount() * LargeListVector.OFFSET_WIDTH); ListVector targetListVector = (ListVector) 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)); // make sure there is enough capacity while (targetVector.getValueCapacity() < newValueCount) { targetVector.reAlloc(); } // append validity buffer BitVectorHelper.concatBits( targetVector.getValidityBuffer(), targetVector.getValueCount(), deltaVector.getValidityBuffer(), deltaVector.getValueCount(), targetVector.getValidityBuffer()); // append offset buffer MemoryUtil.copyMemory( deltaVector.getOffsetBuffer().memoryAddress() + ListVector.OFFSET_WIDTH, targetVector.getOffsetBuffer().memoryAddress() + (targetVector.getValueCount() + 1) * LargeListVector.OFFSET_WIDTH, (long) deltaVector.getValueCount() * ListVector.OFFSET_WIDTH); // increase each offset from the second buffer for (int i = 0; i < deltaVector.getValueCount(); i++) { long oldOffset = targetVector .getOffsetBuffer() .getLong( (long) (targetVector.getValueCount() + 1 + i) * LargeListVector.OFFSET_WIDTH); targetVector .getOffsetBuffer() .setLong( (long) (targetVector.getValueCount() + 1 + i) * LargeListVector.OFFSET_WIDTH, oldOffset + targetListSize); } targetListVector.setLastSet(newValueCount - 1); // append underlying vectors VectorAppender innerAppender = new VectorAppender(targetListVector.getDataVector()); deltaVector.getDataVector().accept(innerAppender, null); targetVector.setValueCount(newValueCount); return targetVector; } @Override public ValueVector visit(FixedSizeListVector deltaVector, Void value) { Preconditions.checkArgument( typeVisitor.equals(deltaVector), "The vector to append must have the same type as the targetVector being appended"); if (deltaVector.getValueCount() == 0) { return targetVector; // optimization, nothing to append, return } FixedSizeListVector targetListVector = (FixedSizeListVector) targetVector; Preconditions.checkArgument( targetListVector.getListSize() == deltaVector.getListSize(), "FixedSizeListVector must have the same list size to append"); int newValueCount = targetVector.getValueCount() + deltaVector.getValueCount(); int targetListSize = targetListVector.getValueCount() * targetListVector.getListSize(); int deltaListSize = deltaVector.getValueCount() * deltaVector.getListSize(); // make sure the underlying vector has value count set targetListVector.getDataVector().setValueCount(targetListSize); deltaVector.getDataVector().setValueCount(deltaListSize); // make sure there is enough capacity while (targetVector.getValueCapacity() < newValueCount) { targetVector.reAlloc(); } // append validity buffer BitVectorHelper.concatBits( targetVector.getValidityBuffer(), targetVector.getValueCount(), deltaVector.getValidityBuffer(), deltaVector.getValueCount(), targetVector.getValidityBuffer()); // append underlying vectors VectorAppender innerAppender = new VectorAppender(targetListVector.getDataVector()); deltaVector.getDataVector().accept(innerAppender, null); targetVector.setValueCount(newValueCount); return targetVector; } @Override public ValueVector visit(NonNullableStructVector deltaVector, Void value) { Preconditions.checkArgument( typeVisitor.equals(deltaVector), "The vector to append must have the same type as the targetVector being appended"); if (deltaVector.getValueCount() == 0) { return targetVector; // optimization, nothing to append, return } NonNullableStructVector targetStructVector = (NonNullableStructVector) targetVector; int newValueCount = targetVector.getValueCount() + deltaVector.getValueCount(); // make sure there is enough capacity while (targetVector.getValueCapacity() < newValueCount) { targetVector.reAlloc(); } // append validity buffer BitVectorHelper.concatBits( targetVector.getValidityBuffer(), targetVector.getValueCount(), deltaVector.getValidityBuffer(), deltaVector.getValueCount(), targetVector.getValidityBuffer()); // append child vectors for (int i = 0; i < targetStructVector.getChildrenFromFields().size(); i++) { ValueVector targetChild = targetStructVector.getVectorById(i); ValueVector deltaChild = deltaVector.getVectorById(i); targetChild.setValueCount(targetStructVector.getValueCount()); deltaChild.setValueCount(deltaVector.getValueCount()); VectorAppender innerAppender = new VectorAppender(targetChild); deltaChild.accept(innerAppender, null); } targetVector.setValueCount(newValueCount); return targetVector; } @Override public ValueVector visit(UnionVector deltaVector, Void value) { // we only make sure that both vectors are union vectors. Preconditions.checkArgument( targetVector.getMinorType() == deltaVector.getMinorType(), "The vector to append must have the same type as the targetVector being appended"); if (deltaVector.getValueCount() == 0) { return targetVector; // optimization, nothing to append, return } UnionVector targetUnionVector = (UnionVector) targetVector; int newValueCount = targetVector.getValueCount() + deltaVector.getValueCount(); // make sure there is enough capacity while (targetUnionVector.getValueCapacity() < newValueCount) { targetUnionVector.reAlloc(); } // append type buffers MemoryUtil.copyMemory( deltaVector.getTypeBufferAddress(), targetUnionVector.getTypeBufferAddress() + targetVector.getValueCount(), deltaVector.getValueCount()); // build the hash set for all types HashSet targetTypes = new HashSet<>(); for (int i = 0; i < targetUnionVector.getValueCount(); i++) { targetTypes.add(targetUnionVector.getTypeValue(i)); } HashSet deltaTypes = new HashSet<>(); for (int i = 0; i < deltaVector.getValueCount(); i++) { deltaTypes.add(deltaVector.getTypeValue(i)); } // append child vectors for (int i = 0; i < Byte.MAX_VALUE; i++) { if (targetTypes.contains(i) || deltaTypes.contains(i)) { ValueVector targetChild = targetUnionVector.getVectorByType(i); if (!targetTypes.contains(i)) { // if the vector type does not exist in the target, it must be newly created // and we must make sure it has enough capacity. while (targetChild.getValueCapacity() < newValueCount) { targetChild.reAlloc(); } } if (deltaTypes.contains(i)) { // append child vectors ValueVector deltaChild = deltaVector.getVectorByType(i); targetChild.setValueCount(targetUnionVector.getValueCount()); deltaChild.setValueCount(deltaVector.getValueCount()); VectorAppender innerAppender = new VectorAppender(targetChild); deltaChild.accept(innerAppender, null); } targetChild.setValueCount(newValueCount); } } targetVector.setValueCount(newValueCount); return targetVector; } @Override public ValueVector visit(DenseUnionVector deltaVector, Void value) { // we only make sure that both vectors are union vectors. Preconditions.checkArgument( targetVector.getMinorType() == deltaVector.getMinorType(), "The vector to append must have the same type as the targetVector being appended"); if (deltaVector.getValueCount() == 0) { return targetVector; // optimization, nothing to append, return } DenseUnionVector targetDenseUnionVector = (DenseUnionVector) targetVector; int newValueCount = targetVector.getValueCount() + deltaVector.getValueCount(); // make sure there is enough capacity while (targetDenseUnionVector.getValueCapacity() < newValueCount) { targetDenseUnionVector.reAlloc(); } // append type buffers MemoryUtil.copyMemory( deltaVector.getTypeBuffer().memoryAddress(), targetDenseUnionVector.getTypeBuffer().memoryAddress() + targetVector.getValueCount(), deltaVector.getValueCount()); // append offset buffers for (int i = 0; i < deltaVector.getValueCount(); i++) { byte typeId = deltaVector.getTypeId(i); ValueVector targetChildVector = targetDenseUnionVector.getVectorByType(typeId); int offsetBase = targetChildVector == null ? 0 : targetChildVector.getValueCount(); int deltaOffset = deltaVector.getOffset(i); long index = (long) (targetVector.getValueCount() + i) * DenseUnionVector.OFFSET_WIDTH; targetVector.getOffsetBuffer().setInt(index, offsetBase + deltaOffset); } // append child vectors for (int i = 0; i <= Byte.MAX_VALUE; i++) { ValueVector targetChildVector = targetDenseUnionVector.getVectorByType((byte) i); ValueVector deltaChildVector = deltaVector.getVectorByType((byte) i); if (targetChildVector == null && deltaChildVector == null) { // the type id is not registered in either vector, we are done. continue; } else if (targetChildVector == null && deltaChildVector != null) { // first register a new child in the target vector targetDenseUnionVector.registerNewTypeId(deltaChildVector.getField()); targetChildVector = targetDenseUnionVector.addVector( (byte) i, deltaChildVector.getField().createVector(targetDenseUnionVector.getAllocator())); // now we have both child vectors not null, we can append them. VectorAppender childAppender = new VectorAppender(targetChildVector); deltaChildVector.accept(childAppender, null); } else if (targetChildVector != null && deltaChildVector == null) { // the value only exists in the target vector, so we are done continue; } else { // both child vectors are non-null // first check vector types TypeEqualsVisitor childTypeVisitor = new TypeEqualsVisitor( targetChildVector, /* check name */ false, /* check meta data*/ false); if (!childTypeVisitor.equals(deltaChildVector)) { throw new IllegalArgumentException( "dense union vectors have different child vector types with type id " + i); } // append child vectors VectorAppender childAppender = new VectorAppender(targetChildVector); deltaChildVector.accept(childAppender, null); } } targetVector.setValueCount(newValueCount); return targetVector; } @Override public ValueVector visit(NullVector deltaVector, Void value) { Preconditions.checkArgument( targetVector.getField().getType().equals(deltaVector.getField().getType()), "The targetVector to append must have the same type as the targetVector being appended"); return targetVector; } @Override public ValueVector visit(ExtensionTypeVector deltaVector, Void value) { ValueVector targetUnderlying = ((ExtensionTypeVector) targetVector).getUnderlyingVector(); VectorAppender underlyingAppender = new VectorAppender(targetUnderlying); 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."); } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/VectorBatchAppender.java000066400000000000000000000025601515451615200332060ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.ValueVector; /** Utility to add vector values in batch. */ public class VectorBatchAppender { /** * Add value vectors in batch. * * @param targetVector the target vector. * @param vectorsToAppend the vectors to append. * @param the vector type. */ public static void batchAppend(V targetVector, V... vectorsToAppend) { VectorAppender appender = new VectorAppender(targetVector); for (V delta : vectorsToAppend) { delta.accept(appender, null); } } } VectorSchemaRootAppender.java000066400000000000000000000067271515451615200341630ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.compare.TypeEqualsVisitor; /** Utility to append {@link org.apache.arrow.vector.VectorSchemaRoot}s with the same schema. */ public class VectorSchemaRootAppender { /** * Appends a number of {@link VectorSchemaRoot}s. * * @param checkSchema if we need to check schema for the vector schema roots. * @param targetRoot the vector schema root to be appended. * @param rootsToAppend the vector schema roots to append. * @throws IllegalArgumentException throws if we need to check schema, and checking schema fails. */ public static void append( boolean checkSchema, VectorSchemaRoot targetRoot, VectorSchemaRoot... rootsToAppend) { // create appenders VectorAppender[] appenders = new VectorAppender[targetRoot.getFieldVectors().size()]; for (int i = 0; i < appenders.length; i++) { appenders[i] = new VectorAppender(targetRoot.getVector(i)); } // create type checkers, if necessary TypeEqualsVisitor[] typeCheckers = null; if (checkSchema) { typeCheckers = new TypeEqualsVisitor[targetRoot.getFieldVectors().size()]; for (int i = 0; i < typeCheckers.length; i++) { typeCheckers[i] = new TypeEqualsVisitor( targetRoot.getVector(i), /* check name */ false, /* check meta data */ false); } } for (VectorSchemaRoot delta : rootsToAppend) { // check schema, if necessary if (checkSchema) { if (delta.getFieldVectors().size() != targetRoot.getFieldVectors().size()) { throw new IllegalArgumentException( "Vector schema roots have different numbers of child vectors."); } for (int i = 0; i < typeCheckers.length; i++) { if (!typeCheckers[i].equals(delta.getVector(i))) { throw new IllegalArgumentException("Vector schema roots have different schemas."); } } } // append child vectors. for (int i = 0; i < appenders.length; i++) { delta.getVector(i).accept(appenders[i], null); } targetRoot.setRowCount(targetRoot.getRowCount() + delta.getRowCount()); } } /** * Appends a number of {@link VectorSchemaRoot}s. This method performs schema checking before * appending data. * * @param targetRoot the vector schema root to be appended. * @param rootsToAppend the vector schema roots to append. * @throws IllegalArgumentException throws if we need to check schema, and checking schema fails. */ public static void append(VectorSchemaRoot targetRoot, VectorSchemaRoot... rootsToAppend) { append(true, targetRoot, rootsToAppend); } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/validate/000077500000000000000000000000001515451615200272715ustar00rootroot00000000000000MetadataV4UnionChecker.java000066400000000000000000000054511515451615200343120ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/validate/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.validate; import java.io.IOException; import java.util.Iterator; import org.apache.arrow.vector.types.MetadataVersion; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; /** * Given a field, checks that no Union fields are present. * *

    This is intended to be used to prevent unions from being read/written with V4 metadata. */ public final class MetadataV4UnionChecker { static boolean isUnion(Field field) { return field.getType().getTypeID() == ArrowType.ArrowTypeID.Union; } static Field check(Field field) { if (isUnion(field)) { return field; } // Naive recursive DFS for (final Field child : field.getChildren()) { final Field result = check(child); if (result != null) { return result; } } return null; } /** * Check the schema, raising an error if an unsupported feature is used (e.g. unions with < V5 * metadata). */ public static void checkForUnion(Iterator fields, MetadataVersion metadataVersion) { if (metadataVersion.toFlatbufID() >= MetadataVersion.V5.toFlatbufID()) { return; } while (fields.hasNext()) { Field union = check(fields.next()); if (union != null) { throw new IllegalArgumentException( "Cannot write union with V4 metadata version, use V5 instead. Found field: " + union); } } } /** * Check the schema, raising an error if an unsupported feature is used (e.g. unions with < V5 * metadata). */ public static void checkRead(Schema schema, MetadataVersion metadataVersion) throws IOException { if (metadataVersion.toFlatbufID() >= MetadataVersion.V5.toFlatbufID()) { return; } for (final Field field : schema.getFields()) { Field union = check(field); if (union != null) { throw new IOException("Cannot read union with V4 metadata version. Found field: " + union); } } } } apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/validate/ValidateUtil.java000066400000000000000000000037201515451615200325250ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.validate; /** Utilities for vector validation. */ public class ValidateUtil { private ValidateUtil() {} /** * Validate the expression. * * @param expression the expression to validate. * @param errorMessage the error message. * @throws ValidateException if the expression evaluates to false. */ public static void validateOrThrow(boolean expression, String errorMessage) { if (!expression) { throw new ValidateException(errorMessage); } } /** * Validate the expression. * * @param expression the expression to validate. * @param errorMessage the error message template. * @param args the error message arguments. * @throws ValidateException if the expression evaluates to false. */ public static void validateOrThrow(boolean expression, String errorMessage, Object... args) { if (!expression) { throw new ValidateException(String.format(errorMessage, args)); } } /** A exception that is thrown when the vector validation fails. */ public static class ValidateException extends RuntimeException { public ValidateException(String message) { super(message); } } } ValidateVectorBufferVisitor.java000066400000000000000000000314551515451615200355130ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/validate/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.validate; import static org.apache.arrow.vector.validate.ValidateUtil.validateOrThrow; import org.apache.arrow.memory.ArrowBuf; 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.BitVector; import org.apache.arrow.vector.ExtensionTypeVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.NullVector; import org.apache.arrow.vector.TypeLayout; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.DenseUnionVector; 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.NonNullableStructVector; import org.apache.arrow.vector.complex.RunEndEncodedVector; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.types.pojo.ArrowType; /** Visitor to validate vector buffers. */ public class ValidateVectorBufferVisitor implements VectorVisitor { private void validateVectorCommon(ValueVector vector) { ArrowType arrowType = vector.getField().getType(); validateOrThrow( vector.getValueCount() >= 0, "Vector valueCount %s is negative.", vector.getValueCapacity()); if (vector instanceof FieldVector) { FieldVector fieldVector = (FieldVector) vector; int typeBufferCount = TypeLayout.getTypeBufferCount(arrowType); 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()); } } } private void validateValidityBuffer(ValueVector vector, int valueCount) { ArrowBuf validityBuffer = vector.getValidityBuffer(); validateOrThrow(validityBuffer != null, "The validity buffer is null."); validateOrThrow( validityBuffer.capacity() * 8 >= valueCount, "Not enough capacity for the validity buffer. Minimum capacity %s, actual capacity %s.", (valueCount + 7) / 8, validityBuffer.capacity()); } private void validateOffsetBuffer(ValueVector vector, long minCapacity) { ArrowBuf offsetBuffer = vector.getOffsetBuffer(); validateOrThrow(offsetBuffer != null, "The offset buffer is null."); validateOrThrow( offsetBuffer.capacity() >= minCapacity, "Not enough capacity for the offset buffer. Minimum capacity %s, actual capacity %s.", minCapacity, offsetBuffer.capacity()); } private void validateFixedWidthDataBuffer(ValueVector vector, int valueCount, int bitWidth) { ArrowBuf dataBuffer = vector.getDataBuffer(); validateOrThrow(dataBuffer != null, "The fixed width data buffer is null."); validateOrThrow( (long) bitWidth * valueCount <= dataBuffer.capacity() * 8L, "Not enough capacity for fixed width data buffer. Minimum capacity %s, actual capacity %s.", ((long) bitWidth * valueCount + 7L) / 8L, dataBuffer.capacity()); } private void validateDataBuffer(ValueVector vector, long minCapacity) { ArrowBuf dataBuffer = vector.getDataBuffer(); validateOrThrow(dataBuffer != null, "The data buffer is null."); validateOrThrow( dataBuffer.capacity() >= minCapacity, "Not enough capacity for data buffer. Minimum capacity %s, actual capacity %s.", minCapacity, dataBuffer.capacity()); } private void validateTypeBuffer(ArrowBuf typeBuf, long minCapacity) { validateOrThrow(typeBuf != null, "The type buffer is null."); validateOrThrow( typeBuf.capacity() >= minCapacity, "Not enough capacity for type buffer. Minimum capacity %s, actual capacity %s.", minCapacity, typeBuf.capacity()); } @Override public Void visit(BaseFixedWidthVector vector, Void value) { int bitWidth = (vector instanceof BitVector) ? 1 : vector.getTypeWidth() * 8; int valueCount = vector.getValueCount(); validateVectorCommon(vector); validateValidityBuffer(vector, valueCount); validateFixedWidthDataBuffer(vector, valueCount, bitWidth); return null; } @Override public Void visit(BaseVariableWidthVector vector, Void value) { int valueCount = vector.getValueCount(); validateVectorCommon(vector); validateValidityBuffer(vector, valueCount); long minOffsetCapacity = valueCount == 0 ? 0L : (long) (valueCount + 1) * BaseVariableWidthVector.OFFSET_WIDTH; validateOffsetBuffer(vector, minOffsetCapacity); int lastOffset = valueCount == 0 ? 0 : vector.getOffsetBuffer().getInt(valueCount * BaseVariableWidthVector.OFFSET_WIDTH); validateDataBuffer(vector, lastOffset); return null; } @Override public Void visit(BaseLargeVariableWidthVector vector, Void value) { int valueCount = vector.getValueCount(); validateVectorCommon(vector); validateValidityBuffer(vector, valueCount); long minOffsetCapacity = valueCount == 0 ? 0L : (long) (valueCount + 1) * BaseLargeVariableWidthVector.OFFSET_WIDTH; validateOffsetBuffer(vector, minOffsetCapacity); long lastOffset = valueCount == 0 ? 0L : vector .getOffsetBuffer() .getLong((long) valueCount * BaseLargeVariableWidthVector.OFFSET_WIDTH); validateDataBuffer(vector, lastOffset); return null; } @Override public Void visit(BaseVariableWidthViewVector vector, Void value) { 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 public Void visit(ListVector vector, Void value) { int valueCount = vector.getValueCount(); validateVectorCommon(vector); validateValidityBuffer(vector, valueCount); long minOffsetCapacity = valueCount == 0 ? 0L : (long) (valueCount + 1) * ListVector.OFFSET_WIDTH; validateOffsetBuffer(vector, minOffsetCapacity); FieldVector dataVector = vector.getDataVector(); int lastOffset = valueCount == 0 ? 0 : vector.getOffsetBuffer().getInt(valueCount * BaseVariableWidthVector.OFFSET_WIDTH); int dataVectorLength = dataVector == null ? 0 : dataVector.getValueCount(); validateOrThrow( dataVectorLength >= lastOffset, "Inner vector does not contain enough elements. Minimum element count %s, actual element count %s", lastOffset + 1, dataVectorLength); if (dataVector != null) { dataVector.accept(this, null); } return null; } @Override public Void visit(FixedSizeListVector vector, Void value) { int valueCount = vector.getValueCount(); validateVectorCommon(vector); validateValidityBuffer(vector, valueCount); FieldVector dataVector = vector.getDataVector(); int dataVectorLength = dataVector == null ? 0 : dataVector.getValueCount(); validateOrThrow( dataVectorLength >= valueCount * vector.getListSize(), "Inner vector does not contain enough elements. Minimum element count %s, actual element count %s.", valueCount * vector.getListSize(), dataVectorLength); if (dataVector != null) { dataVector.accept(this, null); } return null; } @Override public Void visit(LargeListVector vector, Void value) { int valueCount = vector.getValueCount(); validateVectorCommon(vector); validateValidityBuffer(vector, valueCount); long minOffsetCapacity = valueCount == 0 ? 0L : (long) (valueCount + 1) * LargeListVector.OFFSET_WIDTH; validateOffsetBuffer(vector, minOffsetCapacity); FieldVector dataVector = vector.getDataVector(); long lastOffset = valueCount == 0 ? 0 : vector .getOffsetBuffer() .getLong(valueCount * BaseLargeVariableWidthVector.OFFSET_WIDTH); int dataVectorLength = dataVector == null ? 0 : dataVector.getValueCount(); validateOrThrow( dataVectorLength >= lastOffset, "Inner vector does not contain enough elements. Minimum element count %s, actual element count %s", lastOffset + 1, dataVectorLength); if (dataVector != null) { dataVector.accept(this, null); } return null; } @Override public Void visit(NonNullableStructVector vector, Void value) { int valueCount = vector.getValueCount(); validateVectorCommon(vector); validateValidityBuffer(vector, valueCount); for (ValueVector subVector : vector.getChildrenFromFields()) { validateOrThrow( valueCount == subVector.getValueCount(), "Struct vector length not equal to child vector length. Struct vector length %s, child vector length %s", valueCount, subVector.getValueCount()); subVector.accept(this, null); } return null; } @Override public Void visit(UnionVector vector, Void value) { int valueCount = vector.getValueCount(); validateVectorCommon(vector); validateTypeBuffer(vector.getTypeBuffer(), valueCount * UnionVector.TYPE_WIDTH); for (ValueVector subVector : vector.getChildrenFromFields()) { validateOrThrow( valueCount == subVector.getValueCount(), "Union vector length not equal to child vector length. Union vector length %s, child vector length %s", valueCount, subVector.getValueCount()); subVector.accept(this, null); } return null; } @Override public Void visit(DenseUnionVector vector, Void value) { int valueCount = vector.getValueCount(); validateVectorCommon(vector); validateOffsetBuffer(vector, (long) valueCount * DenseUnionVector.OFFSET_WIDTH); validateTypeBuffer(vector.getTypeBuffer(), valueCount * DenseUnionVector.TYPE_WIDTH); for (ValueVector subVector : vector.getChildrenFromFields()) { subVector.accept(this, null); } return null; } @Override public Void visit(NullVector vector, Void value) { return null; } @Override public Void visit(ExtensionTypeVector vector, Void value) { vector.getUnderlyingVector().accept(this, value); return null; } @Override public Void visit(RunEndEncodedVector vector, Void value) { validateVectorCommon(vector); int valueCount = vector.getValueCount(); FieldVector runEndsVector = vector.getRunEndsVector(); if (runEndsVector != null) { validateOrThrow( runEndsVector.getNullCount() == 0, "Run ends vector cannot contain null values"); runEndsVector.accept(this, null); int runCount = runEndsVector.getValueCount(); if (runCount == 0) { validateOrThrow(valueCount == 0, "Run end vector does not contain enough elements"); } else if (runCount > 0) { long lastEnd = ((BaseIntVector) runEndsVector).getValueAsLong(runCount - 1); validateOrThrow( valueCount == lastEnd, "Vector logic length not equal to the last end in run ends vector. Logical length %s, last end %s", valueCount, lastEnd); } } FieldVector valuesVector = vector.getValuesVector(); if (valuesVector != null) { valuesVector.accept(this, null); } return null; } } ValidateVectorDataVisitor.java000066400000000000000000000153111515451615200351440ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/validate/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.validate; import static org.apache.arrow.vector.validate.ValidateUtil.validateOrThrow; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.BaseLargeVariableWidthVector; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.BaseVariableWidthViewVector; import org.apache.arrow.vector.ExtensionTypeVector; import org.apache.arrow.vector.NullVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.DenseUnionVector; 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.NonNullableStructVector; import org.apache.arrow.vector.complex.UnionVector; /** Utility for validating vector data. */ public class ValidateVectorDataVisitor implements VectorVisitor { private void validateOffsetBuffer(ValueVector vector, int valueCount) { if (valueCount == 0) { return; } ArrowBuf offsetBuffer = vector.getOffsetBuffer(); // verify that the values in the offset buffer is non-decreasing int prevValue = offsetBuffer.getInt(0); for (int i = 1; i <= valueCount; i++) { int curValue = offsetBuffer.getInt(i * 4); validateOrThrow( curValue >= 0, "The value at position %s of the offset buffer is negative: %s.", i, curValue); validateOrThrow( curValue >= prevValue, "The values in positions %s and %s of the offset buffer are decreasing: %s, %s.", i - 1, i, prevValue, curValue); prevValue = curValue; } } private void validateLargeOffsetBuffer(ValueVector vector, int valueCount) { if (valueCount == 0) { return; } ArrowBuf offsetBuffer = vector.getOffsetBuffer(); // verify that the values in the large offset buffer is non-decreasing long prevValue = offsetBuffer.getLong(0); for (int i = 1; i <= valueCount; i++) { long curValue = offsetBuffer.getLong((long) i * 8); validateOrThrow( curValue >= 0L, "The value at position %s of the large offset buffer is negative: %s.", i, curValue); validateOrThrow( curValue >= prevValue, "The values in positions %s and %s of the large offset buffer are decreasing: %s, %s.", i - 1, i, prevValue, curValue); prevValue = curValue; } } private void validateTypeBuffer(ArrowBuf typeBuf, int valueCount) { for (int i = 0; i < valueCount; i++) { validateOrThrow( typeBuf.getByte(i) >= 0, "The type id at position %s is negative: %s.", i, typeBuf.getByte(i)); } } @Override public Void visit(BaseFixedWidthVector vector, Void value) { vector.validateScalars(); return null; } @Override public Void visit(BaseVariableWidthVector vector, Void value) { validateOffsetBuffer(vector, vector.getValueCount()); vector.validateScalars(); return null; } @Override public Void visit(BaseLargeVariableWidthVector vector, Void value) { validateLargeOffsetBuffer(vector, vector.getValueCount()); vector.validateScalars(); return null; } @Override public Void visit(BaseVariableWidthViewVector vector, Void value) { vector.validateScalars(); return null; } @Override public Void visit(ListVector vector, Void value) { validateOffsetBuffer(vector, vector.getValueCount()); ValueVector innerVector = vector.getDataVector(); if (innerVector != null) { innerVector.accept(this, null); } return null; } @Override public Void visit(FixedSizeListVector vector, Void value) { validateOffsetBuffer(vector, vector.getValueCount()); ValueVector innerVector = vector.getDataVector(); if (innerVector != null) { innerVector.accept(this, null); } return null; } @Override public Void visit(LargeListVector vector, Void value) { validateLargeOffsetBuffer(vector, vector.getValueCount()); ValueVector innerVector = vector.getDataVector(); if (innerVector != null) { innerVector.accept(this, null); } return null; } @Override public Void visit(NonNullableStructVector vector, Void value) { for (ValueVector subVector : vector.getChildrenFromFields()) { subVector.accept(this, null); } return null; } @Override public Void visit(UnionVector vector, Void value) { validateTypeBuffer(vector.getTypeBuffer(), vector.getValueCount()); for (ValueVector subVector : vector.getChildrenFromFields()) { subVector.accept(this, null); } return null; } @Override public Void visit(DenseUnionVector vector, Void value) { validateTypeBuffer(vector.getTypeBuffer(), vector.getValueCount()); // validate offset buffer for (int i = 0; i < vector.getValueCount(); i++) { int offset = vector.getOffset(i); byte typeId = vector.getTypeId(i); ValueVector subVector = vector.getVectorByType(typeId); validateOrThrow( offset < subVector.getValueCount(), "Dense union vector offset exceeds sub-vector boundary. Vector offset %s, sub vector size %s", offset, subVector.getValueCount()); } for (ValueVector subVector : vector.getChildrenFromFields()) { subVector.accept(this, null); } return null; } @Override public Void visit(NullVector vector, Void value) { ValidateUtil.validateOrThrow( vector.getNullCount() == vector.getValueCount(), "NullVector should have only null entries."); return null; } @Override public Void visit(ExtensionTypeVector vector, Void value) { vector.getUnderlyingVector().accept(this, value); return null; } } ValidateVectorTypeVisitor.java000066400000000000000000000527451515451615200352300ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/validate/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.validate; import static org.apache.arrow.vector.validate.ValidateUtil.validateOrThrow; import org.apache.arrow.vector.BaseFixedWidthVector; 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.DateDayVector; import org.apache.arrow.vector.DateMilliVector; import org.apache.arrow.vector.Decimal256Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; import org.apache.arrow.vector.ExtensionTypeVector; 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.IntervalDayVector; import org.apache.arrow.vector.IntervalMonthDayNanoVector; import org.apache.arrow.vector.IntervalYearVector; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.LargeVarCharVector; 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.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; 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; import org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.IntervalUnit; 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.FieldType; /** Utility to validate vector type information. */ public class ValidateVectorTypeVisitor implements VectorVisitor { private void validateVectorCommon( ValueVector vector, Class expectedArrowType) { validateOrThrow(vector.getField() != null, "Vector field is empty."); validateOrThrow(vector.getField().getFieldType() != null, "Vector field type is empty."); ArrowType arrowType = vector.getField().getFieldType().getType(); validateOrThrow(arrowType != null, "Vector arrow type is empty."); validateOrThrow( expectedArrowType == arrowType.getClass(), "Incorrect arrow type for " + vector.getClass() + " : " + arrowType.toString()); } private void validateIntVector(ValueVector vector, int expectedWidth, boolean expectedSigned) { validateOrThrow( vector.getField().getFieldType().getType() instanceof ArrowType.Int, "Vector %s is not an integer vector.", vector.getClass()); ArrowType.Int intType = (ArrowType.Int) vector.getField().getFieldType().getType(); validateOrThrow( intType.getIsSigned() == expectedSigned, "Expecting bit width %s, actual width %s.", expectedWidth, intType.getBitWidth()); validateOrThrow( intType.getBitWidth() == expectedWidth, "Expecting bit width %s, actual bit width %s.", expectedWidth, intType.getBitWidth()); } private void validateFloatingPointVector( ValueVector vector, FloatingPointPrecision expectedPrecision) { validateOrThrow( vector.getField().getFieldType().getType() instanceof ArrowType.FloatingPoint, "Vector %s is not a floating point vector.", vector.getClass()); ArrowType.FloatingPoint floatType = (ArrowType.FloatingPoint) vector.getField().getFieldType().getType(); validateOrThrow( floatType.getPrecision() == expectedPrecision, "Expecting precision %s, actual precision %s.", expectedPrecision, floatType.getPrecision()); } private void validateDateVector(ValueVector vector, DateUnit expectedDateUnit) { validateOrThrow( vector.getField().getFieldType().getType() instanceof ArrowType.Date, "Vector %s is not a date vector", vector.getClass()); ArrowType.Date dateType = (ArrowType.Date) vector.getField().getFieldType().getType(); validateOrThrow( dateType.getUnit() == expectedDateUnit, "Expecting date unit %s, actual date unit %s.", expectedDateUnit, dateType.getUnit()); } private void validateDecimalVector(ValueVector vector) { validateOrThrow( vector.getField().getFieldType().getType() instanceof ArrowType.Decimal, "Vector %s is not a decimal vector", vector.getClass()); ArrowType.Decimal decimalType = (ArrowType.Decimal) vector.getField().getFieldType().getType(); validateOrThrow( decimalType.getScale() >= 0, "The scale of decimal %s is negative.", decimalType.getScale()); validateOrThrow( decimalType.getScale() <= decimalType.getPrecision(), "The scale of decimal %s is greater than the precision %s.", decimalType.getScale(), decimalType.getPrecision()); switch (decimalType.getBitWidth()) { case DecimalVector.TYPE_WIDTH * 8: validateOrThrow( decimalType.getPrecision() >= 1 && decimalType.getPrecision() <= DecimalVector.MAX_PRECISION, "Invalid precision %s for decimal 128.", decimalType.getPrecision()); break; case Decimal256Vector.TYPE_WIDTH * 8: validateOrThrow( decimalType.getPrecision() >= 1 && decimalType.getPrecision() <= Decimal256Vector.MAX_PRECISION, "Invalid precision %s for decimal 256.", decimalType.getPrecision()); break; default: throw new ValidateUtil.ValidateException( "Only decimal 128 or decimal 256 are supported for decimal types"); } } private void validateTimeVector( ValueVector vector, TimeUnit expectedTimeUnit, int expectedBitWidth) { validateOrThrow( vector.getField().getFieldType().getType() instanceof ArrowType.Time, "Vector %s is not a time vector.", vector.getClass()); ArrowType.Time timeType = (ArrowType.Time) vector.getField().getFieldType().getType(); validateOrThrow( timeType.getUnit() == expectedTimeUnit, "Expecting time unit %s, actual time unit %s.", expectedTimeUnit, timeType.getUnit()); validateOrThrow( timeType.getBitWidth() == expectedBitWidth, "Expecting bit width %s, actual bit width %s.", expectedBitWidth, timeType.getBitWidth()); } private void validateIntervalVector(ValueVector vector, IntervalUnit expectedIntervalUnit) { validateOrThrow( vector.getField().getFieldType().getType() instanceof ArrowType.Interval, "Vector %s is not an interval vector.", vector.getClass()); ArrowType.Interval intervalType = (ArrowType.Interval) vector.getField().getFieldType().getType(); validateOrThrow( intervalType.getUnit() == expectedIntervalUnit, "Expecting interval unit %s, actual date unit %s.", expectedIntervalUnit, intervalType.getUnit()); } private void validateTimeStampVector( ValueVector vector, TimeUnit expectedTimeUnit, boolean expectTZ) { validateOrThrow( vector.getField().getFieldType().getType() instanceof ArrowType.Timestamp, "Vector %s is not a time stamp vector.", vector.getClass()); ArrowType.Timestamp timestampType = (ArrowType.Timestamp) vector.getField().getFieldType().getType(); validateOrThrow( timestampType.getUnit() == expectedTimeUnit, "Expecting time stamp unit %s, actual time stamp unit %s.", expectedTimeUnit, timestampType.getUnit()); if (expectTZ) { validateOrThrow(timestampType.getTimezone() != null, "The time zone should not be null"); } else { validateOrThrow(timestampType.getTimezone() == null, "The time zone should be null"); } } private void validateExtensionTypeVector(ExtensionTypeVector vector) { validateOrThrow( vector.getField().getFieldType().getType() instanceof ArrowType.ExtensionType, "Vector %s is not an extension type vector.", vector.getClass()); validateOrThrow( vector .getField() .getMetadata() .containsKey(ArrowType.ExtensionType.EXTENSION_METADATA_KEY_NAME), "Field %s does not have proper extension type metadata: %s", vector.getField().getName(), vector.getField().getMetadata()); // Validate the storage vector type vector.getUnderlyingVector().accept(this, null); } @Override public Void visit(BaseFixedWidthVector vector, Void value) { if (vector instanceof TinyIntVector) { validateVectorCommon(vector, ArrowType.Int.class); validateIntVector(vector, 8, true); } else if (vector instanceof SmallIntVector) { validateVectorCommon(vector, ArrowType.Int.class); validateIntVector(vector, 16, true); } else if (vector instanceof IntVector) { validateVectorCommon(vector, ArrowType.Int.class); validateIntVector(vector, 32, true); } else if (vector instanceof BigIntVector) { validateVectorCommon(vector, ArrowType.Int.class); validateIntVector(vector, 64, true); } else if (vector instanceof UInt1Vector) { validateVectorCommon(vector, ArrowType.Int.class); validateIntVector(vector, 8, false); } else if (vector instanceof UInt2Vector) { validateVectorCommon(vector, ArrowType.Int.class); validateIntVector(vector, 16, false); } else if (vector instanceof UInt4Vector) { validateVectorCommon(vector, ArrowType.Int.class); validateIntVector(vector, 32, false); } else if (vector instanceof UInt8Vector) { validateVectorCommon(vector, ArrowType.Int.class); validateIntVector(vector, 64, false); } else if (vector instanceof BitVector) { validateVectorCommon(vector, ArrowType.Bool.class); } else if (vector instanceof DecimalVector || vector instanceof Decimal256Vector) { validateVectorCommon(vector, ArrowType.Decimal.class); validateDecimalVector(vector); } else if (vector instanceof DateDayVector) { validateVectorCommon(vector, ArrowType.Date.class); validateDateVector(vector, DateUnit.DAY); } else if (vector instanceof DateMilliVector) { validateVectorCommon(vector, ArrowType.Date.class); validateDateVector(vector, DateUnit.MILLISECOND); } else if (vector instanceof DurationVector) { validateVectorCommon(vector, ArrowType.Duration.class); ArrowType.Duration arrowType = (ArrowType.Duration) vector.getField().getType(); validateOrThrow( ((DurationVector) vector).getUnit() == arrowType.getUnit(), "Different duration time unit for vector and arrow type. Vector time unit %s, type time unit %s.", ((DurationVector) vector).getUnit(), arrowType.getUnit()); } else if (vector instanceof Float4Vector) { validateVectorCommon(vector, ArrowType.FloatingPoint.class); validateFloatingPointVector(vector, FloatingPointPrecision.SINGLE); } else if (vector instanceof Float8Vector) { validateVectorCommon(vector, ArrowType.FloatingPoint.class); validateFloatingPointVector(vector, FloatingPointPrecision.DOUBLE); } else if (vector instanceof IntervalDayVector) { validateVectorCommon(vector, ArrowType.Interval.class); validateIntervalVector(vector, IntervalUnit.DAY_TIME); } else if (vector instanceof IntervalMonthDayNanoVector) { validateVectorCommon(vector, ArrowType.Interval.class); validateIntervalVector(vector, IntervalUnit.MONTH_DAY_NANO); } else if (vector instanceof IntervalYearVector) { validateVectorCommon(vector, ArrowType.Interval.class); validateIntervalVector(vector, IntervalUnit.YEAR_MONTH); } else if (vector instanceof TimeMicroVector) { validateVectorCommon(vector, ArrowType.Time.class); validateTimeVector(vector, TimeUnit.MICROSECOND, 64); } else if (vector instanceof TimeMilliVector) { validateVectorCommon(vector, ArrowType.Time.class); validateTimeVector(vector, TimeUnit.MILLISECOND, 32); } else if (vector instanceof TimeNanoVector) { validateVectorCommon(vector, ArrowType.Time.class); validateTimeVector(vector, TimeUnit.NANOSECOND, 64); } else if (vector instanceof TimeSecVector) { validateVectorCommon(vector, ArrowType.Time.class); validateTimeVector(vector, TimeUnit.SECOND, 32); } else if (vector instanceof TimeStampMicroTZVector) { validateVectorCommon(vector, ArrowType.Timestamp.class); validateTimeStampVector(vector, TimeUnit.MICROSECOND, true); } else if (vector instanceof TimeStampMicroVector) { validateVectorCommon(vector, ArrowType.Timestamp.class); validateTimeStampVector(vector, TimeUnit.MICROSECOND, false); } else if (vector instanceof TimeStampMilliTZVector) { validateVectorCommon(vector, ArrowType.Timestamp.class); validateTimeStampVector(vector, TimeUnit.MILLISECOND, true); } else if (vector instanceof TimeStampMilliVector) { validateVectorCommon(vector, ArrowType.Timestamp.class); validateTimeStampVector(vector, TimeUnit.MILLISECOND, false); } else if (vector instanceof TimeStampNanoTZVector) { validateVectorCommon(vector, ArrowType.Timestamp.class); validateTimeStampVector(vector, TimeUnit.NANOSECOND, true); } else if (vector instanceof TimeStampNanoVector) { validateVectorCommon(vector, ArrowType.Timestamp.class); validateTimeStampVector(vector, TimeUnit.NANOSECOND, false); } else if (vector instanceof TimeStampSecTZVector) { validateVectorCommon(vector, ArrowType.Timestamp.class); validateTimeStampVector(vector, TimeUnit.SECOND, true); } else if (vector instanceof TimeStampSecVector) { validateVectorCommon(vector, ArrowType.Timestamp.class); validateTimeStampVector(vector, TimeUnit.SECOND, false); } else if (vector instanceof FixedSizeBinaryVector) { validateVectorCommon(vector, ArrowType.FixedSizeBinary.class); ArrowType.FixedSizeBinary arrowType = (ArrowType.FixedSizeBinary) vector.getField().getType(); validateOrThrow( arrowType.getByteWidth() > 0, "The byte width of a FixedSizeBinaryVector %s is not positive.", arrowType.getByteWidth()); validateOrThrow( arrowType.getByteWidth() == vector.getTypeWidth(), "Type width mismatch for FixedSizeBinaryVector. Vector type width %s, arrow type width %s.", vector.getTypeWidth(), arrowType.getByteWidth()); } else { throw new IllegalArgumentException( "Unknown type for fixed width vector " + vector.getClass()); } return null; } @Override public Void visit(BaseVariableWidthVector vector, Void value) { if (vector instanceof VarCharVector) { validateVectorCommon(vector, ArrowType.Utf8.class); } else if (vector instanceof VarBinaryVector) { validateVectorCommon(vector, ArrowType.Binary.class); } return null; } @Override public Void visit(BaseLargeVariableWidthVector vector, Void value) { if (vector instanceof LargeVarCharVector) { validateVectorCommon(vector, ArrowType.LargeUtf8.class); } else if (vector instanceof LargeVarBinaryVector) { validateVectorCommon(vector, ArrowType.LargeBinary.class); } return null; } @Override public Void visit(BaseVariableWidthViewVector vector, Void value) { if (vector instanceof ViewVarCharVector) { validateVectorCommon(vector, ArrowType.Utf8View.class); } else if (vector instanceof ViewVarBinaryVector) { validateVectorCommon(vector, ArrowType.BinaryView.class); } return null; } @Override public Void visit(ListVector vector, Void value) { validateVectorCommon(vector, ArrowType.List.class); ValueVector innerVector = vector.getDataVector(); if (innerVector != null) { innerVector.accept(this, null); } return null; } @Override public Void visit(FixedSizeListVector vector, Void value) { validateVectorCommon(vector, ArrowType.FixedSizeList.class); ArrowType.FixedSizeList arrowType = (ArrowType.FixedSizeList) vector.getField().getType(); validateOrThrow( arrowType.getListSize() == vector.getListSize(), "Inconsistent list size for FixedSizeListVector. Vector list size %s, arrow type list size %s.", vector.getListSize(), arrowType.getListSize()); validateOrThrow( arrowType.getListSize() > 0, "The list size %s is not positive.", arrowType.getListSize()); ValueVector innerVector = vector.getDataVector(); if (innerVector != null) { innerVector.accept(this, null); } return null; } @Override public Void visit(LargeListVector vector, Void value) { validateVectorCommon(vector, ArrowType.LargeList.class); ValueVector innerVector = vector.getDataVector(); if (innerVector != null) { innerVector.accept(this, null); } return null; } @Override public Void visit(NonNullableStructVector vector, Void value) { validateVectorCommon(vector, ArrowType.Struct.class); validateOrThrow( vector.getField().getChildren().size() == vector.getChildrenFromFields().size(), "Child field count and child vector count mismatch. Vector child count %s, field child count %s", vector.getChildrenFromFields().size(), vector.getField().getChildren().size()); for (int i = 0; i < vector.getChildrenFromFields().size(); i++) { ValueVector subVector = vector.getChildByOrdinal(i); FieldType subType = vector.getField().getChildren().get(i).getFieldType(); validateOrThrow( subType.equals(subVector.getField().getFieldType()), "Struct vector's field type not equal to the child vector's field type. " + "Struct field type %s, sub-vector field type %s", subType, subVector.getField().getFieldType()); subVector.accept(this, null); } return null; } @Override public Void visit(UnionVector vector, Void value) { validateVectorCommon(vector, ArrowType.Union.class); ArrowType.Union arrowType = (ArrowType.Union) vector.getField().getType(); validateOrThrow( arrowType.getMode() == UnionMode.Sparse, "The union mode of UnionVector must be sparse"); for (ValueVector subVector : vector.getChildrenFromFields()) { subVector.accept(this, null); } return null; } @Override public Void visit(DenseUnionVector vector, Void value) { validateVectorCommon(vector, ArrowType.Union.class); ArrowType.Union arrowType = (ArrowType.Union) vector.getField().getType(); validateOrThrow( arrowType.getMode() == UnionMode.Dense, "The union mode of DenseUnionVector must be dense"); for (ValueVector subVector : vector.getChildrenFromFields()) { subVector.accept(this, null); } return null; } @Override public Void visit(NullVector vector, Void value) { validateVectorCommon(vector, ArrowType.Null.class); return null; } @Override public Void visit(ExtensionTypeVector vector, Void value) { validateExtensionTypeVector(vector); return null; } @Override public Void visit(RunEndEncodedVector vector, Void value) { validateVectorCommon(vector, ArrowType.RunEndEncoded.class); for (ValueVector subVector : vector.getChildrenFromFields()) { if (subVector != null) { subVector.accept(this, null); } } return null; } } ValidateVectorVisitor.java000066400000000000000000000267451515451615200343670ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/main/java/org/apache/arrow/vector/validate/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.validate; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.BaseLargeVariableWidthVector; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.BaseVariableWidthViewVector; import org.apache.arrow.vector.ExtensionTypeVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.NullVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.DenseUnionVector; 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.NonNullableStructVector; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.util.ValueVectorUtility; /** * Visitor to validate vector (without validating data). This visitor could be used for {@link * ValueVector#accept(VectorVisitor, Object)} API, and also users could simply use {@link * ValueVectorUtility#validate(ValueVector)}. */ public class ValidateVectorVisitor implements VectorVisitor { @Override public Void visit(BaseFixedWidthVector 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 public Void visit(BaseVariableWidthVector 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"); } ArrowBuf offsetBuf = vector.getOffsetBuffer(); int minBufferSize = (vector.getValueCount() + 1) * BaseVariableWidthVector.OFFSET_WIDTH; if (offsetBuf.capacity() < minBufferSize) { throw new IllegalArgumentException( String.format( "offsetBuffer too small in vector of type %s" + " and valueCount %s : expected at least %s byte(s), got %s", vector.getField().getType().toString(), vector.getValueCount(), minBufferSize, offsetBuf.capacity())); } int firstOffset = vector.getOffsetBuffer().getInt(0); int lastOffset = vector .getOffsetBuffer() .getInt(vector.getValueCount() * BaseVariableWidthVector.OFFSET_WIDTH); if (firstOffset < 0 || lastOffset < 0) { throw new IllegalArgumentException("Negative offsets in vector"); } int dataExtent = lastOffset - firstOffset; if (dataExtent > 0 && (vector.getDataBuffer().capacity() == 0)) { throw new IllegalArgumentException("dataBuffer capacity is 0"); } if (dataExtent > vector.getDataBuffer().capacity()) { throw new IllegalArgumentException( String.format( "Length spanned by offsets %s larger than" + " dataBuffer capacity %s", dataExtent, vector.getValueCount())); } } return null; } @Override public Void visit(BaseLargeVariableWidthVector left, Void value) { return null; } @Override 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 public Void visit(ListVector vector, Void value) { FieldVector dataVector = vector.getDataVector(); if (vector.getValueCount() > 0) { ArrowBuf offsetBuf = vector.getOffsetBuffer(); int minBufferSize = (vector.getValueCount() + 1) * BaseVariableWidthVector.OFFSET_WIDTH; if (offsetBuf.capacity() < minBufferSize) { throw new IllegalArgumentException( String.format( "offsetBuffer too small in vector of type %s" + " and valueCount %s : expected at least %s byte(s), got %s", vector.getField().getType().toString(), vector.getValueCount(), minBufferSize, offsetBuf.capacity())); } int firstOffset = vector.getOffsetBuffer().getInt(0); int lastOffset = vector .getOffsetBuffer() .getInt(vector.getValueCount() * BaseVariableWidthVector.OFFSET_WIDTH); if (firstOffset < 0 || lastOffset < 0) { throw new IllegalArgumentException("Negative offsets in list vector"); } int dataExtent = lastOffset - firstOffset; if (dataExtent > 0 && (dataVector.getDataBuffer() == null || dataVector.getDataBuffer().capacity() == 0)) { throw new IllegalArgumentException("valueBuffer is null or capacity is 0"); } if (dataExtent > dataVector.getValueCount()) { throw new IllegalArgumentException( String.format( "Length spanned by list offsets (%s) larger than" + " data vector valueCount (length %s)", dataExtent, dataVector.getValueCount())); } } return dataVector.accept(this, null); } @Override public Void visit(LargeListVector vector, Void value) { FieldVector dataVector = vector.getDataVector(); if (vector.getValueCount() > 0) { ArrowBuf offsetBuf = vector.getOffsetBuffer(); long minBufferSize = (vector.getValueCount() + 1) * LargeListVector.OFFSET_WIDTH; if (offsetBuf.capacity() < minBufferSize) { throw new IllegalArgumentException( String.format( "offsetBuffer too small in vector of type %s" + " and valueCount %s : expected at least %s byte(s), got %s", vector.getField().getType().toString(), vector.getValueCount(), minBufferSize, offsetBuf.capacity())); } long firstOffset = vector.getOffsetBuffer().getLong(0); long lastOffset = vector.getOffsetBuffer().getLong(vector.getValueCount() * LargeListVector.OFFSET_WIDTH); if (firstOffset < 0 || lastOffset < 0) { throw new IllegalArgumentException("Negative offsets in list vector"); } long dataExtent = lastOffset - firstOffset; if (dataExtent > 0 && (dataVector.getDataBuffer() == null || dataVector.getDataBuffer().capacity() == 0)) { throw new IllegalArgumentException("valueBuffer is null or capacity is 0"); } if (dataExtent > dataVector.getValueCount()) { throw new IllegalArgumentException( String.format( "Length spanned by list offsets (%s) larger than" + " data vector valueCount (length %s)", dataExtent, dataVector.getValueCount())); } } return dataVector.accept(this, null); } @Override public Void visit(FixedSizeListVector vector, Void value) { FieldVector dataVector = vector.getDataVector(); int valueCount = vector.getValueCount(); int listSize = vector.getListSize(); if (valueCount > 0 && (dataVector.getDataBuffer() == null || dataVector.getDataBuffer().capacity() == 0)) { throw new IllegalArgumentException("valueBuffer is null or capacity is 0"); } if (valueCount * listSize != dataVector.getValueCount()) { throw new IllegalArgumentException( String.format( "data vector valueCount invalid, expect %s, " + "actual is: %s", valueCount * listSize, dataVector.getValueCount())); } return null; } @Override public Void visit(NonNullableStructVector vector, Void value) { List childFields = vector.getField().getChildren(); final int valueCount = vector.getValueCount(); for (int i = 0; i < childFields.size(); i++) { FieldVector child = vector.getChildrenFromFields().get(i); if (child.getValueCount() != valueCount) { throw new IllegalArgumentException( String.format( "struct child vector #%s valueCount is not equals with " + "struct vector, expect %s, actual %s", i, vector.getValueCount(), child.getValueCount())); } if (!childFields.get(i).getType().equals(child.getField().getType())) { throw new IllegalArgumentException( String.format( "struct child vector #%s does not match type: %s vs %s", i, childFields.get(i).getType().toString(), child.getField().getType().toString())); } child.accept(this, null); } return null; } @Override public Void visit(UnionVector vector, Void value) { List childFields = vector.getField().getChildren(); final int valueCount = vector.getValueCount(); for (int i = 0; i < childFields.size(); i++) { FieldVector child = vector.getChildrenFromFields().get(i); if (child.getValueCount() != valueCount) { throw new IllegalArgumentException( String.format( "union child vector #%s valueCount is not equals with union" + " vector, expect %s, actual %s", i, vector.getValueCount(), child.getValueCount())); } if (!childFields.get(i).getType().equals(child.getField().getType())) { throw new IllegalArgumentException( String.format( "union child vector #%s does not match type: %s vs %s", i, childFields.get(i).getType().toString(), child.getField().getType().toString())); } child.accept(this, null); } return null; } @Override public Void visit(DenseUnionVector vector, Void value) { List childFields = vector.getField().getChildren(); for (int i = 0; i < childFields.size(); i++) { FieldVector child = vector.getChildrenFromFields().get(i); if (!childFields.get(i).getType().equals(child.getField().getType())) { throw new IllegalArgumentException( String.format( "union child vector #%s does not match type: %s vs %s", i, childFields.get(i).getType().toString(), child.getField().getType().toString())); } child.accept(this, null); } return null; } @Override public Void visit(NullVector vector, Void value) { return null; } @Override public Void visit(ExtensionTypeVector vector, Void value) { vector.getUnderlyingVector().accept(this, value); return null; } } apache-arrow-java-19.0.0/vector/src/shade/000077500000000000000000000000001515451615200202535ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/shade/LICENSE.txt000066400000000000000000000265051515451615200221060ustar00rootroot00000000000000 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 apache-arrow-java-19.0.0/vector/src/shade/NOTICE.txt000066400000000000000000000002531515451615200217750ustar00rootroot00000000000000Apache Arrow Java Copyright 2016-2025 The Apache Software Foundation This product includes software developed at The Apache Software Foundation (http://www.apache.org/). apache-arrow-java-19.0.0/vector/src/test/000077500000000000000000000000001515451615200201465ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/000077500000000000000000000000001515451615200210675ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/000077500000000000000000000000001515451615200216565ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/000077500000000000000000000000001515451615200230775ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/000077500000000000000000000000001515451615200242315ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/000077500000000000000000000000001515451615200255335ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/DirtyRootAllocator.java000066400000000000000000000032431515451615200322000ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferManager; import org.apache.arrow.memory.RootAllocator; /** * Root allocator that returns buffers pre-filled with a given value.
    * Useful for testing if value vectors are properly zeroing their buffers. */ public class DirtyRootAllocator extends RootAllocator { private final byte fillValue; public DirtyRootAllocator(final long limit, final byte fillValue) { super(limit); this.fillValue = fillValue; } @Override public ArrowBuf buffer(long size) { return buffer(size, null); } @Override public ArrowBuf buffer(long size, BufferManager manager) { ArrowBuf buffer = super.buffer(size, manager); // contaminate the buffer for (int i = 0; i < buffer.capacity(); i++) { buffer.setByte(i, fillValue); } return buffer; } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/ITTestLargeVector.java000066400000000000000000000235521515451615200317170ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import java.math.BigDecimal; import java.nio.charset.StandardCharsets; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.holders.NullableDecimalHolder; import org.junit.jupiter.api.Test; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * Integration test for a vector with a large (more than 2GB) {@link * org.apache.arrow.memory.ArrowBuf} as the data buffer. To run this test, please make sure there is * at least 4GB free memory in the system. */ public class ITTestLargeVector { private static final Logger logger = LoggerFactory.getLogger(ITTestLargeVector.class); @Test public void testLargeLongVector() { logger.trace("Testing large big int vector."); final long bufSize = 4 * 1024 * 1024 * 1024L; final int vecLength = (int) (bufSize / BigIntVector.TYPE_WIDTH); try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); BigIntVector largeVec = new BigIntVector("vec", allocator)) { largeVec.allocateNew(vecLength); logger.trace("Successfully allocated a vector with capacity {}", vecLength); for (int i = 0; i < vecLength; i++) { largeVec.set(i, i * 10L); if ((i + 1) % 10000 == 0) { logger.trace("Successfully written {} values", i + 1); } } logger.trace("Successfully written {} values", vecLength); for (int i = 0; i < vecLength; i++) { long val = largeVec.get(i); assertEquals(i * 10L, val); if ((i + 1) % 10000 == 0) { logger.trace("Successfully read {} values", i + 1); } } logger.trace("Successfully read {} values", vecLength); } logger.trace("Successfully released the large vector."); } @Test public void testLargeIntVector() { logger.trace("Testing large int vector."); final long bufSize = 4 * 1024 * 1024 * 1024L; final int vecLength = (int) (bufSize / IntVector.TYPE_WIDTH); try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); IntVector largeVec = new IntVector("vec", allocator)) { largeVec.allocateNew(vecLength); logger.trace("Successfully allocated a vector with capacity {}", vecLength); for (int i = 0; i < vecLength; i++) { largeVec.set(i, i); if ((i + 1) % 10000 == 0) { logger.trace("Successfully written {} values", i + 1); } } logger.trace("Successfully written {} values", vecLength); for (int i = 0; i < vecLength; i++) { long val = largeVec.get(i); assertEquals(i, val); if ((i + 1) % 10000 == 0) { logger.trace("Successfully read {} values", i + 1); } } logger.trace("Successfully read {} values", vecLength); } logger.trace("Successfully released the large vector."); } @Test public void testLargeDecimalVector() { logger.trace("Testing large decimal vector."); final long bufSize = 4 * 1024 * 1024 * 1024L; final int vecLength = (int) (bufSize / DecimalVector.TYPE_WIDTH); try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); DecimalVector largeVec = new DecimalVector("vec", allocator, 38, 0)) { largeVec.allocateNew(vecLength); logger.trace("Successfully allocated a vector with capacity {}", vecLength); for (int i = 0; i < vecLength; i++) { largeVec.set(i, 0); if ((i + 1) % 10000 == 0) { logger.trace("Successfully written {} values", i + 1); } } logger.trace("Successfully written {} values", vecLength); for (int i = 0; i < vecLength; i++) { ArrowBuf buf = largeVec.get(i); assertEquals(DecimalVector.TYPE_WIDTH, buf.capacity()); assertEquals(0, buf.getLong(0)); assertEquals(0, buf.getLong(8)); if ((i + 1) % 10000 == 0) { logger.trace("Successfully read {} values", i + 1); } } logger.trace("Successfully read {} values", vecLength); // try setting values with a large offset in the buffer largeVec.set(vecLength - 1, 12345L); assertEquals(12345L, largeVec.getObject(vecLength - 1).longValue()); NullableDecimalHolder holder = new NullableDecimalHolder(); holder.buffer = largeVec.valueBuffer; holder.isSet = 1; holder.start = (long) (vecLength - 1) * largeVec.getTypeWidth(); assertTrue(holder.start > Integer.MAX_VALUE); largeVec.set(0, holder); BigDecimal decimal = largeVec.getObject(0); assertEquals(12345L, decimal.longValue()); logger.trace("Successfully setting values from large offsets"); } logger.trace("Successfully released the large vector."); } @Test public void testLargeFixedSizeBinaryVector() { logger.trace("Testing large fixed size binary vector."); final long bufSize = 4 * 1024 * 1024 * 1024L; final int typeWidth = 8; final int vecLength = (int) (bufSize / typeWidth); try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); FixedSizeBinaryVector largeVec = new FixedSizeBinaryVector("vec", allocator, typeWidth)) { largeVec.allocateNew(vecLength); logger.trace("Successfully allocated a vector with capacity {}", vecLength); byte[] value = new byte[] {'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'}; for (int i = 0; i < vecLength; i++) { largeVec.set(i, value); if ((i + 1) % 10000 == 0) { logger.trace("Successfully written {} values", i + 1); } } logger.trace("Successfully written {} values", vecLength); for (int i = 0; i < vecLength; i++) { byte[] buf = largeVec.get(i); assertEquals(typeWidth, buf.length); assertArrayEquals(buf, value); if ((i + 1) % 10000 == 0) { logger.trace("Successfully read {} values", i + 1); } } logger.trace("Successfully read {} values", vecLength); } logger.trace("Successfully released the large vector."); } @Test public void testLargeVarCharVector() { logger.trace("Testing large var char vector."); final long bufSize = 4 * 1024 * 1024 * 1024L; final int vecLength = (int) (bufSize / BaseVariableWidthVector.OFFSET_WIDTH); final String strElement = "a"; try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); VarCharVector largeVec = new VarCharVector("vec", allocator)) { largeVec.allocateNew(vecLength); logger.trace("Successfully allocated a vector with capacity " + vecLength); for (int i = 0; i < vecLength; i++) { largeVec.setSafe(i, strElement.getBytes(StandardCharsets.UTF_8)); if ((i + 1) % 10000 == 0) { logger.trace("Successfully written " + (i + 1) + " values"); } } largeVec.setValueCount(vecLength); assertTrue(largeVec.getOffsetBuffer().readableBytes() > Integer.MAX_VALUE); assertTrue(largeVec.getDataBuffer().readableBytes() < Integer.MAX_VALUE); logger.trace("Successfully written " + vecLength + " values"); for (int i = 0; i < vecLength; i++) { byte[] val = largeVec.get(i); assertEquals(strElement, new String(val, StandardCharsets.UTF_8)); if ((i + 1) % 10000 == 0) { logger.trace("Successfully read " + (i + 1) + " values"); } } logger.trace("Successfully read " + vecLength + " values"); } logger.trace("Successfully released the large vector."); } @Test public void testLargeLargeVarCharVector() { logger.trace("Testing large large var char vector."); final long bufSize = 4 * 1024 * 1024 * 1024L; final int vecLength = (int) (bufSize / BaseLargeVariableWidthVector.OFFSET_WIDTH); final String strElement = "9876543210"; try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); LargeVarCharVector largeVec = new LargeVarCharVector("vec", allocator)) { largeVec.allocateNew(vecLength); logger.trace("Successfully allocated a vector with capacity " + vecLength); for (int i = 0; i < vecLength; i++) { largeVec.setSafe(i, strElement.getBytes(StandardCharsets.UTF_8)); if ((i + 1) % 10000 == 0) { logger.trace("Successfully written " + (i + 1) + " values"); } } largeVec.setValueCount(vecLength); assertTrue(largeVec.getOffsetBuffer().readableBytes() > Integer.MAX_VALUE); assertTrue(largeVec.getDataBuffer().readableBytes() > Integer.MAX_VALUE); logger.trace("Successfully written " + vecLength + " values"); for (int i = 0; i < vecLength; i++) { byte[] val = largeVec.get(i); assertEquals(strElement, new String(val, StandardCharsets.UTF_8)); if ((i + 1) % 10000 == 0) { logger.trace("Successfully read " + (i + 1) + " values"); } } logger.trace("Successfully read " + vecLength + " values"); } logger.trace("Successfully released the large vector."); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestBitVector.java000066400000000000000000000435761515451615200311560ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.stream.IntStream; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.memory.util.hash.MurmurHasher; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; 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 TestBitVector { private static final String EMPTY_SCHEMA_PATH = ""; private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testBitVectorCopyFromSafe() { final int size = 20; try (final BitVector src = new BitVector(EMPTY_SCHEMA_PATH, allocator); final BitVector dst = new BitVector(EMPTY_SCHEMA_PATH, allocator)) { src.allocateNew(size); dst.allocateNew(10); for (int i = 0; i < size; i++) { src.set(i, i % 2); } src.setValueCount(size); for (int i = 0; i < size; i++) { dst.copyFromSafe(i, i, src); } dst.setValueCount(size); for (int i = 0; i < size; i++) { assertEquals(src.getObject(i), dst.getObject(i)); } } } @Test public void testSplitAndTransfer() throws Exception { try (final BitVector sourceVector = new BitVector("bitvector", allocator)) { sourceVector.allocateNew(40); /* populate the bitvector -- 010101010101010101010101..... */ for (int i = 0; i < 40; i++) { if ((i & 1) == 1) { sourceVector.set(i, 1); } else { sourceVector.set(i, 0); } } sourceVector.setValueCount(40); /* check the vector output */ for (int i = 0; i < 40; i++) { int result = sourceVector.get(i); if ((i & 1) == 1) { assertEquals(Integer.toString(1), Integer.toString(result)); } else { assertEquals(Integer.toString(0), Integer.toString(result)); } } try (final BitVector toVector = new BitVector("toVector", allocator)) { final TransferPair transferPair = sourceVector.makeTransferPair(toVector); /* * form test cases such that we cover: * * (1) the start index is exactly where a particular byte starts in the source bit vector * (2) the start index is randomly positioned within a byte in the source bit vector * (2.1) the length is a multiple of 8 * (2.2) the length is not a multiple of 8 */ final int[][] transferLengths = {{0, 8}, {8, 10}, {18, 0}, {18, 8}, {26, 0}, {26, 14}}; for (final int[] transferLength : transferLengths) { final int start = transferLength[0]; final int length = transferLength[1]; transferPair.splitAndTransfer(start, length); /* check the toVector output after doing splitAndTransfer */ for (int i = 0; i < length; i++) { int actual = toVector.get(i); int expected = sourceVector.get(start + i); assertEquals( expected, actual, "different data values not expected --> sourceVector index: " + (start + i) + " toVector index: " + i); } } } } } @Test public void testSplitAndTransfer1() throws Exception { try (final BitVector sourceVector = new BitVector("bitvector", allocator)) { sourceVector.allocateNew(8190); /* populate the bitvector */ for (int i = 0; i < 8190; i++) { sourceVector.set(i, 1); } sourceVector.setValueCount(8190); /* check the vector output */ for (int i = 0; i < 8190; i++) { int result = sourceVector.get(i); assertEquals(Integer.toString(1), Integer.toString(result)); } try (final BitVector toVector = new BitVector("toVector", allocator)) { final TransferPair transferPair = sourceVector.makeTransferPair(toVector); final int[][] transferLengths = {{0, 4095}, {4095, 4095}}; for (final int[] transferLength : transferLengths) { final int start = transferLength[0]; final int length = transferLength[1]; transferPair.splitAndTransfer(start, length); /* check the toVector output after doing splitAndTransfer */ for (int i = 0; i < length; i++) { int actual = toVector.get(i); int expected = sourceVector.get(start + i); assertEquals( expected, actual, "different data values not expected --> sourceVector index: " + (start + i) + " toVector index: " + i); } } } } } @Test public void testSplitAndTransfer2() throws Exception { try (final BitVector sourceVector = new BitVector("bitvector", allocator)) { sourceVector.allocateNew(32); /* populate the bitvector */ for (int i = 0; i < 32; i++) { if ((i & 1) == 1) { sourceVector.set(i, 1); } else { sourceVector.set(i, 0); } } sourceVector.setValueCount(32); /* check the vector output */ for (int i = 0; i < 32; i++) { int result = sourceVector.get(i); if ((i & 1) == 1) { assertEquals(Integer.toString(1), Integer.toString(result)); } else { assertEquals(Integer.toString(0), Integer.toString(result)); } } try (final BitVector toVector = new BitVector("toVector", allocator)) { final TransferPair transferPair = sourceVector.makeTransferPair(toVector); final int[][] transferLengths = { {5, 22}, {5, 24}, {5, 25}, {5, 27}, {0, 31}, {5, 7}, {2, 3} }; for (final int[] transferLength : transferLengths) { final int start = transferLength[0]; final int length = transferLength[1]; transferPair.splitAndTransfer(start, length); /* check the toVector output after doing splitAndTransfer */ for (int i = 0; i < length; i++) { int actual = toVector.get(i); int expected = sourceVector.get(start + i); assertEquals( expected, actual, "different data values not expected --> sourceVector index: " + (start + i) + " toVector index: " + i); } } } } } @Test public void testReallocAfterVectorTransfer1() { try (final BitVector vector = new BitVector(EMPTY_SCHEMA_PATH, allocator)) { vector.allocateNew(4096); int valueCapacity = vector.getValueCapacity(); assertEquals(4096, valueCapacity); for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { vector.setToOne(i); } } for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { assertEquals(1, vector.get(i), "unexpected cleared bit at index: " + i); } else { assertTrue(vector.isNull(i), "unexpected set bit at index: " + i); } } /* trigger first realloc */ vector.setSafeToOne(valueCapacity); assertEquals(valueCapacity * 2, vector.getValueCapacity()); for (int i = valueCapacity; i < valueCapacity * 2; i++) { if ((i & 1) == 1) { vector.setToOne(i); } } for (int i = 0; i < valueCapacity * 2; i++) { if (((i & 1) == 1) || (i == valueCapacity)) { assertEquals(1, vector.get(i), "unexpected cleared bit at index: " + i); } else { assertTrue(vector.isNull(i), "unexpected set bit at index: " + i); } } /* trigger second realloc */ vector.setSafeToOne(valueCapacity * 2); assertEquals(valueCapacity * 4, vector.getValueCapacity()); for (int i = valueCapacity * 2; i < valueCapacity * 4; i++) { if ((i & 1) == 1) { vector.setToOne(i); } } for (int i = 0; i < valueCapacity * 4; i++) { if (((i & 1) == 1) || (i == valueCapacity) || (i == valueCapacity * 2)) { assertEquals(1, vector.get(i), "unexpected cleared bit at index: " + i); } else { assertTrue(vector.isNull(i), "unexpected set bit at index: " + i); } } /* now transfer the vector */ TransferPair transferPair = vector.getTransferPair(allocator); transferPair.transfer(); final BitVector toVector = (BitVector) transferPair.getTo(); assertEquals(valueCapacity * 4, toVector.getValueCapacity()); /* realloc the toVector */ toVector.setSafeToOne(valueCapacity * 4); for (int i = 0; i < toVector.getValueCapacity(); i++) { if (i <= valueCapacity * 4) { if (((i & 1) == 1) || (i == valueCapacity) || (i == valueCapacity * 2) || (i == valueCapacity * 4)) { assertEquals(1, toVector.get(i), "unexpected cleared bit at index: " + i); } else { assertTrue(toVector.isNull(i), "unexpected set bit at index: " + i); } } else { assertTrue(toVector.isNull(i), "unexpected set bit at index: " + i); } } toVector.close(); } } @Test public void testReallocAfterVectorTransfer2() { try (final BitVector vector = new BitVector(EMPTY_SCHEMA_PATH, allocator)) { vector.allocateNew(4096); int valueCapacity = vector.getValueCapacity(); assertEquals(4096, valueCapacity); for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { vector.set(i, 1); } } for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { assertFalse(vector.isNull(i), "unexpected cleared bit at index: " + i); } else { assertTrue(vector.isNull(i), "unexpected set bit at index: " + i); } } /* trigger first realloc */ vector.setSafe(valueCapacity, 1, 1); assertEquals(valueCapacity * 2, vector.getValueCapacity()); for (int i = valueCapacity; i < valueCapacity * 2; i++) { if ((i & 1) == 1) { vector.set(i, 1); } } for (int i = 0; i < valueCapacity * 2; i++) { if (((i & 1) == 1) || (i == valueCapacity)) { assertFalse(vector.isNull(i), "unexpected cleared bit at index: " + i); } else { assertTrue(vector.isNull(i), "unexpected set bit at index: " + i); } } /* trigger second realloc */ vector.setSafe(valueCapacity * 2, 1, 1); assertEquals(valueCapacity * 4, vector.getValueCapacity()); for (int i = valueCapacity * 2; i < valueCapacity * 4; i++) { if ((i & 1) == 1) { vector.set(i, 1); } } for (int i = 0; i < valueCapacity * 4; i++) { if (((i & 1) == 1) || (i == valueCapacity) || (i == valueCapacity * 2)) { assertFalse(vector.isNull(i), "unexpected cleared bit at index: " + i); } else { assertTrue(vector.isNull(i), "unexpected set bit at index: " + i); } } /* now transfer the vector */ TransferPair transferPair = vector.getTransferPair(allocator); transferPair.transfer(); final BitVector toVector = (BitVector) transferPair.getTo(); assertEquals(valueCapacity * 4, toVector.getValueCapacity()); /* realloc the toVector */ toVector.setSafe(valueCapacity * 4, 1, 1); for (int i = 0; i < toVector.getValueCapacity(); i++) { if (i <= valueCapacity * 4) { if (((i & 1) == 1) || (i == valueCapacity) || (i == valueCapacity * 2) || (i == valueCapacity * 4)) { assertFalse(toVector.isNull(i), "unexpected cleared bit at index: " + i); } else { assertTrue(toVector.isNull(i), "unexpected set bit at index: " + i); } } else { assertTrue(toVector.isNull(i), "unexpected set bit at index: " + i); } } toVector.close(); } } @Test public void testBitVector() { // Create a new value vector for 1024 integers try (final BitVector vector = new BitVector(EMPTY_SCHEMA_PATH, allocator)) { vector.allocateNew(1024); vector.setValueCount(1024); // Put and set a few values vector.set(0, 1); vector.set(1, 0); vector.set(100, 0); vector.set(1022, 1); vector.setValueCount(1024); assertEquals(1, vector.get(0)); assertEquals(0, vector.get(1)); assertEquals(0, vector.get(100)); assertEquals(1, vector.get(1022)); assertEquals(1020, vector.getNullCount()); // test setting the same value twice vector.set(0, 1); vector.set(0, 1); vector.set(1, 0); vector.set(1, 0); assertEquals(1, vector.get(0)); assertEquals(0, vector.get(1)); // test toggling the values vector.set(0, 0); vector.set(1, 1); assertEquals(0, vector.get(0)); assertEquals(1, vector.get(1)); // should not change assertEquals(1020, vector.getNullCount()); // Ensure null value assertTrue(vector.isNull(3)); // unset the previously set bits vector.setNull(0); vector.setNull(1); vector.setNull(100); vector.setNull(1022); // this should set all the array to 0 assertEquals(1024, vector.getNullCount()); // set all the array to 1 for (int i = 0; i < 1024; ++i) { assertEquals(1024 - i, vector.getNullCount()); vector.set(i, 1); } assertEquals(0, vector.getNullCount()); vector.allocateNew(1015); vector.setValueCount(1015); // ensure it has been zeroed assertEquals(1015, vector.getNullCount()); vector.set(0, 1); vector.set(1014, 1); // ensure that the last item of the last byte is allocated assertEquals(1013, vector.getNullCount()); vector.zeroVector(); assertEquals(1015, vector.getNullCount()); // set all the array to 1 for (int i = 0; i < 1015; ++i) { assertEquals(1015 - i, vector.getNullCount()); vector.set(i, 1); } assertEquals(0, vector.getNullCount()); } } @Test public void testBitVectorRangeSetAllOnes() { validateRange(1000, 0, 1000); validateRange(1000, 0, 1); validateRange(1000, 1, 2); validateRange(1000, 5, 6); validateRange(1000, 5, 10); validateRange(1000, 5, 150); validateRange(1000, 5, 27); for (int i = 0; i < 8; i++) { for (int j = 0; j < 8; j++) { validateRange(1000, 10 + i, 27 + j); validateRange(1000, i, j); } } } private void validateRange(int length, int start, int count) { String desc = "[" + start + ", " + (start + count) + ") "; try (BitVector bitVector = new BitVector("bits", allocator)) { bitVector.reset(); bitVector.allocateNew(length); bitVector.setRangeToOne(start, count); for (int i = 0; i < start; i++) { assertTrue(bitVector.isNull(i), desc + i); } for (int i = start; i < start + count; i++) { assertEquals(1, bitVector.get(i), desc + i); } for (int i = start + count; i < length; i++) { assertTrue(bitVector.isNull(i), desc + i); } } } @Test public void testBitVectorHashCode() { final int size = 6; try (final BitVector vector = new BitVector(EMPTY_SCHEMA_PATH, allocator)) { ValueVectorDataPopulator.setVector(vector, 0, 1, null, 0, 1, null); int[] hashCodes = new int[size]; IntStream.range(0, size).forEach(i -> hashCodes[i] = vector.hashCode(i)); assertTrue(hashCodes[0] == hashCodes[3]); assertTrue(hashCodes[1] == hashCodes[4]); assertTrue(hashCodes[2] == hashCodes[5]); assertFalse(hashCodes[0] == hashCodes[1]); assertFalse(hashCodes[0] == hashCodes[2]); assertFalse(hashCodes[1] == hashCodes[2]); MurmurHasher hasher = new MurmurHasher(); IntStream.range(0, size).forEach(i -> hashCodes[i] = vector.hashCode(i, hasher)); assertTrue(hashCodes[0] == hashCodes[3]); assertTrue(hashCodes[1] == hashCodes[4]); assertTrue(hashCodes[2] == hashCodes[5]); assertFalse(hashCodes[0] == hashCodes[1]); assertFalse(hashCodes[0] == hashCodes[2]); assertFalse(hashCodes[1] == hashCodes[2]); } } @Test public void testGetTransferPairWithField() { final BitVector fromVector = new BitVector(EMPTY_SCHEMA_PATH, allocator); final TransferPair transferPair = fromVector.getTransferPair(fromVector.getField(), allocator); final BitVector toVector = (BitVector) transferPair.getTo(); // Field inside a new vector created by reusing a field should be the same in memory as the // original field. assertSame(fromVector.getField(), toVector.getField()); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestBitVectorHelper.java000066400000000000000000000261311515451615200323020ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertTrue; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.memory.util.MemoryUtil; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.junit.jupiter.api.Test; public class TestBitVectorHelper { @Test public void testGetNullCount() throws Exception { try (BufferAllocator root = new RootAllocator()) { // test case 1, 1 null value for 0b110 ArrowBuf validityBuffer = root.buffer(3); // we set validity buffer to be 0b10110, but only have 3 items with 1st item is null validityBuffer.setByte(0, 0b10110); // we will only consider 0b110 here, since we only 3 items and only one is null int count = BitVectorHelper.getNullCount(validityBuffer, 3); assertEquals(count, 1); validityBuffer.close(); // test case 2, no null value for 0xFF validityBuffer = root.buffer(8); validityBuffer.setByte(0, 0xFF); count = BitVectorHelper.getNullCount(validityBuffer, 8); assertEquals(0, count); validityBuffer.close(); // test case 3, 1 null value for 0x7F validityBuffer = root.buffer(8); validityBuffer.setByte(0, 0x7F); count = BitVectorHelper.getNullCount(validityBuffer, 8); assertEquals(1, count); validityBuffer.close(); // test case 4, validity buffer has multiple bytes, 11 items validityBuffer = root.buffer(11); validityBuffer.setByte(0, 0b10101010); validityBuffer.setByte(1, 0b01010101); count = BitVectorHelper.getNullCount(validityBuffer, 11); assertEquals(5, count); validityBuffer.close(); } } @Test public void testAllBitsNull() { final int bufferLength = 32 * 1024; try (RootAllocator allocator = new RootAllocator(bufferLength); ArrowBuf validityBuffer = allocator.buffer(bufferLength)) { validityBuffer.setZero(0, bufferLength); int bitLength = 1024; assertTrue(BitVectorHelper.checkAllBitsEqualTo(validityBuffer, bitLength, false)); bitLength = 1027; assertTrue(BitVectorHelper.checkAllBitsEqualTo(validityBuffer, bitLength, false)); validityBuffer.setZero(0, bufferLength); bitLength = 1025; BitVectorHelper.setBit(validityBuffer, 12); assertFalse(BitVectorHelper.checkAllBitsEqualTo(validityBuffer, bitLength, false)); validityBuffer.setZero(0, bufferLength); bitLength = 1025; BitVectorHelper.setBit(validityBuffer, 1024); assertFalse(BitVectorHelper.checkAllBitsEqualTo(validityBuffer, bitLength, false)); validityBuffer.setZero(0, bufferLength); bitLength = 1026; BitVectorHelper.setBit(validityBuffer, 1024); assertFalse(BitVectorHelper.checkAllBitsEqualTo(validityBuffer, bitLength, false)); validityBuffer.setZero(0, bufferLength); bitLength = 1027; BitVectorHelper.setBit(validityBuffer, 1025); assertFalse(BitVectorHelper.checkAllBitsEqualTo(validityBuffer, bitLength, false)); validityBuffer.setZero(0, bufferLength); bitLength = 1031; BitVectorHelper.setBit(validityBuffer, 1029); BitVectorHelper.setBit(validityBuffer, 1030); assertFalse(BitVectorHelper.checkAllBitsEqualTo(validityBuffer, bitLength, false)); } } @Test public void testAllBitsSet() { final int bufferLength = 32 * 1024; try (RootAllocator allocator = new RootAllocator(bufferLength); ArrowBuf validityBuffer = allocator.buffer(bufferLength)) { MemoryUtil.setMemory(validityBuffer.memoryAddress(), bufferLength, (byte) -1); int bitLength = 1024; assertTrue(BitVectorHelper.checkAllBitsEqualTo(validityBuffer, bitLength, true)); bitLength = 1028; assertTrue(BitVectorHelper.checkAllBitsEqualTo(validityBuffer, bitLength, true)); MemoryUtil.setMemory(validityBuffer.memoryAddress(), bufferLength, (byte) -1); bitLength = 1025; BitVectorHelper.unsetBit(validityBuffer, 12); assertFalse(BitVectorHelper.checkAllBitsEqualTo(validityBuffer, bitLength, true)); MemoryUtil.setMemory(validityBuffer.memoryAddress(), bufferLength, (byte) -1); bitLength = 1025; BitVectorHelper.unsetBit(validityBuffer, 1024); assertFalse(BitVectorHelper.checkAllBitsEqualTo(validityBuffer, bitLength, true)); MemoryUtil.setMemory(validityBuffer.memoryAddress(), bufferLength, (byte) -1); bitLength = 1026; BitVectorHelper.unsetBit(validityBuffer, 1024); assertFalse(BitVectorHelper.checkAllBitsEqualTo(validityBuffer, bitLength, true)); MemoryUtil.setMemory(validityBuffer.memoryAddress(), bufferLength, (byte) -1); bitLength = 1027; BitVectorHelper.unsetBit(validityBuffer, 1025); assertFalse(BitVectorHelper.checkAllBitsEqualTo(validityBuffer, bitLength, true)); MemoryUtil.setMemory(validityBuffer.memoryAddress(), bufferLength, (byte) -1); bitLength = 1031; BitVectorHelper.unsetBit(validityBuffer, 1029); BitVectorHelper.unsetBit(validityBuffer, 1030); assertFalse(BitVectorHelper.checkAllBitsEqualTo(validityBuffer, bitLength, true)); } } @Test public void testConcatBits() { try (RootAllocator allocator = new RootAllocator(1024 * 1024)) { try (ArrowBuf buf1 = allocator.buffer(1024); ArrowBuf buf2 = allocator.buffer(1024); ArrowBuf output = allocator.buffer(1024)) { buf1.setZero(0, buf1.capacity()); buf2.setZero(0, buf2.capacity()); final int maxCount = 100; for (int i = 0; i < maxCount; i++) { if (i % 3 == 0) { BitVectorHelper.setBit(buf1, i); BitVectorHelper.setBit(buf2, i); } } // test the case where the number of bits for both sets are multiples of 8. concatAndVerify(buf1, 40, buf2, 48, output); // only the number of bits in the first set is a multiple of 8 concatAndVerify(buf1, 32, buf2, 47, output); // only the number of bits in the second set is a multiple of 8 concatAndVerify(buf1, 31, buf2, 48, output); // neither set has a size that is a multiple of 8 concatAndVerify(buf1, 27, buf2, 52, output); // the remaining bits in the second set is spread in two bytes concatAndVerify(buf1, 31, buf2, 55, output); } } } @Test public void testConcatBitsInPlace() { try (RootAllocator allocator = new RootAllocator(1024 * 1024)) { try (ArrowBuf buf1 = allocator.buffer(1024); ArrowBuf buf2 = allocator.buffer(1024)) { buf1.setZero(0, buf1.capacity()); buf2.setZero(0, buf2.capacity()); final int maxCount = 100; for (int i = 0; i < maxCount; i++) { if (i % 3 == 0) { BitVectorHelper.setBit(buf1, i); BitVectorHelper.setBit(buf2, i); } } // test the case where the number of bits for both sets are multiples of 8. concatAndVerify(buf1, 40, buf2, 48, buf1); // only the number of bits in the first set is a multiple of 8 concatAndVerify(buf1, 32, buf2, 47, buf1); // only the number of bits in the second set is a multiple of 8 concatAndVerify(buf1, 31, buf2, 48, buf1); // neither set has a size that is a multiple of 8 concatAndVerify(buf1, 27, buf2, 52, buf1); // the remaining bits in the second set is spread in two bytes concatAndVerify(buf1, 31, buf2, 55, buf1); } } } @Test public void testLoadValidityBuffer() { try (RootAllocator allocator = new RootAllocator(1024)) { // if the input validity buffer is all null, we should allocate new memory ArrowFieldNode fieldNode = new ArrowFieldNode(1024, 1024); try (ArrowBuf buf = BitVectorHelper.loadValidityBuffer(fieldNode, null, allocator)) { assertEquals(128, allocator.getAllocatedMemory()); for (int i = 0; i < 128; i++) { assertEquals(0, buf.getByte(i)); } } // should also allocate memory if input validity buffer is all not-null fieldNode = new ArrowFieldNode(1024, 0); try (ArrowBuf buf = BitVectorHelper.loadValidityBuffer(fieldNode, null, allocator)) { assertEquals(128, allocator.getAllocatedMemory()); for (int i = 0; i < 128; i++) { assertEquals((byte) 0xff, buf.getByte(i)); } } // should not allocate memory if input validity buffer is not null, even if all values are // null fieldNode = new ArrowFieldNode(1024, 1024); try (ArrowBuf src = allocator.buffer(128); ArrowBuf dst = BitVectorHelper.loadValidityBuffer(fieldNode, src, allocator)) { assertEquals(128, allocator.getAllocatedMemory()); } // ... similarly if all values are not null fieldNode = new ArrowFieldNode(1024, 0); try (ArrowBuf src = allocator.buffer(128); ArrowBuf dst = BitVectorHelper.loadValidityBuffer(fieldNode, src, allocator)) { assertEquals(128, allocator.getAllocatedMemory()); } // mixed case, input should match output int numNulls = 100; fieldNode = new ArrowFieldNode(1024, numNulls); try (ArrowBuf src = allocator.buffer(128)) { src.setZero(0, src.capacity()); for (int i = 0; i < numNulls; i++) { BitVectorHelper.setBit(src, i); } try (ArrowBuf dst = BitVectorHelper.loadValidityBuffer(fieldNode, src, allocator)) { assertEquals(src.memoryAddress(), dst.memoryAddress()); assertEquals(128, allocator.getAllocatedMemory()); for (int i = 0; i < 1024; i++) { assertEquals(BitVectorHelper.get(src, i), BitVectorHelper.get(dst, i)); } } } } } private void concatAndVerify( ArrowBuf buf1, int count1, ArrowBuf buf2, int count2, ArrowBuf output) { BitVectorHelper.concatBits(buf1, count1, buf2, count2, output); int outputIdx = 0; for (int i = 0; i < count1; i++, outputIdx++) { assertEquals(BitVectorHelper.get(output, outputIdx), BitVectorHelper.get(buf1, i)); } for (int i = 0; i < count2; i++, outputIdx++) { assertEquals(BitVectorHelper.get(output, outputIdx), BitVectorHelper.get(buf2, i)); } } } TestBufferOwnershipTransfer.java000066400000000000000000000112031515451615200337710ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertTrue; import java.nio.charset.StandardCharsets; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReferenceManager; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.CallBack; import org.junit.jupiter.api.Test; public class TestBufferOwnershipTransfer { @Test public void testTransferFixedWidth() { BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); BufferAllocator childAllocator1 = allocator.newChildAllocator("child1", 100000, 100000); BufferAllocator childAllocator2 = allocator.newChildAllocator("child2", 100000, 100000); IntVector v1 = new IntVector("v1", childAllocator1); v1.allocateNew(); v1.setValueCount(4095); long totalAllocatedMemory = childAllocator1.getAllocatedMemory(); IntVector v2 = new IntVector("v2", childAllocator2); v1.makeTransferPair(v2).transfer(); assertEquals(0, childAllocator1.getAllocatedMemory()); assertEquals(totalAllocatedMemory, childAllocator2.getAllocatedMemory()); v1.close(); v2.close(); childAllocator1.close(); childAllocator2.close(); allocator.close(); } @Test public void testTransferVariableWidth() { BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); BufferAllocator childAllocator1 = allocator.newChildAllocator("child1", 100000, 100000); BufferAllocator childAllocator2 = allocator.newChildAllocator("child2", 100000, 100000); VarCharVector v1 = new VarCharVector("v1", childAllocator1); v1.allocateNew(); v1.setSafe(4094, "hello world".getBytes(StandardCharsets.UTF_8), 0, 11); v1.setValueCount(4001); VarCharVector v2 = new VarCharVector("v2", childAllocator2); long memoryBeforeTransfer = childAllocator1.getAllocatedMemory(); v1.makeTransferPair(v2).transfer(); assertEquals(0, childAllocator1.getAllocatedMemory()); assertEquals(memoryBeforeTransfer, childAllocator2.getAllocatedMemory()); v1.close(); v2.close(); childAllocator1.close(); childAllocator2.close(); allocator.close(); } private static class Pointer { T value; } private static CallBack newTriggerCallback(final Pointer trigger) { trigger.value = false; return new CallBack() { @Override public void doWork() { trigger.value = true; } }; } @Test public void emptyListTransferShouldNotTriggerSchemaChange() { final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); final Pointer trigger1 = new Pointer<>(); final Pointer trigger2 = new Pointer<>(); final ListVector v1 = new ListVector( "v1", allocator, FieldType.nullable(ArrowType.Null.INSTANCE), newTriggerCallback(trigger1)); final ListVector v2 = new ListVector( "v2", allocator, FieldType.nullable(ArrowType.Null.INSTANCE), newTriggerCallback(trigger2)); try { // since we are working with empty vectors, their internal // buffers will be allocator.EMPTY which use // ReferenceManager.NO_OP instance and transfer() is not // supported v1.makeTransferPair(v2).transfer(); } catch (Exception e) { assertTrue(e instanceof UnsupportedOperationException); assertTrue(e.getMessage().contains(ReferenceManager.NO_OP_ERROR_MESSAGE)); } assertFalse(trigger1.value); assertFalse(trigger2.value); v1.close(); v2.close(); allocator.close(); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestCopyFrom.java000066400000000000000000001135711515451615200310040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.newVector; 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.assertTrue; import java.math.BigDecimal; import java.nio.charset.StandardCharsets; import java.time.Duration; import java.time.Period; import java.util.Objects; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.types.Types.MinorType; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /* * Tested field types: * * NullableInt * NullableBigInt * NullableFloat4 * NullableFloat8 * NullableBit * NullableDecimal * NullableIntervalDay * NullableIntervalYear * NullableSmallInt * NullableTinyInt * NullableVarChar * NullableTimeMicro * NullableTimeMilli * NullableTimeStamp* */ public class TestCopyFrom { private static final String EMPTY_SCHEMA_PATH = ""; private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test /* NullableVarChar */ public void testCopyFromWithNulls() { try (final VarCharVector vector = newVector(VarCharVector.class, EMPTY_SCHEMA_PATH, MinorType.VARCHAR, allocator); final VarCharVector vector2 = newVector(VarCharVector.class, EMPTY_SCHEMA_PATH, MinorType.VARCHAR, allocator)) { vector.allocateNew(); assertTrue(vector.getValueCapacity() >= 1); assertEquals(0, vector.getValueCount()); int initialCapacity = vector.getValueCapacity(); for (int i = 0; i < initialCapacity; i++) { if (i % 3 == 0) { continue; } byte[] b = Integer.toString(i).getBytes(StandardCharsets.UTF_8); vector.setSafe(i, b, 0, b.length); } /* NO reAlloc() should have happened in setSafe() */ int capacity = vector.getValueCapacity(); assertEquals(initialCapacity, capacity); vector.setValueCount(initialCapacity); for (int i = 0; i < initialCapacity; i++) { if (i % 3 == 0) { assertNull(vector.getObject(i)); } else { assertEquals( Integer.toString(i), vector.getObject(i).toString(), "unexpected value at index: " + i); } } vector2.setInitialCapacity(initialCapacity); vector2.allocateNew(); capacity = vector2.getValueCapacity(); assertEquals(initialCapacity, capacity); for (int i = 0; i < initialCapacity; i++) { vector2.copyFromSafe(i, i, vector); if (i % 3 == 0) { assertNull(vector2.getObject(i)); } else { assertEquals( Integer.toString(i), vector2.getObject(i).toString(), "unexpected value at index: " + i); } } /* NO reAlloc() should have happened in copyFrom */ capacity = vector2.getValueCapacity(); assertEquals(initialCapacity, capacity); vector2.setValueCount(initialCapacity); for (int i = 0; i < initialCapacity; i++) { if (i % 3 == 0) { assertNull(vector2.getObject(i)); } else { assertEquals( Integer.toString(i), vector2.getObject(i).toString(), "unexpected value at index: " + i); } } } } @Test /* NullableVarChar */ public void testCopyFromWithNulls1() { try (final VarCharVector vector = newVector(VarCharVector.class, EMPTY_SCHEMA_PATH, MinorType.VARCHAR, allocator); final VarCharVector vector2 = newVector(VarCharVector.class, EMPTY_SCHEMA_PATH, MinorType.VARCHAR, allocator)) { vector.allocateNew(); assertTrue(vector.getValueCapacity() >= 1); assertEquals(0, vector.getValueCount()); int initialCapacity = vector.getValueCapacity(); for (int i = 0; i < initialCapacity; i++) { if (i % 3 == 0) { continue; } byte[] b = Integer.toString(i).getBytes(StandardCharsets.UTF_8); vector.setSafe(i, b, 0, b.length); } /* NO reAlloc() should have happened in setSafe() */ int capacity = vector.getValueCapacity(); assertEquals(initialCapacity, capacity); vector.setValueCount(initialCapacity); for (int i = 0; i < initialCapacity; i++) { if (i % 3 == 0) { assertNull(vector.getObject(i)); } else { assertEquals( Integer.toString(i), vector.getObject(i).toString(), "unexpected value at index: " + i); } } /* set lesser initial capacity than actually needed * to trigger reallocs in copyFromSafe() */ vector2.allocateNew((initialCapacity / 4) * 10, initialCapacity / 4); capacity = vector2.getValueCapacity(); assertTrue(capacity >= initialCapacity / 4); assertTrue(capacity < initialCapacity / 2); for (int i = 0; i < initialCapacity; i++) { vector2.copyFromSafe(i, i, vector); if (i % 3 == 0) { assertNull(vector2.getObject(i)); } else { assertEquals( Integer.toString(i), vector2.getObject(i).toString(), "unexpected value at index: " + i); } } /* 2 reAllocs should have happened in copyFromSafe() */ capacity = vector2.getValueCapacity(); assertTrue(capacity >= initialCapacity); vector2.setValueCount(initialCapacity); for (int i = 0; i < initialCapacity; i++) { if (i % 3 == 0) { assertNull(vector2.getObject(i)); } else { assertEquals( Integer.toString(i), vector2.getObject(i).toString(), "unexpected value at index: " + i); } } } } @Test /* IntVector */ public void testCopyFromWithNulls2() { try (final IntVector vector1 = new IntVector(EMPTY_SCHEMA_PATH, allocator); final IntVector vector2 = new IntVector(EMPTY_SCHEMA_PATH, allocator)) { vector1.allocateNew(); assertTrue(vector1.getValueCapacity() >= vector1.INITIAL_VALUE_ALLOCATION); assertEquals(0, vector1.getValueCount()); int initialCapacity = vector1.getValueCapacity(); for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { continue; } vector1.setSafe(i, 1000 + i); } vector1.setValueCount(initialCapacity); /* No realloc should have happened in setSafe or * setValueCount */ assertEquals(initialCapacity, vector1.getValueCapacity()); assertEquals(initialCapacity, vector1.getValueCount()); for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { assertNull(vector1.getObject(i)); } else { assertEquals(1000 + i, vector1.get(i), "unexpected value at index: " + i); } } /* set lesser initial capacity than actually needed * to trigger reallocs in copyFromSafe() */ vector2.allocateNew(initialCapacity / 4); assertTrue(vector2.getValueCapacity() >= initialCapacity / 4); assertTrue(vector2.getValueCapacity() < initialCapacity / 2); for (int i = 0; i < initialCapacity; i++) { vector2.copyFromSafe(i, i, vector1); } /* 2 realloc should have happened in copyFromSafe() */ assertTrue(vector2.getValueCapacity() >= initialCapacity); vector2.setValueCount(initialCapacity * 2); /* setValueCount() should have done another realloc */ assertEquals(initialCapacity * 2, vector2.getValueCount()); assertTrue(vector2.getValueCapacity() >= initialCapacity * 2); /* check vector data after copy and realloc */ for (int i = 0; i < initialCapacity * 2; i++) { if (((i & 1) == 0) || (i >= initialCapacity)) { assertNull(vector2.getObject(i)); } else { assertEquals(1000 + i, vector2.get(i), "unexpected value at index: " + i); } } } } @Test /* BigIntVector */ public void testCopyFromWithNulls3() { try (final BigIntVector vector1 = new BigIntVector(EMPTY_SCHEMA_PATH, allocator); final BigIntVector vector2 = new BigIntVector(EMPTY_SCHEMA_PATH, allocator)) { vector1.allocateNew(); assertTrue(vector1.getValueCapacity() >= vector1.INITIAL_VALUE_ALLOCATION); assertEquals(0, vector1.getValueCount()); int initialCapacity = vector1.getValueCapacity(); for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { continue; } vector1.setSafe(i, 10000000000L + (long) i); } vector1.setValueCount(initialCapacity); /* No realloc should have happened in setSafe or * setValueCount */ assertEquals(initialCapacity, vector1.getValueCapacity()); assertEquals(initialCapacity, vector1.getValueCount()); for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { assertNull(vector1.getObject(i)); } else { assertEquals(10000000000L + (long) i, vector1.get(i), "unexpected value at index: " + i); } } /* set lesser initial capacity than actually needed * to trigger reallocs in copyFromSafe() */ vector2.allocateNew(initialCapacity / 4); assertTrue(vector2.getValueCapacity() >= initialCapacity / 4); assertTrue(vector2.getValueCapacity() < initialCapacity / 2); for (int i = 0; i < initialCapacity; i++) { vector2.copyFromSafe(i, i, vector1); } /* 2 realloc should have happened in copyFromSafe() */ assertTrue(vector2.getValueCapacity() >= initialCapacity); vector2.setValueCount(initialCapacity * 2); /* setValueCount() should have done another realloc */ assertEquals(initialCapacity * 2, vector2.getValueCount()); assertTrue(vector2.getValueCapacity() >= initialCapacity * 2); /* check vector data after copy and realloc */ for (int i = 0; i < initialCapacity * 2; i++) { if (((i & 1) == 0) || (i >= initialCapacity)) { assertNull(vector2.getObject(i)); } else { assertEquals(10000000000L + (long) i, vector2.get(i), "unexpected value at index: " + i); } } } } @Test /* BitVector */ public void testCopyFromWithNulls4() { try (final BitVector vector1 = new BitVector(EMPTY_SCHEMA_PATH, allocator); final BitVector vector2 = new BitVector(EMPTY_SCHEMA_PATH, allocator)) { vector1.setInitialCapacity(4096); vector1.allocateNew(); assertEquals(4096, vector1.getValueCapacity()); assertEquals(0, vector1.getValueCount()); int counter = 0; for (int i = 0; i < 4096; i++) { if ((i & 1) == 0) { continue; } if ((counter & 1) == 0) { vector1.setSafe(i, 1); } else { vector1.setSafe(i, 0); } counter++; } vector1.setValueCount(4096); /* No realloc should have happened in setSafe or * setValueCount */ assertEquals(4096, vector1.getValueCapacity()); assertEquals(4096, vector1.getValueCount()); counter = 0; for (int i = 0; i < 4096; i++) { if ((i & 1) == 0) { assertNull(vector1.getObject(i)); } else { if ((counter & 1) == 0) { assertTrue(vector1.getObject(i)); } else { assertFalse(vector1.getObject(i)); } counter++; } } /* set lesser initial capacity than actually needed * to trigger reallocs in copyFromSafe() */ vector2.allocateNew(1024); assertEquals(1024, vector2.getValueCapacity()); for (int i = 0; i < 4096; i++) { vector2.copyFromSafe(i, i, vector1); } /* 2 realloc should have happened in copyFromSafe() */ assertEquals(4096, vector2.getValueCapacity()); vector2.setValueCount(8192); /* setValueCount() should have done another realloc */ assertEquals(8192, vector2.getValueCount()); assertEquals(8192, vector2.getValueCapacity()); /* check vector data after copy and realloc */ counter = 0; for (int i = 0; i < 8192; i++) { if (((i & 1) == 0) || (i >= 4096)) { assertNull(vector2.getObject(i)); } else { if ((counter & 1) == 0) { assertTrue(vector2.getObject(i)); } else { assertFalse(vector2.getObject(i)); } counter++; } } } } @Test /* Float4Vector */ public void testCopyFromWithNulls5() { try (final Float4Vector vector1 = new Float4Vector(EMPTY_SCHEMA_PATH, allocator); final Float4Vector vector2 = new Float4Vector(EMPTY_SCHEMA_PATH, allocator)) { vector1.allocateNew(); assertTrue(vector1.getValueCapacity() >= vector1.INITIAL_VALUE_ALLOCATION); assertEquals(0, vector1.getValueCount()); int initialCapacity = vector1.getValueCapacity(); for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { continue; } vector1.setSafe(i, 100.25f + (float) i); } vector1.setValueCount(initialCapacity); /* No realloc should have happened in setSafe or * setValueCount */ assertEquals(initialCapacity, vector1.getValueCapacity()); assertEquals(initialCapacity, vector1.getValueCount()); for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { assertNull(vector1.getObject(i)); } else { assertEquals(100.25f + (float) i, vector1.get(i), 0, "unexpected value at index: " + i); } } /* set lesser initial capacity than actually needed * to trigger reallocs in copyFromSafe() */ vector2.allocateNew(initialCapacity / 4); assertTrue(vector2.getValueCapacity() >= initialCapacity / 4); assertTrue(vector2.getValueCapacity() < initialCapacity / 2); for (int i = 0; i < initialCapacity; i++) { vector2.copyFromSafe(i, i, vector1); } /* 2 realloc should have happened in copyFromSafe() */ assertTrue(vector2.getValueCapacity() >= initialCapacity); vector2.setValueCount(initialCapacity * 2); /* setValueCount() should have done another realloc */ assertEquals(initialCapacity * 2, vector2.getValueCount()); assertTrue(vector2.getValueCapacity() >= initialCapacity * 2); /* check vector data after copy and realloc */ for (int i = 0; i < initialCapacity * 2; i++) { if (((i & 1) == 0) || (i >= initialCapacity)) { assertNull(vector2.getObject(i)); } else { assertEquals(100.25f + i * 1.0f, vector2.get(i), 0, "unexpected value at index: " + i); } } } } @Test /* Float8Vector */ public void testCopyFromWithNulls6() { try (final Float8Vector vector1 = new Float8Vector(EMPTY_SCHEMA_PATH, allocator); final Float8Vector vector2 = new Float8Vector(EMPTY_SCHEMA_PATH, allocator)) { vector1.allocateNew(); assertTrue(vector1.getValueCapacity() >= vector1.INITIAL_VALUE_ALLOCATION); assertEquals(0, vector1.getValueCount()); int initialCapacity = vector1.getValueCapacity(); for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { continue; } vector1.setSafe(i, 123456.7865 + (double) i); } vector1.setValueCount(initialCapacity); /* No realloc should have happened in setSafe or * setValueCount */ assertEquals(initialCapacity, vector1.getValueCapacity()); assertEquals(initialCapacity, vector1.getValueCount()); for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { assertNull(vector1.getObject(i)); } else { assertEquals( 123456.7865 + (double) i, vector1.get(i), 0, "unexpected value at index: " + i); } } /* set lesser initial capacity than actually needed * to trigger reallocs in copyFromSafe() */ vector2.allocateNew(initialCapacity / 4); assertTrue(vector2.getValueCapacity() >= initialCapacity / 4); assertTrue(vector2.getValueCapacity() < initialCapacity / 2); for (int i = 0; i < initialCapacity; i++) { vector2.copyFromSafe(i, i, vector1); } /* 2 realloc should have happened in copyFromSafe() */ assertTrue(vector2.getValueCapacity() >= initialCapacity); vector2.setValueCount(initialCapacity * 2); /* setValueCount() should have done another realloc */ assertEquals(initialCapacity * 2, vector2.getValueCount()); assertTrue(vector2.getValueCapacity() >= initialCapacity * 2); /* check vector data after copy and realloc */ for (int i = 0; i < initialCapacity * 2; i++) { if (((i & 1) == 0) || (i >= initialCapacity)) { assertNull(vector2.getObject(i)); } else { assertEquals( 123456.7865 + (double) i, vector2.get(i), 0, "unexpected value at index: " + i); } } } } @Test /* IntervalDayVector */ public void testCopyFromWithNulls7() { try (final IntervalDayVector vector1 = new IntervalDayVector(EMPTY_SCHEMA_PATH, allocator); final IntervalDayVector vector2 = new IntervalDayVector(EMPTY_SCHEMA_PATH, allocator)) { vector1.allocateNew(); assertTrue(vector1.getValueCapacity() >= vector1.INITIAL_VALUE_ALLOCATION); assertEquals(0, vector1.getValueCount()); int initialCapacity = vector1.getValueCapacity(); final int days = 10; final int milliseconds = 10000; for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { continue; } vector1.setSafe(i, days + i, milliseconds + i); } vector1.setValueCount(initialCapacity); /* No realloc should have happened in setSafe or * setValueCount */ assertEquals(initialCapacity, vector1.getValueCapacity()); assertEquals(initialCapacity, vector1.getValueCount()); for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { assertNull(vector1.getObject(i)); } else { final Duration d = vector1.getObject(i); assertEquals(days + i, d.toDays()); assertEquals(milliseconds + i, d.minusDays(days + i).toMillis()); } } /* set lesser initial capacity than actually needed * to trigger reallocs in copyFromSafe() */ vector2.allocateNew(initialCapacity / 4); assertTrue(vector2.getValueCapacity() >= initialCapacity / 4); assertTrue(vector2.getValueCapacity() < initialCapacity / 2); for (int i = 0; i < initialCapacity; i++) { vector2.copyFromSafe(i, i, vector1); } /* 2 realloc should have happened in copyFromSafe() */ assertTrue(vector2.getValueCapacity() >= initialCapacity); vector2.setValueCount(initialCapacity * 2); /* setValueCount() should have done another realloc */ assertEquals(initialCapacity * 2, vector2.getValueCount()); assertTrue(vector2.getValueCapacity() >= initialCapacity * 2); /* check vector data after copy and realloc */ for (int i = 0; i < initialCapacity * 2; i++) { if (((i & 1) == 0) || (i >= initialCapacity)) { assertNull(vector2.getObject(i)); } else { final Duration d = vector2.getObject(i); assertEquals(days + i, d.toDays()); assertEquals(milliseconds + i, d.minusDays(days + i).toMillis()); } } } } @Test /* IntervalYearVector */ public void testCopyFromWithNulls8() { try (final IntervalYearVector vector1 = new IntervalYearVector(EMPTY_SCHEMA_PATH, allocator); final IntervalYearVector vector2 = new IntervalYearVector(EMPTY_SCHEMA_PATH, allocator)) { vector1.allocateNew(); assertTrue(vector1.getValueCapacity() >= vector1.INITIAL_VALUE_ALLOCATION); assertEquals(0, vector1.getValueCount()); int initialCapacity = vector1.getValueCapacity(); final int interval = 30; /* 2 years 6 months */ final Period[] periods = new Period[4096]; for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { continue; } vector1.setSafe(i, interval + i); final int years = (interval + i) / org.apache.arrow.vector.util.DateUtility.yearsToMonths; final int months = (interval + i) % org.apache.arrow.vector.util.DateUtility.yearsToMonths; periods[i] = Period.ofYears(years).plusMonths(months).normalized(); } vector1.setValueCount(initialCapacity); /* No realloc should have happened in setSafe or * setValueCount */ assertEquals(initialCapacity, vector1.getValueCapacity()); assertEquals(initialCapacity, vector1.getValueCount()); for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { assertNull(vector1.getObject(i)); } else { final Period p = vector1.getObject(i).normalized(); assertEquals(interval + i, vector1.get(i)); assertEquals(periods[i], p); } } /* set lesser initial capacity than actually needed * to trigger reallocs in copyFromSafe() */ vector2.allocateNew(initialCapacity / 4); assertTrue(vector2.getValueCapacity() >= initialCapacity / 4); assertTrue(vector2.getValueCapacity() < initialCapacity / 2); for (int i = 0; i < initialCapacity; i++) { vector2.copyFromSafe(i, i, vector1); } /* 2 realloc should have happened in copyFromSafe() */ assertTrue(vector2.getValueCapacity() >= initialCapacity); vector2.setValueCount(initialCapacity * 2); /* setValueCount() should have done another realloc */ assertEquals(initialCapacity * 2, vector2.getValueCount()); assertTrue(vector2.getValueCapacity() >= initialCapacity * 2); /* check vector data after copy and realloc */ for (int i = 0; i < initialCapacity * 2; i++) { if (((i & 1) == 0) || (i >= initialCapacity)) { assertNull(vector2.getObject(i)); } else { final Period p = vector2.getObject(i).normalized(); assertEquals(periods[i], p); } } } } @Test /* SmallIntVector */ public void testCopyFromWithNulls9() { try (final SmallIntVector vector1 = new SmallIntVector(EMPTY_SCHEMA_PATH, allocator); final SmallIntVector vector2 = new SmallIntVector(EMPTY_SCHEMA_PATH, allocator)) { vector1.allocateNew(); assertTrue(vector1.getValueCapacity() >= vector1.INITIAL_VALUE_ALLOCATION); assertEquals(0, vector1.getValueCount()); int initialCapacity = vector1.getValueCapacity(); final short val = 1000; for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { continue; } vector1.setSafe(i, val + (short) i); } vector1.setValueCount(initialCapacity); /* No realloc should have happened in setSafe or * setValueCount */ assertEquals(initialCapacity, vector1.getValueCapacity()); assertEquals(initialCapacity, vector1.getValueCount()); for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { assertNull(vector1.getObject(i)); } else { assertEquals(val + (short) i, vector1.get(i), "unexpected value at index: " + i); } } /* set lesser initial capacity than actually needed * to trigger reallocs in copyFromSafe() */ vector2.allocateNew(initialCapacity / 4); assertTrue(vector2.getValueCapacity() >= initialCapacity / 4); assertTrue(vector2.getValueCapacity() < initialCapacity / 2); for (int i = 0; i < initialCapacity; i++) { vector2.copyFromSafe(i, i, vector1); } /* 2 realloc should have happened in copyFromSafe() */ assertTrue(vector2.getValueCapacity() >= initialCapacity); vector2.setValueCount(initialCapacity * 2); /* setValueCount() should have done another realloc */ assertEquals(initialCapacity * 2, vector2.getValueCount()); assertTrue(vector2.getValueCapacity() >= initialCapacity * 2); /* check vector data after copy and realloc */ for (int i = 0; i < initialCapacity * 2; i++) { if (((i & 1) == 0) || (i >= initialCapacity)) { assertNull(vector2.getObject(i)); } else { assertEquals(val + (short) i, vector2.get(i), "unexpected value at index: " + i); } } } } @Test /* TimeMicroVector */ public void testCopyFromWithNulls10() { try (final TimeMicroVector vector1 = new TimeMicroVector(EMPTY_SCHEMA_PATH, allocator); final TimeMicroVector vector2 = new TimeMicroVector(EMPTY_SCHEMA_PATH, allocator)) { vector1.allocateNew(); assertTrue(vector1.getValueCapacity() >= vector1.INITIAL_VALUE_ALLOCATION); assertEquals(0, vector1.getValueCount()); int initialCapacity = vector1.getValueCapacity(); final long val = 100485765432L; for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { continue; } vector1.setSafe(i, val + (long) i); } vector1.setValueCount(initialCapacity); /* No realloc should have happened in setSafe or * setValueCount */ assertEquals(initialCapacity, vector1.getValueCapacity()); assertEquals(initialCapacity, vector1.getValueCount()); for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { assertNull(vector1.getObject(i)); } else { assertEquals(val + (long) i, vector1.get(i), "unexpected value at index: " + i); } } /* set lesser initial capacity than actually needed * to trigger reallocs in copyFromSafe() */ vector2.allocateNew(initialCapacity / 4); assertTrue(vector2.getValueCapacity() >= initialCapacity / 4); assertTrue(vector2.getValueCapacity() < initialCapacity / 2); for (int i = 0; i < initialCapacity; i++) { vector2.copyFromSafe(i, i, vector1); } /* 2 realloc should have happened in copyFromSafe() */ assertTrue(vector2.getValueCapacity() >= initialCapacity); vector2.setValueCount(initialCapacity * 2); /* setValueCount() should have done another realloc */ assertEquals(initialCapacity * 2, vector2.getValueCount()); assertTrue(vector2.getValueCapacity() >= initialCapacity * 2); /* check vector data after copy and realloc */ for (int i = 0; i < initialCapacity * 2; i++) { if (((i & 1) == 0) || (i >= initialCapacity)) { assertNull(vector2.getObject(i)); } else { assertEquals(val + (long) i, vector2.get(i), "unexpected value at index: " + i); } } } } @Test /* TimeMilliVector */ public void testCopyFromWithNulls11() { try (final TimeMilliVector vector1 = new TimeMilliVector(EMPTY_SCHEMA_PATH, allocator); final TimeMilliVector vector2 = new TimeMilliVector(EMPTY_SCHEMA_PATH, allocator)) { vector1.allocateNew(); assertTrue(vector1.getValueCapacity() >= vector1.INITIAL_VALUE_ALLOCATION); assertEquals(0, vector1.getValueCount()); int initialCapacity = vector1.getValueCapacity(); final int val = 1000; for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { continue; } vector1.setSafe(i, val + i); } vector1.setValueCount(initialCapacity); /* No realloc should have happened in setSafe or * setValueCount */ assertEquals(initialCapacity, vector1.getValueCapacity()); assertEquals(initialCapacity, vector1.getValueCount()); for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { assertNull(vector1.getObject(i)); } else { assertEquals(val + i, vector1.get(i), "unexpected value at index: " + i); } } /* set lesser initial capacity than actually needed * to trigger reallocs in copyFromSafe() */ vector2.allocateNew(initialCapacity / 4); assertTrue(vector2.getValueCapacity() >= initialCapacity / 4); assertTrue(vector2.getValueCapacity() < initialCapacity / 2); for (int i = 0; i < initialCapacity; i++) { vector2.copyFromSafe(i, i, vector1); } /* 2 realloc should have happened in copyFromSafe() */ assertTrue(vector2.getValueCapacity() >= initialCapacity); vector2.setValueCount(initialCapacity * 2); /* setValueCount() should have done another realloc */ assertEquals(initialCapacity * 2, vector2.getValueCount()); assertTrue(vector2.getValueCapacity() >= initialCapacity * 2); /* check vector data after copy and realloc */ for (int i = 0; i < initialCapacity * 2; i++) { if (((i & 1) == 0) || (i >= initialCapacity)) { assertNull(vector2.getObject(i)); } else { assertEquals(val + i, vector2.get(i), "unexpected value at index: " + i); } } } } @Test /* TinyIntVector */ public void testCopyFromWithNulls12() { try (final TinyIntVector vector1 = new TinyIntVector(EMPTY_SCHEMA_PATH, allocator); final TinyIntVector vector2 = new TinyIntVector(EMPTY_SCHEMA_PATH, allocator)) { vector1.allocateNew(); assertTrue(vector1.getValueCapacity() >= vector1.INITIAL_VALUE_ALLOCATION); assertEquals(0, vector1.getValueCount()); int initialCapacity = vector1.getValueCapacity(); byte val = -128; for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { continue; } vector1.setSafe(i, val); val++; } vector1.setValueCount(initialCapacity); /* No realloc should have happened in setSafe or * setValueCount */ assertEquals(initialCapacity, vector1.getValueCapacity()); assertEquals(initialCapacity, vector1.getValueCount()); val = -128; for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { assertNull(vector1.getObject(i)); } else { assertEquals(val, vector1.get(i), "unexpected value at index: " + i); val++; } } /* set lesser initial capacity than actually needed * to trigger reallocs in copyFromSafe() */ vector2.allocateNew(initialCapacity / 4); assertTrue(vector2.getValueCapacity() >= initialCapacity / 4); for (int i = 0; i < initialCapacity; i++) { vector2.copyFromSafe(i, i, vector1); } /* 2 realloc should have happened in copyFromSafe() */ assertTrue(vector2.getValueCapacity() >= initialCapacity); vector2.setValueCount(initialCapacity * 2); /* setValueCount() should have done another realloc */ assertEquals(initialCapacity * 2, vector2.getValueCount()); assertTrue(vector2.getValueCapacity() >= initialCapacity * 2); /* check vector data after copy and realloc */ val = -128; for (int i = 0; i < initialCapacity * 2; i++) { if (((i & 1) == 0) || (i >= initialCapacity)) { assertNull(vector2.getObject(i)); } else { assertEquals(val, vector2.get(i), "unexpected value at index: " + i); val++; } } } } @Test /* DecimalVector */ public void testCopyFromWithNulls13() { try (final DecimalVector vector1 = new DecimalVector(EMPTY_SCHEMA_PATH, allocator, 30, 16); final DecimalVector vector2 = new DecimalVector(EMPTY_SCHEMA_PATH, allocator, 30, 16)) { vector1.allocateNew(); assertTrue(vector1.getValueCapacity() >= vector1.INITIAL_VALUE_ALLOCATION); assertEquals(0, vector1.getValueCount()); int initialCapacity = vector1.getValueCapacity(); final double baseValue = 104567897654.87654; final BigDecimal[] decimals = new BigDecimal[4096]; for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { continue; } BigDecimal decimal = new BigDecimal(baseValue + (double) i); vector1.setSafe(i, decimal); decimals[i] = decimal; } vector1.setValueCount(initialCapacity); /* No realloc should have happened in setSafe or * setValueCount */ assertEquals(initialCapacity, vector1.getValueCapacity()); assertEquals(initialCapacity, vector1.getValueCount()); for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { assertNull(vector1.getObject(i)); } else { final BigDecimal decimal = vector1.getObject(i); assertEquals(decimals[i], decimal); } } /* set lesser initial capacity than actually needed * to trigger reallocs in copyFromSafe() */ vector2.allocateNew(initialCapacity / 4); assertTrue(vector2.getValueCapacity() >= initialCapacity / 4); assertTrue(vector2.getValueCapacity() < initialCapacity / 2); for (int i = 0; i < initialCapacity; i++) { vector2.copyFromSafe(i, i, vector1); } /* 2 realloc should have happened in copyFromSafe() */ assertTrue(vector2.getValueCapacity() >= initialCapacity); vector2.setValueCount(initialCapacity * 2); /* setValueCount() should have done another realloc */ assertEquals(initialCapacity * 2, vector2.getValueCount()); assertTrue(vector2.getValueCapacity() >= initialCapacity * 2); /* check vector data after copy and realloc */ for (int i = 0; i < initialCapacity * 2; i++) { if (((i & 1) == 0) || (i >= initialCapacity)) { assertNull(vector2.getObject(i)); } else { final BigDecimal decimal = vector2.getObject(i); assertEquals(decimals[i], decimal); } } } } @Test /* TimeStampVector */ public void testCopyFromWithNulls14() { try (final TimeStampVector vector1 = new TimeStampMicroVector(EMPTY_SCHEMA_PATH, allocator); final TimeStampVector vector2 = new TimeStampMicroVector(EMPTY_SCHEMA_PATH, allocator)) { vector1.allocateNew(); assertTrue(vector1.getValueCapacity() >= vector1.INITIAL_VALUE_ALLOCATION); assertEquals(0, vector1.getValueCount()); int initialCapacity = vector1.getValueCapacity(); final long val = 20145678912L; for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { continue; } vector1.setSafe(i, val + (long) i); } vector1.setValueCount(initialCapacity); /* No realloc should have happened in setSafe or * setValueCount */ assertEquals(initialCapacity, vector1.getValueCapacity()); assertEquals(initialCapacity, vector1.getValueCount()); for (int i = 0; i < initialCapacity; i++) { if ((i & 1) == 0) { assertNull(vector1.getObject(i)); } else { assertEquals(val + (long) i, vector1.get(i), "unexpected value at index: " + i); } } /* set lesser initial capacity than actually needed * to trigger reallocs in copyFromSafe() */ vector2.allocateNew(initialCapacity / 4); assertTrue(vector2.getValueCapacity() >= initialCapacity / 4); assertTrue(vector2.getValueCapacity() < initialCapacity / 2); for (int i = 0; i < initialCapacity; i++) { vector2.copyFromSafe(i, i, vector1); } /* 2 realloc should have happened in copyFromSafe() */ assertTrue(vector2.getValueCapacity() >= initialCapacity); vector2.setValueCount(initialCapacity * 2); /* setValueCount() should have done another realloc */ assertEquals(initialCapacity * 2, vector2.getValueCount()); assertTrue(vector2.getValueCapacity() >= initialCapacity * 2); /* check vector data after copy and realloc */ for (int i = 0; i < initialCapacity * 2; i++) { if (((i & 1) == 0) || (i >= initialCapacity)) { assertNull(vector2.getObject(i)); } else { assertEquals(val + (long) i, vector2.get(i), "unexpected value at index: " + i); } } } } @Test // https://issues.apache.org/jira/browse/ARROW-7837 public void testCopySafeArrow7837() { // this test exposes a bug in `handleSafe` where // it reads a stale index and as a result missed a required resize of the value vector. try (VarCharVector vc1 = new VarCharVector("vc1", allocator); VarCharVector vc2 = new VarCharVector("vc2", allocator); ) { // initial size is carefully set in order to force the second 'copyFromSafe' operation // to trigger a reallocation of the vector. vc2.setInitialCapacity(/*valueCount*/ 20, /*density*/ 0.5); vc1.setSafe(0, "1234567890".getBytes(StandardCharsets.UTF_8)); assertFalse(vc1.isNull(0)); assertEquals("1234567890", Objects.requireNonNull(vc1.getObject(0)).toString()); vc2.copyFromSafe(0, 0, vc1); assertFalse(vc2.isNull(0)); assertEquals("1234567890", Objects.requireNonNull(vc2.getObject(0)).toString()); vc2.copyFromSafe(0, 5, vc1); assertTrue(vc2.isNull(1)); assertTrue(vc2.isNull(2)); assertTrue(vc2.isNull(3)); assertTrue(vc2.isNull(4)); assertFalse(vc2.isNull(5)); assertEquals("1234567890", Objects.requireNonNull(vc2.getObject(5)).toString()); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestDecimal256Vector.java000066400000000000000000000407641515451615200322270ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.math.BigDecimal; import java.math.BigInteger; 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; import org.junit.jupiter.api.Test; public class TestDecimal256Vector { private static long[] intValues; static { intValues = new long[60]; for (int i = 0; i < intValues.length / 2; i++) { intValues[i] = 1L << (i + 1); intValues[2 * i] = -1L * (1 << (i + 1)); } } private int scale = 3; private BufferAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testValuesWriteRead() { try (Decimal256Vector decimalVector = TestUtils.newVector( Decimal256Vector.class, "decimal", new ArrowType.Decimal(10, scale, 256), allocator); ) { try (Decimal256Vector oldConstructor = new Decimal256Vector("decimal", allocator, 10, scale); ) { assertEquals(decimalVector.getField().getType(), oldConstructor.getField().getType()); } decimalVector.allocateNew(); BigDecimal[] values = new BigDecimal[intValues.length]; for (int i = 0; i < intValues.length; i++) { BigDecimal decimal = new BigDecimal(BigInteger.valueOf(intValues[i]), scale); values[i] = decimal; decimalVector.setSafe(i, decimal); } decimalVector.setValueCount(intValues.length); for (int i = 0; i < intValues.length; i++) { BigDecimal value = decimalVector.getObject(i); assertEquals(values[i], value, "unexpected data at index: " + i); } } } @Test public void testDecimal256DifferentScaleAndPrecision() { try (Decimal256Vector decimalVector = TestUtils.newVector( Decimal256Vector.class, "decimal", new ArrowType.Decimal(4, 2, 256), allocator)) { decimalVector.allocateNew(); // test Decimal256 with different scale { BigDecimal decimal = new BigDecimal(BigInteger.valueOf(0), 3); UnsupportedOperationException ue = assertThrows( UnsupportedOperationException.class, () -> decimalVector.setSafe(0, decimal)); assertEquals( "BigDecimal scale must equal that in the Arrow vector: 3 != 2", ue.getMessage()); } // test BigDecimal with larger precision than initialized { BigDecimal decimal = new BigDecimal(BigInteger.valueOf(12345), 2); UnsupportedOperationException ue = assertThrows( UnsupportedOperationException.class, () -> decimalVector.setSafe(0, decimal)); assertEquals( "BigDecimal precision cannot be greater than that in the Arrow vector: 5 > 4", ue.getMessage()); } } } @Test public void testWriteBigEndian() { try (Decimal256Vector decimalVector = TestUtils.newVector( Decimal256Vector.class, "decimal", new ArrowType.Decimal(38, 18, 256), allocator); ) { decimalVector.allocateNew(); BigDecimal decimal1 = new BigDecimal("123456789.000000000000000000"); BigDecimal decimal2 = new BigDecimal("11.123456789123456789"); BigDecimal decimal3 = new BigDecimal("1.000000000000000000"); BigDecimal decimal4 = new BigDecimal("0.111111111000000000"); BigDecimal decimal5 = new BigDecimal("987654321.123456789000000000"); BigDecimal decimal6 = new BigDecimal("222222222222.222222222000000000"); BigDecimal decimal7 = new BigDecimal("7777777777777.666666667000000000"); BigDecimal decimal8 = new BigDecimal("1212121212.343434343000000000"); byte[] decimalValue1 = decimal1.unscaledValue().toByteArray(); byte[] decimalValue2 = decimal2.unscaledValue().toByteArray(); byte[] decimalValue3 = decimal3.unscaledValue().toByteArray(); byte[] decimalValue4 = decimal4.unscaledValue().toByteArray(); byte[] decimalValue5 = decimal5.unscaledValue().toByteArray(); byte[] decimalValue6 = decimal6.unscaledValue().toByteArray(); byte[] decimalValue7 = decimal7.unscaledValue().toByteArray(); byte[] decimalValue8 = decimal8.unscaledValue().toByteArray(); decimalVector.setBigEndian(0, decimalValue1); decimalVector.setBigEndian(1, decimalValue2); decimalVector.setBigEndian(2, decimalValue3); decimalVector.setBigEndian(3, decimalValue4); decimalVector.setBigEndian(4, decimalValue5); decimalVector.setBigEndian(5, decimalValue6); decimalVector.setBigEndian(6, decimalValue7); decimalVector.setBigEndian(7, decimalValue8); decimalVector.setValueCount(8); assertEquals(8, decimalVector.getValueCount()); assertEquals(decimal1, decimalVector.getObject(0)); assertEquals(decimal2, decimalVector.getObject(1)); assertEquals(decimal3, decimalVector.getObject(2)); assertEquals(decimal4, decimalVector.getObject(3)); assertEquals(decimal5, decimalVector.getObject(4)); assertEquals(decimal6, decimalVector.getObject(5)); assertEquals(decimal7, decimalVector.getObject(6)); assertEquals(decimal8, decimalVector.getObject(7)); } } @Test public void testLongReadWrite() { try (Decimal256Vector decimalVector = TestUtils.newVector( Decimal256Vector.class, "decimal", new ArrowType.Decimal(38, 0, 256), allocator)) { decimalVector.allocateNew(); long[] longValues = {0L, -2L, Long.MAX_VALUE, Long.MIN_VALUE, 187L}; for (int i = 0; i < longValues.length; ++i) { decimalVector.set(i, longValues[i]); } decimalVector.setValueCount(longValues.length); for (int i = 0; i < longValues.length; ++i) { assertEquals(new BigDecimal(longValues[i]), decimalVector.getObject(i)); } } } @Test public void testBigDecimalReadWrite() { try (Decimal256Vector decimalVector = TestUtils.newVector( Decimal256Vector.class, "decimal", new ArrowType.Decimal(38, 9, 256), allocator); ) { decimalVector.allocateNew(); BigDecimal decimal1 = new BigDecimal("123456789.000000000"); BigDecimal decimal2 = new BigDecimal("11.123456789"); BigDecimal decimal3 = new BigDecimal("1.000000000"); BigDecimal decimal4 = new BigDecimal("-0.111111111"); BigDecimal decimal5 = new BigDecimal("-987654321.123456789"); BigDecimal decimal6 = new BigDecimal("-222222222222.222222222"); BigDecimal decimal7 = new BigDecimal("7777777777777.666666667"); BigDecimal decimal8 = new BigDecimal("1212121212.343434343"); decimalVector.set(0, decimal1); decimalVector.set(1, decimal2); decimalVector.set(2, decimal3); decimalVector.set(3, decimal4); decimalVector.set(4, decimal5); decimalVector.set(5, decimal6); decimalVector.set(6, decimal7); decimalVector.set(7, decimal8); decimalVector.setValueCount(8); assertEquals(8, decimalVector.getValueCount()); assertEquals(decimal1, decimalVector.getObject(0)); assertEquals(decimal2, decimalVector.getObject(1)); assertEquals(decimal3, decimalVector.getObject(2)); assertEquals(decimal4, decimalVector.getObject(3)); assertEquals(decimal5, decimalVector.getObject(4)); assertEquals(decimal6, decimalVector.getObject(5)); assertEquals(decimal7, decimalVector.getObject(6)); assertEquals(decimal8, decimalVector.getObject(7)); } } /** * Test {@link Decimal256Vector#setBigEndian(int, byte[])} which takes BE layout input and stores * in native-endian (NE) layout. Cases to cover: input byte array in different lengths in range * [1-16] and negative values. */ @Test public void decimalBE2NE() { try (Decimal256Vector decimalVector = TestUtils.newVector( Decimal256Vector.class, "decimal", new ArrowType.Decimal(23, 2, 256), allocator)) { decimalVector.allocateNew(); BigInteger[] testBigInts = new BigInteger[] { new BigInteger("0"), new BigInteger("-1"), new BigInteger("23"), new BigInteger("234234"), new BigInteger("-234234234"), new BigInteger("234234234234"), new BigInteger("-56345345345345"), new BigInteger( "2982346298346289346293467923465345634500"), // converts to 16+ byte array new BigInteger( "-389457298347598237459832459823434653600"), // converts to 16+ byte array new BigInteger("-345345"), new BigInteger("754533") }; int insertionIdx = 0; insertionIdx++; // insert a null for (BigInteger val : testBigInts) { decimalVector.setBigEndian(insertionIdx++, val.toByteArray()); } insertionIdx++; // insert a null // insert a zero length buffer decimalVector.setBigEndian(insertionIdx++, new byte[0]); // Try inserting a buffer larger than 33 bytes and expect a failure final int insertionIdxCapture = insertionIdx; IllegalArgumentException ex = assertThrows( IllegalArgumentException.class, () -> decimalVector.setBigEndian(insertionIdxCapture, new byte[33])); assertTrue( ex.getMessage().equals("Invalid decimal value length. Valid length in [1 - 32], got 33")); decimalVector.setValueCount(insertionIdx); // retrieve values and check if they are correct int outputIdx = 0; assertTrue(decimalVector.isNull(outputIdx++)); for (BigInteger expected : testBigInts) { final BigDecimal actual = decimalVector.getObject(outputIdx++); assertEquals(expected, actual.unscaledValue()); } assertTrue(decimalVector.isNull(outputIdx++)); assertEquals(BigInteger.valueOf(0), decimalVector.getObject(outputIdx).unscaledValue()); } } @Test public void setUsingArrowBufOfLEInts() { try (Decimal256Vector decimalVector = TestUtils.newVector( Decimal256Vector.class, "decimal", new ArrowType.Decimal(5, 2, 256), allocator); ArrowBuf buf = allocator.buffer(8); ) { decimalVector.allocateNew(); // add a positive value equivalent to 705.32 int val = 70532; buf.setInt(0, val); decimalVector.setSafe(0, 0, buf, 4); // add a -ve value equivalent to -705.32 val = -70532; buf.setInt(4, val); decimalVector.setSafe(1, 4, buf, 4); decimalVector.setValueCount(2); BigDecimal[] expectedValues = new BigDecimal[] {BigDecimal.valueOf(705.32), BigDecimal.valueOf(-705.32)}; for (int i = 0; i < 2; i++) { BigDecimal value = decimalVector.getObject(i); assertEquals(expectedValues[i], value); } } } @Test public void setUsingArrowLongLEBytes() { try (Decimal256Vector decimalVector = TestUtils.newVector( Decimal256Vector.class, "decimal", new ArrowType.Decimal(18, 0, 256), allocator); ArrowBuf buf = allocator.buffer(16); ) { decimalVector.allocateNew(); long val = Long.MAX_VALUE; buf.setLong(0, val); decimalVector.setSafe(0, 0, buf, 8); val = Long.MIN_VALUE; buf.setLong(8, val); decimalVector.setSafe(1, 8, buf, 8); decimalVector.setValueCount(2); BigDecimal[] expectedValues = new BigDecimal[] {BigDecimal.valueOf(Long.MAX_VALUE), BigDecimal.valueOf(Long.MIN_VALUE)}; for (int i = 0; i < 2; i++) { BigDecimal value = decimalVector.getObject(i); assertEquals(expectedValues[i], value); } } } @Test public void setUsingArrowBufOfBEBytes() { try (Decimal256Vector decimalVector = TestUtils.newVector( Decimal256Vector.class, "decimal", new ArrowType.Decimal(5, 2, 256), allocator); ArrowBuf buf = allocator.buffer(9); ) { BigDecimal[] expectedValues = new BigDecimal[] { BigDecimal.valueOf(705.32), BigDecimal.valueOf(-705.32), BigDecimal.valueOf(705.32) }; verifyWritingArrowBufWithBigEndianBytes(decimalVector, buf, expectedValues, 3); } try (Decimal256Vector decimalVector = TestUtils.newVector( Decimal256Vector.class, "decimal", new ArrowType.Decimal(43, 2, 256), allocator); ArrowBuf buf = allocator.buffer(45); ) { BigDecimal[] expectedValues = new BigDecimal[] { new BigDecimal("29823462983462893462934679234653450000000.63"), new BigDecimal("-2982346298346289346293467923465345.63"), new BigDecimal("2982346298346289346293467923465345.63") }; verifyWritingArrowBufWithBigEndianBytes(decimalVector, buf, expectedValues, 15); } } @Test public void testGetTransferPairWithField() { final Decimal256Vector fromVector = new Decimal256Vector("decimal", allocator, 10, scale); final TransferPair transferPair = fromVector.getTransferPair(fromVector.getField(), allocator); final Decimal256Vector toVector = (Decimal256Vector) transferPair.getTo(); // Field inside a new vector created by reusing a field should be the same in memory as the // original field. 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(); for (int i = 0; i < expectedValues.length; i++) { byte[] bigEndianBytes = expectedValues[i].unscaledValue().toByteArray(); buf.setBytes(length * i, bigEndianBytes, 0, bigEndianBytes.length); decimalVector.setBigEndianSafe(i, length * i, buf, bigEndianBytes.length); } decimalVector.setValueCount(3); for (int i = 0; i < expectedValues.length; i++) { BigDecimal value = decimalVector.getObject(i); assertEquals(expectedValues[i], value); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestDecimalVector.java000066400000000000000000000375551515451615200317760ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; import java.math.BigDecimal; import java.math.BigInteger; 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; import org.junit.jupiter.api.Test; public class TestDecimalVector { private static long[] intValues; static { intValues = new long[60]; for (int i = 0; i < intValues.length / 2; i++) { intValues[i] = 1L << (i + 1); intValues[2 * i] = -1L * (1 << (i + 1)); } } private int scale = 3; private BufferAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testValuesWriteRead() { try (DecimalVector decimalVector = TestUtils.newVector( DecimalVector.class, "decimal", new ArrowType.Decimal(10, scale, 128), allocator); ) { try (DecimalVector oldConstructor = new DecimalVector("decimal", allocator, 10, scale); ) { assertEquals(decimalVector.getField().getType(), oldConstructor.getField().getType()); } decimalVector.allocateNew(); BigDecimal[] values = new BigDecimal[intValues.length]; for (int i = 0; i < intValues.length; i++) { BigDecimal decimal = new BigDecimal(BigInteger.valueOf(intValues[i]), scale); values[i] = decimal; decimalVector.setSafe(i, decimal); } decimalVector.setValueCount(intValues.length); for (int i = 0; i < intValues.length; i++) { BigDecimal value = decimalVector.getObject(i); assertEquals(values[i], value, "unexpected data at index: " + i); } } } @Test public void testBigDecimalDifferentScaleAndPrecision() { try (DecimalVector decimalVector = TestUtils.newVector( DecimalVector.class, "decimal", new ArrowType.Decimal(4, 2, 128), allocator); ) { decimalVector.allocateNew(); // test BigDecimal with different scale boolean hasError = false; try { BigDecimal decimal = new BigDecimal(BigInteger.valueOf(0), 3); decimalVector.setSafe(0, decimal); } catch (UnsupportedOperationException ue) { hasError = true; } finally { assertTrue(hasError); } // test BigDecimal with larger precision than initialized hasError = false; try { BigDecimal decimal = new BigDecimal(BigInteger.valueOf(12345), 2); decimalVector.setSafe(0, decimal); } catch (UnsupportedOperationException ue) { hasError = true; } finally { assertTrue(hasError); } } } @Test public void testWriteBigEndian() { try (DecimalVector decimalVector = TestUtils.newVector( DecimalVector.class, "decimal", new ArrowType.Decimal(38, 9, 128), allocator); ) { decimalVector.allocateNew(); BigDecimal decimal1 = new BigDecimal("123456789.000000000"); BigDecimal decimal2 = new BigDecimal("11.123456789"); BigDecimal decimal3 = new BigDecimal("1.000000000"); BigDecimal decimal4 = new BigDecimal("0.111111111"); BigDecimal decimal5 = new BigDecimal("987654321.123456789"); BigDecimal decimal6 = new BigDecimal("222222222222.222222222"); BigDecimal decimal7 = new BigDecimal("7777777777777.666666667"); BigDecimal decimal8 = new BigDecimal("1212121212.343434343"); byte[] decimalValue1 = decimal1.unscaledValue().toByteArray(); byte[] decimalValue2 = decimal2.unscaledValue().toByteArray(); byte[] decimalValue3 = decimal3.unscaledValue().toByteArray(); byte[] decimalValue4 = decimal4.unscaledValue().toByteArray(); byte[] decimalValue5 = decimal5.unscaledValue().toByteArray(); byte[] decimalValue6 = decimal6.unscaledValue().toByteArray(); byte[] decimalValue7 = decimal7.unscaledValue().toByteArray(); byte[] decimalValue8 = decimal8.unscaledValue().toByteArray(); decimalVector.setBigEndian(0, decimalValue1); decimalVector.setBigEndian(1, decimalValue2); decimalVector.setBigEndian(2, decimalValue3); decimalVector.setBigEndian(3, decimalValue4); decimalVector.setBigEndian(4, decimalValue5); decimalVector.setBigEndian(5, decimalValue6); decimalVector.setBigEndian(6, decimalValue7); decimalVector.setBigEndian(7, decimalValue8); decimalVector.setValueCount(8); assertEquals(8, decimalVector.getValueCount()); assertEquals(decimal1, decimalVector.getObject(0)); assertEquals(decimal2, decimalVector.getObject(1)); assertEquals(decimal3, decimalVector.getObject(2)); assertEquals(decimal4, decimalVector.getObject(3)); assertEquals(decimal5, decimalVector.getObject(4)); assertEquals(decimal6, decimalVector.getObject(5)); assertEquals(decimal7, decimalVector.getObject(6)); assertEquals(decimal8, decimalVector.getObject(7)); } } @Test public void testLongReadWrite() { try (DecimalVector decimalVector = TestUtils.newVector( DecimalVector.class, "decimal", new ArrowType.Decimal(38, 0, 128), allocator)) { decimalVector.allocateNew(); long[] longValues = {0L, -2L, Long.MAX_VALUE, Long.MIN_VALUE, 187L}; for (int i = 0; i < longValues.length; ++i) { decimalVector.set(i, longValues[i]); } decimalVector.setValueCount(longValues.length); for (int i = 0; i < longValues.length; ++i) { assertEquals(new BigDecimal(longValues[i]), decimalVector.getObject(i)); } } } @Test public void testBigDecimalReadWrite() { try (DecimalVector decimalVector = TestUtils.newVector( DecimalVector.class, "decimal", new ArrowType.Decimal(38, 9, 128), allocator); ) { decimalVector.allocateNew(); BigDecimal decimal1 = new BigDecimal("123456789.000000000"); BigDecimal decimal2 = new BigDecimal("11.123456789"); BigDecimal decimal3 = new BigDecimal("1.000000000"); BigDecimal decimal4 = new BigDecimal("-0.111111111"); BigDecimal decimal5 = new BigDecimal("-987654321.123456789"); BigDecimal decimal6 = new BigDecimal("-222222222222.222222222"); BigDecimal decimal7 = new BigDecimal("7777777777777.666666667"); BigDecimal decimal8 = new BigDecimal("1212121212.343434343"); decimalVector.set(0, decimal1); decimalVector.set(1, decimal2); decimalVector.set(2, decimal3); decimalVector.set(3, decimal4); decimalVector.set(4, decimal5); decimalVector.set(5, decimal6); decimalVector.set(6, decimal7); decimalVector.set(7, decimal8); decimalVector.setValueCount(8); assertEquals(8, decimalVector.getValueCount()); assertEquals(decimal1, decimalVector.getObject(0)); assertEquals(decimal2, decimalVector.getObject(1)); assertEquals(decimal3, decimalVector.getObject(2)); assertEquals(decimal4, decimalVector.getObject(3)); assertEquals(decimal5, decimalVector.getObject(4)); assertEquals(decimal6, decimalVector.getObject(5)); assertEquals(decimal7, decimalVector.getObject(6)); assertEquals(decimal8, decimalVector.getObject(7)); } } /** * Test {@link DecimalVector#setBigEndian(int, byte[])} which takes BE layout input and stores in * native-endian (NE) layout. Cases to cover: input byte array in different lengths in range * [1-16] and negative values. */ @Test public void decimalBE2NE() { try (DecimalVector decimalVector = TestUtils.newVector( DecimalVector.class, "decimal", new ArrowType.Decimal(21, 2, 128), allocator)) { decimalVector.allocateNew(); BigInteger[] testBigInts = new BigInteger[] { new BigInteger("0"), new BigInteger("-1"), new BigInteger("23"), new BigInteger("234234"), new BigInteger("-234234234"), new BigInteger("234234234234"), new BigInteger("-56345345345345"), new BigInteger("29823462983462893462934679234653456345"), // converts to 16 byte array new BigInteger("-3894572983475982374598324598234346536"), // converts to 16 byte array new BigInteger("-345345"), new BigInteger("754533") }; int insertionIdx = 0; insertionIdx++; // insert a null for (BigInteger val : testBigInts) { decimalVector.setBigEndian(insertionIdx++, val.toByteArray()); } insertionIdx++; // insert a null // insert a zero length buffer decimalVector.setBigEndian(insertionIdx++, new byte[0]); // Try inserting a buffer larger than 16bytes and expect a failure try { decimalVector.setBigEndian(insertionIdx, new byte[17]); fail("above statement should have failed"); } catch (IllegalArgumentException ex) { assertTrue( ex.getMessage() .equals("Invalid decimal value length. Valid length in [1 - 16], got 17")); } decimalVector.setValueCount(insertionIdx); // retrieve values and check if they are correct int outputIdx = 0; assertTrue(decimalVector.isNull(outputIdx++)); for (BigInteger expected : testBigInts) { final BigDecimal actual = decimalVector.getObject(outputIdx++); assertEquals(expected, actual.unscaledValue()); } assertTrue(decimalVector.isNull(outputIdx++)); assertEquals(BigInteger.valueOf(0), decimalVector.getObject(outputIdx).unscaledValue()); } } @Test public void setUsingArrowBufOfInts() { try (DecimalVector decimalVector = TestUtils.newVector( DecimalVector.class, "decimal", new ArrowType.Decimal(5, 2, 128), allocator); ArrowBuf buf = allocator.buffer(8); ) { decimalVector.allocateNew(); // add a positive value equivalent to 705.32 int val = 70532; buf.setInt(0, val); decimalVector.setSafe(0, 0, buf, 4); // add a -ve value equivalent to -705.32 val = -70532; buf.setInt(4, val); decimalVector.setSafe(1, 4, buf, 4); decimalVector.setValueCount(2); BigDecimal[] expectedValues = new BigDecimal[] {BigDecimal.valueOf(705.32), BigDecimal.valueOf(-705.32)}; for (int i = 0; i < 2; i++) { BigDecimal value = decimalVector.getObject(i); assertEquals(expectedValues[i], value); } } } @Test public void setUsingArrowLongBytes() { try (DecimalVector decimalVector = TestUtils.newVector( DecimalVector.class, "decimal", new ArrowType.Decimal(18, 0, 128), allocator); ArrowBuf buf = allocator.buffer(16); ) { decimalVector.allocateNew(); long val = Long.MAX_VALUE; buf.setLong(0, val); decimalVector.setSafe(0, 0, buf, 8); val = Long.MIN_VALUE; buf.setLong(8, val); decimalVector.setSafe(1, 8, buf, 8); decimalVector.setValueCount(2); BigDecimal[] expectedValues = new BigDecimal[] {BigDecimal.valueOf(Long.MAX_VALUE), BigDecimal.valueOf(Long.MIN_VALUE)}; for (int i = 0; i < 2; i++) { BigDecimal value = decimalVector.getObject(i); assertEquals(expectedValues[i], value); } } } @Test public void setUsingArrowBufOfBEBytes() { try (DecimalVector decimalVector = TestUtils.newVector( DecimalVector.class, "decimal", new ArrowType.Decimal(5, 2, 128), allocator); ArrowBuf buf = allocator.buffer(9); ) { BigDecimal[] expectedValues = new BigDecimal[] { BigDecimal.valueOf(705.32), BigDecimal.valueOf(-705.32), BigDecimal.valueOf(705.32) }; verifyWritingArrowBufWithBigEndianBytes(decimalVector, buf, expectedValues, 3); } try (DecimalVector decimalVector = TestUtils.newVector( DecimalVector.class, "decimal", new ArrowType.Decimal(36, 2, 128), allocator); ArrowBuf buf = allocator.buffer(45); ) { BigDecimal[] expectedValues = new BigDecimal[] { new BigDecimal("2982346298346289346293467923465345.63"), new BigDecimal("-2982346298346289346293467923465345.63"), new BigDecimal("2982346298346289346293467923465345.63") }; verifyWritingArrowBufWithBigEndianBytes(decimalVector, buf, expectedValues, 15); } } @Test public void testGetTransferPairWithField() { final DecimalVector fromVector = new DecimalVector("decimal", allocator, 10, scale); final TransferPair transferPair = fromVector.getTransferPair(fromVector.getField(), allocator); final DecimalVector toVector = (DecimalVector) transferPair.getTo(); // Field inside a new vector created by reusing a field should be the same in memory as the // original field. 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(); for (int i = 0; i < expectedValues.length; i++) { byte[] bigEndianBytes = expectedValues[i].unscaledValue().toByteArray(); buf.setBytes(length * i, bigEndianBytes, 0, bigEndianBytes.length); decimalVector.setBigEndianSafe(i, length * i, buf, bigEndianBytes.length); } decimalVector.setValueCount(3); for (int i = 0; i < expectedValues.length; i++) { BigDecimal value = decimalVector.getObject(i); assertEquals(expectedValues[i], value); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestDenseUnionVector.java000066400000000000000000000623601515451615200324770ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertTrue; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.VectorWithOrdinal; import org.apache.arrow.vector.holders.NullableBigIntHolder; import org.apache.arrow.vector.holders.NullableBitHolder; import org.apache.arrow.vector.holders.NullableFloat4Holder; import org.apache.arrow.vector.holders.NullableIntHolder; import org.apache.arrow.vector.holders.NullableUInt4Holder; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.UnionMode; 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.JsonStringHashMap; import org.apache.arrow.vector.util.Text; 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 TestDenseUnionVector { private static final String EMPTY_SCHEMA_PATH = ""; private BufferAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testDenseUnionVector() throws Exception { final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); uInt4Holder.value = 100; uInt4Holder.isSet = 1; try (DenseUnionVector unionVector = new DenseUnionVector(EMPTY_SCHEMA_PATH, allocator, null, null)) { unionVector.allocateNew(); // write some data byte uint4TypeId = unionVector.registerNewTypeId(Field.nullable("", MinorType.UINT4.getType())); unionVector.setTypeId(0, uint4TypeId); unionVector.setSafe(0, uInt4Holder); unionVector.setTypeId(2, uint4TypeId); unionVector.setSafe(2, uInt4Holder); unionVector.setValueCount(4); // check that what we wrote is correct assertEquals(4, unionVector.getValueCount()); assertEquals(false, unionVector.isNull(0)); assertEquals(100, unionVector.getObject(0)); assertNull(unionVector.getObject(1)); assertEquals(false, unionVector.isNull(2)); assertEquals(100, unionVector.getObject(2)); assertNull(unionVector.getObject(3)); } } @Test public void testSetOffset() { try (DenseUnionVector duv = DenseUnionVector.empty("foo", allocator)) { duv.allocateNew(); byte i32TypeId = duv.registerNewTypeId(Field.notNullable("i32", MinorType.INT.getType())); byte f64TypeId = duv.registerNewTypeId(Field.notNullable("f64", MinorType.FLOAT8.getType())); IntVector i32Vector = ((IntVector) duv.addVector(i32TypeId, new IntVector("i32", allocator))); Float8Vector f64Vector = ((Float8Vector) duv.addVector(f64TypeId, new Float8Vector("f64", allocator))); i32Vector.allocateNew(3); f64Vector.allocateNew(1); duv.setTypeId(0, i32TypeId); duv.setOffset(0, 0); i32Vector.set(0, 42); duv.setTypeId(1, i32TypeId); duv.setOffset(1, 1); i32Vector.set(1, 43); duv.setTypeId(2, f64TypeId); duv.setOffset(2, 0); f64Vector.set(0, 3.14); duv.setTypeId(3, i32TypeId); duv.setOffset(3, 2); i32Vector.set(2, 44); duv.setValueCount(4); assertEquals(42, duv.getObject(0)); assertEquals(43, duv.getObject(1)); assertEquals(3.14, duv.getObject(2)); assertEquals(44, duv.getObject(3)); } } @Test public void testTransfer() throws Exception { try (DenseUnionVector srcVector = new DenseUnionVector(EMPTY_SCHEMA_PATH, allocator, null, null)) { srcVector.allocateNew(); // write some data byte intTypeId = srcVector.registerNewTypeId(Field.nullable("", MinorType.INT.getType())); srcVector.setTypeId(0, intTypeId); srcVector.setSafe(0, newIntHolder(5)); byte bitTypeId = srcVector.registerNewTypeId(Field.nullable("", MinorType.BIT.getType())); srcVector.setTypeId(1, bitTypeId); srcVector.setSafe(1, newBitHolder(false)); srcVector.setTypeId(3, intTypeId); srcVector.setSafe(3, newIntHolder(10)); srcVector.setTypeId(5, bitTypeId); srcVector.setSafe(5, newBitHolder(false)); srcVector.setValueCount(6); try (DenseUnionVector destVector = new DenseUnionVector(EMPTY_SCHEMA_PATH, allocator, null, null)) { TransferPair pair = srcVector.makeTransferPair(destVector); // Creating the transfer should transfer the type of the field at least. assertEquals(srcVector.getField(), destVector.getField()); // transfer pair.transfer(); assertEquals(srcVector.getField(), destVector.getField()); // now check the values are transferred assertEquals(6, destVector.getValueCount()); assertFalse(destVector.isNull(0)); assertEquals(5, destVector.getObject(0)); assertFalse(destVector.isNull(1)); assertEquals(false, destVector.getObject(1)); assertNull(destVector.getObject(2)); assertFalse(destVector.isNull(3)); assertEquals(10, destVector.getObject(3)); assertNull(destVector.getObject(4)); assertFalse(destVector.isNull(5)); assertEquals(false, destVector.getObject(5)); } } } @Test public void testSplitAndTransfer() throws Exception { try (DenseUnionVector sourceVector = new DenseUnionVector(EMPTY_SCHEMA_PATH, allocator, null, null)) { sourceVector.allocateNew(); /* populate the UnionVector */ byte intTypeId = sourceVector.registerNewTypeId(Field.nullable("", MinorType.INT.getType())); sourceVector.setTypeId(0, intTypeId); sourceVector.setSafe(0, newIntHolder(5)); sourceVector.setTypeId(1, intTypeId); sourceVector.setSafe(1, newIntHolder(10)); sourceVector.setTypeId(2, intTypeId); sourceVector.setSafe(2, newIntHolder(15)); sourceVector.setTypeId(3, intTypeId); sourceVector.setSafe(3, newIntHolder(20)); sourceVector.setTypeId(4, intTypeId); sourceVector.setSafe(4, newIntHolder(25)); sourceVector.setTypeId(5, intTypeId); sourceVector.setSafe(5, newIntHolder(30)); sourceVector.setTypeId(6, intTypeId); sourceVector.setSafe(6, newIntHolder(35)); sourceVector.setTypeId(7, intTypeId); sourceVector.setSafe(7, newIntHolder(40)); sourceVector.setTypeId(8, intTypeId); sourceVector.setSafe(8, newIntHolder(45)); sourceVector.setTypeId(9, intTypeId); sourceVector.setSafe(9, newIntHolder(50)); sourceVector.setValueCount(10); /* check the vector output */ assertEquals(10, sourceVector.getValueCount()); assertEquals(false, sourceVector.isNull(0)); assertEquals(5, sourceVector.getObject(0)); assertEquals(false, sourceVector.isNull(1)); assertEquals(10, sourceVector.getObject(1)); assertEquals(false, sourceVector.isNull(2)); assertEquals(15, sourceVector.getObject(2)); assertEquals(false, sourceVector.isNull(3)); assertEquals(20, sourceVector.getObject(3)); assertEquals(false, sourceVector.isNull(4)); assertEquals(25, sourceVector.getObject(4)); assertEquals(false, sourceVector.isNull(5)); assertEquals(30, sourceVector.getObject(5)); assertEquals(false, sourceVector.isNull(6)); assertEquals(35, sourceVector.getObject(6)); assertEquals(false, sourceVector.isNull(7)); assertEquals(40, sourceVector.getObject(7)); assertEquals(false, sourceVector.isNull(8)); assertEquals(45, sourceVector.getObject(8)); assertEquals(false, sourceVector.isNull(9)); assertEquals(50, sourceVector.getObject(9)); try (DenseUnionVector toVector = new DenseUnionVector(EMPTY_SCHEMA_PATH, allocator, null, null)) { toVector.registerNewTypeId(Field.nullable("", MinorType.INT.getType())); final TransferPair transferPair = sourceVector.makeTransferPair(toVector); final int[][] transferLengths = {{0, 3}, {3, 1}, {4, 2}, {6, 1}, {7, 1}, {8, 2}}; for (final int[] transferLength : transferLengths) { final int start = transferLength[0]; final int length = transferLength[1]; transferPair.splitAndTransfer(start, length); /* check the toVector output after doing the splitAndTransfer */ for (int i = 0; i < length; i++) { assertEquals( sourceVector.getObject(start + i), toVector.getObject(i), "Different data at indexes: " + (start + i) + "and " + i); } } } } } @Test public void testSplitAndTransferWithMixedVectors() throws Exception { try (DenseUnionVector sourceVector = new DenseUnionVector(EMPTY_SCHEMA_PATH, allocator, null, null)) { sourceVector.allocateNew(); /* populate the UnionVector */ byte intTypeId = sourceVector.registerNewTypeId(Field.nullable("", MinorType.INT.getType())); sourceVector.setTypeId(0, intTypeId); sourceVector.setSafe(0, newIntHolder(5)); byte float4TypeId = sourceVector.registerNewTypeId(Field.nullable("", MinorType.FLOAT4.getType())); sourceVector.setTypeId(1, float4TypeId); sourceVector.setSafe(1, newFloat4Holder(5.5f)); sourceVector.setTypeId(2, intTypeId); sourceVector.setSafe(2, newIntHolder(10)); sourceVector.setTypeId(3, float4TypeId); sourceVector.setSafe(3, newFloat4Holder(10.5f)); sourceVector.setTypeId(4, intTypeId); sourceVector.setSafe(4, newIntHolder(15)); sourceVector.setTypeId(5, float4TypeId); sourceVector.setSafe(5, newFloat4Holder(15.5f)); sourceVector.setTypeId(6, intTypeId); sourceVector.setSafe(6, newIntHolder(20)); sourceVector.setTypeId(7, float4TypeId); sourceVector.setSafe(7, newFloat4Holder(20.5f)); sourceVector.setTypeId(8, intTypeId); sourceVector.setSafe(8, newIntHolder(30)); sourceVector.setTypeId(9, float4TypeId); sourceVector.setSafe(9, newFloat4Holder(30.5f)); sourceVector.setValueCount(10); /* check the vector output */ assertEquals(10, sourceVector.getValueCount()); assertEquals(false, sourceVector.isNull(0)); assertEquals(5, sourceVector.getObject(0)); assertEquals(false, sourceVector.isNull(1)); assertEquals(5.5f, sourceVector.getObject(1)); assertEquals(false, sourceVector.isNull(2)); assertEquals(10, sourceVector.getObject(2)); assertEquals(false, sourceVector.isNull(3)); assertEquals(10.5f, sourceVector.getObject(3)); assertEquals(false, sourceVector.isNull(4)); assertEquals(15, sourceVector.getObject(4)); assertEquals(false, sourceVector.isNull(5)); assertEquals(15.5f, sourceVector.getObject(5)); assertEquals(false, sourceVector.isNull(6)); assertEquals(20, sourceVector.getObject(6)); assertEquals(false, sourceVector.isNull(7)); assertEquals(20.5f, sourceVector.getObject(7)); assertEquals(false, sourceVector.isNull(8)); assertEquals(30, sourceVector.getObject(8)); assertEquals(false, sourceVector.isNull(9)); assertEquals(30.5f, sourceVector.getObject(9)); try (DenseUnionVector toVector = new DenseUnionVector(EMPTY_SCHEMA_PATH, allocator, null, null)) { toVector.registerNewTypeId(Field.nullable("", MinorType.INT.getType())); toVector.registerNewTypeId(Field.nullable("", MinorType.FLOAT4.getType())); final TransferPair transferPair = sourceVector.makeTransferPair(toVector); final int[][] transferLengths = {{0, 2}, {2, 1}, {3, 2}, {5, 3}, {8, 2}}; for (final int[] transferLength : transferLengths) { final int start = transferLength[0]; final int length = transferLength[1]; transferPair.splitAndTransfer(start, length); /* check the toVector output after doing the splitAndTransfer */ for (int i = 0; i < length; i++) { assertEquals( sourceVector.getObject(start + i), toVector.getObject(i), "Different values at index: " + i); } } } } } @Test public void testSplitAndTransferDuvInStruct() { try (StructVector struct = StructVector.empty("struct", allocator)) { DenseUnionVector duv = struct.addOrGet( "duv", FieldType.notNullable(MinorType.DENSEUNION.getType()), DenseUnionVector.class); byte i32TypeId = duv.registerNewTypeId(Field.notNullable("i32", MinorType.INT.getType())); duv.addVector(i32TypeId, new IntVector("i32", allocator)); struct.setIndexDefined(0); duv.setTypeId(0, i32TypeId); duv.setSafe(0, newIntHolder(42)); struct.setNull(1); struct.setValueCount(2); try (StructVector dest = StructVector.empty("dest", allocator)) { TransferPair pair = struct.makeTransferPair(dest); pair.splitAndTransfer(0, 2); assertEquals(2, dest.getValueCount()); assertFalse(dest.isNull(0)); assertEquals(42, dest.getObject(0).get("duv")); assertTrue(dest.isNull(1)); } } } @Test public void testGetFieldTypeInfo() throws Exception { Map metadata = new HashMap<>(); metadata.put("key1", "value1"); int[] typeIds = new int[2]; typeIds[0] = 0; typeIds[1] = 1; List children = new ArrayList<>(); children.add(new Field("int", FieldType.nullable(MinorType.INT.getType()), null)); children.add(new Field("varchar", FieldType.nullable(MinorType.VARCHAR.getType()), null)); final FieldType fieldType = new FieldType( false, new ArrowType.Union(UnionMode.Dense, typeIds), /* dictionary= */ null, metadata); final Field field = new Field("union", fieldType, children); MinorType minorType = MinorType.DENSEUNION; DenseUnionVector vector = (DenseUnionVector) minorType.getNewVector(field, allocator, null); vector.initializeChildrenFromFields(children); assertEquals(vector.getField(), field); // Union has 2 child vectors assertEquals(2, vector.size()); // Check child field 0 VectorWithOrdinal intChild = vector.getChildVectorWithOrdinal("int"); assertEquals(0, intChild.ordinal); assertEquals(intChild.vector.getField(), children.get(0)); // Check child field 1 VectorWithOrdinal varcharChild = vector.getChildVectorWithOrdinal("varchar"); assertEquals(1, varcharChild.ordinal); assertEquals(varcharChild.vector.getField(), children.get(1)); } @Test public void testGetBufferAddress() throws Exception { try (DenseUnionVector vector = new DenseUnionVector(EMPTY_SCHEMA_PATH, allocator, null, null)) { boolean error = false; vector.allocateNew(); /* populate the UnionVector */ byte intTypeId = vector.registerNewTypeId(Field.nullable("", MinorType.INT.getType())); vector.setTypeId(0, intTypeId); vector.setSafe(0, newIntHolder(5)); byte float4TypeId = vector.registerNewTypeId(Field.nullable("", MinorType.INT.getType())); vector.setTypeId(1, float4TypeId); vector.setSafe(1, newFloat4Holder(5.5f)); vector.setTypeId(2, intTypeId); vector.setSafe(2, newIntHolder(10)); vector.setTypeId(3, float4TypeId); vector.setSafe(3, newFloat4Holder(10.5f)); vector.setValueCount(10); /* check the vector output */ assertEquals(10, vector.getValueCount()); assertEquals(false, vector.isNull(0)); assertEquals(5, vector.getObject(0)); assertEquals(false, vector.isNull(1)); assertEquals(5.5f, vector.getObject(1)); assertEquals(false, vector.isNull(2)); assertEquals(10, vector.getObject(2)); assertEquals(false, vector.isNull(3)); assertEquals(10.5f, vector.getObject(3)); List buffers = vector.getFieldBuffers(); long offsetAddress = vector.getOffsetBufferAddress(); try { vector.getDataBufferAddress(); } catch (UnsupportedOperationException ue) { error = true; } finally { assertTrue(error); } assertEquals(2, buffers.size()); assertEquals(offsetAddress, buffers.get(1).memoryAddress()); } } /** Test adding two struct vectors to the dense union vector. */ @Test public void testMultipleStructs() { FieldType type = new FieldType(true, ArrowType.Struct.INSTANCE, null, null); try (StructVector structVector1 = new StructVector("struct1", allocator, type, null); StructVector structVector2 = new StructVector("struct2", allocator, type, null); DenseUnionVector unionVector = DenseUnionVector.empty("union", allocator)) { // prepare sub vectors // first struct vector: (int, int) IntVector subVector11 = structVector1.addOrGet( "sub11", FieldType.nullable(MinorType.INT.getType()), IntVector.class); subVector11.allocateNew(); ValueVectorDataPopulator.setVector(subVector11, 0, 1); IntVector subVector12 = structVector1.addOrGet( "sub12", FieldType.nullable(MinorType.INT.getType()), IntVector.class); subVector12.allocateNew(); ValueVectorDataPopulator.setVector(subVector12, 0, 10); structVector1.setIndexDefined(0); structVector1.setIndexDefined(1); structVector1.setValueCount(2); // second struct vector: (string, string) VarCharVector subVector21 = structVector2.addOrGet( "sub21", FieldType.nullable(MinorType.VARCHAR.getType()), VarCharVector.class); subVector21.allocateNew(); ValueVectorDataPopulator.setVector(subVector21, "a0"); VarCharVector subVector22 = structVector2.addOrGet( "sub22", FieldType.nullable(MinorType.VARCHAR.getType()), VarCharVector.class); subVector22.allocateNew(); ValueVectorDataPopulator.setVector(subVector22, "b0"); structVector2.setIndexDefined(0); structVector2.setValueCount(1); // register relative types byte typeId1 = unionVector.registerNewTypeId(structVector1.getField()); byte typeId2 = unionVector.registerNewTypeId(structVector2.getField()); assertEquals(0, typeId1); assertEquals(1, typeId2); // add two struct vectors to union vector unionVector.addVector(typeId1, structVector1); unionVector.addVector(typeId2, structVector2); while (unionVector.getValueCapacity() < 3) { unionVector.reAlloc(); } ArrowBuf offsetBuf = unionVector.getOffsetBuffer(); unionVector.setTypeId(0, typeId1); offsetBuf.setInt(0, 0); unionVector.setTypeId(1, typeId2); offsetBuf.setInt(DenseUnionVector.OFFSET_WIDTH, 0); unionVector.setTypeId(2, typeId1); offsetBuf.setInt(DenseUnionVector.OFFSET_WIDTH * 2, 1); unionVector.setValueCount(3); Map value0 = new JsonStringHashMap<>(); value0.put("sub11", 0); value0.put("sub12", 0); assertEquals(value0, unionVector.getObject(0)); Map value1 = new JsonStringHashMap<>(); value1.put("sub21", new Text("a0")); value1.put("sub22", new Text("b0")); assertEquals(value1, unionVector.getObject(1)); Map value2 = new JsonStringHashMap<>(); value2.put("sub11", 1); value2.put("sub12", 10); assertEquals(value2, unionVector.getObject(2)); } } /** Test adding two varchar vectors to the dense union vector. */ @Test public void testMultipleVarChars() { try (VarCharVector childVector1 = new VarCharVector("child1", allocator); VarCharVector childVector2 = new VarCharVector("child2", allocator); DenseUnionVector unionVector = DenseUnionVector.empty("union", allocator)) { // prepare sub vectors ValueVectorDataPopulator.setVector(childVector1, "a0", "a4"); ValueVectorDataPopulator.setVector(childVector2, "b1", "b2"); // register relative types byte typeId1 = unionVector.registerNewTypeId(childVector1.getField()); byte typeId2 = unionVector.registerNewTypeId(childVector2.getField()); assertEquals(0, typeId1); assertEquals(1, typeId2); while (unionVector.getValueCapacity() < 5) { unionVector.reAlloc(); } // add two struct vectors to union vector unionVector.addVector(typeId1, childVector1); unionVector.addVector(typeId2, childVector2); ArrowBuf offsetBuf = unionVector.getOffsetBuffer(); // slot 0 points to child1 unionVector.setTypeId(0, typeId1); offsetBuf.setInt(0, 0); // slot 1 points to child2 unionVector.setTypeId(1, typeId2); offsetBuf.setInt(DenseUnionVector.OFFSET_WIDTH, 0); // slot 2 points to child2 unionVector.setTypeId(2, typeId2); offsetBuf.setInt(DenseUnionVector.OFFSET_WIDTH * 2, 1); // slot 4 points to child1 unionVector.setTypeId(4, typeId1); offsetBuf.setInt(DenseUnionVector.OFFSET_WIDTH * 4, 1); unionVector.setValueCount(5); assertEquals(new Text("a0"), unionVector.getObject(0)); assertEquals(new Text("b1"), unionVector.getObject(1)); assertEquals(new Text("b2"), unionVector.getObject(2)); assertNull(unionVector.getObject(3)); assertEquals(new Text("a4"), unionVector.getObject(4)); } } @Test public void testChildVectorValueCounts() { final NullableIntHolder intHolder = new NullableIntHolder(); intHolder.isSet = 1; final NullableBigIntHolder longHolder = new NullableBigIntHolder(); longHolder.isSet = 1; final NullableFloat4Holder floatHolder = new NullableFloat4Holder(); floatHolder.isSet = 1; try (DenseUnionVector vector = new DenseUnionVector("vector", allocator, null, null)) { vector.allocateNew(); // populate the delta vector with values {7, null, 8L, 9.0f, 10, 12L} while (vector.getValueCapacity() < 6) { vector.reAlloc(); } byte intTypeId = vector.registerNewTypeId(Field.nullable("", Types.MinorType.INT.getType())); vector.setTypeId(0, intTypeId); intHolder.value = 7; vector.setSafe(0, intHolder); byte longTypeId = vector.registerNewTypeId(Field.nullable("", Types.MinorType.BIGINT.getType())); vector.setTypeId(2, longTypeId); longHolder.value = 8L; vector.setSafe(2, longHolder); byte floatTypeId = vector.registerNewTypeId(Field.nullable("", Types.MinorType.FLOAT4.getType())); vector.setTypeId(3, floatTypeId); floatHolder.value = 9.0f; vector.setSafe(3, floatHolder); vector.setTypeId(4, intTypeId); intHolder.value = 10; vector.setSafe(4, intHolder); vector.setTypeId(5, longTypeId); longHolder.value = 12L; vector.setSafe(5, longHolder); vector.setValueCount(6); // verify results IntVector intVector = (IntVector) vector.getVectorByType(intTypeId); assertEquals(2, intVector.getValueCount()); assertEquals(7, intVector.get(0)); assertEquals(10, intVector.get(1)); BigIntVector longVector = (BigIntVector) vector.getVectorByType(longTypeId); assertEquals(2, longVector.getValueCount()); assertEquals(8L, longVector.get(0)); assertEquals(12L, longVector.get(1)); Float4Vector floatVector = (Float4Vector) vector.getVectorByType(floatTypeId); assertEquals(1, floatVector.getValueCount()); assertEquals(9.0f, floatVector.get(0), 0); } } private static NullableIntHolder newIntHolder(int value) { final NullableIntHolder holder = new NullableIntHolder(); holder.isSet = 1; holder.value = value; return holder; } private static NullableBitHolder newBitHolder(boolean value) { final NullableBitHolder holder = new NullableBitHolder(); holder.isSet = 1; holder.value = value ? 1 : 0; return holder; } private static NullableFloat4Holder newFloat4Holder(float value) { final NullableFloat4Holder holder = new NullableFloat4Holder(); holder.isSet = 1; holder.value = value; return holder; } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestDictionaryVector.java000066400000000000000000001405731515451615200325400ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.*; import static org.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector; 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.assertTrue; import static org.junit.jupiter.api.Assertions.fail; import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.HashMap; import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.function.ToIntBiFunction; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.complex.impl.NullableStructWriter; import org.apache.arrow.vector.complex.impl.UnionListWriter; 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.dictionary.ListSubfieldEncoder; import org.apache.arrow.vector.dictionary.StructSubfieldEncoder; import org.apache.arrow.vector.holders.NullableIntHolder; import org.apache.arrow.vector.holders.NullableUInt4Holder; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.DictionaryEncoding; import org.apache.arrow.vector.types.pojo.FieldType; 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; public class TestDictionaryVector { private BufferAllocator allocator; byte[] zero = "foo".getBytes(StandardCharsets.UTF_8); byte[] one = "bar".getBytes(StandardCharsets.UTF_8); byte[] two = "baz".getBytes(StandardCharsets.UTF_8); byte[][] data = new byte[][] {zero, one, two}; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testEncodeStrings() { // Create a new value vector try (final VarCharVector vector = newVarCharVector("foo", allocator); final VarCharVector dictionaryVector = newVarCharVector("dict", allocator); ) { setVector(vector, zero, one, one, two, zero); setVector(dictionaryVector, zero, one, two); Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (final ValueVector encoded = DictionaryEncoder.encode(vector, dictionary)) { // verify indices assertEquals(IntVector.class, encoded.getClass()); IntVector index = ((IntVector) encoded); assertEquals(5, index.getValueCount()); assertEquals(0, index.get(0)); assertEquals(1, index.get(1)); assertEquals(1, index.get(2)); assertEquals(2, index.get(3)); assertEquals(0, index.get(4)); // now run through the decoder and verify we get the original back try (ValueVector decoded = DictionaryEncoder.decode(encoded, dictionary)) { assertEquals(vector.getClass(), decoded.getClass()); assertEquals(vector.getValueCount(), ((VarCharVector) decoded).getValueCount()); for (int i = 0; i < 5; i++) { assertEquals(vector.getObject(i), ((VarCharVector) decoded).getObject(i)); } } } } } @Test public void testEncodeLargeVector() { // Create a new value vector try (final VarCharVector vector = newVarCharVector("foo", allocator); final VarCharVector dictionaryVector = newVarCharVector("dict", allocator); ) { vector.allocateNew(); int count = 10000; for (int i = 0; i < 10000; ++i) { vector.setSafe(i, data[i % 3], 0, data[i % 3].length); } vector.setValueCount(count); setVector(dictionaryVector, zero, one, two); Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (final ValueVector encoded = DictionaryEncoder.encode(vector, dictionary)) { // verify indices assertEquals(IntVector.class, encoded.getClass()); IntVector index = ((IntVector) encoded); assertEquals(count, index.getValueCount()); for (int i = 0; i < count; ++i) { assertEquals(i % 3, index.get(i)); } // now run through the decoder and verify we get the original back try (ValueVector decoded = DictionaryEncoder.decode(encoded, dictionary)) { assertEquals(vector.getClass(), decoded.getClass()); assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < count; ++i) { assertEquals(vector.getObject(i), decoded.getObject(i)); } } } } } @Test public void testEncodeList() { // Create a new value vector try (final ListVector vector = ListVector.empty("vector", allocator); final ListVector dictionaryVector = ListVector.empty("dict", allocator); ) { UnionListWriter writer = vector.getWriter(); writer.allocate(); // set some values writeListVector(writer, new int[] {10, 20}); writeListVector(writer, new int[] {10, 20}); writeListVector(writer, new int[] {10, 20}); writeListVector(writer, new int[] {30, 40, 50}); writeListVector(writer, new int[] {30, 40, 50}); writeListVector(writer, new int[] {10, 20}); writer.setValueCount(6); UnionListWriter dictWriter = dictionaryVector.getWriter(); dictWriter.allocate(); writeListVector(dictWriter, new int[] {10, 20}); writeListVector(dictWriter, new int[] {30, 40, 50}); dictWriter.setValueCount(2); Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (final ValueVector encoded = DictionaryEncoder.encode(vector, dictionary)) { // verify indices assertEquals(IntVector.class, encoded.getClass()); IntVector index = ((IntVector) encoded); assertEquals(6, index.getValueCount()); assertEquals(0, index.get(0)); assertEquals(0, index.get(1)); assertEquals(0, index.get(2)); assertEquals(1, index.get(3)); assertEquals(1, index.get(4)); assertEquals(0, index.get(5)); // now run through the decoder and verify we get the original back try (ValueVector decoded = DictionaryEncoder.decode(encoded, dictionary)) { assertEquals(vector.getClass(), decoded.getClass()); assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < 5; i++) { assertEquals(vector.getObject(i), decoded.getObject(i)); } } } } } @Test public void testEncodeStruct() { // Create a new value vector try (final StructVector vector = StructVector.empty("vector", allocator); final StructVector dictionaryVector = StructVector.empty("dict", allocator); ) { vector.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); vector.addOrGet("f1", FieldType.nullable(new ArrowType.Int(64, true)), BigIntVector.class); dictionaryVector.addOrGet( "f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); dictionaryVector.addOrGet( "f1", FieldType.nullable(new ArrowType.Int(64, true)), BigIntVector.class); NullableStructWriter writer = vector.getWriter(); writer.allocate(); writeStructVector(writer, 1, 10L); writeStructVector(writer, 1, 10L); writeStructVector(writer, 1, 10L); writeStructVector(writer, 2, 20L); writeStructVector(writer, 2, 20L); writeStructVector(writer, 2, 20L); writeStructVector(writer, 1, 10L); writer.setValueCount(7); NullableStructWriter dictWriter = dictionaryVector.getWriter(); dictWriter.allocate(); writeStructVector(dictWriter, 1, 10L); writeStructVector(dictWriter, 2, 20L); dictionaryVector.setValueCount(2); Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (final ValueVector encoded = DictionaryEncoder.encode(vector, dictionary)) { // verify indices assertEquals(IntVector.class, encoded.getClass()); IntVector index = ((IntVector) encoded); assertEquals(7, index.getValueCount()); assertEquals(0, index.get(0)); assertEquals(0, index.get(1)); assertEquals(0, index.get(2)); assertEquals(1, index.get(3)); assertEquals(1, index.get(4)); assertEquals(1, index.get(5)); assertEquals(0, index.get(6)); // now run through the decoder and verify we get the original back try (ValueVector decoded = DictionaryEncoder.decode(encoded, dictionary)) { assertEquals(vector.getClass(), decoded.getClass()); assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < 5; i++) { assertEquals(vector.getObject(i), decoded.getObject(i)); } } } } } @Test public void testEncodeBinaryVector() { // Create a new value vector try (final VarBinaryVector vector = newVarBinaryVector("foo", allocator); final VarBinaryVector dictionaryVector = newVarBinaryVector("dict", allocator)) { setVector(vector, zero, one, one, two, zero); setVector(dictionaryVector, zero, one, two); Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (final ValueVector encoded = DictionaryEncoder.encode(vector, dictionary)) { // verify indices assertEquals(IntVector.class, encoded.getClass()); IntVector index = ((IntVector) encoded); assertEquals(5, index.getValueCount()); assertEquals(0, index.get(0)); assertEquals(1, index.get(1)); assertEquals(1, index.get(2)); assertEquals(2, index.get(3)); assertEquals(0, index.get(4)); // now run through the decoder and verify we get the original back try (VarBinaryVector decoded = (VarBinaryVector) DictionaryEncoder.decode(encoded, dictionary)) { assertEquals(vector.getClass(), decoded.getClass()); assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < 5; i++) { assertTrue(Arrays.equals(vector.getObject(i), decoded.getObject(i))); } } } } } @Test public void testEncodeUnion() { // Create a new value vector try (final UnionVector vector = new UnionVector("vector", allocator, /* field type */ null, /* call-back */ null); final UnionVector dictionaryVector = new UnionVector("dict", allocator, /* field type */ null, /* call-back */ null); ) { final NullableUInt4Holder uintHolder1 = new NullableUInt4Holder(); uintHolder1.value = 10; uintHolder1.isSet = 1; final NullableIntHolder intHolder1 = new NullableIntHolder(); intHolder1.value = 10; intHolder1.isSet = 1; final NullableIntHolder intHolder2 = new NullableIntHolder(); intHolder2.value = 20; intHolder2.isSet = 1; // write data vector.setType(0, Types.MinorType.UINT4); vector.setSafe(0, uintHolder1); vector.setType(1, Types.MinorType.INT); vector.setSafe(1, intHolder1); vector.setType(2, Types.MinorType.INT); vector.setSafe(2, intHolder1); vector.setType(3, Types.MinorType.INT); vector.setSafe(3, intHolder2); vector.setType(4, Types.MinorType.INT); vector.setSafe(4, intHolder2); vector.setValueCount(5); // write dictionary dictionaryVector.setType(0, Types.MinorType.UINT4); dictionaryVector.setSafe(0, uintHolder1); dictionaryVector.setType(1, Types.MinorType.INT); dictionaryVector.setSafe(1, intHolder1); dictionaryVector.setType(2, Types.MinorType.INT); dictionaryVector.setSafe(2, intHolder2); dictionaryVector.setValueCount(3); Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (final ValueVector encoded = DictionaryEncoder.encode(vector, dictionary)) { // verify indices assertEquals(IntVector.class, encoded.getClass()); IntVector index = ((IntVector) encoded); assertEquals(5, index.getValueCount()); assertEquals(0, index.get(0)); assertEquals(1, index.get(1)); assertEquals(1, index.get(2)); assertEquals(2, index.get(3)); assertEquals(2, index.get(4)); // now run through the decoder and verify we get the original back try (ValueVector decoded = DictionaryEncoder.decode(encoded, dictionary)) { assertEquals(vector.getClass(), decoded.getClass()); assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < 5; i++) { assertEquals(vector.getObject(i), decoded.getObject(i)); } } } } } @Test public void testIntEquals() { // test Int try (final IntVector vector1 = new IntVector("int", allocator); final IntVector vector2 = new IntVector("int", allocator)) { Dictionary dict1 = new Dictionary(vector1, new DictionaryEncoding(1L, false, null)); Dictionary dict2 = new Dictionary(vector2, new DictionaryEncoding(1L, false, null)); setVector(vector1, 1, 2, 3); setVector(vector2, 1, 2, 0); assertFalse(dict1.equals(dict2)); vector2.setSafe(2, 3); assertTrue(dict1.equals(dict2)); } } @Test public void testVarcharEquals() { try (final VarCharVector vector1 = new VarCharVector("varchar", allocator); final VarCharVector vector2 = new VarCharVector("varchar", allocator)) { Dictionary dict1 = new Dictionary(vector1, new DictionaryEncoding(1L, false, null)); Dictionary dict2 = new Dictionary(vector2, new DictionaryEncoding(1L, false, null)); setVector(vector1, zero, one, two); setVector(vector2, zero, one, one); assertFalse(dict1.equals(dict2)); vector2.setSafe(2, two, 0, two.length); assertTrue(dict1.equals(dict2)); } } @Test public void testVarBinaryEquals() { try (final VarBinaryVector vector1 = new VarBinaryVector("binary", allocator); final VarBinaryVector vector2 = new VarBinaryVector("binary", allocator)) { Dictionary dict1 = new Dictionary(vector1, new DictionaryEncoding(1L, false, null)); Dictionary dict2 = new Dictionary(vector2, new DictionaryEncoding(1L, false, null)); setVector(vector1, zero, one, two); setVector(vector2, zero, one, one); assertFalse(dict1.equals(dict2)); vector2.setSafe(2, two, 0, two.length); assertTrue(dict1.equals(dict2)); } } @Test public void testListEquals() { try (final ListVector vector1 = ListVector.empty("list", allocator); final ListVector vector2 = ListVector.empty("list", allocator); ) { Dictionary dict1 = new Dictionary(vector1, new DictionaryEncoding(1L, false, null)); Dictionary dict2 = new Dictionary(vector2, new DictionaryEncoding(1L, false, null)); UnionListWriter writer1 = vector1.getWriter(); writer1.allocate(); // set some values writeListVector(writer1, new int[] {1, 2}); writeListVector(writer1, new int[] {3, 4}); writeListVector(writer1, new int[] {5, 6}); writer1.setValueCount(3); UnionListWriter writer2 = vector2.getWriter(); writer2.allocate(); // set some values writeListVector(writer2, new int[] {1, 2}); writeListVector(writer2, new int[] {3, 4}); writeListVector(writer2, new int[] {5, 6}); writer2.setValueCount(3); assertTrue(dict1.equals(dict2)); } } @Test public void testStructEquals() { try (final StructVector vector1 = StructVector.empty("struct", allocator); final StructVector vector2 = StructVector.empty("struct", allocator); ) { vector1.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); vector1.addOrGet("f1", FieldType.nullable(new ArrowType.Int(64, true)), BigIntVector.class); vector2.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); vector2.addOrGet("f1", FieldType.nullable(new ArrowType.Int(64, true)), BigIntVector.class); Dictionary dict1 = new Dictionary(vector1, new DictionaryEncoding(1L, false, null)); Dictionary dict2 = new Dictionary(vector2, new DictionaryEncoding(1L, false, null)); NullableStructWriter writer1 = vector1.getWriter(); writer1.allocate(); writeStructVector(writer1, 1, 10L); writeStructVector(writer1, 2, 20L); writer1.setValueCount(2); NullableStructWriter writer2 = vector2.getWriter(); writer2.allocate(); writeStructVector(writer2, 1, 10L); writeStructVector(writer2, 2, 20L); writer2.setValueCount(2); assertTrue(dict1.equals(dict2)); } } @Test public void testUnionEquals() { try (final UnionVector vector1 = new UnionVector("union", allocator, /* field type */ null, /* call-back */ null); final UnionVector vector2 = new UnionVector("union", allocator, /* field type */ null, /* call-back */ null); ) { final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); uInt4Holder.value = 10; uInt4Holder.isSet = 1; final NullableIntHolder intHolder = new NullableIntHolder(); uInt4Holder.value = 20; uInt4Holder.isSet = 1; vector1.setType(0, Types.MinorType.UINT4); vector1.setSafe(0, uInt4Holder); vector1.setType(2, Types.MinorType.INT); vector1.setSafe(2, intHolder); vector1.setValueCount(3); vector2.setType(0, Types.MinorType.UINT4); vector2.setSafe(0, uInt4Holder); vector2.setType(2, Types.MinorType.INT); vector2.setSafe(2, intHolder); vector2.setValueCount(3); Dictionary dict1 = new Dictionary(vector1, new DictionaryEncoding(1L, false, null)); Dictionary dict2 = new Dictionary(vector2, new DictionaryEncoding(1L, false, null)); assertTrue(dict1.equals(dict2)); } } @Test public void testEncodeWithEncoderInstance() { // Create a new value vector try (final VarCharVector vector = newVarCharVector("vector", allocator); final VarCharVector dictionaryVector = newVarCharVector("dict", allocator); ) { setVector(vector, zero, one, one, two, zero); setVector(dictionaryVector, zero, one, two); Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); DictionaryEncoder encoder = new DictionaryEncoder(dictionary, allocator); try (final ValueVector encoded = encoder.encode(vector)) { // verify indices assertEquals(IntVector.class, encoded.getClass()); IntVector index = ((IntVector) encoded); assertEquals(5, index.getValueCount()); assertEquals(0, index.get(0)); assertEquals(1, index.get(1)); assertEquals(1, index.get(2)); assertEquals(2, index.get(3)); assertEquals(0, index.get(4)); // now run through the decoder and verify we get the original back try (ValueVector decoded = encoder.decode(encoded)) { assertEquals(vector.getClass(), decoded.getClass()); assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < 5; i++) { assertEquals(vector.getObject(i), ((VarCharVector) decoded).getObject(i)); } } } } } @Test public void testEncodeMultiVectors() { // Create a new value vector try (final VarCharVector vector1 = newVarCharVector("vector1", allocator); final VarCharVector vector2 = newVarCharVector("vector2", allocator); final VarCharVector dictionaryVector = newVarCharVector("dict", allocator); ) { setVector(vector1, zero, one, one, two, zero); setVector(vector2, zero, one, one); setVector(dictionaryVector, zero, one, two); Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); DictionaryEncoder encoder = new DictionaryEncoder(dictionary, allocator); try (final ValueVector encoded = encoder.encode(vector1)) { // verify indices assertEquals(IntVector.class, encoded.getClass()); IntVector index = ((IntVector) encoded); assertEquals(5, index.getValueCount()); assertEquals(0, index.get(0)); assertEquals(1, index.get(1)); assertEquals(1, index.get(2)); assertEquals(2, index.get(3)); assertEquals(0, index.get(4)); // now run through the decoder and verify we get the original back try (ValueVector decoded = encoder.decode(encoded)) { assertEquals(vector1.getClass(), decoded.getClass()); assertEquals(vector1.getValueCount(), decoded.getValueCount()); for (int i = 0; i < 5; i++) { assertEquals(vector1.getObject(i), ((VarCharVector) decoded).getObject(i)); } } } try (final ValueVector encoded = encoder.encode(vector2)) { // verify indices assertEquals(IntVector.class, encoded.getClass()); IntVector index = ((IntVector) encoded); assertEquals(3, index.getValueCount()); assertEquals(0, index.get(0)); assertEquals(1, index.get(1)); assertEquals(1, index.get(2)); // now run through the decoder and verify we get the original back try (ValueVector decoded = encoder.decode(encoded)) { assertEquals(vector2.getClass(), decoded.getClass()); assertEquals(vector2.getValueCount(), decoded.getValueCount()); for (int i = 0; i < 3; i++) { assertEquals(vector2.getObject(i), ((VarCharVector) decoded).getObject(i)); } } } } } @Test public void testEncodeListSubField() { // Create a new value vector try (final ListVector vector = ListVector.empty("vector", allocator); final ListVector dictionaryVector = ListVector.empty("dict", allocator); ) { UnionListWriter writer = vector.getWriter(); writer.allocate(); // set some values writeListVector(writer, new int[] {10, 20}); writeListVector(writer, new int[] {10, 20}); writeListVector(writer, new int[] {10, 20}); writeListVector(writer, new int[] {30, 40, 50}); writeListVector(writer, new int[] {30, 40, 50}); writeListVector(writer, new int[] {10, 20}); writer.setValueCount(6); UnionListWriter dictWriter = dictionaryVector.getWriter(); dictWriter.allocate(); writeListVector(dictWriter, new int[] {10, 20, 30, 40, 50}); dictionaryVector.setValueCount(1); Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); ListSubfieldEncoder encoder = new ListSubfieldEncoder(dictionary, allocator); try (final ListVector encoded = (ListVector) encoder.encodeListSubField(vector)) { // verify indices assertEquals(ListVector.class, encoded.getClass()); assertEquals(6, encoded.getValueCount()); int[] realValue1 = convertListToIntArray(encoded.getObject(0)); assertTrue(Arrays.equals(new int[] {0, 1}, realValue1)); int[] realValue2 = convertListToIntArray(encoded.getObject(1)); assertTrue(Arrays.equals(new int[] {0, 1}, realValue2)); int[] realValue3 = convertListToIntArray(encoded.getObject(2)); assertTrue(Arrays.equals(new int[] {0, 1}, realValue3)); int[] realValue4 = convertListToIntArray(encoded.getObject(3)); assertTrue(Arrays.equals(new int[] {2, 3, 4}, realValue4)); int[] realValue5 = convertListToIntArray(encoded.getObject(4)); assertTrue(Arrays.equals(new int[] {2, 3, 4}, realValue5)); int[] realValue6 = convertListToIntArray(encoded.getObject(5)); assertTrue(Arrays.equals(new int[] {0, 1}, realValue6)); // now run through the decoder and verify we get the original back try (ValueVector decoded = encoder.decodeListSubField(encoded)) { assertEquals(vector.getClass(), decoded.getClass()); assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < 5; i++) { assertEquals(vector.getObject(i), decoded.getObject(i)); } } } } } @Test public void testEncodeFixedSizeListSubField() { // Create a new value vector try (final FixedSizeListVector vector = FixedSizeListVector.empty("vector", 2, allocator); final FixedSizeListVector dictionaryVector = FixedSizeListVector.empty("dict", 2, allocator)) { vector.allocateNew(); vector.setValueCount(4); IntVector dataVector = (IntVector) vector.addOrGetVector(FieldType.nullable(Types.MinorType.INT.getType())).getVector(); dataVector.allocateNew(8); dataVector.setValueCount(8); // set value at index 0 vector.setNotNull(0); dataVector.set(0, 10); dataVector.set(1, 20); // set value at index 1 vector.setNotNull(1); dataVector.set(2, 10); dataVector.set(3, 20); // set value at index 2 vector.setNotNull(2); dataVector.set(4, 30); dataVector.set(5, 40); // set value at index 3 vector.setNotNull(3); dataVector.set(6, 10); dataVector.set(7, 20); dictionaryVector.allocateNew(); dictionaryVector.setValueCount(2); IntVector dictDataVector = (IntVector) dictionaryVector .addOrGetVector(FieldType.nullable(Types.MinorType.INT.getType())) .getVector(); dictDataVector.allocateNew(4); dictDataVector.setValueCount(4); dictionaryVector.setNotNull(0); dictDataVector.set(0, 10); dictDataVector.set(1, 20); dictionaryVector.setNotNull(1); dictDataVector.set(2, 30); dictDataVector.set(3, 40); Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); ListSubfieldEncoder encoder = new ListSubfieldEncoder(dictionary, allocator); try (final FixedSizeListVector encoded = (FixedSizeListVector) encoder.encodeListSubField(vector)) { // verify indices assertEquals(FixedSizeListVector.class, encoded.getClass()); assertEquals(4, encoded.getValueCount()); int[] realValue1 = convertListToIntArray(encoded.getObject(0)); assertTrue(Arrays.equals(new int[] {0, 1}, realValue1)); int[] realValue2 = convertListToIntArray(encoded.getObject(1)); assertTrue(Arrays.equals(new int[] {0, 1}, realValue2)); int[] realValue3 = convertListToIntArray(encoded.getObject(2)); assertTrue(Arrays.equals(new int[] {2, 3}, realValue3)); int[] realValue4 = convertListToIntArray(encoded.getObject(3)); assertTrue(Arrays.equals(new int[] {0, 1}, realValue4)); // now run through the decoder and verify we get the original back try (ValueVector decoded = encoder.decodeListSubField(encoded)) { assertEquals(vector.getClass(), decoded.getClass()); assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < 5; i++) { assertEquals(vector.getObject(i), decoded.getObject(i)); } } } } } @Test public void testEncodeStructSubField() { try (final StructVector vector = StructVector.empty("vector", allocator); final VarCharVector dictVector1 = new VarCharVector("f0", allocator); final VarCharVector dictVector2 = new VarCharVector("f1", allocator)) { vector.addOrGet("f0", FieldType.nullable(ArrowType.Utf8.INSTANCE), VarCharVector.class); vector.addOrGet("f1", FieldType.nullable(ArrowType.Utf8.INSTANCE), VarCharVector.class); NullableStructWriter writer = vector.getWriter(); writer.allocate(); // set some values writeStructVector(writer, "aa", "baz"); writeStructVector(writer, "bb", "bar"); writeStructVector(writer, "cc", "foo"); writeStructVector(writer, "aa", "foo"); writeStructVector(writer, "dd", "foo"); writer.setValueCount(5); // initialize dictionaries DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); setVector( dictVector1, "aa".getBytes(StandardCharsets.UTF_8), "bb".getBytes(StandardCharsets.UTF_8), "cc".getBytes(StandardCharsets.UTF_8), "dd".getBytes(StandardCharsets.UTF_8)); setVector( dictVector2, "foo".getBytes(StandardCharsets.UTF_8), "baz".getBytes(StandardCharsets.UTF_8), "bar".getBytes(StandardCharsets.UTF_8)); provider.put(new Dictionary(dictVector1, new DictionaryEncoding(1L, false, null))); provider.put(new Dictionary(dictVector2, new DictionaryEncoding(2L, false, null))); StructSubfieldEncoder encoder = new StructSubfieldEncoder(allocator, provider); Map columnToDictionaryId = new HashMap<>(); columnToDictionaryId.put(0, 1L); columnToDictionaryId.put(1, 2L); try (final StructVector encoded = (StructVector) encoder.encode(vector, columnToDictionaryId)) { // verify indices assertEquals(StructVector.class, encoded.getClass()); assertEquals(5, encoded.getValueCount()); Object[] realValue1 = convertMapValuesToArray(encoded.getObject(0)); assertTrue(Arrays.equals(new Object[] {0, 1}, realValue1)); Object[] realValue2 = convertMapValuesToArray(encoded.getObject(1)); assertTrue(Arrays.equals(new Object[] {1, 2}, realValue2)); Object[] realValue3 = convertMapValuesToArray(encoded.getObject(2)); assertTrue(Arrays.equals(new Object[] {2, 0}, realValue3)); Object[] realValue4 = convertMapValuesToArray(encoded.getObject(3)); assertTrue(Arrays.equals(new Object[] {0, 0}, realValue4)); Object[] realValue5 = convertMapValuesToArray(encoded.getObject(4)); assertTrue(Arrays.equals(new Object[] {3, 0}, realValue5)); // now run through the decoder and verify we get the original back try (ValueVector decoded = encoder.decode(encoded)) { assertEquals(vector.getClass(), decoded.getClass()); assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < 5; i++) { assertEquals(vector.getObject(i), decoded.getObject(i)); } } } } } @Test public void testEncodeStructSubFieldWithCertainColumns() { // in this case, some child vector is encoded and others are not try (final StructVector vector = StructVector.empty("vector", allocator); final VarCharVector dictVector1 = new VarCharVector("f0", allocator)) { vector.addOrGet("f0", FieldType.nullable(ArrowType.Utf8.INSTANCE), VarCharVector.class); vector.addOrGet("f1", FieldType.nullable(ArrowType.Utf8.INSTANCE), VarCharVector.class); NullableStructWriter writer = vector.getWriter(); writer.allocate(); // set some values writeStructVector(writer, "aa", "baz"); writeStructVector(writer, "bb", "bar"); writeStructVector(writer, "cc", "foo"); writeStructVector(writer, "aa", "foo"); writeStructVector(writer, "dd", "foo"); writer.setValueCount(5); // initialize dictionaries DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); setVector( dictVector1, "aa".getBytes(StandardCharsets.UTF_8), "bb".getBytes(StandardCharsets.UTF_8), "cc".getBytes(StandardCharsets.UTF_8), "dd".getBytes(StandardCharsets.UTF_8)); provider.put(new Dictionary(dictVector1, new DictionaryEncoding(1L, false, null))); StructSubfieldEncoder encoder = new StructSubfieldEncoder(allocator, provider); Map columnToDictionaryId = new HashMap<>(); columnToDictionaryId.put(0, 1L); try (final StructVector encoded = (StructVector) encoder.encode(vector, columnToDictionaryId)) { // verify indices assertEquals(StructVector.class, encoded.getClass()); assertEquals(5, encoded.getValueCount()); Object[] realValue1 = convertMapValuesToArray(encoded.getObject(0)); assertTrue(Arrays.equals(new Object[] {0, new Text("baz")}, realValue1)); Object[] realValue2 = convertMapValuesToArray(encoded.getObject(1)); assertTrue(Arrays.equals(new Object[] {1, new Text("bar")}, realValue2)); Object[] realValue3 = convertMapValuesToArray(encoded.getObject(2)); assertTrue(Arrays.equals(new Object[] {2, new Text("foo")}, realValue3)); Object[] realValue4 = convertMapValuesToArray(encoded.getObject(3)); assertTrue(Arrays.equals(new Object[] {0, new Text("foo")}, realValue4)); Object[] realValue5 = convertMapValuesToArray(encoded.getObject(4)); assertTrue(Arrays.equals(new Object[] {3, new Text("foo")}, realValue5)); // now run through the decoder and verify we get the original back try (ValueVector decoded = encoder.decode(encoded)) { assertEquals(vector.getClass(), decoded.getClass()); assertEquals(vector.getValueCount(), decoded.getValueCount()); for (int i = 0; i < 5; i++) { assertEquals(vector.getObject(i), decoded.getObject(i)); } } } } } @Test public void testNoMemoryLeak() { // test no memory leak when encode try (final VarCharVector vector = newVarCharVector("foo", allocator); final VarCharVector dictionaryVector = newVarCharVector("dict", allocator)) { setVector(vector, zero, one, two); setVector(dictionaryVector, zero, one); Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (final ValueVector encoded = DictionaryEncoder.encode(vector, dictionary)) { fail("There should be an exception when encoding"); } catch (Exception e) { assertEquals("Dictionary encoding not defined for value:" + new Text(two), e.getMessage()); } } assertEquals(0, allocator.getAllocatedMemory(), "encode memory leak"); // test no memory leak when decode try (final IntVector indices = newVector(IntVector.class, "", Types.MinorType.INT, allocator); final VarCharVector dictionaryVector = newVarCharVector("dict", allocator)) { setVector(indices, 3); setVector(dictionaryVector, zero, one); Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (final ValueVector decoded = DictionaryEncoder.decode(indices, dictionary, allocator)) { fail("There should be an exception when decoding"); } catch (Exception e) { assertEquals("Provided dictionary does not contain value for index 3", e.getMessage()); } } assertEquals(0, allocator.getAllocatedMemory(), "decode memory leak"); } @Test public void testListNoMemoryLeak() { // Create a new value vector try (final ListVector vector = ListVector.empty("vector", allocator); final ListVector dictionaryVector = ListVector.empty("dict", allocator)) { UnionListWriter writer = vector.getWriter(); writer.allocate(); writeListVector(writer, new int[] {10, 20}); writer.setValueCount(1); UnionListWriter dictWriter = dictionaryVector.getWriter(); dictWriter.allocate(); writeListVector(dictWriter, new int[] {10}); dictionaryVector.setValueCount(1); Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); ListSubfieldEncoder encoder = new ListSubfieldEncoder(dictionary, allocator); try (final ListVector encoded = (ListVector) encoder.encodeListSubField(vector)) { fail("There should be an exception when encoding"); } catch (Exception e) { assertEquals("Dictionary encoding not defined for value:20", e.getMessage()); } } assertEquals(0, allocator.getAllocatedMemory(), "list encode memory leak"); try (final ListVector indices = ListVector.empty("indices", allocator); final ListVector dictionaryVector = ListVector.empty("dict", allocator)) { UnionListWriter writer = indices.getWriter(); writer.allocate(); writeListVector(writer, new int[] {3}); writer.setValueCount(1); UnionListWriter dictWriter = dictionaryVector.getWriter(); dictWriter.allocate(); writeListVector(dictWriter, new int[] {10, 20}); dictionaryVector.setValueCount(1); Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); try (final ValueVector decoded = ListSubfieldEncoder.decodeListSubField(indices, dictionary, allocator)) { fail("There should be an exception when decoding"); } catch (Exception e) { assertEquals("Provided dictionary does not contain value for index 3", e.getMessage()); } } assertEquals(0, allocator.getAllocatedMemory(), "list decode memory leak"); } @Test public void testStructNoMemoryLeak() { try (final StructVector vector = StructVector.empty("vector", allocator); final VarCharVector dictVector1 = new VarCharVector("f0", allocator); final VarCharVector dictVector2 = new VarCharVector("f1", allocator)) { vector.addOrGet("f0", FieldType.nullable(ArrowType.Utf8.INSTANCE), VarCharVector.class); vector.addOrGet("f1", FieldType.nullable(ArrowType.Utf8.INSTANCE), VarCharVector.class); NullableStructWriter writer = vector.getWriter(); writer.allocate(); writeStructVector(writer, "aa", "baz"); writer.setValueCount(1); DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); setVector(dictVector1, "aa".getBytes(StandardCharsets.UTF_8)); setVector(dictVector2, "foo".getBytes(StandardCharsets.UTF_8)); provider.put(new Dictionary(dictVector1, new DictionaryEncoding(1L, false, null))); provider.put(new Dictionary(dictVector2, new DictionaryEncoding(2L, false, null))); StructSubfieldEncoder encoder = new StructSubfieldEncoder(allocator, provider); Map columnToDictionaryId = new HashMap<>(); columnToDictionaryId.put(0, 1L); columnToDictionaryId.put(1, 2L); try (final StructVector encoded = (StructVector) encoder.encode(vector, columnToDictionaryId)) { fail("There should be an exception when encoding"); } catch (Exception e) { assertEquals("Dictionary encoding not defined for value:baz", e.getMessage()); } } assertEquals(0, allocator.getAllocatedMemory(), "struct encode memory leak"); try (final StructVector indices = StructVector.empty("indices", allocator); final VarCharVector dictVector1 = new VarCharVector("f0", allocator); final VarCharVector dictVector2 = new VarCharVector("f1", allocator)) { DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); setVector(dictVector1, "aa".getBytes(StandardCharsets.UTF_8)); setVector(dictVector2, "foo".getBytes(StandardCharsets.UTF_8)); provider.put(new Dictionary(dictVector1, new DictionaryEncoding(1L, false, null))); provider.put(new Dictionary(dictVector2, new DictionaryEncoding(2L, false, null))); ArrowType int32 = new ArrowType.Int(32, true); indices.addOrGet( "f0", new FieldType(true, int32, provider.lookup(1L).getEncoding()), IntVector.class); indices.addOrGet( "f1", new FieldType(true, int32, provider.lookup(2L).getEncoding()), IntVector.class); NullableStructWriter writer = indices.getWriter(); writer.allocate(); writer.start(); writer.integer("f0").writeInt(1); writer.integer("f1").writeInt(3); writer.end(); writer.setValueCount(1); try (final StructVector decode = StructSubfieldEncoder.decode(indices, provider, allocator)) { fail("There should be an exception when decoding"); } catch (Exception e) { assertEquals("Provided dictionary does not contain value for index 3", e.getMessage()); } } assertEquals(0, allocator.getAllocatedMemory(), "struct decode memory leak"); } private void testDictionary( Dictionary dictionary, ToIntBiFunction valGetter) { try (VarCharVector vector = new VarCharVector("vector", allocator)) { setVector(vector, "1", "3", "5", "7", "9"); try (ValueVector encodedVector = DictionaryEncoder.encode(vector, dictionary)) { // verify encoded result assertEquals(vector.getValueCount(), encodedVector.getValueCount()); assertEquals(1, valGetter.applyAsInt(encodedVector, 0)); assertEquals(3, valGetter.applyAsInt(encodedVector, 1)); assertEquals(5, valGetter.applyAsInt(encodedVector, 2)); assertEquals(7, valGetter.applyAsInt(encodedVector, 3)); assertEquals(9, valGetter.applyAsInt(encodedVector, 4)); try (ValueVector decodedVector = DictionaryEncoder.decode(encodedVector, dictionary)) { assertTrue(decodedVector instanceof VarCharVector); assertEquals(vector.getValueCount(), decodedVector.getValueCount()); assertArrayEquals( "1".getBytes(StandardCharsets.UTF_8), ((VarCharVector) decodedVector).get(0)); assertArrayEquals( "3".getBytes(StandardCharsets.UTF_8), ((VarCharVector) decodedVector).get(1)); assertArrayEquals( "5".getBytes(StandardCharsets.UTF_8), ((VarCharVector) decodedVector).get(2)); assertArrayEquals( "7".getBytes(StandardCharsets.UTF_8), ((VarCharVector) decodedVector).get(3)); assertArrayEquals( "9".getBytes(StandardCharsets.UTF_8), ((VarCharVector) decodedVector).get(4)); } } } } @Test public void testDictionaryUInt1() { try (VarCharVector dictionaryVector = new VarCharVector("dict vector", allocator)) { setVector(dictionaryVector, "0", "1", "2", "3", "4", "5", "6", "7", "8", "9"); Dictionary dictionary1 = new Dictionary( dictionaryVector, new DictionaryEncoding( /* id= */ 10L, /* ordered= */ false, /* indexType= */ new ArrowType.Int(/*bitWidth*/ 8, /*isSigned*/ false))); testDictionary(dictionary1, (vector, index) -> ((UInt1Vector) vector).get(index)); } } @Test public void testDictionaryUInt2() { try (VarCharVector dictionaryVector = new VarCharVector("dict vector", allocator)) { setVector(dictionaryVector, "0", "1", "2", "3", "4", "5", "6", "7", "8", "9"); Dictionary dictionary2 = new Dictionary( dictionaryVector, new DictionaryEncoding( /* id= */ 20L, /* ordered= */ false, /* indexType= */ new ArrowType.Int(/* bitWidth= */ 16, /*isSigned*/ false))); testDictionary(dictionary2, (vector, index) -> ((UInt2Vector) vector).get(index)); } } @Test public void testDictionaryUInt4() { try (VarCharVector dictionaryVector = new VarCharVector("dict vector", allocator)) { setVector(dictionaryVector, "0", "1", "2", "3", "4", "5", "6", "7", "8", "9"); Dictionary dictionary4 = new Dictionary( dictionaryVector, new DictionaryEncoding( /* id= */ 30L, /* ordered= */ false, /* indexType= */ new ArrowType.Int(/* bitWidth= */ 32, /*isSigned*/ false))); testDictionary(dictionary4, (vector, index) -> ((UInt4Vector) vector).get(index)); } } @Test public void testDictionaryUInt8() { try (VarCharVector dictionaryVector = new VarCharVector("dict vector", allocator)) { setVector(dictionaryVector, "0", "1", "2", "3", "4", "5", "6", "7", "8", "9"); Dictionary dictionary8 = new Dictionary( dictionaryVector, new DictionaryEncoding( /* id= */ 40L, /* ordered= */ false, /* indexType= */ new ArrowType.Int(/* bitWidth= */ 64, /*isSigned*/ false))); testDictionary(dictionary8, (vector, index) -> (int) ((UInt8Vector) vector).get(index)); } } @Test public void testDictionaryUIntOverflow() { // the size is within the range of UInt1, but outside the range of TinyInt. final int vecLength = 256; try (VarCharVector dictionaryVector = new VarCharVector("dict vector", allocator)) { dictionaryVector.allocateNew(vecLength * 3, vecLength); for (int i = 0; i < vecLength; i++) { dictionaryVector.set(i, String.valueOf(i).getBytes(StandardCharsets.UTF_8)); } dictionaryVector.setValueCount(vecLength); Dictionary dictionary = new Dictionary( dictionaryVector, new DictionaryEncoding( /* id= */ 10L, /* ordered= */ false, /* indexType= */ new ArrowType.Int(/* bitWidth= */ 8, /*isSigned*/ false))); try (VarCharVector vector = new VarCharVector("vector", allocator)) { setVector(vector, "255"); try (UInt1Vector encodedVector = (UInt1Vector) DictionaryEncoder.encode(vector, dictionary)) { // verify encoded result assertEquals(1, encodedVector.getValueCount()); assertEquals(255, encodedVector.getValueAsLong(0)); try (VarCharVector decodedVector = (VarCharVector) DictionaryEncoder.decode(encodedVector, dictionary)) { assertEquals(1, decodedVector.getValueCount()); assertArrayEquals("255".getBytes(StandardCharsets.UTF_8), decodedVector.get(0)); } } } } } private int[] convertListToIntArray(List list) { int[] values = new int[list.size()]; for (int i = 0; i < list.size(); i++) { values[i] = (int) list.get(i); } return values; } private Object[] convertMapValuesToArray(Map map) { Object[] values = new Object[map.size()]; Iterator valueIterator = map.values().iterator(); for (int i = 0; i < map.size(); i++) { values[i] = valueIterator.next(); } return values; } private void writeStructVector(NullableStructWriter writer, String value1, String value2) { byte[] bytes1 = value1.getBytes(StandardCharsets.UTF_8); byte[] bytes2 = value2.getBytes(StandardCharsets.UTF_8); ArrowBuf temp = allocator.buffer(bytes1.length > bytes2.length ? bytes1.length : bytes2.length); writer.start(); temp.setBytes(0, bytes1); writer.varChar("f0").writeVarChar(0, bytes1.length, temp); temp.setBytes(0, bytes2); writer.varChar("f1").writeVarChar(0, bytes2.length, temp); writer.end(); temp.close(); } private void writeStructVector(NullableStructWriter writer, int value1, long value2) { writer.start(); writer.integer("f0").writeInt(value1); writer.bigInt("f1").writeBigInt(value2); writer.end(); } private void writeListVector(UnionListWriter writer, int[] values) { writer.startList(); for (int v : values) { writer.integer().writeInt(v); } writer.endList(); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestDurationVector.java000066400000000000000000000132341515451615200322110ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertNull; import static org.junit.jupiter.api.Assertions.assertSame; import java.time.Duration; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.holders.NullableDurationHolder; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.pojo.ArrowType; 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 TestDurationVector { RootAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() { allocator.close(); } @Test public void testSecBasics() { try (DurationVector secVector = TestUtils.newVector( DurationVector.class, "second", new ArrowType.Duration(TimeUnit.SECOND), allocator)) { secVector.allocateNew(); secVector.setNull(0); secVector.setSafe(1, 1000); secVector.setValueCount(2); assertNull(secVector.getObject(0)); assertEquals(Duration.ofSeconds(1000), secVector.getObject(1)); assertNull(secVector.getAsStringBuilder(0)); assertEquals("PT16M40S", secVector.getAsStringBuilder(1).toString()); // Holder NullableDurationHolder holder = new NullableDurationHolder(); secVector.get(0, holder); assertEquals(0, holder.isSet); secVector.get(1, holder); assertEquals(1, holder.isSet); assertEquals(1000, holder.value); } } @Test public void testMilliBasics() { try (DurationVector milliVector = TestUtils.newVector( DurationVector.class, "nanos", new ArrowType.Duration(TimeUnit.MILLISECOND), allocator)) { milliVector.allocateNew(); milliVector.setNull(0); milliVector.setSafe(1, 1000); milliVector.setValueCount(2); assertNull(milliVector.getObject(0)); assertEquals(Duration.ofSeconds(1), milliVector.getObject(1)); assertNull(milliVector.getAsStringBuilder(0)); assertEquals("PT1S", milliVector.getAsStringBuilder(1).toString()); // Holder NullableDurationHolder holder = new NullableDurationHolder(); milliVector.get(0, holder); assertEquals(0, holder.isSet); milliVector.get(1, holder); assertEquals(1, holder.isSet); assertEquals(1000, holder.value); } } @Test public void testMicroBasics() { try (DurationVector microVector = TestUtils.newVector( DurationVector.class, "micro", new ArrowType.Duration(TimeUnit.MICROSECOND), allocator)) { microVector.allocateNew(); microVector.setNull(0); microVector.setSafe(1, 1000); microVector.setValueCount(2); assertNull(microVector.getObject(0)); assertEquals(Duration.ofMillis(1), microVector.getObject(1)); assertNull(microVector.getAsStringBuilder(0)); assertEquals("PT0.001S", microVector.getAsStringBuilder(1).toString()); // Holder NullableDurationHolder holder = new NullableDurationHolder(); microVector.get(0, holder); assertEquals(0, holder.isSet); microVector.get(1, holder); assertEquals(1, holder.isSet); assertEquals(1000, holder.value); } } @Test public void testNanosBasics() { try (DurationVector nanoVector = TestUtils.newVector( DurationVector.class, "nanos", new ArrowType.Duration(TimeUnit.NANOSECOND), allocator)) { nanoVector.allocateNew(); nanoVector.setNull(0); nanoVector.setSafe(1, 1000000); nanoVector.setValueCount(2); assertNull(nanoVector.getObject(0)); assertEquals(Duration.ofMillis(1), nanoVector.getObject(1)); assertNull(nanoVector.getAsStringBuilder(0)); assertEquals("PT0.001S", nanoVector.getAsStringBuilder(1).toString()); // Holder NullableDurationHolder holder = new NullableDurationHolder(); nanoVector.get(0, holder); assertEquals(0, holder.isSet); nanoVector.get(1, holder); assertEquals(1, holder.isSet); assertEquals(1000000, holder.value); } } @Test public void testGetTransferPairWithField() { final DurationVector fromVector = TestUtils.newVector( DurationVector.class, "nanos", new ArrowType.Duration(TimeUnit.NANOSECOND), allocator); final TransferPair transferPair = fromVector.getTransferPair(fromVector.getField(), allocator); final DurationVector toVector = (DurationVector) transferPair.getTo(); // Field inside a new vector created by reusing a field should be the same in memory as the // original field. assertSame(fromVector.getField(), toVector.getField()); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeBinaryVector.java000066400000000000000000000230731515451615200334650ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; 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 org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.holders.FixedSizeBinaryHolder; import org.apache.arrow.vector.holders.NullableFixedSizeBinaryHolder; import org.apache.arrow.vector.util.ReusableByteArray; 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 TestFixedSizeBinaryVector { private static final int numValues = 123; private static final int typeWidth = 9; private static final int smallDataSize = 6; private static final int largeDataSize = 12; private static byte[][] values; static { values = new byte[numValues][typeWidth]; for (int i = 0; i < numValues; i++) { for (int j = 0; j < typeWidth; j++) { values[i][j] = ((byte) i); } } } private ArrowBuf[] bufs = new ArrowBuf[numValues]; private FixedSizeBinaryHolder[] holders = new FixedSizeBinaryHolder[numValues]; private NullableFixedSizeBinaryHolder[] nullableHolders = new NullableFixedSizeBinaryHolder[numValues]; private static byte[] smallValue; static { smallValue = new byte[smallDataSize]; for (int i = 0; i < smallDataSize; i++) { smallValue[i] = ((byte) i); } } private ArrowBuf smallBuf; private FixedSizeBinaryHolder smallHolder; private NullableFixedSizeBinaryHolder smallNullableHolder; private static byte[] largeValue; static { largeValue = new byte[largeDataSize]; for (int i = 0; i < largeDataSize; i++) { largeValue[i] = ((byte) i); } } private ArrowBuf largeBuf; private FixedSizeBinaryHolder largeHolder; private NullableFixedSizeBinaryHolder largeNullableHolder; private BufferAllocator allocator; private FixedSizeBinaryVector vector; private static void failWithException(String message) throws Exception { throw new Exception(message); } @BeforeEach public void init() throws Exception { allocator = new DirtyRootAllocator(Integer.MAX_VALUE, (byte) 100); vector = new FixedSizeBinaryVector("fixedSizeBinary", allocator, typeWidth); vector.allocateNew(); for (int i = 0; i < numValues; i++) { bufs[i] = allocator.buffer(typeWidth); bufs[i].setBytes(0, values[i]); holders[i] = new FixedSizeBinaryHolder(); holders[i].byteWidth = typeWidth; holders[i].buffer = bufs[i]; nullableHolders[i] = new NullableFixedSizeBinaryHolder(); nullableHolders[i].byteWidth = typeWidth; nullableHolders[i].buffer = bufs[i]; nullableHolders[i].isSet = 1; } smallBuf = allocator.buffer(smallDataSize); smallBuf.setBytes(0, smallValue); smallHolder = new FixedSizeBinaryHolder(); smallHolder.byteWidth = smallDataSize; smallHolder.buffer = smallBuf; smallNullableHolder = new NullableFixedSizeBinaryHolder(); smallNullableHolder.byteWidth = smallDataSize; smallNullableHolder.buffer = smallBuf; largeBuf = allocator.buffer(largeDataSize); largeBuf.setBytes(0, largeValue); largeHolder = new FixedSizeBinaryHolder(); largeHolder.byteWidth = typeWidth; largeHolder.buffer = largeBuf; largeNullableHolder = new NullableFixedSizeBinaryHolder(); largeNullableHolder.byteWidth = typeWidth; largeNullableHolder.buffer = largeBuf; } @AfterEach public void terminate() throws Exception { for (int i = 0; i < numValues; i++) { bufs[i].close(); } smallBuf.close(); largeBuf.close(); vector.close(); allocator.close(); } @Test public void testSetUsingByteArray() { for (int i = 0; i < numValues; i++) { vector.set(i, values[i]); } vector.setValueCount(numValues); for (int i = 0; i < numValues; i++) { assertArrayEquals(values[i], vector.getObject(i)); } } @Test public void testSetUsingNull() { final byte[] value = null; for (int i = 0; i < numValues; i++) { final int index = i; Exception e = assertThrows( NullPointerException.class, () -> { vector.set(index, value); }); assertEquals("expecting a valid byte array", e.getMessage()); } } @Test public void testSetUsingHolder() { for (int i = 0; i < numValues; i++) { vector.set(i, holders[i]); } vector.setValueCount(numValues); for (int i = 0; i < numValues; i++) { assertArrayEquals(values[i], vector.getObject(i)); } } @Test public void testSetUsingNullableHolder() { for (int i = 0; i < numValues; i++) { vector.set(i, nullableHolders[i]); } vector.setValueCount(numValues); for (int i = 0; i < numValues; i++) { assertArrayEquals(values[i], vector.getObject(i)); } } @Test public void testGetUsingNullableHolder() { for (int i = 0; i < numValues; i++) { vector.set(i, holders[i]); } vector.setValueCount(numValues); for (int i = 0; i < numValues; i++) { vector.get(i, nullableHolders[i]); assertEquals(typeWidth, nullableHolders[i].byteWidth); assertTrue(nullableHolders[i].isSet > 0); byte[] actual = new byte[typeWidth]; nullableHolders[i].buffer.getBytes(0, actual, 0, typeWidth); assertArrayEquals(values[i], actual); } } @Test public void testSetWithInvalidInput() throws Exception { String errorMsg = "input data needs to be at least " + typeWidth + " bytes"; // test small inputs, byteWidth matches but value or buffer is too small try { vector.set(0, smallValue); failWithException(errorMsg); } catch (AssertionError | IllegalArgumentException ignore) { } try { vector.set(0, smallHolder); failWithException(errorMsg); } catch (AssertionError | IllegalArgumentException ignore) { } try { vector.set(0, smallNullableHolder); failWithException(errorMsg); } catch (AssertionError | IllegalArgumentException ignore) { } try { vector.set(0, smallBuf); failWithException(errorMsg); } catch (AssertionError | IllegalArgumentException ignore) { } // test large inputs, byteWidth matches but value or buffer is bigger than byteWidth vector.set(0, largeValue); vector.set(0, largeHolder); vector.set(0, largeNullableHolder); vector.set(0, largeBuf); } @Test public void setSetSafeWithInvalidInput() throws Exception { String errorMsg = "input data needs to be at least " + typeWidth + " bytes"; // test small inputs, byteWidth matches but value or buffer is too small try { vector.setSafe(0, smallValue); failWithException(errorMsg); } catch (AssertionError | IllegalArgumentException ignore) { } try { vector.setSafe(0, smallHolder); failWithException(errorMsg); } catch (AssertionError | IllegalArgumentException ignore) { } try { vector.setSafe(0, smallNullableHolder); failWithException(errorMsg); } catch (AssertionError | IllegalArgumentException ignore) { } try { vector.setSafe(0, smallBuf); failWithException(errorMsg); } catch (AssertionError | IllegalArgumentException ignore) { } // test large inputs, byteWidth matches but value or buffer is bigger than byteWidth vector.setSafe(0, largeValue); vector.setSafe(0, largeHolder); vector.setSafe(0, largeNullableHolder); vector.setSafe(0, largeBuf); } @Test public void testGetNull() { vector.setNull(0); assertNull(vector.get(0)); } @Test public void testGetTransferPairWithField() { final FixedSizeBinaryVector fromVector = new FixedSizeBinaryVector("fixedSizeBinary", allocator, typeWidth); final TransferPair transferPair = fromVector.getTransferPair(fromVector.getField(), allocator); final FixedSizeBinaryVector toVector = (FixedSizeBinaryVector) transferPair.getTo(); // Field inside a new vector created by reusing a field should be the same in memory as the // original field. assertSame(fromVector.getField(), toVector.getField()); } @Test public void testGetBytesRepeatedly() { for (int i = 0; i < numValues; i++) { vector.set(i, values[i]); } vector.setValueCount(numValues); ReusableByteArray reusableByteArray = new ReusableByteArray(); for (int i = 0; i < numValues; i++) { // verify results vector.read(i, reusableByteArray); assertArrayEquals(values[i], reusableByteArray.getBuffer()); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java000066400000000000000000000737111515451615200331600ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.math.BigDecimal; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.util.ArrayList; 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; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestFixedSizeListVector { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testIntType() { try (FixedSizeListVector vector = FixedSizeListVector.empty("list", /* size= */ 2, allocator)) { IntVector nested = (IntVector) vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())).getVector(); vector.allocateNew(); for (int i = 0; i < 10; i++) { vector.setNotNull(i); nested.set(i * 2, i); nested.set(i * 2 + 1, i + 10); } vector.setValueCount(10); UnionFixedSizeListReader reader = vector.getReader(); for (int i = 0; i < 10; i++) { reader.setPosition(i); assertTrue(reader.isSet()); assertTrue(reader.next()); assertEquals(i, reader.reader().readInteger().intValue()); assertTrue(reader.next()); assertEquals(i + 10, reader.reader().readInteger().intValue()); assertFalse(reader.next()); assertEquals(Arrays.asList(i, i + 10), reader.readObject()); } } } @Test public void testFloatTypeNullable() { try (FixedSizeListVector vector = FixedSizeListVector.empty("list", /* size= */ 2, allocator)) { Float4Vector nested = (Float4Vector) vector.addOrGetVector(FieldType.nullable(MinorType.FLOAT4.getType())).getVector(); vector.allocateNew(); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector.setNotNull(i); nested.set(i * 2, i + 0.1f); nested.set(i * 2 + 1, i + 10.1f); } } vector.setValueCount(10); UnionFixedSizeListReader reader = vector.getReader(); for (int i = 0; i < 10; i++) { reader.setPosition(i); if (i % 2 == 0) { assertTrue(reader.isSet()); assertTrue(reader.next()); assertEquals(i + 0.1f, reader.reader().readFloat(), 0.00001); assertTrue(reader.next()); assertEquals(i + 10.1f, reader.reader().readFloat(), 0.00001); assertFalse(reader.next()); assertEquals(Arrays.asList(i + 0.1f, i + 10.1f), reader.readObject()); } else { assertFalse(reader.isSet()); assertNull(reader.readObject()); } } } } @Test public void testNestedInList() { try (ListVector vector = ListVector.empty("list", allocator)) { FixedSizeListVector tuples = (FixedSizeListVector) vector.addOrGetVector(FieldType.nullable(new ArrowType.FixedSizeList(2))).getVector(); IntVector innerVector = (IntVector) tuples.addOrGetVector(FieldType.nullable(MinorType.INT.getType())).getVector(); vector.allocateNew(); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { int position = vector.startNewValue(i); for (int j = 0; j < i % 7; j++) { tuples.setNotNull(position + j); innerVector.set((position + j) * 2, j); innerVector.set((position + j) * 2 + 1, j + 1); } vector.endValue(i, i % 7); } } vector.setValueCount(10); UnionListReader reader = vector.getReader(); for (int i = 0; i < 10; i++) { reader.setPosition(i); if (i % 2 == 0) { for (int j = 0; j < i % 7; j++) { assertTrue(reader.next()); FieldReader innerListReader = reader.reader(); for (int k = 0; k < 2; k++) { assertTrue(innerListReader.next()); assertEquals(k + j, innerListReader.reader().readInteger().intValue()); } assertFalse(innerListReader.next()); } assertFalse(reader.next()); } else { assertFalse(reader.isSet()); assertNull(reader.readObject()); } } } } @Test public void testTransferPair() { try (FixedSizeListVector from = new FixedSizeListVector( "from", allocator, new FieldType(true, new ArrowType.FixedSizeList(2), null), null); FixedSizeListVector to = new FixedSizeListVector( "to", allocator, new FieldType(true, new ArrowType.FixedSizeList(2), null), null)) { Float4Vector nested = (Float4Vector) from.addOrGetVector(FieldType.nullable(MinorType.FLOAT4.getType())).getVector(); from.allocateNew(); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { from.setNotNull(i); nested.set(i * 2, i + 0.1f); nested.set(i * 2 + 1, i + 10.1f); } } from.setValueCount(10); TransferPair pair = from.makeTransferPair(to); pair.copyValueSafe(0, 1); pair.copyValueSafe(2, 2); to.copyFromSafe(4, 3, from); to.setValueCount(10); UnionFixedSizeListReader reader = to.getReader(); reader.setPosition(0); assertFalse(reader.isSet()); assertNull(reader.readObject()); reader.setPosition(1); assertTrue(reader.isSet()); assertTrue(reader.next()); assertEquals(0.1f, reader.reader().readFloat(), 0.00001); assertTrue(reader.next()); assertEquals(10.1f, reader.reader().readFloat(), 0.00001); assertFalse(reader.next()); assertEquals(Arrays.asList(0.1f, 10.1f), reader.readObject()); reader.setPosition(2); assertTrue(reader.isSet()); assertTrue(reader.next()); assertEquals(2.1f, reader.reader().readFloat(), 0.00001); assertTrue(reader.next()); assertEquals(12.1f, reader.reader().readFloat(), 0.00001); assertFalse(reader.next()); assertEquals(Arrays.asList(2.1f, 12.1f), reader.readObject()); reader.setPosition(3); assertTrue(reader.isSet()); assertTrue(reader.next()); assertEquals(4.1f, reader.reader().readFloat(), 0.00001); assertTrue(reader.next()); assertEquals(14.1f, reader.reader().readFloat(), 0.00001); assertFalse(reader.next()); assertEquals(Arrays.asList(4.1f, 14.1f), reader.readObject()); for (int i = 4; i < 10; i++) { reader.setPosition(i); assertFalse(reader.isSet()); assertNull(reader.readObject()); } } } @Test public void testTransferEmptyVector() throws Exception { // #43320 try (FixedSizeListVector src = new FixedSizeListVector( "src", allocator, FieldType.nullable(new ArrowType.FixedSizeList(2)), null); FixedSizeListVector dest = new FixedSizeListVector( "dest", allocator, FieldType.nullable(new ArrowType.FixedSizeList(2)), null)) { src.makeTransferPair(dest).transfer(); IntVector els = (IntVector) dest.addOrGetVector(FieldType.nullable(MinorType.INT.getType())).getVector(); dest.allocateNew(); dest.startNewValue(0); els.setSafe(0, 1); els.setSafe(1, 2); dest.setValueCount(1); List expected = new ArrayList<>(2); expected.add(1); expected.add(2); assertEquals(expected, dest.getObject(0)); } } @Test public void testConsistentChildName() throws Exception { try (FixedSizeListVector listVector = FixedSizeListVector.empty("sourceVector", /* size= */ 2, allocator)) { String emptyListStr = listVector.getField().toString(); assertTrue(emptyListStr.contains(ListVector.DATA_VECTOR_NAME)); listVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); String emptyVectorStr = listVector.getField().toString(); assertTrue(emptyVectorStr.contains(ListVector.DATA_VECTOR_NAME)); } } @Test public void testUnionFixedSizeListWriterWithNulls() throws Exception { /* Write to a decimal list vector * each list of size 3 and having its data values alternating between null and a non-null. * Read and verify */ try (final FixedSizeListVector vector = FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { UnionFixedSizeListWriter writer = vector.getWriter(); writer.allocate(); final int valueCount = 100; for (int i = 0; i < valueCount; i++) { writer.startList(); writer.decimal().writeDecimal(new BigDecimal(i)); writer.writeNull(); writer.decimal().writeDecimal(new BigDecimal(i * 3)); writer.endList(); } vector.setValueCount(valueCount); for (int i = 0; i < valueCount; i++) { List values = (List) vector.getObject(i); assertEquals(3, values.size()); assertEquals(new BigDecimal(i), values.get(0)); assertEquals(null, values.get(1)); assertEquals(new BigDecimal(i * 3), values.get(2)); } } } @Test public void testUnionFixedSizeListWriter() throws Exception { try (final FixedSizeListVector vector1 = FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { UnionFixedSizeListWriter writer1 = vector1.getWriter(); writer1.allocate(); int[] values1 = new int[] {1, 2, 3}; int[] values2 = new int[] {4, 5, 6}; int[] values3 = new int[] {7, 8, 9}; // set some values writeListVector(vector1, writer1, values1); writeListVector(vector1, writer1, values2); writeListVector(vector1, writer1, values3); writer1.setValueCount(3); assertEquals(3, vector1.getValueCount()); int[] realValue1 = convertListToIntArray(vector1.getObject(0)); assertTrue(Arrays.equals(values1, realValue1)); int[] realValue2 = convertListToIntArray(vector1.getObject(1)); assertTrue(Arrays.equals(values2, realValue2)); int[] realValue3 = convertListToIntArray(vector1.getObject(2)); assertTrue(Arrays.equals(values3, realValue3)); } } @Test public void testWriteDecimal() throws Exception { try (final FixedSizeListVector vector = FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { UnionFixedSizeListWriter writer = vector.getWriter(); writer.allocate(); final int valueCount = 100; for (int i = 0; i < valueCount; i++) { writer.startList(); writer.decimal().writeDecimal(new BigDecimal(i)); writer.decimal().writeDecimal(new BigDecimal(i * 2)); writer.decimal().writeDecimal(new BigDecimal(i * 3)); writer.endList(); } vector.setValueCount(valueCount); for (int i = 0; i < valueCount; i++) { List values = (List) vector.getObject(i); assertEquals(3, values.size()); assertEquals(new BigDecimal(i), values.get(0)); assertEquals(new BigDecimal(i * 2), values.get(1)); assertEquals(new BigDecimal(i * 3), values.get(2)); } } } @Test public void testDecimalIndexCheck() throws Exception { try (final FixedSizeListVector vector = FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { UnionFixedSizeListWriter writer = vector.getWriter(); writer.allocate(); IllegalStateException e = assertThrows( IllegalStateException.class, () -> { writer.startList(); writer.decimal().writeDecimal(new BigDecimal(1)); writer.decimal().writeDecimal(new BigDecimal(2)); writer.decimal().writeDecimal(new BigDecimal(3)); writer.decimal().writeDecimal(new BigDecimal(4)); writer.endList(); }); assertEquals("values at index 0 is greater than listSize 3", e.getMessage()); } } @Test public void testWriteIllegalData() throws Exception { assertThrows( IllegalStateException.class, () -> { try (final FixedSizeListVector vector1 = FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { UnionFixedSizeListWriter writer1 = vector1.getWriter(); writer1.allocate(); int[] values1 = new int[] {1, 2, 3}; int[] values2 = new int[] {4, 5, 6, 7, 8}; // set some values writeListVector(vector1, writer1, values1); writeListVector(vector1, writer1, values2); writer1.setValueCount(3); assertEquals(3, vector1.getValueCount()); int[] realValue1 = convertListToIntArray(vector1.getObject(0)); assertTrue(Arrays.equals(values1, realValue1)); int[] realValue2 = convertListToIntArray(vector1.getObject(1)); assertTrue(Arrays.equals(values2, realValue2)); } }); } @Test public void testSplitAndTransfer() throws Exception { try (final FixedSizeListVector vector1 = FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { UnionFixedSizeListWriter writer1 = vector1.getWriter(); writer1.allocate(); int[] values1 = new int[] {1, 2, 3}; int[] values2 = new int[] {4, 5, 6}; int[] values3 = new int[] {7, 8, 9}; // set some values writeListVector(vector1, writer1, values1); writeListVector(vector1, writer1, values2); writeListVector(vector1, writer1, values3); writer1.setValueCount(3); TransferPair transferPair = vector1.getTransferPair(allocator); transferPair.splitAndTransfer(0, 2); FixedSizeListVector targetVector = (FixedSizeListVector) transferPair.getTo(); assertEquals(2, targetVector.getValueCount()); int[] realValue1 = convertListToIntArray(targetVector.getObject(0)); assertArrayEquals(values1, realValue1); int[] realValue2 = convertListToIntArray(targetVector.getObject(1)); assertArrayEquals(values2, realValue2); targetVector.clear(); } } @Test public void testZeroWidthVector() { try (final FixedSizeListVector vector1 = FixedSizeListVector.empty("vector", /* size= */ 0, allocator)) { UnionFixedSizeListWriter writer1 = vector1.getWriter(); writer1.allocate(); int[] values1 = new int[] {}; int[] values2 = new int[] {}; int[] values3 = null; int[] values4 = new int[] {}; // set some values writeListVector(vector1, writer1, values1); writeListVector(vector1, writer1, values2); writeListVector(vector1, writer1, values3); writeListVector(vector1, writer1, values4); writer1.setValueCount(4); assertEquals(4, vector1.getValueCount()); int[] realValue1 = convertListToIntArray(vector1.getObject(0)); assertArrayEquals(values1, realValue1); int[] realValue2 = convertListToIntArray(vector1.getObject(1)); assertArrayEquals(values2, realValue2); assertNull(vector1.getObject(2)); int[] realValue4 = convertListToIntArray(vector1.getObject(3)); assertArrayEquals(values4, realValue4); } } @Test public void testVectorWithNulls() { try (final FixedSizeListVector vector1 = FixedSizeListVector.empty("vector", /* size= */ 4, allocator)) { UnionFixedSizeListWriter writer1 = vector1.getWriter(); writer1.allocate(); List values1 = Arrays.asList(null, 1, 2, 3); List values2 = Arrays.asList(4, null, 5, 6); List values3 = null; List values4 = Arrays.asList(7, 8, null, 9); // set some values writeListVector(vector1, writer1, values1); writeListVector(vector1, writer1, values2); writeListVector(vector1, writer1, values3); writeListVector(vector1, writer1, values4); writer1.setValueCount(4); assertEquals(4, vector1.getValueCount()); List realValue1 = vector1.getObject(0); assertEquals(values1, realValue1); List realValue2 = vector1.getObject(1); assertEquals(values2, realValue2); List realValue3 = vector1.getObject(2); assertEquals(values3, realValue3); List realValue4 = vector1.getObject(3); assertEquals(values4, realValue4); } } @Test public void testWriteVarCharHelpers() throws Exception { try (final FixedSizeListVector vector = FixedSizeListVector.empty("vector", /* size= */ 4, allocator)) { UnionFixedSizeListWriter writer = vector.getWriter(); writer.allocate(); writer.startList(); writer.writeVarChar("row1,1"); writer.writeVarChar(new Text("row1,2")); writer.writeNull(); writer.writeNull(); writer.endList(); assertEquals("row1,1", vector.getObject(0).get(0).toString()); assertEquals("row1,2", vector.getObject(0).get(1).toString()); } } @Test public void testWriteLargeVarCharHelpers() throws Exception { try (final FixedSizeListVector vector = FixedSizeListVector.empty("vector", /* size= */ 4, allocator)) { UnionFixedSizeListWriter writer = vector.getWriter(); writer.allocate(); writer.startList(); writer.writeLargeVarChar("row1,1"); writer.writeLargeVarChar(new Text("row1,2")); writer.writeNull(); writer.writeNull(); writer.endList(); assertEquals("row1,1", vector.getObject(0).get(0).toString()); assertEquals("row1,2", vector.getObject(0).get(1).toString()); } } @Test public void testWriteVarBinaryHelpers() throws Exception { try (final FixedSizeListVector vector = FixedSizeListVector.empty("vector", /* size= */ 4, allocator)) { UnionFixedSizeListWriter writer = vector.getWriter(); writer.allocate(); writer.startList(); writer.writeVarBinary("row1,1".getBytes(StandardCharsets.UTF_8)); writer.writeVarBinary( "row1,2".getBytes(StandardCharsets.UTF_8), 0, "row1,2".getBytes(StandardCharsets.UTF_8).length); writer.writeVarBinary(ByteBuffer.wrap("row1,3".getBytes(StandardCharsets.UTF_8))); writer.writeVarBinary( ByteBuffer.wrap("row1,4".getBytes(StandardCharsets.UTF_8)), 0, "row1,4".getBytes(StandardCharsets.UTF_8).length); writer.endList(); assertEquals( "row1,1", new String((byte[]) vector.getObject(0).get(0), StandardCharsets.UTF_8)); assertEquals( "row1,2", new String((byte[]) vector.getObject(0).get(1), StandardCharsets.UTF_8)); assertEquals( "row1,3", new String((byte[]) vector.getObject(0).get(2), StandardCharsets.UTF_8)); assertEquals( "row1,4", new String((byte[]) vector.getObject(0).get(3), StandardCharsets.UTF_8)); } } @Test public void testWriteLargeVarBinaryHelpers() throws Exception { try (final FixedSizeListVector vector = FixedSizeListVector.empty("vector", /* size= */ 4, allocator)) { UnionFixedSizeListWriter writer = vector.getWriter(); writer.allocate(); writer.startList(); writer.writeLargeVarBinary("row1,1".getBytes(StandardCharsets.UTF_8)); writer.writeLargeVarBinary( "row1,2".getBytes(StandardCharsets.UTF_8), 0, "row1,2".getBytes(StandardCharsets.UTF_8).length); writer.writeLargeVarBinary(ByteBuffer.wrap("row1,3".getBytes(StandardCharsets.UTF_8))); writer.writeLargeVarBinary( ByteBuffer.wrap("row1,4".getBytes(StandardCharsets.UTF_8)), 0, "row1,4".getBytes(StandardCharsets.UTF_8).length); writer.endList(); assertEquals( "row1,1", new String((byte[]) vector.getObject(0).get(0), StandardCharsets.UTF_8)); assertEquals( "row1,2", new String((byte[]) vector.getObject(0).get(1), StandardCharsets.UTF_8)); assertEquals( "row1,3", new String((byte[]) vector.getObject(0).get(2), StandardCharsets.UTF_8)); assertEquals( "row1,4", new String((byte[]) vector.getObject(0).get(3), StandardCharsets.UTF_8)); } } @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++) { values[i] = (int) list.get(i); } return values; } private void writeListVector( FixedSizeListVector vector, UnionFixedSizeListWriter writer, int[] values) { writer.startList(); if (values != null) { for (int v : values) { writer.integer().writeInt(v); } } else { vector.setNull(writer.getPosition()); } writer.endList(); } private void writeListVector( FixedSizeListVector vector, UnionFixedSizeListWriter writer, List values) { writer.startList(); if (values != null) { for (Integer v : values) { if (v == null) { writer.writeNull(); } else { writer.integer().writeInt(v); } } } else { vector.setNull(writer.getPosition()); } writer.endList(); } } TestIntervalMonthDayNanoVector.java000066400000000000000000000104171515451615200344110ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.time.Duration; import java.time.Period; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.holders.IntervalMonthDayNanoHolder; import org.apache.arrow.vector.holders.NullableIntervalMonthDayNanoHolder; import org.apache.arrow.vector.types.IntervalUnit; import org.apache.arrow.vector.types.Types; 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; public class TestIntervalMonthDayNanoVector { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testBasics() { try (final IntervalMonthDayNanoVector vector = new IntervalMonthDayNanoVector(/* name= */ "", allocator)) { int valueCount = 100; vector.setInitialCapacity(valueCount); vector.allocateNew(); NullableIntervalMonthDayNanoHolder nullableHolder = new NullableIntervalMonthDayNanoHolder(); nullableHolder.isSet = 1; nullableHolder.months = 2; nullableHolder.days = 20; nullableHolder.nanoseconds = 123; IntervalMonthDayNanoHolder holder = new IntervalMonthDayNanoHolder(); holder.months = Integer.MIN_VALUE; 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(3, holder); nullableHolder.isSet = 0; vector.setSafe(/* index= */ 5, nullableHolder); vector.setValueCount(5); assertEquals("P1M2D PT-0.000000002S ", vector.getAsStringBuilder(0).toString()); assertEquals(null, vector.getAsStringBuilder(1)); assertEquals("P1M2D PT-0.000000003S ", vector.getAsStringBuilder(2).toString()); assertEquals( new PeriodDuration( Period.of(0, Integer.MIN_VALUE, Integer.MIN_VALUE), Duration.ofNanos(Long.MIN_VALUE)), vector.getObject(3)); assertEquals("P2M20D PT0.000000123S ", vector.getAsStringBuilder(4).toString()); assertEquals(null, vector.getObject(5)); vector.get(1, nullableHolder); assertEquals(0, nullableHolder.isSet); vector.get(2, nullableHolder); assertEquals(1, nullableHolder.isSet); assertEquals(1, nullableHolder.months); assertEquals(2, nullableHolder.days); assertEquals(-3, nullableHolder.nanoseconds); IntervalMonthDayNanoVector.getDays(vector.valueBuffer, 2); assertEquals(1, IntervalMonthDayNanoVector.getMonths(vector.valueBuffer, 2)); assertEquals(2, IntervalMonthDayNanoVector.getDays(vector.valueBuffer, 2)); assertEquals(-3, IntervalMonthDayNanoVector.getNanoseconds(vector.valueBuffer, 2)); assertEquals(0, vector.isSet(1)); assertEquals(1, vector.isSet(2)); assertEquals(Types.MinorType.INTERVALMONTHDAYNANO, vector.getMinorType()); ArrowType fieldType = vector.getField().getType(); assertEquals(ArrowType.ArrowTypeID.Interval, fieldType.getTypeID()); ArrowType.Interval intervalType = (ArrowType.Interval) fieldType; assertEquals(IntervalUnit.MONTH_DAY_NANO, intervalType.getUnit()); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestIntervalYearVector.java000066400000000000000000000060321515451615200330270ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertSame; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.types.IntervalUnit; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.ArrowType; 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 TestIntervalYearVector { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testGetAsStringBuilder() { try (final IntervalYearVector vector = new IntervalYearVector("", allocator)) { int valueCount = 100; vector.setInitialCapacity(valueCount); vector.allocateNew(); for (int i = 0; i < valueCount; i++) { vector.set(i, i); } assertEquals("0 years 1 month ", vector.getAsStringBuilder(1).toString()); assertEquals("0 years 10 months ", vector.getAsStringBuilder(10).toString()); assertEquals("1 year 8 months ", vector.getAsStringBuilder(20).toString()); assertEquals("2 years 6 months ", vector.getAsStringBuilder(30).toString()); assertEquals(Types.MinorType.INTERVALYEAR, vector.getMinorType()); ArrowType fieldType = vector.getField().getType(); assertEquals(ArrowType.ArrowTypeID.Interval, fieldType.getTypeID()); ArrowType.Interval intervalType = (ArrowType.Interval) fieldType; assertEquals(IntervalUnit.YEAR_MONTH, intervalType.getUnit()); } } @Test public void testGetTransferPairWithField() { final IntervalYearVector fromVector = new IntervalYearVector("", allocator); final TransferPair transferPair = fromVector.getTransferPair(fromVector.getField(), allocator); final IntervalYearVector toVector = (IntervalYearVector) transferPair.getTo(); // Field inside a new vector created by reusing a field should be the same in memory as the // original field. assertSame(fromVector.getField(), toVector.getField()); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java000066400000000000000000001176441515451615200323240ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; 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; public class TestLargeListVector { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testCopyFrom() throws Exception { try (LargeListVector inVector = LargeListVector.empty("input", allocator); LargeListVector outVector = LargeListVector.empty("output", allocator)) { UnionLargeListWriter writer = inVector.getWriter(); writer.allocate(); // populate input vector with the following records // [1, 2, 3] // null // [] writer.setPosition(0); // optional writer.startList(); writer.bigInt().writeBigInt(1); writer.bigInt().writeBigInt(2); writer.bigInt().writeBigInt(3); writer.endList(); writer.setPosition(2); writer.startList(); writer.endList(); writer.setValueCount(3); // copy values from input to output outVector.allocateNew(); for (int i = 0; i < 3; i++) { outVector.copyFrom(i, i, inVector); } outVector.setValueCount(3); // assert the output vector is correct FieldReader reader = outVector.getReader(); assertTrue(reader.isSet(), "shouldn't be null"); reader.setPosition(1); assertFalse(reader.isSet(), "should be null"); reader.setPosition(2); assertTrue(reader.isSet(), "shouldn't be null"); /* index 0 */ Object result = outVector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Long.valueOf(1), resultSet.get(0)); assertEquals(Long.valueOf(2), resultSet.get(1)); assertEquals(Long.valueOf(3), resultSet.get(2)); /* index 1 */ result = outVector.getObject(1); assertNull(result); /* index 2 */ result = outVector.getObject(2); resultSet = (ArrayList) result; assertEquals(0, resultSet.size()); /* 3+0+0/3 */ assertEquals(1.0D, inVector.getDensity(), 0); } } @Test public void testSetLastSetUsage() throws Exception { try (LargeListVector listVector = LargeListVector.empty("input", allocator)) { /* Explicitly add the dataVector */ MinorType type = MinorType.BIGINT; listVector.addOrGetVector(FieldType.nullable(type.getType())); /* allocate memory */ listVector.allocateNew(); /* get inner buffers; validityBuffer and offsetBuffer */ ArrowBuf validityBuffer = listVector.getValidityBuffer(); ArrowBuf offsetBuffer = listVector.getOffsetBuffer(); /* get the underlying data vector -- BigIntVector */ BigIntVector dataVector = (BigIntVector) listVector.getDataVector(); /* check current lastSet */ assertEquals(-1L, listVector.getLastSet()); int index = 0; int offset; /* write [10, 11, 12] to the list vector at index 0 */ BitVectorHelper.setBit(validityBuffer, index); dataVector.setSafe(0, 1, 10); dataVector.setSafe(1, 1, 11); dataVector.setSafe(2, 1, 12); offsetBuffer.setLong((index + 1) * LargeListVector.OFFSET_WIDTH, 3); index += 1; /* write [13, 14] to the list vector at index 1 */ BitVectorHelper.setBit(validityBuffer, index); dataVector.setSafe(3, 1, 13); dataVector.setSafe(4, 1, 14); offsetBuffer.setLong((index + 1) * LargeListVector.OFFSET_WIDTH, 5); index += 1; /* write [15, 16, 17] to the list vector at index 2 */ BitVectorHelper.setBit(validityBuffer, index); dataVector.setSafe(5, 1, 15); dataVector.setSafe(6, 1, 16); dataVector.setSafe(7, 1, 17); offsetBuffer.setLong((index + 1) * LargeListVector.OFFSET_WIDTH, 8); /* check current lastSet */ assertEquals(-1L, listVector.getLastSet()); /* set lastset and arbitrary valuecount for list vector. * * NOTE: if we don't execute setLastSet() before setLastValueCount(), then * the latter will corrupt the offsetBuffer and thus the accessor will not * retrieve the correct values from underlying dataBuffer. Run the test * by commenting out next line and we should see failures from 5th assert * onwards. This is why doing setLastSet() is important before setValueCount() * once the vector has been loaded. * * Another important thing to remember is the value of lastSet itself. * Even though the listVector has elements till index 2 only, the lastSet should * be set as 3. This is because the offsetBuffer has valid offsets filled till index 3. * If we do setLastSet(2), the offsetBuffer at index 3 will contain incorrect value * after execution of setValueCount(). * * correct state of the listVector * bitvector {1, 1, 1, 0, 0.... } * offsetvector {0, 3, 5, 8, 8, 8.....} * datavector { [10, 11, 12], * [13, 14], * [15, 16, 17] * } * * if we don't do setLastSet() before setValueCount --> incorrect state * bitvector {1, 1, 1, 0, 0.... } * offsetvector {0, 0, 0, 0, 0, 0.....} * datavector { [10, 11, 12], * [13, 14], * [15, 16, 17] * } * * if we do setLastSet(2) before setValueCount --> incorrect state * bitvector {1, 1, 1, 0, 0.... } * offsetvector {0, 3, 5, 5, 5, 5.....} * datavector { [10, 11, 12], * [13, 14], * [15, 16, 17] * } */ listVector.setLastSet(2); listVector.setValueCount(10); /* (3+2+3)/10 */ assertEquals(0.8D, listVector.getDensity(), 0); index = 0; offset = (int) offsetBuffer.getLong(index * LargeListVector.OFFSET_WIDTH); assertEquals(Integer.toString(0), Integer.toString(offset)); Long actual = dataVector.getObject(offset); assertEquals(Long.valueOf(10), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(11), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(12), actual); index++; offset = (int) offsetBuffer.getLong(index * LargeListVector.OFFSET_WIDTH); assertEquals(Integer.toString(3), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(13), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(14), actual); index++; offset = (int) offsetBuffer.getLong(index * LargeListVector.OFFSET_WIDTH); assertEquals(Integer.toString(5), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(15), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(16), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(17), actual); index++; offset = (int) offsetBuffer.getLong(index * LargeListVector.OFFSET_WIDTH); assertEquals(Integer.toString(8), Integer.toString(offset)); actual = dataVector.getObject(offset); assertNull(actual); } } @Test public void testSplitAndTransfer() throws Exception { try (LargeListVector listVector = LargeListVector.empty("sourceVector", allocator)) { /* Explicitly add the dataVector */ MinorType type = MinorType.BIGINT; listVector.addOrGetVector(FieldType.nullable(type.getType())); UnionLargeListWriter listWriter = listVector.getWriter(); /* allocate memory */ listWriter.allocate(); /* populate data */ listWriter.setPosition(0); listWriter.startList(); listWriter.bigInt().writeBigInt(10); listWriter.bigInt().writeBigInt(11); listWriter.bigInt().writeBigInt(12); listWriter.endList(); listWriter.setPosition(1); listWriter.startList(); listWriter.bigInt().writeBigInt(13); listWriter.bigInt().writeBigInt(14); listWriter.endList(); listWriter.setPosition(2); listWriter.startList(); listWriter.bigInt().writeBigInt(15); listWriter.bigInt().writeBigInt(16); listWriter.bigInt().writeBigInt(17); listWriter.bigInt().writeBigInt(18); listWriter.endList(); listWriter.setPosition(3); listWriter.startList(); listWriter.bigInt().writeBigInt(19); listWriter.endList(); listWriter.setPosition(4); listWriter.startList(); listWriter.bigInt().writeBigInt(20); listWriter.bigInt().writeBigInt(21); listWriter.bigInt().writeBigInt(22); listWriter.bigInt().writeBigInt(23); listWriter.endList(); listVector.setValueCount(5); assertEquals(4, listVector.getLastSet()); /* get offset buffer */ final ArrowBuf offsetBuffer = listVector.getOffsetBuffer(); /* get dataVector */ BigIntVector dataVector = (BigIntVector) listVector.getDataVector(); /* check the vector output */ int index = 0; int offset; Long actual; /* index 0 */ assertFalse(listVector.isNull(index)); offset = (int) offsetBuffer.getLong(index * LargeListVector.OFFSET_WIDTH); assertEquals(Integer.toString(0), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(10), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(11), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(12), actual); /* index 1 */ index++; assertFalse(listVector.isNull(index)); offset = (int) offsetBuffer.getLong(index * LargeListVector.OFFSET_WIDTH); assertEquals(Integer.toString(3), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(13), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(14), actual); /* index 2 */ index++; assertFalse(listVector.isNull(index)); offset = (int) offsetBuffer.getLong(index * LargeListVector.OFFSET_WIDTH); assertEquals(Integer.toString(5), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(15), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(16), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(17), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(18), actual); /* index 3 */ index++; assertFalse(listVector.isNull(index)); offset = (int) offsetBuffer.getLong(index * LargeListVector.OFFSET_WIDTH); assertEquals(Integer.toString(9), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(19), actual); /* index 4 */ index++; assertFalse(listVector.isNull(index)); offset = (int) offsetBuffer.getLong(index * LargeListVector.OFFSET_WIDTH); assertEquals(Integer.toString(10), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(20), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(21), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(22), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(23), actual); /* index 5 */ index++; assertTrue(listVector.isNull(index)); offset = (int) offsetBuffer.getLong(index * LargeListVector.OFFSET_WIDTH); assertEquals(Integer.toString(14), Integer.toString(offset)); /* do split and transfer */ try (LargeListVector toVector = LargeListVector.empty("toVector", allocator)) { TransferPair transferPair = listVector.makeTransferPair(toVector); int[][] transferLengths = {{0, 2}, {3, 1}, {4, 1}}; for (final int[] transferLength : transferLengths) { int start = transferLength[0]; int splitLength = transferLength[1]; int dataLength1 = 0; int dataLength2 = 0; int offset1 = 0; int offset2 = 0; transferPair.splitAndTransfer(start, splitLength); /* get offsetBuffer of toVector */ final ArrowBuf toOffsetBuffer = toVector.getOffsetBuffer(); /* get dataVector of toVector */ BigIntVector dataVector1 = (BigIntVector) toVector.getDataVector(); for (int i = 0; i < splitLength; i++) { dataLength1 = (int) offsetBuffer.getLong((start + i + 1) * LargeListVector.OFFSET_WIDTH) - (int) offsetBuffer.getLong((start + i) * LargeListVector.OFFSET_WIDTH); dataLength2 = (int) toOffsetBuffer.getLong((i + 1) * LargeListVector.OFFSET_WIDTH) - (int) toOffsetBuffer.getLong(i * LargeListVector.OFFSET_WIDTH); assertEquals( dataLength1, dataLength2, "Different data lengths at index: " + i + " and start: " + start); offset1 = (int) offsetBuffer.getLong((start + i) * LargeListVector.OFFSET_WIDTH); offset2 = (int) toOffsetBuffer.getLong(i * LargeListVector.OFFSET_WIDTH); for (int j = 0; j < dataLength1; j++) { assertEquals( dataVector.getObject(offset1), dataVector1.getObject(offset2), "Different data at indexes: " + offset1 + " and " + offset2); offset1++; offset2++; } } } } } } @Test public void testNestedLargeListVector() throws Exception { try (LargeListVector listVector = LargeListVector.empty("sourceVector", allocator)) { UnionLargeListWriter listWriter = listVector.getWriter(); /* allocate memory */ listWriter.allocate(); /* the dataVector that backs a listVector will also be a * listVector for this test. */ /* write one or more inner lists at index 0 */ listWriter.setPosition(0); listWriter.startList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(50); listWriter.list().bigInt().writeBigInt(100); listWriter.list().bigInt().writeBigInt(200); listWriter.list().endList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(75); listWriter.list().bigInt().writeBigInt(125); listWriter.list().bigInt().writeBigInt(150); listWriter.list().bigInt().writeBigInt(175); listWriter.list().endList(); listWriter.endList(); /* write one or more inner lists at index 1 */ listWriter.setPosition(1); listWriter.startList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(10); listWriter.list().endList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(15); listWriter.list().bigInt().writeBigInt(20); listWriter.list().endList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(25); listWriter.list().bigInt().writeBigInt(30); listWriter.list().bigInt().writeBigInt(35); listWriter.list().endList(); listWriter.endList(); assertEquals(1, listVector.getLastSet()); listVector.setValueCount(2); assertEquals(2, listVector.getValueCount()); /* get listVector value at index 0 -- the value itself is a listvector */ Object result = listVector.getObject(0); ArrayList> resultSet = (ArrayList>) result; ArrayList list; assertEquals(2, resultSet.size()); /* 2 inner lists at index 0 */ assertEquals(3, resultSet.get(0).size()); /* size of first inner list */ assertEquals(4, resultSet.get(1).size()); /* size of second inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(50), list.get(0)); assertEquals(Long.valueOf(100), list.get(1)); assertEquals(Long.valueOf(200), list.get(2)); list = resultSet.get(1); assertEquals(Long.valueOf(75), list.get(0)); assertEquals(Long.valueOf(125), list.get(1)); assertEquals(Long.valueOf(150), list.get(2)); assertEquals(Long.valueOf(175), list.get(3)); /* get listVector value at index 1 -- the value itself is a listvector */ result = listVector.getObject(1); resultSet = (ArrayList>) result; assertEquals(3, resultSet.size()); /* 3 inner lists at index 1 */ assertEquals(1, resultSet.get(0).size()); /* size of first inner list */ assertEquals(2, resultSet.get(1).size()); /* size of second inner list */ assertEquals(3, resultSet.get(2).size()); /* size of third inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(10), list.get(0)); list = resultSet.get(1); assertEquals(Long.valueOf(15), list.get(0)); assertEquals(Long.valueOf(20), list.get(1)); list = resultSet.get(2); assertEquals(Long.valueOf(25), list.get(0)); assertEquals(Long.valueOf(30), list.get(1)); assertEquals(Long.valueOf(35), list.get(2)); /* check underlying bitVector */ assertFalse(listVector.isNull(0)); assertFalse(listVector.isNull(1)); /* check underlying offsets */ final ArrowBuf offsetBuffer = listVector.getOffsetBuffer(); /* listVector has 2 lists at index 0 and 3 lists at index 1 */ assertEquals(0, offsetBuffer.getLong(0 * LargeListVector.OFFSET_WIDTH)); assertEquals(2, offsetBuffer.getLong(1 * LargeListVector.OFFSET_WIDTH)); assertEquals(5, offsetBuffer.getLong(2 * LargeListVector.OFFSET_WIDTH)); } } @Test public void testNestedLargeListVector1() throws Exception { try (LargeListVector listVector = LargeListVector.empty("sourceVector", allocator)) { MinorType listType = MinorType.LIST; MinorType scalarType = MinorType.BIGINT; listVector.addOrGetVector(FieldType.nullable(listType.getType())); ListVector innerList1 = (ListVector) listVector.getDataVector(); innerList1.addOrGetVector(FieldType.nullable(listType.getType())); ListVector innerList2 = (ListVector) innerList1.getDataVector(); innerList2.addOrGetVector(FieldType.nullable(listType.getType())); ListVector innerList3 = (ListVector) innerList2.getDataVector(); innerList3.addOrGetVector(FieldType.nullable(listType.getType())); ListVector innerList4 = (ListVector) innerList3.getDataVector(); innerList4.addOrGetVector(FieldType.nullable(listType.getType())); ListVector innerList5 = (ListVector) innerList4.getDataVector(); innerList5.addOrGetVector(FieldType.nullable(listType.getType())); ListVector innerList6 = (ListVector) innerList5.getDataVector(); innerList6.addOrGetVector(FieldType.nullable(scalarType.getType())); listVector.setInitialCapacity(128); } } @Test public void testNestedLargeListVector2() throws Exception { try (LargeListVector listVector = LargeListVector.empty("sourceVector", allocator)) { listVector.setInitialCapacity(1); UnionLargeListWriter listWriter = listVector.getWriter(); /* allocate memory */ listWriter.allocate(); /* write one or more inner lists at index 0 */ listWriter.setPosition(0); listWriter.startList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(50); listWriter.list().bigInt().writeBigInt(100); listWriter.list().bigInt().writeBigInt(200); listWriter.list().endList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(75); listWriter.list().bigInt().writeBigInt(125); listWriter.list().endList(); listWriter.endList(); /* write one or more inner lists at index 1 */ listWriter.setPosition(1); listWriter.startList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(15); listWriter.list().bigInt().writeBigInt(20); listWriter.list().endList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(25); listWriter.list().bigInt().writeBigInt(30); listWriter.list().bigInt().writeBigInt(35); listWriter.list().endList(); listWriter.endList(); assertEquals(1, listVector.getLastSet()); listVector.setValueCount(2); assertEquals(2, listVector.getValueCount()); /* get listVector value at index 0 -- the value itself is a listvector */ Object result = listVector.getObject(0); ArrayList> resultSet = (ArrayList>) result; ArrayList list; assertEquals(2, resultSet.size()); /* 2 inner lists at index 0 */ assertEquals(3, resultSet.get(0).size()); /* size of first inner list */ assertEquals(2, resultSet.get(1).size()); /* size of second inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(50), list.get(0)); assertEquals(Long.valueOf(100), list.get(1)); assertEquals(Long.valueOf(200), list.get(2)); list = resultSet.get(1); assertEquals(Long.valueOf(75), list.get(0)); assertEquals(Long.valueOf(125), list.get(1)); /* get listVector value at index 1 -- the value itself is a listvector */ result = listVector.getObject(1); resultSet = (ArrayList>) result; assertEquals(2, resultSet.size()); /* 3 inner lists at index 1 */ assertEquals(2, resultSet.get(0).size()); /* size of first inner list */ assertEquals(3, resultSet.get(1).size()); /* size of second inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(15), list.get(0)); assertEquals(Long.valueOf(20), list.get(1)); list = resultSet.get(1); assertEquals(Long.valueOf(25), list.get(0)); assertEquals(Long.valueOf(30), list.get(1)); assertEquals(Long.valueOf(35), list.get(2)); /* check underlying bitVector */ assertFalse(listVector.isNull(0)); assertFalse(listVector.isNull(1)); /* check underlying offsets */ final ArrowBuf offsetBuffer = listVector.getOffsetBuffer(); /* listVector has 2 lists at index 0 and 3 lists at index 1 */ assertEquals(0, offsetBuffer.getLong(0 * LargeListVector.OFFSET_WIDTH)); assertEquals(2, offsetBuffer.getLong(1 * LargeListVector.OFFSET_WIDTH)); assertEquals(4, offsetBuffer.getLong(2 * LargeListVector.OFFSET_WIDTH)); } } @Test public void testGetBufferAddress() throws Exception { try (LargeListVector listVector = LargeListVector.empty("vector", allocator)) { UnionLargeListWriter listWriter = listVector.getWriter(); boolean error = false; listWriter.allocate(); listWriter.setPosition(0); listWriter.startList(); listWriter.bigInt().writeBigInt(50); listWriter.bigInt().writeBigInt(100); listWriter.bigInt().writeBigInt(200); listWriter.endList(); listWriter.setPosition(1); listWriter.startList(); listWriter.bigInt().writeBigInt(250); listWriter.bigInt().writeBigInt(300); listWriter.endList(); listVector.setValueCount(2); /* check listVector contents */ Object result = listVector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Long.valueOf(50), resultSet.get(0)); assertEquals(Long.valueOf(100), resultSet.get(1)); assertEquals(Long.valueOf(200), resultSet.get(2)); result = listVector.getObject(1); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(250), resultSet.get(0)); assertEquals(Long.valueOf(300), resultSet.get(1)); List buffers = listVector.getFieldBuffers(); long bitAddress = listVector.getValidityBufferAddress(); long offsetAddress = listVector.getOffsetBufferAddress(); try { listVector.getDataBufferAddress(); } catch (UnsupportedOperationException ue) { error = true; } finally { assertTrue(error); } assertEquals(2, buffers.size()); assertEquals(bitAddress, buffers.get(0).memoryAddress()); assertEquals(offsetAddress, buffers.get(1).memoryAddress()); /* (3+2)/2 */ assertEquals(2.5, listVector.getDensity(), 0); } } @Test public void testConsistentChildName() throws Exception { try (LargeListVector listVector = LargeListVector.empty("sourceVector", allocator)) { String emptyListStr = listVector.getField().toString(); assertTrue(emptyListStr.contains(LargeListVector.DATA_VECTOR_NAME)); listVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); String emptyVectorStr = listVector.getField().toString(); assertTrue(emptyVectorStr.contains(LargeListVector.DATA_VECTOR_NAME)); } } @Test public void testSetInitialCapacity() { try (final LargeListVector vector = LargeListVector.empty("", allocator)) { vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); /** * use the default multiplier of 5, 512 * 5 => 2560 * 4 => 10240 bytes => 16KB => 4096 value * capacity. */ vector.setInitialCapacity(512); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 512 * 5); /* use density as 4 */ vector.setInitialCapacity(512, 4); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 512 * 4); /** * inner value capacity we pass to data vector is 512 * 0.1 => 51 For an int vector this is * 204 bytes of memory for data buffer and 7 bytes for validity buffer. and with power of 2 * allocation, we allocate 256 bytes and 8 bytes for the data buffer and validity buffer of * the inner vector. Thus value capacity of inner vector is 64 */ vector.setInitialCapacity(512, 0.1); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 51); /** * inner value capacity we pass to data vector is 512 * 0.01 => 5 For an int vector this is 20 * bytes of memory for data buffer and 1 byte for validity buffer. and with power of 2 * allocation, we allocate 32 bytes and 1 bytes for the data buffer and validity buffer of the * inner vector. Thus value capacity of inner vector is 8 */ vector.setInitialCapacity(512, 0.01); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 5); /** * inner value capacity we pass to data vector is 5 * 0.1 => 0 which is then rounded off to 1. * So we pass value count as 1 to the inner int vector. the offset buffer of the list vector * is allocated for 6 values which is 24 bytes and then rounded off to 32 bytes (8 values) the * validity buffer of the list vector is allocated for 5 values which is 1 byte. This is why * value capacity of the list vector is 7 as we take the min of validity buffer value capacity * and offset buffer value capacity. */ vector.setInitialCapacity(5, 0.1); vector.allocateNew(); assertEquals(7, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 1); } } @Test public void testClearAndReuse() { try (final LargeListVector vector = LargeListVector.empty("list", allocator)) { BigIntVector bigIntVector = (BigIntVector) vector.addOrGetVector(FieldType.nullable(MinorType.BIGINT.getType())).getVector(); vector.setInitialCapacity(10); vector.allocateNew(); vector.startNewValue(0); bigIntVector.setSafe(0, 7); vector.endValue(0, 1); vector.startNewValue(1); bigIntVector.setSafe(1, 8); vector.endValue(1, 1); vector.setValueCount(2); Object result = vector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(Long.valueOf(7), resultSet.get(0)); result = vector.getObject(1); resultSet = (ArrayList) result; assertEquals(Long.valueOf(8), resultSet.get(0)); // Clear and release the buffers to trigger a realloc when adding next value vector.clear(); // The list vector should reuse a buffer when reallocating the offset buffer vector.startNewValue(0); bigIntVector.setSafe(0, 7); vector.endValue(0, 1); vector.startNewValue(1); bigIntVector.setSafe(1, 8); vector.endValue(1, 1); vector.setValueCount(2); result = vector.getObject(0); resultSet = (ArrayList) result; assertEquals(Long.valueOf(7), resultSet.get(0)); result = vector.getObject(1); resultSet = (ArrayList) result; assertEquals(Long.valueOf(8), resultSet.get(0)); } } @Test public void testWriterGetField() { try (final LargeListVector vector = LargeListVector.empty("list", allocator)) { UnionLargeListWriter writer = vector.getWriter(); writer.allocate(); // set some values writer.startList(); writer.integer().writeInt(1); writer.integer().writeInt(2); writer.endList(); vector.setValueCount(2); Field expectedDataField = new Field( BaseRepeatedValueVector.DATA_VECTOR_NAME, FieldType.nullable(new ArrowType.Int(32, true)), null); Field expectedField = new Field( vector.getName(), FieldType.nullable(ArrowType.LargeList.INSTANCE), Arrays.asList(expectedDataField)); assertEquals(expectedField, writer.getField()); } } @Test public void testClose() throws Exception { try (final LargeListVector vector = LargeListVector.empty("list", allocator)) { UnionLargeListWriter writer = vector.getWriter(); writer.allocate(); // set some values writer.startList(); writer.integer().writeInt(1); writer.integer().writeInt(2); writer.endList(); vector.setValueCount(2); assertTrue(vector.getBufferSize() > 0); assertTrue(vector.getDataVector().getBufferSize() > 0); writer.close(); assertEquals(0, vector.getBufferSize()); assertEquals(0, vector.getDataVector().getBufferSize()); } } @Test public void testGetBufferSizeFor() { try (final LargeListVector vector = LargeListVector.empty("list", allocator)) { UnionLargeListWriter writer = vector.getWriter(); writer.allocate(); // set some values writeIntValues(writer, new int[] {1, 2}); writeIntValues(writer, new int[] {3, 4}); writeIntValues(writer, new int[] {5, 6}); writeIntValues(writer, new int[] {7, 8, 9, 10}); writeIntValues(writer, new int[] {11, 12, 13, 14}); writer.setValueCount(5); IntVector dataVector = (IntVector) vector.getDataVector(); int[] indices = new int[] {0, 2, 4, 6, 10, 14}; for (int valueCount = 1; valueCount <= 5; valueCount++) { int validityBufferSize = getValidityBufferSizeFromCount(valueCount); int offsetBufferSize = (valueCount + 1) * LargeListVector.OFFSET_WIDTH; int expectedSize = validityBufferSize + offsetBufferSize + dataVector.getBufferSizeFor(indices[valueCount]); assertEquals(expectedSize, vector.getBufferSizeFor(valueCount)); } } } @Test public void testIsEmpty() { try (final LargeListVector vector = LargeListVector.empty("list", allocator)) { UnionLargeListWriter writer = vector.getWriter(); writer.allocate(); // set values [1,2], null, [], [5,6] writeIntValues(writer, new int[] {1, 2}); writer.setPosition(2); writeIntValues(writer, new int[] {}); writeIntValues(writer, new int[] {5, 6}); writer.setValueCount(4); assertFalse(vector.isEmpty(0)); assertTrue(vector.isNull(1)); assertTrue(vector.isEmpty(1)); assertFalse(vector.isNull(2)); assertTrue(vector.isEmpty(2)); assertFalse(vector.isEmpty(3)); } } @Test public void testTotalCapacity() { final FieldType type = FieldType.nullable(MinorType.INT.getType()); try (final LargeListVector vector = new LargeListVector("list", allocator, type, null)) { // Force the child vector to be allocated based on the type // (this is a bad API: we have to track and repeat the type twice) vector.addOrGetVector(type); // Specify the allocation size but do not actually allocate vector.setInitialTotalCapacity(10, 100); // Finally actually do the allocation vector.allocateNewSafe(); // Note: allocator rounds up and can be greater than the requested allocation. assertTrue(vector.getValueCapacity() >= 10); assertTrue(vector.getDataVector().getValueCapacity() >= 100); } } @Test public void testGetTransferPairWithField() throws Exception { try (final LargeListVector fromVector = LargeListVector.empty("list", allocator)) { UnionLargeListWriter writer = fromVector.getWriter(); writer.allocate(); // set some values writer.startList(); writer.integer().writeInt(1); writer.integer().writeInt(2); writer.endList(); fromVector.setValueCount(2); final TransferPair transferPair = fromVector.getTransferPair(fromVector.getField(), allocator); final LargeListVector toVector = (LargeListVector) transferPair.getTo(); // Field inside a new vector created by reusing a field should be the same in memory as the // original field. assertSame(toVector.getField(), fromVector.getField()); } } @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) { writer.integer().writeInt(v); } writer.endList(); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestLargeListViewVector.java000066400000000000000000002532741515451615200331570ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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; import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.ArrayList; import java.util.Collections; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.BaseLargeRepeatedValueViewVector; import org.apache.arrow.vector.complex.LargeListViewVector; import org.apache.arrow.vector.complex.impl.UnionLargeListViewWriter; 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.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestLargeListViewVector { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testBasicLargeListViewVector() { try (LargeListViewVector largeListViewVector = LargeListViewVector.empty("sourceVector", allocator)) { UnionLargeListViewWriter largeListViewWriter = largeListViewVector.getWriter(); /* allocate memory */ largeListViewWriter.allocate(); /* write the first list at index 0 */ largeListViewWriter.setPosition(0); largeListViewWriter.startListView(); largeListViewWriter.bigInt().writeBigInt(12); largeListViewWriter.bigInt().writeBigInt(-7); largeListViewWriter.bigInt().writeBigInt(25); largeListViewWriter.endListView(); /* the second list at index 1 is null (we are not setting any)*/ /* write the third list at index 2 */ largeListViewWriter.setPosition(2); largeListViewWriter.startListView(); largeListViewWriter.bigInt().writeBigInt(0); largeListViewWriter.bigInt().writeBigInt(-127); largeListViewWriter.bigInt().writeBigInt(127); largeListViewWriter.bigInt().writeBigInt(50); largeListViewWriter.endListView(); /* write the fourth list at index 3 (empty list) */ largeListViewWriter.setPosition(3); largeListViewWriter.startListView(); largeListViewWriter.endListView(); /* write the fifth list at index 4 */ largeListViewWriter.setPosition(4); largeListViewWriter.startListView(); largeListViewWriter.bigInt().writeBigInt(1); largeListViewWriter.bigInt().writeBigInt(2); largeListViewWriter.bigInt().writeBigInt(3); largeListViewWriter.bigInt().writeBigInt(4); largeListViewWriter.endListView(); largeListViewWriter.setValueCount(5); // check value count assertEquals(5, largeListViewVector.getValueCount()); /* get vector at index 0 -- the value is a BigIntVector*/ final ArrowBuf offSetBuffer = largeListViewVector.getOffsetBuffer(); final ArrowBuf sizeBuffer = largeListViewVector.getSizeBuffer(); final FieldVector dataVec = largeListViewVector.getDataVector(); // check offset buffer assertEquals(0, offSetBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offSetBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(7, offSetBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(7, offSetBuffer.getInt(4 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); // check size buffer assertEquals(3, sizeBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(4, sizeBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(4, sizeBuffer.getInt(4 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); // check data vector assertEquals(12, ((BigIntVector) dataVec).get(0)); assertEquals(-7, ((BigIntVector) dataVec).get(1)); assertEquals(25, ((BigIntVector) dataVec).get(2)); assertEquals(0, ((BigIntVector) dataVec).get(3)); assertEquals(-127, ((BigIntVector) dataVec).get(4)); assertEquals(127, ((BigIntVector) dataVec).get(5)); assertEquals(50, ((BigIntVector) dataVec).get(6)); assertEquals(1, ((BigIntVector) dataVec).get(7)); assertEquals(2, ((BigIntVector) dataVec).get(8)); assertEquals(3, ((BigIntVector) dataVec).get(9)); assertEquals(4, ((BigIntVector) dataVec).get(10)); largeListViewVector.validate(); } } @Test public void testImplicitNullVectors() { try (LargeListViewVector largeListViewVector = LargeListViewVector.empty("sourceVector", allocator)) { UnionLargeListViewWriter largeListViewWriter = largeListViewVector.getWriter(); /* allocate memory */ largeListViewWriter.allocate(); final ArrowBuf offSetBuffer = largeListViewVector.getOffsetBuffer(); final ArrowBuf sizeBuffer = largeListViewVector.getSizeBuffer(); /* write the first list at index 0 */ largeListViewWriter.setPosition(0); largeListViewWriter.startListView(); largeListViewWriter.bigInt().writeBigInt(12); largeListViewWriter.bigInt().writeBigInt(-7); largeListViewWriter.bigInt().writeBigInt(25); largeListViewWriter.endListView(); int offSet0 = offSetBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH); int size0 = sizeBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH); // after the first list is written, // the initial offset must be 0, // the size must be 3 (as there are 3 elements in the array), // the lastSet must be 0 since, the first list is written at index 0. assertEquals(0, offSet0); assertEquals(3, size0); largeListViewWriter.setPosition(5); largeListViewWriter.startListView(); // writing the 6th list at index 5, // and the list items from index 1 through 4 are not populated. // but since there is a gap between the 0th and 5th list, in terms // of buffer allocation, the offset and size buffers must be updated // to reflect the implicit null vectors. for (int i = 1; i < 5; i++) { int offSet = offSetBuffer.getInt(i * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH); int size = sizeBuffer.getInt(i * BaseLargeRepeatedValueViewVector.SIZE_WIDTH); // Since the list is not written, the offset and size must equal to child vector's size // i.e., 3, and size should be 0 as the list is not written. // And the last set value is the value currently being written, which is 5. assertEquals(0, offSet); assertEquals(0, size); } largeListViewWriter.bigInt().writeBigInt(12); largeListViewWriter.bigInt().writeBigInt(25); largeListViewWriter.endListView(); int offSet5 = offSetBuffer.getInt(5 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH); int size5 = sizeBuffer.getInt(5 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH); assertEquals(3, offSet5); assertEquals(2, size5); largeListViewWriter.setPosition(10); largeListViewWriter.startListView(); // writing the 11th list at index 10, // and the list items from index 6 through 10 are not populated. // but since there is a gap between the 5th and 11th list, in terms // of buffer allocation, the offset and size buffers must be updated // to reflect the implicit null vectors. for (int i = 6; i < 10; i++) { int offSet = offSetBuffer.getInt(i * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH); int size = sizeBuffer.getInt(i * BaseLargeRepeatedValueViewVector.SIZE_WIDTH); // Since the list is not written, the offset and size must equal to 0 // and size should be 0 as the list is not written. // And the last set value is the value currently being written, which is 10. assertEquals(0, offSet); assertEquals(0, size); } largeListViewWriter.bigInt().writeBigInt(12); largeListViewWriter.endListView(); int offSet11 = offSetBuffer.getInt(10 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH); int size11 = sizeBuffer.getInt(10 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH); assertEquals(5, offSet11); assertEquals(1, size11); largeListViewVector.setValueCount(11); largeListViewVector.validate(); } } @Test public void testNestedLargeListViewVector() throws Exception { try (LargeListViewVector largeListViewVector = LargeListViewVector.empty("sourceVector", allocator)) { UnionLargeListViewWriter largeListViewWriter = largeListViewVector.getWriter(); /* allocate memory */ largeListViewWriter.allocate(); /* the dataVector that backs a largeListViewVector will also be a * largeListViewVector for this test. */ /* write one or more inner lists at index 0 */ largeListViewWriter.setPosition(0); largeListViewWriter.startListView(); largeListViewWriter.listView().startListView(); largeListViewWriter.listView().bigInt().writeBigInt(50); largeListViewWriter.listView().bigInt().writeBigInt(100); largeListViewWriter.listView().bigInt().writeBigInt(200); largeListViewWriter.listView().endListView(); largeListViewWriter.listView().startListView(); largeListViewWriter.listView().bigInt().writeBigInt(75); largeListViewWriter.listView().bigInt().writeBigInt(125); largeListViewWriter.listView().bigInt().writeBigInt(150); largeListViewWriter.listView().bigInt().writeBigInt(175); largeListViewWriter.listView().endListView(); largeListViewWriter.endListView(); /* write one or more inner lists at index 1 */ largeListViewWriter.setPosition(1); largeListViewWriter.startListView(); largeListViewWriter.listView().startListView(); largeListViewWriter.listView().bigInt().writeBigInt(10); largeListViewWriter.listView().endListView(); largeListViewWriter.listView().startListView(); largeListViewWriter.listView().bigInt().writeBigInt(15); largeListViewWriter.listView().bigInt().writeBigInt(20); largeListViewWriter.listView().endListView(); largeListViewWriter.listView().startListView(); largeListViewWriter.listView().bigInt().writeBigInt(25); largeListViewWriter.listView().bigInt().writeBigInt(30); largeListViewWriter.listView().bigInt().writeBigInt(35); largeListViewWriter.listView().endListView(); largeListViewWriter.endListView(); largeListViewVector.setValueCount(2); assertEquals(2, largeListViewVector.getValueCount()); /* get largeListViewVector value at index 0 -- the value itself is a largeListViewVector */ Object result = largeListViewVector.getObject(0); ArrayList> resultSet = (ArrayList>) result; ArrayList list; assertEquals(2, resultSet.size()); /* 2 inner lists at index 0 */ assertEquals(3, resultSet.get(0).size()); /* size of first inner list */ assertEquals(4, resultSet.get(1).size()); /* size of second inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(50), list.get(0)); assertEquals(Long.valueOf(100), list.get(1)); assertEquals(Long.valueOf(200), list.get(2)); list = resultSet.get(1); assertEquals(Long.valueOf(75), list.get(0)); assertEquals(Long.valueOf(125), list.get(1)); assertEquals(Long.valueOf(150), list.get(2)); assertEquals(Long.valueOf(175), list.get(3)); /* get largeListViewVector value at index 1 -- the value itself is a largeListViewVector */ result = largeListViewVector.getObject(1); resultSet = (ArrayList>) result; assertEquals(3, resultSet.size()); /* 3 inner lists at index 1 */ assertEquals(1, resultSet.get(0).size()); /* size of first inner list */ assertEquals(2, resultSet.get(1).size()); /* size of second inner list */ assertEquals(3, resultSet.get(2).size()); /* size of third inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(10), list.get(0)); list = resultSet.get(1); assertEquals(Long.valueOf(15), list.get(0)); assertEquals(Long.valueOf(20), list.get(1)); list = resultSet.get(2); assertEquals(Long.valueOf(25), list.get(0)); assertEquals(Long.valueOf(30), list.get(1)); assertEquals(Long.valueOf(35), list.get(2)); /* check underlying bitVector */ assertFalse(largeListViewVector.isNull(0)); assertFalse(largeListViewVector.isNull(1)); /* check underlying offsets */ final ArrowBuf offsetBuffer = largeListViewVector.getOffsetBuffer(); /* largeListViewVector has 2 lists at index 0 and 3 lists at index 1 */ assertEquals(0, offsetBuffer.getLong(0 * LargeListViewVector.OFFSET_WIDTH)); assertEquals(2, offsetBuffer.getLong(1 * LargeListViewVector.OFFSET_WIDTH)); } } @Test public void testNestedLargeListViewVector1() { try (LargeListViewVector largeListViewVector = LargeListViewVector.empty("sourceVector", allocator)) { MinorType listViewType = MinorType.LARGELISTVIEW; MinorType scalarType = MinorType.BIGINT; largeListViewVector.addOrGetVector(FieldType.nullable(listViewType.getType())); LargeListViewVector innerList1 = (LargeListViewVector) largeListViewVector.getDataVector(); innerList1.addOrGetVector(FieldType.nullable(listViewType.getType())); LargeListViewVector innerList2 = (LargeListViewVector) innerList1.getDataVector(); innerList2.addOrGetVector(FieldType.nullable(listViewType.getType())); LargeListViewVector innerList3 = (LargeListViewVector) innerList2.getDataVector(); innerList3.addOrGetVector(FieldType.nullable(listViewType.getType())); LargeListViewVector innerList4 = (LargeListViewVector) innerList3.getDataVector(); innerList4.addOrGetVector(FieldType.nullable(listViewType.getType())); LargeListViewVector innerList5 = (LargeListViewVector) innerList4.getDataVector(); innerList5.addOrGetVector(FieldType.nullable(listViewType.getType())); LargeListViewVector innerList6 = (LargeListViewVector) innerList5.getDataVector(); innerList6.addOrGetVector(FieldType.nullable(scalarType.getType())); largeListViewVector.setInitialCapacity(128); } } @Test public void testNestedLargeListViewVector2() throws Exception { try (LargeListViewVector largeListViewVector = LargeListViewVector.empty("sourceVector", allocator)) { largeListViewVector.setInitialCapacity(1); UnionLargeListViewWriter largeListViewWriter = largeListViewVector.getWriter(); /* allocate memory */ largeListViewWriter.allocate(); /* write one or more inner lists at index 0 */ largeListViewWriter.setPosition(0); largeListViewWriter.startListView(); largeListViewWriter.listView().startListView(); largeListViewWriter.listView().bigInt().writeBigInt(50); largeListViewWriter.listView().bigInt().writeBigInt(100); largeListViewWriter.listView().bigInt().writeBigInt(200); largeListViewWriter.listView().endListView(); largeListViewWriter.listView().startListView(); largeListViewWriter.listView().bigInt().writeBigInt(75); largeListViewWriter.listView().bigInt().writeBigInt(125); largeListViewWriter.listView().endListView(); largeListViewWriter.endListView(); /* write one or more inner lists at index 1 */ largeListViewWriter.setPosition(1); largeListViewWriter.startListView(); largeListViewWriter.listView().startListView(); largeListViewWriter.listView().bigInt().writeBigInt(15); largeListViewWriter.listView().bigInt().writeBigInt(20); largeListViewWriter.listView().endListView(); largeListViewWriter.listView().startListView(); largeListViewWriter.listView().bigInt().writeBigInt(25); largeListViewWriter.listView().bigInt().writeBigInt(30); largeListViewWriter.listView().bigInt().writeBigInt(35); largeListViewWriter.listView().endListView(); largeListViewWriter.endListView(); largeListViewVector.setValueCount(2); assertEquals(2, largeListViewVector.getValueCount()); /* get largeListViewVector value at index 0 -- the value itself is a largeListViewVector */ Object result = largeListViewVector.getObject(0); ArrayList> resultSet = (ArrayList>) result; ArrayList list; assertEquals(2, resultSet.size()); /* 2 inner lists at index 0 */ assertEquals(3, resultSet.get(0).size()); /* size of first inner list */ assertEquals(2, resultSet.get(1).size()); /* size of second inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(50), list.get(0)); assertEquals(Long.valueOf(100), list.get(1)); assertEquals(Long.valueOf(200), list.get(2)); list = resultSet.get(1); assertEquals(Long.valueOf(75), list.get(0)); assertEquals(Long.valueOf(125), list.get(1)); /* get largeListViewVector value at index 1 -- the value itself is a largeListViewVector */ result = largeListViewVector.getObject(1); resultSet = (ArrayList>) result; assertEquals(2, resultSet.size()); /* 3 inner lists at index 1 */ assertEquals(2, resultSet.get(0).size()); /* size of first inner list */ assertEquals(3, resultSet.get(1).size()); /* size of second inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(15), list.get(0)); assertEquals(Long.valueOf(20), list.get(1)); list = resultSet.get(1); assertEquals(Long.valueOf(25), list.get(0)); assertEquals(Long.valueOf(30), list.get(1)); assertEquals(Long.valueOf(35), list.get(2)); /* check underlying bitVector */ assertFalse(largeListViewVector.isNull(0)); assertFalse(largeListViewVector.isNull(1)); /* check underlying offsets */ final ArrowBuf offsetBuffer = largeListViewVector.getOffsetBuffer(); /* largeListViewVector has 2 lists at index 0 and 3 lists at index 1 */ assertEquals(0, offsetBuffer.getLong(0 * LargeListViewVector.OFFSET_WIDTH)); assertEquals(2, offsetBuffer.getLong(1 * LargeListViewVector.OFFSET_WIDTH)); } } @Test public void testGetBufferAddress() throws Exception { try (LargeListViewVector largeListViewVector = LargeListViewVector.empty("vector", allocator)) { UnionLargeListViewWriter largeListViewWriter = largeListViewVector.getWriter(); boolean error = false; largeListViewWriter.allocate(); largeListViewWriter.setPosition(0); largeListViewWriter.startListView(); largeListViewWriter.bigInt().writeBigInt(50); largeListViewWriter.bigInt().writeBigInt(100); largeListViewWriter.bigInt().writeBigInt(200); largeListViewWriter.endListView(); largeListViewWriter.setPosition(1); largeListViewWriter.startListView(); largeListViewWriter.bigInt().writeBigInt(250); largeListViewWriter.bigInt().writeBigInt(300); largeListViewWriter.endListView(); largeListViewVector.setValueCount(2); /* check largeListViewVector contents */ Object result = largeListViewVector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Long.valueOf(50), resultSet.get(0)); assertEquals(Long.valueOf(100), resultSet.get(1)); assertEquals(Long.valueOf(200), resultSet.get(2)); result = largeListViewVector.getObject(1); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(250), resultSet.get(0)); assertEquals(Long.valueOf(300), resultSet.get(1)); List buffers = largeListViewVector.getFieldBuffers(); long bitAddress = largeListViewVector.getValidityBufferAddress(); long offsetAddress = largeListViewVector.getOffsetBufferAddress(); long sizeAddress = largeListViewVector.getSizeBufferAddress(); try { largeListViewVector.getDataBufferAddress(); } catch (UnsupportedOperationException ue) { error = true; } finally { assertTrue(error); } assertEquals(3, buffers.size()); assertEquals(bitAddress, buffers.get(0).memoryAddress()); assertEquals(offsetAddress, buffers.get(1).memoryAddress()); assertEquals(sizeAddress, buffers.get(2).memoryAddress()); /* (3+2)/2 */ assertEquals(2.5, largeListViewVector.getDensity(), 0); } } /* * Setting up the buffers directly needs to be validated with the base method used in * the LargeListViewVector class where we use the approach of startListView(), * write to the child vector and endListView(). *

    * To support this, we have to consider the following scenarios; *

    * 1. Only using directly buffer-based inserts. * 2. Default list insertion followed by buffer-based inserts. * 3. Buffer-based inserts followed by default list insertion. */ /* Setting up buffers directly would require the following steps to be taken * 0. Allocate buffers in largeListViewVector by calling `allocateNew` method. * 1. Initialize the child vector using `initializeChildrenFromFields` method. * 2. Set values in the child vector. * 3. Set validity, offset and size buffers using `setValidity`, * `setOffset` and `setSize` methods. * 4. Set value count using `setValueCount` method. */ @Test public void testBasicLargeListViewSet() { try (LargeListViewVector largeListViewVector = LargeListViewVector.empty("sourceVector", allocator)) { // Allocate buffers in largeListViewVector by calling `allocateNew` method. largeListViewVector.allocateNew(); // Initialize the child vector using `initializeChildrenFromFields` method. FieldType fieldType = new FieldType(true, new ArrowType.Int(64, true), null, null); Field field = new Field("child-vector", fieldType, null); largeListViewVector.initializeChildrenFromFields(Collections.singletonList(field)); // Set values in the child vector. FieldVector fieldVector = largeListViewVector.getDataVector(); fieldVector.clear(); BigIntVector childVector = (BigIntVector) fieldVector; childVector.allocateNew(7); childVector.set(0, 12); childVector.set(1, -7); childVector.set(2, 25); childVector.set(3, 0); childVector.set(4, -127); childVector.set(5, 127); childVector.set(6, 50); childVector.setValueCount(7); // Set validity, offset and size buffers using `setValidity`, // `setOffset` and `setSize` methods. largeListViewVector.setOffset(0, 0); largeListViewVector.setOffset(1, 3); largeListViewVector.setOffset(2, 3); largeListViewVector.setOffset(3, 7); largeListViewVector.setSize(0, 3); largeListViewVector.setSize(1, 0); largeListViewVector.setSize(2, 4); largeListViewVector.setSize(3, 0); largeListViewVector.setValidity(0, 1); largeListViewVector.setValidity(1, 0); largeListViewVector.setValidity(2, 1); largeListViewVector.setValidity(3, 1); // Set value count using `setValueCount` method. largeListViewVector.setValueCount(4); final ArrowBuf offSetBuffer = largeListViewVector.getOffsetBuffer(); final ArrowBuf sizeBuffer = largeListViewVector.getSizeBuffer(); // check offset buffer assertEquals(0, offSetBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(7, offSetBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); // check size buffer assertEquals(3, sizeBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(4, sizeBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); // check values assertEquals(12, ((BigIntVector) largeListViewVector.getDataVector()).get(0)); assertEquals(-7, ((BigIntVector) largeListViewVector.getDataVector()).get(1)); assertEquals(25, ((BigIntVector) largeListViewVector.getDataVector()).get(2)); assertEquals(0, ((BigIntVector) largeListViewVector.getDataVector()).get(3)); assertEquals(-127, ((BigIntVector) largeListViewVector.getDataVector()).get(4)); assertEquals(127, ((BigIntVector) largeListViewVector.getDataVector()).get(5)); assertEquals(50, ((BigIntVector) largeListViewVector.getDataVector()).get(6)); largeListViewVector.validate(); } } @Test public void testBasicLargeListViewSetNested() { // Expected largeListViewVector // [[[50,100,200],[75,125,150,175]],[[10],[15,20],[25,30,35]]] // Setting child vector // [[50,100,200],[75,125,150,175],[10],[15,20],[25,30,35]] try (LargeListViewVector largeListViewVector = LargeListViewVector.empty("sourceVector", allocator)) { // Allocate buffers in largeListViewVector by calling `allocateNew` method. largeListViewVector.allocateNew(); // Initialize the child vector using `initializeChildrenFromFields` method. FieldType fieldType = new FieldType(true, new ArrowType.LargeListView(), null, null); FieldType childFieldType = new FieldType(true, new ArrowType.Int(64, true), null, null); Field childField = new Field("child-vector", childFieldType, null); List children = new ArrayList<>(); children.add(childField); Field field = new Field("child-vector", fieldType, children); largeListViewVector.initializeChildrenFromFields(Collections.singletonList(field)); // Set values in the child vector. FieldVector fieldVector = largeListViewVector.getDataVector(); fieldVector.clear(); LargeListViewVector childVector = (LargeListViewVector) fieldVector; UnionLargeListViewWriter largeListViewWriter = childVector.getWriter(); largeListViewWriter.allocate(); largeListViewWriter.setPosition(0); largeListViewWriter.startListView(); largeListViewWriter.bigInt().writeBigInt(50); largeListViewWriter.bigInt().writeBigInt(100); largeListViewWriter.bigInt().writeBigInt(200); largeListViewWriter.endListView(); largeListViewWriter.setPosition(1); largeListViewWriter.startListView(); largeListViewWriter.bigInt().writeBigInt(75); largeListViewWriter.bigInt().writeBigInt(125); largeListViewWriter.bigInt().writeBigInt(150); largeListViewWriter.bigInt().writeBigInt(175); largeListViewWriter.endListView(); largeListViewWriter.setPosition(2); largeListViewWriter.startListView(); largeListViewWriter.bigInt().writeBigInt(10); largeListViewWriter.endListView(); largeListViewWriter.startListView(); largeListViewWriter.setPosition(3); largeListViewWriter.bigInt().writeBigInt(15); largeListViewWriter.bigInt().writeBigInt(20); largeListViewWriter.endListView(); largeListViewWriter.startListView(); largeListViewWriter.setPosition(4); largeListViewWriter.bigInt().writeBigInt(25); largeListViewWriter.bigInt().writeBigInt(30); largeListViewWriter.bigInt().writeBigInt(35); largeListViewWriter.endListView(); childVector.setValueCount(5); // Set validity, offset and size buffers using `setValidity`, // `setOffset` and `setSize` methods. largeListViewVector.setValidity(0, 1); largeListViewVector.setValidity(1, 1); largeListViewVector.setOffset(0, 0); largeListViewVector.setOffset(1, 2); largeListViewVector.setSize(0, 2); largeListViewVector.setSize(1, 3); // Set value count using `setValueCount` method. largeListViewVector.setValueCount(2); assertEquals(2, largeListViewVector.getValueCount()); /* get largeListViewVector value at index 0 -- the value itself is a largeListViewVector */ Object result = largeListViewVector.getObject(0); ArrayList> resultSet = (ArrayList>) result; ArrayList list; assertEquals(2, resultSet.size()); /* 2 inner lists at index 0 */ assertEquals(3, resultSet.get(0).size()); /* size of the first inner list */ assertEquals(4, resultSet.get(1).size()); /* size of the second inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(50), list.get(0)); assertEquals(Long.valueOf(100), list.get(1)); assertEquals(Long.valueOf(200), list.get(2)); list = resultSet.get(1); assertEquals(Long.valueOf(75), list.get(0)); assertEquals(Long.valueOf(125), list.get(1)); assertEquals(Long.valueOf(150), list.get(2)); assertEquals(Long.valueOf(175), list.get(3)); /* get largeListViewVector value at index 1 -- the value itself is a largeListViewVector */ result = largeListViewVector.getObject(1); resultSet = (ArrayList>) result; assertEquals(3, resultSet.size()); /* 3 inner lists at index 1 */ assertEquals(1, resultSet.get(0).size()); /* size of the first inner list */ assertEquals(2, resultSet.get(1).size()); /* size of the second inner list */ assertEquals(3, resultSet.get(2).size()); /* size of the third inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(10), list.get(0)); list = resultSet.get(1); assertEquals(Long.valueOf(15), list.get(0)); assertEquals(Long.valueOf(20), list.get(1)); list = resultSet.get(2); assertEquals(Long.valueOf(25), list.get(0)); assertEquals(Long.valueOf(30), list.get(1)); assertEquals(Long.valueOf(35), list.get(2)); /* check underlying bitVector */ assertFalse(largeListViewVector.isNull(0)); assertFalse(largeListViewVector.isNull(1)); final ArrowBuf offSetBuffer = largeListViewVector.getOffsetBuffer(); final ArrowBuf sizeBuffer = largeListViewVector.getSizeBuffer(); // check offset buffer assertEquals(0, offSetBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(2, offSetBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); // check size buffer assertEquals(2, sizeBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(3, sizeBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); largeListViewVector.validate(); } } @Test public void testBasicLargeListViewSetWithListViewWriter() { try (LargeListViewVector largeListViewVector = LargeListViewVector.empty("sourceVector", allocator)) { // Allocate buffers in largeListViewVector by calling `allocateNew` method. largeListViewVector.allocateNew(); // Initialize the child vector using `initializeChildrenFromFields` method. FieldType fieldType = new FieldType(true, new ArrowType.Int(64, true), null, null); Field field = new Field("child-vector", fieldType, null); largeListViewVector.initializeChildrenFromFields(Collections.singletonList(field)); // Set values in the child vector. FieldVector fieldVector = largeListViewVector.getDataVector(); fieldVector.clear(); BigIntVector childVector = (BigIntVector) fieldVector; childVector.allocateNew(7); childVector.set(0, 12); childVector.set(1, -7); childVector.set(2, 25); childVector.set(3, 0); childVector.set(4, -127); childVector.set(5, 127); childVector.set(6, 50); childVector.setValueCount(7); // Set validity, offset and size buffers using `setValidity`, // `setOffset` and `setSize` methods. largeListViewVector.setValidity(0, 1); largeListViewVector.setValidity(1, 0); largeListViewVector.setValidity(2, 1); largeListViewVector.setValidity(3, 1); largeListViewVector.setOffset(0, 0); largeListViewVector.setOffset(1, 3); largeListViewVector.setOffset(2, 3); largeListViewVector.setOffset(3, 7); largeListViewVector.setSize(0, 3); largeListViewVector.setSize(1, 0); largeListViewVector.setSize(2, 4); largeListViewVector.setSize(3, 0); // Set value count using `setValueCount` method. largeListViewVector.setValueCount(4); final ArrowBuf offSetBuffer = largeListViewVector.getOffsetBuffer(); final ArrowBuf sizeBuffer = largeListViewVector.getSizeBuffer(); // check offset buffer assertEquals(0, offSetBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(7, offSetBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); // check size buffer assertEquals(3, sizeBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(4, sizeBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); // check values assertEquals(12, ((BigIntVector) largeListViewVector.getDataVector()).get(0)); assertEquals(-7, ((BigIntVector) largeListViewVector.getDataVector()).get(1)); assertEquals(25, ((BigIntVector) largeListViewVector.getDataVector()).get(2)); assertEquals(0, ((BigIntVector) largeListViewVector.getDataVector()).get(3)); assertEquals(-127, ((BigIntVector) largeListViewVector.getDataVector()).get(4)); assertEquals(127, ((BigIntVector) largeListViewVector.getDataVector()).get(5)); assertEquals(50, ((BigIntVector) largeListViewVector.getDataVector()).get(6)); UnionLargeListViewWriter listViewWriter = largeListViewVector.getWriter(); listViewWriter.setPosition(4); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(121); listViewWriter.bigInt().writeBigInt(-71); listViewWriter.bigInt().writeBigInt(251); listViewWriter.endListView(); largeListViewVector.setValueCount(5); // check offset buffer assertEquals(0, offSetBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(7, offSetBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(7, offSetBuffer.getInt(4 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); // check size buffer assertEquals(3, sizeBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(4, sizeBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(3, sizeBuffer.getInt(4 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); // check values assertEquals(12, ((BigIntVector) largeListViewVector.getDataVector()).get(0)); assertEquals(-7, ((BigIntVector) largeListViewVector.getDataVector()).get(1)); assertEquals(25, ((BigIntVector) largeListViewVector.getDataVector()).get(2)); assertEquals(0, ((BigIntVector) largeListViewVector.getDataVector()).get(3)); assertEquals(-127, ((BigIntVector) largeListViewVector.getDataVector()).get(4)); assertEquals(127, ((BigIntVector) largeListViewVector.getDataVector()).get(5)); assertEquals(50, ((BigIntVector) largeListViewVector.getDataVector()).get(6)); assertEquals(121, ((BigIntVector) largeListViewVector.getDataVector()).get(7)); assertEquals(-71, ((BigIntVector) largeListViewVector.getDataVector()).get(8)); assertEquals(251, ((BigIntVector) largeListViewVector.getDataVector()).get(9)); largeListViewVector.validate(); } } @Test public void testConsistentChildName() throws Exception { try (LargeListViewVector largeListViewVector = LargeListViewVector.empty("sourceVector", allocator)) { String emptyListStr = largeListViewVector.getField().toString(); assertTrue(emptyListStr.contains(LargeListViewVector.DATA_VECTOR_NAME)); largeListViewVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); String emptyVectorStr = largeListViewVector.getField().toString(); assertTrue(emptyVectorStr.contains(LargeListViewVector.DATA_VECTOR_NAME)); } } @Test public void testSetInitialCapacity() { try (final LargeListViewVector vector = LargeListViewVector.empty("", allocator)) { vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); vector.setInitialCapacity(512); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 512); vector.setInitialCapacity(512, 4); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 512 * 4); vector.setInitialCapacity(512, 0.1); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 51); vector.setInitialCapacity(512, 0.01); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 5); vector.setInitialCapacity(5, 0.1); vector.allocateNew(); assertEquals(8, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 1); vector.validate(); } } @Test public void testClearAndReuse() { try (final LargeListViewVector vector = LargeListViewVector.empty("list", allocator)) { BigIntVector bigIntVector = (BigIntVector) vector.addOrGetVector(FieldType.nullable(MinorType.BIGINT.getType())).getVector(); vector.setInitialCapacity(10); vector.allocateNew(); vector.startNewValue(0); bigIntVector.setSafe(0, 7); vector.endValue(0, 1); vector.startNewValue(1); bigIntVector.setSafe(1, 8); vector.endValue(1, 1); vector.setValueCount(2); Object result = vector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(Long.valueOf(7), resultSet.get(0)); result = vector.getObject(1); resultSet = (ArrayList) result; assertEquals(Long.valueOf(8), resultSet.get(0)); // Clear and release the buffers to trigger a realloc when adding next value vector.clear(); // The list vector should reuse a buffer when reallocating the offset buffer vector.startNewValue(0); bigIntVector.setSafe(0, 7); vector.endValue(0, 1); vector.startNewValue(1); bigIntVector.setSafe(1, 8); vector.endValue(1, 1); vector.setValueCount(2); result = vector.getObject(0); resultSet = (ArrayList) result; assertEquals(Long.valueOf(7), resultSet.get(0)); result = vector.getObject(1); resultSet = (ArrayList) result; assertEquals(Long.valueOf(8), resultSet.get(0)); } } @Test public void testWriterGetField() { try (final LargeListViewVector vector = LargeListViewVector.empty("list", allocator)) { UnionLargeListViewWriter writer = vector.getWriter(); writer.allocate(); // set some values writer.startListView(); writer.integer().writeInt(1); writer.integer().writeInt(2); writer.endListView(); vector.setValueCount(2); Field expectedDataField = new Field( BaseLargeRepeatedValueViewVector.DATA_VECTOR_NAME, FieldType.nullable(new ArrowType.Int(32, true)), null); Field expectedField = new Field( vector.getName(), FieldType.nullable(ArrowType.LargeListView.INSTANCE), Collections.singletonList(expectedDataField)); assertEquals(expectedField, writer.getField()); } } @Test public void testClose() throws Exception { try (final LargeListViewVector vector = LargeListViewVector.empty("largelistview", allocator)) { UnionLargeListViewWriter writer = vector.getWriter(); writer.allocate(); // set some values writer.startListView(); writer.integer().writeInt(1); writer.integer().writeInt(2); writer.endListView(); vector.setValueCount(2); assertTrue(vector.getBufferSize() > 0); assertTrue(vector.getDataVector().getBufferSize() > 0); writer.close(); assertEquals(0, vector.getBufferSize()); assertEquals(0, vector.getDataVector().getBufferSize()); vector.validate(); } } @Test public void testGetBufferSizeFor() { try (final LargeListViewVector vector = LargeListViewVector.empty("largelistview", allocator)) { UnionLargeListViewWriter writer = vector.getWriter(); writer.allocate(); // set some values writeIntValues(writer, new int[] {1, 2}); writeIntValues(writer, new int[] {3, 4}); writeIntValues(writer, new int[] {5, 6}); writeIntValues(writer, new int[] {7, 8, 9, 10}); writeIntValues(writer, new int[] {11, 12, 13, 14}); writer.setValueCount(5); IntVector dataVector = (IntVector) vector.getDataVector(); int[] indices = new int[] {0, 2, 4, 6, 10, 14}; for (int valueCount = 1; valueCount <= 5; valueCount++) { int validityBufferSize = getValidityBufferSizeFromCount(valueCount); int offsetBufferSize = valueCount * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH; int sizeBufferSize = valueCount * BaseLargeRepeatedValueViewVector.SIZE_WIDTH; int expectedSize = validityBufferSize + offsetBufferSize + sizeBufferSize + dataVector.getBufferSizeFor(indices[valueCount]); assertEquals(expectedSize, vector.getBufferSizeFor(valueCount)); } vector.validate(); } } @Test public void testIsEmpty() { try (final LargeListViewVector vector = LargeListViewVector.empty("largelistview", allocator)) { UnionLargeListViewWriter writer = vector.getWriter(); writer.allocate(); // set values [1,2], null, [], [5,6] writeIntValues(writer, new int[] {1, 2}); writer.setPosition(2); writeIntValues(writer, new int[] {}); writeIntValues(writer, new int[] {5, 6}); writer.setValueCount(4); assertFalse(vector.isEmpty(0)); assertTrue(vector.isNull(1)); assertTrue(vector.isEmpty(1)); assertFalse(vector.isNull(2)); assertTrue(vector.isEmpty(2)); assertFalse(vector.isEmpty(3)); vector.validate(); } } @Test public void testTotalCapacity() { final FieldType type = FieldType.nullable(MinorType.INT.getType()); try (final LargeListViewVector vector = new LargeListViewVector("largelistview", allocator, type, null)) { // Force the child vector to be allocated based on the type // (this is a bad API: we have to track and repeat the type twice) vector.addOrGetVector(type); // Specify the allocation size but do not actually allocate vector.setInitialTotalCapacity(10, 100); // Finally, actually do the allocation vector.allocateNewSafe(); // Note: allocator rounds up and can be greater than the requested allocation. assertTrue(vector.getValueCapacity() >= 10); assertTrue(vector.getDataVector().getValueCapacity() >= 100); } } @Test public void testSetNull1() { try (LargeListViewVector vector = LargeListViewVector.empty("largelistview", allocator)) { UnionLargeListViewWriter writer = vector.getWriter(); writer.allocate(); writer.setPosition(0); writer.startListView(); writer.bigInt().writeBigInt(10); writer.bigInt().writeBigInt(20); writer.endListView(); vector.setNull(1); writer.setPosition(2); writer.startListView(); writer.bigInt().writeBigInt(30); writer.bigInt().writeBigInt(40); writer.endListView(); vector.setNull(3); vector.setNull(4); writer.setPosition(5); writer.startListView(); writer.bigInt().writeBigInt(50); writer.bigInt().writeBigInt(60); writer.endListView(); vector.setValueCount(6); assertFalse(vector.isNull(0)); assertTrue(vector.isNull(1)); assertFalse(vector.isNull(2)); assertTrue(vector.isNull(3)); assertTrue(vector.isNull(4)); assertFalse(vector.isNull(5)); // validate buffers final ArrowBuf validityBuffer = vector.getValidityBuffer(); final ArrowBuf offsetBuffer = vector.getOffsetBuffer(); final ArrowBuf sizeBuffer = vector.getSizeBuffer(); assertEquals(1, BitVectorHelper.get(validityBuffer, 0)); assertEquals(0, BitVectorHelper.get(validityBuffer, 1)); assertEquals(1, BitVectorHelper.get(validityBuffer, 2)); assertEquals(0, BitVectorHelper.get(validityBuffer, 3)); assertEquals(0, BitVectorHelper.get(validityBuffer, 4)); assertEquals(1, BitVectorHelper.get(validityBuffer, 5)); assertEquals(0, offsetBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(2, offsetBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(4 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(4, offsetBuffer.getInt(5 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(2, sizeBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(2, sizeBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(4 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(2, sizeBuffer.getInt(5 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); // validate values Object result = vector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(10), resultSet.get(0)); assertEquals(Long.valueOf(20), resultSet.get(1)); result = vector.getObject(2); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(30), resultSet.get(0)); assertEquals(Long.valueOf(40), resultSet.get(1)); result = vector.getObject(5); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(50), resultSet.get(0)); assertEquals(Long.valueOf(60), resultSet.get(1)); vector.validate(); } } @Test public void testSetNull2() { try (LargeListViewVector vector = LargeListViewVector.empty("largelistview", allocator)) { // validate setting nulls first and then writing values UnionLargeListViewWriter writer = vector.getWriter(); writer.allocate(); vector.setNull(0); vector.setNull(2); vector.setNull(4); writer.setPosition(1); writer.startListView(); writer.bigInt().writeBigInt(10); writer.bigInt().writeBigInt(20); writer.bigInt().writeBigInt(30); writer.endListView(); writer.setPosition(3); writer.startListView(); writer.bigInt().writeBigInt(40); writer.bigInt().writeBigInt(50); writer.endListView(); writer.setPosition(5); writer.startListView(); writer.bigInt().writeBigInt(60); writer.bigInt().writeBigInt(70); writer.bigInt().writeBigInt(80); writer.endListView(); vector.setValueCount(6); assertTrue(vector.isNull(0)); assertFalse(vector.isNull(1)); assertTrue(vector.isNull(2)); assertFalse(vector.isNull(3)); assertTrue(vector.isNull(4)); assertFalse(vector.isNull(5)); // validate buffers final ArrowBuf validityBuffer = vector.getValidityBuffer(); final ArrowBuf offsetBuffer = vector.getOffsetBuffer(); final ArrowBuf sizeBuffer = vector.getSizeBuffer(); assertEquals(0, BitVectorHelper.get(validityBuffer, 0)); assertEquals(1, BitVectorHelper.get(validityBuffer, 1)); assertEquals(0, BitVectorHelper.get(validityBuffer, 2)); assertEquals(1, BitVectorHelper.get(validityBuffer, 3)); assertEquals(0, BitVectorHelper.get(validityBuffer, 4)); assertEquals(1, BitVectorHelper.get(validityBuffer, 5)); assertEquals(0, offsetBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offsetBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(4 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(5, offsetBuffer.getInt(5 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, sizeBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(3, sizeBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(2, sizeBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(4 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(3, sizeBuffer.getInt(5 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); // validate values Object result = vector.getObject(1); ArrayList resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Long.valueOf(10), resultSet.get(0)); assertEquals(Long.valueOf(20), resultSet.get(1)); assertEquals(Long.valueOf(30), resultSet.get(2)); result = vector.getObject(3); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(40), resultSet.get(0)); assertEquals(Long.valueOf(50), resultSet.get(1)); result = vector.getObject(5); resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Long.valueOf(60), resultSet.get(0)); assertEquals(Long.valueOf(70), resultSet.get(1)); assertEquals(Long.valueOf(80), resultSet.get(2)); vector.validate(); } } @Test public void testSetNull3() { try (LargeListViewVector vector = LargeListViewVector.empty("largelistview", allocator)) { // validate setting values first and then writing nulls UnionLargeListViewWriter writer = vector.getWriter(); writer.allocate(); writer.setPosition(1); writer.startListView(); writer.bigInt().writeBigInt(10); writer.bigInt().writeBigInt(20); writer.bigInt().writeBigInt(30); writer.endListView(); writer.setPosition(3); writer.startListView(); writer.bigInt().writeBigInt(40); writer.bigInt().writeBigInt(50); writer.endListView(); writer.setPosition(5); writer.startListView(); writer.bigInt().writeBigInt(60); writer.bigInt().writeBigInt(70); writer.bigInt().writeBigInt(80); writer.endListView(); vector.setNull(0); vector.setNull(2); vector.setNull(4); vector.setValueCount(6); assertTrue(vector.isNull(0)); assertFalse(vector.isNull(1)); assertTrue(vector.isNull(2)); assertFalse(vector.isNull(3)); assertTrue(vector.isNull(4)); assertFalse(vector.isNull(5)); // validate buffers final ArrowBuf validityBuffer = vector.getValidityBuffer(); final ArrowBuf offsetBuffer = vector.getOffsetBuffer(); final ArrowBuf sizeBuffer = vector.getSizeBuffer(); assertEquals(0, BitVectorHelper.get(validityBuffer, 0)); assertEquals(1, BitVectorHelper.get(validityBuffer, 1)); assertEquals(0, BitVectorHelper.get(validityBuffer, 2)); assertEquals(1, BitVectorHelper.get(validityBuffer, 3)); assertEquals(0, BitVectorHelper.get(validityBuffer, 4)); assertEquals(1, BitVectorHelper.get(validityBuffer, 5)); assertEquals(0, offsetBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offsetBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(4 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(5, offsetBuffer.getInt(5 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, sizeBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(3, sizeBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(2, sizeBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(4 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(3, sizeBuffer.getInt(5 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); // validate values Object result = vector.getObject(1); ArrayList resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Long.valueOf(10), resultSet.get(0)); assertEquals(Long.valueOf(20), resultSet.get(1)); assertEquals(Long.valueOf(30), resultSet.get(2)); result = vector.getObject(3); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(40), resultSet.get(0)); assertEquals(Long.valueOf(50), resultSet.get(1)); result = vector.getObject(5); resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Long.valueOf(60), resultSet.get(0)); assertEquals(Long.valueOf(70), resultSet.get(1)); assertEquals(Long.valueOf(80), resultSet.get(2)); vector.validate(); } } @Test public void testOverWrite1() { try (LargeListViewVector vector = LargeListViewVector.empty("largelistview", allocator)) { UnionLargeListViewWriter writer = vector.getWriter(); writer.allocate(); writer.setPosition(0); writer.startListView(); writer.bigInt().writeBigInt(10); writer.bigInt().writeBigInt(20); writer.bigInt().writeBigInt(30); writer.endListView(); writer.setPosition(1); writer.startListView(); writer.bigInt().writeBigInt(40); writer.bigInt().writeBigInt(50); writer.endListView(); vector.setValueCount(2); writer.setPosition(0); writer.startListView(); writer.bigInt().writeBigInt(60); writer.bigInt().writeBigInt(70); writer.endListView(); writer.setPosition(1); writer.startListView(); writer.bigInt().writeBigInt(80); writer.bigInt().writeBigInt(90); writer.endListView(); vector.setValueCount(2); Object result = vector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(60), resultSet.get(0)); assertEquals(Long.valueOf(70), resultSet.get(1)); result = vector.getObject(1); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(80), resultSet.get(0)); assertEquals(Long.valueOf(90), resultSet.get(1)); vector.validate(); } } @Test public void testOverwriteWithNull() { try (LargeListViewVector vector = LargeListViewVector.empty("largelistview", allocator)) { UnionLargeListViewWriter writer = vector.getWriter(); writer.allocate(); ArrowBuf offsetBuffer = vector.getOffsetBuffer(); ArrowBuf sizeBuffer = vector.getSizeBuffer(); writer.setPosition(0); writer.startListView(); writer.bigInt().writeBigInt(10); writer.bigInt().writeBigInt(20); writer.bigInt().writeBigInt(30); writer.endListView(); writer.setPosition(1); writer.startListView(); writer.bigInt().writeBigInt(40); writer.bigInt().writeBigInt(50); writer.endListView(); vector.setValueCount(2); assertEquals(0, offsetBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offsetBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, sizeBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(2, sizeBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); vector.setNull(0); assertEquals(0, offsetBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, sizeBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); vector.setNull(1); assertEquals(0, offsetBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, sizeBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertTrue(vector.isNull(0)); assertTrue(vector.isNull(1)); writer.setPosition(0); writer.startListView(); writer.bigInt().writeBigInt(60); writer.bigInt().writeBigInt(70); writer.endListView(); assertEquals(0, offsetBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(2, sizeBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); writer.setPosition(1); writer.startListView(); writer.bigInt().writeBigInt(80); writer.bigInt().writeBigInt(90); writer.endListView(); assertEquals(2, offsetBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(2, sizeBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); vector.setValueCount(2); assertFalse(vector.isNull(0)); assertFalse(vector.isNull(1)); Object result = vector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(60), resultSet.get(0)); assertEquals(Long.valueOf(70), resultSet.get(1)); result = vector.getObject(1); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(80), resultSet.get(0)); assertEquals(Long.valueOf(90), resultSet.get(1)); vector.validate(); } } @Test public void testOutOfOrderOffset1() { // [[12, -7, 25], null, [0, -127, 127, 50], [], [50, 12]] try (LargeListViewVector largeListViewVector = LargeListViewVector.empty("largelistview", allocator)) { // Allocate buffers in largeListViewVector by calling `allocateNew` method. largeListViewVector.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); largeListViewVector.initializeChildrenFromFields(Collections.singletonList(field)); // Set values in the child vector. FieldVector fieldVector = largeListViewVector.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. largeListViewVector.setValidity(0, 1); largeListViewVector.setValidity(1, 0); largeListViewVector.setValidity(2, 1); largeListViewVector.setValidity(3, 1); largeListViewVector.setValidity(4, 1); largeListViewVector.setOffset(0, 4); largeListViewVector.setOffset(1, 7); largeListViewVector.setOffset(2, 0); largeListViewVector.setOffset(3, 0); largeListViewVector.setOffset(4, 3); largeListViewVector.setSize(0, 3); largeListViewVector.setSize(1, 0); largeListViewVector.setSize(2, 4); largeListViewVector.setSize(3, 0); largeListViewVector.setSize(4, 2); // Set value count using `setValueCount` method. largeListViewVector.setValueCount(5); final ArrowBuf offSetBuffer = largeListViewVector.getOffsetBuffer(); final ArrowBuf sizeBuffer = largeListViewVector.getSizeBuffer(); // check offset buffer assertEquals(4, offSetBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(7, offSetBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offSetBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offSetBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(4 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); // check size buffer assertEquals(3, sizeBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(4, sizeBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(2, sizeBuffer.getInt(4 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); // check child vector assertEquals(0, ((SmallIntVector) largeListViewVector.getDataVector()).get(0)); assertEquals(-127, ((SmallIntVector) largeListViewVector.getDataVector()).get(1)); assertEquals(127, ((SmallIntVector) largeListViewVector.getDataVector()).get(2)); assertEquals(50, ((SmallIntVector) largeListViewVector.getDataVector()).get(3)); assertEquals(12, ((SmallIntVector) largeListViewVector.getDataVector()).get(4)); assertEquals(-7, ((SmallIntVector) largeListViewVector.getDataVector()).get(5)); assertEquals(25, ((SmallIntVector) largeListViewVector.getDataVector()).get(6)); // check values Object result = largeListViewVector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Short.valueOf("12"), resultSet.get(0)); assertEquals(Short.valueOf("-7"), resultSet.get(1)); assertEquals(Short.valueOf("25"), resultSet.get(2)); assertTrue(largeListViewVector.isNull(1)); result = largeListViewVector.getObject(2); resultSet = (ArrayList) result; assertEquals(4, resultSet.size()); assertEquals(Short.valueOf("0"), resultSet.get(0)); assertEquals(Short.valueOf("-127"), resultSet.get(1)); assertEquals(Short.valueOf("127"), resultSet.get(2)); assertEquals(Short.valueOf("50"), resultSet.get(3)); assertTrue(largeListViewVector.isEmpty(3)); result = largeListViewVector.getObject(4); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Short.valueOf("50"), resultSet.get(0)); assertEquals(Short.valueOf("12"), resultSet.get(1)); largeListViewVector.validate(); } } private int validateSizeBufferAndCalculateMinOffset( int start, int splitLength, ArrowBuf fromOffsetBuffer, ArrowBuf fromSizeBuffer, ArrowBuf toSizeBuffer) { int minOffset = fromOffsetBuffer.getInt((long) start * LargeListViewVector.OFFSET_WIDTH); int fromDataLength; int toDataLength; for (int i = 0; i < splitLength; i++) { fromDataLength = fromSizeBuffer.getInt((long) (start + i) * LargeListViewVector.SIZE_WIDTH); toDataLength = toSizeBuffer.getInt((long) (i) * LargeListViewVector.SIZE_WIDTH); /* validate size */ assertEquals( fromDataLength, toDataLength, "Different data lengths at index: " + i + " and start: " + start); /* calculate minimum offset */ int currentOffset = fromOffsetBuffer.getInt((long) (start + i) * LargeListViewVector.OFFSET_WIDTH); if (currentOffset < minOffset) { minOffset = currentOffset; } } return minOffset; } private void validateOffsetBuffer( int start, int splitLength, ArrowBuf fromOffsetBuffer, ArrowBuf toOffsetBuffer, int minOffset) { int offset1; int offset2; for (int i = 0; i < splitLength; i++) { offset1 = fromOffsetBuffer.getInt((long) (start + i) * LargeListViewVector.OFFSET_WIDTH); offset2 = toOffsetBuffer.getInt((long) (i) * LargeListViewVector.OFFSET_WIDTH); assertEquals( offset1 - minOffset, offset2, "Different offset values at index: " + i + " and start: " + start); } } private void validateDataBuffer( int start, int splitLength, ArrowBuf fromOffsetBuffer, ArrowBuf fromSizeBuffer, BigIntVector fromDataVector, ArrowBuf toOffsetBuffer, BigIntVector toDataVector) { int dataLength; Long fromValue; for (int i = 0; i < splitLength; i++) { dataLength = fromSizeBuffer.getInt((long) (start + i) * LargeListViewVector.SIZE_WIDTH); for (int j = 0; j < dataLength; j++) { fromValue = fromDataVector.getObject( (fromOffsetBuffer.getInt((long) (start + i) * LargeListViewVector.OFFSET_WIDTH) + j)); Long toValue = toDataVector.getObject( (toOffsetBuffer.getInt((long) i * LargeListViewVector.OFFSET_WIDTH) + j)); assertEquals( fromValue, toValue, "Different data values at index: " + i + " and start: " + start); } } } /** * Validate split and transfer of data from fromVector to toVector. Note that this method assumes * that the child vector is BigIntVector. * * @param start start index * @param splitLength length of data to split and transfer * @param fromVector fromVector * @param toVector toVector */ private void validateSplitAndTransfer( TransferPair transferPair, int start, int splitLength, LargeListViewVector fromVector, LargeListViewVector toVector) { transferPair.splitAndTransfer(start, splitLength); /* get offsetBuffer of toVector */ final ArrowBuf toOffsetBuffer = toVector.getOffsetBuffer(); /* get sizeBuffer of toVector */ final ArrowBuf toSizeBuffer = toVector.getSizeBuffer(); /* get dataVector of toVector */ BigIntVector toDataVector = (BigIntVector) toVector.getDataVector(); /* get offsetBuffer of toVector */ final ArrowBuf fromOffsetBuffer = fromVector.getOffsetBuffer(); /* get sizeBuffer of toVector */ final ArrowBuf fromSizeBuffer = fromVector.getSizeBuffer(); /* get dataVector of toVector */ BigIntVector fromDataVector = (BigIntVector) fromVector.getDataVector(); /* validate size buffers */ int minOffset = validateSizeBufferAndCalculateMinOffset( start, splitLength, fromOffsetBuffer, fromSizeBuffer, toSizeBuffer); /* validate offset buffers */ validateOffsetBuffer(start, splitLength, fromOffsetBuffer, toOffsetBuffer, minOffset); /* validate data */ validateDataBuffer( start, splitLength, fromOffsetBuffer, fromSizeBuffer, fromDataVector, toOffsetBuffer, toDataVector); } @Test public void testSplitAndTransfer() throws Exception { try (LargeListViewVector fromVector = LargeListViewVector.empty("sourceVector", allocator)) { /* Explicitly add the dataVector */ MinorType type = MinorType.BIGINT; fromVector.addOrGetVector(FieldType.nullable(type.getType())); UnionLargeListViewWriter listViewWriter = fromVector.getWriter(); /* allocate memory */ listViewWriter.allocate(); /* populate data */ listViewWriter.setPosition(0); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(10); listViewWriter.bigInt().writeBigInt(11); listViewWriter.bigInt().writeBigInt(12); listViewWriter.endListView(); listViewWriter.setPosition(1); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(13); listViewWriter.bigInt().writeBigInt(14); listViewWriter.endListView(); listViewWriter.setPosition(2); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(15); listViewWriter.bigInt().writeBigInt(16); listViewWriter.bigInt().writeBigInt(17); listViewWriter.bigInt().writeBigInt(18); listViewWriter.endListView(); listViewWriter.setPosition(3); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(19); listViewWriter.endListView(); listViewWriter.setPosition(4); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(20); listViewWriter.bigInt().writeBigInt(21); listViewWriter.bigInt().writeBigInt(22); listViewWriter.bigInt().writeBigInt(23); listViewWriter.endListView(); fromVector.setValueCount(5); /* get offset buffer */ final ArrowBuf offsetBuffer = fromVector.getOffsetBuffer(); /* get size buffer */ final ArrowBuf sizeBuffer = fromVector.getSizeBuffer(); /* get dataVector */ BigIntVector dataVector = (BigIntVector) fromVector.getDataVector(); /* check the vector output */ int index = 0; int offset; int size = 0; Long actual; /* index 0 */ assertFalse(fromVector.isNull(index)); offset = offsetBuffer.getInt(index * LargeListViewVector.OFFSET_WIDTH); assertEquals(Integer.toString(0), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(10), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(11), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(12), actual); assertEquals( Integer.toString(3), Integer.toString(sizeBuffer.getInt(index * LargeListViewVector.SIZE_WIDTH))); /* index 1 */ index++; assertFalse(fromVector.isNull(index)); offset = offsetBuffer.getInt(index * LargeListViewVector.OFFSET_WIDTH); assertEquals(Integer.toString(3), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(13), actual); offset++; size++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(14), actual); size++; assertEquals( Integer.toString(size), Integer.toString(sizeBuffer.getInt(index * LargeListViewVector.SIZE_WIDTH))); /* index 2 */ size = 0; index++; assertFalse(fromVector.isNull(index)); offset = offsetBuffer.getInt(index * LargeListViewVector.OFFSET_WIDTH); assertEquals(Integer.toString(5), Integer.toString(offset)); size++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(15), actual); offset++; size++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(16), actual); offset++; size++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(17), actual); offset++; size++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(18), actual); assertEquals( Integer.toString(size), Integer.toString(sizeBuffer.getInt(index * LargeListViewVector.SIZE_WIDTH))); /* index 3 */ size = 0; index++; assertFalse(fromVector.isNull(index)); offset = offsetBuffer.getInt(index * LargeListViewVector.OFFSET_WIDTH); assertEquals(Integer.toString(9), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(19), actual); size++; assertEquals( Integer.toString(size), Integer.toString(sizeBuffer.getInt(index * LargeListViewVector.SIZE_WIDTH))); /* index 4 */ size = 0; index++; assertFalse(fromVector.isNull(index)); offset = offsetBuffer.getInt(index * LargeListViewVector.OFFSET_WIDTH); assertEquals(Integer.toString(10), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(20), actual); offset++; size++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(21), actual); offset++; size++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(22), actual); offset++; size++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(23), actual); size++; assertEquals( Integer.toString(size), Integer.toString(sizeBuffer.getInt(index * LargeListViewVector.SIZE_WIDTH))); /* do split and transfer */ try (LargeListViewVector toVector = LargeListViewVector.empty("toVector", allocator)) { int[][] transferLengths = {{0, 2}, {3, 1}, {4, 1}}; TransferPair transferPair = fromVector.makeTransferPair(toVector); for (final int[] transferLength : transferLengths) { int start = transferLength[0]; int splitLength = transferLength[1]; validateSplitAndTransfer(transferPair, start, splitLength, fromVector, toVector); } } } } @Test public void testGetTransferPairWithField() throws Exception { try (final LargeListViewVector fromVector = LargeListViewVector.empty("listview", allocator)) { UnionLargeListViewWriter writer = fromVector.getWriter(); writer.allocate(); // set some values writer.startListView(); writer.integer().writeInt(1); writer.integer().writeInt(2); writer.endListView(); fromVector.setValueCount(2); final TransferPair transferPair = fromVector.getTransferPair(fromVector.getField(), allocator); final LargeListViewVector toVector = (LargeListViewVector) transferPair.getTo(); // Field inside a new vector created by reusing a field should be the same in memory as the // original field. assertSame(toVector.getField(), fromVector.getField()); } } @Test public void testOutOfOrderOffsetSplitAndTransfer() { // [[12, -7, 25], null, [0, -127, 127, 50], [], [50, 12]] try (LargeListViewVector fromVector = LargeListViewVector.empty("fromVector", allocator)) { // Allocate buffers in LargeListViewVector by calling `allocateNew` method. fromVector.allocateNew(); // Initialize the child vector using `initializeChildrenFromFields` method. FieldType fieldType = new FieldType(true, new ArrowType.Int(64, true), null, null); Field field = new Field("child-vector", fieldType, null); fromVector.initializeChildrenFromFields(Collections.singletonList(field)); // Set values in the child vector. FieldVector fieldVector = fromVector.getDataVector(); fieldVector.clear(); BigIntVector childVector = (BigIntVector) 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. fromVector.setValidity(0, 1); fromVector.setValidity(1, 0); fromVector.setValidity(2, 1); fromVector.setValidity(3, 1); fromVector.setValidity(4, 1); fromVector.setOffset(0, 4); fromVector.setOffset(1, 7); fromVector.setOffset(2, 0); fromVector.setOffset(3, 0); fromVector.setOffset(4, 3); fromVector.setSize(0, 3); fromVector.setSize(1, 0); fromVector.setSize(2, 4); fromVector.setSize(3, 0); fromVector.setSize(4, 2); // Set value count using `setValueCount` method. fromVector.setValueCount(5); final ArrowBuf offSetBuffer = fromVector.getOffsetBuffer(); final ArrowBuf sizeBuffer = fromVector.getSizeBuffer(); // check offset buffer assertEquals(4, offSetBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(7, offSetBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offSetBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offSetBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(4 * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH)); // check size buffer assertEquals(3, sizeBuffer.getInt(0 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(1 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(4, sizeBuffer.getInt(2 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(3 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(2, sizeBuffer.getInt(4 * BaseLargeRepeatedValueViewVector.SIZE_WIDTH)); // check child vector assertEquals(0, ((BigIntVector) fromVector.getDataVector()).get(0)); assertEquals(-127, ((BigIntVector) fromVector.getDataVector()).get(1)); assertEquals(127, ((BigIntVector) fromVector.getDataVector()).get(2)); assertEquals(50, ((BigIntVector) fromVector.getDataVector()).get(3)); assertEquals(12, ((BigIntVector) fromVector.getDataVector()).get(4)); assertEquals(-7, ((BigIntVector) fromVector.getDataVector()).get(5)); assertEquals(25, ((BigIntVector) fromVector.getDataVector()).get(6)); // check values Object result = fromVector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Long.valueOf(12), resultSet.get(0)); assertEquals(Long.valueOf(-7), resultSet.get(1)); assertEquals(Long.valueOf(25), resultSet.get(2)); assertTrue(fromVector.isNull(1)); result = fromVector.getObject(2); resultSet = (ArrayList) result; assertEquals(4, resultSet.size()); assertEquals(Long.valueOf(0), resultSet.get(0)); assertEquals(Long.valueOf(-127), resultSet.get(1)); assertEquals(Long.valueOf(127), resultSet.get(2)); assertEquals(Long.valueOf(50), resultSet.get(3)); assertTrue(fromVector.isEmpty(3)); result = fromVector.getObject(4); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(50), resultSet.get(0)); assertEquals(Long.valueOf(12), resultSet.get(1)); fromVector.validate(); /* do split and transfer */ try (LargeListViewVector toVector = LargeListViewVector.empty("toVector", allocator)) { int[][] transferLengths = {{2, 3}, {0, 1}, {0, 3}}; TransferPair transferPair = fromVector.makeTransferPair(toVector); for (final int[] transferLength : transferLengths) { int start = transferLength[0]; int splitLength = transferLength[1]; validateSplitAndTransfer(transferPair, start, splitLength, fromVector, toVector); } } } } @Test public void testRangeChildVector1() { /* * Non-overlapping ranges * offsets: [0, 2] * sizes: [4, 1] * values: [0, 1, 2, 3] * * vector: [[0, 1, 2, 3], [2]] * */ try (LargeListViewVector largeListViewVector = LargeListViewVector.empty("largelistview", allocator)) { // Allocate buffers in listViewVector by calling `allocateNew` method. largeListViewVector.allocateNew(); // Initialize the child vector using `initializeChildrenFromFields` method. FieldType fieldType = new FieldType(true, new ArrowType.Int(32, true), null, null); Field field = new Field("child-vector", fieldType, null); largeListViewVector.initializeChildrenFromFields(Collections.singletonList(field)); // Set values in the child vector. FieldVector fieldVector = largeListViewVector.getDataVector(); fieldVector.clear(); IntVector childVector = (IntVector) fieldVector; childVector.allocateNew(8); childVector.set(0, 0); childVector.set(1, 1); childVector.set(2, 2); childVector.set(3, 3); childVector.set(4, 4); childVector.set(5, 5); childVector.set(6, 6); childVector.set(7, 7); childVector.setValueCount(8); // Set validity, offset and size buffers using `setValidity`, // `setOffset` and `setSize` methods. largeListViewVector.setValidity(0, 1); largeListViewVector.setValidity(1, 1); largeListViewVector.setOffset(0, 0); largeListViewVector.setOffset(1, 2); largeListViewVector.setSize(0, 4); largeListViewVector.setSize(1, 1); assertEquals(8, largeListViewVector.getDataVector().getValueCount()); largeListViewVector.setValueCount(2); assertEquals(4, largeListViewVector.getDataVector().getValueCount()); IntVector childVector1 = (IntVector) largeListViewVector.getDataVector(); final ArrowBuf dataBuffer = childVector1.getDataBuffer(); final ArrowBuf validityBuffer = childVector1.getValidityBuffer(); // yet the underneath buffer contains the original buffer for (int i = 0; i < validityBuffer.capacity(); i++) { assertEquals(i, dataBuffer.getInt((long) i * IntVector.TYPE_WIDTH)); } } } @Test public void testRangeChildVector2() { /* * Overlapping ranges * offsets: [0, 2] * sizes: [3, 1] * values: [0, 1, 2, 3] * * vector: [[1, 2, 3], [2]] * */ try (LargeListViewVector largeListViewVector = LargeListViewVector.empty("largelistview", allocator)) { // Allocate buffers in listViewVector by calling `allocateNew` method. largeListViewVector.allocateNew(); // Initialize the child vector using `initializeChildrenFromFields` method. FieldType fieldType = new FieldType(true, new ArrowType.Int(32, true), null, null); Field field = new Field("child-vector", fieldType, null); largeListViewVector.initializeChildrenFromFields(Collections.singletonList(field)); // Set values in the child vector. FieldVector fieldVector = largeListViewVector.getDataVector(); fieldVector.clear(); IntVector childVector = (IntVector) fieldVector; childVector.allocateNew(8); childVector.set(0, 0); childVector.set(1, 1); childVector.set(2, 2); childVector.set(3, 3); childVector.set(4, 4); childVector.set(5, 5); childVector.set(6, 6); childVector.set(7, 7); childVector.setValueCount(8); // Set validity, offset and size buffers using `setValidity`, // `setOffset` and `setSize` methods. largeListViewVector.setValidity(0, 1); largeListViewVector.setValidity(1, 1); largeListViewVector.setOffset(0, 1); largeListViewVector.setOffset(1, 2); largeListViewVector.setSize(0, 3); largeListViewVector.setSize(1, 1); assertEquals(8, largeListViewVector.getDataVector().getValueCount()); largeListViewVector.setValueCount(2); assertEquals(4, largeListViewVector.getDataVector().getValueCount()); IntVector childVector1 = (IntVector) largeListViewVector.getDataVector(); final ArrowBuf dataBuffer = childVector1.getDataBuffer(); final ArrowBuf validityBuffer = childVector1.getValidityBuffer(); // yet the underneath buffer contains the original buffer for (int i = 0; i < validityBuffer.capacity(); i++) { assertEquals(i, dataBuffer.getInt((long) i * IntVector.TYPE_WIDTH)); } } } private void writeIntValues(UnionLargeListViewWriter writer, int[] values) { writer.startListView(); for (int v : values) { writer.integer().writeInt(v); } writer.endListView(); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestLargeVarBinaryVector.java000066400000000000000000000125061515451615200332750ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.Objects; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.holders.NullableLargeVarBinaryHolder; import org.apache.arrow.vector.util.ReusableByteArray; 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 TestLargeVarBinaryVector { private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(Integer.MAX_VALUE); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testSetNullableLargeVarBinaryHolder() { try (LargeVarBinaryVector vector = new LargeVarBinaryVector("", allocator)) { vector.allocateNew(100, 10); NullableLargeVarBinaryHolder nullHolder = new NullableLargeVarBinaryHolder(); nullHolder.isSet = 0; NullableLargeVarBinaryHolder binHolder = new NullableLargeVarBinaryHolder(); binHolder.isSet = 1; String str = "hello"; try (ArrowBuf buf = allocator.buffer(16)) { buf.setBytes(0, str.getBytes(StandardCharsets.UTF_8)); binHolder.start = 0; binHolder.end = str.length(); binHolder.buffer = buf; vector.set(0, nullHolder); vector.set(1, binHolder); // verify results assertTrue(vector.isNull(0)); assertEquals( str, new String(Objects.requireNonNull(vector.get(1)), StandardCharsets.UTF_8)); } } } @Test public void testSetNullableLargeVarBinaryHolderSafe() { try (LargeVarBinaryVector vector = new LargeVarBinaryVector("", allocator)) { vector.allocateNew(5, 1); NullableLargeVarBinaryHolder nullHolder = new NullableLargeVarBinaryHolder(); nullHolder.isSet = 0; NullableLargeVarBinaryHolder binHolder = new NullableLargeVarBinaryHolder(); binHolder.isSet = 1; String str = "hello world"; try (ArrowBuf buf = allocator.buffer(16)) { buf.setBytes(0, str.getBytes(StandardCharsets.UTF_8)); binHolder.start = 0; binHolder.end = str.length(); binHolder.buffer = buf; vector.setSafe(0, binHolder); vector.setSafe(1, nullHolder); // verify results assertEquals( str, new String(Objects.requireNonNull(vector.get(0)), StandardCharsets.UTF_8)); assertTrue(vector.isNull(1)); } } } @Test public void testGetBytesRepeatedly() { try (LargeVarBinaryVector vector = new LargeVarBinaryVector("", allocator)) { vector.allocateNew(5, 1); final String str = "hello world"; final String str2 = "foo"; vector.setSafe(0, str.getBytes(StandardCharsets.UTF_8)); vector.setSafe(1, str2.getBytes(StandardCharsets.UTF_8)); // verify results ReusableByteArray reusableByteArray = new ReusableByteArray(); vector.read(0, reusableByteArray); byte[] oldBuffer = reusableByteArray.getBuffer(); assertArrayEquals( str.getBytes(StandardCharsets.UTF_8), Arrays.copyOfRange( reusableByteArray.getBuffer(), 0, (int) reusableByteArray.getLength())); vector.read(1, reusableByteArray); assertArrayEquals( str2.getBytes(StandardCharsets.UTF_8), Arrays.copyOfRange( reusableByteArray.getBuffer(), 0, (int) reusableByteArray.getLength())); // There should not have been any reallocation since the newer value is smaller in length. assertSame(oldBuffer, reusableByteArray.getBuffer()); } } @Test public void testGetTransferPairWithField() { try (BufferAllocator childAllocator1 = allocator.newChildAllocator("child1", 1000000, 1000000); LargeVarBinaryVector v1 = new LargeVarBinaryVector("v1", childAllocator1)) { v1.allocateNew(); v1.setSafe(4094, "hello world".getBytes(StandardCharsets.UTF_8), 0, 11); v1.setValueCount(4001); TransferPair tp = v1.getTransferPair(v1.getField(), allocator); tp.transfer(); LargeVarBinaryVector v2 = (LargeVarBinaryVector) tp.getTo(); assertSame(v1.getField(), v2.getField()); v2.clear(); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestLargeVarCharVector.java000066400000000000000000001006431515451615200327260ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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 static org.junit.jupiter.api.Assertions.fail; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; import java.util.Objects; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.OutOfMemoryException; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.holders.NullableLargeVarCharHolder; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.OversizedAllocationException; import org.apache.arrow.vector.util.Text; 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 TestLargeVarCharVector { private static final byte[] STR1 = "AAAAA1".getBytes(StandardCharsets.UTF_8); private static final byte[] STR2 = "BBBBBBBBB2".getBytes(StandardCharsets.UTF_8); private static final byte[] STR3 = "CCCC3".getBytes(StandardCharsets.UTF_8); private static final byte[] STR4 = "DDDDDDDD4".getBytes(StandardCharsets.UTF_8); private static final byte[] STR5 = "EEE5".getBytes(StandardCharsets.UTF_8); private static final byte[] STR6 = "FFFFF6".getBytes(StandardCharsets.UTF_8); private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(Integer.MAX_VALUE); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testTransfer() { try (BufferAllocator childAllocator1 = allocator.newChildAllocator("child1", 1000000, 1000000); BufferAllocator childAllocator2 = allocator.newChildAllocator("child2", 1000000, 1000000); LargeVarCharVector v1 = new LargeVarCharVector("v1", childAllocator1); LargeVarCharVector v2 = new LargeVarCharVector("v2", childAllocator2); ) { v1.allocateNew(); v1.setSafe(4094, "hello world".getBytes(StandardCharsets.UTF_8), 0, 11); v1.setValueCount(4001); long memoryBeforeTransfer = childAllocator1.getAllocatedMemory(); v1.makeTransferPair(v2).transfer(); assertEquals(0, childAllocator1.getAllocatedMemory()); assertEquals(memoryBeforeTransfer, childAllocator2.getAllocatedMemory()); } } @Test public void testCopyValueSafe() { try (final LargeVarCharVector largeVarCharVector = new LargeVarCharVector("myvector", allocator); final LargeVarCharVector newLargeVarCharVector = new LargeVarCharVector("newvector", allocator)) { largeVarCharVector.allocateNew(10000, 1000); final int valueCount = 500; populateLargeVarcharVector(largeVarCharVector, valueCount, null); final TransferPair tp = largeVarCharVector.makeTransferPair(newLargeVarCharVector); // new vector memory is not pre-allocated, we expect copyValueSafe work fine. for (int i = 0; i < valueCount; i++) { tp.copyValueSafe(i, i); } newLargeVarCharVector.setValueCount(valueCount); for (int i = 0; i < valueCount; i++) { final boolean expectedSet = (i % 3) == 0; if (expectedSet) { assertFalse(largeVarCharVector.isNull(i)); assertFalse(newLargeVarCharVector.isNull(i)); assertArrayEquals(largeVarCharVector.get(i), newLargeVarCharVector.get(i)); } else { assertTrue(newLargeVarCharVector.isNull(i)); } } } } @Test public void testSplitAndTransferNon() { try (final LargeVarCharVector largeVarCharVector = new LargeVarCharVector("myvector", allocator)) { largeVarCharVector.allocateNew(10000, 1000); final int valueCount = 500; populateLargeVarcharVector(largeVarCharVector, valueCount, null); final TransferPair tp = largeVarCharVector.getTransferPair(allocator); try (LargeVarCharVector newLargeVarCharVector = (LargeVarCharVector) tp.getTo()) { tp.splitAndTransfer(0, 0); assertEquals(0, newLargeVarCharVector.getValueCount()); } } } @Test public void testSplitAndTransferAll() { try (final LargeVarCharVector largeVarCharVector = new LargeVarCharVector("myvector", allocator)) { largeVarCharVector.allocateNew(10000, 1000); final int valueCount = 500; populateLargeVarcharVector(largeVarCharVector, valueCount, null); final TransferPair tp = largeVarCharVector.getTransferPair(allocator); try (LargeVarCharVector newLargeVarCharVector = (LargeVarCharVector) tp.getTo()) { tp.splitAndTransfer(0, valueCount); assertEquals(valueCount, newLargeVarCharVector.getValueCount()); } } } @Test public void testInvalidStartIndex() { try (final LargeVarCharVector largeVarCharVector = new LargeVarCharVector("myvector", allocator); final LargeVarCharVector newLargeVarCharVector = new LargeVarCharVector("newvector", allocator)) { largeVarCharVector.allocateNew(10000, 1000); final int valueCount = 500; populateLargeVarcharVector(largeVarCharVector, valueCount, null); final TransferPair tp = largeVarCharVector.makeTransferPair(newLargeVarCharVector); IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () -> tp.splitAndTransfer(valueCount, 10)); assertEquals( "Invalid parameters startIndex: 500, length: 10 for valueCount: 500", e.getMessage()); } } @Test public void testInvalidLength() { try (final LargeVarCharVector largeVarCharVector = new LargeVarCharVector("myvector", allocator); final LargeVarCharVector newLargeVarCharVector = new LargeVarCharVector("newvector", allocator)) { largeVarCharVector.allocateNew(10000, 1000); final int valueCount = 500; populateLargeVarcharVector(largeVarCharVector, valueCount, null); final TransferPair tp = largeVarCharVector.makeTransferPair(newLargeVarCharVector); IllegalArgumentException e = assertThrows( IllegalArgumentException.class, () -> tp.splitAndTransfer(0, valueCount * 2)); assertEquals( "Invalid parameters startIndex: 0, length: 1000 for valueCount: 500", e.getMessage()); } } @Test /* LargeVarCharVector */ public void testSizeOfValueBuffer() { try (final LargeVarCharVector vector = new LargeVarCharVector("", allocator)) { int valueCount = 100; int currentSize = 0; vector.setInitialCapacity(valueCount); vector.allocateNew(); vector.setValueCount(valueCount); for (int i = 0; i < valueCount; i++) { currentSize += i; vector.setSafe(i, new byte[i]); } assertEquals(currentSize, vector.sizeOfValueBuffer()); } } @Test public void testSetLastSetUsage() { final byte[] STR1 = "AAAAA1".getBytes(StandardCharsets.UTF_8); final byte[] STR2 = "BBBBBBBBB2".getBytes(StandardCharsets.UTF_8); final byte[] STR3 = "CCCC3".getBytes(StandardCharsets.UTF_8); final byte[] STR4 = "DDDDDDDD4".getBytes(StandardCharsets.UTF_8); final byte[] STR5 = "EEE5".getBytes(StandardCharsets.UTF_8); final byte[] STR6 = "FFFFF6".getBytes(StandardCharsets.UTF_8); try (final LargeVarCharVector vector = new LargeVarCharVector("myvector", allocator)) { vector.allocateNew(1024 * 10, 1024); setBytes(0, STR1, vector); setBytes(1, STR2, vector); setBytes(2, STR3, vector); setBytes(3, STR4, vector); setBytes(4, STR5, vector); setBytes(5, STR6, vector); /* Check current lastSet */ assertEquals(-1, vector.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); assertArrayEquals(STR5, vector.get(4)); assertArrayEquals(STR6, vector.get(5)); /* * If we don't do setLastSe(5) before setValueCount(), then the latter will corrupt * the value vector by filling in all positions [0,valuecount-1] will empty byte arrays. * Run the test by commenting out next line and we should see incorrect vector output. */ vector.setLastSet(5); vector.setValueCount(20); /* Check current lastSet */ assertEquals(19, vector.getLastSet()); /* Check the vector output again */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); assertArrayEquals(STR5, vector.get(4)); assertArrayEquals(STR6, vector.get(5)); assertEquals(0, vector.getValueLength(6)); assertEquals(0, vector.getValueLength(7)); assertEquals(0, vector.getValueLength(8)); assertEquals(0, vector.getValueLength(9)); assertEquals(0, vector.getValueLength(10)); assertEquals(0, vector.getValueLength(11)); assertEquals(0, vector.getValueLength(12)); assertEquals(0, vector.getValueLength(13)); assertEquals(0, vector.getValueLength(14)); assertEquals(0, vector.getValueLength(15)); assertEquals(0, vector.getValueLength(16)); assertEquals(0, vector.getValueLength(17)); assertEquals(0, vector.getValueLength(18)); assertEquals(0, vector.getValueLength(19)); /* Check offsets */ assertEquals(0, vector.offsetBuffer.getLong(0 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(6, vector.offsetBuffer.getLong(1 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(16, vector.offsetBuffer.getLong(2 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(21, vector.offsetBuffer.getLong(3 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(30, vector.offsetBuffer.getLong(4 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(34, vector.offsetBuffer.getLong(5 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(6 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(7 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(8 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(9 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(10 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(11 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(12 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(13 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(14 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(15 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(16 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(17 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(18 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(19 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); vector.set(19, STR6); assertArrayEquals(STR6, vector.get(19)); assertEquals(40, vector.offsetBuffer.getLong(19 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(46, vector.offsetBuffer.getLong(20 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); } } @Test public void testVectorAllocateNew() { assertThrows( OutOfMemoryException.class, () -> { try (RootAllocator smallAllocator = new RootAllocator(200); LargeVarCharVector vector = new LargeVarCharVector("vec", smallAllocator)) { vector.allocateNew(); } }); } @Test public void testLargeVariableVectorReallocation() { assertThrows( OversizedAllocationException.class, () -> { final LargeVarCharVector vector = new LargeVarCharVector("vector", allocator); // edge case 1: value count = MAX_VALUE_ALLOCATION final long expectedAllocationInBytes = BaseValueVector.MAX_ALLOCATION_SIZE; final int expectedOffsetSize = 10; try { vector.allocateNew(expectedAllocationInBytes, 10); assertTrue(expectedOffsetSize <= vector.getValueCapacity()); assertTrue(expectedAllocationInBytes <= vector.getDataBuffer().capacity()); vector.reAlloc(); assertTrue(expectedOffsetSize * 2 <= vector.getValueCapacity()); assertTrue(expectedAllocationInBytes * 2 <= vector.getDataBuffer().capacity()); } finally { vector.close(); } // common: value count < MAX_VALUE_ALLOCATION try { vector.allocateNew(BaseValueVector.MAX_ALLOCATION_SIZE / 2, 0); vector.reAlloc(); // value allocation reaches to MAX_VALUE_ALLOCATION vector.reAlloc(); // this tests if it overflows } finally { vector.close(); } }); } @Test public void testSplitAndTransfer() { try (final LargeVarCharVector largeVarCharVector = new LargeVarCharVector("myvector", allocator)) { largeVarCharVector.allocateNew(10000, 1000); final int valueCount = 500; final String[] compareArray = new String[valueCount]; populateLargeVarcharVector(largeVarCharVector, valueCount, compareArray); final TransferPair tp = largeVarCharVector.getTransferPair(allocator); try (final LargeVarCharVector newLargeVarCharVector = (LargeVarCharVector) tp.getTo()) { final int[][] startLengths = {{0, 201}, {201, 0}, {201, 200}, {401, 99}}; for (final int[] startLength : startLengths) { final int start = startLength[0]; final int length = startLength[1]; tp.splitAndTransfer(start, length); for (int i = 0; i < length; i++) { final boolean expectedSet = ((start + i) % 3) == 0; if (expectedSet) { final byte[] expectedValue = compareArray[start + i].getBytes(StandardCharsets.UTF_8); assertFalse(newLargeVarCharVector.isNull(i)); assertArrayEquals(expectedValue, newLargeVarCharVector.get(i)); } else { assertTrue(newLargeVarCharVector.isNull(i)); } } } } } } @Test public void testReallocAfterVectorTransfer() { final byte[] STR1 = "AAAAA1".getBytes(StandardCharsets.UTF_8); final byte[] STR2 = "BBBBBBBBB2".getBytes(StandardCharsets.UTF_8); try (final LargeVarCharVector vector = new LargeVarCharVector("vector", allocator)) { /* 4096 values with 10 byte per record */ vector.allocateNew(4096 * 10, 4096); int valueCapacity = vector.getValueCapacity(); assertTrue(valueCapacity >= 4096); /* populate the vector */ for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { vector.set(i, STR1); } else { vector.set(i, STR2); } } /* Check the vector output */ for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { assertArrayEquals(STR1, vector.get(i)); } else { assertArrayEquals(STR2, vector.get(i)); } } /* trigger first realloc */ vector.setSafe(valueCapacity, STR2, 0, STR2.length); assertTrue(vector.getValueCapacity() >= 2 * valueCapacity); while (vector.getByteCapacity() < 10 * vector.getValueCapacity()) { vector.reallocDataBuffer(); } /* populate the remaining vector */ for (int i = valueCapacity; i < vector.getValueCapacity(); i++) { if ((i & 1) == 1) { vector.set(i, STR1); } else { vector.set(i, STR2); } } /* Check the vector output */ valueCapacity = vector.getValueCapacity(); for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { assertArrayEquals(STR1, vector.get(i)); } else { assertArrayEquals(STR2, vector.get(i)); } } /* trigger second realloc */ vector.setSafe(valueCapacity + 10, STR2, 0, STR2.length); assertTrue(vector.getValueCapacity() >= 2 * valueCapacity); while (vector.getByteCapacity() < 10 * vector.getValueCapacity()) { vector.reallocDataBuffer(); } /* populate the remaining vector */ for (int i = valueCapacity; i < vector.getValueCapacity(); i++) { if ((i & 1) == 1) { vector.set(i, STR1); } else { vector.set(i, STR2); } } /* Check the vector output */ valueCapacity = vector.getValueCapacity(); for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { assertArrayEquals(STR1, vector.get(i)); } else { assertArrayEquals(STR2, vector.get(i)); } } /* we are potentially working with 4x the size of vector buffer * that we initially started with. Now let's transfer the vector. */ TransferPair transferPair = vector.getTransferPair(allocator); transferPair.transfer(); try (LargeVarCharVector toVector = (LargeVarCharVector) transferPair.getTo()) { valueCapacity = toVector.getValueCapacity(); for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { assertArrayEquals(STR1, toVector.get(i)); } else { assertArrayEquals(STR2, toVector.get(i)); } } } } } @Test public void testVectorLoadUnload() { try (final LargeVarCharVector vector1 = new LargeVarCharVector("myvector", allocator)) { ValueVectorDataPopulator.setVector(vector1, STR1, STR2, STR3, STR4, STR5, STR6); assertEquals(5, vector1.getLastSet()); vector1.setValueCount(15); assertEquals(14, vector1.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector1.get(0)); assertArrayEquals(STR2, vector1.get(1)); assertArrayEquals(STR3, vector1.get(2)); assertArrayEquals(STR4, vector1.get(3)); assertArrayEquals(STR5, vector1.get(4)); assertArrayEquals(STR6, vector1.get(5)); Field field = vector1.getField(); String fieldName = field.getName(); List fields = new ArrayList<>(); List fieldVectors = new ArrayList<>(); fields.add(field); fieldVectors.add(vector1); Schema schema = new Schema(fields); VectorSchemaRoot schemaRoot1 = new VectorSchemaRoot(schema, fieldVectors, vector1.getValueCount()); VectorUnloader vectorUnloader = new VectorUnloader(schemaRoot1); try (ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch(); VectorSchemaRoot schemaRoot2 = VectorSchemaRoot.create(schema, allocator); ) { VectorLoader vectorLoader = new VectorLoader(schemaRoot2); vectorLoader.load(recordBatch); LargeVarCharVector vector2 = (LargeVarCharVector) schemaRoot2.getVector(fieldName); /* * lastSet would have internally been set by VectorLoader.load() when it invokes * loadFieldBuffers. */ assertEquals(14, vector2.getLastSet()); vector2.setValueCount(25); assertEquals(24, vector2.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector2.get(0)); assertArrayEquals(STR2, vector2.get(1)); assertArrayEquals(STR3, vector2.get(2)); assertArrayEquals(STR4, vector2.get(3)); assertArrayEquals(STR5, vector2.get(4)); assertArrayEquals(STR6, vector2.get(5)); } } } @Test public void testFillEmptiesUsage() { try (final LargeVarCharVector vector = new LargeVarCharVector("myvector", allocator)) { vector.allocateNew(1024 * 10, 1024); setBytes(0, STR1, vector); setBytes(1, STR2, vector); setBytes(2, STR3, vector); setBytes(3, STR4, vector); setBytes(4, STR5, vector); setBytes(5, STR6, vector); /* Check current lastSet */ assertEquals(-1, vector.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); assertArrayEquals(STR5, vector.get(4)); assertArrayEquals(STR6, vector.get(5)); vector.setLastSet(5); /* fill empty byte arrays from index [6, 9] */ vector.fillEmpties(10); /* Check current lastSet */ assertEquals(9, vector.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); assertArrayEquals(STR5, vector.get(4)); assertArrayEquals(STR6, vector.get(5)); assertEquals(0, vector.getValueLength(6)); assertEquals(0, vector.getValueLength(7)); assertEquals(0, vector.getValueLength(8)); assertEquals(0, vector.getValueLength(9)); setBytes(10, STR1, vector); setBytes(11, STR2, vector); vector.setLastSet(11); /* fill empty byte arrays from index [12, 14] */ vector.setValueCount(15); /* Check current lastSet */ assertEquals(14, vector.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); assertArrayEquals(STR5, vector.get(4)); assertArrayEquals(STR6, vector.get(5)); assertEquals(0, vector.getValueLength(6)); assertEquals(0, vector.getValueLength(7)); assertEquals(0, vector.getValueLength(8)); assertEquals(0, vector.getValueLength(9)); assertArrayEquals(STR1, vector.get(10)); assertArrayEquals(STR2, vector.get(11)); assertEquals(0, vector.getValueLength(12)); assertEquals(0, vector.getValueLength(13)); assertEquals(0, vector.getValueLength(14)); /* Check offsets */ assertEquals(0, vector.offsetBuffer.getLong(0 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(6, vector.offsetBuffer.getLong(1 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(16, vector.offsetBuffer.getLong(2 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(21, vector.offsetBuffer.getLong(3 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(30, vector.offsetBuffer.getLong(4 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(34, vector.offsetBuffer.getLong(5 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(6 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(7 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(8 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(9 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getLong(10 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(46, vector.offsetBuffer.getLong(11 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(56, vector.offsetBuffer.getLong(12 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(56, vector.offsetBuffer.getLong(13 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(56, vector.offsetBuffer.getLong(14 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); assertEquals(56, vector.offsetBuffer.getLong(15 * BaseLargeVariableWidthVector.OFFSET_WIDTH)); } } @Test public void testGetBufferAddress1() { try (final LargeVarCharVector vector = new LargeVarCharVector("myvector", allocator)) { ValueVectorDataPopulator.setVector(vector, STR1, STR2, STR3, STR4, STR5, STR6); vector.setValueCount(15); /* check the vector output */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); assertArrayEquals(STR5, vector.get(4)); assertArrayEquals(STR6, vector.get(5)); List buffers = vector.getFieldBuffers(); long bitAddress = vector.getValidityBufferAddress(); long offsetAddress = vector.getOffsetBufferAddress(); long dataAddress = vector.getDataBufferAddress(); assertEquals(3, buffers.size()); assertEquals(bitAddress, buffers.get(0).memoryAddress()); assertEquals(offsetAddress, buffers.get(1).memoryAddress()); assertEquals(dataAddress, buffers.get(2).memoryAddress()); } } @Test public void testSetNullableLargeVarCharHolder() { try (LargeVarCharVector vector = new LargeVarCharVector("", allocator)) { vector.allocateNew(100, 10); NullableLargeVarCharHolder nullHolder = new NullableLargeVarCharHolder(); nullHolder.isSet = 0; NullableLargeVarCharHolder stringHolder = new NullableLargeVarCharHolder(); stringHolder.isSet = 1; String str = "hello"; ArrowBuf buf = allocator.buffer(16); buf.setBytes(0, str.getBytes(StandardCharsets.UTF_8)); stringHolder.start = 0; stringHolder.end = str.length(); stringHolder.buffer = buf; vector.set(0, nullHolder); vector.set(1, stringHolder); // verify results assertTrue(vector.isNull(0)); assertEquals(str, new String(Objects.requireNonNull(vector.get(1)), StandardCharsets.UTF_8)); buf.close(); } } @Test public void testSetNullableLargeVarCharHolderSafe() { try (LargeVarCharVector vector = new LargeVarCharVector("", allocator)) { vector.allocateNew(5, 1); NullableLargeVarCharHolder nullHolder = new NullableLargeVarCharHolder(); nullHolder.isSet = 0; NullableLargeVarCharHolder stringHolder = new NullableLargeVarCharHolder(); stringHolder.isSet = 1; String str = "hello world"; ArrowBuf buf = allocator.buffer(16); buf.setBytes(0, str.getBytes(StandardCharsets.UTF_8)); stringHolder.start = 0; stringHolder.end = str.length(); stringHolder.buffer = buf; vector.setSafe(0, stringHolder); vector.setSafe(1, nullHolder); // verify results assertEquals(str, new String(Objects.requireNonNull(vector.get(0)), StandardCharsets.UTF_8)); assertTrue(vector.isNull(1)); buf.close(); } } @Test public void testGetNullFromLargeVariableWidthVector() { try (final LargeVarCharVector largeVarCharVector = new LargeVarCharVector("largevarcharvec", allocator); final LargeVarBinaryVector largeVarBinaryVector = new LargeVarBinaryVector("largevarbinary", allocator)) { largeVarCharVector.allocateNew(10, 1); largeVarBinaryVector.allocateNew(10, 1); largeVarCharVector.setNull(0); largeVarBinaryVector.setNull(0); assertNull(largeVarCharVector.get(0)); assertNull(largeVarBinaryVector.get(0)); } } @Test public void testLargeVariableWidthVectorNullHashCode() { try (LargeVarCharVector largeVarChVec = new LargeVarCharVector("large var char vector", allocator)) { largeVarChVec.allocateNew(100, 1); largeVarChVec.setValueCount(1); largeVarChVec.set(0, "abc".getBytes(StandardCharsets.UTF_8)); largeVarChVec.setNull(0); assertEquals(0, largeVarChVec.hashCode(0)); } } @Test public void testUnloadLargeVariableWidthVector() { try (final LargeVarCharVector largeVarCharVector = new LargeVarCharVector("var char", allocator)) { largeVarCharVector.allocateNew(5, 2); largeVarCharVector.setValueCount(2); largeVarCharVector.set(0, "abcd".getBytes(StandardCharsets.UTF_8)); List bufs = largeVarCharVector.getFieldBuffers(); assertEquals(3, bufs.size()); ArrowBuf offsetBuf = bufs.get(1); ArrowBuf dataBuf = bufs.get(2); assertEquals(24, offsetBuf.writerIndex()); assertEquals(4, offsetBuf.getLong(8)); assertEquals(4, offsetBuf.getLong(16)); assertEquals(4, dataBuf.writerIndex()); } } @Test public void testNullableType() { try (final LargeVarCharVector vector = new LargeVarCharVector("", allocator)) { vector.setInitialCapacity(512); vector.allocateNew(); assertTrue(vector.getValueCapacity() >= 512); int initialCapacity = vector.getValueCapacity(); try { vector.set(initialCapacity, "foo".getBytes(StandardCharsets.UTF_8)); fail("Expected out of bounds exception"); } catch (Exception e) { // ok } vector.reAlloc(); assertTrue(vector.getValueCapacity() >= 2 * initialCapacity); vector.set(initialCapacity, "foo".getBytes(StandardCharsets.UTF_8)); assertEquals("foo", new String(vector.get(initialCapacity), StandardCharsets.UTF_8)); } } @Test public void testGetTextRepeatedly() { try (final LargeVarCharVector vector = new LargeVarCharVector("myvector", allocator)) { ValueVectorDataPopulator.setVector(vector, STR1, STR2); vector.setValueCount(2); /* check the vector output */ Text text = new Text(); vector.read(0, text); byte[] result = new byte[(int) text.getLength()]; System.arraycopy(text.getBytes(), 0, result, 0, (int) text.getLength()); assertArrayEquals(STR1, result); vector.read(1, text); result = new byte[(int) text.getLength()]; System.arraycopy(text.getBytes(), 0, result, 0, (int) text.getLength()); assertArrayEquals(STR2, text.getBytes()); } } @Test public void testGetTransferPairWithField() { try (BufferAllocator childAllocator1 = allocator.newChildAllocator("child1", 1000000, 1000000); LargeVarCharVector v1 = new LargeVarCharVector("v1", childAllocator1)) { v1.allocateNew(); v1.setSafe(4094, "hello world".getBytes(StandardCharsets.UTF_8), 0, 11); v1.setValueCount(4001); TransferPair tp = v1.getTransferPair(v1.getField(), allocator); tp.transfer(); LargeVarCharVector v2 = (LargeVarCharVector) tp.getTo(); assertSame(v1.getField(), v2.getField()); v2.clear(); } } private void populateLargeVarcharVector( final LargeVarCharVector vector, int valueCount, String[] values) { for (int i = 0; i < valueCount; i += 3) { final String s = String.format("%010d", i); vector.set(i, s.getBytes(StandardCharsets.UTF_8)); if (values != null) { values[i] = s; } } vector.setValueCount(valueCount); } public static void setBytes(int index, byte[] bytes, LargeVarCharVector vector) { final long currentOffset = vector.offsetBuffer.getLong((long) index * BaseLargeVariableWidthVector.OFFSET_WIDTH); BitVectorHelper.setBit(vector.validityBuffer, index); vector.offsetBuffer.setLong( (long) (index + 1) * BaseLargeVariableWidthVector.OFFSET_WIDTH, currentOffset + bytes.length); vector.valueBuffer.setBytes(currentOffset, bytes, 0, bytes.length); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestListVector.java000066400000000000000000001422071515451615200313420ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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; 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.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; 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; public class TestListVector { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testCopyFrom() throws Exception { try (ListVector inVector = ListVector.empty("input", allocator); ListVector outVector = ListVector.empty("output", allocator)) { UnionListWriter writer = inVector.getWriter(); writer.allocate(); // populate input vector with the following records // [1, 2, 3] // null // [] writer.setPosition(0); // optional writer.startList(); writer.bigInt().writeBigInt(1); writer.bigInt().writeBigInt(2); writer.bigInt().writeBigInt(3); writer.endList(); writer.setPosition(2); writer.startList(); writer.endList(); writer.setValueCount(3); // copy values from input to output outVector.allocateNew(); for (int i = 0; i < 3; i++) { outVector.copyFrom(i, i, inVector); } outVector.setValueCount(3); // assert the output vector is correct FieldReader reader = outVector.getReader(); assertTrue(reader.isSet(), "shouldn't be null"); reader.setPosition(1); assertFalse(reader.isSet(), "should be null"); reader.setPosition(2); assertTrue(reader.isSet(), "shouldn't be null"); /* index 0 */ Object result = outVector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Long.valueOf(1), resultSet.get(0)); assertEquals(Long.valueOf(2), resultSet.get(1)); assertEquals(Long.valueOf(3), resultSet.get(2)); /* index 1 */ result = outVector.getObject(1); assertNull(result); /* index 2 */ result = outVector.getObject(2); resultSet = (ArrayList) result; assertEquals(0, resultSet.size()); /* 3+0+0/3 */ assertEquals(1.0D, inVector.getDensity(), 0); } } @Test public void testSetLastSetUsage() throws Exception { try (ListVector listVector = ListVector.empty("input", allocator)) { /* Explicitly add the dataVector */ MinorType type = MinorType.BIGINT; listVector.addOrGetVector(FieldType.nullable(type.getType())); /* allocate memory */ listVector.allocateNew(); /* get inner buffers; validityBuffer and offsetBuffer */ ArrowBuf validityBuffer = listVector.getValidityBuffer(); ArrowBuf offsetBuffer = listVector.getOffsetBuffer(); /* get the underlying data vector -- BigIntVector */ BigIntVector dataVector = (BigIntVector) listVector.getDataVector(); /* check current lastSet */ assertEquals(Integer.toString(-1), Integer.toString(listVector.getLastSet())); int index = 0; int offset; /* write [10, 11, 12] to the list vector at index 0 */ BitVectorHelper.setBit(validityBuffer, index); dataVector.setSafe(0, 1, 10); dataVector.setSafe(1, 1, 11); dataVector.setSafe(2, 1, 12); offsetBuffer.setInt((index + 1) * ListVector.OFFSET_WIDTH, 3); index += 1; /* write [13, 14] to the list vector at index 1 */ BitVectorHelper.setBit(validityBuffer, index); dataVector.setSafe(3, 1, 13); dataVector.setSafe(4, 1, 14); offsetBuffer.setInt((index + 1) * ListVector.OFFSET_WIDTH, 5); index += 1; /* write [15, 16, 17] to the list vector at index 2 */ BitVectorHelper.setBit(validityBuffer, index); dataVector.setSafe(5, 1, 15); dataVector.setSafe(6, 1, 16); dataVector.setSafe(7, 1, 17); offsetBuffer.setInt((index + 1) * ListVector.OFFSET_WIDTH, 8); /* check current lastSet */ assertEquals(Integer.toString(-1), Integer.toString(listVector.getLastSet())); /* set lastset and arbitrary valuecount for list vector. * * NOTE: if we don't execute setLastSet() before setLastValueCount(), then * the latter will corrupt the offsetBuffer and thus the accessor will not * retrieve the correct values from underlying dataBuffer. Run the test * by commenting out next line and we should see failures from 5th assert * onwards. This is why doing setLastSet() is important before setValueCount() * once the vector has been loaded. * * Another important thing to remember is the value of lastSet itself. * Even though the listVector has elements till index 2 only, the lastSet should * be set as 3. This is because the offsetBuffer has valid offsets filled till index 3. * If we do setLastSet(2), the offsetBuffer at index 3 will contain incorrect value * after execution of setValueCount(). * * correct state of the listVector * bitvector {1, 1, 1, 0, 0.... } * offsetvector {0, 3, 5, 8, 8, 8.....} * datavector { [10, 11, 12], * [13, 14], * [15, 16, 17] * } * * if we don't do setLastSet() before setValueCount --> incorrect state * bitvector {1, 1, 1, 0, 0.... } * offsetvector {0, 0, 0, 0, 0, 0.....} * datavector { [10, 11, 12], * [13, 14], * [15, 16, 17] * } * * if we do setLastSet(2) before setValueCount --> incorrect state * bitvector {1, 1, 1, 0, 0.... } * offsetvector {0, 3, 5, 5, 5, 5.....} * datavector { [10, 11, 12], * [13, 14], * [15, 16, 17] * } */ listVector.setLastSet(2); listVector.setValueCount(10); /* (3+2+3)/10 */ assertEquals(0.8D, listVector.getDensity(), 0); index = 0; offset = offsetBuffer.getInt(index * ListVector.OFFSET_WIDTH); assertEquals(Integer.toString(0), Integer.toString(offset)); Long actual = dataVector.getObject(offset); assertEquals(Long.valueOf(10), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(11), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(12), actual); index++; offset = offsetBuffer.getInt(index * ListVector.OFFSET_WIDTH); assertEquals(Integer.toString(3), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(13), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(14), actual); index++; offset = offsetBuffer.getInt(index * ListVector.OFFSET_WIDTH); assertEquals(Integer.toString(5), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(15), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(16), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(17), actual); index++; offset = offsetBuffer.getInt(index * ListVector.OFFSET_WIDTH); assertEquals(Integer.toString(8), Integer.toString(offset)); actual = dataVector.getObject(offset); assertNull(actual); } } @Test public void testSplitAndTransfer() throws Exception { try (ListVector listVector = ListVector.empty("sourceVector", allocator)) { /* Explicitly add the dataVector */ MinorType type = MinorType.BIGINT; listVector.addOrGetVector(FieldType.nullable(type.getType())); UnionListWriter listWriter = listVector.getWriter(); /* allocate memory */ listWriter.allocate(); /* populate data */ listWriter.setPosition(0); listWriter.startList(); listWriter.bigInt().writeBigInt(10); listWriter.bigInt().writeBigInt(11); listWriter.bigInt().writeBigInt(12); listWriter.endList(); listWriter.setPosition(1); listWriter.startList(); listWriter.bigInt().writeBigInt(13); listWriter.bigInt().writeBigInt(14); listWriter.endList(); listWriter.setPosition(2); listWriter.startList(); listWriter.bigInt().writeBigInt(15); listWriter.bigInt().writeBigInt(16); listWriter.bigInt().writeBigInt(17); listWriter.bigInt().writeBigInt(18); listWriter.endList(); listWriter.setPosition(3); listWriter.startList(); listWriter.bigInt().writeBigInt(19); listWriter.endList(); listWriter.setPosition(4); listWriter.startList(); listWriter.bigInt().writeBigInt(20); listWriter.bigInt().writeBigInt(21); listWriter.bigInt().writeBigInt(22); listWriter.bigInt().writeBigInt(23); listWriter.endList(); listVector.setValueCount(5); assertEquals(4, listVector.getLastSet()); /* get offset buffer */ final ArrowBuf offsetBuffer = listVector.getOffsetBuffer(); /* get dataVector */ BigIntVector dataVector = (BigIntVector) listVector.getDataVector(); /* check the vector output */ int index = 0; int offset; Long actual; /* index 0 */ assertFalse(listVector.isNull(index)); offset = offsetBuffer.getInt(index * ListVector.OFFSET_WIDTH); assertEquals(Integer.toString(0), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(10), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(11), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(12), actual); /* index 1 */ index++; assertFalse(listVector.isNull(index)); offset = offsetBuffer.getInt(index * ListVector.OFFSET_WIDTH); assertEquals(Integer.toString(3), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(13), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(14), actual); /* index 2 */ index++; assertFalse(listVector.isNull(index)); offset = offsetBuffer.getInt(index * ListVector.OFFSET_WIDTH); assertEquals(Integer.toString(5), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(15), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(16), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(17), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(18), actual); /* index 3 */ index++; assertFalse(listVector.isNull(index)); offset = offsetBuffer.getInt(index * ListVector.OFFSET_WIDTH); assertEquals(Integer.toString(9), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(19), actual); /* index 4 */ index++; assertFalse(listVector.isNull(index)); offset = offsetBuffer.getInt(index * ListVector.OFFSET_WIDTH); assertEquals(Integer.toString(10), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(20), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(21), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(22), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(23), actual); /* index 5 */ index++; assertTrue(listVector.isNull(index)); offset = offsetBuffer.getInt(index * ListVector.OFFSET_WIDTH); assertEquals(Integer.toString(14), Integer.toString(offset)); /* do split and transfer */ try (ListVector toVector = ListVector.empty("toVector", allocator)) { TransferPair transferPair = listVector.makeTransferPair(toVector); int[][] transferLengths = {{0, 2}, {3, 1}, {4, 1}}; for (final int[] transferLength : transferLengths) { int start = transferLength[0]; int splitLength = transferLength[1]; int dataLength1 = 0; int dataLength2 = 0; int offset1 = 0; int offset2 = 0; transferPair.splitAndTransfer(start, splitLength); /* get offsetBuffer of toVector */ final ArrowBuf toOffsetBuffer = toVector.getOffsetBuffer(); /* get dataVector of toVector */ BigIntVector dataVector1 = (BigIntVector) toVector.getDataVector(); for (int i = 0; i < splitLength; i++) { dataLength1 = offsetBuffer.getInt((start + i + 1) * ListVector.OFFSET_WIDTH) - offsetBuffer.getInt((start + i) * ListVector.OFFSET_WIDTH); dataLength2 = toOffsetBuffer.getInt((i + 1) * ListVector.OFFSET_WIDTH) - toOffsetBuffer.getInt(i * ListVector.OFFSET_WIDTH); assertEquals( dataLength1, dataLength2, "Different data lengths at index: " + i + " and start: " + start); offset1 = offsetBuffer.getInt((start + i) * ListVector.OFFSET_WIDTH); offset2 = toOffsetBuffer.getInt(i * ListVector.OFFSET_WIDTH); for (int j = 0; j < dataLength1; j++) { assertEquals( dataVector.getObject(offset1), dataVector1.getObject(offset2), "Different data at indexes: " + offset1 + " and " + offset2); offset1++; offset2++; } } } } } } @Test public void testNestedListVector() throws Exception { try (ListVector listVector = ListVector.empty("sourceVector", allocator)) { UnionListWriter listWriter = listVector.getWriter(); /* allocate memory */ listWriter.allocate(); /* the dataVector that backs a listVector will also be a * listVector for this test. */ /* write one or more inner lists at index 0 */ listWriter.setPosition(0); listWriter.startList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(50); listWriter.list().bigInt().writeBigInt(100); listWriter.list().bigInt().writeBigInt(200); listWriter.list().endList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(75); listWriter.list().bigInt().writeBigInt(125); listWriter.list().bigInt().writeBigInt(150); listWriter.list().bigInt().writeBigInt(175); listWriter.list().endList(); listWriter.endList(); /* write one or more inner lists at index 1 */ listWriter.setPosition(1); listWriter.startList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(10); listWriter.list().endList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(15); listWriter.list().bigInt().writeBigInt(20); listWriter.list().endList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(25); listWriter.list().bigInt().writeBigInt(30); listWriter.list().bigInt().writeBigInt(35); listWriter.list().endList(); listWriter.endList(); assertEquals(1, listVector.getLastSet()); listVector.setValueCount(2); assertEquals(2, listVector.getValueCount()); /* get listVector value at index 0 -- the value itself is a listvector */ Object result = listVector.getObject(0); ArrayList> resultSet = (ArrayList>) result; ArrayList list; assertEquals(2, resultSet.size()); /* 2 inner lists at index 0 */ assertEquals(3, resultSet.get(0).size()); /* size of first inner list */ assertEquals(4, resultSet.get(1).size()); /* size of second inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(50), list.get(0)); assertEquals(Long.valueOf(100), list.get(1)); assertEquals(Long.valueOf(200), list.get(2)); list = resultSet.get(1); assertEquals(Long.valueOf(75), list.get(0)); assertEquals(Long.valueOf(125), list.get(1)); assertEquals(Long.valueOf(150), list.get(2)); assertEquals(Long.valueOf(175), list.get(3)); /* get listVector value at index 1 -- the value itself is a listvector */ result = listVector.getObject(1); resultSet = (ArrayList>) result; assertEquals(3, resultSet.size()); /* 3 inner lists at index 1 */ assertEquals(1, resultSet.get(0).size()); /* size of first inner list */ assertEquals(2, resultSet.get(1).size()); /* size of second inner list */ assertEquals(3, resultSet.get(2).size()); /* size of third inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(10), list.get(0)); list = resultSet.get(1); assertEquals(Long.valueOf(15), list.get(0)); assertEquals(Long.valueOf(20), list.get(1)); list = resultSet.get(2); assertEquals(Long.valueOf(25), list.get(0)); assertEquals(Long.valueOf(30), list.get(1)); assertEquals(Long.valueOf(35), list.get(2)); /* check underlying bitVector */ assertFalse(listVector.isNull(0)); assertFalse(listVector.isNull(1)); /* check underlying offsets */ final ArrowBuf offsetBuffer = listVector.getOffsetBuffer(); /* listVector has 2 lists at index 0 and 3 lists at index 1 */ assertEquals(0, offsetBuffer.getInt(0 * ListVector.OFFSET_WIDTH)); assertEquals(2, offsetBuffer.getInt(1 * ListVector.OFFSET_WIDTH)); assertEquals(5, offsetBuffer.getInt(2 * ListVector.OFFSET_WIDTH)); } } @Test public void testNestedListVector1() throws Exception { try (ListVector listVector = ListVector.empty("sourceVector", allocator)) { MinorType listType = MinorType.LIST; MinorType scalarType = MinorType.BIGINT; listVector.addOrGetVector(FieldType.nullable(listType.getType())); ListVector innerList1 = (ListVector) listVector.getDataVector(); innerList1.addOrGetVector(FieldType.nullable(listType.getType())); ListVector innerList2 = (ListVector) innerList1.getDataVector(); innerList2.addOrGetVector(FieldType.nullable(listType.getType())); ListVector innerList3 = (ListVector) innerList2.getDataVector(); innerList3.addOrGetVector(FieldType.nullable(listType.getType())); ListVector innerList4 = (ListVector) innerList3.getDataVector(); innerList4.addOrGetVector(FieldType.nullable(listType.getType())); ListVector innerList5 = (ListVector) innerList4.getDataVector(); innerList5.addOrGetVector(FieldType.nullable(listType.getType())); ListVector innerList6 = (ListVector) innerList5.getDataVector(); innerList6.addOrGetVector(FieldType.nullable(scalarType.getType())); listVector.setInitialCapacity(128); } } @Test public void testNestedListVector2() throws Exception { try (ListVector listVector = ListVector.empty("sourceVector", allocator)) { listVector.setInitialCapacity(1); UnionListWriter listWriter = listVector.getWriter(); /* allocate memory */ listWriter.allocate(); /* write one or more inner lists at index 0 */ listWriter.setPosition(0); listWriter.startList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(50); listWriter.list().bigInt().writeBigInt(100); listWriter.list().bigInt().writeBigInt(200); listWriter.list().endList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(75); listWriter.list().bigInt().writeBigInt(125); listWriter.list().endList(); listWriter.endList(); /* write one or more inner lists at index 1 */ listWriter.setPosition(1); listWriter.startList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(15); listWriter.list().bigInt().writeBigInt(20); listWriter.list().endList(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(25); listWriter.list().bigInt().writeBigInt(30); listWriter.list().bigInt().writeBigInt(35); listWriter.list().endList(); listWriter.endList(); assertEquals(1, listVector.getLastSet()); listVector.setValueCount(2); assertEquals(2, listVector.getValueCount()); /* get listVector value at index 0 -- the value itself is a listvector */ Object result = listVector.getObject(0); ArrayList> resultSet = (ArrayList>) result; ArrayList list; assertEquals(2, resultSet.size()); /* 2 inner lists at index 0 */ assertEquals(3, resultSet.get(0).size()); /* size of first inner list */ assertEquals(2, resultSet.get(1).size()); /* size of second inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(50), list.get(0)); assertEquals(Long.valueOf(100), list.get(1)); assertEquals(Long.valueOf(200), list.get(2)); list = resultSet.get(1); assertEquals(Long.valueOf(75), list.get(0)); assertEquals(Long.valueOf(125), list.get(1)); /* get listVector value at index 1 -- the value itself is a listvector */ result = listVector.getObject(1); resultSet = (ArrayList>) result; assertEquals(2, resultSet.size()); /* 3 inner lists at index 1 */ assertEquals(2, resultSet.get(0).size()); /* size of first inner list */ assertEquals(3, resultSet.get(1).size()); /* size of second inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(15), list.get(0)); assertEquals(Long.valueOf(20), list.get(1)); list = resultSet.get(1); assertEquals(Long.valueOf(25), list.get(0)); assertEquals(Long.valueOf(30), list.get(1)); assertEquals(Long.valueOf(35), list.get(2)); /* check underlying bitVector */ assertFalse(listVector.isNull(0)); assertFalse(listVector.isNull(1)); /* check underlying offsets */ final ArrowBuf offsetBuffer = listVector.getOffsetBuffer(); /* listVector has 2 lists at index 0 and 3 lists at index 1 */ assertEquals(0, offsetBuffer.getInt(0 * ListVector.OFFSET_WIDTH)); assertEquals(2, offsetBuffer.getInt(1 * ListVector.OFFSET_WIDTH)); assertEquals(4, offsetBuffer.getInt(2 * ListVector.OFFSET_WIDTH)); } } @Test public void testGetBufferAddress() throws Exception { try (ListVector listVector = ListVector.empty("vector", allocator)) { UnionListWriter listWriter = listVector.getWriter(); boolean error = false; listWriter.allocate(); listWriter.setPosition(0); listWriter.startList(); listWriter.bigInt().writeBigInt(50); listWriter.bigInt().writeBigInt(100); listWriter.bigInt().writeBigInt(200); listWriter.endList(); listWriter.setPosition(1); listWriter.startList(); listWriter.bigInt().writeBigInt(250); listWriter.bigInt().writeBigInt(300); listWriter.endList(); listVector.setValueCount(2); /* check listVector contents */ Object result = listVector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Long.valueOf(50), resultSet.get(0)); assertEquals(Long.valueOf(100), resultSet.get(1)); assertEquals(Long.valueOf(200), resultSet.get(2)); result = listVector.getObject(1); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(250), resultSet.get(0)); assertEquals(Long.valueOf(300), resultSet.get(1)); List buffers = listVector.getFieldBuffers(); long bitAddress = listVector.getValidityBufferAddress(); long offsetAddress = listVector.getOffsetBufferAddress(); try { listVector.getDataBufferAddress(); } catch (UnsupportedOperationException ue) { error = true; } finally { assertTrue(error); } assertEquals(2, buffers.size()); assertEquals(bitAddress, buffers.get(0).memoryAddress()); assertEquals(offsetAddress, buffers.get(1).memoryAddress()); /* (3+2)/2 */ assertEquals(2.5, listVector.getDensity(), 0); } } @Test public void testConsistentChildName() throws Exception { try (ListVector listVector = ListVector.empty("sourceVector", allocator)) { String emptyListStr = listVector.getField().toString(); assertTrue(emptyListStr.contains(ListVector.DATA_VECTOR_NAME)); listVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); String emptyVectorStr = listVector.getField().toString(); assertTrue(emptyVectorStr.contains(ListVector.DATA_VECTOR_NAME)); } } @Test public void testSetInitialCapacity() { try (final ListVector vector = ListVector.empty("", allocator)) { vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); /* * use the default multiplier of 5, * 512 * 5 => 2560 * 4 => 10240 bytes => 16KB => 4096 value capacity. */ vector.setInitialCapacity(512); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 512 * 5); /* use density as 4 */ vector.setInitialCapacity(512, 4); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 512 * 4); /* * inner value capacity we pass to data vector is 512 * 0.1 => 51 * For an int vector this is 204 bytes of memory for data buffer * and 7 bytes for validity buffer. * and with power of 2 allocation, we allocate 256 bytes and 8 bytes * for the data buffer and validity buffer of the inner vector. Thus * value capacity of inner vector is 64 */ vector.setInitialCapacity(512, 0.1); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 51); /* * inner value capacity we pass to data vector is 512 * 0.01 => 5 * For an int vector this is 20 bytes of memory for data buffer * and 1 byte for validity buffer. * and with power of 2 allocation, we allocate 32 bytes and 1 bytes * for the data buffer and validity buffer of the inner vector. Thus * value capacity of inner vector is 8 */ vector.setInitialCapacity(512, 0.01); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 5); /* * inner value capacity we pass to data vector is 5 * 0.1 => 0 * which is then rounded off to 1. So we pass value count as 1 * to the inner int vector. * the offset buffer of the list vector is allocated for 6 values * which is 24 bytes and then rounded off to 32 bytes (8 values) * the validity buffer of the list vector is allocated for 5 * values which is 1 byte. This is why value capacity of the list * vector is 7 as we take the min of validity buffer value capacity * and offset buffer value capacity. */ vector.setInitialCapacity(5, 0.1); vector.allocateNew(); assertEquals(7, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 1); } } @Test public void testClearAndReuse() { try (final ListVector vector = ListVector.empty("list", allocator)) { BigIntVector bigIntVector = (BigIntVector) vector.addOrGetVector(FieldType.nullable(MinorType.BIGINT.getType())).getVector(); vector.setInitialCapacity(10); vector.allocateNew(); vector.startNewValue(0); bigIntVector.setSafe(0, 7); vector.endValue(0, 1); vector.startNewValue(1); bigIntVector.setSafe(1, 8); vector.endValue(1, 1); vector.setValueCount(2); Object result = vector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(Long.valueOf(7), resultSet.get(0)); result = vector.getObject(1); resultSet = (ArrayList) result; assertEquals(Long.valueOf(8), resultSet.get(0)); // Clear and release the buffers to trigger a realloc when adding next value vector.clear(); // The list vector should reuse a buffer when reallocating the offset buffer vector.startNewValue(0); bigIntVector.setSafe(0, 7); vector.endValue(0, 1); vector.startNewValue(1); bigIntVector.setSafe(1, 8); vector.endValue(1, 1); vector.setValueCount(2); result = vector.getObject(0); resultSet = (ArrayList) result; assertEquals(Long.valueOf(7), resultSet.get(0)); result = vector.getObject(1); resultSet = (ArrayList) result; assertEquals(Long.valueOf(8), resultSet.get(0)); } } @Test public void testWriterGetField() { try (final ListVector vector = ListVector.empty("list", allocator)) { UnionListWriter writer = vector.getWriter(); writer.allocate(); // set some values writer.startList(); writer.integer().writeInt(1); writer.integer().writeInt(2); writer.endList(); vector.setValueCount(2); Field expectedDataField = new Field( BaseRepeatedValueVector.DATA_VECTOR_NAME, FieldType.nullable(new ArrowType.Int(32, true)), null); Field expectedField = new Field( vector.getName(), FieldType.nullable(ArrowType.List.INSTANCE), Arrays.asList(expectedDataField)); assertEquals(expectedField, writer.getField()); } } @Test public void testWriterGetTimestampMilliTZField() { try (final ListVector vector = ListVector.empty("list", allocator)) { org.apache.arrow.vector.complex.writer.FieldWriter writer = vector.getWriter(); writer.allocate(); writer.startList(); writer.timeStampMilliTZ().writeTimeStampMilliTZ(1000L); writer.timeStampMilliTZ().writeTimeStampMilliTZ(2000L); writer.endList(); vector.setValueCount(1); Field expectedDataField = new Field( BaseRepeatedValueVector.DATA_VECTOR_NAME, FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), null); Field expectedField = new Field( vector.getName(), FieldType.nullable(ArrowType.List.INSTANCE), Arrays.asList(expectedDataField)); assertEquals(expectedField, writer.getField()); } } @Test public void testWriterUsingHolderGetTimestampMilliTZField() { try (final ListVector vector = ListVector.empty("list", allocator)) { org.apache.arrow.vector.complex.writer.FieldWriter 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(ArrowType.List.INSTANCE), Arrays.asList(expectedDataField)); assertEquals(expectedField, writer.getField()); } } @Test public void testWriterGetDurationField() { try (final ListVector vector = ListVector.empty("list", allocator)) { org.apache.arrow.vector.complex.writer.FieldWriter 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(ArrowType.List.INSTANCE), Arrays.asList(expectedDataField)); assertEquals(expectedField, writer.getField()); } } @Test public void testWriterGetFixedSizeBinaryField() throws Exception { // Adapted from: TestComplexWriter.java:fixedSizeBinaryWriters // test values int numValues = 10; int byteWidth = 9; byte[][] values = new byte[numValues][byteWidth]; for (int i = 0; i < numValues; i++) { for (int j = 0; j < byteWidth; j++) { values[i][j] = ((byte) i); } } ArrowBuf[] bufs = new ArrowBuf[numValues]; for (int i = 0; i < numValues; i++) { bufs[i] = allocator.buffer(byteWidth); bufs[i].setBytes(0, values[i]); } try (final ListVector vector = ListVector.empty("list", allocator)) { org.apache.arrow.vector.complex.writer.FieldWriter writer = vector.getWriter(); writer.allocate(); FixedSizeBinaryHolder binHolder = new FixedSizeBinaryHolder(); binHolder.byteWidth = byteWidth; writer.startList(); for (int i = 0; i < numValues; i++) { binHolder.buffer = bufs[i]; writer.fixedSizeBinary().write(binHolder); } // Writing with a different byteWidth should throw // Note just reusing the last buffer value since that won't matter here anyway binHolder.byteWidth = 3; IllegalArgumentException ex = assertThrows( IllegalArgumentException.class, () -> writer.fixedSizeBinary().write(binHolder)); assertEquals("holder.byteWidth: 3 not equal to vector byteWidth: 9", ex.getMessage()); writer.endList(); vector.setValueCount(1); Field expectedDataField = new Field( BaseRepeatedValueVector.DATA_VECTOR_NAME, FieldType.nullable(new ArrowType.FixedSizeBinary(byteWidth)), null); Field expectedField = new Field( vector.getName(), FieldType.nullable(ArrowType.List.INSTANCE), Arrays.asList(expectedDataField)); assertEquals(expectedField, writer.getField()); } AutoCloseables.close(bufs); } @Test public void testClose() throws Exception { try (final ListVector vector = ListVector.empty("list", allocator)) { UnionListWriter writer = vector.getWriter(); writer.allocate(); // set some values writer.startList(); writer.integer().writeInt(1); writer.integer().writeInt(2); writer.endList(); vector.setValueCount(2); assertTrue(vector.getBufferSize() > 0); assertTrue(vector.getDataVector().getBufferSize() > 0); writer.close(); assertEquals(0, vector.getBufferSize()); assertEquals(0, vector.getDataVector().getBufferSize()); } } @Test public void testGetBufferSizeFor() { try (final ListVector vector = ListVector.empty("list", allocator)) { UnionListWriter writer = vector.getWriter(); writer.allocate(); // set some values writeIntValues(writer, new int[] {1, 2}); writeIntValues(writer, new int[] {3, 4}); writeIntValues(writer, new int[] {5, 6}); writeIntValues(writer, new int[] {7, 8, 9, 10}); writeIntValues(writer, new int[] {11, 12, 13, 14}); writer.setValueCount(5); IntVector dataVector = (IntVector) vector.getDataVector(); int[] indices = new int[] {0, 2, 4, 6, 10, 14}; for (int valueCount = 1; valueCount <= 5; valueCount++) { int validityBufferSize = getValidityBufferSizeFromCount(valueCount); int offsetBufferSize = (valueCount + 1) * BaseRepeatedValueVector.OFFSET_WIDTH; int expectedSize = validityBufferSize + offsetBufferSize + dataVector.getBufferSizeFor(indices[valueCount]); assertEquals(expectedSize, vector.getBufferSizeFor(valueCount)); } } } @Test public void testIsEmpty() { try (final ListVector vector = ListVector.empty("list", allocator)) { UnionListWriter writer = vector.getWriter(); writer.allocate(); // set values [1,2], null, [], [5,6] writeIntValues(writer, new int[] {1, 2}); writer.setPosition(2); writeIntValues(writer, new int[] {}); writeIntValues(writer, new int[] {5, 6}); writer.setValueCount(4); assertFalse(vector.isEmpty(0)); assertTrue(vector.isNull(1)); assertTrue(vector.isEmpty(1)); assertFalse(vector.isNull(2)); assertTrue(vector.isEmpty(2)); assertFalse(vector.isEmpty(3)); } } @Test public void testTotalCapacity() { final FieldType type = FieldType.nullable(MinorType.INT.getType()); try (final ListVector vector = new ListVector("list", allocator, type, null)) { // Force the child vector to be allocated based on the type // (this is a bad API: we have to track and repeat the type twice) vector.addOrGetVector(type); // Specify the allocation size but do not actually allocate vector.setInitialTotalCapacity(10, 100); // Finally actually do the allocation vector.allocateNewSafe(); // Note: allocator rounds up and can be greater than the requested allocation. assertTrue(vector.getValueCapacity() >= 10); assertTrue(vector.getDataVector().getValueCapacity() >= 100); } } @Test public void testGetTransferPairWithField() { try (ListVector fromVector = ListVector.empty("input", allocator)) { UnionListWriter writer = fromVector.getWriter(); writer.allocate(); writer.setPosition(0); // optional writer.startList(); writer.bigInt().writeBigInt(1); writer.bigInt().writeBigInt(2); writer.bigInt().writeBigInt(3); writer.endList(); writer.setValueCount(1); final TransferPair transferPair = fromVector.getTransferPair(fromVector.getField(), allocator); final ListVector toVector = (ListVector) transferPair.getTo(); // Field inside a new vector created by reusing a field should be the same in memory as the // original field. assertSame(toVector.getField(), fromVector.getField()); } } @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) { writer.integer().writeInt(v); } writer.endList(); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestListViewVector.java000066400000000000000000002547011515451615200322000ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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; import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; 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.BaseRepeatedValueViewVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.ListViewVector; import org.apache.arrow.vector.complex.impl.UnionListViewWriter; import org.apache.arrow.vector.holders.DurationHolder; import org.apache.arrow.vector.holders.TimeStampMilliTZHolder; 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.TransferPair; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestListViewVector { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testBasicListViewVector() { try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", allocator)) { UnionListViewWriter listViewWriter = listViewVector.getWriter(); /* allocate memory */ listViewWriter.allocate(); /* write the first list at index 0 */ listViewWriter.setPosition(0); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(12); listViewWriter.bigInt().writeBigInt(-7); listViewWriter.bigInt().writeBigInt(25); listViewWriter.endListView(); /* the second list at index 1 is null (we are not setting any)*/ /* write the third list at index 2 */ listViewWriter.setPosition(2); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(0); listViewWriter.bigInt().writeBigInt(-127); listViewWriter.bigInt().writeBigInt(127); listViewWriter.bigInt().writeBigInt(50); listViewWriter.endListView(); /* write the fourth list at index 3 (empty list) */ listViewWriter.setPosition(3); listViewWriter.startListView(); listViewWriter.endListView(); /* write the fifth list at index 4 */ listViewWriter.setPosition(4); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(1); listViewWriter.bigInt().writeBigInt(2); listViewWriter.bigInt().writeBigInt(3); listViewWriter.bigInt().writeBigInt(4); listViewWriter.endListView(); listViewWriter.setValueCount(5); // check value count assertEquals(5, listViewVector.getValueCount()); /* get vector at index 0 -- the value is a BigIntVector*/ final ArrowBuf offSetBuffer = listViewVector.getOffsetBuffer(); final ArrowBuf sizeBuffer = listViewVector.getSizeBuffer(); final FieldVector dataVec = listViewVector.getDataVector(); // check offset buffer assertEquals(0, offSetBuffer.getInt(0 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offSetBuffer.getInt(1 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(2 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(7, offSetBuffer.getInt(3 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(7, offSetBuffer.getInt(4 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); // check size buffer assertEquals(3, sizeBuffer.getInt(0 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(1 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(4, sizeBuffer.getInt(2 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(3 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(4, sizeBuffer.getInt(4 * BaseRepeatedValueViewVector.SIZE_WIDTH)); // check data vector assertEquals(12, ((BigIntVector) dataVec).get(0)); assertEquals(-7, ((BigIntVector) dataVec).get(1)); assertEquals(25, ((BigIntVector) dataVec).get(2)); assertEquals(0, ((BigIntVector) dataVec).get(3)); assertEquals(-127, ((BigIntVector) dataVec).get(4)); assertEquals(127, ((BigIntVector) dataVec).get(5)); assertEquals(50, ((BigIntVector) dataVec).get(6)); assertEquals(1, ((BigIntVector) dataVec).get(7)); assertEquals(2, ((BigIntVector) dataVec).get(8)); assertEquals(3, ((BigIntVector) dataVec).get(9)); assertEquals(4, ((BigIntVector) dataVec).get(10)); listViewVector.validate(); } } @Test public void testImplicitNullVectors() { try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", allocator)) { UnionListViewWriter listViewWriter = listViewVector.getWriter(); /* allocate memory */ listViewWriter.allocate(); final ArrowBuf offSetBuffer = listViewVector.getOffsetBuffer(); final ArrowBuf sizeBuffer = listViewVector.getSizeBuffer(); /* write the first list at index 0 */ listViewWriter.setPosition(0); listViewWriter.startList(); listViewWriter.bigInt().writeBigInt(12); listViewWriter.bigInt().writeBigInt(-7); listViewWriter.bigInt().writeBigInt(25); listViewWriter.endListView(); int offSet0 = offSetBuffer.getInt(0 * BaseRepeatedValueViewVector.OFFSET_WIDTH); int size0 = sizeBuffer.getInt(0 * BaseRepeatedValueViewVector.SIZE_WIDTH); // after the first list is written, // the initial offset must be 0, // the size must be 3 (as there are 3 elements in the array), // the lastSet must be 0 since, the first list is written at index 0. assertEquals(0, offSet0); assertEquals(3, size0); listViewWriter.setPosition(5); listViewWriter.startListView(); // writing the 6th list at index 5, // and the list items from index 1 through 4 are not populated. // but since there is a gap between the 0th and 5th list, in terms // of buffer allocation, the offset and size buffers must be updated // to reflect the implicit null vectors. for (int i = 1; i < 5; i++) { int offSet = offSetBuffer.getInt(i * BaseRepeatedValueViewVector.OFFSET_WIDTH); int size = sizeBuffer.getInt(i * BaseRepeatedValueViewVector.SIZE_WIDTH); // Since the list is not written, the offset and size must equal to child vector's size // i.e., 3, and size should be 0 as the list is not written. // And the last set value is the value currently being written, which is 5. assertEquals(0, offSet); assertEquals(0, size); } listViewWriter.bigInt().writeBigInt(12); listViewWriter.bigInt().writeBigInt(25); listViewWriter.endListView(); int offSet5 = offSetBuffer.getInt(5 * BaseRepeatedValueViewVector.OFFSET_WIDTH); int size5 = sizeBuffer.getInt(5 * BaseRepeatedValueViewVector.SIZE_WIDTH); assertEquals(3, offSet5); assertEquals(2, size5); listViewWriter.setPosition(10); listViewWriter.startListView(); // writing the 11th list at index 10, // and the list items from index 6 through 10 are not populated. // but since there is a gap between the 5th and 11th list, in terms // of buffer allocation, the offset and size buffers must be updated // to reflect the implicit null vectors. for (int i = 6; i < 10; i++) { int offSet = offSetBuffer.getInt(i * BaseRepeatedValueViewVector.OFFSET_WIDTH); int size = sizeBuffer.getInt(i * BaseRepeatedValueViewVector.SIZE_WIDTH); // Since the list is not written, the offset and size must equal to 0 // and size should be 0 as the list is not written. // And the last set value is the value currently being written, which is 10. assertEquals(0, offSet); assertEquals(0, size); } listViewWriter.bigInt().writeBigInt(12); listViewWriter.endListView(); int offSet11 = offSetBuffer.getInt(10 * BaseRepeatedValueViewVector.OFFSET_WIDTH); int size11 = sizeBuffer.getInt(10 * BaseRepeatedValueViewVector.SIZE_WIDTH); assertEquals(5, offSet11); assertEquals(1, size11); listViewVector.setValueCount(11); listViewVector.validate(); } } @Test public void testNestedListViewVector() { try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", allocator)) { UnionListViewWriter listViewWriter = listViewVector.getWriter(); /* allocate memory */ listViewWriter.allocate(); /* the dataVector that backs a listVector will also be a * listVector for this test. */ /* write one or more inner lists at index 0 */ listViewWriter.setPosition(0); listViewWriter.startListView(); listViewWriter.listView().startListView(); listViewWriter.listView().bigInt().writeBigInt(50); listViewWriter.listView().bigInt().writeBigInt(100); listViewWriter.listView().bigInt().writeBigInt(200); listViewWriter.listView().endListView(); listViewWriter.listView().startListView(); listViewWriter.listView().bigInt().writeBigInt(75); listViewWriter.listView().bigInt().writeBigInt(125); listViewWriter.listView().bigInt().writeBigInt(150); listViewWriter.listView().bigInt().writeBigInt(175); listViewWriter.listView().endListView(); listViewWriter.endListView(); /* write one or more inner lists at index 1 */ listViewWriter.setPosition(1); listViewWriter.startListView(); listViewWriter.listView().startListView(); listViewWriter.listView().bigInt().writeBigInt(10); listViewWriter.listView().endListView(); listViewWriter.listView().startListView(); listViewWriter.listView().bigInt().writeBigInt(15); listViewWriter.listView().bigInt().writeBigInt(20); listViewWriter.listView().endListView(); listViewWriter.listView().startListView(); listViewWriter.listView().bigInt().writeBigInt(25); listViewWriter.listView().bigInt().writeBigInt(30); listViewWriter.listView().bigInt().writeBigInt(35); listViewWriter.listView().endListView(); listViewWriter.endListView(); listViewVector.setValueCount(2); // [[[50,100,200],[75,125,150,175]], [[10],[15,20],[25,30,35]]] assertEquals(2, listViewVector.getValueCount()); /* get listViewVector value at index 0 -- the value itself is a listViewVector */ Object result = listViewVector.getObject(0); ArrayList> resultSet = (ArrayList>) result; ArrayList list; assertEquals(2, resultSet.size()); /* 2 inner lists at index 0 */ assertEquals(3, resultSet.get(0).size()); /* size of the first inner list */ assertEquals(4, resultSet.get(1).size()); /* size of the second inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(50), list.get(0)); assertEquals(Long.valueOf(100), list.get(1)); assertEquals(Long.valueOf(200), list.get(2)); list = resultSet.get(1); assertEquals(Long.valueOf(75), list.get(0)); assertEquals(Long.valueOf(125), list.get(1)); assertEquals(Long.valueOf(150), list.get(2)); assertEquals(Long.valueOf(175), list.get(3)); /* get listViewVector value at index 1 -- the value itself is a listViewVector */ result = listViewVector.getObject(1); resultSet = (ArrayList>) result; assertEquals(3, resultSet.size()); /* 3 inner lists at index 1 */ assertEquals(1, resultSet.get(0).size()); /* size of the first inner list */ assertEquals(2, resultSet.get(1).size()); /* size of the second inner list */ assertEquals(3, resultSet.get(2).size()); /* size of the third inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(10), list.get(0)); list = resultSet.get(1); assertEquals(Long.valueOf(15), list.get(0)); assertEquals(Long.valueOf(20), list.get(1)); list = resultSet.get(2); assertEquals(Long.valueOf(25), list.get(0)); assertEquals(Long.valueOf(30), list.get(1)); assertEquals(Long.valueOf(35), list.get(2)); /* check underlying bitVector */ assertFalse(listViewVector.isNull(0)); assertFalse(listViewVector.isNull(1)); final ArrowBuf offSetBuffer = listViewVector.getOffsetBuffer(); final ArrowBuf sizeBuffer = listViewVector.getSizeBuffer(); // check offset buffer assertEquals(0, offSetBuffer.getInt(0 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(2, offSetBuffer.getInt(1 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); // check size buffer assertEquals(2, sizeBuffer.getInt(0 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(3, sizeBuffer.getInt(1 * BaseRepeatedValueViewVector.SIZE_WIDTH)); listViewVector.validate(); } } @Test public void testNestedListVector() throws Exception { try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", allocator)) { MinorType listType = MinorType.LISTVIEW; MinorType scalarType = MinorType.BIGINT; listViewVector.addOrGetVector(FieldType.nullable(listType.getType())); ListViewVector innerList1 = (ListViewVector) listViewVector.getDataVector(); innerList1.addOrGetVector(FieldType.nullable(listType.getType())); ListViewVector innerList2 = (ListViewVector) innerList1.getDataVector(); innerList2.addOrGetVector(FieldType.nullable(listType.getType())); ListViewVector innerList3 = (ListViewVector) innerList2.getDataVector(); innerList3.addOrGetVector(FieldType.nullable(listType.getType())); ListViewVector innerList4 = (ListViewVector) innerList3.getDataVector(); innerList4.addOrGetVector(FieldType.nullable(listType.getType())); ListViewVector innerList5 = (ListViewVector) innerList4.getDataVector(); innerList5.addOrGetVector(FieldType.nullable(listType.getType())); ListViewVector innerList6 = (ListViewVector) innerList5.getDataVector(); innerList6.addOrGetVector(FieldType.nullable(scalarType.getType())); listViewVector.setInitialCapacity(128); listViewVector.validate(); } } private void setValuesInBuffer(int[] bufValues, ArrowBuf buffer, long bufWidth) { for (int i = 0; i < bufValues.length; i++) { buffer.setInt(i * bufWidth, bufValues[i]); } } /* * Setting up the buffers directly needs to be validated with the base method used in * the ListVector class where we use the approach of startListView(), * write to the child vector and endListView(). *

    * To support this, we have to consider the following scenarios; *

    * 1. Only using directly buffer-based inserts. * 2. Default list insertion followed by buffer-based inserts. * 3. Buffer-based inserts followed by default list insertion. */ /* Setting up buffers directly would require the following steps to be taken * 0. Allocate buffers in listViewVector by calling `allocateNew` method. * 1. Initialize the child vector using `initializeChildrenFromFields` method. * 2. Set values in the child vector. * 3. Set validity, offset and size buffers using `setValidity`, * `setOffset` and `setSize` methods. * 4. Set value count using `setValueCount` method. */ @Test public void testBasicListViewSet() { try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", 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(64, 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(); BigIntVector childVector = (BigIntVector) fieldVector; childVector.allocateNew(7); childVector.set(0, 12); childVector.set(1, -7); childVector.set(2, 25); childVector.set(3, 0); childVector.set(4, -127); childVector.set(5, 127); childVector.set(6, 50); childVector.setValueCount(7); // Set validity, offset and size buffers using `setValidity`, // `setOffset` and `setSize` methods. listViewVector.setOffset(0, 0); listViewVector.setOffset(1, 3); listViewVector.setOffset(2, 3); listViewVector.setOffset(3, 7); listViewVector.setSize(0, 3); listViewVector.setSize(1, 0); listViewVector.setSize(2, 4); listViewVector.setSize(3, 0); listViewVector.setValidity(0, 1); listViewVector.setValidity(1, 0); listViewVector.setValidity(2, 1); listViewVector.setValidity(3, 1); // Set value count using `setValueCount` method. listViewVector.setValueCount(4); final ArrowBuf offSetBuffer = listViewVector.getOffsetBuffer(); final ArrowBuf sizeBuffer = listViewVector.getSizeBuffer(); // check offset buffer assertEquals(0, offSetBuffer.getInt(0 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(1 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(2 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(7, offSetBuffer.getInt(3 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); // check size buffer assertEquals(3, sizeBuffer.getInt(0 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(1 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(4, sizeBuffer.getInt(2 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(3 * BaseRepeatedValueViewVector.SIZE_WIDTH)); // check values assertEquals(12, ((BigIntVector) listViewVector.getDataVector()).get(0)); assertEquals(-7, ((BigIntVector) listViewVector.getDataVector()).get(1)); assertEquals(25, ((BigIntVector) listViewVector.getDataVector()).get(2)); assertEquals(0, ((BigIntVector) listViewVector.getDataVector()).get(3)); assertEquals(-127, ((BigIntVector) listViewVector.getDataVector()).get(4)); assertEquals(127, ((BigIntVector) listViewVector.getDataVector()).get(5)); assertEquals(50, ((BigIntVector) listViewVector.getDataVector()).get(6)); listViewVector.validate(); } } @Test public void testBasicListViewSetNested() { // Expected listview // [[[50,100,200],[75,125,150,175]],[[10],[15,20],[25,30,35]]] // Setting child vector // [[50,100,200],[75,125,150,175],[10],[15,20],[25,30,35]] try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", 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.ListView(), null, null); FieldType childFieldType = new FieldType(true, new ArrowType.Int(64, true), null, null); Field childField = new Field("child-vector", childFieldType, null); List children = new ArrayList<>(); children.add(childField); Field field = new Field("child-vector", fieldType, children); listViewVector.initializeChildrenFromFields(Collections.singletonList(field)); // Set values in the child vector. FieldVector fieldVector = listViewVector.getDataVector(); fieldVector.clear(); ListViewVector childVector = (ListViewVector) fieldVector; UnionListViewWriter listViewWriter = childVector.getWriter(); listViewWriter.allocate(); listViewWriter.setPosition(0); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(50); listViewWriter.bigInt().writeBigInt(100); listViewWriter.bigInt().writeBigInt(200); listViewWriter.endListView(); listViewWriter.setPosition(1); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(75); listViewWriter.bigInt().writeBigInt(125); listViewWriter.bigInt().writeBigInt(150); listViewWriter.bigInt().writeBigInt(175); listViewWriter.endListView(); listViewWriter.setPosition(2); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(10); listViewWriter.endListView(); listViewWriter.startListView(); listViewWriter.setPosition(3); listViewWriter.bigInt().writeBigInt(15); listViewWriter.bigInt().writeBigInt(20); listViewWriter.endListView(); listViewWriter.startListView(); listViewWriter.setPosition(4); listViewWriter.bigInt().writeBigInt(25); listViewWriter.bigInt().writeBigInt(30); listViewWriter.bigInt().writeBigInt(35); listViewWriter.endListView(); childVector.setValueCount(5); // Set validity, offset and size buffers using `setValidity`, // `setOffset` and `setSize` methods. listViewVector.setValidity(0, 1); listViewVector.setValidity(1, 1); listViewVector.setOffset(0, 0); listViewVector.setOffset(1, 2); listViewVector.setSize(0, 2); listViewVector.setSize(1, 3); // Set value count using `setValueCount` method. listViewVector.setValueCount(2); assertEquals(2, listViewVector.getValueCount()); /* get listViewVector value at index 0 -- the value itself is a listViewVector */ Object result = listViewVector.getObject(0); ArrayList> resultSet = (ArrayList>) result; ArrayList list; assertEquals(2, resultSet.size()); /* 2 inner lists at index 0 */ assertEquals(3, resultSet.get(0).size()); /* size of the first inner list */ assertEquals(4, resultSet.get(1).size()); /* size of the second inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(50), list.get(0)); assertEquals(Long.valueOf(100), list.get(1)); assertEquals(Long.valueOf(200), list.get(2)); list = resultSet.get(1); assertEquals(Long.valueOf(75), list.get(0)); assertEquals(Long.valueOf(125), list.get(1)); assertEquals(Long.valueOf(150), list.get(2)); assertEquals(Long.valueOf(175), list.get(3)); /* get listViewVector value at index 1 -- the value itself is a listViewVector */ result = listViewVector.getObject(1); resultSet = (ArrayList>) result; assertEquals(3, resultSet.size()); /* 3 inner lists at index 1 */ assertEquals(1, resultSet.get(0).size()); /* size of the first inner list */ assertEquals(2, resultSet.get(1).size()); /* size of the second inner list */ assertEquals(3, resultSet.get(2).size()); /* size of the third inner list */ list = resultSet.get(0); assertEquals(Long.valueOf(10), list.get(0)); list = resultSet.get(1); assertEquals(Long.valueOf(15), list.get(0)); assertEquals(Long.valueOf(20), list.get(1)); list = resultSet.get(2); assertEquals(Long.valueOf(25), list.get(0)); assertEquals(Long.valueOf(30), list.get(1)); assertEquals(Long.valueOf(35), list.get(2)); /* check underlying bitVector */ assertFalse(listViewVector.isNull(0)); assertFalse(listViewVector.isNull(1)); final ArrowBuf offSetBuffer = listViewVector.getOffsetBuffer(); final ArrowBuf sizeBuffer = listViewVector.getSizeBuffer(); // check offset buffer assertEquals(0, offSetBuffer.getInt(0 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(2, offSetBuffer.getInt(1 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); // check size buffer assertEquals(2, sizeBuffer.getInt(0 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(3, sizeBuffer.getInt(1 * BaseRepeatedValueViewVector.SIZE_WIDTH)); listViewVector.validate(); } } @Test public void testBasicListViewSetWithListViewWriter() { try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", 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(64, 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(); BigIntVector childVector = (BigIntVector) fieldVector; childVector.allocateNew(7); childVector.set(0, 12); childVector.set(1, -7); childVector.set(2, 25); childVector.set(3, 0); childVector.set(4, -127); childVector.set(5, 127); childVector.set(6, 50); 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.setOffset(0, 0); listViewVector.setOffset(1, 3); listViewVector.setOffset(2, 3); listViewVector.setOffset(3, 7); listViewVector.setSize(0, 3); listViewVector.setSize(1, 0); listViewVector.setSize(2, 4); listViewVector.setSize(3, 0); // Set value count using `setValueCount` method. listViewVector.setValueCount(4); final ArrowBuf offSetBuffer = listViewVector.getOffsetBuffer(); final ArrowBuf sizeBuffer = listViewVector.getSizeBuffer(); // check offset buffer assertEquals(0, offSetBuffer.getInt(0 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(1 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(2 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(7, offSetBuffer.getInt(3 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); // check size buffer assertEquals(3, sizeBuffer.getInt(0 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(1 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(4, sizeBuffer.getInt(2 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(3 * BaseRepeatedValueViewVector.SIZE_WIDTH)); // check values assertEquals(12, ((BigIntVector) listViewVector.getDataVector()).get(0)); assertEquals(-7, ((BigIntVector) listViewVector.getDataVector()).get(1)); assertEquals(25, ((BigIntVector) listViewVector.getDataVector()).get(2)); assertEquals(0, ((BigIntVector) listViewVector.getDataVector()).get(3)); assertEquals(-127, ((BigIntVector) listViewVector.getDataVector()).get(4)); assertEquals(127, ((BigIntVector) listViewVector.getDataVector()).get(5)); assertEquals(50, ((BigIntVector) listViewVector.getDataVector()).get(6)); UnionListViewWriter listViewWriter = listViewVector.getWriter(); listViewWriter.setPosition(4); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(121); listViewWriter.bigInt().writeBigInt(-71); listViewWriter.bigInt().writeBigInt(251); listViewWriter.endListView(); listViewVector.setValueCount(5); // check offset buffer assertEquals(0, offSetBuffer.getInt(0 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(1 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(2 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(7, offSetBuffer.getInt(3 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(7, offSetBuffer.getInt(4 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); // check size buffer assertEquals(3, sizeBuffer.getInt(0 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(1 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(4, sizeBuffer.getInt(2 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(3 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(3, sizeBuffer.getInt(4 * BaseRepeatedValueViewVector.SIZE_WIDTH)); // check values assertEquals(12, ((BigIntVector) listViewVector.getDataVector()).get(0)); assertEquals(-7, ((BigIntVector) listViewVector.getDataVector()).get(1)); assertEquals(25, ((BigIntVector) listViewVector.getDataVector()).get(2)); assertEquals(0, ((BigIntVector) listViewVector.getDataVector()).get(3)); assertEquals(-127, ((BigIntVector) listViewVector.getDataVector()).get(4)); assertEquals(127, ((BigIntVector) listViewVector.getDataVector()).get(5)); assertEquals(50, ((BigIntVector) listViewVector.getDataVector()).get(6)); assertEquals(121, ((BigIntVector) listViewVector.getDataVector()).get(7)); assertEquals(-71, ((BigIntVector) listViewVector.getDataVector()).get(8)); assertEquals(251, ((BigIntVector) listViewVector.getDataVector()).get(9)); listViewVector.validate(); } } @Test public void testGetBufferAddress() throws Exception { try (ListViewVector listViewVector = ListViewVector.empty("vector", allocator)) { UnionListViewWriter listViewWriter = listViewVector.getWriter(); boolean error = false; listViewWriter.allocate(); listViewWriter.setPosition(0); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(50); listViewWriter.bigInt().writeBigInt(100); listViewWriter.bigInt().writeBigInt(200); listViewWriter.endListView(); listViewWriter.setPosition(1); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(250); listViewWriter.bigInt().writeBigInt(300); listViewWriter.endListView(); listViewVector.setValueCount(2); /* check listVector contents */ Object result = listViewVector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Long.valueOf(50), resultSet.get(0)); assertEquals(Long.valueOf(100), resultSet.get(1)); assertEquals(Long.valueOf(200), resultSet.get(2)); result = listViewVector.getObject(1); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(250), resultSet.get(0)); assertEquals(Long.valueOf(300), resultSet.get(1)); List buffers = listViewVector.getFieldBuffers(); long bitAddress = listViewVector.getValidityBufferAddress(); long offsetAddress = listViewVector.getOffsetBufferAddress(); long sizeAddress = listViewVector.getSizeBufferAddress(); try { listViewVector.getDataBufferAddress(); } catch (UnsupportedOperationException ue) { error = true; } finally { assertTrue(error); } assertEquals(3, buffers.size()); assertEquals(bitAddress, buffers.get(0).memoryAddress()); assertEquals(offsetAddress, buffers.get(1).memoryAddress()); assertEquals(sizeAddress, buffers.get(2).memoryAddress()); /* (3+2)/2 */ assertEquals(2.5, listViewVector.getDensity(), 0); listViewVector.validate(); } } @Test public void testConsistentChildName() throws Exception { try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", allocator)) { String emptyListStr = listViewVector.getField().toString(); assertTrue(emptyListStr.contains(ListVector.DATA_VECTOR_NAME)); listViewVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); String emptyVectorStr = listViewVector.getField().toString(); assertTrue(emptyVectorStr.contains(ListVector.DATA_VECTOR_NAME)); } } @Test public void testSetInitialCapacity() { try (final ListViewVector vector = ListViewVector.empty("", allocator)) { vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); vector.setInitialCapacity(512); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 512); vector.setInitialCapacity(512, 4); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 512 * 4); vector.setInitialCapacity(512, 0.1); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 51); vector.setInitialCapacity(512, 0.01); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 5); vector.setInitialCapacity(5, 0.1); vector.allocateNew(); assertEquals(8, vector.getValueCapacity()); assertTrue(vector.getDataVector().getValueCapacity() >= 1); vector.validate(); } } @Test public void testClearAndReuse() { try (final ListViewVector vector = ListViewVector.empty("listview", allocator)) { BigIntVector bigIntVector = (BigIntVector) vector.addOrGetVector(FieldType.nullable(MinorType.BIGINT.getType())).getVector(); vector.setInitialCapacity(10); vector.allocateNew(); vector.startNewValue(0); bigIntVector.setSafe(0, 7); vector.endValue(0, 1); vector.startNewValue(1); bigIntVector.setSafe(1, 8); vector.endValue(1, 1); vector.setValueCount(2); Object result = vector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(Long.valueOf(7), resultSet.get(0)); result = vector.getObject(1); resultSet = (ArrayList) result; assertEquals(Long.valueOf(8), resultSet.get(0)); // Clear and release the buffers to trigger a realloc when adding next value vector.clear(); // The list vector should reuse a buffer when reallocating the offset buffer vector.startNewValue(0); bigIntVector.setSafe(0, 7); vector.endValue(0, 1); vector.startNewValue(1); bigIntVector.setSafe(1, 8); vector.endValue(1, 1); vector.setValueCount(2); result = vector.getObject(0); resultSet = (ArrayList) result; assertEquals(Long.valueOf(7), resultSet.get(0)); result = vector.getObject(1); resultSet = (ArrayList) result; assertEquals(Long.valueOf(8), resultSet.get(0)); vector.validate(); } } @Test public void testWriterGetField() { // adopted from ListVector test cases try (final ListViewVector vector = ListViewVector.empty("listview", allocator)) { UnionListViewWriter writer = vector.getWriter(); writer.allocate(); // set some values writer.startListView(); writer.integer().writeInt(1); writer.integer().writeInt(2); writer.endListView(); vector.setValueCount(2); Field expectedDataField = new Field( BaseRepeatedValueVector.DATA_VECTOR_NAME, FieldType.nullable(new ArrowType.Int(32, true)), null); Field expectedField = new Field( vector.getName(), FieldType.nullable(ArrowType.ListView.INSTANCE), Arrays.asList(expectedDataField)); assertEquals(expectedField, writer.getField()); vector.validate(); } } @Test public void testWriterUsingHolderGetTimestampMilliTZField() { // adopted from ListVector test cases try (final ListViewVector vector = ListViewVector.empty("listview", allocator)) { org.apache.arrow.vector.complex.writer.FieldWriter writer = vector.getWriter(); writer.allocate(); TimeStampMilliTZHolder holder = new TimeStampMilliTZHolder(); holder.timezone = "SomeFakeTimeZone"; writer.startListView(); 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.endListView(); 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(ArrowType.ListView.INSTANCE), Arrays.asList(expectedDataField)); assertEquals(expectedField, writer.getField()); vector.validate(); } } @Test public void testWriterGetDurationField() { // adopted from ListVector test cases try (final ListViewVector vector = ListViewVector.empty("listview", allocator)) { org.apache.arrow.vector.complex.writer.FieldWriter writer = vector.getWriter(); writer.allocate(); DurationHolder durationHolder = new DurationHolder(); durationHolder.unit = TimeUnit.MILLISECOND; writer.startListView(); 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.endListView(); 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(ArrowType.ListView.INSTANCE), Arrays.asList(expectedDataField)); assertEquals(expectedField, writer.getField()); vector.validate(); } } @Test public void testClose() throws Exception { try (final ListViewVector vector = ListViewVector.empty("listview", allocator)) { UnionListViewWriter writer = vector.getWriter(); writer.allocate(); // set some values writer.startListView(); writer.integer().writeInt(1); writer.integer().writeInt(2); writer.endListView(); vector.setValueCount(2); assertTrue(vector.getBufferSize() > 0); assertTrue(vector.getDataVector().getBufferSize() > 0); writer.close(); assertEquals(0, vector.getBufferSize()); assertEquals(0, vector.getDataVector().getBufferSize()); vector.validate(); } } @Test public void testGetBufferSizeFor() { try (final ListViewVector vector = ListViewVector.empty("listview", allocator)) { UnionListViewWriter writer = vector.getWriter(); writer.allocate(); // set some values writeIntValues(writer, new int[] {1, 2}); writeIntValues(writer, new int[] {3, 4}); writeIntValues(writer, new int[] {5, 6}); writeIntValues(writer, new int[] {7, 8, 9, 10}); writeIntValues(writer, new int[] {11, 12, 13, 14}); writer.setValueCount(5); IntVector dataVector = (IntVector) vector.getDataVector(); int[] indices = new int[] {0, 2, 4, 6, 10, 14}; for (int valueCount = 1; valueCount <= 5; valueCount++) { int validityBufferSize = getValidityBufferSizeFromCount(valueCount); int offsetBufferSize = valueCount * BaseRepeatedValueViewVector.OFFSET_WIDTH; int sizeBufferSize = valueCount * BaseRepeatedValueViewVector.SIZE_WIDTH; int expectedSize = validityBufferSize + offsetBufferSize + sizeBufferSize + dataVector.getBufferSizeFor(indices[valueCount]); assertEquals(expectedSize, vector.getBufferSizeFor(valueCount)); } vector.validate(); } } @Test public void testIsEmpty() { try (final ListViewVector vector = ListViewVector.empty("listview", allocator)) { UnionListViewWriter writer = vector.getWriter(); writer.allocate(); // set values [1,2], null, [], [5,6] writeIntValues(writer, new int[] {1, 2}); writer.setPosition(2); writeIntValues(writer, new int[] {}); writeIntValues(writer, new int[] {5, 6}); writer.setValueCount(4); assertFalse(vector.isEmpty(0)); assertTrue(vector.isNull(1)); assertTrue(vector.isEmpty(1)); assertFalse(vector.isNull(2)); assertTrue(vector.isEmpty(2)); assertFalse(vector.isEmpty(3)); vector.validate(); } } @Test public void testTotalCapacity() { // adopted from ListVector test cases final FieldType type = FieldType.nullable(MinorType.INT.getType()); try (final ListViewVector vector = new ListViewVector("listview", allocator, type, null)) { // Force the child vector to be allocated based on the type // (this is a bad API: we have to track and repeat the type twice) vector.addOrGetVector(type); // Specify the allocation size but do not allocate vector.setInitialTotalCapacity(10, 100); // Finally, actually do the allocation vector.allocateNewSafe(); // Note: allocator rounds up and can be greater than the requested allocation. assertTrue(vector.getValueCapacity() >= 10); assertTrue(vector.getDataVector().getValueCapacity() >= 100); vector.validate(); } } @Test public void testSetNull1() { try (ListViewVector vector = ListViewVector.empty("listview", allocator)) { UnionListViewWriter writer = vector.getWriter(); writer.allocate(); writer.setPosition(0); writer.startListView(); writer.bigInt().writeBigInt(10); writer.bigInt().writeBigInt(20); writer.endListView(); vector.setNull(1); writer.setPosition(2); writer.startListView(); writer.bigInt().writeBigInt(30); writer.bigInt().writeBigInt(40); writer.endListView(); vector.setNull(3); vector.setNull(4); writer.setPosition(5); writer.startListView(); writer.bigInt().writeBigInt(50); writer.bigInt().writeBigInt(60); writer.endListView(); vector.setValueCount(6); assertFalse(vector.isNull(0)); assertTrue(vector.isNull(1)); assertFalse(vector.isNull(2)); assertTrue(vector.isNull(3)); assertTrue(vector.isNull(4)); assertFalse(vector.isNull(5)); // validate buffers final ArrowBuf validityBuffer = vector.getValidityBuffer(); final ArrowBuf offsetBuffer = vector.getOffsetBuffer(); final ArrowBuf sizeBuffer = vector.getSizeBuffer(); assertEquals(1, BitVectorHelper.get(validityBuffer, 0)); assertEquals(0, BitVectorHelper.get(validityBuffer, 1)); assertEquals(1, BitVectorHelper.get(validityBuffer, 2)); assertEquals(0, BitVectorHelper.get(validityBuffer, 3)); assertEquals(0, BitVectorHelper.get(validityBuffer, 4)); assertEquals(1, BitVectorHelper.get(validityBuffer, 5)); assertEquals(0, offsetBuffer.getInt(0 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(1 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(2, offsetBuffer.getInt(2 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(3 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(4 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(4, offsetBuffer.getInt(5 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(2, sizeBuffer.getInt(0 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(1 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(2, sizeBuffer.getInt(2 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(3 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(4 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(2, sizeBuffer.getInt(5 * BaseRepeatedValueViewVector.SIZE_WIDTH)); // validate values Object result = vector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(10), resultSet.get(0)); assertEquals(Long.valueOf(20), resultSet.get(1)); result = vector.getObject(2); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(30), resultSet.get(0)); assertEquals(Long.valueOf(40), resultSet.get(1)); result = vector.getObject(5); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(50), resultSet.get(0)); assertEquals(Long.valueOf(60), resultSet.get(1)); vector.validate(); } } @Test public void testSetNull2() { try (ListViewVector vector = ListViewVector.empty("listview", allocator)) { // validate setting nulls first and then writing values UnionListViewWriter writer = vector.getWriter(); writer.allocate(); vector.setNull(0); vector.setNull(2); vector.setNull(4); writer.setPosition(1); writer.startListView(); writer.bigInt().writeBigInt(10); writer.bigInt().writeBigInt(20); writer.bigInt().writeBigInt(30); writer.endListView(); writer.setPosition(3); writer.startListView(); writer.bigInt().writeBigInt(40); writer.bigInt().writeBigInt(50); writer.endListView(); writer.setPosition(5); writer.startListView(); writer.bigInt().writeBigInt(60); writer.bigInt().writeBigInt(70); writer.bigInt().writeBigInt(80); writer.endListView(); vector.setValueCount(6); assertTrue(vector.isNull(0)); assertFalse(vector.isNull(1)); assertTrue(vector.isNull(2)); assertFalse(vector.isNull(3)); assertTrue(vector.isNull(4)); assertFalse(vector.isNull(5)); // validate buffers final ArrowBuf validityBuffer = vector.getValidityBuffer(); final ArrowBuf offsetBuffer = vector.getOffsetBuffer(); final ArrowBuf sizeBuffer = vector.getSizeBuffer(); assertEquals(0, BitVectorHelper.get(validityBuffer, 0)); assertEquals(1, BitVectorHelper.get(validityBuffer, 1)); assertEquals(0, BitVectorHelper.get(validityBuffer, 2)); assertEquals(1, BitVectorHelper.get(validityBuffer, 3)); assertEquals(0, BitVectorHelper.get(validityBuffer, 4)); assertEquals(1, BitVectorHelper.get(validityBuffer, 5)); assertEquals(0, offsetBuffer.getInt(0 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(1 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(2 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offsetBuffer.getInt(3 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(4 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(5, offsetBuffer.getInt(5 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, sizeBuffer.getInt(0 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(3, sizeBuffer.getInt(1 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(2 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(2, sizeBuffer.getInt(3 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(4 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(3, sizeBuffer.getInt(5 * BaseRepeatedValueViewVector.SIZE_WIDTH)); // validate values Object result = vector.getObject(1); ArrayList resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Long.valueOf(10), resultSet.get(0)); assertEquals(Long.valueOf(20), resultSet.get(1)); assertEquals(Long.valueOf(30), resultSet.get(2)); result = vector.getObject(3); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(40), resultSet.get(0)); assertEquals(Long.valueOf(50), resultSet.get(1)); result = vector.getObject(5); resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Long.valueOf(60), resultSet.get(0)); assertEquals(Long.valueOf(70), resultSet.get(1)); assertEquals(Long.valueOf(80), resultSet.get(2)); vector.validate(); } } @Test public void testSetNull3() { try (ListViewVector vector = ListViewVector.empty("listview", allocator)) { // validate setting values first and then writing nulls UnionListViewWriter writer = vector.getWriter(); writer.allocate(); writer.setPosition(1); writer.startListView(); writer.bigInt().writeBigInt(10); writer.bigInt().writeBigInt(20); writer.bigInt().writeBigInt(30); writer.endListView(); writer.setPosition(3); writer.startListView(); writer.bigInt().writeBigInt(40); writer.bigInt().writeBigInt(50); writer.endListView(); writer.setPosition(5); writer.startListView(); writer.bigInt().writeBigInt(60); writer.bigInt().writeBigInt(70); writer.bigInt().writeBigInt(80); writer.endListView(); vector.setNull(0); vector.setNull(2); vector.setNull(4); vector.setValueCount(6); assertTrue(vector.isNull(0)); assertFalse(vector.isNull(1)); assertTrue(vector.isNull(2)); assertFalse(vector.isNull(3)); assertTrue(vector.isNull(4)); assertFalse(vector.isNull(5)); // validate buffers final ArrowBuf validityBuffer = vector.getValidityBuffer(); final ArrowBuf offsetBuffer = vector.getOffsetBuffer(); final ArrowBuf sizeBuffer = vector.getSizeBuffer(); assertEquals(0, BitVectorHelper.get(validityBuffer, 0)); assertEquals(1, BitVectorHelper.get(validityBuffer, 1)); assertEquals(0, BitVectorHelper.get(validityBuffer, 2)); assertEquals(1, BitVectorHelper.get(validityBuffer, 3)); assertEquals(0, BitVectorHelper.get(validityBuffer, 4)); assertEquals(1, BitVectorHelper.get(validityBuffer, 5)); assertEquals(0, offsetBuffer.getInt(0 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(1 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(2 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offsetBuffer.getInt(3 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(4 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(5, offsetBuffer.getInt(5 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, sizeBuffer.getInt(0 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(3, sizeBuffer.getInt(1 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(2 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(2, sizeBuffer.getInt(3 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(4 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(3, sizeBuffer.getInt(5 * BaseRepeatedValueViewVector.SIZE_WIDTH)); // validate values Object result = vector.getObject(1); ArrayList resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Long.valueOf(10), resultSet.get(0)); assertEquals(Long.valueOf(20), resultSet.get(1)); assertEquals(Long.valueOf(30), resultSet.get(2)); result = vector.getObject(3); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(40), resultSet.get(0)); assertEquals(Long.valueOf(50), resultSet.get(1)); result = vector.getObject(5); resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Long.valueOf(60), resultSet.get(0)); assertEquals(Long.valueOf(70), resultSet.get(1)); assertEquals(Long.valueOf(80), resultSet.get(2)); vector.validate(); } } @Test public void testOverWrite1() { try (ListViewVector vector = ListViewVector.empty("listview", allocator)) { UnionListViewWriter writer = vector.getWriter(); writer.allocate(); writer.setPosition(0); writer.startListView(); writer.bigInt().writeBigInt(10); writer.bigInt().writeBigInt(20); writer.bigInt().writeBigInt(30); writer.endListView(); writer.setPosition(1); writer.startListView(); writer.bigInt().writeBigInt(40); writer.bigInt().writeBigInt(50); writer.endListView(); vector.setValueCount(2); writer.setPosition(0); writer.startListView(); writer.bigInt().writeBigInt(60); writer.bigInt().writeBigInt(70); writer.endListView(); writer.setPosition(1); writer.startListView(); writer.bigInt().writeBigInt(80); writer.bigInt().writeBigInt(90); writer.endListView(); vector.setValueCount(2); Object result = vector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(60), resultSet.get(0)); assertEquals(Long.valueOf(70), resultSet.get(1)); result = vector.getObject(1); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(80), resultSet.get(0)); assertEquals(Long.valueOf(90), resultSet.get(1)); vector.validate(); } } @Test public void testOverwriteWithNull() { try (ListViewVector vector = ListViewVector.empty("listview", allocator)) { UnionListViewWriter writer = vector.getWriter(); writer.allocate(); ArrowBuf offsetBuffer = vector.getOffsetBuffer(); ArrowBuf sizeBuffer = vector.getSizeBuffer(); writer.setPosition(0); writer.startListView(); writer.bigInt().writeBigInt(10); writer.bigInt().writeBigInt(20); writer.bigInt().writeBigInt(30); writer.endListView(); writer.setPosition(1); writer.startListView(); writer.bigInt().writeBigInt(40); writer.bigInt().writeBigInt(50); writer.endListView(); vector.setValueCount(2); assertEquals(0, offsetBuffer.getInt(0 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offsetBuffer.getInt(1 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, sizeBuffer.getInt(0 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(2, sizeBuffer.getInt(1 * BaseRepeatedValueViewVector.SIZE_WIDTH)); vector.setNull(0); assertEquals(0, offsetBuffer.getInt(0 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, sizeBuffer.getInt(0 * BaseRepeatedValueViewVector.SIZE_WIDTH)); vector.setNull(1); assertEquals(0, offsetBuffer.getInt(0 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, sizeBuffer.getInt(0 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertTrue(vector.isNull(0)); assertTrue(vector.isNull(1)); writer.setPosition(0); writer.startListView(); writer.bigInt().writeBigInt(60); writer.bigInt().writeBigInt(70); writer.endListView(); assertEquals(0, offsetBuffer.getInt(0 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(2, sizeBuffer.getInt(0 * BaseRepeatedValueViewVector.SIZE_WIDTH)); writer.setPosition(1); writer.startListView(); writer.bigInt().writeBigInt(80); writer.bigInt().writeBigInt(90); writer.endListView(); assertEquals(2, offsetBuffer.getInt(1 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(2, sizeBuffer.getInt(1 * BaseRepeatedValueViewVector.SIZE_WIDTH)); vector.setValueCount(2); assertFalse(vector.isNull(0)); assertFalse(vector.isNull(1)); Object result = vector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(60), resultSet.get(0)); assertEquals(Long.valueOf(70), resultSet.get(1)); result = vector.getObject(1); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(80), resultSet.get(0)); assertEquals(Long.valueOf(90), resultSet.get(1)); vector.validate(); } } @Test public void testOutOfOrderOffset1() { // [[12, -7, 25], null, [0, -127, 127, 50], [], [50, 12]] try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { initializeListViewVectorAsInSpecification(listViewVector); final ArrowBuf offSetBuffer = listViewVector.getOffsetBuffer(); final ArrowBuf sizeBuffer = listViewVector.getSizeBuffer(); // check offset buffer assertEquals(4, offSetBuffer.getInt(0 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(7, offSetBuffer.getInt(1 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offSetBuffer.getInt(2 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offSetBuffer.getInt(3 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(4 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); // check size buffer assertEquals(3, sizeBuffer.getInt(0 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(1 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(4, sizeBuffer.getInt(2 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(3 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(2, sizeBuffer.getInt(4 * BaseRepeatedValueViewVector.SIZE_WIDTH)); // check child vector assertEquals(0, ((SmallIntVector) listViewVector.getDataVector()).get(0)); assertEquals(-127, ((SmallIntVector) listViewVector.getDataVector()).get(1)); assertEquals(127, ((SmallIntVector) listViewVector.getDataVector()).get(2)); assertEquals(50, ((SmallIntVector) listViewVector.getDataVector()).get(3)); assertEquals(12, ((SmallIntVector) listViewVector.getDataVector()).get(4)); assertEquals(-7, ((SmallIntVector) listViewVector.getDataVector()).get(5)); assertEquals(25, ((SmallIntVector) listViewVector.getDataVector()).get(6)); // check values Object result = listViewVector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Short.valueOf("12"), resultSet.get(0)); assertEquals(Short.valueOf("-7"), resultSet.get(1)); assertEquals(Short.valueOf("25"), resultSet.get(2)); assertTrue(listViewVector.isNull(1)); result = listViewVector.getObject(2); resultSet = (ArrayList) result; assertEquals(4, resultSet.size()); assertEquals(Short.valueOf("0"), resultSet.get(0)); assertEquals(Short.valueOf("-127"), resultSet.get(1)); assertEquals(Short.valueOf("127"), resultSet.get(2)); assertEquals(Short.valueOf("50"), resultSet.get(3)); assertTrue(listViewVector.isEmpty(3)); result = listViewVector.getObject(4); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Short.valueOf("50"), resultSet.get(0)); assertEquals(Short.valueOf("12"), resultSet.get(1)); listViewVector.validate(); } } private int validateSizeBufferAndCalculateMinOffset( int start, int splitLength, ArrowBuf fromOffsetBuffer, ArrowBuf fromSizeBuffer, ArrowBuf toSizeBuffer) { int minOffset = fromOffsetBuffer.getInt((long) start * ListViewVector.OFFSET_WIDTH); int fromDataLength; int toDataLength; for (int i = 0; i < splitLength; i++) { fromDataLength = fromSizeBuffer.getInt((long) (start + i) * ListViewVector.SIZE_WIDTH); toDataLength = toSizeBuffer.getInt((long) (i) * ListViewVector.SIZE_WIDTH); /* validate size */ assertEquals( fromDataLength, toDataLength, "Different data lengths at index: " + i + " and start: " + start); /* calculate minimum offset */ int currentOffset = fromOffsetBuffer.getInt((long) (start + i) * ListViewVector.OFFSET_WIDTH); if (currentOffset < minOffset) { minOffset = currentOffset; } } return minOffset; } private void validateOffsetBuffer( int start, int splitLength, ArrowBuf fromOffsetBuffer, ArrowBuf toOffsetBuffer, int minOffset) { int offset1; int offset2; for (int i = 0; i < splitLength; i++) { offset1 = fromOffsetBuffer.getInt((long) (start + i) * ListViewVector.OFFSET_WIDTH); offset2 = toOffsetBuffer.getInt((long) (i) * ListViewVector.OFFSET_WIDTH); assertEquals( offset1 - minOffset, offset2, "Different offset values at index: " + i + " and start: " + start); } } private void validateDataBuffer( int start, int splitLength, ArrowBuf fromOffsetBuffer, ArrowBuf fromSizeBuffer, BigIntVector fromDataVector, ArrowBuf toOffsetBuffer, BigIntVector toDataVector) { int dataLength; Long fromValue; for (int i = 0; i < splitLength; i++) { dataLength = fromSizeBuffer.getInt((long) (start + i) * ListViewVector.SIZE_WIDTH); for (int j = 0; j < dataLength; j++) { fromValue = fromDataVector.getObject( (fromOffsetBuffer.getInt((long) (start + i) * ListViewVector.OFFSET_WIDTH) + j)); Long toValue = toDataVector.getObject( (toOffsetBuffer.getInt((long) i * ListViewVector.OFFSET_WIDTH) + j)); assertEquals( fromValue, toValue, "Different data values at index: " + i + " and start: " + start); } } } /** * Validate split and transfer of data from fromVector to toVector. Note that this method assumes * that the child vector is BigIntVector. * * @param start start index * @param splitLength length of data to split and transfer * @param fromVector fromVector * @param toVector toVector */ private void validateSplitAndTransfer( TransferPair transferPair, int start, int splitLength, ListViewVector fromVector, ListViewVector toVector) { transferPair.splitAndTransfer(start, splitLength); /* get offsetBuffer of toVector */ final ArrowBuf toOffsetBuffer = toVector.getOffsetBuffer(); /* get sizeBuffer of toVector */ final ArrowBuf toSizeBuffer = toVector.getSizeBuffer(); /* get dataVector of toVector */ BigIntVector toDataVector = (BigIntVector) toVector.getDataVector(); /* get offsetBuffer of toVector */ final ArrowBuf fromOffsetBuffer = fromVector.getOffsetBuffer(); /* get sizeBuffer of toVector */ final ArrowBuf fromSizeBuffer = fromVector.getSizeBuffer(); /* get dataVector of toVector */ BigIntVector fromDataVector = (BigIntVector) fromVector.getDataVector(); /* validate size buffers */ int minOffset = validateSizeBufferAndCalculateMinOffset( start, splitLength, fromOffsetBuffer, fromSizeBuffer, toSizeBuffer); /* validate offset buffers */ validateOffsetBuffer(start, splitLength, fromOffsetBuffer, toOffsetBuffer, minOffset); /* validate data */ validateDataBuffer( start, splitLength, fromOffsetBuffer, fromSizeBuffer, fromDataVector, toOffsetBuffer, toDataVector); } @Test public void testSplitAndTransfer() throws Exception { try (ListViewVector fromVector = ListViewVector.empty("sourceVector", allocator)) { /* Explicitly add the dataVector */ MinorType type = MinorType.BIGINT; fromVector.addOrGetVector(FieldType.nullable(type.getType())); UnionListViewWriter listViewWriter = fromVector.getWriter(); /* allocate memory */ listViewWriter.allocate(); /* populate data */ listViewWriter.setPosition(0); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(10); listViewWriter.bigInt().writeBigInt(11); listViewWriter.bigInt().writeBigInt(12); listViewWriter.endListView(); listViewWriter.setPosition(1); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(13); listViewWriter.bigInt().writeBigInt(14); listViewWriter.endListView(); listViewWriter.setPosition(2); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(15); listViewWriter.bigInt().writeBigInt(16); listViewWriter.bigInt().writeBigInt(17); listViewWriter.bigInt().writeBigInt(18); listViewWriter.endListView(); listViewWriter.setPosition(3); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(19); listViewWriter.endListView(); listViewWriter.setPosition(4); listViewWriter.startListView(); listViewWriter.bigInt().writeBigInt(20); listViewWriter.bigInt().writeBigInt(21); listViewWriter.bigInt().writeBigInt(22); listViewWriter.bigInt().writeBigInt(23); listViewWriter.endListView(); fromVector.setValueCount(5); /* get offset buffer */ final ArrowBuf offsetBuffer = fromVector.getOffsetBuffer(); /* get size buffer */ final ArrowBuf sizeBuffer = fromVector.getSizeBuffer(); /* get dataVector */ BigIntVector dataVector = (BigIntVector) fromVector.getDataVector(); /* check the vector output */ int index = 0; int offset; int size = 0; Long actual; /* index 0 */ assertFalse(fromVector.isNull(index)); offset = offsetBuffer.getInt(index * ListViewVector.OFFSET_WIDTH); assertEquals(Integer.toString(0), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(10), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(11), actual); offset++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(12), actual); assertEquals( Integer.toString(3), Integer.toString(sizeBuffer.getInt(index * ListViewVector.SIZE_WIDTH))); /* index 1 */ index++; assertFalse(fromVector.isNull(index)); offset = offsetBuffer.getInt(index * ListViewVector.OFFSET_WIDTH); assertEquals(Integer.toString(3), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(13), actual); offset++; size++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(14), actual); size++; assertEquals( Integer.toString(size), Integer.toString(sizeBuffer.getInt(index * ListViewVector.SIZE_WIDTH))); /* index 2 */ size = 0; index++; assertFalse(fromVector.isNull(index)); offset = offsetBuffer.getInt(index * ListViewVector.OFFSET_WIDTH); assertEquals(Integer.toString(5), Integer.toString(offset)); size++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(15), actual); offset++; size++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(16), actual); offset++; size++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(17), actual); offset++; size++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(18), actual); assertEquals( Integer.toString(size), Integer.toString(sizeBuffer.getInt(index * ListViewVector.SIZE_WIDTH))); /* index 3 */ size = 0; index++; assertFalse(fromVector.isNull(index)); offset = offsetBuffer.getInt(index * ListViewVector.OFFSET_WIDTH); assertEquals(Integer.toString(9), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(19), actual); size++; assertEquals( Integer.toString(size), Integer.toString(sizeBuffer.getInt(index * ListViewVector.SIZE_WIDTH))); /* index 4 */ size = 0; index++; assertFalse(fromVector.isNull(index)); offset = offsetBuffer.getInt(index * ListViewVector.OFFSET_WIDTH); assertEquals(Integer.toString(10), Integer.toString(offset)); actual = dataVector.getObject(offset); assertEquals(Long.valueOf(20), actual); offset++; size++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(21), actual); offset++; size++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(22), actual); offset++; size++; actual = dataVector.getObject(offset); assertEquals(Long.valueOf(23), actual); size++; assertEquals( Integer.toString(size), Integer.toString(sizeBuffer.getInt(index * ListViewVector.SIZE_WIDTH))); /* do split and transfer */ try (ListViewVector toVector = ListViewVector.empty("toVector", allocator)) { int[][] transferLengths = {{0, 2}, {3, 1}, {4, 1}}; TransferPair transferPair = fromVector.makeTransferPair(toVector); for (final int[] transferLength : transferLengths) { int start = transferLength[0]; int splitLength = transferLength[1]; validateSplitAndTransfer(transferPair, start, splitLength, fromVector, toVector); } } } } @Test public void testOutOfOrderOffsetSplitAndTransfer() { // [[12, -7, 25], null, [0, -127, 127, 50], [], [50, 12]] try (ListViewVector fromVector = ListViewVector.empty("fromVector", allocator)) { // Allocate buffers in listViewVector by calling `allocateNew` method. fromVector.allocateNew(); // Initialize the child vector using `initializeChildrenFromFields` method. FieldType fieldType = new FieldType(true, new ArrowType.Int(64, true), null, null); Field field = new Field("child-vector", fieldType, null); fromVector.initializeChildrenFromFields(Collections.singletonList(field)); // Set values in the child vector. FieldVector fieldVector = fromVector.getDataVector(); fieldVector.clear(); BigIntVector childVector = (BigIntVector) 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. fromVector.setValidity(0, 1); fromVector.setValidity(1, 0); fromVector.setValidity(2, 1); fromVector.setValidity(3, 1); fromVector.setValidity(4, 1); fromVector.setOffset(0, 4); fromVector.setOffset(1, 7); fromVector.setOffset(2, 0); fromVector.setOffset(3, 0); fromVector.setOffset(4, 3); fromVector.setSize(0, 3); fromVector.setSize(1, 0); fromVector.setSize(2, 4); fromVector.setSize(3, 0); fromVector.setSize(4, 2); // Set value count using `setValueCount` method. fromVector.setValueCount(5); final ArrowBuf offSetBuffer = fromVector.getOffsetBuffer(); final ArrowBuf sizeBuffer = fromVector.getSizeBuffer(); // check offset buffer assertEquals(4, offSetBuffer.getInt(0 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(7, offSetBuffer.getInt(1 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offSetBuffer.getInt(2 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(0, offSetBuffer.getInt(3 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); assertEquals(3, offSetBuffer.getInt(4 * BaseRepeatedValueViewVector.OFFSET_WIDTH)); // check size buffer assertEquals(3, sizeBuffer.getInt(0 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(1 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(4, sizeBuffer.getInt(2 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(0, sizeBuffer.getInt(3 * BaseRepeatedValueViewVector.SIZE_WIDTH)); assertEquals(2, sizeBuffer.getInt(4 * BaseRepeatedValueViewVector.SIZE_WIDTH)); // check child vector assertEquals(0, ((BigIntVector) fromVector.getDataVector()).get(0)); assertEquals(-127, ((BigIntVector) fromVector.getDataVector()).get(1)); assertEquals(127, ((BigIntVector) fromVector.getDataVector()).get(2)); assertEquals(50, ((BigIntVector) fromVector.getDataVector()).get(3)); assertEquals(12, ((BigIntVector) fromVector.getDataVector()).get(4)); assertEquals(-7, ((BigIntVector) fromVector.getDataVector()).get(5)); assertEquals(25, ((BigIntVector) fromVector.getDataVector()).get(6)); // check values Object result = fromVector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); assertEquals(Long.valueOf(12), resultSet.get(0)); assertEquals(Long.valueOf(-7), resultSet.get(1)); assertEquals(Long.valueOf(25), resultSet.get(2)); assertTrue(fromVector.isNull(1)); result = fromVector.getObject(2); resultSet = (ArrayList) result; assertEquals(4, resultSet.size()); assertEquals(Long.valueOf(0), resultSet.get(0)); assertEquals(Long.valueOf(-127), resultSet.get(1)); assertEquals(Long.valueOf(127), resultSet.get(2)); assertEquals(Long.valueOf(50), resultSet.get(3)); assertTrue(fromVector.isEmpty(3)); result = fromVector.getObject(4); resultSet = (ArrayList) result; assertEquals(2, resultSet.size()); assertEquals(Long.valueOf(50), resultSet.get(0)); assertEquals(Long.valueOf(12), resultSet.get(1)); fromVector.validate(); /* do split and transfer */ try (ListViewVector toVector = ListViewVector.empty("toVector", allocator)) { int[][] transferLengths = {{2, 3}, {0, 1}, {0, 3}}; TransferPair transferPair = fromVector.makeTransferPair(toVector); for (final int[] transferLength : transferLengths) { int start = transferLength[0]; int splitLength = transferLength[1]; validateSplitAndTransfer(transferPair, start, splitLength, fromVector, toVector); } } } } @Test public void testRangeChildVector1() { /* * Non-overlapping ranges * offsets: [0, 2] * sizes: [4, 1] * values: [0, 1, 2, 3] * * vector: [[0, 1, 2, 3], [2]] * */ 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(32, 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(); IntVector childVector = (IntVector) fieldVector; childVector.allocateNew(8); childVector.set(0, 0); childVector.set(1, 1); childVector.set(2, 2); childVector.set(3, 3); childVector.set(4, 4); childVector.set(5, 5); childVector.set(6, 6); childVector.set(7, 7); childVector.setValueCount(8); // Set validity, offset and size buffers using `setValidity`, // `setOffset` and `setSize` methods. listViewVector.setValidity(0, 1); listViewVector.setValidity(1, 1); listViewVector.setOffset(0, 0); listViewVector.setOffset(1, 2); listViewVector.setSize(0, 4); listViewVector.setSize(1, 1); assertEquals(8, listViewVector.getDataVector().getValueCount()); listViewVector.setValueCount(2); assertEquals(4, listViewVector.getDataVector().getValueCount()); IntVector childVector1 = (IntVector) listViewVector.getDataVector(); final ArrowBuf dataBuffer = childVector1.getDataBuffer(); final ArrowBuf validityBuffer = childVector1.getValidityBuffer(); // yet the underneath buffer contains the original buffer for (int i = 0; i < validityBuffer.capacity(); i++) { assertEquals(i, dataBuffer.getInt((long) i * IntVector.TYPE_WIDTH)); } } } @Test public void testRangeChildVector2() { /* * Overlapping ranges * offsets: [0, 2] * sizes: [3, 1] * values: [0, 1, 2, 3] * * vector: [[1, 2, 3], [2]] * */ 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(32, 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(); IntVector childVector = (IntVector) fieldVector; childVector.allocateNew(8); childVector.set(0, 0); childVector.set(1, 1); childVector.set(2, 2); childVector.set(3, 3); childVector.set(4, 4); childVector.set(5, 5); childVector.set(6, 6); childVector.set(7, 7); childVector.setValueCount(8); // Set validity, offset and size buffers using `setValidity`, // `setOffset` and `setSize` methods. listViewVector.setValidity(0, 1); listViewVector.setValidity(1, 1); listViewVector.setOffset(0, 1); listViewVector.setOffset(1, 2); listViewVector.setSize(0, 3); listViewVector.setSize(1, 1); assertEquals(8, listViewVector.getDataVector().getValueCount()); listViewVector.setValueCount(2); assertEquals(4, listViewVector.getDataVector().getValueCount()); IntVector childVector1 = (IntVector) listViewVector.getDataVector(); final ArrowBuf dataBuffer = childVector1.getDataBuffer(); final ArrowBuf validityBuffer = childVector1.getValidityBuffer(); // yet the underneath buffer contains the original buffer for (int i = 0; i < validityBuffer.capacity(); i++) { assertEquals(i, dataBuffer.getInt((long) i * IntVector.TYPE_WIDTH)); } } } @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) { writer.integer().writeInt(v); } 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()); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java000066400000000000000000001704321515451615200311450ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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; import org.apache.arrow.vector.complex.StructVector; 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; public class TestMapVector { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } public T getResultKey(Map resultStruct) { assertTrue(resultStruct.containsKey(MapVector.KEY_NAME)); return resultStruct.get(MapVector.KEY_NAME); } public T getResultValue(Map resultStruct) { assertTrue(resultStruct.containsKey(MapVector.VALUE_NAME)); return resultStruct.get(MapVector.VALUE_NAME); } @Test public void testBasicOperation() { int count = 5; try (MapVector mapVector = MapVector.empty("map", allocator, false)) { mapVector.allocateNew(); UnionMapWriter mapWriter = mapVector.getWriter(); for (int i = 0; i < count; i++) { mapWriter.startMap(); for (int j = 0; j < i + 1; j++) { mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(j); mapWriter.value().integer().writeInt(j); mapWriter.endEntry(); } mapWriter.endMap(); } mapWriter.setValueCount(count); UnionMapReader mapReader = mapVector.getReader(); for (int i = 0; i < count; i++) { mapReader.setPosition(i); for (int j = 0; j < i + 1; j++) { mapReader.next(); assertEquals(j, mapReader.key().readLong().longValue(), "record: " + i); assertEquals(j, mapReader.value().readInteger().intValue()); } } } } @Test public void testBasicOperationNulls() { int count = 6; try (MapVector mapVector = MapVector.empty("map", allocator, false)) { mapVector.allocateNew(); UnionMapWriter mapWriter = mapVector.getWriter(); for (int i = 0; i < count; i++) { // i == 1 is a NULL if (i != 1) { mapWriter.setPosition(i); mapWriter.startMap(); // i == 3 is an empty map if (i != 3) { for (int j = 0; j < i + 1; j++) { mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(j); // i == 5 maps to a NULL value if (i != 5) { mapWriter.value().integer().writeInt(j); } mapWriter.endEntry(); } } mapWriter.endMap(); } } mapWriter.setValueCount(count); UnionMapReader mapReader = mapVector.getReader(); for (int i = 0; i < count; i++) { mapReader.setPosition(i); if (i == 1) { assertFalse(mapReader.isSet()); } else { if (i == 3) { JsonStringArrayList result = (JsonStringArrayList) mapReader.readObject(); assertTrue(result.isEmpty()); } else { for (int j = 0; j < i + 1; j++) { mapReader.next(); assertEquals(j, mapReader.key().readLong().longValue(), "record: " + i); if (i == 5) { assertFalse(mapReader.value().isSet()); } else { assertEquals(j, mapReader.value().readInteger().intValue()); } } } } } } } @Test public void testCopyFrom() throws Exception { try (MapVector inVector = MapVector.empty("input", allocator, false); MapVector outVector = MapVector.empty("output", allocator, false)) { UnionMapWriter writer = inVector.getWriter(); writer.allocate(); // populate input vector with the following records // {1 -> 11, 2 -> 22, 3 -> 33} // null // {2 -> null} writer.setPosition(0); // optional writer.startMap(); writer.startEntry(); writer.key().bigInt().writeBigInt(1); writer.value().bigInt().writeBigInt(11); writer.endEntry(); writer.startEntry(); writer.key().bigInt().writeBigInt(2); writer.value().bigInt().writeBigInt(22); writer.endEntry(); writer.startEntry(); writer.key().bigInt().writeBigInt(3); writer.value().bigInt().writeBigInt(33); writer.endEntry(); writer.endMap(); writer.setPosition(2); writer.startMap(); writer.startEntry(); writer.key().bigInt().writeBigInt(2); writer.endEntry(); writer.endMap(); writer.setValueCount(3); // copy values from input to output outVector.allocateNew(); for (int i = 0; i < 3; i++) { outVector.copyFrom(i, i, inVector); } outVector.setValueCount(3); // assert the output vector is correct FieldReader reader = outVector.getReader(); assertTrue(reader.isSet(), "shouldn't be null"); reader.setPosition(1); assertFalse(reader.isSet(), "should be null"); reader.setPosition(2); assertTrue(reader.isSet(), "shouldn't be null"); /* index 0 */ Object result = outVector.getObject(0); ArrayList resultSet = (ArrayList) result; assertEquals(3, resultSet.size()); Map resultStruct = (Map) resultSet.get(0); assertEquals(1L, getResultKey(resultStruct)); assertEquals(11L, getResultValue(resultStruct)); resultStruct = (Map) resultSet.get(1); assertEquals(2L, getResultKey(resultStruct)); assertEquals(22L, getResultValue(resultStruct)); resultStruct = (Map) resultSet.get(2); assertEquals(3L, getResultKey(resultStruct)); assertEquals(33L, getResultValue(resultStruct)); /* index 1 */ result = outVector.getObject(1); assertNull(result); /* index 2 */ result = outVector.getObject(2); resultSet = (ArrayList) result; assertEquals(1, resultSet.size()); resultStruct = (Map) resultSet.get(0); assertEquals(2L, getResultKey(resultStruct)); assertFalse(resultStruct.containsKey(MapVector.VALUE_NAME)); } } @Test public void testSplitAndTransfer() throws Exception { try (MapVector mapVector = MapVector.empty("sourceVector", allocator, false)) { /* Explicitly add the map child vectors */ FieldType type = new FieldType(false, ArrowType.Struct.INSTANCE, null, null); AddOrGetResult addResult = mapVector.addOrGetVector(type); FieldType keyType = new FieldType(false, MinorType.BIGINT.getType(), null, null); FieldType valueType = FieldType.nullable(MinorType.FLOAT8.getType()); addResult.getVector().addOrGet(MapVector.KEY_NAME, keyType, BigIntVector.class); addResult.getVector().addOrGet(MapVector.VALUE_NAME, valueType, Float8Vector.class); UnionMapWriter mapWriter = mapVector.getWriter(); /* allocate memory */ mapWriter.allocate(); /* populate data */ mapWriter.setPosition(0); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(10); mapWriter.value().float8().writeFloat8(1.0); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(11); mapWriter.value().float8().writeFloat8(1.1); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(12); mapWriter.value().float8().writeFloat8(1.2); mapWriter.endEntry(); mapWriter.endMap(); mapWriter.setPosition(1); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(13); mapWriter.value().float8().writeFloat8(1.3); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(14); mapWriter.value().float8().writeFloat8(1.4); mapWriter.endEntry(); mapWriter.endMap(); mapWriter.setPosition(2); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(15); mapWriter.value().float8().writeFloat8(1.5); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(16); mapWriter.value().float8().writeFloat8(1.6); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(17); mapWriter.value().float8().writeFloat8(1.7); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(18); mapWriter.value().float8().writeFloat8(1.8); mapWriter.endEntry(); mapWriter.endMap(); mapWriter.setPosition(3); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(19); mapWriter.value().float8().writeFloat8(1.9); mapWriter.endEntry(); mapWriter.endMap(); mapWriter.setPosition(4); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(20); mapWriter.value().float8().writeFloat8(2.0); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(21); mapWriter.value().float8().writeFloat8(2.1); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(22); mapWriter.value().float8().writeFloat8(2.2); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(23); mapWriter.value().float8().writeFloat8(2.3); mapWriter.endEntry(); mapWriter.endMap(); mapVector.setValueCount(5); assertEquals(4, mapVector.getLastSet()); /* get offset buffer */ final ArrowBuf offsetBuffer = mapVector.getOffsetBuffer(); /* get dataVector */ StructVector dataVector = (StructVector) mapVector.getDataVector(); /* check the vector output */ int index = 0; int offset; Map result; /* index 0 */ assertFalse(mapVector.isNull(index)); offset = offsetBuffer.getInt(index * MapVector.OFFSET_WIDTH); assertEquals(Integer.toString(0), Integer.toString(offset)); result = dataVector.getObject(offset); assertEquals(10L, getResultKey(result)); assertEquals(1.0, getResultValue(result)); offset++; result = dataVector.getObject(offset); assertEquals(11L, getResultKey(result)); assertEquals(1.1, getResultValue(result)); offset++; result = dataVector.getObject(offset); assertEquals(12L, getResultKey(result)); assertEquals(1.2, getResultValue(result)); /* index 1 */ index++; assertFalse(mapVector.isNull(index)); offset = offsetBuffer.getInt(index * MapVector.OFFSET_WIDTH); assertEquals(Integer.toString(3), Integer.toString(offset)); result = dataVector.getObject(offset); assertEquals(13L, getResultKey(result)); assertEquals(1.3, getResultValue(result)); offset++; result = dataVector.getObject(offset); assertEquals(14L, getResultKey(result)); assertEquals(1.4, getResultValue(result)); /* index 2 */ index++; assertFalse(mapVector.isNull(index)); offset = offsetBuffer.getInt(index * MapVector.OFFSET_WIDTH); assertEquals(Integer.toString(5), Integer.toString(offset)); result = dataVector.getObject(offset); assertEquals(15L, getResultKey(result)); assertEquals(1.5, getResultValue(result)); offset++; result = dataVector.getObject(offset); assertEquals(16L, getResultKey(result)); assertEquals(1.6, getResultValue(result)); offset++; result = dataVector.getObject(offset); assertEquals(17L, getResultKey(result)); assertEquals(1.7, getResultValue(result)); offset++; result = dataVector.getObject(offset); assertEquals(18L, getResultKey(result)); assertEquals(1.8, getResultValue(result)); /* index 3 */ index++; assertFalse(mapVector.isNull(index)); offset = offsetBuffer.getInt(index * MapVector.OFFSET_WIDTH); assertEquals(Integer.toString(9), Integer.toString(offset)); result = dataVector.getObject(offset); assertEquals(19L, getResultKey(result)); assertEquals(1.9, getResultValue(result)); /* index 4 */ index++; assertFalse(mapVector.isNull(index)); offset = offsetBuffer.getInt(index * MapVector.OFFSET_WIDTH); assertEquals(Integer.toString(10), Integer.toString(offset)); result = dataVector.getObject(offset); assertEquals(20L, getResultKey(result)); assertEquals(2.0, getResultValue(result)); offset++; result = dataVector.getObject(offset); assertEquals(21L, getResultKey(result)); assertEquals(2.1, getResultValue(result)); offset++; result = dataVector.getObject(offset); assertEquals(22L, getResultKey(result)); assertEquals(2.2, getResultValue(result)); offset++; result = dataVector.getObject(offset); assertEquals(23L, getResultKey(result)); assertEquals(2.3, getResultValue(result)); /* index 5 */ index++; assertTrue(mapVector.isNull(index)); offset = offsetBuffer.getInt(index * MapVector.OFFSET_WIDTH); assertEquals(Integer.toString(14), Integer.toString(offset)); /* do split and transfer */ try (MapVector toVector = MapVector.empty("toVector", allocator, false)) { TransferPair transferPair = mapVector.makeTransferPair(toVector); int[][] transferLengths = {{0, 2}, {3, 1}, {4, 1}}; for (final int[] transferLength : transferLengths) { int start = transferLength[0]; int splitLength = transferLength[1]; int dataLength1 = 0; int dataLength2 = 0; int offset1 = 0; int offset2 = 0; transferPair.splitAndTransfer(start, splitLength); /* get offsetBuffer of toVector */ final ArrowBuf toOffsetBuffer = toVector.getOffsetBuffer(); /* get dataVector of toVector */ StructVector dataVector1 = (StructVector) toVector.getDataVector(); for (int i = 0; i < splitLength; i++) { dataLength1 = offsetBuffer.getInt((start + i + 1) * MapVector.OFFSET_WIDTH) - offsetBuffer.getInt((start + i) * MapVector.OFFSET_WIDTH); dataLength2 = toOffsetBuffer.getInt((i + 1) * MapVector.OFFSET_WIDTH) - toOffsetBuffer.getInt(i * MapVector.OFFSET_WIDTH); assertEquals( dataLength1, dataLength2, "Different data lengths at index: " + i + " and start: " + start); offset1 = offsetBuffer.getInt((start + i) * MapVector.OFFSET_WIDTH); offset2 = toOffsetBuffer.getInt(i * MapVector.OFFSET_WIDTH); for (int j = 0; j < dataLength1; j++) { assertEquals( dataVector.getObject(offset1), dataVector1.getObject(offset2), "Different data at indexes: " + offset1 + " and " + offset2); offset1++; offset2++; } } } } } } @Test public void testMapWithListValue() throws Exception { try (MapVector mapVector = MapVector.empty("sourceVector", allocator, false)) { UnionMapWriter mapWriter = mapVector.getWriter(); ListWriter valueWriter; /* allocate memory */ mapWriter.allocate(); /* the dataVector that backs a listVector will also be a * listVector for this test. */ /* write one or more maps index 0 */ mapWriter.setPosition(0); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(1); valueWriter = mapWriter.value().list(); valueWriter.startList(); valueWriter.bigInt().writeBigInt(50); valueWriter.bigInt().writeBigInt(100); valueWriter.bigInt().writeBigInt(200); valueWriter.endList(); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(2); valueWriter = mapWriter.value().list(); valueWriter.startList(); valueWriter.bigInt().writeBigInt(75); valueWriter.bigInt().writeBigInt(125); valueWriter.bigInt().writeBigInt(150); valueWriter.bigInt().writeBigInt(175); valueWriter.endList(); mapWriter.endEntry(); mapWriter.endMap(); /* write one or more maps at index 1 */ mapWriter.setPosition(1); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(3); valueWriter = mapWriter.value().list(); valueWriter.startList(); valueWriter.bigInt().writeBigInt(10); valueWriter.endList(); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(4); valueWriter = mapWriter.value().list(); valueWriter.startList(); valueWriter.bigInt().writeBigInt(15); valueWriter.bigInt().writeBigInt(20); valueWriter.endList(); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(5); valueWriter = mapWriter.value().list(); valueWriter.startList(); valueWriter.bigInt().writeBigInt(25); valueWriter.bigInt().writeBigInt(30); valueWriter.bigInt().writeBigInt(35); valueWriter.endList(); mapWriter.endEntry(); mapWriter.endMap(); assertEquals(1, mapVector.getLastSet()); mapWriter.setValueCount(2); assertEquals(2, mapVector.getValueCount()); // Get mapVector element at index 0 Object result = mapVector.getObject(0); ArrayList resultSet = (ArrayList) result; // 2 map entries at index 0 assertEquals(2, resultSet.size()); // First Map entry Map resultStruct = (Map) resultSet.get(0); assertEquals(1L, getResultKey(resultStruct)); ArrayList list = (ArrayList) getResultValue(resultStruct); assertEquals(3, list.size()); // value is a list with 3 elements assertEquals(Long.valueOf(50), list.get(0)); assertEquals(Long.valueOf(100), list.get(1)); assertEquals(Long.valueOf(200), list.get(2)); // Second Map entry resultStruct = (Map) resultSet.get(1); list = (ArrayList) getResultValue(resultStruct); assertEquals(4, list.size()); // value is a list with 4 elements assertEquals(Long.valueOf(75), list.get(0)); assertEquals(Long.valueOf(125), list.get(1)); assertEquals(Long.valueOf(150), list.get(2)); assertEquals(Long.valueOf(175), list.get(3)); // Get mapVector element at index 1 result = mapVector.getObject(1); resultSet = (ArrayList) result; // First Map entry resultStruct = (Map) resultSet.get(0); assertEquals(3L, getResultKey(resultStruct)); list = (ArrayList) getResultValue(resultStruct); assertEquals(1, list.size()); // value is a list with 1 element assertEquals(Long.valueOf(10), list.get(0)); // Second Map entry resultStruct = (Map) resultSet.get(1); assertEquals(4L, getResultKey(resultStruct)); list = (ArrayList) getResultValue(resultStruct); assertEquals(2, list.size()); // value is a list with 1 element assertEquals(Long.valueOf(15), list.get(0)); assertEquals(Long.valueOf(20), list.get(1)); // Third Map entry resultStruct = (Map) resultSet.get(2); assertEquals(5L, getResultKey(resultStruct)); list = (ArrayList) getResultValue(resultStruct); assertEquals(3, list.size()); // value is a list with 1 element assertEquals(Long.valueOf(25), list.get(0)); assertEquals(Long.valueOf(30), list.get(1)); assertEquals(Long.valueOf(35), list.get(2)); /* check underlying bitVector */ assertFalse(mapVector.isNull(0)); assertFalse(mapVector.isNull(1)); /* check underlying offsets */ final ArrowBuf offsetBuffer = mapVector.getOffsetBuffer(); /* mapVector has 2 entries at index 0 and 3 entries at index 1 */ assertEquals(0, offsetBuffer.getInt(0 * MapVector.OFFSET_WIDTH)); assertEquals(2, offsetBuffer.getInt(1 * MapVector.OFFSET_WIDTH)); assertEquals(5, offsetBuffer.getInt(2 * MapVector.OFFSET_WIDTH)); } } @Test public void testMapWithMapValue() throws Exception { try (MapVector mapVector = MapVector.empty("sourceVector", allocator, false)) { UnionMapWriter mapWriter = mapVector.getWriter(); MapWriter valueWriter; // we are essentially writing Map> // populate map vector with the following four records // [ // null, // [1:[50: 100, 200:400], 2:[75: 175, 150: 250]], // [3:[10: 20], 4:[15: 20], 5:[25: 30, 35: null]], // [8:[15: 30, 10: 20]] // ] /* write null at index 0 */ mapWriter.setPosition(0); mapWriter.writeNull(); /* write one or more maps at index 1 */ mapWriter.setPosition(1); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(1); valueWriter = mapWriter.value().map(false); valueWriter.startMap(); writeEntry(valueWriter, 50, 100L); writeEntry(valueWriter, 200, 400L); valueWriter.endMap(); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(2); valueWriter = mapWriter.value().map(false); valueWriter.startMap(); writeEntry(valueWriter, 75, 175L); writeEntry(valueWriter, 150, 250L); valueWriter.endMap(); mapWriter.endEntry(); mapWriter.endMap(); /* write one or more maps at index 2 */ mapWriter.setPosition(2); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(3); valueWriter = mapWriter.value().map(true); valueWriter.startMap(); writeEntry(valueWriter, 10, 20L); valueWriter.endMap(); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(4); valueWriter = mapWriter.value().map(false); valueWriter.startMap(); writeEntry(valueWriter, 15, 20L); valueWriter.endMap(); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(5); valueWriter = mapWriter.value().map(false); valueWriter.startMap(); writeEntry(valueWriter, 25, 30L); writeEntry(valueWriter, 35, (Long) null); valueWriter.endMap(); mapWriter.endEntry(); mapWriter.endMap(); /* write one or more maps at index 3 */ mapWriter.setPosition(3); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(8); valueWriter = mapWriter.value().map(); valueWriter.startMap(); writeEntry(valueWriter, 15, 30L); writeEntry(valueWriter, 10, 20L); valueWriter.endMap(); mapWriter.endEntry(); mapWriter.endMap(); assertEquals(3, mapVector.getLastSet()); mapWriter.setValueCount(4); assertEquals(4, mapVector.getValueCount()); // Get mapVector element at index 0 Object result = mapVector.getObject(0); assertNull(result); // Get mapVector element at index 1 result = mapVector.getObject(1); ArrayList resultSet = (ArrayList) result; // 2 map entries at index 0 assertEquals(2, resultSet.size()); // First Map entry Map resultStruct = (Map) resultSet.get(0); assertEquals(1L, getResultKey(resultStruct)); ArrayList> list = (ArrayList>) getResultValue(resultStruct); assertEquals(2, list.size()); // value is a list of 2 two maps Map innerMap = list.get(0); assertEquals(50L, getResultKey(innerMap)); assertEquals(100L, getResultValue(innerMap)); innerMap = list.get(1); assertEquals(200L, getResultKey(innerMap)); assertEquals(400L, getResultValue(innerMap)); // Second Map entry resultStruct = (Map) resultSet.get(1); assertEquals(2L, getResultKey(resultStruct)); list = (ArrayList>) getResultValue(resultStruct); assertEquals(2, list.size()); // value is a list of two maps innerMap = list.get(0); assertEquals(75L, getResultKey(innerMap)); assertEquals(175L, getResultValue(innerMap)); innerMap = list.get(1); assertEquals(150L, getResultKey(innerMap)); assertEquals(250L, getResultValue(innerMap)); // Get mapVector element at index 2 result = mapVector.getObject(2); resultSet = (ArrayList) result; // 3 map entries at index 1 assertEquals(3, resultSet.size()); // First Map entry resultStruct = (Map) resultSet.get(0); assertEquals(3L, getResultKey(resultStruct)); list = (ArrayList>) getResultValue(resultStruct); assertEquals(1, list.size()); // value is a list of maps with 1 element innerMap = list.get(0); assertEquals(10L, getResultKey(innerMap)); assertEquals(20L, getResultValue(innerMap)); // Second Map entry resultStruct = (Map) resultSet.get(1); assertEquals(4L, getResultKey(resultStruct)); list = (ArrayList>) getResultValue(resultStruct); assertEquals(1, list.size()); // value is a list of maps with 1 element innerMap = list.get(0); assertEquals(15L, getResultKey(innerMap)); assertEquals(20L, getResultValue(innerMap)); // Third Map entry resultStruct = (Map) resultSet.get(2); assertEquals(5L, getResultKey(resultStruct)); list = (ArrayList>) getResultValue(resultStruct); assertEquals(2, list.size()); // value is a list of maps with 2 elements innerMap = list.get(0); assertEquals(25L, getResultKey(innerMap)); assertEquals(30L, getResultValue(innerMap)); innerMap = list.get(1); assertEquals(35L, getResultKey(innerMap)); assertNull(innerMap.get(MapVector.VALUE_NAME)); // Get mapVector element at index 3 result = mapVector.getObject(3); resultSet = (ArrayList) result; // only 1 map entry at index 3 assertEquals(1, resultSet.size()); resultStruct = (Map) resultSet.get(0); assertEquals(8L, getResultKey(resultStruct)); list = (ArrayList>) getResultValue(resultStruct); assertEquals(2, list.size()); // value is a list of 2 maps innerMap = list.get(0); assertEquals(15L, getResultKey(innerMap)); assertEquals(30L, getResultValue(innerMap)); innerMap = list.get(1); assertEquals(10L, getResultKey(innerMap)); assertEquals(20L, getResultValue(innerMap)); /* check underlying bitVector */ assertTrue(mapVector.isNull(0)); assertFalse(mapVector.isNull(1)); assertFalse(mapVector.isNull(2)); assertFalse(mapVector.isNull(3)); /* check underlying offsets */ final ArrowBuf offsetBuffer = mapVector.getOffsetBuffer(); // mapVector has 0 entries at index 0, 2 entries at index 1, 3 entries at index 2, // and 1 entry at index 3 assertEquals(0, offsetBuffer.getInt(0 * MapVector.OFFSET_WIDTH)); assertEquals(0, offsetBuffer.getInt(1 * MapVector.OFFSET_WIDTH)); assertEquals(2, offsetBuffer.getInt(2 * MapVector.OFFSET_WIDTH)); assertEquals(5, offsetBuffer.getInt(3 * MapVector.OFFSET_WIDTH)); assertEquals(6, offsetBuffer.getInt(4 * MapVector.OFFSET_WIDTH)); } } @Test public void testMapWithMapKeyAndMapValue() throws Exception { try (MapVector mapVector = MapVector.empty("sourceVector", allocator, false)) { UnionMapWriter mapWriter = mapVector.getWriter(); MapWriter keyWriter; MapWriter valueWriter; // we are essentially writing Map, Map> // populate map vector with the following two records // [ // [[5: 10, 20: 40]:[50: 100, 200: 400], [50: 100]:[75: 175, 150: 250]], // [[1: 2]:[10: 20], [30: 40]:[15: 20], [50: 60, 70: null]:[25: 30, 35: null], [5: null]: // null] // ] mapWriter.setPosition(0); mapWriter.startMap(); mapWriter.startEntry(); keyWriter = mapWriter.key().map(false); keyWriter.startMap(); writeEntry(keyWriter, 5, 10); writeEntry(keyWriter, 20, 40); keyWriter.endMap(); valueWriter = mapWriter.value().map(false); valueWriter.startMap(); writeEntry(valueWriter, 50, 100L); writeEntry(valueWriter, 200, 400L); valueWriter.endMap(); mapWriter.endEntry(); mapWriter.startEntry(); keyWriter = mapWriter.key().map(false); keyWriter.startMap(); writeEntry(keyWriter, 50, 100); keyWriter.endMap(); valueWriter = mapWriter.value().map(false); valueWriter.startMap(); writeEntry(valueWriter, 75, 175L); writeEntry(valueWriter, 150, 250L); valueWriter.endMap(); mapWriter.endEntry(); mapWriter.endMap(); /* write one or more maps at index 1 */ mapWriter.setPosition(1); mapWriter.startMap(); mapWriter.startEntry(); keyWriter = mapWriter.key().map(false); keyWriter.startMap(); writeEntry(keyWriter, 1, 2); keyWriter.endMap(); valueWriter = mapWriter.value().map(true); valueWriter.startMap(); writeEntry(valueWriter, 10, 20L); valueWriter.endMap(); mapWriter.endEntry(); mapWriter.startEntry(); keyWriter = mapWriter.key().map(false); keyWriter.startMap(); writeEntry(keyWriter, 30, 40); keyWriter.endMap(); valueWriter = mapWriter.value().map(false); valueWriter.startMap(); writeEntry(valueWriter, 15, 20L); valueWriter.endMap(); mapWriter.endEntry(); mapWriter.startEntry(); keyWriter = mapWriter.key().map(false); keyWriter.startMap(); writeEntry(keyWriter, 50, 60); writeEntry(keyWriter, 70, (Integer) null); keyWriter.endMap(); valueWriter = mapWriter.value().map(false); valueWriter.startMap(); writeEntry(valueWriter, 25, 30L); writeEntry(valueWriter, 35, (Long) null); valueWriter.endMap(); mapWriter.endEntry(); mapWriter.startEntry(); keyWriter = mapWriter.key().map(false); keyWriter.startMap(); writeEntry(keyWriter, 5, (Integer) null); keyWriter.endMap(); valueWriter = mapWriter.value().map(false); valueWriter.writeNull(); mapWriter.endEntry(); mapWriter.endMap(); assertEquals(1, mapVector.getLastSet()); mapWriter.setValueCount(2); assertEquals(2, mapVector.getValueCount()); // Get mapVector element at index 0 Object result = mapVector.getObject(0); ArrayList resultSet = (ArrayList) result; // 2 map entries at index 0 assertEquals(2, resultSet.size()); // First Map entry Map>> resultStruct = (Map>>) resultSet.get(0); ArrayList> list = getResultKey(resultStruct); assertEquals(2, list.size()); // key is a list of 2 two maps Map innerMap = list.get(0); assertEquals(5, getResultKey(innerMap)); assertEquals(10, getResultValue(innerMap)); innerMap = list.get(1); assertEquals(20, getResultKey(innerMap)); assertEquals(40, getResultValue(innerMap)); list = getResultValue(resultStruct); assertEquals(2, list.size()); // value is a list of 2 two maps innerMap = list.get(0); assertEquals(50L, getResultKey(innerMap)); assertEquals(100L, getResultValue(innerMap)); innerMap = list.get(1); assertEquals(200L, getResultKey(innerMap)); assertEquals(400L, getResultValue(innerMap)); // Second Map entry resultStruct = (Map>>) resultSet.get(1); list = getResultKey(resultStruct); assertEquals(1, list.size()); // key is a list of 1 two map innerMap = list.get(0); assertEquals(50, getResultKey(innerMap)); assertEquals(100, getResultValue(innerMap)); list = getResultValue(resultStruct); assertEquals(2, list.size()); // value is a list of two maps innerMap = list.get(0); assertEquals(75L, getResultKey(innerMap)); assertEquals(175L, getResultValue(innerMap)); innerMap = list.get(1); assertEquals(150L, getResultKey(innerMap)); assertEquals(250L, getResultValue(innerMap)); // Get mapVector element at index 1 result = mapVector.getObject(1); resultSet = (ArrayList) result; // 4 map entries at index 1 assertEquals(4, resultSet.size()); // First Map entry resultStruct = (Map>>) resultSet.get(0); list = getResultKey(resultStruct); assertEquals(1, list.size()); // key is a list of 1 map innerMap = list.get(0); assertEquals(1, getResultKey(innerMap)); assertEquals(2, getResultValue(innerMap)); list = getResultValue(resultStruct); assertEquals(1, list.size()); // value is a list of maps with 1 element innerMap = list.get(0); assertEquals(10L, getResultKey(innerMap)); assertEquals(20L, getResultValue(innerMap)); // Second Map entry resultStruct = (Map>>) resultSet.get(1); list = getResultKey(resultStruct); assertEquals(1, list.size()); // key is a list of 1 map innerMap = list.get(0); assertEquals(30, getResultKey(innerMap)); assertEquals(40, getResultValue(innerMap)); list = getResultValue(resultStruct); assertEquals(1, list.size()); // value is a list of maps with 1 element innerMap = list.get(0); assertEquals(15L, getResultKey(innerMap)); assertEquals(20L, getResultValue(innerMap)); // Third Map entry resultStruct = (Map>>) resultSet.get(2); list = getResultKey(resultStruct); assertEquals(2, list.size()); // key is a list of two maps innerMap = list.get(0); assertEquals(50, getResultKey(innerMap)); assertEquals(60, getResultValue(innerMap)); innerMap = list.get(1); assertEquals(70, getResultKey(innerMap)); assertNull(innerMap.get(MapVector.VALUE_NAME)); list = getResultValue(resultStruct); assertEquals(2, list.size()); // value is a list of maps with 2 elements innerMap = list.get(0); assertEquals(25L, getResultKey(innerMap)); assertEquals(30L, getResultValue(innerMap)); innerMap = list.get(1); assertEquals(35L, getResultKey(innerMap)); assertNull(innerMap.get(MapVector.VALUE_NAME)); // Fourth Map entry resultStruct = (Map>>) resultSet.get(3); list = getResultKey(resultStruct); assertEquals(1, list.size()); // key is a list of two maps innerMap = list.get(0); assertEquals(5, getResultKey(innerMap)); assertNull(innerMap.get(MapVector.VALUE_NAME)); assertNull(resultStruct.get(MapVector.VALUE_NAME)); /* check underlying bitVector */ assertFalse(mapVector.isNull(0)); assertFalse(mapVector.isNull(1)); /* check underlying offsets */ final ArrowBuf offsetBuffer = mapVector.getOffsetBuffer(); /* mapVector has 2 entries at index 0 and 4 entries at index 1 */ assertEquals(0, offsetBuffer.getInt(0)); assertEquals(2, offsetBuffer.getInt(MapVector.OFFSET_WIDTH)); assertEquals(6, offsetBuffer.getInt(2 * MapVector.OFFSET_WIDTH)); } } private void writeEntry(MapWriter writer, long key, Long value) { writer.startEntry(); writer.key().bigInt().writeBigInt(key); if (value != null) { writer.value().bigInt().writeBigInt(value); } writer.endEntry(); } private void writeEntry(MapWriter writer, int key, Integer value) { writer.startEntry(); writer.key().integer().writeInt(key); if (value != null) { writer.value().integer().writeInt(value); } writer.endEntry(); } @Test public void testClearAndReuse() { try (final MapVector vector = MapVector.empty("map", allocator, false)) { vector.allocateNew(); UnionMapWriter mapWriter = vector.getWriter(); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(1); mapWriter.value().integer().writeInt(11); mapWriter.endEntry(); mapWriter.endMap(); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(2); mapWriter.value().integer().writeInt(22); mapWriter.endEntry(); mapWriter.endMap(); mapWriter.setValueCount(2); Object result = vector.getObject(0); ArrayList resultSet = (ArrayList) result; Map resultStruct = (Map) resultSet.get(0); assertEquals(1L, getResultKey(resultStruct)); assertEquals(11, getResultValue(resultStruct)); result = vector.getObject(1); resultSet = (ArrayList) result; resultStruct = (Map) resultSet.get(0); assertEquals(2L, getResultKey(resultStruct)); assertEquals(22, getResultValue(resultStruct)); // Clear and release the buffers to trigger a realloc when adding next value vector.clear(); mapWriter = new UnionMapWriter(vector); // The map vector should reuse a buffer when reallocating the offset buffer mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(3); mapWriter.value().integer().writeInt(33); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(4); mapWriter.value().integer().writeInt(44); mapWriter.endEntry(); mapWriter.endMap(); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(5); mapWriter.value().integer().writeInt(55); mapWriter.endEntry(); mapWriter.endMap(); mapWriter.setValueCount(2); result = vector.getObject(0); resultSet = (ArrayList) result; resultStruct = (Map) resultSet.get(0); assertEquals(3L, getResultKey(resultStruct)); assertEquals(33, getResultValue(resultStruct)); resultStruct = (Map) resultSet.get(1); assertEquals(4L, getResultKey(resultStruct)); assertEquals(44, getResultValue(resultStruct)); result = vector.getObject(1); resultSet = (ArrayList) result; resultStruct = (Map) resultSet.get(0); assertEquals(5L, getResultKey(resultStruct)); assertEquals(55, getResultValue(resultStruct)); } } @Test public void testGetTransferPair() { try (MapVector mapVector = MapVector.empty("mapVector", allocator, false)) { FieldType type = new FieldType(false, ArrowType.Struct.INSTANCE, null, null); AddOrGetResult addResult = mapVector.addOrGetVector(type); FieldType keyType = new FieldType(false, MinorType.BIGINT.getType(), null, null); FieldType valueType = FieldType.nullable(MinorType.FLOAT8.getType()); addResult.getVector().addOrGet(MapVector.KEY_NAME, keyType, BigIntVector.class); addResult.getVector().addOrGet(MapVector.VALUE_NAME, valueType, Float8Vector.class); mapVector.allocateNew(); mapVector.setValueCount(0); assertEquals(-1, mapVector.getLastSet()); TransferPair tp = mapVector.getTransferPair(mapVector.getName(), allocator, null); tp.transfer(); ValueVector vector = tp.getTo(); assertSame(vector.getClass(), mapVector.getClass()); vector.clear(); } } @Test public void testGetTransferPairWithField() { try (MapVector mapVector = MapVector.empty("mapVector", allocator, false)) { FieldType type = new FieldType(false, ArrowType.Struct.INSTANCE, null, null); AddOrGetResult addResult = mapVector.addOrGetVector(type); FieldType keyType = new FieldType(false, MinorType.BIGINT.getType(), null, null); FieldType valueType = FieldType.nullable(MinorType.FLOAT8.getType()); addResult.getVector().addOrGet(MapVector.KEY_NAME, keyType, BigIntVector.class); addResult.getVector().addOrGet(MapVector.VALUE_NAME, valueType, Float8Vector.class); mapVector.allocateNew(); mapVector.setValueCount(0); assertEquals(-1, mapVector.getLastSet()); TransferPair tp = mapVector.getTransferPair(mapVector.getField(), allocator); tp.transfer(); MapVector toVector = (MapVector) tp.getTo(); assertSame(toVector.getField(), mapVector.getField()); toVector.clear(); } } @Test public void testGetTransferPairWithFieldAndCallBack() { SchemaChangeCallBack callBack = new SchemaChangeCallBack(); try (MapVector mapVector = MapVector.empty("mapVector", allocator, false)) { FieldType type = new FieldType(false, ArrowType.Struct.INSTANCE, null, null); AddOrGetResult addResult = mapVector.addOrGetVector(type); FieldType keyType = new FieldType(false, MinorType.BIGINT.getType(), null, null); FieldType valueType = FieldType.nullable(MinorType.FLOAT8.getType()); addResult.getVector().addOrGet(MapVector.KEY_NAME, keyType, BigIntVector.class); addResult.getVector().addOrGet(MapVector.VALUE_NAME, valueType, Float8Vector.class); mapVector.allocateNew(); mapVector.setValueCount(0); assertEquals(-1, mapVector.getLastSet()); TransferPair tp = mapVector.getTransferPair(mapVector.getField(), allocator, callBack); tp.transfer(); MapVector toVector = (MapVector) tp.getTo(); assertSame(toVector.getField(), mapVector.getField()); toVector.clear(); } } @Test public void testMakeTransferPairPreserveNullability() { Field intField = new Field("int", FieldType.notNullable(MinorType.INT.getType()), null); List fields = Collections.singletonList(intField); Field structField = new Field("struct", FieldType.notNullable(ArrowType.Struct.INSTANCE), fields); Field structField2 = new Field("struct", FieldType.notNullable(ArrowType.Struct.INSTANCE), fields); FieldVector vec = structField.createVector(allocator); TransferPair tp = vec.getTransferPair(structField2, allocator); tp.transfer(); FieldVector res = (FieldVector) tp.getTo(); 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); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestNullCheckingForGet.java000066400000000000000000000062341515451615200327200ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import java.lang.reflect.Field; import java.net.URLClassLoader; import org.junit.jupiter.api.Test; /** Test cases for {@link NullCheckingForGet}. */ public class TestNullCheckingForGet { /** * Get a copy of the current class loader. * * @return the newly created class loader. */ private ClassLoader copyClassLoader() { ClassLoader curClassLoader = this.getClass().getClassLoader(); if (curClassLoader instanceof URLClassLoader) { // for Java 1.8 return new URLClassLoader(((URLClassLoader) curClassLoader).getURLs(), null); } // for Java 1.9 and Java 11. return null; } /** * Get the value of flag {@link NullCheckingForGet#NULL_CHECKING_ENABLED}. * * @param classLoader the class loader from which to get the flag value. * @return value of the flag. */ private boolean getFlagValue(ClassLoader classLoader) throws Exception { Class clazz = classLoader.loadClass("org.apache.arrow.vector.NullCheckingForGet"); Field field = clazz.getField("NULL_CHECKING_ENABLED"); return (Boolean) field.get(null); } /** * Ensure the flag for null checking is enabled by default. This will protect users from JVM * crashes. */ @Test public void testDefaultValue() throws Exception { ClassLoader classLoader = copyClassLoader(); if (classLoader != null) { boolean nullCheckingEnabled = getFlagValue(classLoader); assertTrue(nullCheckingEnabled); } } /** * Test setting the null checking flag by the system property. * * @throws Exception if loading class {@link NullCheckingForGet#NULL_CHECKING_ENABLED} fails. */ @Test public void testEnableSysProperty() throws Exception { String sysProperty = System.getProperty("arrow.enable_null_check_for_get"); System.setProperty("arrow.enable_null_check_for_get", "false"); ClassLoader classLoader = copyClassLoader(); if (classLoader != null) { boolean nullCheckingEnabled = getFlagValue(classLoader); assertFalse(nullCheckingEnabled); } // restore system property if (sysProperty != null) { System.setProperty("arrow.enable_null_check_for_get", sysProperty); } else { System.clearProperty("arrow.enable_null_check_for_get"); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestOpaqueExtensionType.java000066400000000000000000000166111515451615200332340ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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.assertNotEquals; 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.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.ByteBuffer; import java.util.Collections; import java.util.stream.Stream; 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.InvalidExtensionMetadataException; import org.apache.arrow.vector.extension.OpaqueType; import org.apache.arrow.vector.extension.OpaqueVector; import org.apache.arrow.vector.ipc.ArrowStreamReader; import org.apache.arrow.vector.ipc.ArrowStreamWriter; 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.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.MethodSource; import org.junit.jupiter.params.provider.ValueSource; class TestOpaqueExtensionType { BufferAllocator allocator; @BeforeEach void beforeEach() { allocator = new RootAllocator(); } @AfterEach void afterEach() { allocator.close(); } @ParameterizedTest @ValueSource( strings = { "{\"type_name\": \"\", \"vendor_name\": \"\"}", "{\"type_name\": \"\", \"vendor_name\": \"\", \"extra_field\": 42}", "{\"type_name\": \"array\", \"vendor_name\": \"postgresql\"}", "{\"type_name\": \"foo.bar\", \"vendor_name\": \"postgresql\"}", }) void testDeserializeValid(String serialized) { ArrowType storageType = Types.MinorType.NULL.getType(); OpaqueType type = new OpaqueType(storageType, "", ""); assertDoesNotThrow(() -> type.deserialize(storageType, serialized)); } @ParameterizedTest @ValueSource( strings = { "", "{\"type_name\": \"\"}", "{\"vendor_name\": \"\"}", "{\"type_name\": null, \"vendor_name\": \"\"}", "{\"type_name\": \"\", \"vendor_name\": null}", "{\"type_name\": 42, \"vendor_name\": \"\"}", "{\"type_name\": \"\", \"vendor_name\": 42}", "{\"type_name\": \"\", \"vendor_name\": \"\"", }) void testDeserializeInvalid(String serialized) { ArrowType storageType = Types.MinorType.NULL.getType(); OpaqueType type = new OpaqueType(storageType, "", ""); assertThrows( InvalidExtensionMetadataException.class, () -> type.deserialize(storageType, serialized)); } @ParameterizedTest @MethodSource("storageType") void testRoundTrip(ArrowType storageType) { OpaqueType type = new OpaqueType(storageType, "foo", "bar"); assertEquals(storageType, type.storageType()); assertEquals("foo", type.typeName()); if (storageType.isComplex()) { assertThrows( UnsupportedOperationException.class, () -> type.getNewVector("name", FieldType.nullable(type), allocator)); } else { assertDoesNotThrow(() -> type.getNewVector("name", FieldType.nullable(type), allocator)) .close(); } String serialized = assertDoesNotThrow(type::serialize); OpaqueType holder = new OpaqueType(Types.MinorType.NULL.getType(), "", ""); OpaqueType deserialized = (OpaqueType) holder.deserialize(storageType, serialized); assertEquals(type, deserialized); assertNotEquals(holder, deserialized); } @ParameterizedTest @MethodSource("storageType") void testIpcRoundTrip(ArrowType storageType) { OpaqueType.ensureRegistered(); OpaqueType type = new OpaqueType(storageType, "foo", "bar"); Schema schema = new Schema(Collections.singletonList(Field.nullable("unknown", type))); byte[] serialized = schema.serializeAsMessage(); Schema deseralized = Schema.deserializeMessage(ByteBuffer.wrap(serialized)); assertEquals(schema, deseralized); } @Test void testVectorType() throws IOException { OpaqueType.ensureRegistered(); ArrowType storageType = Types.MinorType.VARBINARY.getType(); OpaqueType type = new OpaqueType(storageType, "foo", "bar"); try (FieldVector vector = type.getNewVector("field", FieldType.nullable(type), allocator)) { OpaqueVector opaque = assertInstanceOf(OpaqueVector.class, vector); assertEquals("field", opaque.getField().getName()); assertEquals(type, opaque.getField().getType()); VarBinaryVector binary = assertInstanceOf(VarBinaryVector.class, opaque.getUnderlyingVector()); binary.setSafe(0, new byte[] {0, 1, 2, 3}); binary.setNull(1); opaque.setValueCount(2); ByteArrayOutputStream baos = new ByteArrayOutputStream(); try (VectorSchemaRoot root = new VectorSchemaRoot(Collections.singletonList(opaque)); 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(2, root.getRowCount()); assertEquals(new Schema(Collections.singletonList(opaque.getField())), root.getSchema()); OpaqueVector actual = assertInstanceOf(OpaqueVector.class, root.getVector("field")); assertFalse(actual.isNull(0)); assertTrue(actual.isNull(1)); assertArrayEquals(new byte[] {0, 1, 2, 3}, (byte[]) actual.getObject(0)); assertNull(actual.getObject(1)); } } } static Stream storageType() { return Stream.of( Types.MinorType.NULL.getType(), Types.MinorType.BIGINT.getType(), Types.MinorType.BIT.getType(), Types.MinorType.VARBINARY.getType(), Types.MinorType.VARCHAR.getType(), Types.MinorType.LIST.getType(), new ArrowType.Decimal(12, 4, 128)); } } TestOutOfMemoryForValueVector.java000066400000000000000000000052011515451615200342310ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertThrows; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.OutOfMemoryException; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** This class tests cases where we expect to receive {@link OutOfMemoryException}. */ public class TestOutOfMemoryForValueVector { private static final String EMPTY_SCHEMA_PATH = ""; private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(200); // Start with low memory limit } @Test public void variableWidthVectorAllocateNew() { assertThrows( OutOfMemoryException.class, () -> { try (VarCharVector vector = new VarCharVector(EMPTY_SCHEMA_PATH, allocator)) { vector.allocateNew(); } }); } @Test public void variableWidthVectorAllocateNewCustom() { assertThrows( OutOfMemoryException.class, () -> { try (VarCharVector vector = new VarCharVector(EMPTY_SCHEMA_PATH, allocator)) { vector.allocateNew(2342, 234); } }); } @Test public void fixedWidthVectorAllocateNew() { assertThrows( OutOfMemoryException.class, () -> { try (IntVector vector = new IntVector(EMPTY_SCHEMA_PATH, allocator)) { vector.allocateNew(); } }); } @Test public void fixedWidthVectorAllocateNewCustom() { assertThrows( OutOfMemoryException.class, () -> { try (IntVector vector = new IntVector(EMPTY_SCHEMA_PATH, allocator)) { vector.allocateNew(2342); } }); } @AfterEach public void terminate() { allocator.close(); } } TestOversizedAllocationForValueVector.java000066400000000000000000000130041515451615200357640ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.memory.util.LargeMemoryUtil.checkedCastToInt; 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 org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.util.OversizedAllocationException; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** * This class tests that OversizedAllocationException occurs when a large memory is allocated for a * vector. Typically, arrow allows the allocation of the size of at most Integer.MAX_VALUE, but this * might cause OOM in tests. Thus, the max allocation size is limited to 1 KB in this class. Please * see the surefire option in pom.xml. */ public class TestOversizedAllocationForValueVector { private static final String EMPTY_SCHEMA_PATH = ""; private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testFixedVectorReallocation() { assertThrows( OversizedAllocationException.class, () -> { final UInt4Vector vector = new UInt4Vector(EMPTY_SCHEMA_PATH, allocator); // edge case 1: buffer size = max value capacity final int expectedValueCapacity = checkedCastToInt(BaseValueVector.MAX_ALLOCATION_SIZE / 4); try { vector.allocateNew(expectedValueCapacity); assertEquals(expectedValueCapacity, vector.getValueCapacity()); vector.reAlloc(); assertEquals(expectedValueCapacity * 2, vector.getValueCapacity()); } finally { vector.close(); } // common case: value count < max value capacity try { vector.allocateNew(checkedCastToInt(BaseValueVector.MAX_ALLOCATION_SIZE / 8)); vector.reAlloc(); // value allocation reaches to MAX_VALUE_ALLOCATION vector.reAlloc(); // this should throw an IOOB } finally { vector.close(); } }); } @Test public void testBitVectorReallocation() { assertThrows( OversizedAllocationException.class, () -> { final BitVector vector = new BitVector(EMPTY_SCHEMA_PATH, allocator); // edge case 1: buffer size ~ max value capacity final int expectedValueCapacity = 1 << 29; try { vector.allocateNew(expectedValueCapacity); assertEquals(expectedValueCapacity, vector.getValueCapacity()); vector.reAlloc(); assertEquals(expectedValueCapacity * 2, vector.getValueCapacity()); } finally { vector.close(); } // common: value count < MAX_VALUE_ALLOCATION try { vector.allocateNew(expectedValueCapacity); for (int i = 0; i < 3; i++) { vector.reAlloc(); // expand buffer size } assertEquals(Integer.MAX_VALUE, vector.getValueCapacity()); vector.reAlloc(); // buffer size ~ max allocation assertEquals(Integer.MAX_VALUE, vector.getValueCapacity()); vector.reAlloc(); // overflow } finally { vector.close(); } }); } @Test public void testVariableVectorReallocation() { assertThrows( OversizedAllocationException.class, () -> { final VarCharVector vector = new VarCharVector(EMPTY_SCHEMA_PATH, allocator); // edge case 1: value count = MAX_VALUE_ALLOCATION final long expectedAllocationInBytes = BaseValueVector.MAX_ALLOCATION_SIZE; final int expectedOffsetSize = 10; try { vector.allocateNew(expectedAllocationInBytes, 10); assertTrue(expectedOffsetSize <= vector.getValueCapacity()); assertTrue(expectedAllocationInBytes <= vector.getDataBuffer().capacity()); vector.reAlloc(); assertTrue(expectedOffsetSize * 2 <= vector.getValueCapacity()); assertTrue(expectedAllocationInBytes * 2 <= vector.getDataBuffer().capacity()); } finally { vector.close(); } // common: value count < MAX_VALUE_ALLOCATION try { vector.allocateNew(BaseValueVector.MAX_ALLOCATION_SIZE / 2, 0); vector.reAlloc(); // value allocation reaches to MAX_VALUE_ALLOCATION vector.reAlloc(); // this tests if it overflows } finally { vector.close(); } }); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestPeriodDuration.java000066400000000000000000000104551515451615200321730ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertNotEquals; import java.time.Duration; import java.time.LocalDate; import java.time.LocalDateTime; import java.time.Period; import java.time.temporal.ChronoUnit; import org.junit.jupiter.api.Test; public class TestPeriodDuration { @Test public void testBasics() { PeriodDuration pd1 = new PeriodDuration(Period.of(1, 2, 3), Duration.ofNanos(123)); PeriodDuration pdEq1 = new PeriodDuration(Period.of(1, 2, 3), Duration.ofNanos(123)); PeriodDuration pd2 = new PeriodDuration(Period.of(1, 2, 3), Duration.ofNanos(12)); PeriodDuration pd3 = new PeriodDuration(Period.of(-1, -2, -3), Duration.ofNanos(-123)); assertEquals(pd1, pdEq1); assertEquals(pd1.hashCode(), pdEq1.hashCode()); assertNotEquals(pd1, pd2); assertNotEquals(pd1.hashCode(), pd2.hashCode()); assertNotEquals(pd1, pd3); assertNotEquals(pd1.hashCode(), pd3.hashCode()); } @Test public void testToISO8601IntervalString() { assertEquals("P0D", new PeriodDuration(Period.ZERO, Duration.ZERO).toISO8601IntervalString()); assertEquals( "P1Y2M3D", new PeriodDuration(Period.of(1, 2, 3), Duration.ZERO).toISO8601IntervalString()); assertEquals( "PT0.000000123S", new PeriodDuration(Period.ZERO, Duration.ofNanos(123)).toISO8601IntervalString()); assertEquals( "PT1.000000123S", new PeriodDuration(Period.ZERO, Duration.ofSeconds(1).withNanos(123)) .toISO8601IntervalString()); assertEquals( "PT1H1.000000123S", new PeriodDuration(Period.ZERO, Duration.ofSeconds(3601).withNanos(123)) .toISO8601IntervalString()); assertEquals( "PT24H1M1.000000123S", new PeriodDuration(Period.ZERO, Duration.ofSeconds(86461).withNanos(123)) .toISO8601IntervalString()); assertEquals( "P1Y2M3DT24H1M1.000000123S", new PeriodDuration(Period.of(1, 2, 3), Duration.ofSeconds(86461).withNanos(123)) .toISO8601IntervalString()); assertEquals( "P-1Y-2M-3D", new PeriodDuration(Period.of(-1, -2, -3), Duration.ZERO).toISO8601IntervalString()); assertEquals( "PT-0.000000123S", new PeriodDuration(Period.ZERO, Duration.ofNanos(-123)).toISO8601IntervalString()); assertEquals( "PT-24H-1M-0.999999877S", new PeriodDuration(Period.ZERO, Duration.ofSeconds(-86461).withNanos(123)) .toISO8601IntervalString()); assertEquals( "P-1Y-2M-3DT-0.999999877S", new PeriodDuration(Period.of(-1, -2, -3), Duration.ofSeconds(-1).withNanos(123)) .toISO8601IntervalString()); } @Test public void testTemporalAccessor() { LocalDate date = LocalDate.of(2024, 1, 2); PeriodDuration pd1 = new PeriodDuration(Period.ofYears(1), Duration.ZERO); assertEquals(LocalDate.of(2025, 1, 2), pd1.addTo(date)); LocalDateTime dateTime = LocalDateTime.of(2024, 1, 2, 3, 4); PeriodDuration pd2 = new PeriodDuration(Period.ZERO, Duration.ofMinutes(1)); assertEquals(LocalDateTime.of(2024, 1, 2, 3, 3), pd2.subtractFrom(dateTime)); PeriodDuration pd3 = new PeriodDuration(Period.of(1, 2, 3), Duration.ofSeconds(86461).withNanos(123)); assertEquals(pd3.get(ChronoUnit.YEARS), 1); assertEquals(pd3.get(ChronoUnit.MONTHS), 2); assertEquals(pd3.get(ChronoUnit.DAYS), 3); assertEquals(pd3.get(ChronoUnit.SECONDS), 86461); assertEquals(pd3.get(ChronoUnit.NANOS), 123); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestRunEndEncodedVector.java000066400000000000000000000317411515451615200331040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.List; import java.util.function.Function; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.compare.Range; import org.apache.arrow.vector.compare.RangeEqualsVisitor; import org.apache.arrow.vector.complex.RunEndEncodedVector; import org.apache.arrow.vector.types.Types; 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.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 TestRunEndEncodedVector { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testInitializeChildrenFromFields() { final FieldType valueType = FieldType.notNullable(Types.MinorType.BIGINT.getType()); final FieldType runEndType = FieldType.notNullable(Types.MinorType.INT.getType()); final Field valueField = new Field("value", valueType, null); final Field runEndField = new Field("ree", runEndType, null); try (RunEndEncodedVector reeVector = RunEndEncodedVector.empty("empty", allocator)) { reeVector.initializeChildrenFromFields(List.of(runEndField, valueField)); reeVector.validate(); } } /** Create REE vector with constant value. */ @Test public void testConstantValueVector() { final Field runEndEncodedField = createBigIntRunEndEncodedField("constant"); int logicalValueCount = 100; // constant vector try (RunEndEncodedVector reeVector = new RunEndEncodedVector(runEndEncodedField, allocator, null)) { Long value = 65536L; setConstantVector(reeVector, value, logicalValueCount); assertEquals(logicalValueCount, reeVector.getValueCount()); for (int i = 0; i < logicalValueCount; i++) { assertEquals(value, reeVector.getObject(i)); } } // constant null vector try (RunEndEncodedVector reeVector = new RunEndEncodedVector(runEndEncodedField, allocator, null)) { setConstantVector(reeVector, null, logicalValueCount); assertEquals(logicalValueCount, reeVector.getValueCount()); // Null count is always 0 for run-end encoded array assertEquals(0, reeVector.getNullCount()); for (int i = 0; i < logicalValueCount; i++) { assertTrue(reeVector.isNull(i)); assertNull(reeVector.getObject(i)); } } } @Test public void testBasicRunEndEncodedVector() { try (RunEndEncodedVector reeVector = new RunEndEncodedVector(createBigIntRunEndEncodedField("basic"), allocator, null)) { // Create REE vector representing: // [null, 2, 2, null, null, null, 4, 4, 4, 4, null, null, null, null, null]. int runCount = 5; final int logicalValueCount = setBasicVector(reeVector, runCount, i -> i % 2 == 0 ? null : i + 1, i -> i + 1); assertEquals(15, reeVector.getValueCount()); checkBasic(runCount, reeVector); // test index out of bound assertThrows(IndexOutOfBoundsException.class, () -> reeVector.getObject(-1)); assertThrows(IndexOutOfBoundsException.class, () -> reeVector.getObject(logicalValueCount)); } } private static void checkBasic(int runCount, RunEndEncodedVector reeVector) { int index = 0; for (int run = 0; run < runCount; run++) { long expectedRunValue = (long) run + 1; for (int j = 0; j <= run; j++) { if (run % 2 == 0) { assertNull(reeVector.getObject(index)); } else { assertEquals(expectedRunValue, reeVector.getObject(index)); } index++; } } } @Test public void testRangeCompare() { // test compare same constant vector RunEndEncodedVector constantVector = new RunEndEncodedVector(createBigIntRunEndEncodedField("constant"), allocator, null); int logicalValueCount = 15; setConstantVector(constantVector, 1L, logicalValueCount); assertTrue( constantVector.accept( new RangeEqualsVisitor(constantVector, constantVector), new Range(0, 0, logicalValueCount))); assertTrue( constantVector.accept( new RangeEqualsVisitor(constantVector, constantVector), new Range(1, 1, 14))); assertTrue( constantVector.accept( new RangeEqualsVisitor(constantVector, constantVector), new Range(1, 2, 13))); // 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 = new RunEndEncodedVector(createBigIntRunEndEncodedField("basic"), allocator, null); setBasicVector(reeVector, 5, i -> i + 1, i -> i + 1); assertTrue( reeVector.accept( new RangeEqualsVisitor(reeVector, reeVector), new Range(0, 0, logicalValueCount))); assertTrue( reeVector.accept( new RangeEqualsVisitor(reeVector, reeVector), new Range(2, 2, logicalValueCount - 2))); assertFalse( reeVector.accept( new RangeEqualsVisitor(reeVector, reeVector), new Range(1, 2, logicalValueCount - 2))); assertFalse( reeVector.accept( new RangeEqualsVisitor(reeVector, constantVector), new Range(0, 0, logicalValueCount))); // Create REE vector representing: [2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 5]. RunEndEncodedVector reeVector2 = new RunEndEncodedVector(createBigIntRunEndEncodedField("basic"), allocator, null); setBasicVector(reeVector2, 4, i -> i + 2, i -> i + 2); assertTrue( reeVector.accept( new RangeEqualsVisitor(reeVector, reeVector2), new Range(1, 0, logicalValueCount - 1))); constantVector.close(); reeVector.close(); reeVector2.close(); } private static Field createBigIntRunEndEncodedField(String fieldName) { final FieldType valueType = FieldType.notNullable(Types.MinorType.BIGINT.getType()); final FieldType runEndType = FieldType.notNullable(Types.MinorType.INT.getType()); final Field valueField = new Field("value", valueType, null); final Field runEndField = new Field("ree", runEndType, null); return new Field( fieldName, FieldType.notNullable(RunEndEncoded.INSTANCE), List.of(runEndField, valueField)); } private static void setConstantVector( RunEndEncodedVector constantVector, Long value, long logicalValueCount) { setBasicVector(constantVector, 1, i -> value, i -> logicalValueCount); } private static int setBasicVector( RunEndEncodedVector reeVector, int runCount, Function runValueSupplier, Function runLengthSupplier) { reeVector.allocateNew(); reeVector.setInitialCapacity(runCount); int end = 0; for (int i = 0; i < runCount; i++) { Long runValue = runValueSupplier.apply((long) i); if (runValue == null) { reeVector.getValuesVector().setNull(i); } else { ((BigIntVector) reeVector.getValuesVector()).set(i, runValue); } Long runLength = runLengthSupplier.apply((long) i); assert runLength != null && runLength > 0; end += runLength; ((IntVector) reeVector.getRunEndsVector()).set(i, end); } final int logicalValueCount = end; reeVector.getValuesVector().setValueCount(runCount); reeVector.getRunEndsVector().setValueCount(runCount); reeVector.setValueCount(logicalValueCount); return logicalValueCount; } @Test public void testTransfer() { // constant vector try (RunEndEncodedVector reeVector = new RunEndEncodedVector(createBigIntRunEndEncodedField("constant"), allocator, null)) { Long value = 65536L; int logicalValueCount = 100; setConstantVector(reeVector, value, logicalValueCount); assertEquals(logicalValueCount, reeVector.getValueCount()); for (int i = 0; i < logicalValueCount; i++) { assertEquals(value, reeVector.getObject(i)); } TransferPair transferPair = reeVector.getTransferPair(allocator); transferPair.transfer(); assertEquals(0, reeVector.getValueCount()); assertEquals(0, reeVector.getValuesVector().getValueCount()); assertEquals(0, reeVector.getRunEndsVector().getValueCount()); try (RunEndEncodedVector toVector = (RunEndEncodedVector) transferPair.getTo()) { assertEquals(logicalValueCount, toVector.getValueCount()); for (int i = 0; i < logicalValueCount; i++) { assertEquals(value, toVector.getObject(i)); } } } // basic run end encoded vector try (RunEndEncodedVector reeVector = new RunEndEncodedVector(createBigIntRunEndEncodedField("basic"), allocator, null)) { // Create REE vector representing: // [null, 2, 2, null, null, null, 4, 4, 4, 4, null, null, null, null, null]. int runCount = 5; final int logicalValueCount = setBasicVector(reeVector, runCount, i -> i % 2 == 0 ? null : i + 1, i -> i + 1); assertEquals(15, reeVector.getValueCount()); checkBasic(runCount, reeVector); TransferPair transferPair = reeVector.getTransferPair(allocator); transferPair.transfer(); assertEquals(0, reeVector.getValueCount()); assertEquals(0, reeVector.getValuesVector().getValueCount()); assertEquals(0, reeVector.getRunEndsVector().getValueCount()); try (RunEndEncodedVector toVector = (RunEndEncodedVector) transferPair.getTo()) { assertEquals(logicalValueCount, toVector.getValueCount()); checkBasic(runCount, toVector); } } } @Test public void testSplitAndTransfer() { // test compare same constant vector try (RunEndEncodedVector constantVector = new RunEndEncodedVector(createBigIntRunEndEncodedField("constant"), allocator, null)) { int logicalValueCount = 15; setConstantVector(constantVector, 1L, logicalValueCount); try (RunEndEncodedVector toVector = RunEndEncodedVector.empty("constant", allocator)) { TransferPair transferPair = constantVector.makeTransferPair(toVector); int startIndex = 1; int transferLength = 10; transferPair.splitAndTransfer(startIndex, transferLength); toVector.validate(); assertEquals(transferLength, toVector.getValueCount()); assertTrue( constantVector.accept( new RangeEqualsVisitor(constantVector, toVector), new Range(1, 0, transferLength))); } } try (RunEndEncodedVector reeVector = new RunEndEncodedVector(createBigIntRunEndEncodedField("ree"), allocator, null)) { setBasicVector(reeVector, 5, i -> i + 1, i -> i + 1); int[][] transferConfigs = {{0, 0}, {0, 1}, {0, 9}, {1, 0}, {1, 10}, {1, 14}}; try (RunEndEncodedVector toVector = RunEndEncodedVector.empty("ree", allocator)) { TransferPair transferPair = reeVector.makeTransferPair(toVector); for (final int[] transferConfig : transferConfigs) { int startIndex = transferConfig[0]; int transferLength = transferConfig[1]; transferPair.splitAndTransfer(startIndex, transferLength); toVector.validate(); assertEquals(transferLength, toVector.getValueCount()); assertTrue( reeVector.accept( new RangeEqualsVisitor(reeVector, toVector), new Range(startIndex, 0, transferLength))); } } } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestSplitAndTransfer.java000066400000000000000000001051441515451615200324660ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.Arrays.asList; 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.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.LargeListViewVector; 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.UnionVector; 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; 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 TestSplitAndTransfer { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void terminate() throws Exception { allocator.close(); } private void populateVarcharVector( final VarCharVector vector, int valueCount, String[] compareArray) { for (int i = 0; i < valueCount; i += 3) { final String s = String.format("%010d", i); vector.set(i, s.getBytes(StandardCharsets.UTF_8)); if (compareArray != null) { compareArray[i] = s; } } vector.setValueCount(valueCount); } private void populateBaseVariableWidthViewVector( final BaseVariableWidthViewVector vector, int valueCount, String[] compareArray) { for (int i = 0; i < valueCount; i += 3) { final String s = String.format("%010d", i); vector.set(i, s.getBytes(StandardCharsets.UTF_8)); if (compareArray != null) { compareArray[i] = s; } } vector.setValueCount(valueCount); } private void populateIntVector(final IntVector vector, int valueCount) { for (int i = 0; i < valueCount; i++) { vector.set(i, i); } vector.setValueCount(valueCount); } private void populateDenseUnionVector(final DenseUnionVector vector, int valueCount) { VarCharVector varCharVector = vector.addOrGet("varchar", FieldType.nullable(new ArrowType.Utf8()), VarCharVector.class); BigIntVector intVector = vector.addOrGet("int", FieldType.nullable(new ArrowType.Int(64, true)), BigIntVector.class); for (int i = 0; i < valueCount; i++) { vector.setTypeId(i, (byte) (i % 2)); if (i % 2 == 0) { final String s = String.format("%010d", i); varCharVector.setSafe(i / 2, s.getBytes(StandardCharsets.UTF_8)); } else { intVector.setSafe(i / 2, i); } } vector.setValueCount(valueCount); } @Test public void testWithEmptyVector() { // MapVector use TransferImpl from ListVector ListVector listVector = ListVector.empty("", allocator); TransferPair transferPair = listVector.getTransferPair(allocator); transferPair.splitAndTransfer(0, 0); assertEquals(0, transferPair.getTo().getValueCount()); // BaseFixedWidthVector IntVector intVector = new IntVector("", allocator); transferPair = intVector.getTransferPair(allocator); transferPair.splitAndTransfer(0, 0); assertEquals(0, transferPair.getTo().getValueCount()); // BaseVariableWidthVector VarCharVector varCharVector = new VarCharVector("", allocator); transferPair = varCharVector.getTransferPair(allocator); transferPair.splitAndTransfer(0, 0); assertEquals(0, transferPair.getTo().getValueCount()); // BaseVariableWidthViewVector: ViewVarCharVector ViewVarCharVector viewVarCharVector = new ViewVarCharVector("", allocator); transferPair = viewVarCharVector.getTransferPair(allocator); transferPair.splitAndTransfer(0, 0); assertEquals(0, transferPair.getTo().getValueCount()); // BaseVariableWidthVector: ViewVarBinaryVector ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("", allocator); transferPair = viewVarBinaryVector.getTransferPair(allocator); transferPair.splitAndTransfer(0, 0); assertEquals(0, transferPair.getTo().getValueCount()); // BaseLargeVariableWidthVector LargeVarCharVector largeVarCharVector = new LargeVarCharVector("", allocator); transferPair = largeVarCharVector.getTransferPair(allocator); transferPair.splitAndTransfer(0, 0); assertEquals(0, transferPair.getTo().getValueCount()); // StructVector StructVector structVector = StructVector.empty("", allocator); transferPair = structVector.getTransferPair(allocator); transferPair.splitAndTransfer(0, 0); assertEquals(0, transferPair.getTo().getValueCount()); // FixedSizeListVector FixedSizeListVector fixedSizeListVector = FixedSizeListVector.empty("", 0, allocator); transferPair = fixedSizeListVector.getTransferPair(allocator); transferPair.splitAndTransfer(0, 0); assertEquals(0, transferPair.getTo().getValueCount()); // FixedSizeBinaryVector FixedSizeBinaryVector fixedSizeBinaryVector = new FixedSizeBinaryVector("", allocator, 4); transferPair = fixedSizeBinaryVector.getTransferPair(allocator); transferPair.splitAndTransfer(0, 0); assertEquals(0, transferPair.getTo().getValueCount()); // UnionVector UnionVector unionVector = UnionVector.empty("", allocator); transferPair = unionVector.getTransferPair(allocator); transferPair.splitAndTransfer(0, 0); assertEquals(0, transferPair.getTo().getValueCount()); // DenseUnionVector DenseUnionVector duv = DenseUnionVector.empty("", allocator); transferPair = duv.getTransferPair(allocator); transferPair.splitAndTransfer(0, 0); assertEquals(0, transferPair.getTo().getValueCount()); // non empty from vector // BaseFixedWidthVector IntVector fromIntVector = new IntVector("", allocator); fromIntVector.allocateNew(100); populateIntVector(fromIntVector, 100); transferPair = fromIntVector.getTransferPair(allocator); IntVector toIntVector = (IntVector) transferPair.getTo(); transferPair.splitAndTransfer(0, 0); assertEquals(0, toIntVector.getValueCount()); transferPair.splitAndTransfer(50, 0); assertEquals(0, toIntVector.getValueCount()); transferPair.splitAndTransfer(100, 0); assertEquals(0, toIntVector.getValueCount()); fromIntVector.clear(); toIntVector.clear(); // DenseUnionVector DenseUnionVector fromDuv = DenseUnionVector.empty("", allocator); populateDenseUnionVector(fromDuv, 100); transferPair = fromDuv.getTransferPair(allocator); DenseUnionVector toDUV = (DenseUnionVector) transferPair.getTo(); transferPair.splitAndTransfer(0, 0); assertEquals(0, toDUV.getValueCount()); transferPair.splitAndTransfer(50, 0); assertEquals(0, toDUV.getValueCount()); transferPair.splitAndTransfer(100, 0); assertEquals(0, toDUV.getValueCount()); fromDuv.clear(); toDUV.clear(); } @Test public void testWithNullVector() { int valueCount = 123; int startIndex = 10; NullVector fromNullVector = new NullVector("nullVector"); fromNullVector.setValueCount(valueCount); TransferPair transferPair = fromNullVector.getTransferPair(fromNullVector.getAllocator()); transferPair.splitAndTransfer(startIndex, valueCount - startIndex); NullVector toNullVector = (NullVector) transferPair.getTo(); assertEquals(valueCount - startIndex, toNullVector.getValueCount()); // no allocations to clear for NullVector } @Test public void testWithZeroVector() { ZeroVector fromZeroVector = new ZeroVector("zeroVector"); TransferPair transferPair = fromZeroVector.getTransferPair(fromZeroVector.getAllocator()); transferPair.splitAndTransfer(0, 0); ZeroVector toZeroVector = (ZeroVector) transferPair.getTo(); assertEquals(0, toZeroVector.getValueCount()); // no allocations to clear for ZeroVector } @Test public void testListVectorWithEmptyMapVector() { // List not null>> int valueCount = 1; List children = new ArrayList<>(); children.add(new Field("key", FieldType.notNullable(new ArrowType.Utf8()), null)); children.add(new Field("value", FieldType.nullable(new ArrowType.Utf8()), null)); Field structField = new Field("entries", FieldType.notNullable(ArrowType.Struct.INSTANCE), children); Field mapField = new Field("element", FieldType.notNullable(new ArrowType.Map(false)), asList(structField)); Field listField = new Field("list", FieldType.nullable(new ArrowType.List()), asList(mapField)); ListVector fromListVector = (ListVector) listField.createVector(allocator); fromListVector.allocateNew(); fromListVector.setValueCount(valueCount); // child vector is empty MapVector dataVector = (MapVector) fromListVector.getDataVector(); dataVector.allocateNew(); // unset capacity to mimic observed failure mode dataVector.getOffsetBuffer().capacity(0); TransferPair transferPair = fromListVector.getTransferPair(fromListVector.getAllocator()); transferPair.splitAndTransfer(0, valueCount); ListVector toListVector = (ListVector) transferPair.getTo(); assertEquals(valueCount, toListVector.getValueCount()); fromListVector.clear(); toListVector.clear(); } @Test /* VarCharVector */ public void test() throws Exception { try (final VarCharVector varCharVector = new VarCharVector("myvector", allocator)) { varCharVector.allocateNew(10000, 1000); final int valueCount = 500; final String[] compareArray = new String[valueCount]; populateVarcharVector(varCharVector, valueCount, compareArray); final TransferPair tp = varCharVector.getTransferPair(allocator); final VarCharVector newVarCharVector = (VarCharVector) tp.getTo(); final int[][] startLengths = {{0, 201}, {201, 0}, {201, 200}, {401, 99}}; for (final int[] startLength : startLengths) { final int start = startLength[0]; final int length = startLength[1]; tp.splitAndTransfer(start, length); for (int i = 0; i < length; i++) { final boolean expectedSet = ((start + i) % 3) == 0; if (expectedSet) { final byte[] expectedValue = compareArray[start + i].getBytes(StandardCharsets.UTF_8); assertFalse(newVarCharVector.isNull(i)); assertArrayEquals(expectedValue, newVarCharVector.get(i)); } else { assertTrue(newVarCharVector.isNull(i)); } } newVarCharVector.clear(); } } } private void testView(BaseVariableWidthViewVector vector) { vector.allocateNew(10000, 1000); final int valueCount = 500; final String[] compareArray = new String[valueCount]; populateBaseVariableWidthViewVector(vector, valueCount, compareArray); final TransferPair tp = vector.getTransferPair(allocator); final BaseVariableWidthViewVector newVector = (BaseVariableWidthViewVector) tp.getTo(); ; final int[][] startLengths = {{0, 201}, {201, 0}, {201, 200}, {401, 99}}; for (final int[] startLength : startLengths) { final int start = startLength[0]; final int length = startLength[1]; tp.splitAndTransfer(start, length); for (int i = 0; i < length; i++) { final boolean expectedSet = ((start + i) % 3) == 0; if (expectedSet) { final byte[] expectedValue = compareArray[start + i].getBytes(StandardCharsets.UTF_8); assertFalse(newVector.isNull(i)); assertArrayEquals(expectedValue, newVector.get(i)); } else { assertTrue(newVector.isNull(i)); } } newVector.clear(); } } @Test public void testUtf8View() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { testView(viewVarCharVector); } } @Test public void testBinaryView() throws Exception { try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator)) { testView(viewVarBinaryVector); } } @Test public void testMemoryConstrainedTransfer() { try (final VarCharVector varCharVector = new VarCharVector("myvector", allocator)) { allocator.setLimit(32768); /* set limit of 32KB */ varCharVector.allocateNew(10000, 1000); final int valueCount = 1000; populateVarcharVector(varCharVector, valueCount, null); final TransferPair tp = varCharVector.getTransferPair(allocator); final VarCharVector newVarCharVector = (VarCharVector) tp.getTo(); final int[][] startLengths = {{0, 700}, {700, 299}}; for (final int[] startLength : startLengths) { final int start = startLength[0]; final int length = startLength[1]; tp.splitAndTransfer(start, length); newVarCharVector.clear(); } } } private void testMemoryConstrainedTransferInViews(BaseVariableWidthViewVector vector) { // Here we have the target vector being transferred with a long string // hence, the data buffer will be allocated. // The default data buffer allocation takes // BaseVariableWidthViewVector.INITIAL_VIEW_VALUE_ALLOCATION * // BaseVariableWidthViewVector.ELEMENT_SIZE // set limit = BaseVariableWidthViewVector.INITIAL_VIEW_VALUE_ALLOCATION * // BaseVariableWidthViewVector.ELEMENT_SIZE final int setLimit = BaseVariableWidthViewVector.INITIAL_VIEW_VALUE_ALLOCATION * BaseVariableWidthViewVector.ELEMENT_SIZE; allocator.setLimit(setLimit); vector.allocateNew(16000, 1000); final int valueCount = 1000; populateBaseVariableWidthViewVector(vector, valueCount, null); final TransferPair tp = vector.getTransferPair(allocator); final BaseVariableWidthViewVector newVector = (BaseVariableWidthViewVector) tp.getTo(); final int[][] startLengths = {{0, 700}, {700, 299}}; for (final int[] startLength : startLengths) { final int start = startLength[0]; final int length = startLength[1]; tp.splitAndTransfer(start, length); newVector.clear(); } } @Test public void testMemoryConstrainedTransferInUtf8Views() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { testMemoryConstrainedTransferInViews(viewVarCharVector); } } @Test public void testMemoryConstrainedTransferInBinaryViews() { try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator)) { testMemoryConstrainedTransferInViews(viewVarBinaryVector); } } @Test public void testTransfer() { try (final VarCharVector varCharVector = new VarCharVector("myvector", allocator)) { varCharVector.allocateNew(10000, 1000); final int valueCount = 500; final String[] compareArray = new String[valueCount]; populateVarcharVector(varCharVector, valueCount, compareArray); final TransferPair tp = varCharVector.getTransferPair(allocator); final VarCharVector newVarCharVector = (VarCharVector) tp.getTo(); tp.transfer(); assertEquals(0, varCharVector.valueCount); assertEquals(valueCount, newVarCharVector.valueCount); for (int i = 0; i < valueCount; i++) { final boolean expectedSet = (i % 3) == 0; if (expectedSet) { final byte[] expectedValue = compareArray[i].getBytes(StandardCharsets.UTF_8); assertFalse(newVarCharVector.isNull(i)); assertArrayEquals(expectedValue, newVarCharVector.get(i)); } else { assertTrue(newVarCharVector.isNull(i)); } } newVarCharVector.clear(); } } private void testTransferInViews(BaseVariableWidthViewVector vector) { vector.allocateNew(16000, 1000); final int valueCount = 500; final String[] compareArray = new String[valueCount]; populateBaseVariableWidthViewVector(vector, valueCount, compareArray); final TransferPair tp = vector.getTransferPair(allocator); final BaseVariableWidthViewVector newVector = (BaseVariableWidthViewVector) tp.getTo(); tp.transfer(); assertEquals(0, vector.valueCount); assertEquals(valueCount, newVector.valueCount); for (int i = 0; i < valueCount; i++) { final boolean expectedSet = (i % 3) == 0; if (expectedSet) { final byte[] expectedValue = compareArray[i].getBytes(StandardCharsets.UTF_8); assertFalse(newVector.isNull(i)); assertArrayEquals(expectedValue, newVector.get(i)); } else { assertTrue(newVector.isNull(i)); } } newVector.clear(); } @Test public void testTransferInUtf8Views() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { testTransferInViews(viewVarCharVector); } } @Test public void testTransferInBinaryViews() { try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator)) { testTransferInViews(viewVarBinaryVector); } } @Test public void testCopyValueSafe() { try (final VarCharVector varCharVector = new VarCharVector("myvector", allocator); final VarCharVector newVarCharVector = new VarCharVector("newvector", allocator)) { varCharVector.allocateNew(10000, 1000); final int valueCount = 500; populateVarcharVector(varCharVector, valueCount, null); final TransferPair tp = varCharVector.makeTransferPair(newVarCharVector); // new vector memory is not pre-allocated, we expect copyValueSafe work fine. for (int i = 0; i < valueCount; i++) { tp.copyValueSafe(i, i); } newVarCharVector.setValueCount(valueCount); for (int i = 0; i < valueCount; i++) { final boolean expectedSet = (i % 3) == 0; if (expectedSet) { assertFalse(varCharVector.isNull(i)); assertFalse(newVarCharVector.isNull(i)); assertArrayEquals(varCharVector.get(i), newVarCharVector.get(i)); } else { assertTrue(newVarCharVector.isNull(i)); } } } } @Test public void testSplitAndTransferNon() { try (final VarCharVector varCharVector = new VarCharVector("myvector", allocator)) { varCharVector.allocateNew(10000, 1000); final int valueCount = 500; populateVarcharVector(varCharVector, valueCount, null); final TransferPair tp = varCharVector.getTransferPair(allocator); VarCharVector newVarCharVector = (VarCharVector) tp.getTo(); tp.splitAndTransfer(0, 0); assertEquals(0, newVarCharVector.getValueCount()); newVarCharVector.clear(); } } private void testSplitAndTransferNonInViews(BaseVariableWidthViewVector vector) { vector.allocateNew(16000, 1000); final int valueCount = 500; populateBaseVariableWidthViewVector(vector, valueCount, null); final TransferPair tp = vector.getTransferPair(allocator); BaseVariableWidthViewVector newVector = (BaseVariableWidthViewVector) tp.getTo(); tp.splitAndTransfer(0, 0); assertEquals(0, newVector.getValueCount()); newVector.clear(); } @Test public void testSplitAndTransferNonInUtf8Views() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { testSplitAndTransferNonInViews(viewVarCharVector); } } @Test public void testSplitAndTransferNonInBinaryViews() { try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator)) { testSplitAndTransferNonInViews(viewVarBinaryVector); } } @Test public void testSplitAndTransferAll() { try (final VarCharVector varCharVector = new VarCharVector("myvector", allocator)) { varCharVector.allocateNew(10000, 1000); final int valueCount = 500; populateVarcharVector(varCharVector, valueCount, null); final TransferPair tp = varCharVector.getTransferPair(allocator); VarCharVector newVarCharVector = (VarCharVector) tp.getTo(); tp.splitAndTransfer(0, valueCount); assertEquals(valueCount, newVarCharVector.getValueCount()); newVarCharVector.clear(); } } private void testSplitAndTransferAllInViews(BaseVariableWidthViewVector vector) { vector.allocateNew(16000, 1000); final int valueCount = 500; populateBaseVariableWidthViewVector(vector, valueCount, null); final TransferPair tp = vector.getTransferPair(allocator); BaseVariableWidthViewVector newViewVarCharVector = (BaseVariableWidthViewVector) tp.getTo(); tp.splitAndTransfer(0, valueCount); assertEquals(valueCount, newViewVarCharVector.getValueCount()); newViewVarCharVector.clear(); } @Test public void testSplitAndTransferAllInUtf8Views() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { testSplitAndTransferAllInViews(viewVarCharVector); } } @Test public void testSplitAndTransferAllInBinaryViews() { try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator)) { testSplitAndTransferAllInViews(viewVarBinaryVector); } } @Test public void testInvalidStartIndex() { try (final VarCharVector varCharVector = new VarCharVector("myvector", allocator); final VarCharVector newVarCharVector = new VarCharVector("newvector", allocator)) { varCharVector.allocateNew(10000, 1000); final int valueCount = 500; populateVarcharVector(varCharVector, valueCount, null); final TransferPair tp = varCharVector.makeTransferPair(newVarCharVector); IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () -> tp.splitAndTransfer(valueCount, 10)); assertEquals( "Invalid parameters startIndex: 500, length: 10 for valueCount: 500", e.getMessage()); newVarCharVector.clear(); } } private void testInvalidStartIndexInViews( BaseVariableWidthViewVector vector, BaseVariableWidthViewVector newVector) { vector.allocateNew(16000, 1000); final int valueCount = 500; populateBaseVariableWidthViewVector(vector, valueCount, null); final TransferPair tp = vector.makeTransferPair(newVector); IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () -> tp.splitAndTransfer(valueCount, 10)); assertEquals( "Invalid parameters startIndex: 500, length: 10 for valueCount: 500", e.getMessage()); newVector.clear(); } @Test public void testInvalidStartIndexInUtf8Views() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator); final ViewVarCharVector newViewVarCharVector = new ViewVarCharVector("newvector", allocator)) { testInvalidStartIndexInViews(viewVarCharVector, newViewVarCharVector); } } @Test public void testInvalidStartIndexInBinaryViews() { try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator); final ViewVarBinaryVector newViewVarBinaryVector = new ViewVarBinaryVector("newvector", allocator)) { testInvalidStartIndexInViews(viewVarBinaryVector, newViewVarBinaryVector); } } @Test public void testInvalidLength() { try (final VarCharVector varCharVector = new VarCharVector("myvector", allocator); final VarCharVector newVarCharVector = new VarCharVector("newvector", allocator)) { varCharVector.allocateNew(10000, 1000); final int valueCount = 500; populateVarcharVector(varCharVector, valueCount, null); final TransferPair tp = varCharVector.makeTransferPair(newVarCharVector); IllegalArgumentException e = assertThrows( IllegalArgumentException.class, () -> tp.splitAndTransfer(0, valueCount * 2)); assertEquals( "Invalid parameters startIndex: 0, length: 1000 for valueCount: 500", e.getMessage()); newVarCharVector.clear(); } } private void testInvalidLengthInViews( BaseVariableWidthViewVector vector, BaseVariableWidthViewVector newVector) { vector.allocateNew(16000, 1000); final int valueCount = 500; populateBaseVariableWidthViewVector(vector, valueCount, null); final TransferPair tp = vector.makeTransferPair(newVector); IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () -> tp.splitAndTransfer(0, valueCount * 2)); assertEquals( "Invalid parameters startIndex: 0, length: 1000 for valueCount: 500", e.getMessage()); newVector.clear(); } @Test public void testInvalidLengthInUtf8Views() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator); final ViewVarCharVector newViewVarCharVector = new ViewVarCharVector("newvector", allocator)) { testInvalidLengthInViews(viewVarCharVector, newViewVarCharVector); } } @Test public void testInvalidLengthInBinaryViews() { try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator); final ViewVarBinaryVector newViewVarBinaryVector = new ViewVarBinaryVector("newvector", allocator)) { testInvalidLengthInViews(viewVarBinaryVector, newViewVarBinaryVector); } } @Test public void testZeroStartIndexAndLength() { try (final VarCharVector varCharVector = new VarCharVector("myvector", allocator); final VarCharVector newVarCharVector = new VarCharVector("newvector", allocator)) { varCharVector.allocateNew(0, 0); final int valueCount = 0; populateVarcharVector(varCharVector, valueCount, null); final TransferPair tp = varCharVector.makeTransferPair(newVarCharVector); tp.splitAndTransfer(0, 0); assertEquals(valueCount, newVarCharVector.getValueCount()); newVarCharVector.clear(); } } private void testZeroStartIndexAndLengthInViews( BaseVariableWidthViewVector vector, BaseVariableWidthViewVector newVector) { vector.allocateNew(0, 0); final int valueCount = 0; populateBaseVariableWidthViewVector(vector, valueCount, null); final TransferPair tp = vector.makeTransferPair(newVector); tp.splitAndTransfer(0, 0); assertEquals(valueCount, newVector.getValueCount()); newVector.clear(); } @Test public void testZeroStartIndexAndLengthInUtf8Views() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator); final ViewVarCharVector newViewVarCharVector = new ViewVarCharVector("newvector", allocator)) { testZeroStartIndexAndLengthInViews(viewVarCharVector, newViewVarCharVector); } } @Test public void testZeroStartIndexAndLengthInBinaryViews() { try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator); final ViewVarBinaryVector newViewVarBinaryVector = new ViewVarBinaryVector("newvector", allocator)) { testZeroStartIndexAndLengthInViews(viewVarBinaryVector, newViewVarBinaryVector); } } @Test public void testZeroLength() { try (final VarCharVector varCharVector = new VarCharVector("myvector", allocator); final VarCharVector newVarCharVector = new VarCharVector("newvector", allocator)) { varCharVector.allocateNew(10000, 1000); final int valueCount = 500; populateVarcharVector(varCharVector, valueCount, null); final TransferPair tp = varCharVector.makeTransferPair(newVarCharVector); tp.splitAndTransfer(500, 0); assertEquals(0, newVarCharVector.getValueCount()); newVarCharVector.clear(); } } private void testZeroLengthInViews( BaseVariableWidthViewVector vector, BaseVariableWidthViewVector newVector) { vector.allocateNew(16000, 1000); final int valueCount = 500; populateBaseVariableWidthViewVector(vector, valueCount, null); final TransferPair tp = vector.makeTransferPair(newVector); tp.splitAndTransfer(500, 0); assertEquals(0, newVector.getValueCount()); newVector.clear(); } @Test public void testZeroLengthInUtf8Views() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator); final ViewVarCharVector newViewVarCharVector = new ViewVarCharVector("newvector", allocator)) { testZeroLengthInViews(viewVarCharVector, newViewVarCharVector); } } @Test public void testZeroLengthInBinaryViews() { try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator); final ViewVarBinaryVector newViewVarBinaryVector = new ViewVarBinaryVector("newvector", allocator)) { testZeroLengthInViews(viewVarBinaryVector, newViewVarBinaryVector); } } @Test public void testUnionVectorZeroStartIndexAndLength() { try (final UnionVector unionVector = UnionVector.empty("myvector", allocator); final UnionVector newUnionVector = UnionVector.empty("newvector", allocator)) { unionVector.allocateNew(); final int valueCount = 0; unionVector.setValueCount(valueCount); final TransferPair tp = unionVector.makeTransferPair(newUnionVector); tp.splitAndTransfer(0, 0); assertEquals(valueCount, newUnionVector.getValueCount()); newUnionVector.clear(); } } @Test public void testFixedWidthVectorZeroStartIndexAndLength() { try (final IntVector intVector = new IntVector("myvector", allocator); final IntVector newIntVector = new IntVector("newvector", allocator)) { intVector.allocateNew(0); final int valueCount = 0; intVector.setValueCount(valueCount); final TransferPair tp = intVector.makeTransferPair(newIntVector); tp.splitAndTransfer(0, 0); assertEquals(valueCount, newIntVector.getValueCount()); newIntVector.clear(); } } @Test public void testBitVectorZeroStartIndexAndLength() { try (final BitVector bitVector = new BitVector("myvector", allocator); final BitVector newBitVector = new BitVector("newvector", allocator)) { bitVector.allocateNew(0); final int valueCount = 0; bitVector.setValueCount(valueCount); final TransferPair tp = bitVector.makeTransferPair(newBitVector); tp.splitAndTransfer(0, 0); assertEquals(valueCount, newBitVector.getValueCount()); newBitVector.clear(); } } @Test public void testFixedSizeListVectorZeroStartIndexAndLength() { try (final FixedSizeListVector listVector = FixedSizeListVector.empty("list", 4, allocator); final FixedSizeListVector newListVector = FixedSizeListVector.empty("newList", 4, allocator)) { listVector.allocateNew(); final int valueCount = 0; listVector.setValueCount(valueCount); final TransferPair tp = listVector.makeTransferPair(newListVector); tp.splitAndTransfer(0, 0); assertEquals(valueCount, newListVector.getValueCount()); newListVector.clear(); } } @Test public void testListVectorZeroStartIndexAndLength() { try (final ListVector listVector = ListVector.empty("list", allocator); final ListVector newListVector = ListVector.empty("newList", allocator)) { listVector.allocateNew(); final int valueCount = 0; listVector.setValueCount(valueCount); final TransferPair tp = listVector.makeTransferPair(newListVector); tp.splitAndTransfer(0, 0); assertEquals(valueCount, newListVector.getValueCount()); newListVector.clear(); } } @Test public void testLargeListViewVectorZeroStartIndexAndLength() { try (final LargeListViewVector listVector = LargeListViewVector.empty("largelistview", allocator); final LargeListViewVector newListVector = LargeListViewVector.empty("newList", allocator)) { listVector.allocateNew(); final int valueCount = 0; listVector.setValueCount(valueCount); final TransferPair tp = listVector.makeTransferPair(newListVector); tp.splitAndTransfer(0, 0); assertEquals(valueCount, newListVector.getValueCount()); newListVector.clear(); } } @Test public void testStructVectorZeroStartIndexAndLength() { Map metadata = new HashMap<>(); metadata.put("k1", "v1"); FieldType type = new FieldType(true, Struct.INSTANCE, null, metadata); try (final StructVector structVector = new StructVector("structvec", allocator, type, null); final StructVector newStructVector = new StructVector("newStructvec", allocator, type, null)) { structVector.allocateNew(); final int valueCount = 0; structVector.setValueCount(valueCount); final TransferPair tp = structVector.makeTransferPair(newStructVector); tp.splitAndTransfer(0, 0); assertEquals(valueCount, newStructVector.getValueCount()); newStructVector.clear(); } } @Test public void testMapVectorZeroStartIndexAndLength() { Map metadata = new HashMap<>(); metadata.put("k1", "v1"); FieldType type = new FieldType(true, new ArrowType.Map(false), null, metadata); try (final MapVector mapVector = new MapVector("mapVec", allocator, type, null); final MapVector newMapVector = new MapVector("newMapVec", allocator, type, null)) { mapVector.allocateNew(); final int valueCount = 0; mapVector.setValueCount(valueCount); final TransferPair tp = mapVector.makeTransferPair(newMapVector); tp.splitAndTransfer(0, 0); assertEquals(valueCount, newMapVector.getValueCount()); newMapVector.clear(); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java000066400000000000000000000427361515451615200317210ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertNotEquals; 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 java.util.ArrayList; 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; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.UnionVector; 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; 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 TestStructVector { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testFieldMetadata() throws Exception { Map metadata = new HashMap<>(); metadata.put("k1", "v1"); FieldType type = new FieldType(true, Struct.INSTANCE, null, metadata); try (StructVector vector = new StructVector("struct", allocator, type, null)) { assertEquals(vector.getField().getMetadata(), type.getMetadata()); } } @Test public void testMakeTransferPair() { try (final StructVector s1 = StructVector.empty("s1", allocator); final StructVector s2 = StructVector.empty("s2", allocator)) { s1.addOrGet("struct_child", FieldType.nullable(MinorType.INT.getType()), IntVector.class); s1.makeTransferPair(s2); final FieldVector child = s1.getChild("struct_child"); final FieldVector toChild = s2.addOrGet("struct_child", child.getField().getFieldType(), child.getClass()); assertEquals(0, toChild.getValueCapacity()); assertEquals(0, toChild.getDataBuffer().capacity()); assertEquals(0, toChild.getValidityBuffer().capacity()); } } @Test public void testAllocateAfterReAlloc() throws Exception { Map metadata = new HashMap<>(); metadata.put("k1", "v1"); FieldType type = new FieldType(true, Struct.INSTANCE, null, metadata); try (StructVector vector = new StructVector("struct", allocator, type, null)) { MinorType childtype = MinorType.INT; vector.addOrGet("intchild", FieldType.nullable(childtype.getType()), IntVector.class); /* * Allocate the default size, and then, reAlloc. This should double the allocation. */ vector.allocateNewSafe(); // Initial allocation vector.reAlloc(); // Double the allocation size of self, and all children. long savedValidityBufferCapacity = vector.getValidityBuffer().capacity(); int savedValueCapacity = vector.getValueCapacity(); /* * Clear and allocate again. */ vector.clear(); vector.allocateNewSafe(); /* * Verify that the buffer sizes haven't changed. */ assertEquals(vector.getValidityBuffer().capacity(), savedValidityBufferCapacity); assertEquals(vector.getValueCapacity(), savedValueCapacity); } } @Test public void testReadNullValue() { Map metadata = new HashMap<>(); metadata.put("k1", "v1"); FieldType type = new FieldType(true, Struct.INSTANCE, null, metadata); try (StructVector vector = new StructVector("struct", allocator, type, null)) { MinorType childtype = MinorType.INT; vector.addOrGet("intchild", FieldType.nullable(childtype.getType()), IntVector.class); vector.setValueCount(2); IntVector intVector = (IntVector) vector.getChild("intchild"); intVector.setSafe(0, 100); vector.setIndexDefined(0); intVector.setNull(1); vector.setNull(1); ComplexHolder holder = new ComplexHolder(); vector.get(0, holder); assertNotEquals(0, holder.isSet); assertNotNull(holder.reader); vector.get(1, holder); assertEquals(0, holder.isSet); assertNull(holder.reader); } } @Test public void testGetPrimitiveVectors() { FieldType type = new FieldType(true, Struct.INSTANCE, null, null); try (StructVector vector = new StructVector("struct", allocator, type, null)) { // add list vector vector.addOrGet("list", FieldType.nullable(MinorType.LIST.getType()), ListVector.class); ListVector listVector = vector.addOrGetList("list"); listVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); // add union vector vector.addOrGet("union", FieldType.nullable(MinorType.UNION.getType()), UnionVector.class); 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(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.EXTENSIONTYPE, primitiveVectors.get(3).getMinorType()); assertEquals(MinorType.VARCHAR, primitiveVectors.get(4).getMinorType()); assertEquals(MinorType.EXTENSIONTYPE, primitiveVectors.get(5).getMinorType()); } } @Test public void testAddOrGetComplexChildVectors() { FieldType type = new FieldType(true, Struct.INSTANCE, null, null); try (StructVector vector = new StructVector("struct", allocator, type, null)) { vector.addOrGetList("list"); vector.addOrGetFixedSizeList("fixedList", 2); vector.addOrGetUnion("union"); vector.addOrGetStruct("struct"); vector.addOrGetMap("map", true); List children = vector.getChildrenFromFields(); assertEquals(5, children.size()); assertEquals(MinorType.LIST, children.get(0).getMinorType()); assertEquals(MinorType.FIXED_SIZE_LIST, children.get(1).getMinorType()); assertEquals(MinorType.UNION, children.get(2).getMinorType()); assertEquals(MinorType.STRUCT, children.get(3).getMinorType()); assertEquals(MinorType.MAP, children.get(4).getMinorType()); } } @Test public void testAddChildVectorsWithDuplicatedFieldNamesForConflictPolicyAppend() { final FieldType type = new FieldType(true, Struct.INSTANCE, null, null); try (StructVector vector = new StructVector( "struct", allocator, type, null, AbstractStructVector.ConflictPolicy.CONFLICT_APPEND, true)) { final List initFields = new ArrayList<>(); // Add a bit more fields to test against stability of the internal field // ordering mechanism of StructVector initFields.add(Field.nullable("varchar1", MinorType.VARCHAR.getType())); initFields.add(Field.nullable("int1", MinorType.INT.getType())); initFields.add(Field.nullable("varchar2", MinorType.VARCHAR.getType())); initFields.add(Field.nullable("int2", MinorType.INT.getType())); initFields.add(Field.nullable("varchar3", MinorType.VARCHAR.getType())); initFields.add(Field.nullable("int3", MinorType.INT.getType())); initFields.add(Field.nullable("uncertain-type", MinorType.INT.getType())); // To ensure duplicated field names don't mess up the original field order // in the struct vector initFields.add(Field.nullable("varchar1", MinorType.VARCHAR.getType())); initFields.add(Field.nullable("varchar2", MinorType.VARCHAR.getType())); initFields.add(Field.nullable("varchar3", MinorType.VARCHAR.getType())); initFields.add(Field.nullable("uncertain-type", MinorType.VARCHAR.getType())); vector.initializeChildrenFromFields(initFields); List children = vector.getChildrenFromFields(); assertEquals(11, children.size()); assertEquals("varchar1", children.get(0).getName()); assertEquals("int1", children.get(1).getName()); assertEquals("varchar2", children.get(2).getName()); assertEquals("int2", children.get(3).getName()); assertEquals("varchar3", children.get(4).getName()); assertEquals("int3", children.get(5).getName()); assertEquals("uncertain-type", children.get(6).getName()); assertEquals("varchar1", children.get(7).getName()); assertEquals("varchar2", children.get(8).getName()); assertEquals("varchar3", children.get(9).getName()); assertEquals("uncertain-type", children.get(10).getName()); assertEquals(MinorType.VARCHAR, children.get(0).getMinorType()); assertEquals(MinorType.INT, children.get(1).getMinorType()); assertEquals(MinorType.VARCHAR, children.get(2).getMinorType()); assertEquals(MinorType.INT, children.get(3).getMinorType()); assertEquals(MinorType.VARCHAR, children.get(4).getMinorType()); assertEquals(MinorType.INT, children.get(5).getMinorType()); assertEquals(MinorType.INT, children.get(6).getMinorType()); assertEquals(MinorType.VARCHAR, children.get(7).getMinorType()); assertEquals(MinorType.VARCHAR, children.get(8).getMinorType()); assertEquals(MinorType.VARCHAR, children.get(9).getMinorType()); assertEquals(MinorType.VARCHAR, children.get(10).getMinorType()); } } @Test public void testAddChildVectorsWithDuplicatedFieldNamesForConflictPolicyReplace() { final FieldType type = new FieldType(true, Struct.INSTANCE, null, null); try (StructVector vector = new StructVector( "struct", allocator, type, null, AbstractStructVector.ConflictPolicy.CONFLICT_REPLACE, true)) { final List initFields = new ArrayList<>(); // Add a bit more fields to test against stability of the internal field // ordering mechanism of StructVector initFields.add(Field.nullable("varchar1", MinorType.VARCHAR.getType())); initFields.add(Field.nullable("int1", MinorType.INT.getType())); initFields.add(Field.nullable("varchar2", MinorType.VARCHAR.getType())); initFields.add(Field.nullable("int2", MinorType.INT.getType())); initFields.add(Field.nullable("varchar3", MinorType.VARCHAR.getType())); initFields.add(Field.nullable("int3", MinorType.INT.getType())); initFields.add(Field.nullable("uncertain-type", MinorType.INT.getType())); // To ensure duplicated field names don't mess up the original field order // in the struct vector initFields.add(Field.nullable("varchar1", MinorType.VARCHAR.getType())); initFields.add(Field.nullable("varchar2", MinorType.VARCHAR.getType())); initFields.add(Field.nullable("varchar3", MinorType.VARCHAR.getType())); initFields.add(Field.nullable("uncertain-type", MinorType.VARCHAR.getType())); vector.initializeChildrenFromFields(initFields); List children = vector.getChildrenFromFields(); assertEquals(7, children.size()); assertEquals("varchar1", children.get(0).getName()); assertEquals("int1", children.get(1).getName()); assertEquals("varchar2", children.get(2).getName()); assertEquals("int2", children.get(3).getName()); assertEquals("varchar3", children.get(4).getName()); assertEquals("int3", children.get(5).getName()); assertEquals("uncertain-type", children.get(6).getName()); assertEquals(MinorType.VARCHAR, children.get(0).getMinorType()); assertEquals(MinorType.INT, children.get(1).getMinorType()); assertEquals(MinorType.VARCHAR, children.get(2).getMinorType()); assertEquals(MinorType.INT, children.get(3).getMinorType()); assertEquals(MinorType.VARCHAR, children.get(4).getMinorType()); assertEquals(MinorType.INT, children.get(5).getMinorType()); assertEquals(MinorType.VARCHAR, children.get(6).getMinorType()); } } @Test public void testTypedGetters() { try (final StructVector s1 = StructVector.empty("s1", allocator)) { s1.addOrGet("struct_child", FieldType.nullable(MinorType.INT.getType()), IntVector.class); assertEquals(IntVector.class, s1.getChild("struct_child", IntVector.class).getClass()); assertEquals(IntVector.class, s1.getVectorById(0, IntVector.class).getClass()); } } @Test public void testGetTransferPair() { try (final StructVector fromVector = simpleStructVector("s1", allocator)) { TransferPair tp = fromVector.getTransferPair(fromVector.getField(), allocator); final StructVector toVector = (StructVector) tp.getTo(); // Field inside a new vector created by reusing a field should be the same in memory as the // original field. assertSame(toVector.getField(), fromVector.getField()); toVector.clear(); } } @Test public void testGetTransferPairWithFieldAndCallBack() { SchemaChangeCallBack callBack = new SchemaChangeCallBack(); try (final StructVector fromVector = simpleStructVector("s1", allocator)) { TransferPair tp = fromVector.getTransferPair(fromVector.getField(), allocator, callBack); final StructVector toVector = (StructVector) tp.getTo(); // Field inside a new vector created by reusing a field should be the same in memory as the // original field. assertSame(toVector.getField(), fromVector.getField()); toVector.clear(); } } @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"; StructVector structVector = StructVector.empty(name, allocator); final int size = 6; // number of structs NullableStructWriter structWriter = structVector.getWriter(); structVector.addOrGet( INT_COL, FieldType.nullable(Types.MinorType.INT.getType()), IntVector.class); structVector.addOrGet( FLT_COL, FieldType.nullable(Types.MinorType.INT.getType()), IntVector.class); structVector.allocateNew(); IntWriter intWriter = structWriter.integer(INT_COL); Float8Writer float8Writer = structWriter.float8(FLT_COL); for (int i = 0; i < size; i++) { structWriter.setPosition(i); structWriter.start(); intWriter.writeInt(i); float8Writer.writeFloat8(i * .1); structWriter.end(); } structWriter.setValueCount(size); return structVector; } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestTypeLayout.java000066400000000000000000000155471515451615200313710ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.Random; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.IntervalUnit; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.UnionMode; 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; public class TestTypeLayout { private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(Integer.MAX_VALUE); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testTypeBufferCount() { ArrowType type = new ArrowType.Int(8, true); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.Union(UnionMode.Sparse, new int[2]); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.Union(UnionMode.Dense, new int[1]); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.Struct(); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.Timestamp(TimeUnit.MILLISECOND, null); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.List(); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.FixedSizeList(5); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.Map(false); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.Decimal(10, 10, 128); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.Decimal(10, 10, 256); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.FixedSizeBinary(5); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.Bool(); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.Binary(); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.Utf8(); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.Null(); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.Date(DateUnit.DAY); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.Time(TimeUnit.MILLISECOND, 32); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.Interval(IntervalUnit.DAY_TIME); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); type = new ArrowType.Duration(TimeUnit.MILLISECOND); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); } 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 } return sb.toString(); } @Test public void testTypeBufferCountInVectorsWithVariadicBuffers() { // empty vector try (ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { ArrowType type = viewVarCharVector.getMinorType().getType(); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); } // vector with long strings try (ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { viewVarCharVector.allocateNew(32, 6); viewVarCharVector.setSafe(0, generateRandomString(8).getBytes()); viewVarCharVector.setSafe(1, generateRandomString(12).getBytes()); viewVarCharVector.setSafe(2, generateRandomString(14).getBytes()); viewVarCharVector.setSafe(3, generateRandomString(18).getBytes()); viewVarCharVector.setSafe(4, generateRandomString(22).getBytes()); viewVarCharVector.setSafe(5, generateRandomString(24).getBytes()); viewVarCharVector.setValueCount(6); ArrowType type = viewVarCharVector.getMinorType().getType(); assertEquals( TypeLayout.getTypeBufferCount(type), TypeLayout.getTypeLayout(type).getBufferLayouts().size()); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestUnionVector.java000066400000000000000000000456361515451615200315270ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.complex.VectorWithOrdinal; import org.apache.arrow.vector.complex.impl.UnionWriter; import org.apache.arrow.vector.holders.NullableBitHolder; import org.apache.arrow.vector.holders.NullableFloat4Holder; import org.apache.arrow.vector.holders.NullableIntHolder; import org.apache.arrow.vector.holders.NullableUInt4Holder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.UnionMode; 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.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestUnionVector { private static final String EMPTY_SCHEMA_PATH = ""; private BufferAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testUnionVector() throws Exception { final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); uInt4Holder.value = 100; uInt4Holder.isSet = 1; try (UnionVector unionVector = new UnionVector( EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) { unionVector.allocateNew(); // write some data unionVector.setType(0, MinorType.UINT4); unionVector.setSafe(0, uInt4Holder); unionVector.setType(2, MinorType.UINT4); unionVector.setSafe(2, uInt4Holder); unionVector.setValueCount(4); // check that what we wrote is correct assertEquals(4, unionVector.getValueCount()); assertEquals(false, unionVector.isNull(0)); assertEquals(100, unionVector.getObject(0)); assertNull(unionVector.getObject(1)); assertEquals(false, unionVector.isNull(2)); assertEquals(100, unionVector.getObject(2)); assertNull(unionVector.getObject(3)); } } @Test public void testUnionVectorMapValue() throws Exception { try (UnionVector unionVector = new UnionVector( EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) { unionVector.allocateNew(); UnionWriter writer = (UnionWriter) unionVector.getWriter(); // populate map vector with the following two records // [ // null, // [[1: 2], [3: 4], [5: null]] // ] writer.setPosition(0); writer.writeNull(); writer.setPosition(1); writer.startMap(); writer.startEntry(); writer.key().integer().writeInt(1); writer.value().integer().writeInt(2); writer.endEntry(); writer.startEntry(); writer.key().integer().writeInt(3); writer.value().integer().writeInt(4); writer.endEntry(); writer.startEntry(); writer.key().integer().writeInt(5); writer.endEntry(); writer.endMap(); unionVector.setValueCount(2); // check that what we wrote is correct assertEquals(2, unionVector.getValueCount()); // first entry assertNull(unionVector.getObject(0)); // second entry List> resultList = (List>) unionVector.getObject(1); assertEquals(3, resultList.size()); Map resultMap = resultList.get(0); assertEquals(1, (int) resultMap.get(MapVector.KEY_NAME)); assertEquals(2, (int) resultMap.get(MapVector.VALUE_NAME)); resultMap = resultList.get(1); assertEquals(3, (int) resultMap.get(MapVector.KEY_NAME)); assertEquals(4, (int) resultMap.get(MapVector.VALUE_NAME)); resultMap = resultList.get(2); assertEquals(5, (int) resultMap.get(MapVector.KEY_NAME)); assertNull(resultMap.get(MapVector.VALUE_NAME)); } } @Test public void testTransfer() throws Exception { try (UnionVector srcVector = new UnionVector( EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) { srcVector.allocateNew(); // write some data srcVector.setType(0, MinorType.INT); srcVector.setSafe(0, newIntHolder(5)); srcVector.setType(1, MinorType.BIT); srcVector.setSafe(1, newBitHolder(false)); srcVector.setType(3, MinorType.INT); srcVector.setSafe(3, newIntHolder(10)); srcVector.setType(5, MinorType.BIT); srcVector.setSafe(5, newBitHolder(false)); srcVector.setValueCount(6); try (UnionVector destVector = new UnionVector( EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) { TransferPair pair = srcVector.makeTransferPair(destVector); // Creating the transfer should transfer the type of the field at least. assertEquals(srcVector.getField(), destVector.getField()); // transfer pair.transfer(); assertEquals(srcVector.getField(), destVector.getField()); // now check the values are transferred assertEquals(6, destVector.getValueCount()); assertFalse(destVector.isNull(0)); assertEquals(5, destVector.getObject(0)); assertFalse(destVector.isNull(1)); assertEquals(false, destVector.getObject(1)); assertNull(destVector.getObject(2)); assertFalse(destVector.isNull(3)); assertEquals(10, destVector.getObject(3)); assertNull(destVector.getObject(4)); assertFalse(destVector.isNull(5)); assertEquals(false, destVector.getObject(5)); } } } @Test public void testSplitAndTransfer() throws Exception { try (UnionVector sourceVector = new UnionVector( EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) { sourceVector.allocateNew(); /* populate the UnionVector */ sourceVector.setType(0, MinorType.INT); sourceVector.setSafe(0, newIntHolder(5)); sourceVector.setType(1, MinorType.INT); sourceVector.setSafe(1, newIntHolder(10)); sourceVector.setType(2, MinorType.INT); sourceVector.setSafe(2, newIntHolder(15)); sourceVector.setType(3, MinorType.INT); sourceVector.setSafe(3, newIntHolder(20)); sourceVector.setType(4, MinorType.INT); sourceVector.setSafe(4, newIntHolder(25)); sourceVector.setType(5, MinorType.INT); sourceVector.setSafe(5, newIntHolder(30)); sourceVector.setType(6, MinorType.INT); sourceVector.setSafe(6, newIntHolder(35)); sourceVector.setType(7, MinorType.INT); sourceVector.setSafe(7, newIntHolder(40)); sourceVector.setType(8, MinorType.INT); sourceVector.setSafe(8, newIntHolder(45)); sourceVector.setType(9, MinorType.INT); sourceVector.setSafe(9, newIntHolder(50)); sourceVector.setValueCount(10); /* check the vector output */ assertEquals(10, sourceVector.getValueCount()); assertEquals(false, sourceVector.isNull(0)); assertEquals(5, sourceVector.getObject(0)); assertEquals(false, sourceVector.isNull(1)); assertEquals(10, sourceVector.getObject(1)); assertEquals(false, sourceVector.isNull(2)); assertEquals(15, sourceVector.getObject(2)); assertEquals(false, sourceVector.isNull(3)); assertEquals(20, sourceVector.getObject(3)); assertEquals(false, sourceVector.isNull(4)); assertEquals(25, sourceVector.getObject(4)); assertEquals(false, sourceVector.isNull(5)); assertEquals(30, sourceVector.getObject(5)); assertEquals(false, sourceVector.isNull(6)); assertEquals(35, sourceVector.getObject(6)); assertEquals(false, sourceVector.isNull(7)); assertEquals(40, sourceVector.getObject(7)); assertEquals(false, sourceVector.isNull(8)); assertEquals(45, sourceVector.getObject(8)); assertEquals(false, sourceVector.isNull(9)); assertEquals(50, sourceVector.getObject(9)); try (UnionVector toVector = new UnionVector( EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) { final TransferPair transferPair = sourceVector.makeTransferPair(toVector); final int[][] transferLengths = {{0, 3}, {3, 1}, {4, 2}, {6, 1}, {7, 1}, {8, 2}}; for (final int[] transferLength : transferLengths) { final int start = transferLength[0]; final int length = transferLength[1]; transferPair.splitAndTransfer(start, length); /* check the toVector output after doing the splitAndTransfer */ for (int i = 0; i < length; i++) { assertEquals( sourceVector.getObject(start + i), toVector.getObject(i), "Different data at indexes: " + (start + i) + "and " + i); } } } } } @Test public void testSplitAndTransferWithMixedVectors() throws Exception { try (UnionVector sourceVector = new UnionVector( EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) { sourceVector.allocateNew(); /* populate the UnionVector */ sourceVector.setType(0, MinorType.INT); sourceVector.setSafe(0, newIntHolder(5)); sourceVector.setType(1, MinorType.FLOAT4); sourceVector.setSafe(1, newFloat4Holder(5.5f)); sourceVector.setType(2, MinorType.INT); sourceVector.setSafe(2, newIntHolder(10)); sourceVector.setType(3, MinorType.FLOAT4); sourceVector.setSafe(3, newFloat4Holder(10.5f)); sourceVector.setType(4, MinorType.INT); sourceVector.setSafe(4, newIntHolder(15)); sourceVector.setType(5, MinorType.FLOAT4); sourceVector.setSafe(5, newFloat4Holder(15.5f)); sourceVector.setType(6, MinorType.INT); sourceVector.setSafe(6, newIntHolder(20)); sourceVector.setType(7, MinorType.FLOAT4); sourceVector.setSafe(7, newFloat4Holder(20.5f)); sourceVector.setType(8, MinorType.INT); sourceVector.setSafe(8, newIntHolder(30)); sourceVector.setType(9, MinorType.FLOAT4); sourceVector.setSafe(9, newFloat4Holder(30.5f)); sourceVector.setValueCount(10); /* check the vector output */ assertEquals(10, sourceVector.getValueCount()); assertEquals(false, sourceVector.isNull(0)); assertEquals(5, sourceVector.getObject(0)); assertEquals(false, sourceVector.isNull(1)); assertEquals(5.5f, sourceVector.getObject(1)); assertEquals(false, sourceVector.isNull(2)); assertEquals(10, sourceVector.getObject(2)); assertEquals(false, sourceVector.isNull(3)); assertEquals(10.5f, sourceVector.getObject(3)); assertEquals(false, sourceVector.isNull(4)); assertEquals(15, sourceVector.getObject(4)); assertEquals(false, sourceVector.isNull(5)); assertEquals(15.5f, sourceVector.getObject(5)); assertEquals(false, sourceVector.isNull(6)); assertEquals(20, sourceVector.getObject(6)); assertEquals(false, sourceVector.isNull(7)); assertEquals(20.5f, sourceVector.getObject(7)); assertEquals(false, sourceVector.isNull(8)); assertEquals(30, sourceVector.getObject(8)); assertEquals(false, sourceVector.isNull(9)); assertEquals(30.5f, sourceVector.getObject(9)); try (UnionVector toVector = new UnionVector( EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) { final TransferPair transferPair = sourceVector.makeTransferPair(toVector); final int[][] transferLengths = {{0, 2}, {2, 1}, {3, 2}, {5, 3}, {8, 2}}; for (final int[] transferLength : transferLengths) { final int start = transferLength[0]; final int length = transferLength[1]; transferPair.splitAndTransfer(start, length); /* check the toVector output after doing the splitAndTransfer */ for (int i = 0; i < length; i++) { assertEquals( sourceVector.getObject(start + i), toVector.getObject(i), "Different values at index: " + i); } } } } } @Test public void testGetFieldTypeInfo() throws Exception { Map metadata = new HashMap<>(); metadata.put("key1", "value1"); int[] typeIds = new int[2]; typeIds[0] = MinorType.INT.ordinal(); typeIds[1] = MinorType.VARCHAR.ordinal(); List children = new ArrayList<>(); children.add(new Field("int", FieldType.nullable(MinorType.INT.getType()), null)); children.add(new Field("varchar", FieldType.nullable(MinorType.VARCHAR.getType()), null)); final FieldType fieldType = new FieldType( false, new ArrowType.Union(UnionMode.Sparse, typeIds), /* dictionary= */ null, metadata); final Field field = new Field("union", fieldType, children); MinorType minorType = MinorType.UNION; UnionVector vector = (UnionVector) minorType.getNewVector(field, allocator, null); vector.initializeChildrenFromFields(children); assertTrue(vector.getField().equals(field)); // Union has 2 child vectors assertEquals(2, vector.size()); // Check child field 0 VectorWithOrdinal intChild = vector.getChildVectorWithOrdinal("int"); assertEquals(0, intChild.ordinal); assertEquals(intChild.vector.getField(), children.get(0)); // Check child field 1 VectorWithOrdinal varcharChild = vector.getChildVectorWithOrdinal("varchar"); assertEquals(1, varcharChild.ordinal); assertEquals(varcharChild.vector.getField(), children.get(1)); } @Test public void testGetBufferAddress() throws Exception { try (UnionVector vector = new UnionVector( EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) { boolean error = false; vector.allocateNew(); /* populate the UnionVector */ vector.setType(0, MinorType.INT); vector.setSafe(0, newIntHolder(5)); vector.setType(1, MinorType.FLOAT4); vector.setSafe(1, newFloat4Holder(5.5f)); vector.setType(2, MinorType.INT); vector.setSafe(2, newIntHolder(10)); vector.setType(3, MinorType.FLOAT4); vector.setSafe(3, newFloat4Holder(10.5f)); vector.setValueCount(10); /* check the vector output */ assertEquals(10, vector.getValueCount()); assertEquals(false, vector.isNull(0)); assertEquals(5, vector.getObject(0)); assertEquals(false, vector.isNull(1)); assertEquals(5.5f, vector.getObject(1)); assertEquals(false, vector.isNull(2)); assertEquals(10, vector.getObject(2)); assertEquals(false, vector.isNull(3)); assertEquals(10.5f, vector.getObject(3)); List buffers = vector.getFieldBuffers(); try { vector.getOffsetBufferAddress(); } catch (UnsupportedOperationException ue) { error = true; } finally { assertTrue(error); error = false; } try { vector.getDataBufferAddress(); } catch (UnsupportedOperationException ue) { error = true; } finally { assertTrue(error); } assertEquals(1, buffers.size()); } } @Test public void testSetGetNull() { try (UnionVector srcVector = new UnionVector( EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) { srcVector.allocateNew(); final NullableIntHolder holder = new NullableIntHolder(); holder.isSet = 1; holder.value = 5; // write some data srcVector.setType(0, MinorType.INT); srcVector.setSafe(0, holder); assertFalse(srcVector.isNull(0)); holder.isSet = 0; srcVector.setSafe(0, holder); assertNull(srcVector.getObject(0)); } } @Test public void testCreateNewVectorWithoutTypeExceptionThrown() { try (UnionVector vector = new UnionVector( EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) { IllegalArgumentException e1 = assertThrows(IllegalArgumentException.class, () -> vector.getTimeStampMilliTZVector()); assertEquals( "No TimeStampMilliTZ present. Provide ArrowType argument to create a new vector", e1.getMessage()); IllegalArgumentException e2 = assertThrows(IllegalArgumentException.class, () -> vector.getDurationVector()); assertEquals( "No Duration present. Provide ArrowType argument to create a new vector", e2.getMessage()); IllegalArgumentException e3 = assertThrows(IllegalArgumentException.class, () -> vector.getFixedSizeBinaryVector()); assertEquals( "No FixedSizeBinary present. Provide ArrowType argument to create a new vector", e3.getMessage()); IllegalArgumentException e4 = assertThrows(IllegalArgumentException.class, () -> vector.getDecimalVector()); assertEquals( "No Decimal present. Provide ArrowType argument to create a new vector", e4.getMessage()); } } private static NullableIntHolder newIntHolder(int value) { final NullableIntHolder holder = new NullableIntHolder(); holder.isSet = 1; holder.value = value; return holder; } private static NullableBitHolder newBitHolder(boolean value) { final NullableBitHolder holder = new NullableBitHolder(); holder.isSet = 1; holder.value = value ? 1 : 0; return holder; } private static NullableFloat4Holder newFloat4Holder(float value) { final NullableFloat4Holder holder = new NullableFloat4Holder(); holder.isSet = 1; holder.value = value; return holder; } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestUtils.java000066400000000000000000000067411515451615200303460ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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 { public static VarCharVector newVarCharVector(String name, BufferAllocator allocator) { return (VarCharVector) FieldType.nullable(new ArrowType.Utf8()).createNewSingleVector(name, allocator, null); } public static ViewVarCharVector newViewVarCharVector(String name, BufferAllocator allocator) { return (ViewVarCharVector) FieldType.nullable(new ArrowType.Utf8View()).createNewSingleVector(name, allocator, null); } public static VarBinaryVector newVarBinaryVector(String name, BufferAllocator allocator) { return (VarBinaryVector) FieldType.nullable(new ArrowType.Binary()).createNewSingleVector(name, allocator, null); } public static ViewVarBinaryVector newViewVarBinaryVector(String name, BufferAllocator allocator) { return (ViewVarBinaryVector) FieldType.nullable(new ArrowType.BinaryView()).createNewSingleVector(name, allocator, null); } public static T newVector( Class c, String name, ArrowType type, BufferAllocator allocator) { return c.cast(FieldType.nullable(type).createNewSingleVector(name, allocator, null)); } 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; } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestUuidType.java000066400000000000000000000223511515451615200310110ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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()); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestUuidVector.java000066400000000000000000000574321515451615200313420ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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)); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java000066400000000000000000004327431515451615200315120ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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; import static org.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector; 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.nio.ByteBuffer; import java.nio.charset.Charset; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.memory.rounding.DefaultRoundingPolicy; import org.apache.arrow.memory.util.ArrowBufPointer; import org.apache.arrow.memory.util.CommonUtil; import org.apache.arrow.vector.compare.Range; import org.apache.arrow.vector.compare.RangeEqualsVisitor; import org.apache.arrow.vector.compare.VectorEqualsVisitor; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.FixedSizeListVector; 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.StructVector; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.complex.impl.NullableStructWriter; import org.apache.arrow.vector.complex.impl.UnionLargeListViewWriter; 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; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.UnionMode; 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.OversizedAllocationException; import org.apache.arrow.vector.util.ReusableByteArray; import org.apache.arrow.vector.util.Text; 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 TestValueVector { private static final String EMPTY_SCHEMA_PATH = ""; private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } private static final Charset utf8Charset = StandardCharsets.UTF_8; private static final byte[] STR1 = "AAAAA1".getBytes(utf8Charset); private static final byte[] STR2 = "BBBBBBBBB2".getBytes(utf8Charset); private static final byte[] STR3 = "CCCC3".getBytes(utf8Charset); private static final byte[] STR4 = "DDDDDDDD4".getBytes(utf8Charset); 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) / 9); private static final int MAX_VALUE_COUNT_8BYTE = (int) (MAX_VALUE_COUNT / 2); @AfterEach public void terminate() throws Exception { allocator.close(); } /* * Tests for Fixed-Width vectors * * Covered types as of now * * -- UInt4Vector * -- IntVector * -- Float4Vector * -- Float8Vector * * -- UInt4Vector * -- IntVector * -- Float4Vector * * TODO: * * -- SmallIntVector * -- BigIntVector * -- TinyIntVector */ @Test /* UInt4Vector */ public void testFixedType1() { // Create a new value vector for 1024 integers. try (final UInt4Vector vector = new UInt4Vector(EMPTY_SCHEMA_PATH, allocator)) { boolean error = false; vector.allocateNew(1024); int initialCapacity = vector.getValueCapacity(); assertTrue(initialCapacity >= 1024); // Put and set a few values vector.setSafe(0, 100); vector.setSafe(1, 101); vector.setSafe(100, 102); vector.setSafe(1022, 103); vector.setSafe(1023, 104); assertEquals(100, vector.get(0)); assertEquals(101, vector.get(1)); assertEquals(102, vector.get(100)); assertEquals(103, vector.get(1022)); assertEquals(104, vector.get(1023)); try { vector.set(initialCapacity, 10000); } catch (IndexOutOfBoundsException ie) { error = true; } finally { assertTrue(error); error = false; } try { vector.get(initialCapacity); } catch (IndexOutOfBoundsException ie) { error = true; } finally { assertTrue(error); error = false; } /* this should trigger a realloc() */ vector.setSafe(initialCapacity, 10000); /* underlying buffer should now be able to store double the number of values */ assertTrue(vector.getValueCapacity() >= 2 * initialCapacity); /* check vector data after realloc */ assertEquals(100, vector.get(0)); assertEquals(101, vector.get(1)); assertEquals(102, vector.get(100)); assertEquals(103, vector.get(1022)); assertEquals(104, vector.get(1023)); assertEquals(10000, vector.get(initialCapacity)); /* reset the vector */ int capacityBeforeReset = vector.getValueCapacity(); vector.reset(); /* capacity shouldn't change after reset */ assertEquals(capacityBeforeReset, vector.getValueCapacity()); /* vector data should have been zeroed out */ for (int i = 0; i < capacityBeforeReset; i++) { // TODO: test vector.get(i) is 0 after unsafe get added assertTrue(vector.isNull(i), "non-zero data not expected at index: " + i); } } } @Test public void testNoOverFlowWithUINT() { try (final UInt8Vector uInt8Vector = new UInt8Vector("uint8", allocator); final UInt4Vector uInt4Vector = new UInt4Vector("uint4", allocator); final UInt1Vector uInt1Vector = new UInt1Vector("uint1", allocator)) { long[] longValues = new long[] {Long.MIN_VALUE, Long.MAX_VALUE, -1L}; uInt8Vector.allocateNew(3); uInt8Vector.setValueCount(3); for (int i = 0; i < longValues.length; i++) { uInt8Vector.set(i, longValues[i]); long readValue = uInt8Vector.getObjectNoOverflow(i).longValue(); assertEquals(readValue, longValues[i]); } int[] intValues = new int[] {Integer.MIN_VALUE, Integer.MAX_VALUE, -1}; uInt4Vector.allocateNew(3); uInt4Vector.setValueCount(3); for (int i = 0; i < intValues.length; i++) { uInt4Vector.set(i, intValues[i]); int actualValue = (int) UInt4Vector.getNoOverflow(uInt4Vector.getDataBuffer(), i); assertEquals(intValues[i], actualValue); } byte[] byteValues = new byte[] {Byte.MIN_VALUE, Byte.MAX_VALUE, -1}; uInt1Vector.allocateNew(3); uInt1Vector.setValueCount(3); for (int i = 0; i < byteValues.length; i++) { uInt1Vector.set(i, byteValues[i]); byte actualValue = (byte) UInt1Vector.getNoOverflow(uInt1Vector.getDataBuffer(), i); assertEquals(byteValues[i], actualValue); } } } @Test /* IntVector */ public void testFixedType2() { try (final IntVector intVector = new IntVector(EMPTY_SCHEMA_PATH, allocator)) { boolean error = false; int initialCapacity = 16; /* we should not throw exception for these values of capacity */ intVector.setInitialCapacity(MAX_VALUE_COUNT - 1); intVector.setInitialCapacity(MAX_VALUE_COUNT); try { intVector.setInitialCapacity(MAX_VALUE_COUNT * 2); } catch (OversizedAllocationException oe) { error = true; } finally { assertTrue(error); error = false; } intVector.setInitialCapacity(initialCapacity); /* no memory allocation has happened yet so capacity of underlying buffer should be 0 */ assertEquals(0, intVector.getValueCapacity()); /* allocate 64 bytes (16 * 4) */ intVector.allocateNew(); /* underlying buffer should be able to store 16 values */ assertTrue(intVector.getValueCapacity() >= initialCapacity); initialCapacity = intVector.getValueCapacity(); /* populate the vector */ int j = 1; for (int i = 0; i < initialCapacity; i += 2) { intVector.set(i, j); j++; } try { intVector.set(initialCapacity, j); } catch (IndexOutOfBoundsException ie) { error = true; } finally { assertTrue(error); error = false; } /* check vector contents */ j = 1; for (int i = 0; i < initialCapacity; i += 2) { assertEquals(j, intVector.get(i), "unexpected value at index: " + i); j++; } try { intVector.get(initialCapacity); } catch (IndexOutOfBoundsException ie) { error = true; } finally { assertTrue(error); error = false; } /* this should trigger a realloc() */ intVector.setSafe(initialCapacity, j); /* underlying buffer should now be able to store double the number of values */ assertTrue(intVector.getValueCapacity() >= initialCapacity * 2); /* vector data should still be intact after realloc */ j = 1; for (int i = 0; i <= initialCapacity; i += 2) { assertEquals(j, intVector.get(i), "unexpected value at index: " + i); j++; } /* reset the vector */ int capacityBeforeRealloc = intVector.getValueCapacity(); intVector.reset(); /* capacity shouldn't change after reset */ assertEquals(capacityBeforeRealloc, intVector.getValueCapacity()); /* vector data should have been zeroed out */ for (int i = 0; i < capacityBeforeRealloc; i++) { assertTrue(intVector.isNull(i), "non-zero data not expected at index: " + i); } } } @Test /* VarCharVector */ public void testSizeOfValueBuffer() { try (final VarCharVector vector = new VarCharVector(EMPTY_SCHEMA_PATH, allocator)) { int valueCount = 100; int currentSize = 0; vector.setInitialCapacity(valueCount); vector.allocateNew(); vector.setValueCount(valueCount); for (int i = 0; i < valueCount; i++) { currentSize += i; vector.setSafe(i, new byte[i]); } assertEquals(currentSize, vector.sizeOfValueBuffer()); } } @Test public void testFixedFloat2() { try (final Float2Vector floatVector = new Float2Vector(EMPTY_SCHEMA_PATH, allocator)) { boolean error = false; int initialCapacity = 16; /* we should not throw exception for these values of capacity */ floatVector.setInitialCapacity(MAX_VALUE_COUNT - 1); floatVector.setInitialCapacity(MAX_VALUE_COUNT); try { floatVector.setInitialCapacity(MAX_VALUE_COUNT * 4); } catch (OversizedAllocationException oe) { error = true; } finally { assertTrue(error); error = false; } floatVector.setInitialCapacity(initialCapacity); /* no memory allocation has happened yet so capacity of underlying buffer should be 0 */ assertEquals(0, floatVector.getValueCapacity()); /* allocate 32 bytes (16 * 2) */ floatVector.allocateNew(); /* underlying buffer should be able to store 16 values */ assertTrue(floatVector.getValueCapacity() >= initialCapacity); initialCapacity = floatVector.getValueCapacity(); floatVector.zeroVector(); /* populate the floatVector */ floatVector.set(0, (short) 0x101c); // Float16.toFloat16(+0.00050163269043f) floatVector.set(2, (short) 0x901c); // Float16.toFloat16(-0.00050163269043f) floatVector.set(4, (short) 0x101d); // Float16.toFloat16(+0.000502109527588f) floatVector.set(6, (short) 0x901d); // Float16.toFloat16(-0.000502109527588f) floatVector.set(8, (short) 0x121c); // Float16.toFloat16(+0.00074577331543f) floatVector.set(10, (short) 0x921c); // Float16.toFloat16(-0.00074577331543f) floatVector.set(12, (short) 0x501c); // Float16.toFloat16(+32.875f) floatVector.set(14, (short) 0xd01c); // Float16.toFloat16(-32.875f) try { floatVector.set(initialCapacity, (short) 0x141c); } catch (IndexOutOfBoundsException ie) { error = true; } finally { assertTrue(error); error = false; } /* check vector contents */ assertEquals((short) 0x101c, floatVector.get(0)); assertEquals((short) 0x901c, floatVector.get(2)); assertEquals((short) 0x101d, floatVector.get(4)); assertEquals((short) 0x901d, floatVector.get(6)); assertEquals((short) 0x121c, floatVector.get(8)); assertEquals((short) 0x921c, floatVector.get(10)); assertEquals((short) 0x501c, floatVector.get(12)); assertEquals((short) 0xd01c, floatVector.get(14)); try { floatVector.get(initialCapacity); } catch (IndexOutOfBoundsException ie) { error = true; } finally { assertTrue(error); } /* this should trigger a realloc() */ floatVector.setSafe(initialCapacity, (short) 0x141c); // Float16.toFloat16(+0.00100326538086f) /* underlying buffer should now be able to store double the number of values */ assertTrue(floatVector.getValueCapacity() >= initialCapacity * 2); /* vector data should still be intact after realloc */ assertEquals((short) 0x101c, floatVector.get(0)); assertEquals((short) 0x901c, floatVector.get(2)); assertEquals((short) 0x101d, floatVector.get(4)); assertEquals((short) 0x901d, floatVector.get(6)); assertEquals((short) 0x121c, floatVector.get(8)); assertEquals((short) 0x921c, floatVector.get(10)); assertEquals((short) 0x501c, floatVector.get(12)); assertEquals((short) 0xd01c, floatVector.get(14)); assertEquals((short) 0x141c, floatVector.get(initialCapacity)); /* reset the vector */ int capacityBeforeReset = floatVector.getValueCapacity(); floatVector.reset(); /* capacity shouldn't change after reset */ assertEquals(capacityBeforeReset, floatVector.getValueCapacity()); /* vector data should be zeroed out */ for (int i = 0; i < capacityBeforeReset; i++) { assertTrue(floatVector.isNull(i), "non-zero data not expected at index: " + i); } } } @Test public void testFixedFloat2WithPossibleTruncate() { try (final Float2Vector floatVector = new Float2Vector(EMPTY_SCHEMA_PATH, allocator)) { boolean error = false; int initialCapacity = 16; /* we should not throw exception for these values of capacity */ floatVector.setInitialCapacity(MAX_VALUE_COUNT - 1); floatVector.setInitialCapacity(MAX_VALUE_COUNT); try { floatVector.setInitialCapacity(MAX_VALUE_COUNT * 4); } catch (OversizedAllocationException oe) { error = true; } finally { assertTrue(error); error = false; } floatVector.setInitialCapacity(initialCapacity); /* no memory allocation has happened yet so capacity of underlying buffer should be 0 */ assertEquals(0, floatVector.getValueCapacity()); /* allocate 32 bytes (16 * 2) */ floatVector.allocateNew(); /* underlying buffer should be able to store 16 values */ assertTrue(floatVector.getValueCapacity() >= initialCapacity); initialCapacity = floatVector.getValueCapacity(); floatVector.zeroVector(); /* populate the floatVector */ floatVector.set(0, (short) 0x101c); // Float16.toFloat16(+0.00050163269043f) floatVector.set(2, (short) 0x901c); // Float16.toFloat16(-0.00050163269043f) floatVector.set(4, (short) 0x101d); // Float16.toFloat16(+0.000502109527588f) floatVector.setWithPossibleTruncate(6, 2049.0f); // in f32=2049.000000, out f16=2048 floatVector.setWithPossibleTruncate(8, 4098.0f); // in f32=4098.000000, out f16=4096 floatVector.setWithPossibleTruncate(10, 8196.0f); // in f32=8196.000000, out f16=8192 floatVector.setWithPossibleTruncate(12, 16392.0f); // in f32=16392.000000, out f16=16384 floatVector.setWithPossibleTruncate(14, 32784.0f); // in f32=32784.000000, out f16=32768 try { floatVector.setWithPossibleTruncate( initialCapacity, 1.618034f); // in f32=1.618034, out f16=1.6181641 } catch (IndexOutOfBoundsException ie) { error = true; } finally { assertTrue(error); error = false; } /* check vector contents */ assertEquals((short) 0x101c, floatVector.get(0)); assertEquals((short) 0x901c, floatVector.get(2)); assertEquals((short) 0x101d, floatVector.get(4)); assertEquals(2048.0f, floatVector.getValueAsFloat(6), 0); assertEquals(4096.0f, floatVector.getValueAsFloat(8), 0); assertEquals(8192.0f, floatVector.getValueAsFloat(10), 0); assertEquals(16384.0f, floatVector.getValueAsDouble(12), 0); assertEquals(32768.0f, floatVector.getValueAsDouble(14), 0); try { floatVector.get(initialCapacity); } catch (IndexOutOfBoundsException ie) { error = true; } finally { assertTrue(error); } /* this should trigger a realloc() */ floatVector.setSafeWithPossibleTruncate( initialCapacity, 1.618034f); // in f32=1.618034, out f16=1.6181641 /* underlying buffer should now be able to store double the number of values */ assertTrue(floatVector.getValueCapacity() >= initialCapacity * 2); /* vector data should still be intact after realloc */ assertEquals((short) 0x101c, floatVector.get(0)); assertEquals((short) 0x901c, floatVector.get(2)); assertEquals((short) 0x101d, floatVector.get(4)); assertEquals(2048.0f, floatVector.getValueAsFloat(6), 0); assertEquals(4096.0f, floatVector.getValueAsFloat(8), 0); assertEquals(8192.0f, floatVector.getValueAsFloat(10), 0); assertEquals(16384.0f, floatVector.getValueAsDouble(12), 0); assertEquals(32768.0f, floatVector.getValueAsDouble(14), 0); assertEquals(1.6181641f, floatVector.getValueAsDouble(initialCapacity), 0); /* reset the vector */ int capacityBeforeReset = floatVector.getValueCapacity(); floatVector.reset(); /* capacity shouldn't change after reset */ assertEquals(capacityBeforeReset, floatVector.getValueCapacity()); /* vector data should be zeroed out */ for (int i = 0; i < capacityBeforeReset; i++) { assertTrue(floatVector.isNull(i), "non-zero data not expected at index: " + i); } } } @Test /* Float4Vector */ public void testFixedType3() { try (final Float4Vector floatVector = new Float4Vector(EMPTY_SCHEMA_PATH, allocator)) { boolean error = false; int initialCapacity = 16; /* we should not throw exception for these values of capacity */ floatVector.setInitialCapacity(MAX_VALUE_COUNT - 1); floatVector.setInitialCapacity(MAX_VALUE_COUNT); try { floatVector.setInitialCapacity(MAX_VALUE_COUNT * 2); } catch (OversizedAllocationException oe) { error = true; } finally { assertTrue(error); error = false; } floatVector.setInitialCapacity(initialCapacity); /* no memory allocation has happened yet so capacity of underlying buffer should be 0 */ assertEquals(0, floatVector.getValueCapacity()); /* allocate 64 bytes (16 * 4) */ floatVector.allocateNew(); /* underlying buffer should be able to store 16 values */ assertTrue(floatVector.getValueCapacity() >= initialCapacity); initialCapacity = floatVector.getValueCapacity(); floatVector.zeroVector(); /* populate the floatVector */ floatVector.set(0, 1.5f); floatVector.set(2, 2.5f); floatVector.set(4, 3.3f); floatVector.set(6, 4.8f); floatVector.set(8, 5.6f); floatVector.set(10, 6.6f); floatVector.set(12, 7.8f); floatVector.set(14, 8.5f); try { floatVector.set(initialCapacity, 9.5f); } catch (IndexOutOfBoundsException ie) { error = true; } finally { assertTrue(error); error = false; } /* check vector contents */ assertEquals(1.5f, floatVector.get(0), 0); assertEquals(2.5f, floatVector.get(2), 0); assertEquals(3.3f, floatVector.get(4), 0); assertEquals(4.8f, floatVector.get(6), 0); assertEquals(5.6f, floatVector.get(8), 0); assertEquals(6.6f, floatVector.get(10), 0); assertEquals(7.8f, floatVector.get(12), 0); assertEquals(8.5f, floatVector.get(14), 0); try { floatVector.get(initialCapacity); } catch (IndexOutOfBoundsException ie) { error = true; } finally { assertTrue(error); error = false; } /* this should trigger a realloc() */ floatVector.setSafe(initialCapacity, 9.5f); /* underlying buffer should now be able to store double the number of values */ assertTrue(floatVector.getValueCapacity() >= initialCapacity * 2); /* vector data should still be intact after realloc */ assertEquals(1.5f, floatVector.get(0), 0); assertEquals(2.5f, floatVector.get(2), 0); assertEquals(3.3f, floatVector.get(4), 0); assertEquals(4.8f, floatVector.get(6), 0); assertEquals(5.6f, floatVector.get(8), 0); assertEquals(6.6f, floatVector.get(10), 0); assertEquals(7.8f, floatVector.get(12), 0); assertEquals(8.5f, floatVector.get(14), 0); assertEquals(9.5f, floatVector.get(initialCapacity), 0); /* reset the vector */ int capacityBeforeReset = floatVector.getValueCapacity(); floatVector.reset(); /* capacity shouldn't change after reset */ assertEquals(capacityBeforeReset, floatVector.getValueCapacity()); /* vector data should be zeroed out */ for (int i = 0; i < capacityBeforeReset; i++) { assertTrue(floatVector.isNull(i), "non-zero data not expected at index: " + i); } } } @Test /* Float8Vector */ public void testFixedType4() { try (final Float8Vector floatVector = new Float8Vector(EMPTY_SCHEMA_PATH, allocator)) { boolean error = false; int initialCapacity = 16; /* we should not throw exception for these values of capacity */ floatVector.setInitialCapacity(MAX_VALUE_COUNT_8BYTE - 1); floatVector.setInitialCapacity(MAX_VALUE_COUNT_8BYTE); try { floatVector.setInitialCapacity(MAX_VALUE_COUNT_8BYTE * 2); } catch (OversizedAllocationException oe) { error = true; } finally { assertTrue(error); error = false; } floatVector.setInitialCapacity(initialCapacity); /* no memory allocation has happened yet so capacity of underlying buffer should be 0 */ assertEquals(0, floatVector.getValueCapacity()); /* allocate 128 bytes (16 * 8) */ floatVector.allocateNew(); /* underlying buffer should be able to store 16 values */ assertTrue(floatVector.getValueCapacity() >= initialCapacity); initialCapacity = floatVector.getValueCapacity(); /* populate the vector */ floatVector.set(0, 1.55); floatVector.set(2, 2.53); floatVector.set(4, 3.36); floatVector.set(6, 4.82); floatVector.set(8, 5.67); floatVector.set(10, 6.67); floatVector.set(12, 7.87); floatVector.set(14, 8.56); try { floatVector.set(initialCapacity, 9.53); } catch (IndexOutOfBoundsException ie) { error = true; } finally { assertTrue(error); error = false; } /* check floatVector contents */ assertEquals(1.55, floatVector.get(0), 0); assertEquals(2.53, floatVector.get(2), 0); assertEquals(3.36, floatVector.get(4), 0); assertEquals(4.82, floatVector.get(6), 0); assertEquals(5.67, floatVector.get(8), 0); assertEquals(6.67, floatVector.get(10), 0); assertEquals(7.87, floatVector.get(12), 0); assertEquals(8.56, floatVector.get(14), 0); try { floatVector.get(initialCapacity); } catch (IndexOutOfBoundsException ie) { error = true; } finally { assertTrue(error); error = false; } /* this should trigger a realloc() */ floatVector.setSafe(initialCapacity, 9.53); /* underlying buffer should now be able to store double the number of values */ assertTrue(floatVector.getValueCapacity() >= initialCapacity * 2); /* vector data should still be intact after realloc */ assertEquals(1.55, floatVector.get(0), 0); assertEquals(2.53, floatVector.get(2), 0); assertEquals(3.36, floatVector.get(4), 0); assertEquals(4.82, floatVector.get(6), 0); assertEquals(5.67, floatVector.get(8), 0); assertEquals(6.67, floatVector.get(10), 0); assertEquals(7.87, floatVector.get(12), 0); assertEquals(8.56, floatVector.get(14), 0); assertEquals(9.53, floatVector.get(initialCapacity), 0); /* reset the vector */ int capacityBeforeReset = floatVector.getValueCapacity(); floatVector.reset(); /* capacity shouldn't change after reset */ assertEquals(capacityBeforeReset, floatVector.getValueCapacity()); /* vector data should be zeroed out */ for (int i = 0; i < capacityBeforeReset; i++) { assertTrue(floatVector.isNull(i), "non-zero data not expected at index: " + i); } } } @Test /* UInt4Vector */ public void testNullableFixedType1() { // Create a new value vector for 1024 integers. try (final UInt4Vector vector = newVector( UInt4Vector.class, EMPTY_SCHEMA_PATH, new ArrowType.Int(32, false), allocator); ) { boolean error = false; int initialCapacity = 1024; vector.setInitialCapacity(initialCapacity); /* no memory allocation has happened yet */ assertEquals(0, vector.getValueCapacity()); vector.allocateNew(); assertTrue(vector.getValueCapacity() >= initialCapacity); initialCapacity = vector.getValueCapacity(); // Put and set a few values vector.set(0, 100); vector.set(1, 101); vector.set(100, 102); vector.set(initialCapacity - 2, 103); vector.set(initialCapacity - 1, 104); /* check vector contents */ assertEquals(100, vector.get(0)); assertEquals(101, vector.get(1)); assertEquals(102, vector.get(100)); assertEquals(103, vector.get(initialCapacity - 2)); assertEquals(104, vector.get(initialCapacity - 1)); /* check unset bits/null values */ for (int i = 2, j = 101; i <= 99 || j <= initialCapacity - 3; i++, j++) { if (i <= 99) { assertTrue(vector.isNull(i)); } if (j <= initialCapacity - 3) { assertTrue(vector.isNull(j)); } } try { vector.set(initialCapacity, 10000); } catch (IndexOutOfBoundsException ie) { error = true; } finally { assertTrue(error); error = false; } try { vector.get(initialCapacity); } catch (IndexOutOfBoundsException ie) { error = true; } finally { assertTrue(error); error = false; } /* should trigger a realloc of the underlying bitvector and valuevector */ vector.setSafe(initialCapacity, 10000); /* check new capacity */ assertTrue(vector.getValueCapacity() >= initialCapacity * 2); /* vector contents should still be intact after realloc */ assertEquals(100, vector.get(0)); assertEquals(101, vector.get(1)); assertEquals(102, vector.get(100)); assertEquals(103, vector.get(initialCapacity - 2)); assertEquals(104, vector.get(initialCapacity - 1)); assertEquals(10000, vector.get(initialCapacity)); /* check unset bits/null values */ for (int i = 2, j = 101; i < 99 || j < initialCapacity - 3; i++, j++) { if (i <= 99) { assertTrue(vector.isNull(i)); } if (j <= initialCapacity - 3) { assertTrue(vector.isNull(j)); } } /* reset the vector */ int capacityBeforeReset = vector.getValueCapacity(); vector.reset(); /* capacity shouldn't change after reset */ assertEquals(capacityBeforeReset, vector.getValueCapacity()); /* vector data should be zeroed out */ for (int i = 0; i < capacityBeforeReset; i++) { assertTrue(vector.isNull(i), "non-null data not expected at index: " + i); } } } @Test /* Float4Vector */ public void testNullableFixedType2() { // Create a new value vector for 1024 integers try (final Float4Vector vector = newVector(Float4Vector.class, EMPTY_SCHEMA_PATH, MinorType.FLOAT4, allocator); ) { boolean error = false; int initialCapacity = 16; vector.setInitialCapacity(initialCapacity); /* no memory allocation has happened yet */ assertEquals(0, vector.getValueCapacity()); vector.allocateNew(); assertTrue(vector.getValueCapacity() >= initialCapacity); initialCapacity = vector.getValueCapacity(); /* populate the vector */ vector.set(0, 100.5f); vector.set(2, 201.5f); vector.set(4, 300.3f); vector.set(6, 423.8f); vector.set(8, 555.6f); vector.set(10, 66.6f); vector.set(12, 78.8f); vector.set(14, 89.5f); try { vector.set(initialCapacity, 90.5f); } catch (IndexOutOfBoundsException ie) { error = true; } finally { assertTrue(error); error = false; } /* check vector contents */ assertEquals(100.5f, vector.get(0), 0); assertTrue(vector.isNull(1)); assertEquals(201.5f, vector.get(2), 0); assertTrue(vector.isNull(3)); assertEquals(300.3f, vector.get(4), 0); assertTrue(vector.isNull(5)); assertEquals(423.8f, vector.get(6), 0); assertTrue(vector.isNull(7)); assertEquals(555.6f, vector.get(8), 0); assertTrue(vector.isNull(9)); assertEquals(66.6f, vector.get(10), 0); assertTrue(vector.isNull(11)); assertEquals(78.8f, vector.get(12), 0); assertTrue(vector.isNull(13)); assertEquals(89.5f, vector.get(14), 0); assertTrue(vector.isNull(15)); try { vector.get(initialCapacity); } catch (IndexOutOfBoundsException ie) { error = true; } finally { assertTrue(error); error = false; } /* this should trigger a realloc() */ vector.setSafe(initialCapacity, 90.5f); /* underlying buffer should now be able to store double the number of values */ assertTrue(vector.getValueCapacity() >= 2 * initialCapacity); /* vector data should still be intact after realloc */ assertEquals(100.5f, vector.get(0), 0); assertTrue(vector.isNull(1)); assertEquals(201.5f, vector.get(2), 0); assertTrue(vector.isNull(3)); assertEquals(300.3f, vector.get(4), 0); assertTrue(vector.isNull(5)); assertEquals(423.8f, vector.get(6), 0); assertTrue(vector.isNull(7)); assertEquals(555.6f, vector.get(8), 0); assertTrue(vector.isNull(9)); assertEquals(66.6f, vector.get(10), 0); assertTrue(vector.isNull(11)); assertEquals(78.8f, vector.get(12), 0); assertTrue(vector.isNull(13)); assertEquals(89.5f, vector.get(14), 0); assertTrue(vector.isNull(15)); /* reset the vector */ int capacityBeforeReset = vector.getValueCapacity(); vector.reset(); /* capacity shouldn't change after reset */ assertEquals(capacityBeforeReset, vector.getValueCapacity()); /* vector data should be zeroed out */ for (int i = 0; i < capacityBeforeReset; i++) { assertTrue(vector.isNull(i), "non-null data not expected at index: " + i); } } } @Test /* IntVector */ public void testNullableFixedType3() { // Create a new value vector for 1024 integers try (final IntVector vector = newVector(IntVector.class, EMPTY_SCHEMA_PATH, MinorType.INT, allocator)) { int initialCapacity = 1024; /* no memory allocation has happened yet so capacity of underlying buffer should be 0 */ assertEquals(0, vector.getValueCapacity()); /* allocate space for 4KB data (1024 * 4) */ vector.allocateNew(initialCapacity); /* underlying buffer should be able to store 1024 values */ assertTrue(vector.getValueCapacity() >= initialCapacity); initialCapacity = vector.getValueCapacity(); vector.set(0, 1); vector.set(1, 2); vector.set(100, 3); vector.set(1022, 4); vector.set(1023, 5); /* check vector contents */ int j = 1; for (int i = 0; i <= 1023; i++) { if ((i >= 2 && i <= 99) || (i >= 101 && i <= 1021)) { assertTrue(vector.isNull(i), "non-null data not expected at index: " + i); } else { assertFalse(vector.isNull(i), "null data not expected at index: " + i); assertEquals(j, vector.get(i), "unexpected value at index: " + i); j++; } } vector.setValueCount(1024); List buffers = vector.getFieldBuffers(); assertEquals(2, buffers.size()); ArrowBuf validityVectorBuf = buffers.get(0); /* bitvector tracks 1024 integers --> 1024 bits --> 128 bytes */ assertTrue(validityVectorBuf.readableBytes() >= 128); assertEquals(3, validityVectorBuf.getByte(0)); // 1st and second bit defined for (int i = 1; i < 12; i++) { assertEquals(0, validityVectorBuf.getByte(i)); // nothing defined until 100 } assertEquals(16, validityVectorBuf.getByte(12)); // 100th bit is defined (12 * 8 + 4) for (int i = 13; i < 127; i++) { assertEquals(0, validityVectorBuf.getByte(i)); // nothing defined between 100th and 1022nd } assertEquals(-64, validityVectorBuf.getByte(127)); // 1022nd and 1023rd bit defined /* this should trigger a realloc() */ vector.setSafe(initialCapacity, 6); /* underlying buffer should now be able to store double the number of values */ assertTrue(vector.getValueCapacity() >= 2 * initialCapacity); /* vector data should still be intact after realloc */ j = 1; for (int i = 0; i < (initialCapacity * 2); i++) { if ((i > 1023 && i != initialCapacity) || (i >= 2 && i <= 99) || (i >= 101 && i <= 1021)) { assertTrue(vector.isNull(i), "non-null data not expected at index: " + i); } else { assertFalse(vector.isNull(i), "null data not expected at index: " + i); assertEquals(j, vector.get(i), "unexpected value at index: " + i); j++; } } /* reset the vector */ int capacityBeforeReset = vector.getValueCapacity(); vector.reset(); /* capacity shouldn't change after reset */ assertEquals(capacityBeforeReset, vector.getValueCapacity()); /* vector data should have been zeroed out */ for (int i = 0; i < capacityBeforeReset; i++) { assertTrue(vector.isNull(i), "non-null data not expected at index: " + i); } vector.allocateNew(initialCapacity * 4); // vector has been erased for (int i = 0; i < initialCapacity * 4; i++) { assertTrue(vector.isNull(i), "non-null data not expected at index: " + i); } } } @Test /* IntVector */ public void testNullableFixedType4() { try (final IntVector vector = newVector(IntVector.class, EMPTY_SCHEMA_PATH, MinorType.INT, allocator)) { /* no memory allocation has happened yet */ assertEquals(0, vector.getValueCapacity()); vector.allocateNew(); int valueCapacity = vector.getValueCapacity(); assertEquals(vector.INITIAL_VALUE_ALLOCATION, valueCapacity); int baseValue = 20000; for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { vector.set(i, baseValue + i); } } for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { assertFalse(vector.isNull(i), "unexpected null value at index: " + i); assertEquals((baseValue + i), vector.get(i), "unexpected value at index: " + i); } else { assertTrue(vector.isNull(i), "unexpected non-null value at index: " + i); } } vector.setSafe(valueCapacity, 20000000); assertTrue(vector.getValueCapacity() >= valueCapacity * 2); for (int i = 0; i < vector.getValueCapacity(); i++) { if (i == valueCapacity) { assertFalse(vector.isNull(i), "unexpected null value at index: " + i); assertEquals(20000000, vector.get(i), "unexpected value at index: " + i); } else if (i < valueCapacity) { if ((i & 1) == 1) { assertFalse(vector.isNull(i), "unexpected null value at index: " + i); assertEquals((baseValue + i), vector.get(i), "unexpected value at index: " + i); } } else { assertTrue(vector.isNull(i), "unexpected non-null value at index: " + i); } } vector.zeroVector(); for (int i = 0; i < vector.getValueCapacity(); i += 2) { vector.set(i, baseValue + i); } for (int i = 0; i < vector.getValueCapacity(); i++) { if (i % 2 == 0) { assertFalse(vector.isNull(i), "unexpected null value at index: " + i); assertEquals((baseValue + i), vector.get(i), "unexpected value at index: " + i); } else { assertTrue(vector.isNull(i), "unexpected non-null value at index: " + i); } } int valueCapacityBeforeRealloc = vector.getValueCapacity(); vector.setSafe(valueCapacityBeforeRealloc + 1000, 400000000); assertTrue(vector.getValueCapacity() >= valueCapacity * 4); for (int i = 0; i < vector.getValueCapacity(); i++) { if (i == (valueCapacityBeforeRealloc + 1000)) { assertFalse(vector.isNull(i), "unexpected null value at index: " + i); assertEquals(400000000, vector.get(i), "unexpected value at index: " + i); } else if (i < valueCapacityBeforeRealloc && (i % 2) == 0) { assertFalse(vector.isNull(i), "unexpected null value at index: " + i); assertEquals(baseValue + i, vector.get(i), "unexpected value at index: " + i); } else { assertTrue(vector.isNull(i), "unexpected non-null value at index: " + i); } } /* reset the vector */ int valueCapacityBeforeReset = vector.getValueCapacity(); vector.reset(); /* capacity shouldn't change after reset */ assertEquals(valueCapacityBeforeReset, vector.getValueCapacity()); /* vector data should be zeroed out */ for (int i = 0; i < valueCapacityBeforeReset; i++) { assertTrue(vector.isNull(i), "non-null data not expected at index: " + i); } } } /* * Tests for Variable Width Vectors * * Covered types as of now * * -- VarCharVector * -- VarBinaryVector * * TODO: * * -- VarCharVector * -- VarBinaryVector */ /** * ARROW-7831: this checks that a slice taken off a buffer is still readable after that buffer's * allocator is closed. */ @Test /* VarCharVector */ public void testSplitAndTransfer1() { try (final VarCharVector targetVector = newVarCharVector("split-target", allocator)) { try (final VarCharVector sourceVector = newVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { sourceVector.allocateNew(1024 * 10, 1024); sourceVector.set(0, STR1); sourceVector.set(1, STR2); sourceVector.set(2, STR3); sourceVector.setValueCount(3); final long allocatedMem = allocator.getAllocatedMemory(); final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); final int offsetRefCnt = sourceVector.getOffsetBuffer().refCnt(); final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); // split and transfer with slice starting at the beginning: this should not allocate // anything new sourceVector.splitAndTransferTo(0, 2, targetVector); assertEquals(allocatedMem, allocator.getAllocatedMemory()); // The validity and offset buffers are sliced from a same buffer.See // BaseFixedWidthVector#allocateBytes. // Therefore, the refcnt of the validity buffer is increased once since the startIndex is 0. // The refcnt of the // offset buffer is increased as well for the same reason. This amounts to a total of 2. assertEquals(validityRefCnt + 2, sourceVector.getValidityBuffer().refCnt()); assertEquals(offsetRefCnt + 2, sourceVector.getOffsetBuffer().refCnt()); assertEquals(dataRefCnt + 1, sourceVector.getDataBuffer().refCnt()); } assertArrayEquals(STR1, targetVector.get(0)); assertArrayEquals(STR2, targetVector.get(1)); } } /** * ARROW-7831: this checks that a vector that got sliced is still readable after the slice's * allocator got closed. */ @Test /* VarCharVector */ public void testSplitAndTransfer2() { try (final VarCharVector sourceVector = newVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { try (final VarCharVector targetVector = newVarCharVector("split-target", allocator)) { sourceVector.allocateNew(1024 * 10, 1024); sourceVector.set(0, STR1); sourceVector.set(1, STR2); sourceVector.set(2, STR3); sourceVector.setValueCount(3); final long allocatedMem = allocator.getAllocatedMemory(); final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); final int offsetRefCnt = sourceVector.getOffsetBuffer().refCnt(); final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); // split and transfer with slice starting at the beginning: this should not allocate // anything new sourceVector.splitAndTransferTo(0, 2, targetVector); assertEquals(allocatedMem, allocator.getAllocatedMemory()); // The validity and offset buffers are sliced from a same buffer.See // BaseFixedWidthVector#allocateBytes. // Therefore, the refcnt of the validity buffer is increased once since the startIndex is 0. // The refcnt of the // offset buffer is increased as well for the same reason. This amounts to a total of 2. assertEquals(validityRefCnt + 2, sourceVector.getValidityBuffer().refCnt()); assertEquals(offsetRefCnt + 2, sourceVector.getOffsetBuffer().refCnt()); assertEquals(dataRefCnt + 1, sourceVector.getDataBuffer().refCnt()); } assertArrayEquals(STR1, sourceVector.get(0)); assertArrayEquals(STR2, sourceVector.get(1)); assertArrayEquals(STR3, sourceVector.get(2)); } } /** * ARROW-7831: this checks an offset splitting optimization, in the case where all the values up * to the start of the slice are null/empty, which avoids allocation for the offset buffer. */ @Test /* VarCharVector */ public void testSplitAndTransfer3() { try (final VarCharVector targetVector = newVarCharVector("split-target", allocator); final VarCharVector sourceVector = newVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { sourceVector.allocateNew(1024 * 10, 1024); sourceVector.set(0, new byte[0]); sourceVector.setNull(1); sourceVector.set(2, STR1); sourceVector.set(3, STR2); sourceVector.set(4, STR3); sourceVector.setValueCount(5); final long allocatedMem = allocator.getAllocatedMemory(); final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); final int offsetRefCnt = sourceVector.getOffsetBuffer().refCnt(); final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); sourceVector.splitAndTransferTo(2, 2, targetVector); // because the offset starts at 0 since the first 2 values are empty/null, the allocation only // consists in // the size needed for the validity buffer final long validitySize = DefaultRoundingPolicy.DEFAULT_ROUNDING_POLICY.getRoundedSize( getValidityBufferSizeFromCount(2)); assertEquals(allocatedMem + validitySize, allocator.getAllocatedMemory()); // The validity and offset buffers are sliced from a same buffer.See // BaseFixedWidthVector#allocateBytes. // Since values up to the startIndex are empty/null, the offset buffer doesn't need to be // reallocated and // therefore its refcnt is increased by 1. assertEquals(validityRefCnt + 1, sourceVector.getValidityBuffer().refCnt()); assertEquals(offsetRefCnt + 1, sourceVector.getOffsetBuffer().refCnt()); assertEquals(dataRefCnt + 1, sourceVector.getDataBuffer().refCnt()); assertArrayEquals(STR1, targetVector.get(0)); assertArrayEquals(STR2, targetVector.get(1)); } } /** * ARROW-7831: ensures that data is transferred from one allocator to another in case of 0-index * start special cases. */ @Test /* VarCharVector */ public void testSplitAndTransfer4() { try (final BufferAllocator targetAllocator = allocator.newChildAllocator("target-alloc", 256, 256); final VarCharVector targetVector = newVarCharVector("split-target", targetAllocator)) { try (final BufferAllocator sourceAllocator = allocator.newChildAllocator("source-alloc", 256, 256); final VarCharVector sourceVector = newVarCharVector(EMPTY_SCHEMA_PATH, sourceAllocator)) { sourceVector.allocateNew(50, 3); sourceVector.set(0, STR1); sourceVector.set(1, STR2); sourceVector.set(2, STR3); sourceVector.setValueCount(3); final long allocatedMem = allocator.getAllocatedMemory(); final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); final int offsetRefCnt = sourceVector.getOffsetBuffer().refCnt(); final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); // split and transfer with slice starting at the beginning: this should not allocate // anything new sourceVector.splitAndTransferTo(0, 2, targetVector); assertEquals(allocatedMem, allocator.getAllocatedMemory()); // Unlike testSplitAndTransfer1 where the buffers originated from the same allocator, the // refcnts of each // buffers for this test should be the same as what the source allocator ended up with. assertEquals(validityRefCnt, sourceVector.getValidityBuffer().refCnt()); assertEquals(offsetRefCnt, sourceVector.getOffsetBuffer().refCnt()); assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); } assertArrayEquals(STR1, targetVector.get(0)); assertArrayEquals(STR2, targetVector.get(1)); } } @Test /* VarCharVector */ public void testNullableVarType1() { // Create a new value vector for 1024 integers. try (final VarCharVector vector = newVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { vector.allocateNew(1024 * 10, 1024); vector.set(0, STR1); vector.set(1, STR2); vector.set(2, STR3); vector.setSafe(3, STR3, 1, STR3.length - 1); vector.setSafe(4, STR3, 2, STR3.length - 2); ByteBuffer str3ByteBuffer = ByteBuffer.wrap(STR3); vector.setSafe(5, str3ByteBuffer, 1, STR3.length - 1); vector.setSafe(6, str3ByteBuffer, 2, STR3.length - 2); // Set with convenience function Text txt = new Text("foo"); vector.setSafe(7, txt); // Check the sample strings. assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(Arrays.copyOfRange(STR3, 1, STR3.length), vector.get(3)); assertArrayEquals(Arrays.copyOfRange(STR3, 2, STR3.length), vector.get(4)); assertArrayEquals(Arrays.copyOfRange(STR3, 1, STR3.length), vector.get(5)); assertArrayEquals(Arrays.copyOfRange(STR3, 2, STR3.length), vector.get(6)); // Check returning a Text object assertEquals(txt, vector.getObject(7)); // Ensure null value throws. assertNull(vector.get(8)); } } @Test public void testGetTextRepeatedly() { try (final VarCharVector vector = new VarCharVector("myvector", allocator)) { ValueVectorDataPopulator.setVector(vector, STR1, STR2); vector.setValueCount(2); /* check the vector output */ Text text = new Text(); vector.read(0, text); assertArrayEquals(STR1, text.getBytes()); vector.read(1, text); assertArrayEquals(STR2, text.getBytes()); } } @Test /* VarBinaryVector */ public void testNullableVarType2() { // Create a new value vector for 1024 integers. try (final VarBinaryVector vector = newVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { vector.allocateNew(1024 * 10, 1024); vector.set(0, STR1); vector.set(1, STR2); vector.set(2, STR3); vector.setSafe(3, STR3, 1, STR3.length - 1); vector.setSafe(4, STR3, 2, STR3.length - 2); ByteBuffer str3ByteBuffer = ByteBuffer.wrap(STR3); vector.setSafe(5, str3ByteBuffer, 1, STR3.length - 1); vector.setSafe(6, str3ByteBuffer, 2, STR3.length - 2); // Check the sample strings. assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(Arrays.copyOfRange(STR3, 1, STR3.length), vector.get(3)); assertArrayEquals(Arrays.copyOfRange(STR3, 2, STR3.length), vector.get(4)); assertArrayEquals(Arrays.copyOfRange(STR3, 1, STR3.length), vector.get(5)); assertArrayEquals(Arrays.copyOfRange(STR3, 2, STR3.length), vector.get(6)); // Ensure null value throws. assertNull(vector.get(7)); } } @Test public void testReallocateCheckSuccess() { assertThrows( OversizedAllocationException.class, () -> { // Create a new value vector for 1024 integers. try (final VarBinaryVector vector = newVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { vector.allocateNew(1024 * 10, 1024); vector.set(0, STR1); // Check the sample strings. assertArrayEquals(STR1, vector.get(0)); // update the index offset to a larger one ArrowBuf offsetBuf = vector.getOffsetBuffer(); offsetBuf.setInt(VarBinaryVector.OFFSET_WIDTH, Integer.MAX_VALUE - 5); vector.setValueLengthSafe(1, 6); } }); } @Test public void testGetBytesRepeatedly() { try (VarBinaryVector vector = new VarBinaryVector("", allocator)) { vector.allocateNew(5, 1); final String str = "hello world"; final String str2 = "foo"; vector.setSafe(0, str.getBytes(StandardCharsets.UTF_8)); vector.setSafe(1, str2.getBytes(StandardCharsets.UTF_8)); // verify results ReusableByteArray reusableByteArray = new ReusableByteArray(); vector.read(0, reusableByteArray); assertArrayEquals( str.getBytes(StandardCharsets.UTF_8), Arrays.copyOfRange( reusableByteArray.getBuffer(), 0, (int) reusableByteArray.getLength())); byte[] oldBuffer = reusableByteArray.getBuffer(); vector.read(1, reusableByteArray); assertArrayEquals( str2.getBytes(StandardCharsets.UTF_8), Arrays.copyOfRange( reusableByteArray.getBuffer(), 0, (int) reusableByteArray.getLength())); // There should not have been any reallocation since the newer value is smaller in length. assertSame(oldBuffer, reusableByteArray.getBuffer()); } } /* * generic tests * * -- lastSet() and setValueCount() * -- fillEmpties() * -- VectorLoader and VectorUnloader * -- some realloc tests * * TODO: * * The realloc() related tests below should be moved up and we need to * add realloc related tests (edge cases) for more vector types. */ @Test /* Float8Vector */ public void testReallocAfterVectorTransfer1() { try (final Float8Vector vector = new Float8Vector(EMPTY_SCHEMA_PATH, allocator)) { int initialCapacity = 4096; /* use the default capacity; 4096*8 => 32KB */ vector.setInitialCapacity(initialCapacity); vector.allocateNew(); assertTrue(vector.getValueCapacity() >= initialCapacity); initialCapacity = vector.getValueCapacity(); double baseValue = 100.375; for (int i = 0; i < initialCapacity; i++) { vector.setSafe(i, baseValue + (double) i); } /* the above setSafe calls should not have triggered a realloc as * we are within the capacity. check the vector contents */ assertEquals(initialCapacity, vector.getValueCapacity()); for (int i = 0; i < initialCapacity; i++) { double value = vector.get(i); assertEquals(baseValue + (double) i, value, 0); } /* this should trigger a realloc */ vector.setSafe(initialCapacity, baseValue + (double) initialCapacity); assertTrue(vector.getValueCapacity() >= initialCapacity * 2); int capacityAfterRealloc1 = vector.getValueCapacity(); for (int i = initialCapacity + 1; i < capacityAfterRealloc1; i++) { vector.setSafe(i, baseValue + (double) i); } for (int i = 0; i < capacityAfterRealloc1; i++) { double value = vector.get(i); assertEquals(baseValue + (double) i, value, 0); } /* this should trigger a realloc */ vector.setSafe(capacityAfterRealloc1, baseValue + (double) capacityAfterRealloc1); assertTrue(vector.getValueCapacity() >= initialCapacity * 4); int capacityAfterRealloc2 = vector.getValueCapacity(); for (int i = capacityAfterRealloc1 + 1; i < capacityAfterRealloc2; i++) { vector.setSafe(i, baseValue + (double) i); } for (int i = 0; i < capacityAfterRealloc2; i++) { double value = vector.get(i); assertEquals(baseValue + (double) i, value, 0); } /* at this point we are working with a 128KB buffer data for this * vector. now let's transfer this vector */ TransferPair transferPair = vector.getTransferPair(allocator); transferPair.transfer(); Float8Vector toVector = (Float8Vector) transferPair.getTo(); /* now let's realloc the toVector */ toVector.reAlloc(); assertTrue(toVector.getValueCapacity() >= initialCapacity * 8); for (int i = 0; i < toVector.getValueCapacity(); i++) { if (i < capacityAfterRealloc2) { assertEquals(baseValue + (double) i, toVector.get(i), 0); } else { assertTrue(toVector.isNull(i)); } } toVector.close(); } } @Test /* Float8Vector */ public void testReallocAfterVectorTransfer2() { try (final Float8Vector vector = new Float8Vector(EMPTY_SCHEMA_PATH, allocator)) { int initialCapacity = 4096; vector.allocateNew(initialCapacity); assertTrue(vector.getValueCapacity() >= initialCapacity); initialCapacity = vector.getValueCapacity(); double baseValue = 100.375; for (int i = 0; i < initialCapacity; i++) { vector.setSafe(i, baseValue + (double) i); } /* the above setSafe calls should not have triggered a realloc as * we are within the capacity. check the vector contents */ assertEquals(initialCapacity, vector.getValueCapacity()); for (int i = 0; i < initialCapacity; i++) { double value = vector.get(i); assertEquals(baseValue + (double) i, value, 0); } /* this should trigger a realloc */ vector.setSafe(initialCapacity, baseValue + (double) initialCapacity); assertTrue(vector.getValueCapacity() >= initialCapacity * 2); int capacityAfterRealloc1 = vector.getValueCapacity(); for (int i = initialCapacity + 1; i < capacityAfterRealloc1; i++) { vector.setSafe(i, baseValue + (double) i); } for (int i = 0; i < capacityAfterRealloc1; i++) { double value = vector.get(i); assertEquals(baseValue + (double) i, value, 0); } /* this should trigger a realloc */ vector.setSafe(capacityAfterRealloc1, baseValue + (double) capacityAfterRealloc1); assertTrue(vector.getValueCapacity() >= initialCapacity * 4); int capacityAfterRealloc2 = vector.getValueCapacity(); for (int i = capacityAfterRealloc1 + 1; i < capacityAfterRealloc2; i++) { vector.setSafe(i, baseValue + (double) i); } for (int i = 0; i < capacityAfterRealloc2; i++) { double value = vector.get(i); assertEquals(baseValue + (double) i, value, 0); } /* at this point we are working with a 128KB buffer data for this * vector. now let's transfer this vector */ TransferPair transferPair = vector.getTransferPair(allocator); transferPair.transfer(); Float8Vector toVector = (Float8Vector) transferPair.getTo(); /* check toVector contents before realloc */ for (int i = 0; i < toVector.getValueCapacity(); i++) { assertFalse(toVector.isNull(i), "unexpected null value at index: " + i); double value = toVector.get(i); assertEquals(baseValue + (double) i, value, 0, "unexpected value at index: " + i); } /* now let's realloc the toVector and check contents again */ toVector.reAlloc(); assertTrue(toVector.getValueCapacity() >= initialCapacity * 8); for (int i = 0; i < toVector.getValueCapacity(); i++) { if (i < capacityAfterRealloc2) { assertFalse(toVector.isNull(i), "unexpected null value at index: " + i); double value = toVector.get(i); assertEquals(baseValue + (double) i, value, 0, "unexpected value at index: " + i); } else { assertTrue(toVector.isNull(i), "unexpected non-null value at index: " + i); } } toVector.close(); } } @Test /* VarCharVector */ public void testReallocAfterVectorTransfer3() { try (final VarCharVector vector = new VarCharVector(EMPTY_SCHEMA_PATH, allocator)) { /* 4096 values with 10 byte per record */ vector.allocateNew(4096 * 10, 4096); int valueCapacity = vector.getValueCapacity(); assertTrue(valueCapacity >= 4096); /* populate the vector */ for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { vector.set(i, STR1); } else { vector.set(i, STR2); } } /* Check the vector output */ for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { assertArrayEquals(STR1, vector.get(i)); } else { assertArrayEquals(STR2, vector.get(i)); } } /* trigger first realloc */ vector.setSafe(valueCapacity, STR2, 0, STR2.length); assertTrue(vector.getValueCapacity() >= 2 * valueCapacity); while (vector.getByteCapacity() < 10 * vector.getValueCapacity()) { vector.reallocDataBuffer(); } /* populate the remaining vector */ for (int i = valueCapacity; i < vector.getValueCapacity(); i++) { if ((i & 1) == 1) { vector.set(i, STR1); } else { vector.set(i, STR2); } } /* Check the vector output */ valueCapacity = vector.getValueCapacity(); for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { assertArrayEquals(STR1, vector.get(i)); } else { assertArrayEquals(STR2, vector.get(i)); } } /* trigger second realloc */ vector.setSafe(valueCapacity + 10, STR2, 0, STR2.length); assertTrue(vector.getValueCapacity() >= 2 * valueCapacity); while (vector.getByteCapacity() < 10 * vector.getValueCapacity()) { vector.reallocDataBuffer(); } /* populate the remaining vector */ for (int i = valueCapacity; i < vector.getValueCapacity(); i++) { if ((i & 1) == 1) { vector.set(i, STR1); } else { vector.set(i, STR2); } } /* Check the vector output */ valueCapacity = vector.getValueCapacity(); for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { assertArrayEquals(STR1, vector.get(i)); } else { assertArrayEquals(STR2, vector.get(i)); } } /* we are potentially working with 4x the size of vector buffer * that we initially started with. Now let's transfer the vector. */ TransferPair transferPair = vector.getTransferPair(allocator); transferPair.transfer(); VarCharVector toVector = (VarCharVector) transferPair.getTo(); valueCapacity = toVector.getValueCapacity(); for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { assertArrayEquals(STR1, toVector.get(i)); } else { assertArrayEquals(STR2, toVector.get(i)); } } toVector.close(); } } @Test /* IntVector */ public void testReallocAfterVectorTransfer4() { try (final IntVector vector = new IntVector(EMPTY_SCHEMA_PATH, allocator)) { /* 4096 values */ vector.allocateNew(4096); int valueCapacity = vector.getValueCapacity(); assertTrue(valueCapacity >= 4096); /* populate the vector */ for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 0) { vector.set(i, 1000 + i); } } /* Check the vector output */ for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 0) { assertEquals(1000 + i, vector.get(i)); } else { assertTrue(vector.isNull(i)); } } /* trigger first realloc */ vector.setSafe(valueCapacity, 10000000); assertTrue(vector.getValueCapacity() >= valueCapacity * 2); /* populate the remaining vector */ for (int i = valueCapacity; i < vector.getValueCapacity(); i++) { if ((i & 1) == 0) { vector.set(i, 1000 + i); } } /* Check the vector output */ valueCapacity = vector.getValueCapacity(); for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 0) { assertEquals(1000 + i, vector.get(i)); } else { assertTrue(vector.isNull(i)); } } /* trigger second realloc */ vector.setSafe(valueCapacity, 10000000); assertTrue(vector.getValueCapacity() >= valueCapacity * 2); /* populate the remaining vector */ for (int i = valueCapacity; i < vector.getValueCapacity(); i++) { if ((i & 1) == 0) { vector.set(i, 1000 + i); } } /* Check the vector output */ valueCapacity = vector.getValueCapacity(); for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 0) { assertEquals(1000 + i, vector.get(i)); } else { assertTrue(vector.isNull(i)); } } /* we are potentially working with 4x the size of vector buffer * that we initially started with. Now let's transfer the vector. */ TransferPair transferPair = vector.getTransferPair(allocator); transferPair.transfer(); IntVector toVector = (IntVector) transferPair.getTo(); /* value capacity of source and target vectors should be same after * the transfer. */ assertEquals(valueCapacity, toVector.getValueCapacity()); for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 0) { assertEquals(1000 + i, toVector.get(i)); } else { assertTrue(toVector.isNull(i)); } } toVector.close(); } } @Test public void testReAllocFixedWidthVector() { // Create a new value vector for 1024 integers try (final Float4Vector vector = newVector(Float4Vector.class, EMPTY_SCHEMA_PATH, MinorType.FLOAT4, allocator)) { vector.allocateNew(1024); assertTrue(vector.getValueCapacity() >= 1024); int initialCapacity = vector.getValueCapacity(); // Put values in indexes that fall within the initial allocation vector.setSafe(0, 100.1f); vector.setSafe(100, 102.3f); vector.setSafe(1023, 104.5f); // Now try to put values in space that falls beyond the initial allocation vector.setSafe(2000, 105.5f); // Check valueCapacity is more than initial allocation assertTrue(vector.getValueCapacity() >= 2 * initialCapacity); assertEquals(100.1f, vector.get(0), 0); assertEquals(102.3f, vector.get(100), 0); assertEquals(104.5f, vector.get(1023), 0); assertEquals(105.5f, vector.get(2000), 0); // Set the valueCount to be more than valueCapacity of current allocation. This is possible // for ValueVectors // as we don't call setSafe for null values, but we do call setValueCount when all values are // inserted into the // vector vector.setValueCount(vector.getValueCapacity() + 200); } } @Test public void testReAllocVariableWidthVector() { try (final VarCharVector vector = newVector(VarCharVector.class, EMPTY_SCHEMA_PATH, MinorType.VARCHAR, allocator)) { vector.setInitialCapacity(4095); vector.allocateNew(); int initialCapacity = vector.getValueCapacity(); assertTrue(initialCapacity >= 4095); /* Put values in indexes that fall within the initial allocation */ vector.setSafe(0, STR1, 0, STR1.length); vector.setSafe(initialCapacity - 1, STR2, 0, STR2.length); /* the above set calls should NOT have triggered a realloc */ assertEquals(initialCapacity, vector.getValueCapacity()); /* Now try to put values in space that falls beyond the initial allocation */ vector.setSafe(initialCapacity + 200, STR3, 0, STR3.length); /* Check valueCapacity is more than initial allocation */ assertTrue(initialCapacity * 2 <= vector.getValueCapacity()); assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(initialCapacity - 1)); assertArrayEquals(STR3, vector.get(initialCapacity + 200)); // Set the valueCount to be more than valueCapacity of current allocation. This is possible // for ValueVectors // as we don't call setSafe for null values, but we do call setValueCount when the current // batch is processed. vector.setValueCount(vector.getValueCapacity() + 200); } } @Test public void testFillEmptiesNotOverfill() { try (final VarCharVector vector = newVector(VarCharVector.class, EMPTY_SCHEMA_PATH, MinorType.VARCHAR, allocator)) { vector.setInitialCapacity(4095); vector.allocateNew(); int initialCapacity = vector.getValueCapacity(); assertTrue(initialCapacity >= 4095); vector.setSafe(4094, "hello".getBytes(StandardCharsets.UTF_8), 0, 5); /* the above set method should NOT have triggered a realloc */ assertEquals(initialCapacity, vector.getValueCapacity()); long bufSizeBefore = vector.getFieldBuffers().get(1).capacity(); vector.setValueCount(initialCapacity); assertEquals(bufSizeBefore, vector.getFieldBuffers().get(1).capacity()); assertEquals(initialCapacity, vector.getValueCapacity()); } } @Test public void testSetSafeWithArrowBufNoExcessAllocs() { final int numValues = BaseFixedWidthVector.INITIAL_VALUE_ALLOCATION * 2; final byte[] valueBytes = "hello world".getBytes(StandardCharsets.UTF_8); final int valueBytesLength = valueBytes.length; final int isSet = 1; try (final VarCharVector fromVector = newVector(VarCharVector.class, EMPTY_SCHEMA_PATH, MinorType.VARCHAR, allocator); final VarCharVector toVector = newVector(VarCharVector.class, EMPTY_SCHEMA_PATH, MinorType.VARCHAR, allocator)) { /* * Populate the from vector with 'numValues' with byte-arrays, each of size 'valueBytesLength'. */ fromVector.setInitialCapacity(numValues); fromVector.allocateNew(); for (int i = 0; i < numValues; ++i) { fromVector.setSafe(i, valueBytes, 0 /*start*/, valueBytesLength); } fromVector.setValueCount(numValues); ArrowBuf fromDataBuffer = fromVector.getDataBuffer(); assertTrue(numValues * valueBytesLength <= fromDataBuffer.capacity()); /* * Copy the entries one-by-one from 'fromVector' to 'toVector', but use the setSafe with * ArrowBuf API (instead of setSafe with byte-array). */ toVector.setInitialCapacity(numValues); toVector.allocateNew(); for (int i = 0; i < numValues; i++) { int start = fromVector.getStartOffset(i); int end = fromVector.getStartOffset(i + 1); toVector.setSafe(i, isSet, start, end, fromDataBuffer); } /* * Since the 'fromVector' and 'toVector' have the same initial capacity, and were populated * with the same varchar elements, the allocations and hence, the final capacity should be * the same. */ assertEquals(fromDataBuffer.capacity(), toVector.getDataBuffer().capacity()); } } @Test public void testCopyFromWithNulls() { try (final VarCharVector vector = newVector(VarCharVector.class, EMPTY_SCHEMA_PATH, MinorType.VARCHAR, allocator); final VarCharVector vector2 = newVector(VarCharVector.class, EMPTY_SCHEMA_PATH, MinorType.VARCHAR, allocator)) { vector.setInitialCapacity(4095); vector.allocateNew(); int capacity = vector.getValueCapacity(); assertTrue(capacity >= 4095); for (int i = 0; i < capacity; i++) { if (i % 3 == 0) { continue; } byte[] b = Integer.toString(i).getBytes(StandardCharsets.UTF_8); vector.setSafe(i, b, 0, b.length); } /* NO reAlloc() should have happened in setSafe() */ assertEquals(capacity, vector.getValueCapacity()); vector.setValueCount(capacity); for (int i = 0; i < capacity; i++) { if (i % 3 == 0) { assertNull(vector.getObject(i)); } else { assertEquals( Integer.toString(i), vector.getObject(i).toString(), "unexpected value at index: " + i); } } vector2.setInitialCapacity(4095); vector2.allocateNew(); int capacity2 = vector2.getValueCapacity(); assertEquals(capacity2, capacity); for (int i = 0; i < capacity; i++) { vector2.copyFromSafe(i, i, vector); if (i % 3 == 0) { assertNull(vector2.getObject(i)); } else { assertEquals( Integer.toString(i), vector2.getObject(i).toString(), "unexpected value at index: " + i); } } /* NO reAlloc() should have happened in copyFrom */ assertEquals(capacity, vector2.getValueCapacity()); vector2.setValueCount(capacity); for (int i = 0; i < capacity; i++) { if (i % 3 == 0) { assertNull(vector2.getObject(i)); } else { assertEquals( Integer.toString(i), vector2.getObject(i).toString(), "unexpected value at index: " + i); } } } } @Test public void testCopyFromWithNulls1() { try (final VarCharVector vector = newVector(VarCharVector.class, EMPTY_SCHEMA_PATH, MinorType.VARCHAR, allocator); final VarCharVector vector2 = newVector(VarCharVector.class, EMPTY_SCHEMA_PATH, MinorType.VARCHAR, allocator)) { vector.setInitialCapacity(4095); vector.allocateNew(); int capacity = vector.getValueCapacity(); assertTrue(capacity >= 4095); for (int i = 0; i < capacity; i++) { if (i % 3 == 0) { continue; } byte[] b = Integer.toString(i).getBytes(StandardCharsets.UTF_8); vector.setSafe(i, b, 0, b.length); } /* NO reAlloc() should have happened in setSafe() */ assertEquals(capacity, vector.getValueCapacity()); vector.setValueCount(capacity); for (int i = 0; i < capacity; i++) { if (i % 3 == 0) { assertNull(vector.getObject(i)); } else { assertEquals( Integer.toString(i), vector.getObject(i).toString(), "unexpected value at index: " + i); } } /* set lesser initial capacity than actually needed * to trigger reallocs in copyFromSafe() */ vector2.allocateNew(1024 * 10, 1024); int capacity2 = vector2.getValueCapacity(); assertTrue(capacity2 >= 1024); assertTrue(capacity2 <= capacity); for (int i = 0; i < capacity; i++) { vector2.copyFromSafe(i, i, vector); if (i % 3 == 0) { assertNull(vector2.getObject(i)); } else { assertEquals( Integer.toString(i), vector2.getObject(i).toString(), "unexpected value at index: " + i); } } /* 2 reAllocs should have happened in copyFromSafe() */ assertEquals(capacity, vector2.getValueCapacity()); vector2.setValueCount(capacity); for (int i = 0; i < capacity; i++) { if (i % 3 == 0) { assertNull(vector2.getObject(i)); } else { assertEquals( Integer.toString(i), vector2.getObject(i).toString(), "unexpected value at index: " + i); } } } } @Test public void testSetLastSetUsage() { try (final VarCharVector vector = new VarCharVector("myvector", allocator)) { vector.allocateNew(1024 * 10, 1024); setBytes(0, STR1, vector); setBytes(1, STR2, vector); setBytes(2, STR3, vector); setBytes(3, STR4, vector); setBytes(4, STR5, vector); setBytes(5, STR6, vector); /* Check current lastSet */ assertEquals(-1, vector.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); assertArrayEquals(STR5, vector.get(4)); assertArrayEquals(STR6, vector.get(5)); /* * If we don't do setLastSe(5) before setValueCount(), then the latter will corrupt * the value vector by filling in all positions [0,valuecount-1] will empty byte arrays. * Run the test by commenting out next line and we should see incorrect vector output. */ vector.setLastSet(5); vector.setValueCount(20); /* Check current lastSet */ assertEquals(19, vector.getLastSet()); /* Check the vector output again */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); assertArrayEquals(STR5, vector.get(4)); assertArrayEquals(STR6, vector.get(5)); assertEquals(0, vector.getValueLength(6)); assertEquals(0, vector.getValueLength(7)); assertEquals(0, vector.getValueLength(8)); assertEquals(0, vector.getValueLength(9)); assertEquals(0, vector.getValueLength(10)); assertEquals(0, vector.getValueLength(11)); assertEquals(0, vector.getValueLength(12)); assertEquals(0, vector.getValueLength(13)); assertEquals(0, vector.getValueLength(14)); assertEquals(0, vector.getValueLength(15)); assertEquals(0, vector.getValueLength(16)); assertEquals(0, vector.getValueLength(17)); assertEquals(0, vector.getValueLength(18)); assertEquals(0, vector.getValueLength(19)); /* Check offsets */ assertEquals(0, vector.offsetBuffer.getInt(0 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(6, vector.offsetBuffer.getInt(1 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(16, vector.offsetBuffer.getInt(2 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(21, vector.offsetBuffer.getInt(3 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(30, vector.offsetBuffer.getInt(4 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(34, vector.offsetBuffer.getInt(5 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(6 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(7 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(8 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(9 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(10 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(11 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(12 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(13 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(14 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(15 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(16 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(17 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(18 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(19 * BaseVariableWidthVector.OFFSET_WIDTH)); vector.set(19, STR6); assertArrayEquals(STR6, vector.get(19)); assertEquals(40, vector.offsetBuffer.getInt(19 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(46, vector.offsetBuffer.getInt(20 * BaseVariableWidthVector.OFFSET_WIDTH)); } } @Test public void testVectorLoadUnload() { try (final VarCharVector vector1 = new VarCharVector("myvector", allocator)) { setVector(vector1, STR1, STR2, STR3, STR4, STR5, STR6); assertEquals(5, vector1.getLastSet()); vector1.setValueCount(15); assertEquals(14, vector1.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector1.get(0)); assertArrayEquals(STR2, vector1.get(1)); assertArrayEquals(STR3, vector1.get(2)); assertArrayEquals(STR4, vector1.get(3)); assertArrayEquals(STR5, vector1.get(4)); assertArrayEquals(STR6, vector1.get(5)); Field field = vector1.getField(); String fieldName = field.getName(); List fields = new ArrayList<>(); List fieldVectors = new ArrayList<>(); fields.add(field); fieldVectors.add(vector1); Schema schema = new Schema(fields); VectorSchemaRoot schemaRoot1 = new VectorSchemaRoot(schema, fieldVectors, vector1.getValueCount()); VectorUnloader vectorUnloader = new VectorUnloader(schemaRoot1); try (ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch(); BufferAllocator finalVectorsAllocator = allocator.newChildAllocator("new vector", 0, Long.MAX_VALUE); VectorSchemaRoot schemaRoot2 = VectorSchemaRoot.create(schema, finalVectorsAllocator); ) { VectorLoader vectorLoader = new VectorLoader(schemaRoot2); vectorLoader.load(recordBatch); VarCharVector vector2 = (VarCharVector) schemaRoot2.getVector(fieldName); /* * lastSet would have internally been set by VectorLoader.load() when it invokes * loadFieldBuffers. */ assertEquals(14, vector2.getLastSet()); vector2.setValueCount(25); assertEquals(24, vector2.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector2.get(0)); assertArrayEquals(STR2, vector2.get(1)); assertArrayEquals(STR3, vector2.get(2)); assertArrayEquals(STR4, vector2.get(3)); assertArrayEquals(STR5, vector2.get(4)); assertArrayEquals(STR6, vector2.get(5)); } } } @Test public void testFillEmptiesUsage() { try (final VarCharVector vector = new VarCharVector("myvector", allocator)) { vector.allocateNew(1024 * 10, 1024); setBytes(0, STR1, vector); setBytes(1, STR2, vector); setBytes(2, STR3, vector); setBytes(3, STR4, vector); setBytes(4, STR5, vector); setBytes(5, STR6, vector); /* Check current lastSet */ assertEquals(-1, vector.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); assertArrayEquals(STR5, vector.get(4)); assertArrayEquals(STR6, vector.get(5)); vector.setLastSet(5); /* fill empty byte arrays from index [6, 9] */ vector.fillEmpties(10); /* Check current lastSet */ assertEquals(9, vector.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); assertArrayEquals(STR5, vector.get(4)); assertArrayEquals(STR6, vector.get(5)); assertEquals(0, vector.getValueLength(6)); assertEquals(0, vector.getValueLength(7)); assertEquals(0, vector.getValueLength(8)); assertEquals(0, vector.getValueLength(9)); setBytes(10, STR1, vector); setBytes(11, STR2, vector); vector.setLastSet(11); /* fill empty byte arrays from index [12, 14] */ vector.setValueCount(15); /* Check current lastSet */ assertEquals(14, vector.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); assertArrayEquals(STR5, vector.get(4)); assertArrayEquals(STR6, vector.get(5)); assertEquals(0, vector.getValueLength(6)); assertEquals(0, vector.getValueLength(7)); assertEquals(0, vector.getValueLength(8)); assertEquals(0, vector.getValueLength(9)); assertArrayEquals(STR1, vector.get(10)); assertArrayEquals(STR2, vector.get(11)); assertEquals(0, vector.getValueLength(12)); assertEquals(0, vector.getValueLength(13)); assertEquals(0, vector.getValueLength(14)); /* Check offsets */ assertEquals(0, vector.offsetBuffer.getInt(0 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(6, vector.offsetBuffer.getInt(1 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(16, vector.offsetBuffer.getInt(2 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(21, vector.offsetBuffer.getInt(3 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(30, vector.offsetBuffer.getInt(4 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(34, vector.offsetBuffer.getInt(5 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(6 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(7 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(8 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(9 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(40, vector.offsetBuffer.getInt(10 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(46, vector.offsetBuffer.getInt(11 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(56, vector.offsetBuffer.getInt(12 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(56, vector.offsetBuffer.getInt(13 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(56, vector.offsetBuffer.getInt(14 * BaseVariableWidthVector.OFFSET_WIDTH)); assertEquals(56, vector.offsetBuffer.getInt(15 * BaseVariableWidthVector.OFFSET_WIDTH)); } } @Test /* VarCharVector */ public void testGetBufferAddress1() { try (final VarCharVector vector = new VarCharVector("myvector", allocator)) { setVector(vector, STR1, STR2, STR3, STR4, STR5, STR6); vector.setValueCount(15); /* check the vector output */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); assertArrayEquals(STR5, vector.get(4)); assertArrayEquals(STR6, vector.get(5)); List buffers = vector.getFieldBuffers(); long bitAddress = vector.getValidityBufferAddress(); long offsetAddress = vector.getOffsetBufferAddress(); long dataAddress = vector.getDataBufferAddress(); assertEquals(3, buffers.size()); assertEquals(bitAddress, buffers.get(0).memoryAddress()); assertEquals(offsetAddress, buffers.get(1).memoryAddress()); assertEquals(dataAddress, buffers.get(2).memoryAddress()); } } @Test /* IntVector */ public void testGetBufferAddress2() { try (final IntVector vector = new IntVector("myvector", allocator)) { boolean error = false; vector.allocateNew(16); /* populate the vector */ for (int i = 0; i < 16; i += 2) { vector.set(i, i + 10); } /* check the vector output */ for (int i = 0; i < 16; i += 2) { assertEquals(i + 10, vector.get(i)); } List buffers = vector.getFieldBuffers(); long bitAddress = vector.getValidityBufferAddress(); long dataAddress = vector.getDataBufferAddress(); try { vector.getOffsetBufferAddress(); } catch (UnsupportedOperationException ue) { error = true; } finally { assertTrue(error); } assertEquals(2, buffers.size()); assertEquals(bitAddress, buffers.get(0).memoryAddress()); assertEquals(dataAddress, buffers.get(1).memoryAddress()); } } @Test public void testMultipleClose() { BufferAllocator vectorAllocator = allocator.newChildAllocator("vector_allocator", 0, Long.MAX_VALUE); IntVector vector = newVector(IntVector.class, EMPTY_SCHEMA_PATH, MinorType.INT, vectorAllocator); vector.close(); vectorAllocator.close(); vector.close(); vectorAllocator.close(); } /* this method is used by the tests to bypass the vector set methods that manipulate * lastSet. The method is to test the lastSet property and that's why we load the vector * in a way that lastSet is not set automatically. */ public static void setBytes(int index, byte[] bytes, VarCharVector vector) { final int currentOffset = vector.offsetBuffer.getInt(index * BaseVariableWidthVector.OFFSET_WIDTH); BitVectorHelper.setBit(vector.validityBuffer, index); vector.offsetBuffer.setInt( (index + 1) * BaseVariableWidthVector.OFFSET_WIDTH, currentOffset + bytes.length); vector.valueBuffer.setBytes(currentOffset, bytes, 0, bytes.length); } @Test /* VarCharVector */ public void testSetInitialCapacity() { try (final VarCharVector vector = new VarCharVector(EMPTY_SCHEMA_PATH, allocator)) { /* use the default 8 data bytes on average per element */ int defaultCapacity = BaseValueVector.INITIAL_VALUE_ALLOCATION - 1; vector.setInitialCapacity(defaultCapacity); vector.allocateNew(); assertEquals(defaultCapacity, vector.getValueCapacity()); assertEquals( CommonUtil.nextPowerOfTwo(defaultCapacity * 8), vector.getDataBuffer().capacity()); vector.setInitialCapacity(defaultCapacity, 1); vector.allocateNew(); assertEquals(defaultCapacity, vector.getValueCapacity()); assertEquals(CommonUtil.nextPowerOfTwo(defaultCapacity), vector.getDataBuffer().capacity()); vector.setInitialCapacity(defaultCapacity, 0.1); vector.allocateNew(); assertEquals(defaultCapacity, vector.getValueCapacity()); assertEquals( CommonUtil.nextPowerOfTwo((int) (defaultCapacity * 0.1)), vector.getDataBuffer().capacity()); vector.setInitialCapacity(defaultCapacity, 0.01); vector.allocateNew(); assertEquals(defaultCapacity, vector.getValueCapacity()); assertEquals( CommonUtil.nextPowerOfTwo((int) (defaultCapacity * 0.01)), vector.getDataBuffer().capacity()); vector.setInitialCapacity(5, 0.01); vector.allocateNew(); assertEquals(5, vector.getValueCapacity()); assertEquals(2, vector.getDataBuffer().capacity()); } } @Test public void testDefaultAllocNewAll() { int defaultCapacity = BaseValueVector.INITIAL_VALUE_ALLOCATION; int expectedSize; long beforeSize; try (BufferAllocator childAllocator = allocator.newChildAllocator("defaultAllocs", 0, Long.MAX_VALUE); final IntVector intVector = new IntVector(EMPTY_SCHEMA_PATH, childAllocator); final BigIntVector bigIntVector = new BigIntVector(EMPTY_SCHEMA_PATH, childAllocator); final BitVector bitVector = new BitVector(EMPTY_SCHEMA_PATH, childAllocator); final DecimalVector decimalVector = new DecimalVector(EMPTY_SCHEMA_PATH, childAllocator, 38, 6); final VarCharVector varCharVector = new VarCharVector(EMPTY_SCHEMA_PATH, childAllocator)) { // verify that the wastage is within bounds for IntVector. beforeSize = childAllocator.getAllocatedMemory(); intVector.allocateNew(); assertTrue(intVector.getValueCapacity() >= defaultCapacity); expectedSize = (defaultCapacity * IntVector.TYPE_WIDTH) + getValidityBufferSizeFromCount(defaultCapacity); assertTrue(childAllocator.getAllocatedMemory() - beforeSize <= expectedSize * 1.05); // verify that the wastage is within bounds for BigIntVector. beforeSize = childAllocator.getAllocatedMemory(); bigIntVector.allocateNew(); assertTrue(bigIntVector.getValueCapacity() >= defaultCapacity); expectedSize = (defaultCapacity * bigIntVector.TYPE_WIDTH) + getValidityBufferSizeFromCount(defaultCapacity); assertTrue(childAllocator.getAllocatedMemory() - beforeSize <= expectedSize * 1.05); // verify that the wastage is within bounds for DecimalVector. beforeSize = childAllocator.getAllocatedMemory(); decimalVector.allocateNew(); assertTrue(decimalVector.getValueCapacity() >= defaultCapacity); expectedSize = (defaultCapacity * decimalVector.TYPE_WIDTH) + getValidityBufferSizeFromCount(defaultCapacity); assertTrue(childAllocator.getAllocatedMemory() - beforeSize <= expectedSize * 1.05); // verify that the wastage is within bounds for VarCharVector. // var char vector have an offsets array that is 1 less than defaultCapacity beforeSize = childAllocator.getAllocatedMemory(); varCharVector.allocateNew(); assertTrue(varCharVector.getValueCapacity() >= defaultCapacity - 1); expectedSize = (defaultCapacity * VarCharVector.OFFSET_WIDTH) + getValidityBufferSizeFromCount(defaultCapacity) + defaultCapacity * 8; // wastage should be less than 5%. assertTrue(childAllocator.getAllocatedMemory() - beforeSize <= expectedSize * 1.05); // verify that the wastage is within bounds for BitVector. beforeSize = childAllocator.getAllocatedMemory(); bitVector.allocateNew(); assertTrue(bitVector.getValueCapacity() >= defaultCapacity); expectedSize = getValidityBufferSizeFromCount(defaultCapacity) * 2; assertTrue(childAllocator.getAllocatedMemory() - beforeSize <= expectedSize * 1.05); } } @Test public void testSetNullableVarCharHolder() { try (VarCharVector vector = new VarCharVector("", allocator)) { vector.allocateNew(100, 10); NullableVarCharHolder nullHolder = new NullableVarCharHolder(); nullHolder.isSet = 0; NullableVarCharHolder stringHolder = new NullableVarCharHolder(); stringHolder.isSet = 1; String str = "hello"; ArrowBuf buf = allocator.buffer(16); buf.setBytes(0, str.getBytes(StandardCharsets.UTF_8)); stringHolder.start = 0; stringHolder.end = str.length(); stringHolder.buffer = buf; vector.set(0, nullHolder); vector.set(1, stringHolder); // verify results assertTrue(vector.isNull(0)); assertEquals(str, new String(vector.get(1), StandardCharsets.UTF_8)); buf.close(); } } @Test public void testSetNullableVarCharHolderSafe() { try (VarCharVector vector = new VarCharVector("", allocator)) { vector.allocateNew(5, 1); NullableVarCharHolder nullHolder = new NullableVarCharHolder(); nullHolder.isSet = 0; NullableVarCharHolder stringHolder = new NullableVarCharHolder(); stringHolder.isSet = 1; String str = "hello world"; ArrowBuf buf = allocator.buffer(16); buf.setBytes(0, str.getBytes(StandardCharsets.UTF_8)); stringHolder.start = 0; stringHolder.end = str.length(); stringHolder.buffer = buf; vector.setSafe(0, stringHolder); vector.setSafe(1, nullHolder); // verify results assertEquals(str, new String(vector.get(0), StandardCharsets.UTF_8)); assertTrue(vector.isNull(1)); buf.close(); } } @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)) { vector.allocateNew(100, 10); NullableVarBinaryHolder nullHolder = new NullableVarBinaryHolder(); nullHolder.isSet = 0; NullableVarBinaryHolder binHolder = new NullableVarBinaryHolder(); binHolder.isSet = 1; String str = "hello"; ArrowBuf buf = allocator.buffer(16); buf.setBytes(0, str.getBytes(StandardCharsets.UTF_8)); binHolder.start = 0; binHolder.end = str.length(); binHolder.buffer = buf; vector.set(0, nullHolder); vector.set(1, binHolder); // verify results assertTrue(vector.isNull(0)); assertEquals(str, new String(vector.get(1), StandardCharsets.UTF_8)); buf.close(); } } @Test public void testSetNullableVarBinaryHolderSafe() { try (VarBinaryVector vector = new VarBinaryVector("", allocator)) { vector.allocateNew(5, 1); NullableVarBinaryHolder nullHolder = new NullableVarBinaryHolder(); nullHolder.isSet = 0; NullableVarBinaryHolder binHolder = new NullableVarBinaryHolder(); binHolder.isSet = 1; String str = "hello world"; ArrowBuf buf = allocator.buffer(16); buf.setBytes(0, str.getBytes(StandardCharsets.UTF_8)); binHolder.start = 0; binHolder.end = str.length(); binHolder.buffer = buf; vector.setSafe(0, binHolder); vector.setSafe(1, nullHolder); // verify results assertEquals(str, new String(vector.get(0), StandardCharsets.UTF_8)); assertTrue(vector.isNull(1)); buf.close(); } } @Test public void testGetPointerFixedWidth() { final int vectorLength = 100; try (IntVector vec1 = new IntVector("vec1", allocator); IntVector vec2 = new IntVector("vec2", allocator)) { vec1.allocateNew(vectorLength); vec2.allocateNew(vectorLength); for (int i = 0; i < vectorLength; i++) { if (i % 10 == 0) { vec1.setNull(i); vec2.setNull(i); } else { vec1.set(i, i * 1234); vec2.set(i, i * 1234); } } ArrowBufPointer ptr1 = new ArrowBufPointer(); ArrowBufPointer ptr2 = new ArrowBufPointer(); for (int i = 0; i < vectorLength; i++) { vec1.getDataPointer(i, ptr1); vec2.getDataPointer(i, ptr2); if (i % 10 == 0) { assertNull(ptr1.getBuf()); assertNull(ptr2.getBuf()); } assertTrue(ptr1.equals(ptr2)); assertTrue(ptr2.equals(ptr2)); } } } @Test public void testGetPointerVariableWidth() { final String[] sampleData = new String[] {"abc", "123", "def", null, "hello", "aaaaa", "world", "2019", null, "0717"}; try (VarCharVector vec1 = new VarCharVector("vec1", allocator); VarCharVector vec2 = new VarCharVector("vec2", allocator)) { vec1.allocateNew(sampleData.length * 10, sampleData.length); vec2.allocateNew(sampleData.length * 10, sampleData.length); for (int i = 0; i < sampleData.length; i++) { String str = sampleData[i]; if (str != null) { vec1.set(i, sampleData[i].getBytes(StandardCharsets.UTF_8)); vec2.set(i, sampleData[i].getBytes(StandardCharsets.UTF_8)); } else { vec1.setNull(i); vec2.setNull(i); } } ArrowBufPointer ptr1 = new ArrowBufPointer(); ArrowBufPointer ptr2 = new ArrowBufPointer(); for (int i = 0; i < sampleData.length; i++) { vec1.getDataPointer(i, ptr1); vec2.getDataPointer(i, ptr2); assertTrue(ptr1.equals(ptr2)); assertTrue(ptr2.equals(ptr2)); } } } @Test public void testGetNullFromVariableWidthVector() { try (final VarCharVector varCharVector = new VarCharVector("varcharvec", allocator); final VarBinaryVector varBinaryVector = new VarBinaryVector("varbinary", allocator)) { varCharVector.allocateNew(10, 1); varBinaryVector.allocateNew(10, 1); varCharVector.setNull(0); varBinaryVector.setNull(0); assertNull(varCharVector.get(0)); assertNull(varBinaryVector.get(0)); } } @Test public void testZeroVectorEquals() { try (final ZeroVector vector1 = new ZeroVector("vector"); final ZeroVector vector2 = new ZeroVector("vector")) { VectorEqualsVisitor visitor = new VectorEqualsVisitor(); assertTrue(visitor.vectorEquals(vector1, vector2)); } } @Test public void testZeroVectorNotEquals() { try (final IntVector intVector = new IntVector("int", allocator); final ZeroVector zeroVector = new ZeroVector("zero"); final ZeroVector zeroVector1 = new ZeroVector("zero1")) { VectorEqualsVisitor zeroVisitor = new VectorEqualsVisitor(); assertFalse(zeroVisitor.vectorEquals(intVector, zeroVector)); VectorEqualsVisitor intVisitor = new VectorEqualsVisitor(); assertFalse(intVisitor.vectorEquals(zeroVector, intVector)); VectorEqualsVisitor twoZeroVisitor = new VectorEqualsVisitor(); // they are not equal because of distinct names assertFalse(twoZeroVisitor.vectorEquals(zeroVector, zeroVector1)); } } @Test public void testBitVectorEquals() { try (final BitVector vector1 = new BitVector("bit", allocator); final BitVector vector2 = new BitVector("bit", allocator)) { setVector(vector1, 0, 1, 0); setVector(vector2, 1, 1, 0); VectorEqualsVisitor visitor = new VectorEqualsVisitor(); assertFalse(visitor.vectorEquals(vector1, vector2)); vector1.set(0, 1); assertTrue(visitor.vectorEquals(vector1, vector2)); } } @Test public void testIntVectorEqualsWithNull() { try (final IntVector vector1 = new IntVector("int", allocator); final IntVector vector2 = new IntVector("int", allocator)) { setVector(vector1, 1, 2); setVector(vector2, 1, null); VectorEqualsVisitor visitor = new VectorEqualsVisitor(); assertFalse(visitor.vectorEquals(vector1, vector2)); } } @Test public void testIntVectorEquals() { try (final IntVector vector1 = new IntVector("int", allocator); final IntVector vector2 = new IntVector("int", allocator)) { setVector(vector1, 1, 2, 3); setVector(vector2, 1, 2, null); VectorEqualsVisitor visitor = new VectorEqualsVisitor(); assertFalse(visitor.vectorEquals(vector1, vector2)); vector2.setValueCount(3); vector2.setSafe(2, 2); assertFalse(vector1.equals(vector2)); vector2.setSafe(2, 3); assertTrue(visitor.vectorEquals(vector1, vector2)); } } @Test public void testDecimalVectorEquals() { try (final DecimalVector vector1 = new DecimalVector("decimal", allocator, 3, 3); final DecimalVector vector2 = new DecimalVector("decimal", allocator, 3, 3); final DecimalVector vector3 = new DecimalVector("decimal", allocator, 3, 2)) { setVector(vector1, 100L, 200L); setVector(vector2, 100L, 200L); setVector(vector3, 100L, 200L); VectorEqualsVisitor visitor1 = new VectorEqualsVisitor(); VectorEqualsVisitor visitor2 = new VectorEqualsVisitor(); assertTrue(visitor1.vectorEquals(vector1, vector2)); assertFalse(visitor2.vectorEquals(vector1, vector3)); } } @Test public void testVarcharVectorEqualsWithNull() { try (final VarCharVector vector1 = new VarCharVector("varchar", allocator); final VarCharVector vector2 = new VarCharVector("varchar", allocator)) { setVector(vector1, STR1, STR2); setVector(vector2, STR1, null); VectorEqualsVisitor visitor = new VectorEqualsVisitor(); assertFalse(visitor.vectorEquals(vector1, vector2)); } } @Test public void testVarcharVectorEquals() { try (final VarCharVector vector1 = new VarCharVector("varchar", allocator); final VarCharVector vector2 = new VarCharVector("varchar", allocator)) { setVector(vector1, STR1, STR2, STR3); setVector(vector2, STR1, STR2); VectorEqualsVisitor visitor = new VectorEqualsVisitor(); assertFalse(visitor.vectorEquals(vector1, vector2)); vector2.setSafe(2, STR3, 0, STR3.length); vector2.setValueCount(3); assertTrue(visitor.vectorEquals(vector1, vector2)); } } @Test public void testVarBinaryVectorEquals() { try (final VarBinaryVector vector1 = new VarBinaryVector("binary", allocator); final VarBinaryVector vector2 = new VarBinaryVector("binary", allocator)) { setVector(vector1, STR1, STR2, STR3); setVector(vector2, STR1, STR2); VectorEqualsVisitor visitor = new VectorEqualsVisitor(); assertFalse(visitor.vectorEquals(vector1, vector2)); vector2.setSafe(2, STR3, 0, STR3.length); vector2.setValueCount(3); assertTrue(visitor.vectorEquals(vector1, vector2)); } } @Test public void testListVectorEqualsWithNull() { try (final ListVector vector1 = ListVector.empty("list", allocator); final ListVector vector2 = ListVector.empty("list", allocator); ) { UnionListWriter writer1 = vector1.getWriter(); writer1.allocate(); // set some values writeListVector(writer1, new int[] {1, 2}); writeListVector(writer1, new int[] {3, 4}); writeListVector(writer1, new int[] {}); writer1.setValueCount(3); UnionListWriter writer2 = vector2.getWriter(); writer2.allocate(); // set some values writeListVector(writer2, new int[] {1, 2}); writeListVector(writer2, new int[] {3, 4}); writer2.setValueCount(3); VectorEqualsVisitor visitor = new VectorEqualsVisitor(); assertFalse(visitor.vectorEquals(vector1, vector2)); } } @Test public void testListViewVectorEqualsWithNull() { try (final ListViewVector vector1 = ListViewVector.empty("listview", allocator); final ListViewVector vector2 = ListViewVector.empty("listview", allocator); ) { UnionListViewWriter writer1 = vector1.getWriter(); writer1.allocate(); // set some values writeListViewVector(writer1, new int[] {1, 2}); writeListViewVector(writer1, new int[] {3, 4}); writeListViewVector(writer1, new int[] {}); writer1.setValueCount(3); UnionListViewWriter writer2 = vector2.getWriter(); writer2.allocate(); // set some values writeListViewVector(writer2, new int[] {1, 2}); writeListViewVector(writer2, new int[] {3, 4}); writer2.setValueCount(3); VectorEqualsVisitor visitor = new VectorEqualsVisitor(); assertFalse(visitor.vectorEquals(vector1, vector2)); } } @Test public void testLargeListViewVectorEqualsWithNull() { try (final LargeListViewVector vector1 = LargeListViewVector.empty("largelistview", allocator); final LargeListViewVector vector2 = LargeListViewVector.empty("largelistview", allocator); ) { UnionLargeListViewWriter writer1 = vector1.getWriter(); writer1.allocate(); // set some values writeLargeListViewVector(writer1, new int[] {1, 2}); writeLargeListViewVector(writer1, new int[] {3, 4}); writeLargeListViewVector(writer1, new int[] {}); writer1.setValueCount(3); UnionLargeListViewWriter writer2 = vector2.getWriter(); writer2.allocate(); // set some values writeLargeListViewVector(writer2, new int[] {1, 2}); writeLargeListViewVector(writer2, new int[] {3, 4}); writer2.setValueCount(3); VectorEqualsVisitor visitor = new VectorEqualsVisitor(); assertFalse(visitor.vectorEquals(vector1, vector2)); } } @Test public void testListVectorEquals() { try (final ListVector vector1 = ListVector.empty("list", allocator); final ListVector vector2 = ListVector.empty("list", allocator); ) { UnionListWriter writer1 = vector1.getWriter(); writer1.allocate(); // set some values writeListVector(writer1, new int[] {1, 2}); writeListVector(writer1, new int[] {3, 4}); writeListVector(writer1, new int[] {5, 6}); writer1.setValueCount(3); UnionListWriter writer2 = vector2.getWriter(); writer2.allocate(); // set some values writeListVector(writer2, new int[] {1, 2}); writeListVector(writer2, new int[] {3, 4}); writer2.setValueCount(2); VectorEqualsVisitor visitor = new VectorEqualsVisitor(); assertFalse(visitor.vectorEquals(vector1, vector2)); writeListVector(writer2, new int[] {5, 6}); writer2.setValueCount(3); assertTrue(visitor.vectorEquals(vector1, vector2)); } } @Test public void testListViewVectorEquals() { try (final ListViewVector vector1 = ListViewVector.empty("listview", allocator); final ListViewVector vector2 = ListViewVector.empty("listview", allocator); ) { UnionListViewWriter writer1 = vector1.getWriter(); writer1.allocate(); // set some values writeListViewVector(writer1, new int[] {1, 2}); writeListViewVector(writer1, new int[] {3, 4}); writeListViewVector(writer1, new int[] {5, 6}); writer1.setValueCount(3); UnionListViewWriter writer2 = vector2.getWriter(); writer2.allocate(); // set some values writeListViewVector(writer2, new int[] {1, 2}); writeListViewVector(writer2, new int[] {3, 4}); writer2.setValueCount(2); VectorEqualsVisitor visitor = new VectorEqualsVisitor(); assertFalse(visitor.vectorEquals(vector1, vector2)); writeListViewVector(writer2, new int[] {5, 6}); writer2.setValueCount(3); assertTrue(visitor.vectorEquals(vector1, vector2)); } } @Test public void testLargeListViewVectorEquals() { try (final LargeListViewVector vector1 = LargeListViewVector.empty("largelistview", allocator); final LargeListViewVector vector2 = LargeListViewVector.empty("largelistview", allocator); ) { UnionLargeListViewWriter writer1 = vector1.getWriter(); writer1.allocate(); // set some values writeLargeListViewVector(writer1, new int[] {1, 2}); writeLargeListViewVector(writer1, new int[] {3, 4}); writeLargeListViewVector(writer1, new int[] {5, 6}); writer1.setValueCount(3); UnionLargeListViewWriter writer2 = vector2.getWriter(); writer2.allocate(); // set some values writeLargeListViewVector(writer2, new int[] {1, 2}); writeLargeListViewVector(writer2, new int[] {3, 4}); writer2.setValueCount(2); VectorEqualsVisitor visitor = new VectorEqualsVisitor(); assertFalse(visitor.vectorEquals(vector1, vector2)); writeLargeListViewVector(writer2, new int[] {5, 6}); writer2.setValueCount(3); assertTrue(visitor.vectorEquals(vector1, vector2)); } } @Test public void testListVectorSetNull() { try (final ListVector vector = ListVector.empty("list", allocator)) { UnionListWriter writer = vector.getWriter(); writer.allocate(); writeListVector(writer, new int[] {1, 2}); writeListVector(writer, new int[] {3, 4}); writeListVector(writer, new int[] {5, 6}); vector.setNull(3); vector.setNull(4); vector.setNull(5); writer.setValueCount(6); assertEquals(vector.getObject(0), Arrays.asList(1, 2)); assertEquals(vector.getObject(1), Arrays.asList(3, 4)); assertEquals(vector.getObject(2), Arrays.asList(5, 6)); assertTrue(vector.isNull(3)); assertTrue(vector.isNull(4)); assertTrue(vector.isNull(5)); } } @Test public void testListViewVectorSetNull() { try (final ListViewVector vector = ListViewVector.empty("listview", allocator)) { UnionListViewWriter writer = vector.getWriter(); writer.allocate(); writeListViewVector(writer, new int[] {1, 2}); writeListViewVector(writer, new int[] {3, 4}); writeListViewVector(writer, new int[] {5, 6}); vector.setNull(3); vector.setNull(4); vector.setNull(5); writer.setValueCount(6); assertEquals(vector.getObject(0), Arrays.asList(1, 2)); assertEquals(vector.getObject(1), Arrays.asList(3, 4)); assertEquals(vector.getObject(2), Arrays.asList(5, 6)); assertTrue(vector.isNull(3)); assertTrue(vector.isNull(4)); assertTrue(vector.isNull(5)); } } @Test public void testLargeListViewVectorSetNull() { try (final LargeListViewVector vector = LargeListViewVector.empty("largelistview", allocator)) { UnionLargeListViewWriter writer = vector.getWriter(); writer.allocate(); writeLargeListViewVector(writer, new int[] {1, 2}); writeLargeListViewVector(writer, new int[] {3, 4}); writeLargeListViewVector(writer, new int[] {5, 6}); vector.setNull(3); vector.setNull(4); vector.setNull(5); writer.setValueCount(6); assertEquals(vector.getObject(0), Arrays.asList(1, 2)); assertEquals(vector.getObject(1), Arrays.asList(3, 4)); assertEquals(vector.getObject(2), Arrays.asList(5, 6)); assertTrue(vector.isNull(3)); assertTrue(vector.isNull(4)); assertTrue(vector.isNull(5)); } } @Test public void testStructVectorEqualsWithNull() { try (final StructVector vector1 = StructVector.empty("struct", allocator); final StructVector vector2 = StructVector.empty("struct", allocator); ) { vector1.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); vector1.addOrGet("f1", FieldType.nullable(new ArrowType.Int(64, true)), BigIntVector.class); vector2.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); vector2.addOrGet("f1", FieldType.nullable(new ArrowType.Int(64, true)), BigIntVector.class); NullableStructWriter writer1 = vector1.getWriter(); writer1.allocate(); writeStructVector(writer1, 1, 10L); writeStructVector(writer1, 2, 20L); writeStructVector(writer1, 3, 30L); writer1.setValueCount(3); NullableStructWriter writer2 = vector2.getWriter(); writer2.allocate(); writeStructVector(writer2, 1, 10L); writeStructVector(writer2, 3, 30L); writer2.setValueCount(3); VectorEqualsVisitor visitor = new VectorEqualsVisitor(); assertFalse(visitor.vectorEquals(vector1, vector2)); } } @Test public void testStructVectorEquals() { try (final StructVector vector1 = StructVector.empty("struct", allocator); final StructVector vector2 = StructVector.empty("struct", allocator); ) { vector1.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); vector1.addOrGet("f1", FieldType.nullable(new ArrowType.Int(64, true)), BigIntVector.class); vector2.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); vector2.addOrGet("f1", FieldType.nullable(new ArrowType.Int(64, true)), BigIntVector.class); NullableStructWriter writer1 = vector1.getWriter(); writer1.allocate(); writeStructVector(writer1, 1, 10L); writeStructVector(writer1, 2, 20L); writeStructVector(writer1, 3, 30L); writer1.setValueCount(3); NullableStructWriter writer2 = vector2.getWriter(); writer2.allocate(); writeStructVector(writer2, 1, 10L); writeStructVector(writer2, 2, 20L); writer2.setValueCount(2); VectorEqualsVisitor visitor = new VectorEqualsVisitor(); assertFalse(visitor.vectorEquals(vector1, vector2)); writeStructVector(writer2, 3, 30L); writer2.setValueCount(3); assertTrue(visitor.vectorEquals(vector1, vector2)); } } @Test public void testStructVectorEqualsWithDiffChild() { try (final StructVector vector1 = StructVector.empty("struct", allocator); final StructVector vector2 = StructVector.empty("struct", allocator); ) { vector1.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); vector1.addOrGet("f1", FieldType.nullable(new ArrowType.Int(64, true)), BigIntVector.class); vector2.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); vector2.addOrGet("f10", FieldType.nullable(new ArrowType.Int(64, true)), BigIntVector.class); NullableStructWriter writer1 = vector1.getWriter(); writer1.allocate(); writeStructVector(writer1, 1, 10L); writeStructVector(writer1, 2, 20L); writer1.setValueCount(2); NullableStructWriter writer2 = vector2.getWriter(); writer2.allocate(); writeStructVector(writer2, 1, 10L); writeStructVector(writer2, 2, 20L); writer2.setValueCount(2); VectorEqualsVisitor visitor = new VectorEqualsVisitor(); assertFalse(visitor.vectorEquals(vector1, vector2)); } } @Test public void testStructVectorAcceptsDenseUnionChild() { Field childField = new Field( "child", FieldType.notNullable(new ArrowType.Union(UnionMode.Dense, new int[] {})), Collections.emptyList()); Field structField = new Field( "struct", FieldType.notNullable(ArrowType.Struct.INSTANCE), Collections.singletonList(childField)); try (FieldVector structVec = structField.createVector(allocator)) { assertEquals(structField, structVec.getField()); } } @Test public void testUnionVectorEquals() { try (final UnionVector vector1 = new UnionVector("union", allocator, /* field type */ null, /* call-back */ null); final UnionVector vector2 = new UnionVector("union", allocator, /* field type */ null, /* call-back */ null); ) { final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); uInt4Holder.value = 10; uInt4Holder.isSet = 1; final NullableIntHolder intHolder = new NullableIntHolder(); uInt4Holder.value = 20; uInt4Holder.isSet = 1; vector1.setType(0, Types.MinorType.UINT4); vector1.setSafe(0, uInt4Holder); vector1.setType(1, Types.MinorType.INT); vector1.setSafe(1, intHolder); vector1.setValueCount(2); vector2.setType(0, Types.MinorType.UINT4); vector2.setSafe(0, uInt4Holder); vector2.setType(1, Types.MinorType.INT); vector2.setSafe(1, intHolder); vector2.setValueCount(2); VectorEqualsVisitor visitor = new VectorEqualsVisitor(); assertTrue(visitor.vectorEquals(vector1, vector2)); } } @Test public void testEqualsWithIndexOutOfRange() { assertThrows( IllegalArgumentException.class, () -> { try (final IntVector vector1 = new IntVector("int", allocator); final IntVector vector2 = new IntVector("int", allocator)) { setVector(vector1, 1, 2); setVector(vector2, 1, 2); assertTrue(new RangeEqualsVisitor(vector1, vector2).rangeEquals(new Range(2, 3, 1))); } }); } @Test public void testFixedWidthVectorNullHashCode() { try (IntVector intVec = new IntVector("int vector", allocator)) { intVec.allocateNew(1); intVec.setValueCount(1); intVec.set(0, 100); intVec.setNull(0); assertEquals(0, intVec.hashCode(0)); } } @Test public void testVariableWidthVectorNullHashCode() { try (VarCharVector varChVec = new VarCharVector("var char vector", allocator)) { varChVec.allocateNew(100, 1); varChVec.setValueCount(1); varChVec.set(0, "abc".getBytes(StandardCharsets.UTF_8)); varChVec.setNull(0); assertEquals(0, varChVec.hashCode(0)); } } @Test public void testUnionNullHashCode() { try (UnionVector srcVector = new UnionVector( EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) { srcVector.allocateNew(); final NullableIntHolder holder = new NullableIntHolder(); holder.isSet = 0; // write some data srcVector.setType(0, MinorType.INT); srcVector.setSafe(0, holder); assertEquals(0, srcVector.hashCode(0)); } } @Test public void testToString() { try (final IntVector intVector = new IntVector("intVector", allocator); final ListVector listVector = ListVector.empty("listVector", allocator); final StructVector structVector = StructVector.empty("structVector", allocator)) { // validate intVector toString assertEquals("[]", intVector.toString()); intVector.setValueCount(3); intVector.setSafe(0, 1); intVector.setSafe(1, 2); intVector.setSafe(2, 3); assertEquals("[1, 2, 3]", intVector.toString()); // validate intVector with plenty values intVector.setValueCount(100); for (int i = 0; i < 100; i++) { intVector.setSafe(i, i); } assertEquals( "[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, ... 90, 91, 92, 93, 94, 95, 96, 97, 98, 99]", intVector.toString()); // validate listVector toString listVector.allocateNewSafe(); listVector.initializeChildrenFromFields( Collections.singletonList(Field.nullable("child", ArrowType.Utf8.INSTANCE))); VarCharVector dataVector = (VarCharVector) listVector.getDataVector(); listVector.startNewValue(0); dataVector.setSafe(0, "aaa".getBytes(StandardCharsets.UTF_8)); dataVector.setSafe(1, "bbb".getBytes(StandardCharsets.UTF_8)); listVector.endValue(0, 2); listVector.startNewValue(1); dataVector.setSafe(2, "ccc".getBytes(StandardCharsets.UTF_8)); dataVector.setSafe(3, "ddd".getBytes(StandardCharsets.UTF_8)); listVector.endValue(1, 2); listVector.setValueCount(2); assertEquals("[[\"aaa\",\"bbb\"], [\"ccc\",\"ddd\"]]", listVector.toString()); // validate structVector toString structVector.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); structVector.addOrGet( "f1", FieldType.nullable(new ArrowType.Int(64, true)), BigIntVector.class); NullableStructWriter structWriter = structVector.getWriter(); structWriter.allocate(); writeStructVector(structWriter, 1, 10L); writeStructVector(structWriter, 2, 20L); structWriter.setValueCount(2); assertEquals("[{\"f0\":1,\"f1\":10}, {\"f0\":2,\"f1\":20}]", structVector.toString()); } } @Test public void testUInt1VectorToString() { try (final UInt1Vector uInt1Vector = new UInt1Vector("uInt1Vector", allocator)) { setVector(uInt1Vector, (byte) 0xff); assertEquals("[255]", uInt1Vector.toString()); } } @Test public void testUInt2VectorToString() { try (final UInt2Vector uInt2Vector = new UInt2Vector("uInt2Vector", allocator)) { setVector(uInt2Vector, (char) 0xffff); assertEquals("[65535]", uInt2Vector.toString()); } } @Test public void testUInt4VectorToString() { try (final UInt4Vector uInt4Vector = new UInt4Vector("uInt4Vector", allocator)) { setVector(uInt4Vector, 0xffffffff); assertEquals("[4294967295]", uInt4Vector.toString()); } } @Test public void testUInt8VectorToString() { try (final UInt8Vector uInt8Vector = new UInt8Vector("uInt8Vector", allocator)) { setVector(uInt8Vector, 0xffffffffffffffffL); assertEquals("[18446744073709551615]", uInt8Vector.toString()); } } @Test public void testUnloadVariableWidthVector() { try (final VarCharVector varCharVector = new VarCharVector("var char", allocator)) { varCharVector.allocateNew(5, 2); varCharVector.setValueCount(2); varCharVector.set(0, "abcd".getBytes(StandardCharsets.UTF_8)); List bufs = varCharVector.getFieldBuffers(); assertEquals(3, bufs.size()); ArrowBuf offsetBuf = bufs.get(1); ArrowBuf dataBuf = bufs.get(2); assertEquals(12, offsetBuf.writerIndex()); assertEquals(4, offsetBuf.getInt(4)); assertEquals(4, offsetBuf.getInt(8)); assertEquals(4, dataBuf.writerIndex()); } } private void writeStructVector(NullableStructWriter writer, int value1, long value2) { writer.start(); writer.integer("f0").writeInt(value1); writer.bigInt("f1").writeBigInt(value2); writer.end(); } private void writeListVector(UnionListWriter writer, int[] values) { writer.startList(); for (int v : values) { writer.integer().writeInt(v); } writer.endList(); } private void writeListViewVector(UnionListViewWriter writer, int[] values) { writer.startListView(); for (int v : values) { writer.integer().writeInt(v); } writer.endListView(); } private void writeLargeListViewVector(UnionLargeListViewWriter writer, int[] values) { writer.startListView(); for (int v : values) { writer.integer().writeInt(v); } writer.endListView(); } @Test public void testVariableVectorGetEndOffset() { try (final VarCharVector vector1 = new VarCharVector("v1", allocator); final VarBinaryVector vector2 = new VarBinaryVector("v2", allocator)) { setVector(vector1, STR1, null, STR2); setVector(vector2, STR1, STR2, STR3); assertEquals(0, vector1.getStartOffset(0)); assertEquals(STR1.length, vector1.getEndOffset(0)); assertEquals(STR1.length, vector1.getStartOffset(1)); assertEquals(STR1.length, vector1.getEndOffset(1)); assertEquals(STR1.length, vector1.getStartOffset(2)); assertEquals(STR1.length + STR2.length, vector1.getEndOffset(2)); assertEquals(0, vector2.getStartOffset(0)); assertEquals(STR1.length, vector2.getEndOffset(0)); assertEquals(STR1.length, vector2.getStartOffset(1)); assertEquals(STR1.length + STR2.length, vector2.getEndOffset(1)); assertEquals(STR1.length + STR2.length, vector2.getStartOffset(2)); assertEquals(STR1.length + STR2.length + STR3.length, vector2.getEndOffset(2)); } } @Test public void testEmptyBufBehavior() { final int valueCount = 10; try (final IntVector vector = new IntVector("v", allocator)) { assertEquals(1, vector.getDataBuffer().refCnt()); assertEquals(1, vector.getValidityBuffer().refCnt()); assertEquals(0, vector.getDataBuffer().capacity()); assertEquals(0, vector.getValidityBuffer().capacity()); vector.allocateNew(valueCount); assertEquals(2, vector.getDataBuffer().refCnt()); assertEquals(2, vector.getValidityBuffer().refCnt()); assertEquals(56, vector.getDataBuffer().capacity()); assertEquals(8, vector.getValidityBuffer().capacity()); vector.close(); assertEquals(1, vector.getDataBuffer().refCnt()); assertEquals(1, vector.getValidityBuffer().refCnt()); assertEquals(0, vector.getDataBuffer().capacity()); assertEquals(0, vector.getValidityBuffer().capacity()); } try (final VarCharVector vector = new VarCharVector("v", allocator)) { assertEquals(1, vector.getDataBuffer().refCnt()); assertEquals(1, vector.getValidityBuffer().refCnt()); assertEquals(1, vector.getOffsetBuffer().refCnt()); assertEquals(0, vector.getDataBuffer().capacity()); assertEquals(0, vector.getValidityBuffer().capacity()); assertEquals(0, vector.getOffsetBuffer().capacity()); vector.allocateNew(valueCount); assertEquals(1, vector.getDataBuffer().refCnt()); assertEquals(2, vector.getValidityBuffer().refCnt()); assertEquals(2, vector.getOffsetBuffer().refCnt()); assertEquals(32768, vector.getDataBuffer().capacity()); assertEquals(8, vector.getValidityBuffer().capacity()); assertEquals(56, vector.getOffsetBuffer().capacity()); vector.close(); assertEquals(1, vector.getDataBuffer().refCnt()); assertEquals(1, vector.getValidityBuffer().refCnt()); assertEquals(1, vector.getOffsetBuffer().refCnt()); assertEquals(0, vector.getDataBuffer().capacity()); assertEquals(0, vector.getValidityBuffer().capacity()); assertEquals(0, vector.getOffsetBuffer().capacity()); } try (final ListVector vector = ListVector.empty("v", allocator)) { assertEquals(1, vector.getValidityBuffer().refCnt()); assertEquals(1, vector.getOffsetBuffer().refCnt()); assertEquals(0, vector.getValidityBuffer().capacity()); assertEquals(0, vector.getOffsetBuffer().capacity()); vector.setValueCount(valueCount); vector.allocateNewSafe(); assertEquals(1, vector.getValidityBuffer().refCnt()); assertEquals(1, vector.getOffsetBuffer().refCnt()); assertEquals(512, vector.getValidityBuffer().capacity()); assertEquals(16384, vector.getOffsetBuffer().capacity()); vector.close(); assertEquals(1, vector.getValidityBuffer().refCnt()); assertEquals(1, vector.getOffsetBuffer().refCnt()); assertEquals(0, vector.getValidityBuffer().capacity()); assertEquals(0, vector.getOffsetBuffer().capacity()); } try (final FixedSizeListVector vector = FixedSizeListVector.empty("v", 2, allocator)) { assertEquals(1, vector.getValidityBuffer().refCnt()); assertEquals(0, vector.getValidityBuffer().capacity()); vector.setValueCount(10); vector.allocateNewSafe(); assertEquals(1, vector.getValidityBuffer().refCnt()); assertEquals(512, vector.getValidityBuffer().capacity()); vector.close(); assertEquals(1, vector.getValidityBuffer().refCnt()); assertEquals(0, vector.getValidityBuffer().capacity()); } try (final StructVector vector = StructVector.empty("v", allocator)) { assertEquals(1, vector.getValidityBuffer().refCnt()); assertEquals(0, vector.getValidityBuffer().capacity()); vector.setValueCount(valueCount); vector.allocateNewSafe(); assertEquals(1, vector.getValidityBuffer().refCnt()); assertEquals(512, vector.getValidityBuffer().capacity()); vector.close(); assertEquals(1, vector.getValidityBuffer().refCnt()); assertEquals(0, vector.getValidityBuffer().capacity()); } try (final UnionVector vector = UnionVector.empty("v", allocator)) { assertEquals(1, vector.getTypeBuffer().refCnt()); assertEquals(0, vector.getTypeBuffer().capacity()); vector.setValueCount(10); vector.allocateNewSafe(); assertEquals(1, vector.getTypeBuffer().refCnt()); assertEquals(4096, vector.getTypeBuffer().capacity()); vector.close(); assertEquals(1, vector.getTypeBuffer().refCnt()); assertEquals(0, vector.getTypeBuffer().capacity()); } try (final DenseUnionVector vector = DenseUnionVector.empty("v", allocator)) { assertEquals(1, vector.getTypeBuffer().refCnt()); assertEquals(1, vector.getOffsetBuffer().refCnt()); assertEquals(0, vector.getTypeBuffer().capacity()); assertEquals(0, vector.getOffsetBuffer().capacity()); vector.setValueCount(valueCount); vector.allocateNew(); assertEquals(1, vector.getTypeBuffer().refCnt()); assertEquals(1, vector.getOffsetBuffer().refCnt()); assertEquals(4096, vector.getTypeBuffer().capacity()); assertEquals(16384, vector.getOffsetBuffer().capacity()); vector.close(); assertEquals(1, vector.getTypeBuffer().refCnt()); assertEquals(1, vector.getOffsetBuffer().refCnt()); assertEquals(0, vector.getTypeBuffer().capacity()); assertEquals(0, vector.getOffsetBuffer().capacity()); } } @Test public void testSetGetUInt1() { try (UInt1Vector vector = new UInt1Vector("vector", allocator)) { vector.allocateNew(2); vector.setWithPossibleTruncate(0, UInt1Vector.MAX_UINT1); vector.setUnsafeWithPossibleTruncate(1, UInt1Vector.MAX_UINT1); vector.setValueCount(2); assertEquals(UInt1Vector.MAX_UINT1 & UInt1Vector.PROMOTION_MASK, vector.getValueAsLong(0)); assertEquals(UInt1Vector.MAX_UINT1 & UInt1Vector.PROMOTION_MASK, vector.getValueAsLong(1)); } } @Test public void testSetGetUInt2() { try (UInt2Vector vector = new UInt2Vector("vector", allocator)) { vector.allocateNew(2); vector.setWithPossibleTruncate(0, UInt2Vector.MAX_UINT2); vector.setUnsafeWithPossibleTruncate(1, UInt2Vector.MAX_UINT2); vector.setValueCount(2); assertEquals(UInt2Vector.MAX_UINT2, vector.getValueAsLong(0)); assertEquals(UInt2Vector.MAX_UINT2, vector.getValueAsLong(1)); } } @Test public void testSetGetUInt4() { try (UInt4Vector vector = new UInt4Vector("vector", allocator)) { vector.allocateNew(2); vector.setWithPossibleTruncate(0, UInt4Vector.MAX_UINT4); vector.setUnsafeWithPossibleTruncate(1, UInt4Vector.MAX_UINT4); vector.setValueCount(2); long expected = UInt4Vector.MAX_UINT4 & UInt4Vector.PROMOTION_MASK; assertEquals(expected, vector.getValueAsLong(0)); assertEquals(expected, vector.getValueAsLong(1)); } } @Test public void testSplitAndTransferFixedWithVector1() { RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); try (BufferAllocator child = allocator.newChildAllocator("child", 0, Long.MAX_VALUE)) { try (IntVector vector = new IntVector("vector", child)) { vector.setSafe(0, 1); vector.setSafe(1, 2); vector.setSafe(2, 3); vector.setValueCount(3); TransferPair transferPair = vector.getTransferPair(allocator); transferPair.splitAndTransfer(0, 1); try (IntVector target = (IntVector) transferPair.getTo()) { // no-op try-with-resource assertEquals(1, target.get(0)); } } } } @Test public void testSplitAndTransferFixedWithVector2() { IntVector target; try (BufferAllocator child = allocator.newChildAllocator("child", 0, Long.MAX_VALUE)) { try (IntVector vector = new IntVector("source", child)) { vector.setSafe(0, 1); vector.setSafe(1, 2); vector.setSafe(2, 3); vector.setValueCount(3); TransferPair transferPair = vector.getTransferPair(allocator); transferPair.splitAndTransfer(0, 1); target = (IntVector) transferPair.getTo(); assertEquals(1, target.get(0)); } } target.close(); } @Test public void testVectorLoadUnloadOnNonVariadicVectors() { try (final IntVector vector1 = new IntVector("myvector", allocator)) { setVector(vector1, 1, 2, 3, 4, 5, 6); vector1.setValueCount(15); /* Check the vector output */ assertEquals(1, vector1.get(0)); assertEquals(2, vector1.get(1)); assertEquals(3, vector1.get(2)); assertEquals(4, vector1.get(3)); assertEquals(5, vector1.get(4)); assertEquals(6, vector1.get(5)); Field field = vector1.getField(); String fieldName = field.getName(); List fields = new ArrayList<>(); List fieldVectors = new ArrayList<>(); fields.add(field); fieldVectors.add(vector1); Schema schema = new Schema(fields); VectorSchemaRoot schemaRoot1 = new VectorSchemaRoot(schema, fieldVectors, vector1.getValueCount()); VectorUnloader vectorUnloader = new VectorUnloader(schemaRoot1); try (ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch(); BufferAllocator finalVectorsAllocator = allocator.newChildAllocator("new vector", 0, Long.MAX_VALUE); VectorSchemaRoot schemaRoot2 = VectorSchemaRoot.create(schema, finalVectorsAllocator); ) { // validating recordBatch doesn't contain an output for variadicBufferCounts assertTrue(recordBatch.getVariadicBufferCounts().isEmpty()); VectorLoader vectorLoader = new VectorLoader(schemaRoot2); vectorLoader.load(recordBatch); IntVector vector2 = (IntVector) schemaRoot2.getVector(fieldName); vector2.setValueCount(25); /* Check the vector output */ assertEquals(1, vector2.get(0)); assertEquals(2, vector2.get(1)); assertEquals(3, vector2.get(2)); assertEquals(4, vector2.get(3)); assertEquals(5, vector2.get(4)); assertEquals(6, vector2.get(5)); } } } @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)); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestValueVectorIterable.java000066400000000000000000000747261515451615200331650ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.hamcrest.MatcherAssert.assertThat; import java.math.BigDecimal; import java.math.RoundingMode; import java.time.Duration; import java.time.LocalDateTime; import java.time.Period; import java.time.ZoneOffset; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import org.apache.arrow.memory.BufferAllocator; 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.ListViewVector; import org.apache.arrow.vector.complex.NonNullableStructVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.impl.NullableStructWriter; import org.apache.arrow.vector.complex.impl.UnionLargeListWriter; import org.apache.arrow.vector.complex.impl.UnionListViewWriter; import org.apache.arrow.vector.complex.impl.UnionListWriter; 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.types.pojo.FieldType; import org.apache.arrow.vector.util.Text; import org.hamcrest.collection.IsIterableContainingInOrder; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestValueVectorIterable { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testBigIntVectorIterable() { try (final BigIntVector bigIntVector = new BigIntVector("bigInt", allocator)) { bigIntVector.allocateNew(3); bigIntVector.setSafe(0, 6); bigIntVector.setSafe(1, 2); bigIntVector.setSafe(2, 19); bigIntVector.setValueCount(3); assertThat( bigIntVector.getValueIterable(), IsIterableContainingInOrder.contains(6L, 2L, 19L)); } } @Test public void testBitVectorIterable() { try (final BitVector bitVector = new BitVector("bit", allocator)) { bitVector.allocateNew(3); bitVector.setSafe(0, 0); bitVector.setNull(1); bitVector.setSafe(2, 1); bitVector.setValueCount(3); assertThat( bitVector.getValueIterable(), IsIterableContainingInOrder.contains(false, null, true)); } } @Test public void testDateDayVectorIterable() { try (final DateDayVector dateDayVector = new DateDayVector("dateDay", allocator)) { dateDayVector.allocateNew(3); dateDayVector.setSafe(0, 30000); dateDayVector.setNull(1); dateDayVector.setSafe(2, 555); dateDayVector.setValueCount(3); assertThat( dateDayVector.getValueIterable(), IsIterableContainingInOrder.contains(30000, null, 555)); } } @Test public void testDateMilliVectorIterable() { try (final DateMilliVector dateMilliVector = new DateMilliVector("dateMilli", allocator)) { dateMilliVector.allocateNew(3); dateMilliVector.setSafe(0, 30000L); dateMilliVector.setNull(1); dateMilliVector.setSafe(2, 555L); dateMilliVector.setValueCount(3); final LocalDateTime value1 = LocalDateTime.ofEpochSecond(30L, 0, ZoneOffset.ofHours(0)); final LocalDateTime value3 = LocalDateTime.ofEpochSecond(0L, 555000000, ZoneOffset.ofHours(0)); assertThat( dateMilliVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testDecimal256VectorIterable() { try (final Decimal256Vector decimal256Vector = new Decimal256Vector("decimal256", allocator, 8, 2)) { decimal256Vector.allocateNew(3); decimal256Vector.setSafe(0, 30000L); decimal256Vector.setNull(1); decimal256Vector.setSafe(2, 555L); decimal256Vector.setValueCount(3); final BigDecimal value1 = new BigDecimal(300L).setScale(2, RoundingMode.HALF_UP); final BigDecimal value3 = new BigDecimal(555L).scaleByPowerOfTen(-2).setScale(2, RoundingMode.HALF_UP); assertThat( decimal256Vector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testDecimalVectorIterable() { try (final DecimalVector decimalVector = new DecimalVector("decimalDay", allocator, 8, 2)) { decimalVector.allocateNew(3); decimalVector.setSafe(0, 30000); decimalVector.setNull(1); decimalVector.setSafe(2, 555); decimalVector.setValueCount(3); final BigDecimal value1 = new BigDecimal(300L).setScale(2, RoundingMode.HALF_UP); final BigDecimal value3 = new BigDecimal(555L).scaleByPowerOfTen(-2).setScale(2, RoundingMode.HALF_UP); assertThat( decimalVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testDurationVectorIterable() { try (final DurationVector durationVector = new DurationVector( Field.nullablePrimitive("duration", new ArrowType.Duration(TimeUnit.MILLISECOND)), allocator)) { durationVector.allocateNew(3); durationVector.setSafe(0, 30000); durationVector.setNull(1); durationVector.setSafe(2, 555); durationVector.setValueCount(3); final Duration value1 = Duration.ofMillis(30000); final Duration value3 = Duration.ofMillis(555); assertThat( durationVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testFixedSizeBinaryVectorIterable() { try (final FixedSizeBinaryVector fixedSizeBinaryVector = new FixedSizeBinaryVector("binary", allocator, 4)) { final byte[] value1 = new byte[] {0, 0, 0, 1}; final byte[] value3 = new byte[] {1, 0, 0, 0}; fixedSizeBinaryVector.allocateNew(3); fixedSizeBinaryVector.setSafe(0, value1); fixedSizeBinaryVector.setNull(1); fixedSizeBinaryVector.setSafe(2, value3); fixedSizeBinaryVector.setValueCount(3); assertThat( fixedSizeBinaryVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testFloat2VectorIterable() { try (final Float2Vector float2Vector = new Float2Vector("float2", allocator)) { float2Vector.allocateNew(3); float2Vector.setSafe(0, (short) 7777); float2Vector.setNull(1); float2Vector.setSafe(2, (short) 5); float2Vector.setValueCount(3); assertThat( float2Vector.getValueIterable(), IsIterableContainingInOrder.contains((short) 7777, null, (short) 5)); } } @Test public void testFloat4VectorIterable() { try (final Float4Vector float4Vector = new Float4Vector("float4", allocator)) { float4Vector.allocateNew(3); float4Vector.setSafe(0, 16.32f); float4Vector.setNull(1); float4Vector.setSafe(2, -10.75f); float4Vector.setValueCount(3); assertThat( float4Vector.getValueIterable(), IsIterableContainingInOrder.contains(16.32f, null, -10.75f)); } } @Test public void testFloat8VectorIterable() { try (final Float8Vector float8Vector = new Float8Vector("float8", allocator)) { float8Vector.allocateNew(3); float8Vector.setSafe(0, 16.32); float8Vector.setNull(1); float8Vector.setSafe(2, -10.75); float8Vector.setValueCount(3); assertThat( float8Vector.getValueIterable(), IsIterableContainingInOrder.contains(16.32, null, -10.75)); } } @Test public void testIntVectorIterable() { try (final IntVector intVector = new IntVector("int", allocator)) { intVector.allocateNew(3); intVector.setSafe(0, 78); intVector.setNull(1); intVector.setSafe(2, -93); intVector.setValueCount(3); assertThat(intVector.getValueIterable(), IsIterableContainingInOrder.contains(78, null, -93)); } } @Test public void testIntervalDayVectorIterable() { try (final IntervalDayVector intervalDayVector = new IntervalDayVector("intervalDay", allocator)) { intervalDayVector.allocateNew(3); intervalDayVector.setSafe(0, 63, 0); intervalDayVector.setNull(1); intervalDayVector.setSafe(2, 555, 0); intervalDayVector.setValueCount(3); final Duration value1 = Duration.ofDays(63); final Duration value3 = Duration.ofDays(555); assertThat( intervalDayVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testIntervalMonthDayNanoVectorIterable() { try (final IntervalMonthDayNanoVector intervalMonthDayNanoVector = new IntervalMonthDayNanoVector("intervalMonthDayNano", allocator)) { intervalMonthDayNanoVector.allocateNew(3); intervalMonthDayNanoVector.setSafe(0, 3, 4, 0); intervalMonthDayNanoVector.setNull(1); intervalMonthDayNanoVector.setSafe(2, 7, 18, 0); intervalMonthDayNanoVector.setValueCount(3); final PeriodDuration value1 = new PeriodDuration(Period.of(0, 3, 4), Duration.ofSeconds(0)); final PeriodDuration value3 = new PeriodDuration(Period.of(0, 7, 18), Duration.ofSeconds(0)); assertThat( intervalMonthDayNanoVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testIntervalYearVectorIterable() { try (final IntervalYearVector intervalYearVector = new IntervalYearVector("intervalYear", allocator)) { intervalYearVector.allocateNew(3); intervalYearVector.setSafe(0, 3); intervalYearVector.setNull(1); intervalYearVector.setSafe(2, 17); intervalYearVector.setValueCount(3); final Period value1 = Period.ofMonths(3); final Period value3 = Period.ofMonths(17); assertThat( intervalYearVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testLargeVarBinaryVectorIterable() { try (final LargeVarBinaryVector largeVarBinaryVector = new LargeVarBinaryVector("largeVarBinary", allocator)) { final byte[] value1 = new byte[] {0, 0, 0, 1}; final byte[] value3 = new byte[] {1, 0, 0}; largeVarBinaryVector.allocateNew(3); largeVarBinaryVector.setSafe(0, value1); largeVarBinaryVector.setNull(1); largeVarBinaryVector.setSafe(2, value3); largeVarBinaryVector.setValueCount(3); assertThat( largeVarBinaryVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testLargeVarCharVectorIterable() { try (final LargeVarCharVector largeVarCharVector = new LargeVarCharVector("largeVarChar", allocator)) { final Text value1 = new Text("hello"); final Text value3 = new Text("worlds"); largeVarCharVector.allocateNew(3); largeVarCharVector.setSafe(0, value1); largeVarCharVector.setNull(1); largeVarCharVector.setSafe(2, value3); largeVarCharVector.setValueCount(3); assertThat( largeVarCharVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testNullVectorIterable() { try (final NullVector nullVector = new NullVector("null", 3)) { assertThat( nullVector.getValueIterable(), IsIterableContainingInOrder.contains(null, null, null)); } } @Test public void testSmallIntVectorIterable() { try (final SmallIntVector smallIntVector = new SmallIntVector("smallInt", allocator)) { smallIntVector.allocateNew(3); smallIntVector.setSafe(0, 78); smallIntVector.setNull(1); smallIntVector.setSafe(2, -93); smallIntVector.setValueCount(3); assertThat( smallIntVector.getValueIterable(), IsIterableContainingInOrder.contains((short) 78, null, (short) -93)); } } @Test public void testTimeMicroVectorIterable() { try (final TimeMicroVector timeMicroVector = new TimeMicroVector("timeMicro", allocator)) { timeMicroVector.allocateNew(3); timeMicroVector.setSafe(0, 70000); timeMicroVector.setNull(1); timeMicroVector.setSafe(2, 555); timeMicroVector.setValueCount(3); assertThat( timeMicroVector.getValueIterable(), IsIterableContainingInOrder.contains(70000L, null, 555L)); } } @Test public void testTimeMilliVectorIterable() { try (final TimeMilliVector timeMilliVector = new TimeMilliVector("timeMilli", allocator)) { timeMilliVector.allocateNew(3); timeMilliVector.setSafe(0, 70000); timeMilliVector.setNull(1); timeMilliVector.setSafe(2, 555); timeMilliVector.setValueCount(3); final LocalDateTime value1 = LocalDateTime.ofEpochSecond(70L, 0, ZoneOffset.ofHours(0)); final LocalDateTime value3 = LocalDateTime.ofEpochSecond(0L, 555000000, ZoneOffset.ofHours(0)); assertThat( timeMilliVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testTimeNanoVectorIterable() { try (final TimeNanoVector timeNanoVector = new TimeNanoVector("timeNano", allocator)) { timeNanoVector.allocateNew(3); timeNanoVector.setSafe(0, 70000); timeNanoVector.setNull(1); timeNanoVector.setSafe(2, 555); timeNanoVector.setValueCount(3); assertThat( timeNanoVector.getValueIterable(), IsIterableContainingInOrder.contains(70000L, null, 555L)); } } @Test public void testTimeSecVectorIterable() { try (final TimeSecVector timeSecVector = new TimeSecVector("timeSec", allocator)) { timeSecVector.allocateNew(3); timeSecVector.setSafe(0, 70000); timeSecVector.setNull(1); timeSecVector.setSafe(2, 555); timeSecVector.setValueCount(3); assertThat( timeSecVector.getValueIterable(), IsIterableContainingInOrder.contains(70000, null, 555)); } } @Test public void testTimeStampMicroTZVectorIterable() { try (final TimeStampMicroTZVector timeStampMicroTzVector = new TimeStampMicroTZVector("timeStampMicroTZ", allocator, "UTC")) { timeStampMicroTzVector.allocateNew(3); timeStampMicroTzVector.setSafe(0, 70000); timeStampMicroTzVector.setNull(1); timeStampMicroTzVector.setSafe(2, 555); timeStampMicroTzVector.setValueCount(3); assertThat( timeStampMicroTzVector.getValueIterable(), IsIterableContainingInOrder.contains(70000L, null, 555L)); } } @Test public void testTimeStampMicroVectorIterable() { try (final TimeStampMicroVector timeStampMicroVector = new TimeStampMicroVector("timeStampMicro", allocator)) { timeStampMicroVector.allocateNew(3); timeStampMicroVector.setSafe(0, 70000); timeStampMicroVector.setNull(1); timeStampMicroVector.setSafe(2, 555); timeStampMicroVector.setValueCount(3); final LocalDateTime value1 = LocalDateTime.ofEpochSecond(0L, 70000000, ZoneOffset.ofHours(0)); final LocalDateTime value3 = LocalDateTime.ofEpochSecond(0L, 555000, ZoneOffset.ofHours(0)); assertThat( timeStampMicroVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testTimeStampMilliTZVectorIterable() { try (final TimeStampMilliTZVector timeStampMilliTzVector = new TimeStampMilliTZVector("timeStampMilliTZ", allocator, "UTC")) { timeStampMilliTzVector.allocateNew(3); timeStampMilliTzVector.setSafe(0, 70000); timeStampMilliTzVector.setNull(1); timeStampMilliTzVector.setSafe(2, 555); timeStampMilliTzVector.setValueCount(3); assertThat( timeStampMilliTzVector.getValueIterable(), IsIterableContainingInOrder.contains(70000L, null, 555L)); } } @Test public void testTimeStampMilliVectorIterable() { try (final TimeStampMilliVector timeStampMilliVector = new TimeStampMilliVector("timeStampMilli", allocator)) { timeStampMilliVector.allocateNew(3); timeStampMilliVector.setSafe(0, 70000); timeStampMilliVector.setNull(1); timeStampMilliVector.setSafe(2, 555); timeStampMilliVector.setValueCount(3); final LocalDateTime value1 = LocalDateTime.ofEpochSecond(70L, 0, ZoneOffset.ofHours(0)); final LocalDateTime value3 = LocalDateTime.ofEpochSecond(0L, 555000000, ZoneOffset.ofHours(0)); assertThat( timeStampMilliVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testTimeStampNanoTZVectorIterable() { try (final TimeStampNanoTZVector timeStampNanoTzVector = new TimeStampNanoTZVector("timeStampNanoTZ", allocator, "UTC")) { timeStampNanoTzVector.allocateNew(3); timeStampNanoTzVector.setSafe(0, 70000); timeStampNanoTzVector.setNull(1); timeStampNanoTzVector.setSafe(2, 555); timeStampNanoTzVector.setValueCount(3); assertThat( timeStampNanoTzVector.getValueIterable(), IsIterableContainingInOrder.contains(70000L, null, 555L)); } } @Test public void testTimeStampNanoVectorIterable() { try (final TimeStampNanoVector timeStampNanoVector = new TimeStampNanoVector("timeStampNano", allocator)) { timeStampNanoVector.allocateNew(3); timeStampNanoVector.setSafe(0, 70000); timeStampNanoVector.setNull(1); timeStampNanoVector.setSafe(2, 555); timeStampNanoVector.setValueCount(3); final LocalDateTime value1 = LocalDateTime.ofEpochSecond(0L, 70000, ZoneOffset.ofHours(0)); final LocalDateTime value3 = LocalDateTime.ofEpochSecond(0L, 555, ZoneOffset.ofHours(0)); assertThat( timeStampNanoVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testTimeStampSecTZVectorIterable() { try (final TimeStampSecTZVector timeStampSecTzVector = new TimeStampSecTZVector("timeStampSecTZ", allocator, "UTC")) { timeStampSecTzVector.allocateNew(3); timeStampSecTzVector.setSafe(0, 70000); timeStampSecTzVector.setNull(1); timeStampSecTzVector.setSafe(2, 555); timeStampSecTzVector.setValueCount(3); assertThat( timeStampSecTzVector.getValueIterable(), IsIterableContainingInOrder.contains(70000L, null, 555L)); } } @Test public void testTimeStampSecVectorIterable() { try (final TimeStampSecVector timeStampSecVector = new TimeStampSecVector("timeStampSec", allocator)) { timeStampSecVector.allocateNew(3); timeStampSecVector.setSafe(0, 70000); timeStampSecVector.setNull(1); timeStampSecVector.setSafe(2, 555); timeStampSecVector.setValueCount(3); final LocalDateTime value1 = LocalDateTime.ofEpochSecond(70000L, 0, ZoneOffset.ofHours(0)); final LocalDateTime value3 = LocalDateTime.ofEpochSecond(555L, 0, ZoneOffset.ofHours(0)); assertThat( timeStampSecVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testTinyIntVectorIterable() { try (final TinyIntVector tinyIntVector = new TinyIntVector("tinyInt", allocator)) { tinyIntVector.allocateNew(3); tinyIntVector.setSafe(0, 8); tinyIntVector.setNull(1); tinyIntVector.setSafe(2, -17); tinyIntVector.setValueCount(3); assertThat( tinyIntVector.getValueIterable(), IsIterableContainingInOrder.contains((byte) 8, null, (byte) -17)); } } @Test public void testUInt1VectorIterable() { try (final UInt1Vector uint1Vector = new UInt1Vector("uint1", allocator)) { uint1Vector.allocateNew(3); uint1Vector.setSafe(0, 8); uint1Vector.setNull(1); uint1Vector.setSafe(2, 101); uint1Vector.setValueCount(3); assertThat( uint1Vector.getValueIterable(), IsIterableContainingInOrder.contains((byte) 8, null, (byte) 101)); } } @Test public void testUInt2VectorIterable() { try (final UInt2Vector uint2Vector = new UInt2Vector("uint2", allocator)) { uint2Vector.allocateNew(3); uint2Vector.setSafe(0, 78); uint2Vector.setNull(1); uint2Vector.setSafe(2, 3456); uint2Vector.setValueCount(3); assertThat( uint2Vector.getValueIterable(), IsIterableContainingInOrder.contains((char) 78, null, (char) 3456)); } } @Test public void testUInt4VectorIterable() { try (final UInt4Vector uint4Vector = new UInt4Vector("uint4", allocator)) { uint4Vector.allocateNew(3); uint4Vector.setSafe(0, 78); uint4Vector.setNull(1); uint4Vector.setSafe(2, 3456); uint4Vector.setValueCount(3); assertThat( uint4Vector.getValueIterable(), IsIterableContainingInOrder.contains(78, null, 3456)); } } @Test public void testUInt8VectorIterable() { try (final UInt8Vector uint8Vector = new UInt8Vector("uint8", allocator)) { uint8Vector.allocateNew(3); uint8Vector.setSafe(0, 6); uint8Vector.setSafe(1, 2); uint8Vector.setSafe(2, 19); uint8Vector.setValueCount(3); assertThat(uint8Vector.getValueIterable(), IsIterableContainingInOrder.contains(6L, 2L, 19L)); } } @Test public void testVarBinaryVectorIterable() { try (final VarBinaryVector varBinaryVector = new VarBinaryVector("varBinary", allocator)) { final byte[] value1 = new byte[] {0, 0, 0, 1}; final byte[] value3 = new byte[] {1, 0, 0}; varBinaryVector.allocateNew(3); varBinaryVector.setSafe(0, value1); varBinaryVector.setNull(1); varBinaryVector.setSafe(2, value3); varBinaryVector.setValueCount(3); assertThat( varBinaryVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testVarCharVectorIterable() { try (final VarCharVector varCharVector = new VarCharVector("varChar", allocator)) { final Text value1 = new Text("hello"); final Text value3 = new Text("worlds"); varCharVector.allocateNew(3); varCharVector.setSafe(0, value1); varCharVector.setNull(1); varCharVector.setSafe(2, value3); varCharVector.setValueCount(3); assertThat( varCharVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testViewVarBinaryVectorIterable() { try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("viewVarBinary", allocator)) { final byte[] value1 = new byte[] {0, 0, 0, 1}; final byte[] value3 = new byte[] {1, 0, 0}; viewVarBinaryVector.allocateNew(3); viewVarBinaryVector.setSafe(0, value1); viewVarBinaryVector.setNull(1); viewVarBinaryVector.setSafe(2, value3); viewVarBinaryVector.setValueCount(3); assertThat( viewVarBinaryVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testViewVarCharVectorIterable() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("viewVarChar", allocator)) { final Text value1 = new Text("hello"); final Text value3 = new Text("worlds"); viewVarCharVector.allocateNew(3); viewVarCharVector.setSafe(0, value1); viewVarCharVector.setNull(1); viewVarCharVector.setSafe(2, value3); viewVarCharVector.setValueCount(3); assertThat( viewVarCharVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } @Test public void testFIxedSizeListVectorIterable() { try (final FixedSizeListVector fixedSizeListVector = new FixedSizeListVector( "fixedSizeList", allocator, FieldType.nullable(new ArrowType.FixedSizeList(3)), null)) { final IntVector listVector = (IntVector) fixedSizeListVector .addOrGetVector(FieldType.nullable(Types.MinorType.INT.getType())) .getVector(); listVector.setSafe(0, 1); listVector.setSafe(1, 2); listVector.setSafe(2, 3); listVector.setSafe(3, 4); listVector.setSafe(4, 5); listVector.setSafe(5, 6); listVector.setValueCount(6); fixedSizeListVector.setValueCount(2); fixedSizeListVector.setNotNull(0); fixedSizeListVector.setNotNull(1); final List list1 = new ArrayList<>(); list1.add(1); list1.add(2); list1.add(3); final List list2 = new ArrayList<>(); list2.add(4); list2.add(5); list2.add(6); assertThat( fixedSizeListVector.getValueIterable(), IsIterableContainingInOrder.contains(list1, list2)); } } @Test public void testLargeListVectorIterable() { try (final LargeListVector largeListVector = LargeListVector.empty("largeList", allocator)) { UnionLargeListWriter writer = largeListVector.getWriter(); writer.allocate(); writer.startList(); writer.integer().writeInt(1); writer.integer().writeInt(2); writer.integer().writeInt(3); writer.endList(); writer.startList(); writer.integer().writeInt(4); writer.integer().writeInt(5); writer.endList(); largeListVector.setValueCount(2); final List list1 = new ArrayList<>(); list1.add(1); list1.add(2); list1.add(3); final List list2 = new ArrayList<>(); list2.add(4); list2.add(5); assertThat( largeListVector.getValueIterable(), IsIterableContainingInOrder.contains(list1, list2)); } } @Test public void testListVectorIterable() { try (final ListVector listVector = ListVector.empty("largeList", allocator)) { UnionListWriter writer = listVector.getWriter(); writer.allocate(); writer.startList(); writer.integer().writeInt(1); writer.integer().writeInt(2); writer.integer().writeInt(3); writer.endList(); writer.startList(); writer.integer().writeInt(4); writer.integer().writeInt(5); writer.endList(); listVector.setValueCount(2); final List list1 = new ArrayList<>(); list1.add(1); list1.add(2); list1.add(3); final List list2 = new ArrayList<>(); list2.add(4); list2.add(5); assertThat(listVector.getValueIterable(), IsIterableContainingInOrder.contains(list1, list2)); } } @Test public void testListViewVectorIterable() { try (final ListViewVector listViewVector = ListViewVector.empty("largeList", allocator)) { UnionListViewWriter writer = listViewVector.getWriter(); writer.allocate(); writer.startList(); writer.integer().writeInt(1); writer.integer().writeInt(2); writer.integer().writeInt(3); writer.endList(); writer.startList(); writer.integer().writeInt(4); writer.integer().writeInt(5); writer.endList(); listViewVector.setValueCount(2); final List list1 = new ArrayList<>(); list1.add(1); list1.add(2); list1.add(3); final List list2 = new ArrayList<>(); list2.add(4); list2.add(5); assertThat( listViewVector.getValueIterable(), IsIterableContainingInOrder.contains(list1, list2)); } } @Test public void testNonNullableStructVectorIterable() { try (final NonNullableStructVector nonNullableStructVector = NonNullableStructVector.empty("nonNullableStruct", allocator)) { nonNullableStructVector.setValueCount(2); IntVector key1Vector = nonNullableStructVector.addOrGet( "key1", FieldType.notNullable(new ArrowType.Int(32, true)), IntVector.class); IntVector key2Vector = nonNullableStructVector.addOrGet( "key2", FieldType.notNullable(new ArrowType.Int(32, true)), IntVector.class); key1Vector.setSafe(0, 1); key1Vector.setSafe(1, 3); key2Vector.setSafe(0, 2); key2Vector.setSafe(1, 4); key1Vector.setValueCount(2); key2Vector.setValueCount(2); final Map value1 = new HashMap<>(); value1.put("key1", 1); value1.put("key2", 2); final Map value2 = new HashMap<>(); value2.put("key1", 3); value2.put("key2", 4); assertThat( nonNullableStructVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, value2)); } } @Test public void testStructVectorIterable() { try (final StructVector structVector = StructVector.empty("struct", allocator)) { structVector.addOrGetList("struct"); NullableStructWriter structWriter = structVector.getWriter(); structWriter.setPosition(0); structWriter.start(); structWriter.integer("key1").writeInt(1); structWriter.integer("key2").writeInt(2); structWriter.end(); structWriter.setPosition(2); structWriter.start(); structWriter.integer("key1").writeInt(3); structWriter.integer("key2").writeInt(4); structWriter.end(); structWriter.setValueCount(3); final Map value1 = new HashMap<>(); value1.put("key1", 1); value1.put("key2", 2); final Map value3 = new HashMap<>(); value3.put("key1", 3); value3.put("key2", 4); assertThat( structVector.getValueIterable(), IsIterableContainingInOrder.contains(value1, null, value3)); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestVarCharListVector.java000066400000000000000000000051731515451615200326110ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.nio.charset.StandardCharsets; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.impl.UnionListWriter; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.FieldType; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestVarCharListVector { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testVarCharListWithNulls() { byte[] bytes = "a".getBytes(StandardCharsets.UTF_8); try (ListVector vector = new ListVector( "VarList", allocator, FieldType.nullable(Types.MinorType.VARCHAR.getType()), null); ArrowBuf tempBuf = allocator.buffer(bytes.length)) { UnionListWriter writer = vector.getWriter(); writer.allocate(); // populate input vector with the following records // ["a"] // null // ["b"] writer.setPosition(0); // optional writer.startList(); tempBuf.setBytes(0, bytes); writer.writeVarChar(0, bytes.length, tempBuf); writer.endList(); writer.setPosition(2); writer.startList(); bytes = "b".getBytes(StandardCharsets.UTF_8); tempBuf.setBytes(0, bytes); writer.writeVarChar(0, bytes.length, tempBuf); writer.endList(); writer.setValueCount(2); assertEquals(2, vector.getValueCount()); assertEquals(2, vector.getDataVector().getValueCount()); } } } TestVariableWidthViewVector.java000066400000000000000000003226731515451615200337370ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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; import static org.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector; 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.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Random; import java.util.function.BiConsumer; import java.util.function.Function; 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.rounding.DefaultRoundingPolicy; import org.apache.arrow.memory.util.ArrowBufPointer; import org.apache.arrow.memory.util.CommonUtil; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.holders.NullableViewVarBinaryHolder; import org.apache.arrow.vector.holders.NullableViewVarCharHolder; import org.apache.arrow.vector.holders.ValueHolder; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; 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; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; public class TestVariableWidthViewVector { // short string (length <= 12) private static final byte[] STR0 = "0123456".getBytes(StandardCharsets.UTF_8); // short string (length <= 12) private static final byte[] STR1 = "012345678912".getBytes(StandardCharsets.UTF_8); // long string (length > 12) private static final byte[] STR2 = "0123456789123".getBytes(StandardCharsets.UTF_8); // long string (length > 12) private static final byte[] STR3 = "01234567891234567".getBytes(StandardCharsets.UTF_8); // short string (length <= 12) private static final byte[] STR4 = "01234567".getBytes(StandardCharsets.UTF_8); // short string (length <= 12) private static final byte[] STR5 = "A1234A".getBytes(StandardCharsets.UTF_8); // short string (length <= 12) private static final byte[] STR6 = "B1234567B".getBytes(StandardCharsets.UTF_8); // long string (length > 12) private static final byte[] STR7 = "K01234567891234567K".getBytes(StandardCharsets.UTF_8); // long string (length > 12) private static final byte[] STR8 = "M012345678912345678M".getBytes(StandardCharsets.UTF_8); private static final String EMPTY_SCHEMA_PATH = ""; private BufferAllocator allocator; private Random random; @BeforeEach public void prepare() { allocator = new RootAllocator(Integer.MAX_VALUE); random = new Random(); } @AfterEach public void shutdown() { allocator.close(); } public static void setBytes(int index, byte[] bytes, ViewVarCharVector vector) { BitVectorHelper.setBit(vector.validityBuffer, index); vector.setBytes(index, bytes, 0, bytes.length); } @Test public void testInlineAllocation() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { viewVarCharVector.allocateNew(48, 3); final int valueCount = 3; viewVarCharVector.set(0, STR0); viewVarCharVector.set(1, STR1); viewVarCharVector.set(2, STR4); viewVarCharVector.setValueCount(valueCount); byte[] view1 = viewVarCharVector.get(0); byte[] view2 = viewVarCharVector.get(1); byte[] view3 = viewVarCharVector.get(2); assertNotNull(view1); assertNotNull(view2); assertNotNull(view3); String str1 = new String(STR0, StandardCharsets.UTF_8); String str2 = new String(STR1, StandardCharsets.UTF_8); String str3 = new String(STR4, StandardCharsets.UTF_8); assertEquals(new String(view1, StandardCharsets.UTF_8), str1); assertEquals(new String(view2, StandardCharsets.UTF_8), str2); assertEquals(new String(view3, StandardCharsets.UTF_8), str3); assertTrue(viewVarCharVector.dataBuffers.isEmpty()); assertEquals( new String( Objects.requireNonNull(viewVarCharVector.getObject(0)).getBuffer(), StandardCharsets.UTF_8), str1); assertEquals( new String( Objects.requireNonNull(viewVarCharVector.getObject(1)).getBuffer(), StandardCharsets.UTF_8), str2); assertEquals( new String( Objects.requireNonNull(viewVarCharVector.getObject(2)).getBuffer(), StandardCharsets.UTF_8), str3); } } @Test public void testDataBufferBasedAllocationInSameBuffer() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { viewVarCharVector.allocateNew(48, 4); final int valueCount = 4; String str4 = TestUtils.generateRandomString(34); viewVarCharVector.set(0, STR1); viewVarCharVector.set(1, STR2); viewVarCharVector.set(2, STR3); viewVarCharVector.set(3, str4.getBytes(StandardCharsets.UTF_8)); viewVarCharVector.setValueCount(valueCount); byte[] view1 = viewVarCharVector.get(0); byte[] view2 = viewVarCharVector.get(1); byte[] view3 = viewVarCharVector.get(2); byte[] view4 = viewVarCharVector.get(3); assertNotNull(view1); assertNotNull(view2); assertNotNull(view3); assertNotNull(view4); String str1 = new String(STR1, StandardCharsets.UTF_8); String str2 = new String(STR2, StandardCharsets.UTF_8); String str3 = new String(STR3, StandardCharsets.UTF_8); assertEquals(new String(view1, StandardCharsets.UTF_8), str1); assertEquals(new String(view2, StandardCharsets.UTF_8), str2); assertEquals(new String(view3, StandardCharsets.UTF_8), str3); assertEquals(new String(view4, StandardCharsets.UTF_8), str4); assertEquals(1, viewVarCharVector.dataBuffers.size()); assertEquals( new String( Objects.requireNonNull(viewVarCharVector.getObject(0)).getBuffer(), StandardCharsets.UTF_8), str1); assertEquals( new String( Objects.requireNonNull(viewVarCharVector.getObject(1)).getBuffer(), StandardCharsets.UTF_8), str2); assertEquals( new String( Objects.requireNonNull(viewVarCharVector.getObject(2)).getBuffer(), StandardCharsets.UTF_8), str3); assertEquals( new String( Objects.requireNonNull(viewVarCharVector.getObject(3)).getBuffer(), StandardCharsets.UTF_8), str4); } } @Test public void testDataBufferBasedAllocationInOtherBuffer() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { viewVarCharVector.allocateNew(48, 4); final int valueCount = 4; String str4 = TestUtils.generateRandomString(35); viewVarCharVector.set(0, STR1); viewVarCharVector.set(1, STR2); viewVarCharVector.set(2, STR3); viewVarCharVector.set(3, str4.getBytes(StandardCharsets.UTF_8)); viewVarCharVector.setValueCount(valueCount); byte[] view1 = viewVarCharVector.get(0); byte[] view2 = viewVarCharVector.get(1); byte[] view3 = viewVarCharVector.get(2); byte[] view4 = viewVarCharVector.get(3); assertNotNull(view1); assertNotNull(view2); assertNotNull(view3); assertNotNull(view4); String str1 = new String(STR1, StandardCharsets.UTF_8); String str2 = new String(STR2, StandardCharsets.UTF_8); String str3 = new String(STR3, StandardCharsets.UTF_8); assertEquals(new String(view1, StandardCharsets.UTF_8), str1); assertEquals(new String(view2, StandardCharsets.UTF_8), str2); assertEquals(new String(view3, StandardCharsets.UTF_8), str3); assertEquals(new String(view4, StandardCharsets.UTF_8), str4); assertEquals(2, viewVarCharVector.dataBuffers.size()); assertEquals( new String( Objects.requireNonNull(viewVarCharVector.getObject(0)).getBuffer(), StandardCharsets.UTF_8), str1); assertEquals( new String( Objects.requireNonNull(viewVarCharVector.getObject(1)).getBuffer(), StandardCharsets.UTF_8), str2); assertEquals( new String( Objects.requireNonNull(viewVarCharVector.getObject(2)).getBuffer(), StandardCharsets.UTF_8), str3); assertEquals( new String( Objects.requireNonNull(viewVarCharVector.getObject(3)).getBuffer(), StandardCharsets.UTF_8), str4); } } @Test public void testSetSafe() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { viewVarCharVector.allocateNew(1, 1); 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 Map expected = new HashMap<>(); for (byte[] string : strings) { int cap = viewVarCharVector.getValueCapacity(); expected.put(cap, string); viewVarCharVector.setSafe(cap, string); } int nullIndex = viewVarCharVector.getValueCapacity(); viewVarCharVector.setNull(nullIndex); int valueCount = nullIndex + 1; viewVarCharVector.setValueCount(valueCount); assertEquals(viewVarCharVector.getNullCount(), valueCount - strings.size()); assertEquals(128, viewVarCharVector.getValueCapacity()); assertEquals(2, viewVarCharVector.dataBuffers.size()); for (int i = 0; i < viewVarCharVector.getValueCapacity(); i++) { if (expected.containsKey(i)) { assertArrayEquals(expected.get(i), viewVarCharVector.get(i)); } else { assertNull(viewVarCharVector.get(i)); } } } } @Test public void testMixedAllocation() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { viewVarCharVector.allocateNew(128, 6); final int valueCount = 6; String str4 = TestUtils.generateRandomString(35); String str6 = TestUtils.generateRandomString(40); viewVarCharVector.set(0, STR1); viewVarCharVector.set(1, STR2); viewVarCharVector.set(2, STR3); viewVarCharVector.set(3, str4.getBytes(StandardCharsets.UTF_8)); viewVarCharVector.set(4, STR1); viewVarCharVector.set(5, str6.getBytes(StandardCharsets.UTF_8)); viewVarCharVector.setValueCount(valueCount); byte[] view1 = viewVarCharVector.get(0); byte[] view2 = viewVarCharVector.get(1); byte[] view3 = viewVarCharVector.get(2); byte[] view4 = viewVarCharVector.get(3); byte[] view5 = viewVarCharVector.get(4); byte[] view6 = viewVarCharVector.get(5); assertNotNull(view1); assertNotNull(view2); assertNotNull(view3); assertNotNull(view4); assertNotNull(view5); assertNotNull(view6); String str1 = new String(STR1, StandardCharsets.UTF_8); String str2 = new String(STR2, StandardCharsets.UTF_8); String str3 = new String(STR3, StandardCharsets.UTF_8); assertEquals(new String(view1, StandardCharsets.UTF_8), str1); assertEquals(new String(view2, StandardCharsets.UTF_8), str2); assertEquals(new String(view3, StandardCharsets.UTF_8), str3); assertEquals(new String(view4, StandardCharsets.UTF_8), str4); assertEquals(new String(view5, StandardCharsets.UTF_8), str1); assertEquals(new String(view6, StandardCharsets.UTF_8), str6); assertEquals(1, viewVarCharVector.dataBuffers.size()); assertEquals( new String( Objects.requireNonNull(viewVarCharVector.getObject(0)).getBuffer(), StandardCharsets.UTF_8), str1); assertEquals( new String( Objects.requireNonNull(viewVarCharVector.getObject(1)).getBuffer(), StandardCharsets.UTF_8), str2); assertEquals( new String( Objects.requireNonNull(viewVarCharVector.getObject(2)).getBuffer(), StandardCharsets.UTF_8), str3); assertEquals( new String( Objects.requireNonNull(viewVarCharVector.getObject(3)).getBuffer(), StandardCharsets.UTF_8), str4); assertEquals( new String( Objects.requireNonNull(viewVarCharVector.getObject(4)).getBuffer(), StandardCharsets.UTF_8), str1); assertEquals( new String( Objects.requireNonNull(viewVarCharVector.getObject(5)).getBuffer(), StandardCharsets.UTF_8), str6); } } @Test public void testSetNullableViewVarCharHolder() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { viewVarCharVector.allocateNew(0, 0); final List strings = List.of(STR0, STR1, STR2, STR3, STR4, STR5); NullableViewVarCharHolder stringHolder = new NullableViewVarCharHolder(); // set not null int size = strings.size(); for (int i = 0; i < size; i++) { setAndCheck(viewVarCharVector, i, strings.get(i), stringHolder); } // set null setAndCheck(viewVarCharVector, 6, null, stringHolder); // copy by holder // len < 12 copyAndCheck(viewVarCharVector, stringHolder, 0, 7); // len > 12 copyAndCheck(viewVarCharVector, stringHolder, 2, 8); // null copyAndCheck(viewVarCharVector, stringHolder, 6, 9); // test overwrite for (int i = 0; i < size; i++) { setAndCheck(viewVarCharVector, i, strings.get(size - i - 1), stringHolder); } String longString = TestUtils.generateRandomString(128); setAndCheck(viewVarCharVector, 6, longString.getBytes(), stringHolder); } } @Test public void testSetNullableViewVarBinaryHolder() { try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator)) { viewVarBinaryVector.allocateNew(0, 0); final List strings = List.of(STR0, STR1, STR2, STR3, STR4, STR5); NullableViewVarBinaryHolder holder = new NullableViewVarBinaryHolder(); // set not null int size = strings.size(); for (int i = 0; i < size; i++) { setAndCheck(viewVarBinaryVector, i, strings.get(i), holder); } // set null setAndCheck(viewVarBinaryVector, 6, null, holder); // copy by holder // len < 12 copyAndCheck(viewVarBinaryVector, holder, 0, 7); // len > 12 copyAndCheck(viewVarBinaryVector, holder, 2, 8); // null copyAndCheck(viewVarBinaryVector, holder, 6, 9); // test overwrite for (int i = 0; i < size; i++) { setAndCheck(viewVarBinaryVector, i, strings.get(size - i - 1), holder); } String longString = TestUtils.generateRandomString(128); setAndCheck(viewVarBinaryVector, 6, longString.getBytes(), holder); } } private static void copyAndCheck( BaseVariableWidthViewVector vector, ValueHolder holder, int fromIndex, int toIndex) { if (vector instanceof ViewVarCharVector) { ViewVarCharVector viewVarCharVector = (ViewVarCharVector) vector; NullableViewVarCharHolder stringHolder = (NullableViewVarCharHolder) holder; viewVarCharVector.get(fromIndex, stringHolder); viewVarCharVector.setSafe(toIndex, stringHolder); } if (vector instanceof ViewVarBinaryVector) { ViewVarBinaryVector viewVarBinaryVector = (ViewVarBinaryVector) vector; NullableViewVarBinaryHolder binaryHolder = (NullableViewVarBinaryHolder) holder; viewVarBinaryVector.get(fromIndex, binaryHolder); viewVarBinaryVector.setSafe(toIndex, binaryHolder); } assertArrayEquals(vector.get(fromIndex), vector.get(toIndex)); } private void setAndCheck( ViewVarCharVector vector, int index, byte[] str, NullableViewVarCharHolder stringHolder) { ArrowBuf buf = null; if (null == str) { stringHolder.isSet = 0; } else { buf = allocator.buffer(str.length); buf.setBytes(0, str); stringHolder.isSet = 1; stringHolder.start = 0; stringHolder.end = str.length; stringHolder.buffer = buf; } vector.setSafe(index, stringHolder); // verify results assertArrayEquals(str, vector.get(index)); AutoCloseables.closeNoChecked(buf); } private void setAndCheck( ViewVarBinaryVector vector, int index, byte[] str, NullableViewVarBinaryHolder binaryHolder) { ArrowBuf buf = null; if (null == str) { binaryHolder.isSet = 0; } else { buf = allocator.buffer(str.length); buf.setBytes(0, str); binaryHolder.isSet = 1; binaryHolder.start = 0; binaryHolder.end = str.length; binaryHolder.buffer = buf; } vector.setSafe(index, binaryHolder); // verify results assertArrayEquals(str, vector.get(index)); AutoCloseables.closeNoChecked(buf); } @Test public void testAllocationIndexOutOfBounds() { assertThrows( IndexOutOfBoundsException.class, () -> { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { viewVarCharVector.allocateNew(32, 3); final int valueCount = 3; viewVarCharVector.set(0, STR1); viewVarCharVector.set(1, STR2); viewVarCharVector.set(2, STR2); viewVarCharVector.setValueCount(valueCount); } }); } @Test public void testSizeOfViewBufferElements() { try (final ViewVarCharVector vector = new ViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { int valueCount = 100; int currentSize = 0; vector.setInitialCapacity(valueCount); vector.allocateNew(); vector.setValueCount(valueCount); for (int i = 0; i < valueCount; i++) { currentSize += i; vector.setSafe(i, new byte[i]); } assertEquals(currentSize, vector.sizeOfViewBufferElements()); } } @Test public void testNullableVarType1() { // Create a new value vector for 1024 integers. try (final ViewVarCharVector vector = newViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { vector.allocateNew(1024 * 10, 1024); vector.set(0, STR1); vector.set(1, STR2); vector.set(2, STR3); vector.setSafe(3, STR3, 1, STR3.length - 1); vector.setSafe(4, STR3, 2, STR3.length - 2); ByteBuffer str3ByteBuffer = ByteBuffer.wrap(STR3); vector.setSafe(5, str3ByteBuffer, 1, STR3.length - 1); vector.setSafe(6, str3ByteBuffer, 2, STR3.length - 2); // Set with convenience function Text txt = new Text("foo"); vector.setSafe(7, txt.getBytes(), 0, (int) txt.getLength()); // Check the sample strings. assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(Arrays.copyOfRange(STR3, 1, STR3.length), vector.get(3)); assertArrayEquals(Arrays.copyOfRange(STR3, 2, STR3.length), vector.get(4)); assertArrayEquals(Arrays.copyOfRange(STR3, 1, STR3.length), vector.get(5)); assertArrayEquals(Arrays.copyOfRange(STR3, 2, STR3.length), vector.get(6)); // Check returning a Text object assertEquals(txt, vector.getObject(7)); // Ensure null value throws. assertNull(vector.get(8)); } } @Test public void testGetTextRepeatedly() { try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { ValueVectorDataPopulator.setVector(vector, STR1, STR2); vector.setValueCount(2); /* check the vector output */ Text text = new Text(); vector.read(0, text); assertArrayEquals(STR1, text.getBytes()); vector.read(1, text); assertArrayEquals(STR2, text.getBytes()); } } @Test public void testNullableVarType2() { try (final ViewVarBinaryVector vector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { vector.allocateNew(1024 * 10, 1024); vector.set(0, STR1); vector.set(1, STR2); vector.set(2, STR3); vector.setSafe(3, STR3, 1, STR3.length - 1); vector.setSafe(4, STR3, 2, STR3.length - 2); ByteBuffer str3ByteBuffer = ByteBuffer.wrap(STR3); vector.setSafe(5, str3ByteBuffer, 1, STR3.length - 1); vector.setSafe(6, str3ByteBuffer, 2, STR3.length - 2); // Check the sample strings. assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(Arrays.copyOfRange(STR3, 1, STR3.length), vector.get(3)); assertArrayEquals(Arrays.copyOfRange(STR3, 2, STR3.length), vector.get(4)); assertArrayEquals(Arrays.copyOfRange(STR3, 1, STR3.length), vector.get(5)); assertArrayEquals(Arrays.copyOfRange(STR3, 2, STR3.length), vector.get(6)); // Ensure null value throws. assertNull(vector.get(7)); } } @Test public void testGetBytesRepeatedly() { try (ViewVarBinaryVector vector = new ViewVarBinaryVector("", allocator)) { vector.allocateNew(5, 1); final String str = "hello world!!!"; final String str2 = "foo"; vector.setSafe(0, str.getBytes(StandardCharsets.UTF_8)); vector.setSafe(1, str2.getBytes(StandardCharsets.UTF_8)); // verify results ReusableByteArray reusableByteArray = new ReusableByteArray(); vector.read(0, reusableByteArray); assertArrayEquals( str.getBytes(StandardCharsets.UTF_8), Arrays.copyOfRange( reusableByteArray.getBuffer(), 0, (int) reusableByteArray.getLength())); byte[] oldBuffer = reusableByteArray.getBuffer(); vector.read(1, reusableByteArray); assertArrayEquals( str2.getBytes(StandardCharsets.UTF_8), Arrays.copyOfRange( reusableByteArray.getBuffer(), 0, (int) reusableByteArray.getLength())); // There should not have been any reallocation since the newer value is smaller in length. assertSame(oldBuffer, reusableByteArray.getBuffer()); } } @Test public void testReAllocVariableWidthViewVector() { try (final ViewVarCharVector vector = newVector( ViewVarCharVector.class, EMPTY_SCHEMA_PATH, Types.MinorType.VIEWVARCHAR, allocator)) { final int capacityLimit = 4095; final int overLimitIndex = 200; vector.setInitialCapacity(capacityLimit); vector.allocateNew(); int initialCapacity = vector.getValueCapacity(); assertTrue(initialCapacity >= capacityLimit); /* Put values in indexes that fall within the initial allocation */ vector.setSafe(0, STR1, 0, STR1.length); vector.setSafe(initialCapacity - 1, STR2, 0, STR2.length); /* the set calls above should NOT have triggered a realloc */ assertEquals(initialCapacity, vector.getValueCapacity()); /* Now try to put values in space that falls beyond the initial allocation */ vector.setSafe(initialCapacity + overLimitIndex, STR3, 0, STR3.length); /* Check valueCapacity is more than initial allocation */ assertTrue(initialCapacity * 2 <= vector.getValueCapacity()); assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(initialCapacity - 1)); assertArrayEquals(STR3, vector.get(initialCapacity + overLimitIndex)); // Set the valueCount to be more than valueCapacity of current allocation. This is possible // for ValueVectors // as we don't call setSafe for null values, but we do call setValueCount when the current // batch is processed. vector.setValueCount(vector.getValueCapacity() + overLimitIndex); } } @Test public void testSetSafeWithArrowBufNoExcessAllocs() { final int numValues = BaseVariableWidthViewVector.INITIAL_VALUE_ALLOCATION * 2; final byte[] valueBytes = "hello world!!!".getBytes(StandardCharsets.UTF_8); final int valueBytesLength = valueBytes.length; final int isSet = 1; try (final ViewVarCharVector fromVector = newVector( ViewVarCharVector.class, EMPTY_SCHEMA_PATH, Types.MinorType.VIEWVARCHAR, allocator); final ViewVarCharVector toVector = newVector( ViewVarCharVector.class, EMPTY_SCHEMA_PATH, Types.MinorType.VIEWVARCHAR, allocator)) { /* * Populate the `fromVector` with `numValues` with byte-arrays, each of size `valueBytesLength`. */ fromVector.setInitialCapacity(numValues); fromVector.allocateNew(); for (int i = 0; i < numValues; ++i) { fromVector.setSafe(i, valueBytes, 0 /*start*/, valueBytesLength); } fromVector.setValueCount(numValues); ArrowBuf fromDataBuffer = fromVector.getDataBuffer(); assertTrue(numValues * valueBytesLength <= fromDataBuffer.capacity()); /* * Copy the entries one-by-one from 'fromVector' to 'toVector', but use the setSafe with * ArrowBuf API (instead of setSafe with byte-array). */ toVector.setInitialCapacity(numValues); toVector.allocateNew(); for (int i = 0; i < numValues; i++) { int start = fromVector.getTotalValueLengthUpToIndex(i); // across variable // width implementations int end = fromVector.getTotalValueLengthUpToIndex(i + 1); toVector.setSafe(i, isSet, start, end, fromDataBuffer); } /* * Since the 'fromVector' and 'toVector' have the same initial capacity, and were populated * with the same varchar elements, the allocations and hence, the final capacity should be * the same. */ assertEquals(fromDataBuffer.capacity(), toVector.getDataBuffer().capacity()); } } @Test public void testSetLastSetUsage() { try (final ViewVarCharVector vector = new ViewVarCharVector("myvector", allocator)) { vector.allocateNew(1024 * 10, 1024); setBytes(0, STR1, vector); setBytes(1, STR2, vector); setBytes(2, STR3, vector); setBytes(3, STR4, vector); /* Check current lastSet */ assertEquals(-1, vector.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); /* * If we don't do setLastSe(3) before setValueCount(), then the latter will corrupt * the value vector by filling in all positions [0,valuecount-1] will empty byte arrays. * Run the test by commenting on the next line, and we should see incorrect vector output. */ vector.setLastSet(3); vector.setValueCount(20); /* Check current lastSet */ assertEquals(19, vector.getLastSet()); /* Check the vector output again */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); assertEquals(0, vector.getValueLength(4)); assertEquals(0, vector.getValueLength(5)); assertEquals(0, vector.getValueLength(6)); assertEquals(0, vector.getValueLength(7)); assertEquals(0, vector.getValueLength(8)); assertEquals(0, vector.getValueLength(9)); assertEquals(0, vector.getValueLength(10)); assertEquals(0, vector.getValueLength(11)); assertEquals(0, vector.getValueLength(12)); assertEquals(0, vector.getValueLength(13)); assertEquals(0, vector.getValueLength(14)); assertEquals(0, vector.getValueLength(15)); assertEquals(0, vector.getValueLength(16)); assertEquals(0, vector.getValueLength(17)); assertEquals(0, vector.getValueLength(18)); assertEquals(0, vector.getValueLength(19)); } } @Test public void testFillEmptiesUsage() { try (final ViewVarCharVector vector = new ViewVarCharVector("myvector", allocator)) { vector.allocateNew(1024 * 10, 1024); setBytes(0, STR1, vector); setBytes(1, STR2, vector); setBytes(2, STR3, vector); setBytes(3, STR4, vector); /* Check current lastSet */ assertEquals(-1, vector.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); vector.setLastSet(3); /* fill empty byte arrays from index [4, 9] */ vector.fillEmpties(10); /* Check current lastSet */ assertEquals(9, vector.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); assertEquals(0, vector.getValueLength(4)); assertEquals(0, vector.getValueLength(5)); assertEquals(0, vector.getValueLength(6)); assertEquals(0, vector.getValueLength(7)); assertEquals(0, vector.getValueLength(8)); assertEquals(0, vector.getValueLength(9)); setBytes(10, STR1, vector); setBytes(11, STR2, vector); vector.setLastSet(11); /* fill empty byte arrays from index [12, 14] */ vector.setValueCount(15); /* Check current lastSet */ assertEquals(14, vector.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); assertEquals(0, vector.getValueLength(4)); assertEquals(0, vector.getValueLength(5)); assertEquals(0, vector.getValueLength(6)); assertEquals(0, vector.getValueLength(7)); assertEquals(0, vector.getValueLength(8)); assertEquals(0, vector.getValueLength(9)); assertArrayEquals(STR1, vector.get(10)); assertArrayEquals(STR2, vector.get(11)); assertEquals(0, vector.getValueLength(12)); assertEquals(0, vector.getValueLength(13)); assertEquals(0, vector.getValueLength(14)); } } @Test public void testGetBufferAddress1() { try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { setVector(vector, STR1, STR2, STR3, STR4); vector.setValueCount(15); /* check the vector output */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR3, vector.get(2)); assertArrayEquals(STR4, vector.get(3)); List buffers = vector.getFieldBuffers(); long bitAddress = vector.getValidityBufferAddress(); long dataAddress = vector.getDataBufferAddress(); assertEquals(3, buffers.size()); assertEquals(bitAddress, buffers.get(0).memoryAddress()); assertEquals(dataAddress, buffers.get(1).memoryAddress()); } } @Test public void testSetInitialCapacityInViews() { try (final ViewVarCharVector vector = new ViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { /* use the default 16 data bytes on average per element */ final int viewSize = BaseVariableWidthViewVector.ELEMENT_SIZE; int defaultCapacity = BaseVariableWidthViewVector.INITIAL_VIEW_VALUE_ALLOCATION / viewSize; vector.setInitialCapacity(defaultCapacity); vector.allocateNew(); assertEquals(defaultCapacity, vector.getValueCapacity()); assertEquals( CommonUtil.nextPowerOfTwo(defaultCapacity * viewSize), vector.getDataBuffer().capacity()); double density = 4.0; final int valueCount = 5; vector.setInitialCapacity(valueCount, density); vector.allocateNew(); assertEquals(8, vector.getValueCapacity()); assertEquals(128, vector.getDataBuffer().capacity()); int initialDataBufferSize = (int) (valueCount * density); // making sure a databuffer is allocated vector.set(4, "01234567890123456".getBytes(StandardCharsets.UTF_8)); assertEquals(vector.dataBuffers.size(), 1); ArrowBuf dataBuf = vector.dataBuffers.get(0); try (ArrowBuf tempBuf = vector.allocator.buffer(initialDataBufferSize)) { // replicating a new buffer allocation process when a new buffer is added to the // data buffer when inserting an element with length > 12 assertEquals(tempBuf.capacity(), dataBuf.capacity()); } } } @Test public void testGetPointerVariableWidthViews() { final String[] sampleData = new String[] { "abc", "1234567890123", "def", null, "hello world java", "aaaaa", "world", "2019", null, "0717" }; try (ViewVarCharVector vec1 = new ViewVarCharVector("vec1", allocator); ViewVarCharVector vec2 = new ViewVarCharVector("vec2", allocator)) { vec1.allocateNew((long) sampleData.length * 16, sampleData.length); vec2.allocateNew((long) sampleData.length * 16, sampleData.length); for (int i = 0; i < sampleData.length; i++) { String str = sampleData[i]; if (str != null) { vec1.set(i, sampleData[i].getBytes(StandardCharsets.UTF_8)); vec2.set(i, sampleData[i].getBytes(StandardCharsets.UTF_8)); } else { vec1.setNull(i); vec2.setNull(i); } } ArrowBufPointer ptr1 = new ArrowBufPointer(); ArrowBufPointer ptr2 = new ArrowBufPointer(); for (int i = 0; i < sampleData.length; i++) { vec1.getDataPointer(i, ptr1); vec2.getDataPointer(i, ptr2); assertTrue(ptr1.equals(ptr2)); assertTrue(ptr2.equals(ptr2)); } } } @Test public void testGetNullFromVariableWidthViewVector() { try (final ViewVarCharVector varCharViewVector = new ViewVarCharVector("viewvarcharvec", allocator); final ViewVarBinaryVector varBinaryViewVector = new ViewVarBinaryVector("viewvarbinary", allocator)) { varCharViewVector.allocateNew(16, 1); varBinaryViewVector.allocateNew(16, 1); varCharViewVector.setNull(0); varBinaryViewVector.setNull(0); assertNull(varCharViewVector.get(0)); assertNull(varBinaryViewVector.get(0)); } } @Test public void testVariableWidthViewVectorNullHashCode() { try (ViewVarCharVector viewVarChar = new ViewVarCharVector("view var char vector", allocator)) { viewVarChar.allocateNew(100, 1); viewVarChar.setValueCount(1); viewVarChar.set(0, "abc".getBytes(StandardCharsets.UTF_8)); viewVarChar.setNull(0); assertEquals(0, viewVarChar.hashCode(0)); } } @Test public void testUnloadVariableWidthViewVector() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("view var char", allocator)) { viewVarCharVector.allocateNew(16, 2); viewVarCharVector.setValueCount(2); viewVarCharVector.set(0, "abcd".getBytes(StandardCharsets.UTF_8)); List bufs = viewVarCharVector.getFieldBuffers(); assertEquals(2, bufs.size()); ArrowBuf viewBuf = bufs.get(1); assertEquals(32, viewBuf.writerIndex()); final String longString = "012345678901234"; viewVarCharVector.set(1, longString.getBytes(StandardCharsets.UTF_8)); bufs = viewVarCharVector.getFieldBuffers(); assertEquals(3, bufs.size()); ArrowBuf referenceBuf = bufs.get(2); assertEquals(longString.length(), referenceBuf.writerIndex()); } } @Test public void testUnSupportedOffSet() { // offset is not a feature required in ViewVarCharVector assertThrows( UnsupportedOperationException.class, () -> { try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { setVector(vector, STR1, STR2); vector.setValueCount(2); /* check the vector output */ assertArrayEquals(STR1, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); vector.getOffsetBuffer(); } }); } private void validateViewBuffer( int index, ViewVarCharVector vector, byte[] expectedData, int expectedBufId, int expectedOffSet) { final ArrowBuf viewBuffer = vector.viewBuffer; int writePosition = index * BaseVariableWidthViewVector.ELEMENT_SIZE; final int prefixBufWidth = BaseVariableWidthViewVector.PREFIX_WIDTH; final int lengthBufWidth = BaseVariableWidthViewVector.LENGTH_WIDTH; int length = viewBuffer.getInt(writePosition); // validate length of the view assertEquals(expectedData.length, length); byte[] prefixBytes = new byte[prefixBufWidth]; viewBuffer.getBytes(writePosition + lengthBufWidth, prefixBytes); // validate the prefix byte[] expectedPrefixBytes = new byte[prefixBufWidth]; System.arraycopy(expectedData, 0, expectedPrefixBytes, 0, prefixBufWidth); assertArrayEquals(expectedPrefixBytes, prefixBytes); if (length > 12) { /// validate bufId int bufId = viewBuffer.getInt(writePosition + lengthBufWidth + prefixBufWidth); assertEquals(expectedBufId, bufId); // validate offset int offset = viewBuffer.getInt( writePosition + lengthBufWidth + prefixBufWidth + BaseVariableWidthViewVector.BUF_INDEX_WIDTH); assertEquals(expectedOffSet, offset); } // validate retrieved data assertArrayEquals(expectedData, vector.get(index)); } @Test public void testOverwriteShortFromLongString() { /*NA: not applicable */ // Overwriting at the beginning of the buffer. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(16, 1); // set short string vector.set(0, STR0); vector.setValueCount(1); assertEquals(0, vector.dataBuffers.size()); assertArrayEquals(STR0, vector.get(0)); validateViewBuffer(0, vector, STR0, /*NA*/ -1, /*NA*/ -1); // set long string vector.set(0, STR3); vector.setValueCount(1); assertEquals(1, vector.dataBuffers.size()); assertArrayEquals(STR3, vector.get(0)); validateViewBuffer(0, vector, STR3, 0, 0); } // Overwriting in the middle of the buffer when existing buffers are all shorts. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(48, 3); // set short string 1 vector.set(0, STR0); // set short string 2 vector.set(1, STR5); // set short string 3 vector.set(2, STR6); vector.setValueCount(3); // overwrite index 1 with a long string vector.set(1, STR7); vector.setValueCount(3); validateViewBuffer(0, vector, STR0, /*NA*/ -1, /*NA*/ -1); validateViewBuffer(1, vector, STR7, 0, 0); validateViewBuffer(2, vector, STR6, /*NA*/ -1, /*NA*/ -1); } // Overwriting in the middle of the buffer with a mix of short and long strings. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(80, 5); // set short string 1 vector.set(0, STR0); // set long string 1 vector.set(1, STR3); // set short string 2 vector.set(2, STR5); // set short string 3 vector.set(3, STR6); // set long string 2 vector.set(4, STR7); vector.setValueCount(5); // overwrite index 2 with a long string vector.set(2, STR8); vector.setValueCount(5); validateViewBuffer(0, vector, STR0, /*NA*/ -1, /*NA*/ -1); validateViewBuffer(1, vector, STR3, 0, 0); // Since we did overwrite index 2 with STR8, and as we are using append-only approach, // it will be appended to the data buffer. // Thus, it will be stored in the dataBuffer in order i.e. [STR3, STR7, STR8]. validateViewBuffer(2, vector, STR8, 0, STR3.length + STR7.length); validateViewBuffer(3, vector, STR6, /*NA*/ -1, /*NA*/ -1); validateViewBuffer(4, vector, STR7, 0, STR3.length); } // Overwriting in the middle of the buffer with a mix of short and long strings. // Here the short string is overwritten with a long string, and its length is larger than // the remaining capacity of the existing data buffer. // This would allocate a new buffer in the data buffers. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(80, 5); // set short string 1 vector.set(0, STR0); // set long string 1 vector.set(1, STR3); // set short string 2 vector.set(2, STR5); // set short string 3 vector.set(3, STR6); // set long string 2 vector.set(4, STR7); vector.setValueCount(5); // overwrite index 2 with a long string 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. final ArrowBuf currentDataBuf = vector.dataBuffers.get(0); final long remainingCapacity = currentDataBuf.capacity() - currentDataBuf.writerIndex(); assertTrue(remainingCapacity < longStringBytes.length); vector.set(2, longStringBytes); vector.setValueCount(5); validateViewBuffer(0, vector, STR0, /*NA*/ -1, /*NA*/ -1); validateViewBuffer(1, vector, STR3, 0, 0); // overwritten long string will be stored in the new data buffer. validateViewBuffer(2, vector, longStringBytes, 1, 0); validateViewBuffer(3, vector, STR6, /*NA*/ -1, /*NA*/ -1); validateViewBuffer(4, vector, STR7, 0, STR3.length); } } @Test public void testOverwriteLongFromShortString() { // Overwriting at the beginning of the buffer. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(16, 1); // set short string vector.set(0, STR3); vector.setValueCount(1); // set long string vector.set(0, STR0); vector.setValueCount(1); validateViewBuffer(0, vector, STR0, /*NA*/ -1, /*NA*/ -1); } // Overwriting in the middle of the buffer when existing buffers are all longs. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(48, 3); // set long string 1 vector.set(0, STR3); // set long string 2 vector.set(1, STR8); // set long string 3 vector.set(2, STR7); vector.setValueCount(3); // overwrite index 1 with a short string vector.set(1, STR6); vector.setValueCount(3); validateViewBuffer(0, vector, STR3, 0, 0); validateViewBuffer(1, vector, STR6, /*NA*/ -1, /*NA*/ -1); // since the append-only approach is used, // STR8 will still be in the first data buffer in dataBuffers. validateViewBuffer(2, vector, STR7, 0, STR3.length + STR8.length); } // Overwriting in the middle of the buffer with a mix of short and long strings. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(80, 5); // set long string 1 vector.set(0, STR3); // set short string 1 vector.set(1, STR5); // set long string 2 vector.set(2, STR7); // set long string 3 vector.set(3, STR8); // set short string 2 vector.set(4, STR6); vector.setValueCount(5); // overwrite index 2 with a short string vector.set(2, STR0); vector.setValueCount(5); validateViewBuffer(0, vector, STR3, 0, 0); validateViewBuffer(1, vector, STR5, /*NA*/ -1, /*NA*/ -1); validateViewBuffer(2, vector, STR0, /*NA*/ -1, /*NA*/ -1); // since the append-only approach is used, // STR7 will still be in the first data buffer in dataBuffers. validateViewBuffer(3, vector, STR8, 0, STR3.length + STR7.length); validateViewBuffer(4, vector, STR6, /*NA*/ -1, /*NA*/ -1); } } @Test public void testOverwriteLongFromAShorterLongString() { // Overwriting at the beginning of the buffer. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(16, 1); // set long string vector.set(0, STR7); vector.setValueCount(1); // set shorter long string, since append-only approach is used and the remaining capacity // is not enough to store the new string; a new buffer will be allocated. final ArrowBuf currentDataBuf = vector.dataBuffers.get(0); final long remainingCapacity = currentDataBuf.capacity() - currentDataBuf.writerIndex(); assertTrue(remainingCapacity < STR3.length); // set shorter long string vector.set(0, STR3); vector.setValueCount(1); validateViewBuffer(0, vector, STR3, 1, 0); } // Overwriting in the middle of the buffer when existing buffers are all longs. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { // extra memory is allocated vector.allocateNew(128, 3); // set long string 1 vector.set(0, STR3); // set long string 2 vector.set(1, STR8); // set long string 3 vector.set(2, STR7); vector.setValueCount(3); // overwrite index 1 with a shorter long string // Since append-only approach is used // and 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(); assertTrue(remainingCapacity > STR2.length); vector.set(1, STR2); vector.setValueCount(3); validateViewBuffer(0, vector, STR3, 0, 0); // since the append-only approach is used, // STR8 will still be in the first data buffer in dataBuffers. validateViewBuffer(1, vector, STR2, 0, STR3.length + STR8.length + STR7.length); validateViewBuffer(2, vector, STR7, 0, STR3.length + STR8.length); } // Overwriting in the middle of the buffer with a mix of short and long strings. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(128, 5); // set long string 1 vector.set(0, STR3); // set short string 1 vector.set(1, STR5); // set long string 2 vector.set(2, STR7); // set long string 3 vector.set(3, STR8); // set short string 2 vector.set(4, STR6); vector.setValueCount(5); // overwrite index 2 with a shorter long string // Since append-only approach is used // and 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(); assertTrue(remainingCapacity > STR2.length); vector.set(2, STR2); vector.setValueCount(5); validateViewBuffer(0, vector, STR3, 0, 0); validateViewBuffer(1, vector, STR5, /*NA*/ -1, /*NA*/ -1); // since the append-only approach is used, // STR7 will still be in the first data buffer in dataBuffers. validateViewBuffer(2, vector, STR2, 0, STR3.length + STR7.length + STR8.length); validateViewBuffer(3, vector, STR8, 0, STR3.length + STR7.length); validateViewBuffer(4, vector, STR6, /*NA*/ -1, /*NA*/ -1); } } @Test public void testOverwriteLongFromALongerLongString() { // Overwriting at the beginning of the buffer. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(16, 1); // set long string vector.set(0, STR3); vector.setValueCount(1); // set longer long string, since append-only approach is used and the remaining capacity // is not enough to store the new string; a new buffer will be allocated. final ArrowBuf currentDataBuf = vector.dataBuffers.get(0); final long remainingCapacity = currentDataBuf.capacity() - currentDataBuf.writerIndex(); assertTrue(remainingCapacity < STR7.length); // set longer long string vector.set(0, STR7); vector.setValueCount(1); validateViewBuffer(0, vector, STR7, 1, 0); } // Overwriting in the middle of the buffer when existing buffers are all longs. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { // extra memory is allocated vector.allocateNew(48, 3); // set long string 1 vector.set(0, STR3); // set long string 2 vector.set(1, STR8); // set long string 3 vector.set(2, STR7); vector.setValueCount(3); // overwrite index 1 with a longer long string // the remaining capacity is not enough to store in the same data buffer // 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 = TestUtils.generateRandomString(35); byte[] longerStringBytes = longerString.getBytes(StandardCharsets.UTF_8); assertTrue(remainingCapacity < longerStringBytes.length); vector.set(1, longerStringBytes); vector.setValueCount(3); validateViewBuffer(0, vector, STR3, 0, 0); validateViewBuffer(1, vector, longerStringBytes, 1, 0); // since the append-only approach is used, // STR8 will still be in the first data buffer in dataBuffers. validateViewBuffer(2, vector, STR7, 0, STR3.length + STR8.length); } // Overwriting in the middle of the buffer with a mix of short and long strings. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(128, 5); // set long string 1 vector.set(0, STR3); // set short string 1 vector.set(1, STR5); // set long string 2 vector.set(2, STR7); // set long string 3 vector.set(3, STR2); // set short string 2 vector.set(4, STR6); vector.setValueCount(5); // overwrite index 2 with a longer long string // 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 = TestUtils.generateRandomString(24); byte[] longerStringBytes = longerString.getBytes(StandardCharsets.UTF_8); assertTrue(remainingCapacity > longerStringBytes.length); vector.set(2, longerStringBytes); vector.setValueCount(5); validateViewBuffer(0, vector, STR3, 0, 0); validateViewBuffer(1, vector, STR5, /*NA*/ -1, /*NA*/ -1); // since the append-only approach is used, // STR7 will still be in the first data buffer in dataBuffers. validateViewBuffer(2, vector, longerStringBytes, 0, STR3.length + STR7.length + STR2.length); validateViewBuffer(3, vector, STR2, 0, STR3.length + STR7.length); validateViewBuffer(4, vector, STR6, /*NA*/ -1, /*NA*/ -1); } } @Test public void testSafeOverwriteShortFromLongString() { /*NA: not applicable */ // Overwriting at the beginning of the buffer. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(16, 1); // set short string vector.setSafe(0, STR0); vector.setValueCount(1); assertEquals(0, vector.dataBuffers.size()); assertArrayEquals(STR0, vector.get(0)); // set long string vector.setSafe(0, STR3); vector.setValueCount(1); assertEquals(1, vector.dataBuffers.size()); assertArrayEquals(STR3, vector.get(0)); } // Overwriting in the middle of the buffer when existing buffers are all shorts. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(16, 3); // set short string 1 vector.setSafe(0, STR0); // set short string 2 vector.setSafe(1, STR5); // set short string 3 vector.setSafe(2, STR6); vector.setValueCount(3); // overwrite index 1 with a long string vector.setSafe(1, STR7); vector.setValueCount(3); assertArrayEquals(STR0, vector.get(0)); assertArrayEquals(STR7, vector.get(1)); assertArrayEquals(STR6, vector.get(2)); } // Overwriting in the middle of the buffer with a mix of short and long strings. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(16, 5); // set short string 1 vector.setSafe(0, STR0); // set long string 1 vector.setSafe(1, STR3); // set short string 2 vector.setSafe(2, STR5); // set short string 3 vector.setSafe(3, STR6); // set long string 2 vector.setSafe(4, STR7); vector.setValueCount(5); // overwrite index 2 with a long string vector.setSafe(2, STR8); vector.setValueCount(5); assertArrayEquals(STR0, vector.get(0)); assertArrayEquals(STR3, vector.get(1)); assertArrayEquals(STR8, vector.get(2)); assertArrayEquals(STR6, vector.get(3)); assertArrayEquals(STR7, vector.get(4)); } // Overwriting in the middle of the buffer with a mix of short and long strings. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(16, 5); // set short string 1 vector.setSafe(0, STR0); // set long string 1 vector.setSafe(1, STR3); // set short string 2 vector.setSafe(2, STR5); // set short string 3 vector.setSafe(3, STR6); // set long string 2 vector.setSafe(4, STR7); vector.setValueCount(5); // overwrite index 2 with a long string String longString = TestUtils.generateRandomString(128); byte[] longStringBytes = longString.getBytes(StandardCharsets.UTF_8); vector.setSafe(2, longStringBytes); vector.setValueCount(5); assertArrayEquals(STR0, vector.get(0)); assertArrayEquals(STR3, vector.get(1)); assertArrayEquals(longStringBytes, vector.get(2)); assertArrayEquals(STR6, vector.get(3)); assertArrayEquals(STR7, vector.get(4)); } } @Test public void testSafeOverwriteLongFromShortString() { // Overwriting at the beginning of the buffer. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(16, 1); // set short string vector.setSafe(0, STR3); vector.setValueCount(1); // set long string vector.setSafe(0, STR0); vector.setValueCount(1); assertArrayEquals(STR0, vector.get(0)); } // Overwriting in the middle of the buffer when existing buffers are all longs. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(16, 3); // set long string 1 vector.setSafe(0, STR3); // set long string 2 vector.setSafe(1, STR8); // set long string 3 vector.setSafe(2, STR7); vector.setValueCount(3); // overwrite index 1 with a short string vector.setSafe(1, STR6); vector.setValueCount(3); assertArrayEquals(STR3, vector.get(0)); assertArrayEquals(STR6, vector.get(1)); assertArrayEquals(STR7, vector.get(2)); } // Overwriting in the middle of the buffer with a mix of short and long strings. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(16, 5); // set long string 1 vector.setSafe(0, STR3); // set short string 1 vector.setSafe(1, STR5); // set long string 2 vector.setSafe(2, STR7); // set long string 3 vector.setSafe(3, STR8); // set short string 2 vector.setSafe(4, STR6); vector.setValueCount(5); // overwrite index 2 with a short string vector.setSafe(2, STR0); vector.setValueCount(5); assertArrayEquals(STR3, vector.get(0)); assertArrayEquals(STR5, vector.get(1)); assertArrayEquals(STR0, vector.get(2)); assertArrayEquals(STR8, vector.get(3)); assertArrayEquals(STR6, vector.get(4)); } } @Test public void testSafeOverwriteLongFromAShorterLongString() { // Overwriting at the beginning of the buffer. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(16, 1); // set long string vector.setSafe(0, STR7); vector.setValueCount(1); // set shorter long string vector.setSafe(0, STR3); vector.setValueCount(1); assertArrayEquals(STR3, vector.get(0)); } // Overwriting in the middle of the buffer when existing buffers are all longs. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { // extra memory is allocated vector.allocateNew(16, 3); // set long string 1 vector.setSafe(0, STR3); // set long string 2 vector.setSafe(1, STR8); // set long string 3 vector.setSafe(2, STR7); vector.setValueCount(3); // overwrite index 1 with a shorter long string vector.setSafe(1, STR2); vector.setValueCount(3); assertArrayEquals(STR3, vector.get(0)); assertArrayEquals(STR2, vector.get(1)); assertArrayEquals(STR7, vector.get(2)); } // Overwriting in the middle of the buffer with a mix of short and long strings. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(16, 5); // set long string 1 vector.setSafe(0, STR3); // set short string 1 vector.setSafe(1, STR5); // set long string 2 vector.setSafe(2, STR7); // set long string 3 vector.setSafe(3, STR8); // set short string 2 vector.setSafe(4, STR6); vector.setValueCount(5); // overwrite index 2 with a shorter long string vector.setSafe(2, STR2); vector.setValueCount(5); assertArrayEquals(STR3, vector.get(0)); assertArrayEquals(STR5, vector.get(1)); assertArrayEquals(STR2, vector.get(2)); assertArrayEquals(STR8, vector.get(3)); assertArrayEquals(STR6, vector.get(4)); } } @Test public void testSafeOverwriteLongFromALongerLongString() { // Overwriting at the beginning of the buffer. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(16, 1); // set long string vector.setSafe(0, STR3); vector.setValueCount(1); // set longer long string vector.setSafe(0, STR7); vector.setValueCount(1); assertArrayEquals(STR7, vector.get(0)); } // Overwriting in the middle of the buffer when existing buffers are all longs. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { // extra memory is allocated vector.allocateNew(16, 3); // set long string 1 vector.setSafe(0, STR3); // set long string 2 vector.setSafe(1, STR8); // set long string 3 vector.setSafe(2, STR7); vector.setValueCount(3); String longerString = TestUtils.generateRandomString(35); byte[] longerStringBytes = longerString.getBytes(StandardCharsets.UTF_8); vector.setSafe(1, longerStringBytes); vector.setValueCount(3); assertArrayEquals(STR3, vector.get(0)); assertArrayEquals(longerStringBytes, vector.get(1)); assertArrayEquals(STR7, vector.get(2)); } // Overwriting in the middle of the buffer with a mix of short and long strings. try (final ViewVarCharVector vector = new ViewVarCharVector("myviewvector", allocator)) { vector.allocateNew(16, 5); // set long string 1 vector.setSafe(0, STR3); // set short string 1 vector.setSafe(1, STR5); // set long string 2 vector.setSafe(2, STR7); // set long string 3 vector.setSafe(3, STR2); // set short string 2 vector.setSafe(4, STR6); vector.setValueCount(5); String longerString = TestUtils.generateRandomString(24); byte[] longerStringBytes = longerString.getBytes(StandardCharsets.UTF_8); vector.setSafe(2, longerStringBytes); vector.setValueCount(5); assertArrayEquals(STR3, vector.get(0)); assertArrayEquals(STR5, vector.get(1)); assertArrayEquals(longerStringBytes, vector.get(2)); assertArrayEquals(STR2, vector.get(3)); assertArrayEquals(STR6, vector.get(4)); } } @Test public void testVectorLoadUnloadInLine() { try (final ViewVarCharVector vector1 = new ViewVarCharVector("myvector", allocator)) { setVector(vector1, STR0, STR1, STR4, STR5, STR6); assertEquals(4, vector1.getLastSet()); vector1.setValueCount(15); assertEquals(14, vector1.getLastSet()); /* Check the vector output */ assertArrayEquals(STR0, vector1.get(0)); assertArrayEquals(STR1, vector1.get(1)); assertArrayEquals(STR4, vector1.get(2)); assertArrayEquals(STR5, vector1.get(3)); assertArrayEquals(STR6, vector1.get(4)); Field field = vector1.getField(); String fieldName = field.getName(); List fields = new ArrayList<>(); List fieldVectors = new ArrayList<>(); fields.add(field); fieldVectors.add(vector1); Schema schema = new Schema(fields); VectorSchemaRoot schemaRoot1 = new VectorSchemaRoot(schema, fieldVectors, vector1.getValueCount()); VectorUnloader vectorUnloader = new VectorUnloader(schemaRoot1); try (ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch(); BufferAllocator finalVectorsAllocator = allocator.newChildAllocator("new vector", 0, Long.MAX_VALUE); VectorSchemaRoot schemaRoot2 = VectorSchemaRoot.create(schema, finalVectorsAllocator); ) { VectorLoader vectorLoader = new VectorLoader(schemaRoot2); vectorLoader.load(recordBatch); ViewVarCharVector vector2 = (ViewVarCharVector) schemaRoot2.getVector(fieldName); /* * lastSet would have internally been set by VectorLoader.load() when it invokes * loadFieldBuffers. */ assertEquals(14, vector2.getLastSet()); vector2.setValueCount(25); assertEquals(24, vector2.getLastSet()); /* Check the vector output */ assertArrayEquals(STR0, vector2.get(0)); assertArrayEquals(STR1, vector2.get(1)); assertArrayEquals(STR4, vector2.get(2)); assertArrayEquals(STR5, vector2.get(3)); assertArrayEquals(STR6, vector2.get(4)); } } } @Test public void testVectorLoadUnload() { try (final ViewVarCharVector vector1 = new ViewVarCharVector("myvector", allocator)) { setVector(vector1, STR1, STR2, STR3, STR4, STR5, STR6); assertEquals(5, vector1.getLastSet()); vector1.setValueCount(15); assertEquals(14, vector1.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector1.get(0)); assertArrayEquals(STR2, vector1.get(1)); assertArrayEquals(STR3, vector1.get(2)); assertArrayEquals(STR4, vector1.get(3)); assertArrayEquals(STR5, vector1.get(4)); assertArrayEquals(STR6, vector1.get(5)); Field field = vector1.getField(); String fieldName = field.getName(); List fields = new ArrayList<>(); List fieldVectors = new ArrayList<>(); fields.add(field); fieldVectors.add(vector1); Schema schema = new Schema(fields); VectorSchemaRoot schemaRoot1 = new VectorSchemaRoot(schema, fieldVectors, vector1.getValueCount()); VectorUnloader vectorUnloader = new VectorUnloader(schemaRoot1); try (ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch(); BufferAllocator finalVectorsAllocator = allocator.newChildAllocator("new vector", 0, Long.MAX_VALUE); VectorSchemaRoot schemaRoot2 = VectorSchemaRoot.create(schema, finalVectorsAllocator); ) { VectorLoader vectorLoader = new VectorLoader(schemaRoot2); vectorLoader.load(recordBatch); ViewVarCharVector vector2 = (ViewVarCharVector) schemaRoot2.getVector(fieldName); /* * lastSet would have internally been set by VectorLoader.load() when it invokes * loadFieldBuffers. */ assertEquals(14, vector2.getLastSet()); vector2.setValueCount(25); assertEquals(24, vector2.getLastSet()); /* Check the vector output */ assertArrayEquals(STR1, vector2.get(0)); assertArrayEquals(STR2, vector2.get(1)); assertArrayEquals(STR3, vector2.get(2)); assertArrayEquals(STR4, vector2.get(3)); assertArrayEquals(STR5, vector2.get(4)); assertArrayEquals(STR6, vector2.get(5)); } } } static Stream vectorCreatorProvider() { return Stream.of( Arguments.of( (Function) (allocator -> newVector( ViewVarBinaryVector.class, EMPTY_SCHEMA_PATH, Types.MinorType.VIEWVARBINARY, allocator))), Arguments.of( (Function) (allocator -> newVector( ViewVarCharVector.class, EMPTY_SCHEMA_PATH, Types.MinorType.VIEWVARCHAR, allocator)))); } @ParameterizedTest @MethodSource({"vectorCreatorProvider"}) public void testCopyFromWithNulls( Function vectorCreator) { try (final BaseVariableWidthViewVector vector = vectorCreator.apply(allocator); final BaseVariableWidthViewVector vector2 = vectorCreator.apply(allocator)) { final int initialCapacity = 1024; vector.setInitialCapacity(initialCapacity); vector.allocateNew(); int capacity = vector.getValueCapacity(); assertTrue(capacity >= initialCapacity); // setting number of values such that we have enough space in the initial allocation // to avoid re-allocation. This is to test copyFrom() without re-allocation. final int numberOfValues = initialCapacity / 2 / ViewVarCharVector.ELEMENT_SIZE; final String prefixString = TestUtils.generateRandomString(12); for (int i = 0; i < numberOfValues; i++) { if (i % 3 == 0) { // null values vector.setNull(i); } else if (i % 3 == 1) { // short strings byte[] b = Integer.toString(i).getBytes(StandardCharsets.UTF_8); vector.set(i, b, 0, b.length); } else { // long strings byte[] b = (i + prefixString).getBytes(StandardCharsets.UTF_8); vector.set(i, b, 0, b.length); } } assertEquals(capacity, vector.getValueCapacity()); vector.setValueCount(numberOfValues); for (int i = 0; i < numberOfValues; i++) { if (i % 3 == 0) { assertNull(vector.getObject(i)); } else if (i % 3 == 1) { assertArrayEquals( Integer.toString(i).getBytes(StandardCharsets.UTF_8), vector.get(i), "unexpected value at index: " + i); } else { assertArrayEquals( (i + prefixString).getBytes(StandardCharsets.UTF_8), vector.get(i), "unexpected value at index: " + i); } } vector2.setInitialCapacity(initialCapacity); vector2.allocateNew(); int capacity2 = vector2.getValueCapacity(); assertEquals(capacity2, capacity); for (int i = 0; i < numberOfValues; i++) { vector2.copyFrom(i, i, vector); if (i % 3 == 0) { assertNull(vector2.getObject(i)); } else if (i % 3 == 1) { assertArrayEquals( Integer.toString(i).getBytes(StandardCharsets.UTF_8), vector2.get(i), "unexpected value at index: " + i); } else { assertArrayEquals( (i + prefixString).getBytes(StandardCharsets.UTF_8), vector2.get(i), "unexpected value at index: " + i); } } assertEquals(capacity, vector2.getValueCapacity()); vector2.setValueCount(numberOfValues); for (int i = 0; i < numberOfValues; i++) { if (i % 3 == 0) { assertNull(vector2.getObject(i)); } else if (i % 3 == 1) { assertArrayEquals( Integer.toString(i).getBytes(StandardCharsets.UTF_8), vector2.get(i), "unexpected value at index: " + i); } else { assertArrayEquals( (i + prefixString).getBytes(StandardCharsets.UTF_8), vector2.get(i), "unexpected value at index: " + i); } } } } @ParameterizedTest @MethodSource("vectorCreatorProvider") public void testCopyFromSafeWithNulls( Function vectorCreator) { try (final BaseVariableWidthViewVector vector = vectorCreator.apply(allocator); final BaseVariableWidthViewVector vector2 = vectorCreator.apply(allocator)) { final int initialCapacity = 4096; vector.setInitialCapacity(initialCapacity); vector.allocateNew(); int capacity = vector.getValueCapacity(); assertTrue(capacity >= initialCapacity); final int numberOfValues = initialCapacity / ViewVarCharVector.ELEMENT_SIZE; final String prefixString = TestUtils.generateRandomString(12); for (int i = 0; i < numberOfValues; i++) { if (i % 3 == 0) { // null values vector.setNull(i); } else if (i % 3 == 1) { // short strings byte[] b = Integer.toString(i).getBytes(StandardCharsets.UTF_8); vector.setSafe(i, b, 0, b.length); } else { // long strings byte[] b = (i + prefixString).getBytes(StandardCharsets.UTF_8); vector.setSafe(i, b, 0, b.length); } } /* NO reAlloc() should have happened in setSafe() */ assertEquals(capacity, vector.getValueCapacity()); vector.setValueCount(numberOfValues); for (int i = 0; i < numberOfValues; i++) { if (i % 3 == 0) { assertNull(vector.getObject(i)); } else if (i % 3 == 1) { assertArrayEquals( Integer.toString(i).getBytes(StandardCharsets.UTF_8), vector.get(i), "unexpected value at index: " + i); } else { assertArrayEquals( (i + prefixString).getBytes(StandardCharsets.UTF_8), vector.get(i), "unexpected value at index: " + i); } } vector2.setInitialCapacity(initialCapacity); vector2.allocateNew(); int capacity2 = vector2.getValueCapacity(); assertEquals(capacity2, capacity); for (int i = 0; i < numberOfValues; i++) { vector2.copyFromSafe(i, i, vector); if (i % 3 == 0) { assertNull(vector2.getObject(i)); } else if (i % 3 == 1) { assertArrayEquals( Integer.toString(i).getBytes(StandardCharsets.UTF_8), vector2.get(i), "unexpected value at index: " + i); } else { assertArrayEquals( (i + prefixString).getBytes(StandardCharsets.UTF_8), vector2.get(i), "unexpected value at index: " + i); } } /* NO reAlloc() should have happened in setSafe() */ assertEquals(capacity, vector2.getValueCapacity()); vector2.setValueCount(numberOfValues); for (int i = 0; i < numberOfValues; i++) { if (i % 3 == 0) { assertNull(vector2.getObject(i)); } else if (i % 3 == 1) { assertArrayEquals( Integer.toString(i).getBytes(StandardCharsets.UTF_8), vector2.get(i), "unexpected value at index: " + i); } else { assertArrayEquals( (i + prefixString).getBytes(StandardCharsets.UTF_8), vector2.get(i), "unexpected value at index: " + i); } } // make it reallocate int valueCapacity = vector2.getValueCapacity(); for (int i = 0; i < numberOfValues; i++) { int thisIndex = i + valueCapacity; vector2.copyFromSafe(i, thisIndex, vector); if (i % 3 == 0) { assertNull(vector2.getObject(thisIndex)); } else if (i % 3 == 1) { assertArrayEquals( Integer.toString(i).getBytes(StandardCharsets.UTF_8), vector2.get(thisIndex), "unexpected value at index: " + i); } else { assertArrayEquals( (i + prefixString).getBytes(StandardCharsets.UTF_8), vector2.get(thisIndex), "unexpected value at index: " + i); } } // test target vector with different initialCapacity try (final BaseVariableWidthViewVector vector3 = vectorCreator.apply(allocator)) { vector3.setInitialCapacity(16); vector3.allocateNew(); for (int i = 0; i < numberOfValues; i++) { vector3.copyFromSafe(i, i, vector); if (i % 3 == 0) { assertNull(vector3.getObject(i)); } else { assertArrayEquals(vector.get(i), vector3.get(i)); } } } // test overwrite a used vector by copy try (final BaseVariableWidthViewVector targetVector = vectorCreator.apply(allocator)) { targetVector.setInitialCapacity(initialCapacity); targetVector.allocateNew(); // source vector: null, short, long... // target vector: long, null, short... for (int i = 0; i < numberOfValues; i++) { if (i % 3 == 0) { // long strings byte[] b = (i + prefixString).getBytes(StandardCharsets.UTF_8); targetVector.set(i, b, 0, b.length); } else if (i % 3 == 1) { // null values targetVector.setNull(i); } else { // short strings byte[] b = Integer.toString(i).getBytes(StandardCharsets.UTF_8); targetVector.set(i, b, 0, b.length); } } targetVector.setValueCount(numberOfValues); for (int i = 0; i < numberOfValues; i++) { targetVector.copyFromSafe(i, i, vector); if (i % 3 == 0) { assertNull(targetVector.getObject(i)); } else { assertArrayEquals(vector.get(i), targetVector.get(i)); } } } } } public byte[] generateRandomBinaryData(int size) { byte[] binaryData = new byte[size]; random.nextBytes(binaryData); return binaryData; } private byte[][] generateBinaryDataArray(int size, int length) { byte[][] binaryDataArray = new byte[size][]; for (int i = 0; i < size; i++) { binaryDataArray[i] = generateRandomBinaryData(length); } return binaryDataArray; } private void testSplitAndTransferOnSlicedBufferHelper( BaseVariableWidthViewVector targetVector, BaseVariableWidthViewVector sourceVector, int startIndex, int length, byte[][] data) { sourceVector.allocateNew(1024 * 10, 1024); for (int i = 0; i < data.length; i++) { sourceVector.set(i, data[i]); } sourceVector.setValueCount(data.length); final long allocatedMem = allocator.getAllocatedMemory(); final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); sourceVector.splitAndTransferTo(startIndex, length, targetVector); // we allocate view and data buffers for the target vector assertTrue(allocatedMem < allocator.getAllocatedMemory()); // The validity buffer is sliced from the same buffer.See // BaseFixedWidthViewVector#allocateBytes. // Therefore, the refcnt of the validity buffer is increased once since the startIndex is 0. assertEquals(validityRefCnt + 1, sourceVector.getValidityBuffer().refCnt()); // since the new view buffer is allocated, the refcnt is the same as the source vector. assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); } /** * ARROW-7831: this checks a slice taken off a buffer is still readable after that buffer's * allocator is closed. With short strings. */ @Test public void testSplitAndTransferWithShortStringOnSlicedBuffer() { final byte[][] data = new byte[][] {STR4, STR5, STR6}; final int startIndex = 0; final int length = 2; BiConsumer validateVector = (targetVector, expectedData) -> { IntStream.range(startIndex, length) .forEach(i -> assertArrayEquals(expectedData[i], targetVector.get(i - startIndex))); }; try (final ViewVarCharVector targetVector = newViewVarCharVector("split-target", allocator)) { try (final ViewVarCharVector sourceVector = newViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { testSplitAndTransferOnSlicedBufferHelper( targetVector, sourceVector, startIndex, length, data); } validateVector.accept(targetVector, data); } final byte[][] binaryData = generateBinaryDataArray(3, 10); try (final ViewVarBinaryVector targetVector = newViewVarBinaryVector("split-target", allocator)) { try (final ViewVarBinaryVector sourceVector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { testSplitAndTransferOnSlicedBufferHelper( targetVector, sourceVector, startIndex, length, binaryData); } validateVector.accept(targetVector, binaryData); } } /** * ARROW-7831: this checks a slice taken off a buffer is still readable after that buffer's * allocator is closed. With a long string included. */ @Test public void testSplitAndTransferWithLongStringsOnSlicedBuffer() { final byte[][] data = new byte[][] {STR2, STR5, STR6}; final int startIndex = 0; final int length = 2; BiConsumer validateVector = (targetVector, expectedData) -> { IntStream.range(startIndex, length) .forEach(i -> assertArrayEquals(expectedData[i], targetVector.get(i - startIndex))); }; try (final ViewVarCharVector targetVector = newViewVarCharVector("split-target", allocator)) { try (final ViewVarCharVector sourceVector = newViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { testSplitAndTransferOnSlicedBufferHelper( targetVector, sourceVector, startIndex, length, data); } validateVector.accept(targetVector, data); } final byte[][] binaryData = generateBinaryDataArray(3, 18); try (final ViewVarBinaryVector targetVector = newViewVarBinaryVector("split-target", allocator)) { try (final ViewVarBinaryVector sourceVector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { testSplitAndTransferOnSlicedBufferHelper( targetVector, sourceVector, startIndex, length, binaryData); } validateVector.accept(targetVector, binaryData); } } private void testSplitAndTransferOnSlicedVectorHelper( BaseVariableWidthViewVector sourceVector, BaseVariableWidthViewVector targetVector, int startIndex, int length, byte[][] data) { sourceVector.allocateNew(1024 * 10, 1024); for (int i = 0; i < data.length; i++) { sourceVector.set(i, data[i]); } sourceVector.setValueCount(data.length); final long allocatedMem = allocator.getAllocatedMemory(); final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); sourceVector.splitAndTransferTo(startIndex, length, targetVector); // we allocate view and data buffers for the target vector assertTrue(allocatedMem < allocator.getAllocatedMemory()); // The validity buffer is sliced from the same buffer.See // BaseFixedWidthViewVector#allocateBytes. // Therefore, the refcnt of the validity buffer is increased once since the startIndex is 0. assertEquals(validityRefCnt + 1, sourceVector.getValidityBuffer().refCnt()); // since the new view buffer is allocated, the refcnt is the same as the source vector. assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); for (int i = startIndex; i < length; i++) { assertArrayEquals(data[i], targetVector.get(i - startIndex)); } } /** * ARROW-7831: this checks a vector that got sliced is still readable after the slice's allocator * got closed. With short strings. */ @Test public void testSplitAndTransferWithShortStringsOnSlicedVector() { byte[][] data = new byte[][] {STR4, STR5, STR6}; final int startIndex = 0; final int length = 2; BiConsumer validateVector = (sourceVector, expectedData) -> { IntStream.range(startIndex, length) .forEach(i -> assertArrayEquals(expectedData[i], sourceVector.get(i))); }; try (final ViewVarCharVector sourceVector = newViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { try (final ViewVarCharVector targetVector = newViewVarCharVector("split-target", allocator)) { testSplitAndTransferOnSlicedVectorHelper( sourceVector, targetVector, startIndex, length, data); } validateVector.accept(sourceVector, data); } byte[][] binaryData = generateBinaryDataArray(3, 10); try (final ViewVarBinaryVector sourceVector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { try (final ViewVarBinaryVector targetVector = newViewVarBinaryVector("split-target", allocator)) { testSplitAndTransferOnSlicedVectorHelper( sourceVector, targetVector, startIndex, length, binaryData); } validateVector.accept(sourceVector, binaryData); } } /** * ARROW-7831: this checks a vector that got sliced is still readable after the slice's allocator * got closed. With a long string included. */ @Test public void testSplitAndTransferWithLongStringsOnSlicedVector() { final byte[][] data = new byte[][] {STR2, STR5, STR6}; final int startIndex = 0; final int length = 2; BiConsumer validateVector = (sourceVector, expectedData) -> { IntStream.range(startIndex, length) .forEach(i -> assertArrayEquals(expectedData[i], sourceVector.get(i))); }; try (final ViewVarCharVector sourceVector = newViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { try (final ViewVarCharVector targetVector = newViewVarCharVector("split-target", allocator)) { testSplitAndTransferOnSlicedVectorHelper( sourceVector, targetVector, startIndex, length, data); } validateVector.accept(sourceVector, data); } final byte[][] binaryData = generateBinaryDataArray(3, 20); try (final ViewVarBinaryVector sourceVector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { try (final ViewVarBinaryVector targetVector = newViewVarBinaryVector("split-target", allocator)) { testSplitAndTransferOnSlicedVectorHelper( sourceVector, targetVector, startIndex, length, binaryData); } validateVector.accept(sourceVector, binaryData); } } private void testSplitAndTransferOnValiditySplitHelper( BaseVariableWidthViewVector targetVector, BaseVariableWidthViewVector sourceVector, int startIndex, int length, byte[][] data) { sourceVector.allocateNew(1024 * 10, 1024); sourceVector.set(0, new byte[0]); sourceVector.setNull(1); for (int i = 0; i < data.length; i++) { if (data[i] == null) { sourceVector.setNull(i); } else { sourceVector.set(i, data[i]); } } sourceVector.setValueCount(data.length); final long allocatedMem = allocator.getAllocatedMemory(); final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); sourceVector.splitAndTransferTo(startIndex, length, targetVector); // the allocation only consists in the size needed for the validity buffer final long validitySize = DefaultRoundingPolicy.DEFAULT_ROUNDING_POLICY.getRoundedSize( 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. // Since values up to the startIndex are empty/null validity refcnt should not change. assertEquals(validityRefCnt, sourceVector.getValidityBuffer().refCnt()); // since the new view buffer is allocated, the refcnt is the same as the source vector. assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); for (int i = startIndex; i < startIndex + length; i++) { assertArrayEquals(data[i], targetVector.get(i - startIndex)); } for (int i = 0; i < data.length; i++) { if (data[i] == null) { assertTrue(sourceVector.isNull(i)); } else { assertArrayEquals(data[i], sourceVector.get(i)); } } } /** * ARROW-7831: this checks a validity splitting where the validity buffer is sliced from the same * buffer. In the case where all the values up to the start of the slice are null/empty. With * short strings. */ @Test public void testSplitAndTransferWithShortStringsOnValiditySplit() { final byte[][] data = new byte[][] {new byte[0], null, STR4, STR5, STR6}; final int startIndex = 2; final int length = 2; try (final ViewVarCharVector targetVector = newViewVarCharVector("split-target", allocator); final ViewVarCharVector sourceVector = newViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { testSplitAndTransferOnValiditySplitHelper( targetVector, sourceVector, startIndex, length, data); } final byte[][] binaryData = generateBinaryDataArray(5, 10); binaryData[0] = new byte[0]; binaryData[1] = null; try (final ViewVarBinaryVector targetVector = newViewVarBinaryVector("split-target", allocator); final ViewVarBinaryVector sourceVector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { testSplitAndTransferOnValiditySplitHelper( targetVector, sourceVector, startIndex, length, binaryData); } } /** * ARROW-7831: this checks a validity splitting where the validity buffer is sliced from the same * buffer. In the case where all the values up to the start of the slice are null/empty. With long * strings. */ @Test public void testSplitAndTransferWithLongStringsOnValiditySplit() { final byte[][] data = new byte[][] {new byte[0], null, STR1, STR2, STR3}; final int startIndex = 2; final int length = 2; try (final ViewVarCharVector targetVector = newViewVarCharVector("split-target", allocator); final ViewVarCharVector sourceVector = newViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { testSplitAndTransferOnValiditySplitHelper( targetVector, sourceVector, startIndex, length, data); } final byte[][] binaryData = generateBinaryDataArray(5, 18); binaryData[0] = new byte[0]; binaryData[1] = null; try (final ViewVarBinaryVector targetVector = newViewVarBinaryVector("split-target", allocator); final ViewVarBinaryVector sourceVector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { testSplitAndTransferOnValiditySplitHelper( targetVector, sourceVector, startIndex, length, binaryData); } } private void testSplitAndTransferOnAllocatorToAllocator( BaseVariableWidthViewVector targetVector, BaseVariableWidthViewVector sourceVector, int startIndex, int length, byte[][] data) { sourceVector.allocateNew(50, data.length); for (int i = 0; i < data.length; i++) { sourceVector.set(i, data[i]); } sourceVector.setValueCount(data.length); final long allocatedMem = allocator.getAllocatedMemory(); final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); sourceVector.splitAndTransferTo(startIndex, length, targetVector); if (sourceVector.getDataBuffers().isEmpty()) { // no extra allocation as strings are all inline assertEquals(allocatedMem, allocator.getAllocatedMemory()); } else { // extra allocation as some strings are not inline assertTrue(allocatedMem < allocator.getAllocatedMemory()); } // the refcnts of each buffer for this test should be the same as what // the source allocator ended up with. assertEquals(validityRefCnt, sourceVector.getValidityBuffer().refCnt()); // since the new view buffer is allocated, the refcnt is the same as the source vector. assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); for (int i = 0; i < data.length; i++) { assertArrayEquals(data[i], sourceVector.get(i)); } } /** * ARROW-7831: ensures that data is transferred from one allocator to another in case of 0-index * start special cases. With short strings. */ @Test public void testSplitAndTransferWithShortStringsOnAllocatorToAllocator() { final int maxAllocation = 512; final byte[][] data = new byte[][] {STR4, STR5, STR6}; final int startIndex = 0; final int length = 2; BiConsumer validateVector = (targetVector, expectedData) -> { IntStream.range(startIndex, length) .forEach(i -> assertArrayEquals(expectedData[i], targetVector.get(i - startIndex))); }; try (final BufferAllocator targetAllocator = allocator.newChildAllocator("target-alloc", 256, maxAllocation); final ViewVarCharVector targetVector = newViewVarCharVector("split-target", targetAllocator)) { try (final BufferAllocator sourceAllocator = allocator.newChildAllocator("source-alloc", 256, maxAllocation); final ViewVarCharVector sourceVector = newViewVarCharVector(EMPTY_SCHEMA_PATH, sourceAllocator)) { testSplitAndTransferOnAllocatorToAllocator( targetVector, sourceVector, startIndex, length, data); } validateVector.accept(targetVector, data); } final byte[][] binaryData = generateBinaryDataArray(3, 10); try (final BufferAllocator targetAllocator = allocator.newChildAllocator("target-alloc", 256, maxAllocation); final ViewVarBinaryVector targetVector = newViewVarBinaryVector("split-target", targetAllocator)) { try (final BufferAllocator sourceAllocator = allocator.newChildAllocator("source-alloc", 256, maxAllocation); final ViewVarBinaryVector sourceVector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, sourceAllocator)) { testSplitAndTransferOnAllocatorToAllocator( targetVector, sourceVector, startIndex, length, binaryData); } validateVector.accept(targetVector, binaryData); } } /** * ARROW-7831: ensures that data is transferred from one allocator to another in case of 0-index * start special cases. With long strings. */ @Test public void testSplitAndTransferWithLongStringsOnAllocatorToAllocator() { final int initialReservation = 1024; // Here we have the target vector being transferred with a long string // hence, the data buffer will be allocated. // The default data buffer allocation takes // BaseVariableWidthViewVector.INITIAL_VIEW_VALUE_ALLOCATION * // BaseVariableWidthViewVector.ELEMENT_SIZE final byte[][] data = new byte[][] {STR1, STR2, STR3}; final int startIndex = 0; final int length = 2; BiConsumer validateVector = (targetVector, expectedData) -> { IntStream.range(startIndex, length) .forEach(i -> assertArrayEquals(expectedData[i], targetVector.get(i - startIndex))); }; final int maxAllocation = initialReservation + BaseVariableWidthViewVector.INITIAL_VIEW_VALUE_ALLOCATION * BaseVariableWidthViewVector.ELEMENT_SIZE; try (final BufferAllocator targetAllocator = allocator.newChildAllocator("target-alloc", initialReservation, maxAllocation); final ViewVarCharVector targetVector = newViewVarCharVector("split-target", targetAllocator)) { try (final BufferAllocator sourceAllocator = allocator.newChildAllocator("source-alloc", initialReservation, maxAllocation); final ViewVarCharVector sourceVector = newViewVarCharVector(EMPTY_SCHEMA_PATH, sourceAllocator)) { testSplitAndTransferOnAllocatorToAllocator( targetVector, sourceVector, startIndex, length, data); } validateVector.accept(targetVector, data); } final byte[][] binaryData = generateBinaryDataArray(3, 18); try (final BufferAllocator targetAllocator = allocator.newChildAllocator("target-alloc", initialReservation, maxAllocation); final ViewVarBinaryVector targetVector = newViewVarBinaryVector("split-target", targetAllocator)) { try (final BufferAllocator sourceAllocator = allocator.newChildAllocator("source-alloc", initialReservation, maxAllocation); final ViewVarBinaryVector sourceVector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, sourceAllocator)) { testSplitAndTransferOnAllocatorToAllocator( targetVector, sourceVector, startIndex, length, binaryData); } validateVector.accept(targetVector, binaryData); } } private void testReallocAfterVectorTransferHelper( BaseVariableWidthViewVector vector, byte[] str1, byte[] str2) { /* 4096 values with 16 bytes per record */ final int bytesPerRecord = 32; vector.allocateNew(4096 * bytesPerRecord, 4096); int valueCapacity = vector.getValueCapacity(); assertTrue(valueCapacity >= 4096); /* populate the vector */ for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { vector.set(i, str1); } else { vector.set(i, str2); } } /* Check the vector output */ for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { assertArrayEquals(str1, vector.get(i)); } else { assertArrayEquals(str2, vector.get(i)); } } /* trigger first realloc */ vector.setSafe(valueCapacity, str2, 0, str2.length); assertTrue(vector.getValueCapacity() >= 2 * valueCapacity); while (vector.getByteCapacity() < bytesPerRecord * vector.getValueCapacity()) { vector.reallocViewBuffer(); vector.reallocViewDataBuffer(); } /* populate the remaining vector */ for (int i = valueCapacity; i < vector.getValueCapacity(); i++) { if ((i & 1) == 1) { vector.set(i, str1); } else { vector.set(i, str2); } } /* Check the vector output */ valueCapacity = vector.getValueCapacity(); for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { assertArrayEquals(str1, vector.get(i)); } else { assertArrayEquals(str2, vector.get(i)); } } /* trigger second realloc */ vector.setSafe(valueCapacity + bytesPerRecord, str2, 0, str2.length); assertTrue(vector.getValueCapacity() >= 2 * valueCapacity); while (vector.getByteCapacity() < bytesPerRecord * vector.getValueCapacity()) { vector.reallocViewBuffer(); vector.reallocViewDataBuffer(); } /* populate the remaining vector */ for (int i = valueCapacity; i < vector.getValueCapacity(); i++) { if ((i & 1) == 1) { vector.set(i, str1); } else { vector.set(i, str2); } } /* Check the vector output */ valueCapacity = vector.getValueCapacity(); for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { assertArrayEquals(str1, vector.get(i)); } else { assertArrayEquals(str2, vector.get(i)); } } /* We are potentially working with 4x the size of vector buffer * that we initially started with. * Now let's transfer the vector. */ TransferPair transferPair = vector.getTransferPair(allocator); transferPair.transfer(); BaseVariableWidthViewVector toVector = (BaseVariableWidthViewVector) transferPair.getTo(); valueCapacity = toVector.getValueCapacity(); for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { assertArrayEquals(str1, toVector.get(i)); } else { assertArrayEquals(str2, toVector.get(i)); } } toVector.close(); } @Test public void testReallocAfterVectorTransfer() { try (final ViewVarCharVector vector = new ViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { testReallocAfterVectorTransferHelper(vector, STR1, STR2); } try (final ViewVarBinaryVector vector = new ViewVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { testReallocAfterVectorTransferHelper( vector, generateRandomBinaryData(12), generateRandomBinaryData(13)); } } private void testSplitAndTransferWithMultipleDataBuffersHelper( BaseVariableWidthViewVector sourceVector, BaseVariableWidthViewVector targetVector, int startIndex, int length, byte[][] data) { sourceVector.allocateNew(48, 4); for (int i = 0; i < data.length; i++) { sourceVector.set(i, data[i]); } sourceVector.setValueCount(data.length); // we should have multiple data buffers assertTrue(sourceVector.getDataBuffers().size() > 1); final long allocatedMem = allocator.getAllocatedMemory(); final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); // split and transfer with slice starting at the beginning: // this should not allocate anything new sourceVector.splitAndTransferTo(startIndex, length, targetVector); // we allocate view and data buffers for the target vector assertTrue(allocatedMem < allocator.getAllocatedMemory()); // the refcnts of each buffer for this test should be the same as what // the source allocator ended up with. assertEquals(validityRefCnt, sourceVector.getValidityBuffer().refCnt()); // since the new view buffer is allocated, the refcnt is the same as the source vector. assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); for (int i = 0; i < data.length; i++) { assertArrayEquals(data[i], sourceVector.get(i)); } } /** * ARROW-7831: ensures that data is transferred from one allocator to another in case of 0-index * start special cases. With long strings and multiple data buffers. Check multi-data buffer * source copying */ @Test public void testSplitAndTransferWithMultipleDataBuffers() { 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; BiConsumer validateVector = (targetVector, expectedData) -> { IntStream.range(startIndex, length) .forEach(i -> assertArrayEquals(expectedData[i], targetVector.get(i - startIndex))); }; try (final ViewVarCharVector targetVector = new ViewVarCharVector("target", allocator)) { try (final ViewVarCharVector sourceVector = new ViewVarCharVector("source", allocator)) { testSplitAndTransferWithMultipleDataBuffersHelper( sourceVector, targetVector, startIndex, length, data); } validateVector.accept(targetVector, data); } try (final ViewVarBinaryVector targetVector = new ViewVarBinaryVector("target", allocator)) { try (final ViewVarBinaryVector sourceVector = new ViewVarBinaryVector("source", allocator)) { testSplitAndTransferWithMultipleDataBuffersHelper( sourceVector, targetVector, startIndex, length, data); } validateVector.accept(targetVector, data); } } @Test public void testVectorLoadUnloadOnMixedTypes() { try (final IntVector vector1 = new IntVector("myvector", allocator); final ViewVarCharVector vector2 = new ViewVarCharVector("myviewvector", allocator)) { final int valueCount = 15; setVector(vector1, 1, 2, 3, 4, 5, 6); vector1.setValueCount(valueCount); setVector(vector2, STR1, STR2, STR3, STR4, STR5, STR6); vector1.setValueCount(valueCount); /* Check the vector output */ assertEquals(1, vector1.get(0)); assertEquals(2, vector1.get(1)); assertEquals(3, vector1.get(2)); assertEquals(4, vector1.get(3)); assertEquals(5, vector1.get(4)); assertEquals(6, vector1.get(5)); Field field1 = vector1.getField(); String fieldName1 = field1.getName(); Field field2 = vector2.getField(); String fieldName2 = field2.getName(); List fields = new ArrayList<>(2); List fieldVectors = new ArrayList<>(2); fields.add(field1); fields.add(field2); fieldVectors.add(vector1); fieldVectors.add(vector2); Schema schema = new Schema(fields); VectorSchemaRoot schemaRoot1 = new VectorSchemaRoot(schema, fieldVectors, valueCount); VectorUnloader vectorUnloader = new VectorUnloader(schemaRoot1); try (ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch(); BufferAllocator finalVectorsAllocator = allocator.newChildAllocator("new vector", 0, Long.MAX_VALUE); VectorSchemaRoot schemaRoot2 = VectorSchemaRoot.create(schema, finalVectorsAllocator); ) { // validating recordBatch contains an output for variadicBufferCounts assertFalse(recordBatch.getVariadicBufferCounts().isEmpty()); assertEquals(1, recordBatch.getVariadicBufferCounts().size()); VectorLoader vectorLoader = new VectorLoader(schemaRoot2); vectorLoader.load(recordBatch); IntVector vector3 = (IntVector) schemaRoot2.getVector(fieldName1); vector3.setValueCount(25); /* Check the vector output */ assertEquals(1, vector3.get(0)); assertEquals(2, vector3.get(1)); assertEquals(3, vector3.get(2)); assertEquals(4, vector3.get(3)); assertEquals(5, vector3.get(4)); assertEquals(6, vector3.get(5)); ViewVarCharVector vector4 = (ViewVarCharVector) schemaRoot2.getVector(fieldName2); vector4.setValueCount(25); /* Check the vector output */ assertArrayEquals(STR1, vector4.get(0)); assertArrayEquals(STR2, vector4.get(1)); assertArrayEquals(STR3, vector4.get(2)); assertArrayEquals(STR4, vector4.get(3)); assertArrayEquals(STR5, vector4.get(4)); assertArrayEquals(STR6, vector4.get(5)); } } } @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")); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestVectorAlloc.java000066400000000000000000000163621515451615200314630ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertTrue; import java.util.Arrays; import java.util.Collections; import org.apache.arrow.memory.AllocationListener; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.memory.rounding.DefaultRoundingPolicy; import org.apache.arrow.memory.rounding.RoundingPolicy; import org.apache.arrow.vector.types.TimeUnit; 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.Decimal; import org.apache.arrow.vector.types.pojo.ArrowType.Duration; import org.apache.arrow.vector.types.pojo.ArrowType.FixedSizeBinary; import org.apache.arrow.vector.types.pojo.ArrowType.Timestamp; 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; public class TestVectorAlloc { private BufferAllocator rootAllocator; private BufferAllocator policyAllocator; @BeforeEach public void init() { rootAllocator = new RootAllocator(Long.MAX_VALUE); policyAllocator = new RootAllocator(AllocationListener.NOOP, Integer.MAX_VALUE, new CustomPolicy()); } @AfterEach public void terminate() throws Exception { rootAllocator.close(); policyAllocator.close(); } private static Field field(String name, ArrowType type) { return new Field(name, new FieldType(true, type, null), Collections.emptyList()); } @Test public void testVectorAllocWithField() { Schema schema = new Schema( Arrays.asList( field("TINYINT", MinorType.TINYINT.getType()), field("SMALLINT", MinorType.SMALLINT.getType()), field("INT", MinorType.INT.getType()), field("BIGINT", MinorType.BIGINT.getType()), field("UINT1", MinorType.UINT1.getType()), field("UINT2", MinorType.UINT2.getType()), field("UINT4", MinorType.UINT4.getType()), field("UINT8", MinorType.UINT8.getType()), field("FLOAT4", MinorType.FLOAT4.getType()), field("FLOAT8", MinorType.FLOAT8.getType()), field("UTF8", MinorType.VARCHAR.getType()), field("VARBINARY", MinorType.VARBINARY.getType()), field("BIT", MinorType.BIT.getType()), field("DECIMAL", new Decimal(38, 5, 128)), field("FIXEDSIZEBINARY", new FixedSizeBinary(50)), field("DATEDAY", MinorType.DATEDAY.getType()), field("DATEMILLI", MinorType.DATEMILLI.getType()), field("TIMESEC", MinorType.TIMESEC.getType()), field("TIMEMILLI", MinorType.TIMEMILLI.getType()), field("TIMEMICRO", MinorType.TIMEMICRO.getType()), field("TIMENANO", MinorType.TIMENANO.getType()), field("TIMESTAMPSEC", MinorType.TIMESTAMPSEC.getType()), field("TIMESTAMPMILLI", MinorType.TIMESTAMPMILLI.getType()), field("TIMESTAMPMICRO", MinorType.TIMESTAMPMICRO.getType()), field("TIMESTAMPNANO", MinorType.TIMESTAMPNANO.getType()), field("TIMESTAMPSECTZ", new Timestamp(TimeUnit.SECOND, "PST")), field("TIMESTAMPMILLITZ", new Timestamp(TimeUnit.MILLISECOND, "PST")), field("TIMESTAMPMICROTZ", new Timestamp(TimeUnit.MICROSECOND, "PST")), field("TIMESTAMPNANOTZ", new Timestamp(TimeUnit.NANOSECOND, "PST")), field("INTERVALDAY", MinorType.INTERVALDAY.getType()), field("INTERVALYEAR", MinorType.INTERVALYEAR.getType()), field("DURATION", new Duration(TimeUnit.MILLISECOND)))); try (BufferAllocator allocator = rootAllocator.newChildAllocator("child", 0, Long.MAX_VALUE)) { for (Field field : schema.getFields()) { try (FieldVector vector = field.createVector(allocator)) { assertEquals( vector.getMinorType(), Types.getMinorTypeForArrowType(field.getFieldType().getType())); vector.allocateNew(); } } } } private static final int CUSTOM_SEGMENT_SIZE = 200; /** A custom rounding policy that rounds the size to the next multiple of 200. */ private static class CustomPolicy implements RoundingPolicy { @Override public long getRoundedSize(long requestSize) { return (requestSize + CUSTOM_SEGMENT_SIZE - 1) / CUSTOM_SEGMENT_SIZE * CUSTOM_SEGMENT_SIZE; } } @Test public void testFixedWidthVectorAllocation() { try (IntVector vec1 = new IntVector("vec", policyAllocator); IntVector vec2 = new IntVector("vec", rootAllocator)) { assertTrue(vec1.getAllocator().getRoundingPolicy() instanceof CustomPolicy); vec1.allocateNew(50); long totalCapacity = vec1.getValidityBuffer().capacity() + vec1.getDataBuffer().capacity(); // the total capacity must be a multiple of the segment size assertTrue(totalCapacity % CUSTOM_SEGMENT_SIZE == 0); assertTrue(vec2.getAllocator().getRoundingPolicy() instanceof DefaultRoundingPolicy); vec2.allocateNew(50); totalCapacity = vec2.getValidityBuffer().capacity() + vec2.getDataBuffer().capacity(); // the total capacity must be a power of two assertEquals(0, totalCapacity & (totalCapacity - 1)); } } @Test public void testVariableWidthVectorAllocation() { try (VarCharVector vec1 = new VarCharVector("vec", policyAllocator); VarCharVector vec2 = new VarCharVector("vec", rootAllocator)) { assertTrue(vec1.getAllocator().getRoundingPolicy() instanceof CustomPolicy); vec1.allocateNew(50); long totalCapacity = vec1.getValidityBuffer().capacity() + vec1.getOffsetBuffer().capacity(); // the total capacity must be a multiple of the segment size assertTrue(totalCapacity % CUSTOM_SEGMENT_SIZE == 0); assertTrue(vec2.getAllocator().getRoundingPolicy() instanceof DefaultRoundingPolicy); vec2.allocateNew(50); totalCapacity = vec2.getValidityBuffer().capacity() + vec2.getOffsetBuffer().capacity(); // the total capacity must be a power of two assertEquals(0, totalCapacity & (totalCapacity - 1)); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestVectorReAlloc.java000066400000000000000000000405631515451615200317520ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; 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; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.complex.impl.NullableStructWriter; import org.apache.arrow.vector.complex.impl.UnionFixedSizeListWriter; import org.apache.arrow.vector.complex.impl.UnionListWriter; import org.apache.arrow.vector.holders.NullableIntHolder; 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.DataSizeRoundingUtil; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestVectorReAlloc { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testFixedType() { try (final UInt4Vector vector = new UInt4Vector("", allocator)) { vector.setInitialCapacity(512); vector.allocateNew(); assertTrue(vector.getValueCapacity() >= 512); int initialCapacity = vector.getValueCapacity(); try { vector.set(initialCapacity, 0); fail("Expected out of bounds exception"); } catch (Exception e) { // ok } vector.reAlloc(); assertTrue(vector.getValueCapacity() >= 2 * initialCapacity); vector.set(initialCapacity, 100); assertEquals(100, vector.get(initialCapacity)); } } @Test public void testNullableType() { try (final VarCharVector vector = new VarCharVector("", allocator)) { vector.setInitialCapacity(512); vector.allocateNew(); assertTrue(vector.getValueCapacity() >= 512); int initialCapacity = vector.getValueCapacity(); try { vector.set(initialCapacity, "foo".getBytes(StandardCharsets.UTF_8)); fail("Expected out of bounds exception"); } catch (Exception e) { // ok } vector.reAlloc(); assertTrue(vector.getValueCapacity() >= 2 * initialCapacity); vector.set(initialCapacity, "foo".getBytes(StandardCharsets.UTF_8)); assertEquals(new String(vector.get(initialCapacity), StandardCharsets.UTF_8), "foo"); } } @Test public void testListType() { try (final ListVector vector = ListVector.empty("", allocator)) { vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); vector.setInitialCapacity(512); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); try { vector.getInnerValueCountAt(2014); fail("Expected out of bounds exception"); } catch (Exception e) { // ok } vector.reAlloc(); assertEquals(1024, vector.getValueCapacity()); assertEquals(0, vector.getOffsetBuffer().getInt(2014 * ListVector.OFFSET_WIDTH)); } } @Test public void testStructType() { try (final StructVector vector = StructVector.empty("", allocator)) { vector.addOrGet("", FieldType.nullable(MinorType.INT.getType()), IntVector.class); vector.setInitialCapacity(512); vector.allocateNew(); assertEquals(512, vector.getValueCapacity()); try { vector.getObject(513); fail("Expected out of bounds exception"); } catch (Exception e) { // ok } vector.reAlloc(); assertEquals(1024, vector.getValueCapacity()); assertNull(vector.getObject(513)); } } @Test public void testVariableWidthTypeSetNullValues() { // Test ARROW-11223 bug is fixed try (final BaseVariableWidthVector v1 = new VarCharVector("var1", allocator)) { v1.setInitialCapacity(512); v1.allocateNew(); int numNullValues1 = v1.getValueCapacity() + 1; for (int i = 0; i < numNullValues1; i++) { v1.setNull(i); } assertTrue(v1.getBufferSizeFor(numNullValues1) > 0); } try (final BaseLargeVariableWidthVector v2 = new LargeVarCharVector("var2", allocator)) { v2.setInitialCapacity(512); v2.allocateNew(); int numNullValues2 = v2.getValueCapacity() + 1; for (int i = 0; i < numNullValues2; i++) { v2.setNull(i); } assertTrue(v2.getBufferSizeFor(numNullValues2) > 0); } } @Test public void testFixedAllocateAfterReAlloc() throws Exception { try (final IntVector vector = new IntVector("", allocator)) { /* * Allocate the default size, and then, reAlloc. This should double the allocation. */ vector.allocateNewSafe(); // Initial allocation vector.reAlloc(); // Double the allocation size. int savedValueCapacity = vector.getValueCapacity(); /* * Clear and allocate again. */ vector.clear(); vector.allocateNewSafe(); /* * Verify that the buffer sizes haven't changed. */ assertEquals(vector.getValueCapacity(), savedValueCapacity); } } @Test public void testVariableAllocateAfterReAlloc() throws Exception { try (final VarCharVector vector = new VarCharVector("", allocator)) { /* * Allocate the default size, and then, reAlloc. This should double the allocation. */ vector.allocateNewSafe(); // Initial allocation vector.reAlloc(); // Double the allocation size. int savedValueCapacity = vector.getValueCapacity(); long savedValueBufferSize = vector.valueBuffer.capacity(); /* * Clear and allocate again. */ vector.clear(); vector.allocateNewSafe(); /* * Verify that the buffer sizes haven't changed. */ assertEquals(vector.getValueCapacity(), savedValueCapacity); assertEquals(vector.valueBuffer.capacity(), savedValueBufferSize); } } @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)) { /* * Allocate the default size, and then, reAlloc. This should double the allocation. */ vector.allocateNewSafe(); // Initial allocation vector.reAlloc(); // Double the allocation size. int savedValueCapacity = vector.getValueCapacity(); long savedValueBufferSize = vector.valueBuffer.capacity(); /* * Clear and allocate again. */ vector.clear(); vector.allocateNewSafe(); /* * Verify that the buffer sizes haven't changed. */ assertEquals(vector.getValueCapacity(), savedValueCapacity); assertEquals(vector.valueBuffer.capacity(), savedValueBufferSize); } } @Test public void testVarCharAllocateNew() throws Exception { final int count = 6000; try (final VarCharVector vector = new VarCharVector("", allocator)) { vector.allocateNew(count); // verify that the validity buffer and value buffer have capacity for at least 'count' // elements. assertTrue( vector.getValidityBuffer().capacity() >= DataSizeRoundingUtil.divideBy8Ceil(count)); assertTrue( vector.getOffsetBuffer().capacity() >= (count + 1) * BaseVariableWidthVector.OFFSET_WIDTH); } } @Test public void testLargeVarCharAllocateNew() throws Exception { final int count = 6000; try (final LargeVarCharVector vector = new LargeVarCharVector("", allocator)) { vector.allocateNew(count); // verify that the validity buffer and value buffer have capacity for at least 'count' // elements. assertTrue( vector.getValidityBuffer().capacity() >= DataSizeRoundingUtil.divideBy8Ceil(count)); assertTrue( vector.getOffsetBuffer().capacity() >= (count + 1) * BaseLargeVariableWidthVector.OFFSET_WIDTH); } } @Test public void testVarCharAllocateNewUsingHelper() throws Exception { final int count = 6000; try (final VarCharVector vector = new VarCharVector("", allocator)) { AllocationHelper.allocateNew(vector, count); // verify that the validity buffer and value buffer have capacity for at least 'count' // elements. assertTrue( vector.getValidityBuffer().capacity() >= DataSizeRoundingUtil.divideBy8Ceil(count)); assertTrue( vector.getOffsetBuffer().capacity() >= (count + 1) * BaseVariableWidthVector.OFFSET_WIDTH); } } @Test public void testLargeVarCharAllocateNewUsingHelper() throws Exception { final int count = 6000; try (final LargeVarCharVector vector = new LargeVarCharVector("", allocator)) { AllocationHelper.allocateNew(vector, count); // verify that the validity buffer and value buffer have capacity for at least 'count' // elements. assertTrue( vector.getValidityBuffer().capacity() >= DataSizeRoundingUtil.divideBy8Ceil(count)); assertTrue( vector.getOffsetBuffer().capacity() >= (count + 1) * BaseLargeVariableWidthVector.OFFSET_WIDTH); } } @Test public void testFixedRepeatedClearAndSet() throws Exception { try (final IntVector vector = new IntVector("", allocator)) { vector.allocateNewSafe(); // Initial allocation vector.clear(); // clear vector. vector.setSafe(0, 10); int savedValueCapacity = vector.getValueCapacity(); for (int i = 0; i < 1024; ++i) { vector.clear(); // clear vector. vector.setSafe(0, 10); } // should be deterministic, and not cause a run-away increase in capacity. assertEquals(vector.getValueCapacity(), savedValueCapacity); } } @Test public void testVariableRepeatedClearAndSet() throws Exception { try (final VarCharVector vector = new VarCharVector("", allocator)) { vector.allocateNewSafe(); // Initial allocation vector.clear(); // clear vector. vector.setSafe(0, "hello world".getBytes(StandardCharsets.UTF_8)); int savedValueCapacity = vector.getValueCapacity(); for (int i = 0; i < 1024; ++i) { vector.clear(); // clear vector. vector.setSafe(0, "hello world".getBytes(StandardCharsets.UTF_8)); } // should be deterministic, and not cause a run-away increase in capacity. assertEquals(vector.getValueCapacity(), savedValueCapacity); } } @Test public void testRepeatedValueVectorClearAndSet() throws Exception { try (final ListVector vector = new ListVector("", allocator, FieldType.nullable(MinorType.INT.getType()), null)) { vector.allocateNewSafe(); // Initial allocation UnionListWriter writer = vector.getWriter(); vector.clear(); // clear vector. writer.setPosition(0); // optional writer.startList(); writer.writeInt(0); writer.endList(); int savedValueCapacity = vector.getValueCapacity(); for (int i = 0; i < 1024; ++i) { vector.clear(); // clear vector. writer.setPosition(0); // optional writer.startList(); writer.writeInt(i); writer.endList(); } // should be deterministic, and not cause a run-away increase in capacity. assertEquals(vector.getValueCapacity(), savedValueCapacity); } } @Test public void testStructVectorClearAndSet() throws Exception { try (final StructVector vector = StructVector.empty("v", allocator)) { vector.allocateNewSafe(); // Initial allocation NullableStructWriter writer = vector.getWriter(); vector.clear(); // clear vector. writer.setPosition(0); // optional writer.start(); writer.integer("int").writeInt(0); writer.end(); int savedValueCapacity = vector.getValueCapacity(); for (int i = 0; i < 1024; ++i) { vector.clear(); // clear vector. writer.setPosition(0); // optional writer.start(); writer.integer("int").writeInt(i); writer.end(); } // should be deterministic, and not cause a run-away increase in capacity. assertEquals(vector.getValueCapacity(), savedValueCapacity); } } @Test public void testFixedSizeListVectorClearAndSet() { try (final FixedSizeListVector vector = new FixedSizeListVector( "", allocator, FieldType.nullable(new ArrowType.FixedSizeList(2)), null)) { vector.allocateNewSafe(); // Initial allocation UnionFixedSizeListWriter writer = vector.getWriter(); vector.clear(); // clear vector. writer.setPosition(0); // optional writer.startList(); writer.writeInt(0); writer.writeInt(1); writer.endList(); int savedValueCapacity = vector.getValueCapacity(); for (int i = 0; i < 1024; ++i) { vector.clear(); // clear vector. writer.setPosition(0); // optional writer.startList(); writer.writeInt(i); writer.writeInt(i + 1); writer.endList(); } // should be deterministic, and not cause a run-away increase in capacity. assertEquals(vector.getValueCapacity(), savedValueCapacity); } } @Test public void testUnionVectorClearAndSet() { try (final UnionVector vector = new UnionVector("", allocator, /* field type */ null, /* call-back */ null)) { vector.allocateNewSafe(); // Initial allocation NullableIntHolder holder = new NullableIntHolder(); holder.isSet = 1; holder.value = 1; vector.clear(); // clear vector. vector.setType(0, MinorType.INT); vector.setSafe(0, holder); int savedValueCapacity = vector.getValueCapacity(); for (int i = 0; i < 1024; ++i) { vector.clear(); // clear vector. vector.setType(0, MinorType.INT); vector.setSafe(0, holder); } // should be deterministic, and not cause a run-away increase in capacity. assertEquals(vector.getValueCapacity(), savedValueCapacity); } } @Test public void testDenseUnionVectorClearAndSet() { try (final DenseUnionVector vector = new DenseUnionVector("", allocator, null, null)) { vector.allocateNewSafe(); // Initial allocation NullableIntHolder holder = new NullableIntHolder(); holder.isSet = 1; holder.value = 1; byte intTypeId = vector.registerNewTypeId(Field.nullable("", MinorType.INT.getType())); vector.clear(); vector.setTypeId(0, intTypeId); vector.setSafe(0, holder); int savedValueCapacity = vector.getValueCapacity(); for (int i = 0; i < 1024; ++i) { vector.clear(); vector.setTypeId(0, intTypeId); vector.setSafe(0, holder); } // should be deterministic, and not cause a run-away increase in capacity. assertEquals(vector.getValueCapacity(), savedValueCapacity); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestVectorReset.java000066400000000000000000000155551515451615200315160ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertTrue; import java.nio.charset.StandardCharsets; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.ListViewVector; import org.apache.arrow.vector.complex.NonNullableStructVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType.FixedSizeList; import org.apache.arrow.vector.types.pojo.ArrowType.Int; import org.apache.arrow.vector.types.pojo.FieldType; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestVectorReset { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void terminate() throws Exception { allocator.close(); } private void resetVectorAndVerify(ValueVector vector, ArrowBuf[] bufs) { long[] sizeBefore = new long[bufs.length]; for (int i = 0; i < bufs.length; i++) { sizeBefore[i] = bufs[i].capacity(); } vector.reset(); for (int i = 0; i < bufs.length; i++) { assertEquals(sizeBefore[i], bufs[i].capacity()); verifyBufferZeroed(bufs[i]); } assertEquals(0, vector.getValueCount()); } private void verifyBufferZeroed(ArrowBuf buf) { for (int i = 0; i < buf.capacity(); i++) { assertTrue((byte) 0 == buf.getByte(i)); } } @Test public void testFixedTypeReset() { try (final UInt4Vector vector = new UInt4Vector("UInt4", allocator)) { vector.allocateNewSafe(); vector.setNull(0); vector.setValueCount(1); resetVectorAndVerify(vector, vector.getBuffers(false)); } } @Test public void testVariableTypeReset() { try (final VarCharVector vector = new VarCharVector("VarChar", allocator)) { vector.allocateNewSafe(); vector.set(0, "a".getBytes(StandardCharsets.UTF_8)); vector.setLastSet(0); vector.setValueCount(1); resetVectorAndVerify(vector, vector.getBuffers(false)); assertEquals(-1, vector.getLastSet()); } } @Test public void testVariableViewTypeReset() { try (final ViewVarCharVector vector = new ViewVarCharVector("ViewVarChar", allocator)) { vector.allocateNewSafe(); vector.set(0, "a".getBytes(StandardCharsets.UTF_8)); vector.setLastSet(0); vector.setValueCount(1); resetVectorAndVerify(vector, vector.getBuffers(false)); assertEquals(-1, vector.getLastSet()); } } @Test public void testLargeVariableTypeReset() { try (final LargeVarCharVector vector = new LargeVarCharVector("LargeVarChar", allocator)) { vector.allocateNewSafe(); vector.set(0, "a".getBytes(StandardCharsets.UTF_8)); vector.setLastSet(0); vector.setValueCount(1); resetVectorAndVerify(vector, vector.getBuffers(false)); assertEquals(-1, vector.getLastSet()); } } @Test public void testListTypeReset() { try (final ListVector variableList = new ListVector( "VarList", allocator, FieldType.nullable(MinorType.INT.getType()), null); final FixedSizeListVector fixedList = new FixedSizeListVector( "FixedList", allocator, FieldType.nullable(new FixedSizeList(2)), null); final ListViewVector variableViewList = new ListViewVector( "VarListView", allocator, FieldType.nullable(MinorType.INT.getType()), null)) { // ListVector variableList.allocateNewSafe(); variableList.startNewValue(0); variableList.endValue(0, 0); variableList.setValueCount(1); resetVectorAndVerify(variableList, variableList.getBuffers(false)); assertEquals(-1, variableList.getLastSet()); // FixedSizeListVector fixedList.allocateNewSafe(); fixedList.setNull(0); fixedList.setValueCount(1); resetVectorAndVerify(fixedList, fixedList.getBuffers(false)); // ListViewVector variableViewList.allocateNewSafe(); variableViewList.startNewValue(0); variableViewList.endValue(0, 0); variableViewList.setValueCount(1); resetVectorAndVerify(variableViewList, variableViewList.getBuffers(false)); } } @Test public void testStructTypeReset() { try (final NonNullableStructVector nonNullableStructVector = new NonNullableStructVector( "Struct", allocator, FieldType.nullable(MinorType.INT.getType()), null); final StructVector structVector = new StructVector( "NullableStruct", allocator, FieldType.nullable(MinorType.INT.getType()), null)) { // NonNullableStructVector nonNullableStructVector.allocateNewSafe(); IntVector structChild = nonNullableStructVector.addOrGet( "child", FieldType.nullable(new Int(32, true)), IntVector.class); structChild.setNull(0); nonNullableStructVector.setValueCount(1); resetVectorAndVerify(nonNullableStructVector, nonNullableStructVector.getBuffers(false)); // StructVector structVector.allocateNewSafe(); structVector.setNull(0); structVector.setValueCount(1); resetVectorAndVerify(structVector, structVector.getBuffers(false)); } } @Test public void testUnionTypeReset() { try (final UnionVector vector = new UnionVector("Union", allocator, /* field type */ null, /* call-back */ null); final IntVector dataVector = new IntVector("Int", allocator)) { vector.getBufferSize(); vector.allocateNewSafe(); dataVector.allocateNewSafe(); vector.addVector(dataVector); dataVector.setNull(0); vector.setValueCount(1); resetVectorAndVerify(vector, vector.getBuffers(false)); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java000066400000000000000000000273711515451615200324770ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.assertNotEquals; 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.List; import java.util.stream.Collectors; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.impl.UnionListWriter; 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.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestVectorSchemaRoot { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void terminate() { allocator.close(); } @Test public void testResetRowCount() { final int size = 20; try (final BitVector vec1 = new BitVector("bit", allocator); final IntVector vec2 = new IntVector("int", allocator)) { VectorSchemaRoot vsr = VectorSchemaRoot.of(vec1, vec2); vsr.allocateNew(); assertEquals(0, vsr.getRowCount()); for (int i = 0; i < size; i++) { vec1.setSafe(i, i % 2); vec2.setSafe(i, i); } vsr.setRowCount(size); checkCount(vec1, vec2, vsr, size); vsr.allocateNew(); checkCount(vec1, vec2, vsr, 0); for (int i = 0; i < size; i++) { vec1.setSafe(i, i % 2); vec2.setSafe(i, i); } vsr.setRowCount(size); checkCount(vec1, vec2, vsr, size); vsr.clear(); checkCount(vec1, vec2, vsr, 0); } } private void checkCount(BitVector vec1, IntVector vec2, VectorSchemaRoot vsr, int count) { assertEquals(vec1.getValueCount(), count); assertEquals(vec2.getValueCount(), count); assertEquals(vsr.getRowCount(), count); } private VectorSchemaRoot createBatch() { 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<>(); Field childField = new Field("varCharCol", varCharType, null); List childFields = new ArrayList<>(); childFields.add(childField); schemaFields.add(new Field("listCol", listType, childFields)); Schema schema = new Schema(schemaFields); VectorSchemaRoot schemaRoot = VectorSchemaRoot.create(schema, allocator); // get and allocate the vector ListVector vector = (ListVector) schemaRoot.getVector("listCol"); vector.allocateNew(); // write data to the vector UnionListWriter writer = vector.getWriter(); writer.setPosition(0); // write data vector(0) writer.startList(); // write data vector(0)(0) writer.list().startList(); // According to the schema above, the list element should have varchar type. // When we write a big int, the original writer cannot handle this, so the writer will // be promoted, and the vector structure will be different from the schema. writer.list().bigInt().writeBigInt(0); writer.list().bigInt().writeBigInt(1); writer.list().endList(); // write data vector(0)(1) writer.list().startList(); writer.list().float8().writeFloat8(3.0D); writer.list().float8().writeFloat8(7.0D); writer.list().endList(); // finish data vector(0) writer.endList(); writer.setPosition(1); // write data vector(1) writer.startList(); // write data vector(1)(0) writer.list().startList(); writer.list().integer().writeInt(3); writer.list().integer().writeInt(2); writer.list().endList(); // finish data vector(1) writer.endList(); vector.setValueCount(2); return schemaRoot; } @Test public void testAddVector() { 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(1, intVector3); assertEquals(3, newRecordBatch.getFieldVectors().size()); assertEquals(intVector3, newRecordBatch.getFieldVectors().get(1)); original.close(); newRecordBatch.close(); } } @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); final IntVector intVector2 = new IntVector("intVector2", allocator); final IntVector intVector3 = new IntVector("intVector3", allocator); ) { VectorSchemaRoot original = new VectorSchemaRoot(Arrays.asList(intVector1, intVector2, intVector3)); assertEquals(3, original.getFieldVectors().size()); VectorSchemaRoot newRecordBatch = original.removeVector(0); assertEquals(2, newRecordBatch.getFieldVectors().size()); assertEquals(intVector2, newRecordBatch.getFieldVectors().get(0)); assertEquals(intVector3, newRecordBatch.getFieldVectors().get(1)); original.close(); newRecordBatch.close(); } } @Test public void testSlice() { try (final IntVector intVector = new IntVector("intVector", allocator); final Float4Vector float4Vector = new Float4Vector("float4Vector", allocator)) { final int numRows = 10; intVector.setValueCount(numRows); float4Vector.setValueCount(numRows); for (int i = 0; i < numRows; i++) { intVector.setSafe(i, i); float4Vector.setSafe(i, i + 0.1f); } final VectorSchemaRoot original = new VectorSchemaRoot(Arrays.asList(intVector, float4Vector)); for (int sliceIndex = 0; sliceIndex < numRows; sliceIndex++) { for (int sliceLength = 0; sliceIndex + sliceLength <= numRows; sliceLength++) { try (VectorSchemaRoot slice = original.slice(sliceIndex, sliceLength)) { assertEquals(sliceLength, slice.getRowCount()); // validate data final IntVector childIntVector = (IntVector) slice.getFieldVectors().get(0); final Float4Vector childFloatVector = (Float4Vector) slice.getFieldVectors().get(1); for (int i = 0; i < sliceLength; i++) { final int originalIndex = i + sliceIndex; assertEquals(originalIndex, childIntVector.get(i)); assertEquals(originalIndex + 0.1f, childFloatVector.get(i), 0); } } } } original.close(); } } @Test public void testSliceWithInvalidParam() { assertThrows( IllegalArgumentException.class, () -> { try (final IntVector intVector = new IntVector("intVector", allocator); final Float4Vector float4Vector = new Float4Vector("float4Vector", allocator)) { intVector.setValueCount(10); float4Vector.setValueCount(10); for (int i = 0; i < 10; i++) { intVector.setSafe(i, i); float4Vector.setSafe(i, i + 0.1f); } final VectorSchemaRoot original = new VectorSchemaRoot(Arrays.asList(intVector, float4Vector)); original.slice(0, 20); } }); } @Test public void testEquals() { try (final IntVector intVector1 = new IntVector("intVector1", allocator); final IntVector intVector2 = new IntVector("intVector2", allocator); final IntVector intVector3 = new IntVector("intVector3", allocator); ) { intVector1.setValueCount(5); for (int i = 0; i < 5; i++) { intVector1.set(i, i); } VectorSchemaRoot root1 = new VectorSchemaRoot(Arrays.asList(intVector1, intVector2, intVector3)); VectorSchemaRoot root2 = new VectorSchemaRoot(Arrays.asList(intVector1, intVector2)); VectorSchemaRoot root3 = new VectorSchemaRoot(Arrays.asList(intVector1, intVector2, intVector3)); assertFalse(root1.equals(root2)); assertTrue(root1.equals(root3)); root1.close(); root2.close(); root3.close(); } } @Test public void testApproxEquals() { try (final Float4Vector float4Vector1 = new Float4Vector("floatVector", allocator); final Float4Vector float4Vector2 = new Float4Vector("floatVector", allocator); final Float4Vector float4Vector3 = new Float4Vector("floatVector", allocator); ) { float4Vector1.setValueCount(5); float4Vector2.setValueCount(5); float4Vector3.setValueCount(5); final float epsilon = 1.0E-6f; for (int i = 0; i < 5; i++) { float4Vector1.set(i, i); float4Vector2.set(i, i + epsilon * 2); float4Vector3.set(i, i + epsilon / 2); } VectorSchemaRoot root1 = new VectorSchemaRoot(Arrays.asList(float4Vector1)); VectorSchemaRoot root2 = new VectorSchemaRoot(Arrays.asList(float4Vector2)); VectorSchemaRoot root3 = new VectorSchemaRoot(Arrays.asList(float4Vector3)); assertFalse(root1.approxEquals(root2)); assertTrue(root1.approxEquals(root3)); root1.close(); root2.close(); root3.close(); } } @Test public void testSchemaSync() { // create vector schema root try (VectorSchemaRoot schemaRoot = createBatch()) { Schema newSchema = new Schema( schemaRoot.getFieldVectors().stream() .map(vec -> vec.getField()) .collect(Collectors.toList())); assertNotEquals(newSchema, schemaRoot.getSchema()); assertTrue(schemaRoot.syncSchema()); assertEquals(newSchema, schemaRoot.getSchema()); // no schema update this time. assertFalse(schemaRoot.syncSchema()); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadLoad.java000066400000000000000000000326771515451615200324620ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 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; import java.io.IOException; import java.util.ArrayList; import java.util.Collections; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.complex.NonNullableStructVector; import org.apache.arrow.vector.complex.impl.ComplexWriterImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.complex.writer.BaseWriter.ComplexWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.ListWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter; import org.apache.arrow.vector.complex.writer.BigIntWriter; import org.apache.arrow.vector.complex.writer.IntWriter; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; 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.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; public class TestVectorUnloadLoad { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testNullCodec() { final Schema schema = new Schema(Collections.emptyList()); try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { root.setRowCount(1); final VectorUnloader unloader = new VectorUnloader(root, /*includeNulls*/ true, /*codec*/ null, /*alignBuffers*/ true); unloader.getRecordBatch().close(); } } @Test public void testUnloadLoad() throws IOException { int count = 10000; Schema schema; try (BufferAllocator originalVectorsAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); NonNullableStructVector parent = NonNullableStructVector.empty("parent", originalVectorsAllocator)) { // write some data ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); IntWriter intWriter = rootWriter.integer("int"); BigIntWriter bigIntWriter = rootWriter.bigInt("bigInt"); for (int i = 0; i < count; i++) { intWriter.setPosition(i); intWriter.writeInt(i); bigIntWriter.setPosition(i); bigIntWriter.writeBigInt(i); } writer.setValueCount(count); // unload it FieldVector root = parent.getChild("root"); schema = new Schema(root.getField().getChildren()); VectorUnloader vectorUnloader = newVectorUnloader(root); try (ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch(); BufferAllocator finalVectorsAllocator = allocator.newChildAllocator("final vectors", 0, Integer.MAX_VALUE); VectorSchemaRoot newRoot = VectorSchemaRoot.create(schema, finalVectorsAllocator); ) { // load it VectorLoader vectorLoader = new VectorLoader(newRoot); vectorLoader.load(recordBatch); FieldReader intReader = newRoot.getVector("int").getReader(); FieldReader bigIntReader = newRoot.getVector("bigInt").getReader(); for (int i = 0; i < count; i++) { intReader.setPosition(i); assertEquals(i, intReader.readInteger().intValue()); bigIntReader.setPosition(i); assertEquals(i, bigIntReader.readLong().longValue()); } } } } @Test public void testUnloadLoadAddPadding() throws IOException { int count = 10000; Schema schema; try (BufferAllocator originalVectorsAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); NonNullableStructVector parent = NonNullableStructVector.empty("parent", originalVectorsAllocator)) { // write some data ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); ListWriter list = rootWriter.list("list"); IntWriter intWriter = list.integer(); for (int i = 0; i < count; i++) { list.setPosition(i); list.startList(); for (int j = 0; j < i % 4 + 1; j++) { intWriter.writeInt(i); } list.endList(); } writer.setValueCount(count); // unload it FieldVector root = parent.getChild("root"); schema = new Schema(root.getField().getChildren()); VectorUnloader vectorUnloader = newVectorUnloader(root); try (ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch(); BufferAllocator finalVectorsAllocator = allocator.newChildAllocator("final vectors", 0, Integer.MAX_VALUE); VectorSchemaRoot newRoot = VectorSchemaRoot.create(schema, finalVectorsAllocator); ) { List oldBuffers = recordBatch.getBuffers(); List newBuffers = new ArrayList<>(); for (ArrowBuf oldBuffer : oldBuffers) { long l = oldBuffer.readableBytes(); if (l % 64 != 0) { // pad l = l + 64 - l % 64; } ArrowBuf newBuffer = allocator.buffer(l); for (long i = oldBuffer.readerIndex(); i < oldBuffer.writerIndex(); i++) { newBuffer.setByte(i - oldBuffer.readerIndex(), oldBuffer.getByte(i)); } newBuffer.readerIndex(0); newBuffer.writerIndex(l); newBuffers.add(newBuffer); } try (ArrowRecordBatch newBatch = new ArrowRecordBatch(recordBatch.getLength(), recordBatch.getNodes(), newBuffers); ) { // load it VectorLoader vectorLoader = new VectorLoader(newRoot); vectorLoader.load(newBatch); FieldReader reader = newRoot.getVector("list").getReader(); for (int i = 0; i < count; i++) { reader.setPosition(i); List expected = new ArrayList<>(); for (int j = 0; j < i % 4 + 1; j++) { expected.add(i); } assertEquals(expected, reader.readObject()); } } for (ArrowBuf newBuf : newBuffers) { newBuf.getReferenceManager().release(); } } } } /** * The validity buffer can be empty if: - all values are defined. - all values are null. * * @throws IOException on error */ @Test public void testLoadValidityBuffer() throws IOException { Schema schema = new Schema( asList( new Field( "intDefined", FieldType.nullable(new ArrowType.Int(32, true)), Collections.emptyList()), new Field( "intNull", FieldType.nullable(new ArrowType.Int(32, true)), Collections.emptyList()))); int count = 10; ArrowBuf[] values = new ArrowBuf[4]; for (int i = 0; i < 4; i += 2) { ArrowBuf buf1 = allocator.buffer(getValidityBufferSizeFromCount(count)); ArrowBuf buf2 = allocator.buffer(count * 4); // integers buf1.setZero(0, buf1.capacity()); buf2.setZero(0, buf2.capacity()); values[i] = buf1; values[i + 1] = buf2; for (int j = 0; j < count; j++) { if (i == 2) { BitVectorHelper.unsetBit(buf1, j); } else { BitVectorHelper.setBit(buf1, j); } buf2.setInt(j * 4, j); } buf1.writerIndex((int) Math.ceil(count / 8)); buf2.writerIndex(count * 4); } /* * values[0] - validity buffer for first vector * values[1] - data buffer for first vector * values[2] - validity buffer for second vector * values[3] - data buffer for second vector */ try (ArrowRecordBatch recordBatch = new ArrowRecordBatch( count, asList(new ArrowFieldNode(count, 0), new ArrowFieldNode(count, count)), asList(values[0], values[1], values[2], values[3])); BufferAllocator finalVectorsAllocator = allocator.newChildAllocator("final vectors", 0, Integer.MAX_VALUE); VectorSchemaRoot newRoot = VectorSchemaRoot.create(schema, finalVectorsAllocator); ) { // load it VectorLoader vectorLoader = new VectorLoader(newRoot); vectorLoader.load(recordBatch); IntVector intDefinedVector = (IntVector) newRoot.getVector("intDefined"); IntVector intNullVector = (IntVector) newRoot.getVector("intNull"); for (int i = 0; i < count; i++) { assertFalse(intDefinedVector.isNull(i), "#" + i); assertEquals(i, intDefinedVector.get(i), "#" + i); assertTrue(intNullVector.isNull(i), "#" + i); } intDefinedVector.setSafe(count + 10, 1234); assertTrue(intDefinedVector.isNull(count + 1)); // empty slots should still default to unset intDefinedVector.setSafe(count + 1, 789); assertFalse(intDefinedVector.isNull(count + 1)); assertEquals(789, intDefinedVector.get(count + 1)); assertTrue(intDefinedVector.isNull(count)); assertTrue(intDefinedVector.isNull(count + 2)); assertTrue(intDefinedVector.isNull(count + 3)); assertTrue(intDefinedVector.isNull(count + 4)); assertTrue(intDefinedVector.isNull(count + 5)); assertTrue(intDefinedVector.isNull(count + 6)); assertTrue(intDefinedVector.isNull(count + 7)); assertTrue(intDefinedVector.isNull(count + 8)); assertTrue(intDefinedVector.isNull(count + 9)); assertFalse(intDefinedVector.isNull(count + 10)); assertEquals(1234, intDefinedVector.get(count + 10)); } finally { for (ArrowBuf arrowBuf : values) { arrowBuf.getReferenceManager().release(); } } } @Test public void testUnloadLoadDuplicates() throws IOException { int count = 10; Schema schema = new Schema( asList( new Field( "duplicate", FieldType.nullable(new ArrowType.Int(32, true)), Collections.emptyList()), new Field( "duplicate", FieldType.nullable(new ArrowType.Int(32, true)), Collections.emptyList()))); try (BufferAllocator originalVectorsAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); ) { List sources = new ArrayList<>(); for (Field field : schema.getFields()) { FieldVector vector = field.createVector(originalVectorsAllocator); vector.allocateNew(); sources.add(vector); IntVector intVector = (IntVector) vector; for (int i = 0; i < count; i++) { intVector.set(i, i); } intVector.setValueCount(count); } try (VectorSchemaRoot root = new VectorSchemaRoot(schema.getFields(), sources, count)) { VectorUnloader vectorUnloader = new VectorUnloader(root); try (ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch(); BufferAllocator finalVectorsAllocator = allocator.newChildAllocator("final vectors", 0, Integer.MAX_VALUE); VectorSchemaRoot newRoot = VectorSchemaRoot.create(schema, finalVectorsAllocator); ) { // load it VectorLoader vectorLoader = new VectorLoader(newRoot); vectorLoader.load(recordBatch); List targets = newRoot.getFieldVectors(); assertEquals(sources.size(), targets.size()); for (int k = 0; k < sources.size(); k++) { IntVector src = (IntVector) sources.get(k); IntVector tgt = (IntVector) targets.get(k); assertEquals(src.getValueCount(), tgt.getValueCount()); for (int i = 0; i < count; i++) { assertEquals(src.get(i), tgt.get(i)); } } } } } } public static VectorUnloader newVectorUnloader(FieldVector root) { Schema schema = new Schema(root.getField().getChildren()); int valueCount = root.getValueCount(); List fields = root.getChildrenFromFields(); VectorSchemaRoot vsr = new VectorSchemaRoot(schema.getFields(), fields, valueCount); return new VectorUnloader(vsr); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/compare/000077500000000000000000000000001515451615200271615ustar00rootroot00000000000000TestRangeEqualsVisitor.java000066400000000000000000001252351515451615200344040ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/compare/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compare; import static org.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; 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; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.LargeVarCharVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.ZeroVector; import org.apache.arrow.vector.compare.util.ValueEpsilonEqualizers; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.FixedSizeListVector; 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; import org.apache.arrow.vector.complex.impl.UnionFixedSizeListWriter; import org.apache.arrow.vector.complex.impl.UnionLargeListViewWriter; import org.apache.arrow.vector.complex.impl.UnionListViewWriter; import org.apache.arrow.vector.complex.impl.UnionListWriter; import org.apache.arrow.vector.holders.NullableBigIntHolder; import org.apache.arrow.vector.holders.NullableFloat4Holder; import org.apache.arrow.vector.holders.NullableFloat8Holder; import org.apache.arrow.vector.holders.NullableIntHolder; 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; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; public class TestRangeEqualsVisitor { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } private static final Charset utf8Charset = Charset.forName("UTF-8"); private static final byte[] STR1 = "AAAAA1".getBytes(utf8Charset); private static final byte[] STR2 = "BBBBBBBBB2".getBytes(utf8Charset); private static final byte[] STR3 = "CCCC3".getBytes(utf8Charset); private static final byte[] STR4 = "12345678901234A".getBytes(utf8Charset); private static final byte[] STR5 = "A2345678901234ABC".getBytes(utf8Charset); private static final byte[] STR6 = "AB45678901234ABCD".getBytes(utf8Charset); @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testIntVectorEqualsWithNull() { try (final IntVector vector1 = new IntVector("int", allocator); final IntVector vector2 = new IntVector("int", allocator)) { setVector(vector1, 1, 2); setVector(vector2, 1, null); assertFalse(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testEqualsWithTypeChange() { try (final IntVector vector1 = new IntVector("vector", allocator); final IntVector vector2 = new IntVector("vector", allocator); final BigIntVector vector3 = new BigIntVector("vector", allocator)) { setVector(vector1, 1, 2); setVector(vector2, 1, 2); RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector1, vector2); Range range = new Range(0, 0, 2); assertTrue(vector1.accept(visitor, range)); // visitor left vector changed, will reset and check type again assertFalse(vector3.accept(visitor, range)); } } @Test public void testBaseFixedWidthVectorRangeEqual() { try (final IntVector vector1 = new IntVector("int", allocator); final IntVector vector2 = new IntVector("int", allocator)) { setVector(vector1, 1, 2, 3, 4, 5); setVector(vector2, 11, 2, 3, 4, 55); RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector1, vector2); assertTrue(visitor.rangeEquals(new Range(1, 1, 3))); } } @Test public void testBaseVariableVectorRangeEquals() { try (final VarCharVector vector1 = new VarCharVector("varchar", allocator); final VarCharVector vector2 = new VarCharVector("varchar", allocator)) { setVector(vector1, STR1, STR2, STR3, STR2, STR1); setVector(vector2, STR1, STR2, STR3, STR2, STR1); RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector1, vector2); assertTrue(visitor.rangeEquals(new Range(1, 1, 3))); } } @Test public void testBaseVariableViewVectorRangeEquals() { try (final ViewVarCharVector vector1 = new ViewVarCharVector("varchar", allocator); final ViewVarCharVector vector2 = new ViewVarCharVector("varchar", allocator)) { setVector(vector1, STR1, STR2, STR4, STR3, STR2, STR5, STR1, STR6, STR1, STR2, STR4); setVector(vector2, STR1, STR2, STR4, STR3, STR2, STR5, STR1, STR6, STR1, STR2, STR4); RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector1, vector2); // inclusion of long string in the middle assertTrue(visitor.rangeEquals(new Range(1, 1, 3))); assertFalse(visitor.rangeEquals(new Range(0, 1, 4))); // inclusion of long string at the start assertTrue(visitor.rangeEquals(new Range(2, 2, 4))); assertFalse(visitor.rangeEquals(new Range(2, 5, 4))); // inclusion of long string at the end assertTrue(visitor.rangeEquals(new Range(4, 4, 4))); // unequal range assertTrue(visitor.rangeEquals(new Range(8, 0, 3))); assertFalse(visitor.rangeEquals(new Range(4, 5, 3))); // checking the same ranges when nulls are set vector1.setNull(1); vector2.setNull(1); vector1.setNull(3); vector2.setNull(3); vector1.setNull(5); vector2.setNull(5); vector1.setNull(9); vector2.setNull(9); // inclusion of long string in the middle assertTrue(visitor.rangeEquals(new Range(1, 1, 3))); assertFalse(visitor.rangeEquals(new Range(0, 1, 4))); // inclusion of long string at the start assertTrue(visitor.rangeEquals(new Range(2, 2, 4))); assertFalse(visitor.rangeEquals(new Range(2, 5, 4))); // inclusion of long string at the end assertTrue(visitor.rangeEquals(new Range(4, 4, 4))); // unequal range assertTrue(visitor.rangeEquals(new Range(8, 0, 3))); assertFalse(visitor.rangeEquals(new Range(4, 5, 3))); } } @Test public void testListVectorWithDifferentChild() { try (final ListVector vector1 = ListVector.empty("list", allocator); final ListVector vector2 = ListVector.empty("list", allocator); ) { vector1.allocateNew(); vector1.initializeChildrenFromFields( Arrays.asList(Field.nullable("child", new ArrowType.Int(32, true)))); vector2.allocateNew(); vector2.initializeChildrenFromFields( Arrays.asList(Field.nullable("child", new ArrowType.Int(64, true)))); RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector1, vector2); assertFalse(visitor.rangeEquals(new Range(0, 0, 0))); } } @Test public void testListViewVectorWithDifferentChild() { try (final ListViewVector vector1 = ListViewVector.empty("listview", allocator); final ListViewVector vector2 = ListViewVector.empty("listview", allocator); ) { vector1.allocateNew(); vector1.initializeChildrenFromFields( Arrays.asList(Field.nullable("child", new ArrowType.Int(32, true)))); vector2.allocateNew(); vector2.initializeChildrenFromFields( Arrays.asList(Field.nullable("child", new ArrowType.Int(64, true)))); RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector1, vector2); assertFalse(visitor.rangeEquals(new Range(0, 0, 0))); } } @Test public void testLargeListViewVectorWithDifferentChild() { try (final LargeListViewVector vector1 = LargeListViewVector.empty("largelistview", allocator); final LargeListViewVector vector2 = LargeListViewVector.empty("largelistview", allocator); ) { vector1.allocateNew(); vector1.initializeChildrenFromFields( Arrays.asList(Field.nullable("child", new ArrowType.Int(32, true)))); vector2.allocateNew(); vector2.initializeChildrenFromFields( Arrays.asList(Field.nullable("child", new ArrowType.Int(64, true)))); RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector1, vector2); assertFalse(visitor.rangeEquals(new Range(0, 0, 0))); } } @Test public void testListVectorRangeEquals() { try (final ListVector vector1 = ListVector.empty("list", allocator); final ListVector vector2 = ListVector.empty("list", allocator); ) { UnionListWriter writer1 = vector1.getWriter(); writer1.allocate(); // set some values writeListVector(writer1, new int[] {1, 2}); writeListVector(writer1, new int[] {3, 4}); writeListVector(writer1, new int[] {5, 6}); writeListVector(writer1, new int[] {7, 8}); writeListVector(writer1, new int[] {9, 10}); writer1.setValueCount(5); UnionListWriter writer2 = vector2.getWriter(); writer2.allocate(); // set some values writeListVector(writer2, new int[] {0, 0}); writeListVector(writer2, new int[] {3, 4}); writeListVector(writer2, new int[] {5, 6}); writeListVector(writer2, new int[] {7, 8}); writeListVector(writer2, new int[] {0, 0}); writer2.setValueCount(5); RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector1, vector2); assertTrue(visitor.rangeEquals(new Range(1, 1, 3))); } } @Test public void testListViewVectorRangeEquals() { try (final ListViewVector vector1 = ListViewVector.empty("listview", allocator); final ListViewVector vector2 = ListViewVector.empty("listview", allocator); ) { UnionListViewWriter writer1 = vector1.getWriter(); writer1.allocate(); // set some values writeListViewVector(writer1, new int[] {1, 2}); writeListViewVector(writer1, new int[] {3, 4}); writeListViewVector(writer1, new int[] {5, 6}); writeListViewVector(writer1, new int[] {7, 8}); writeListViewVector(writer1, new int[] {9, 10}); writer1.setValueCount(5); UnionListViewWriter writer2 = vector2.getWriter(); writer2.allocate(); // set some values writeListViewVector(writer2, new int[] {0, 0}); writeListViewVector(writer2, new int[] {3, 4}); writeListViewVector(writer2, new int[] {5, 6}); writeListViewVector(writer2, new int[] {7, 8}); writeListViewVector(writer2, new int[] {0, 0}); writer2.setValueCount(5); RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector1, vector2); assertTrue(visitor.rangeEquals(new Range(1, 1, 3))); } } @Test public void testLargeListViewVectorRangeEquals() { try (final LargeListViewVector vector1 = LargeListViewVector.empty("largelistview", allocator); final LargeListViewVector vector2 = LargeListViewVector.empty("largelistview", allocator); ) { UnionLargeListViewWriter writer1 = vector1.getWriter(); writer1.allocate(); // set some values writeLargeListViewVector(writer1, new int[] {1, 2}); writeLargeListViewVector(writer1, new int[] {3, 4}); writeLargeListViewVector(writer1, new int[] {5, 6}); writeLargeListViewVector(writer1, new int[] {7, 8}); writeLargeListViewVector(writer1, new int[] {9, 10}); writer1.setValueCount(5); UnionLargeListViewWriter writer2 = vector2.getWriter(); writer2.allocate(); // set some values writeLargeListViewVector(writer2, new int[] {0, 0}); writeLargeListViewVector(writer2, new int[] {3, 4}); writeLargeListViewVector(writer2, new int[] {5, 6}); writeLargeListViewVector(writer2, new int[] {7, 8}); writeLargeListViewVector(writer2, new int[] {0, 0}); writer2.setValueCount(5); RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector1, vector2); assertTrue(visitor.rangeEquals(new Range(1, 1, 3))); } } @Test public void testBitVectorRangeEquals() { try (final BitVector vector1 = new BitVector("v1", allocator); final BitVector vector2 = new BitVector("v2", allocator); ) { boolean[] v1 = new boolean[] {true, false, true, true, true}; boolean[] v2 = new boolean[] {false, true, true, true, false}; vector1.setValueCount(5); for (int i = 0; i < 5; i++) { vector1.set(i, v1[i] ? 1 : 0); } vector2.setValueCount(5); for (int i = 0; i < 5; i++) { vector2.set(i, v2[i] ? 1 : 0); } RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector1, vector2); assertTrue(visitor.compareBaseFixedWidthVectors(new Range(1, 0, 4))); assertFalse(visitor.compareBaseFixedWidthVectors(new Range(0, 0, 5))); } } @Test public void testFixedSizeListVectorRangeEquals() { try (final FixedSizeListVector vector1 = FixedSizeListVector.empty("list", 2, allocator); final FixedSizeListVector vector2 = FixedSizeListVector.empty("list", 2, allocator); ) { UnionFixedSizeListWriter writer1 = vector1.getWriter(); writer1.allocate(); // set some values writeFixedSizeListVector(writer1, new int[] {1, 2}); writeFixedSizeListVector(writer1, new int[] {3, 4}); writeFixedSizeListVector(writer1, new int[] {5, 6}); writeFixedSizeListVector(writer1, new int[] {7, 8}); writeFixedSizeListVector(writer1, new int[] {9, 10}); writer1.setValueCount(5); UnionFixedSizeListWriter writer2 = vector2.getWriter(); writer2.allocate(); // set some values writeFixedSizeListVector(writer2, new int[] {0, 0}); writeFixedSizeListVector(writer2, new int[] {3, 4}); writeFixedSizeListVector(writer2, new int[] {5, 6}); writeFixedSizeListVector(writer2, new int[] {7, 8}); writeFixedSizeListVector(writer2, new int[] {0, 0}); writer2.setValueCount(5); RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector1, vector2); assertTrue(visitor.rangeEquals(new Range(1, 1, 3))); assertFalse(visitor.rangeEquals(new Range(0, 0, 5))); } } @Test public void testLargeVariableWidthVectorRangeEquals() { try (final LargeVarCharVector vector1 = new LargeVarCharVector("vector1", allocator); final LargeVarCharVector vector2 = new LargeVarCharVector("vector2", allocator)) { setVector(vector1, "aaa", "bbb", "ccc", null, "ddd"); setVector(vector2, "ccc", "aaa", "bbb", null, "ddd"); RangeEqualsVisitor visitor = new RangeEqualsVisitor( vector1, vector2, (v1, v2) -> new TypeEqualsVisitor(v2, /*check name*/ false, /*check metadata*/ false) .equals(v1)); assertFalse( visitor.rangeEquals(new Range(/*left start*/ 0, /*right start*/ 0, /*length*/ 1))); assertTrue(visitor.rangeEquals(new Range(/*left start*/ 0, /*right start*/ 1, /*length*/ 1))); assertFalse( visitor.rangeEquals(new Range(/*left start*/ 0, /*right start*/ 0, /*length*/ 3))); assertTrue(visitor.rangeEquals(new Range(/*left start*/ 0, /*right start*/ 1, /*length*/ 2))); assertTrue(visitor.rangeEquals(new Range(/*left start*/ 3, /*right start*/ 3, /*length*/ 1))); assertTrue(visitor.rangeEquals(new Range(/*left start*/ 3, /*right start*/ 3, /*length*/ 2))); assertFalse( visitor.rangeEquals(new Range(/*left start*/ 2, /*right start*/ 2, /*length*/ 2))); } } @Test public void testStructVectorRangeEquals() { try (final StructVector vector1 = StructVector.empty("struct", allocator); final StructVector vector2 = StructVector.empty("struct", allocator); ) { vector1.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); vector1.addOrGet("f1", FieldType.nullable(new ArrowType.Int(64, true)), BigIntVector.class); vector2.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); vector2.addOrGet("f1", FieldType.nullable(new ArrowType.Int(64, true)), BigIntVector.class); NullableStructWriter writer1 = vector1.getWriter(); writer1.allocate(); writeStructVector(writer1, 0, 0L); writeStructVector(writer1, 1, 10L); writeStructVector(writer1, 2, 20L); writeStructVector(writer1, 3, 30L); writeStructVector(writer1, 4, 40L); writeStructVector(writer1, 5, 50L); writer1.setValueCount(6); NullableStructWriter writer2 = vector2.getWriter(); writer2.allocate(); writeStructVector(writer2, 0, 0L); writeStructVector(writer2, 2, 20L); writeStructVector(writer2, 3, 30L); writeStructVector(writer2, 4, 40L); writeStructVector(writer2, 0, 0L); writer2.setValueCount(5); RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector1, vector2); assertTrue(visitor.rangeEquals(new Range(2, 1, 3))); // different nullability but same values vector1.setNull(3); assertFalse(visitor.rangeEquals(new Range(2, 1, 3))); // both null and same values vector2.setNull(2); assertTrue(visitor.rangeEquals(new Range(2, 1, 3))); // both not null but different values assertFalse(visitor.rangeEquals(new Range(2, 1, 4))); // both null but different values vector1.setNull(5); vector2.setNull(4); assertTrue(visitor.rangeEquals(new Range(2, 1, 4))); } } @Test public void testUnionVectorRangeEquals() { try (final UnionVector vector1 = new UnionVector("union", allocator, /* field type */ null, /* call-back */ null); final UnionVector vector2 = new UnionVector("union", allocator, /* field type */ null, /* call-back */ null); ) { final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); uInt4Holder.value = 10; uInt4Holder.isSet = 1; final NullableIntHolder intHolder = new NullableIntHolder(); uInt4Holder.value = 20; uInt4Holder.isSet = 1; vector1.setType(0, Types.MinorType.UINT4); vector1.setSafe(0, uInt4Holder); vector1.setType(1, Types.MinorType.INT); vector1.setSafe(1, intHolder); vector1.setType(2, Types.MinorType.INT); vector1.setSafe(2, intHolder); vector1.setValueCount(3); vector2.setType(0, Types.MinorType.UINT4); vector2.setSafe(0, uInt4Holder); vector2.setType(1, Types.MinorType.INT); vector2.setSafe(1, intHolder); vector2.setType(2, Types.MinorType.INT); vector2.setSafe(2, intHolder); vector2.setValueCount(3); RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector1, vector2); assertTrue(visitor.rangeEquals(new Range(1, 1, 2))); } } /** * Test comparing two union vectors. The two vectors are different in total, but have a range with * equal values. */ @Test public void testUnionVectorSubRangeEquals() { try (final UnionVector vector1 = new UnionVector("union", allocator, null, null); final UnionVector vector2 = new UnionVector("union", allocator, null, null); ) { final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); uInt4Holder.value = 10; uInt4Holder.isSet = 1; final NullableIntHolder intHolder = new NullableIntHolder(); intHolder.value = 20; intHolder.isSet = 1; vector1.setType(0, Types.MinorType.UINT4); vector1.setSafe(0, uInt4Holder); vector1.setType(1, Types.MinorType.INT); vector1.setSafe(1, intHolder); vector1.setType(2, Types.MinorType.INT); vector1.setSafe(2, intHolder); vector1.setType(3, Types.MinorType.INT); vector1.setSafe(3, intHolder); vector1.setValueCount(4); vector2.setType(0, Types.MinorType.UINT4); vector2.setSafe(0, uInt4Holder); vector2.setType(1, Types.MinorType.INT); vector2.setSafe(1, intHolder); vector2.setType(2, Types.MinorType.INT); vector2.setSafe(2, intHolder); vector2.setType(3, Types.MinorType.UINT4); vector2.setSafe(3, uInt4Holder); vector2.setValueCount(4); RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector1, vector2); assertFalse(visitor.rangeEquals(new Range(0, 0, 4))); assertTrue(visitor.rangeEquals(new Range(1, 1, 2))); } } @Test public void testDenseUnionVectorEquals() { final NullableIntHolder intHolder = new NullableIntHolder(); intHolder.isSet = 1; intHolder.value = 100; final NullableBigIntHolder bigIntHolder = new NullableBigIntHolder(); bigIntHolder.isSet = 1; bigIntHolder.value = 200L; final NullableFloat4Holder float4Holder = new NullableFloat4Holder(); float4Holder.isSet = 1; float4Holder.value = 400F; final NullableFloat8Holder float8Holder = new NullableFloat8Holder(); float8Holder.isSet = 1; float8Holder.value = 800D; try (DenseUnionVector vector1 = new DenseUnionVector("vector1", allocator, null, null); DenseUnionVector vector2 = new DenseUnionVector("vector2", allocator, null, null)) { vector1.allocateNew(); vector2.allocateNew(); // populate vector1: {100, 200L, null, 400F, 800D} byte intTypeId = vector1.registerNewTypeId(Field.nullable("int", Types.MinorType.INT.getType())); byte longTypeId = vector1.registerNewTypeId(Field.nullable("long", Types.MinorType.BIGINT.getType())); byte floatTypeId = vector1.registerNewTypeId(Field.nullable("float", Types.MinorType.FLOAT4.getType())); byte doubleTypeId = vector1.registerNewTypeId(Field.nullable("double", Types.MinorType.FLOAT8.getType())); vector1.setTypeId(0, intTypeId); vector1.setSafe(0, intHolder); vector1.setTypeId(1, longTypeId); vector1.setSafe(1, bigIntHolder); vector1.setTypeId(3, floatTypeId); vector1.setSafe(3, float4Holder); vector1.setTypeId(4, doubleTypeId); vector1.setSafe(4, float8Holder); vector1.setValueCount(5); // populate vector2: {400F, null, 200L, null, 400F, 800D, 100} intTypeId = vector2.registerNewTypeId(Field.nullable("int", Types.MinorType.INT.getType())); longTypeId = vector2.registerNewTypeId(Field.nullable("long", Types.MinorType.BIGINT.getType())); floatTypeId = vector2.registerNewTypeId(Field.nullable("float", Types.MinorType.FLOAT4.getType())); doubleTypeId = vector2.registerNewTypeId(Field.nullable("double", Types.MinorType.FLOAT8.getType())); vector2.setTypeId(0, floatTypeId); vector2.setSafe(0, float4Holder); vector2.setTypeId(2, longTypeId); vector2.setSafe(2, bigIntHolder); vector2.setTypeId(4, floatTypeId); vector2.setSafe(4, float4Holder); vector2.setTypeId(5, doubleTypeId); vector2.setSafe(5, float8Holder); vector2.setTypeId(6, intTypeId); vector2.setSafe(6, intHolder); vector2.setValueCount(7); // compare ranges TypeEqualsVisitor typeVisitor = new TypeEqualsVisitor(vector2, /* check name */ false, /* check meta data */ true); RangeEqualsVisitor equalsVisitor = new RangeEqualsVisitor(vector1, vector2, (left, right) -> typeVisitor.equals(left)); // different ranges {100, 200L} != {400F, null} assertFalse(equalsVisitor.rangeEquals(new Range(0, 0, 2))); // different ranges without null {100, 200L} != {400F, null} assertFalse(equalsVisitor.rangeEquals(new Range(3, 5, 2))); // equal ranges {200L, null, 400F, 800D} assertTrue(equalsVisitor.rangeEquals(new Range(1, 2, 4))); // equal ranges without null {400F, 800D} assertTrue(equalsVisitor.rangeEquals(new Range(3, 4, 2))); // equal ranges with only null {null} assertTrue(equalsVisitor.rangeEquals(new Range(2, 3, 1))); // equal ranges with single element {100} assertTrue(equalsVisitor.rangeEquals(new Range(0, 6, 1))); // different ranges with single element {100} != {200L} assertFalse(equalsVisitor.rangeEquals(new Range(0, 2, 1))); } } @Disabled @Test public void testEqualsWithOutTypeCheck() { try (final IntVector intVector = new IntVector("int", allocator); final ZeroVector zeroVector = new ZeroVector("zero")) { assertTrue(VectorEqualsVisitor.vectorEquals(intVector, zeroVector, null)); assertTrue(VectorEqualsVisitor.vectorEquals(zeroVector, intVector, null)); } } @Test public void testFloat4ApproxEquals() { try (final Float4Vector vector1 = new Float4Vector("float", allocator); final Float4Vector vector2 = new Float4Vector("float", allocator); final Float4Vector vector3 = new Float4Vector("float", allocator)) { final float epsilon = 1.0E-6f; setVector(vector1, 1.1f, 2.2f); setVector(vector2, 1.1f + epsilon / 2, 2.2f + epsilon / 2); setVector(vector3, 1.1f + epsilon * 2, 2.2f + epsilon * 2); Range range = new Range(0, 0, vector1.getValueCount()); ApproxEqualsVisitor visitor12 = new ApproxEqualsVisitor(vector1, vector2, epsilon, epsilon); assertTrue(visitor12.rangeEquals(range)); ApproxEqualsVisitor visitor13 = new ApproxEqualsVisitor(vector1, vector3, epsilon, epsilon); assertFalse(visitor13.rangeEquals(range)); } } @Test public void testFloat8ApproxEquals() { try (final Float8Vector vector1 = new Float8Vector("float", allocator); final Float8Vector vector2 = new Float8Vector("float", allocator); final Float8Vector vector3 = new Float8Vector("float", 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); Range range = new Range(0, 0, vector1.getValueCount()); assertTrue(new ApproxEqualsVisitor(vector1, vector2, epsilon, epsilon).rangeEquals(range)); assertFalse(new ApproxEqualsVisitor(vector1, vector3, epsilon, epsilon).rangeEquals(range)); } } @Test public void testStructVectorApproxEquals() { try (final StructVector right = StructVector.empty("struct", allocator); final StructVector left1 = StructVector.empty("struct", allocator); final StructVector left2 = StructVector.empty("struct", allocator)) { right.addOrGet( "f0", FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), Float4Vector.class); right.addOrGet( "f1", FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), Float8Vector.class); left1.addOrGet( "f0", FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), Float4Vector.class); left1.addOrGet( "f1", FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), Float8Vector.class); left2.addOrGet( "f0", FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), Float4Vector.class); left2.addOrGet( "f1", FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), Float8Vector.class); final float epsilon = 1.0E-6f; NullableStructWriter rightWriter = right.getWriter(); rightWriter.allocate(); writeStructVector(rightWriter, 1.1f, 2.2); writeStructVector(rightWriter, 2.02f, 4.04); rightWriter.setValueCount(2); NullableStructWriter leftWriter1 = left1.getWriter(); leftWriter1.allocate(); writeStructVector(leftWriter1, 1.1f + epsilon / 2, 2.2 + epsilon / 2); writeStructVector(leftWriter1, 2.02f - epsilon / 2, 4.04 - epsilon / 2); leftWriter1.setValueCount(2); NullableStructWriter leftWriter2 = left2.getWriter(); leftWriter2.allocate(); writeStructVector(leftWriter2, 1.1f + epsilon * 2, 2.2 + epsilon * 2); writeStructVector(leftWriter2, 2.02f - epsilon * 2, 4.04 - epsilon * 2); leftWriter2.setValueCount(2); Range range = new Range(0, 0, right.getValueCount()); assertTrue(new ApproxEqualsVisitor(left1, right, epsilon, epsilon).rangeEquals(range)); assertFalse(new ApproxEqualsVisitor(left2, right, epsilon, epsilon).rangeEquals(range)); } } @Test public void testUnionVectorApproxEquals() { try (final UnionVector right = new UnionVector("union", allocator, /* field type */ null, /* call-back */ null); final UnionVector left1 = new UnionVector("union", allocator, /* field type */ null, /* call-back */ null); final UnionVector left2 = new UnionVector("union", allocator, /* field type */ null, /* call-back */ null); ) { final NullableFloat4Holder float4Holder = new NullableFloat4Holder(); float4Holder.value = 1.01f; float4Holder.isSet = 1; final NullableFloat8Holder float8Holder = new NullableFloat8Holder(); float8Holder.value = 2.02f; float8Holder.isSet = 1; final float epsilon = 1.0E-6f; right.setType(0, Types.MinorType.FLOAT4); right.setSafe(0, float4Holder); right.setType(1, Types.MinorType.FLOAT8); right.setSafe(1, float8Holder); right.setValueCount(2); float4Holder.value += epsilon / 2; float8Holder.value += epsilon / 2; left1.setType(0, Types.MinorType.FLOAT4); left1.setSafe(0, float4Holder); left1.setType(1, Types.MinorType.FLOAT8); left1.setSafe(1, float8Holder); left1.setValueCount(2); float4Holder.value += epsilon * 2; float8Holder.value += epsilon * 2; left2.setType(0, Types.MinorType.FLOAT4); left2.setSafe(0, float4Holder); left2.setType(1, Types.MinorType.FLOAT8); left2.setSafe(1, float8Holder); left2.setValueCount(2); Range range = new Range(0, 0, right.getValueCount()); assertTrue(new ApproxEqualsVisitor(left1, right, epsilon, epsilon).rangeEquals(range)); assertFalse(new ApproxEqualsVisitor(left2, right, epsilon, epsilon).rangeEquals(range)); } } @Test public void testDenseUnionVectorApproxEquals() { final NullableFloat4Holder float4Holder = new NullableFloat4Holder(); float4Holder.isSet = 1; final NullableFloat8Holder float8Holder = new NullableFloat8Holder(); float8Holder.isSet = 1; final float floatEpsilon = 0.02F; final double doubleEpsilon = 0.02; try (final DenseUnionVector vector1 = new DenseUnionVector("vector1", allocator, null, null); final DenseUnionVector vector2 = new DenseUnionVector("vector2", allocator, null, null); final DenseUnionVector vector3 = new DenseUnionVector("vector2", allocator, null, null)) { vector1.allocateNew(); vector2.allocateNew(); vector3.allocateNew(); // populate vector1: {1.0f, 2.0D} byte floatTypeId = vector1.registerNewTypeId(Field.nullable("float", Types.MinorType.FLOAT4.getType())); byte doubleTypeId = vector1.registerNewTypeId(Field.nullable("double", Types.MinorType.FLOAT8.getType())); float4Holder.value = 1.0f; vector1.setTypeId(0, floatTypeId); vector1.setSafe(0, float4Holder); float8Holder.value = 2.0D; vector1.setTypeId(1, doubleTypeId); vector1.setSafe(1, float8Holder); vector1.setValueCount(2); // populate vector2: {1.01f, 2.01D} floatTypeId = vector2.registerNewTypeId(Field.nullable("float", Types.MinorType.FLOAT4.getType())); doubleTypeId = vector2.registerNewTypeId(Field.nullable("double", Types.MinorType.FLOAT8.getType())); float4Holder.value = 1.01f; vector2.setTypeId(0, floatTypeId); vector2.setSafe(0, float4Holder); float8Holder.value = 2.01D; vector2.setTypeId(1, doubleTypeId); vector2.setSafe(1, float8Holder); vector2.setValueCount(2); // populate vector3: {1.05f, 2.05D} floatTypeId = vector3.registerNewTypeId(Field.nullable("float", Types.MinorType.FLOAT4.getType())); doubleTypeId = vector3.registerNewTypeId(Field.nullable("double", Types.MinorType.FLOAT8.getType())); float4Holder.value = 1.05f; vector3.setTypeId(0, floatTypeId); vector3.setSafe(0, float4Holder); float8Holder.value = 2.05D; vector3.setTypeId(1, doubleTypeId); vector3.setSafe(1, float8Holder); vector3.setValueCount(2); // verify comparison results Range range = new Range(0, 0, 2); // compare vector1 and vector2 ApproxEqualsVisitor approxEqualsVisitor = new ApproxEqualsVisitor( vector1, vector2, new ValueEpsilonEqualizers.Float4EpsilonEqualizer(floatEpsilon), new ValueEpsilonEqualizers.Float8EpsilonEqualizer(doubleEpsilon), (v1, v2) -> new TypeEqualsVisitor(v2, /* check name */ false, /* check meta */ true) .equals(v1)); assertTrue(approxEqualsVisitor.rangeEquals(range)); // compare vector1 and vector3 approxEqualsVisitor = new ApproxEqualsVisitor( vector1, vector3, new ValueEpsilonEqualizers.Float4EpsilonEqualizer(floatEpsilon), new ValueEpsilonEqualizers.Float8EpsilonEqualizer(doubleEpsilon), (v1, v2) -> new TypeEqualsVisitor(v2, /* check name */ false, /* check meta */ true) .equals(v1)); assertFalse(approxEqualsVisitor.rangeEquals(range)); } } @Test public void testListVectorApproxEquals() { try (final ListVector right = ListVector.empty("list", allocator); final ListVector left1 = ListVector.empty("list", allocator); final ListVector left2 = ListVector.empty("list", allocator); ) { final float epsilon = 1.0E-6f; UnionListWriter rightWriter = right.getWriter(); rightWriter.allocate(); writeListVector(rightWriter, new double[] {1, 2}); writeListVector(rightWriter, new double[] {1.01, 2.02}); rightWriter.setValueCount(2); UnionListWriter leftWriter1 = left1.getWriter(); leftWriter1.allocate(); writeListVector(leftWriter1, new double[] {1, 2}); writeListVector(leftWriter1, new double[] {1.01 + epsilon / 2, 2.02 - epsilon / 2}); leftWriter1.setValueCount(2); UnionListWriter leftWriter2 = left2.getWriter(); leftWriter2.allocate(); writeListVector(leftWriter2, new double[] {1, 2}); writeListVector(leftWriter2, new double[] {1.01 + epsilon * 2, 2.02 - epsilon * 2}); leftWriter2.setValueCount(2); Range range = new Range(0, 0, right.getValueCount()); assertTrue(new ApproxEqualsVisitor(left1, right, epsilon, epsilon).rangeEquals(range)); assertFalse(new ApproxEqualsVisitor(left2, right, epsilon, epsilon).rangeEquals(range)); } } @Test public void testListViewVectorApproxEquals() { try (final ListViewVector right = ListViewVector.empty("listview", allocator); final ListViewVector left1 = ListViewVector.empty("listview", allocator); final ListViewVector left2 = ListViewVector.empty("listview", allocator); ) { final float epsilon = 1.0E-6f; UnionListViewWriter rightWriter = right.getWriter(); rightWriter.allocate(); writeListViewVector(rightWriter, new double[] {1, 2}); writeListViewVector(rightWriter, new double[] {1.01, 2.02}); rightWriter.setValueCount(2); UnionListViewWriter leftWriter1 = left1.getWriter(); leftWriter1.allocate(); writeListViewVector(leftWriter1, new double[] {1, 2}); writeListViewVector(leftWriter1, new double[] {1.01 + epsilon / 2, 2.02 - epsilon / 2}); leftWriter1.setValueCount(2); UnionListViewWriter leftWriter2 = left2.getWriter(); leftWriter2.allocate(); writeListViewVector(leftWriter2, new double[] {1, 2}); writeListViewVector(leftWriter2, new double[] {1.01 + epsilon * 2, 2.02 - epsilon * 2}); leftWriter2.setValueCount(2); Range range = new Range(0, 0, right.getValueCount()); assertTrue(new ApproxEqualsVisitor(left1, right, epsilon, epsilon).rangeEquals(range)); assertFalse(new ApproxEqualsVisitor(left2, right, epsilon, epsilon).rangeEquals(range)); } } @Test public void testLargeListViewVectorApproxEquals() { try (final LargeListViewVector right = LargeListViewVector.empty("largelistview", allocator); final LargeListViewVector left1 = LargeListViewVector.empty("largelistview", allocator); final LargeListViewVector left2 = LargeListViewVector.empty("largelistview", allocator); ) { final float epsilon = 1.0E-6f; UnionLargeListViewWriter rightWriter = right.getWriter(); rightWriter.allocate(); writeLargeListViewVector(rightWriter, new double[] {1, 2}); writeLargeListViewVector(rightWriter, new double[] {1.01, 2.02}); rightWriter.setValueCount(2); UnionLargeListViewWriter leftWriter1 = left1.getWriter(); leftWriter1.allocate(); writeLargeListViewVector(leftWriter1, new double[] {1, 2}); writeLargeListViewVector(leftWriter1, new double[] {1.01 + epsilon / 2, 2.02 - epsilon / 2}); leftWriter1.setValueCount(2); UnionLargeListViewWriter leftWriter2 = left2.getWriter(); leftWriter2.allocate(); writeLargeListViewVector(leftWriter2, new double[] {1, 2}); writeLargeListViewVector(leftWriter2, new double[] {1.01 + epsilon * 2, 2.02 - epsilon * 2}); leftWriter2.setValueCount(2); Range range = new Range(0, 0, right.getValueCount()); assertTrue(new ApproxEqualsVisitor(left1, right, epsilon, epsilon).rangeEquals(range)); assertFalse(new ApproxEqualsVisitor(left2, right, epsilon, epsilon).rangeEquals(range)); } } @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); writer.bigInt("f1").writeBigInt(value2); writer.end(); } private void writeStructVector(NullableStructWriter writer, float value1, double value2) { writer.start(); writer.float4("f0").writeFloat4(value1); writer.float8("f1").writeFloat8(value2); writer.end(); } private void writeListVector(UnionListWriter writer, int[] values) { writer.startList(); for (int v : values) { writer.integer().writeInt(v); } writer.endList(); } private void writeListViewVector(UnionListViewWriter writer, int[] values) { writer.startListView(); for (int v : values) { writer.integer().writeInt(v); } writer.endListView(); } private void writeLargeListViewVector(UnionLargeListViewWriter writer, int[] values) { writer.startListView(); for (int v : values) { writer.integer().writeInt(v); } writer.endListView(); } private void writeFixedSizeListVector(UnionFixedSizeListWriter writer, int[] values) { writer.startList(); for (int v : values) { writer.integer().writeInt(v); } writer.endList(); } private void writeListVector(UnionListWriter writer, double[] values) { writer.startList(); for (double v : values) { writer.float8().writeFloat8(v); } writer.endList(); } private void writeListViewVector(UnionListViewWriter writer, double[] values) { writer.startListView(); for (double v : values) { writer.float8().writeFloat8(v); } writer.endListView(); } private void writeLargeListViewVector(UnionLargeListViewWriter writer, double[] values) { writer.startListView(); for (double v : values) { writer.float8().writeFloat8(v); } writer.endListView(); } } TestTypeEqualsVisitor.java000066400000000000000000000253301515451615200342640ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/compare/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.compare; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.charset.StandardCharsets; import java.util.HashMap; import java.util.Map; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.IntVector; 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.LargeListViewVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.ListViewVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.UnionVector; 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.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestTypeEqualsVisitor { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testTypeEqualsWithName() { try (final IntVector right = new IntVector("int", allocator); final IntVector left1 = new IntVector("int", allocator); final IntVector left2 = new IntVector("int2", allocator)) { TypeEqualsVisitor visitor = new TypeEqualsVisitor(right); assertTrue(visitor.equals(left1)); assertFalse(visitor.equals(left2)); } } @Test public void testTypeEqualsWithMetadata() { Map metadata = new HashMap<>(); metadata.put("key1", "value1"); FieldType typeWithoutMeta = new FieldType(true, new ArrowType.Int(32, true), null, null); FieldType typeWithMeta = new FieldType(true, new ArrowType.Int(32, true), null, metadata); try (IntVector right = (IntVector) typeWithoutMeta.createNewSingleVector("int", allocator, null); IntVector left1 = (IntVector) typeWithoutMeta.createNewSingleVector("int", allocator, null); IntVector left2 = (IntVector) typeWithMeta.createNewSingleVector("int", allocator, null)) { TypeEqualsVisitor visitor = new TypeEqualsVisitor(right); assertTrue(visitor.equals(left1)); assertFalse(visitor.equals(left2)); } } @Test public void testListTypeEquals() { try (final ListVector right = ListVector.empty("list", allocator); final ListVector left1 = ListVector.empty("list", allocator); final ListVector left2 = ListVector.empty("list", allocator)) { right.addOrGetVector(FieldType.nullable(new ArrowType.Utf8())); left1.addOrGetVector(FieldType.nullable(new ArrowType.Utf8())); left2.addOrGetVector(FieldType.nullable(new ArrowType.FixedSizeBinary(2))); TypeEqualsVisitor visitor = new TypeEqualsVisitor(right); assertTrue(visitor.equals(left1)); assertFalse(visitor.equals(left2)); } } @Test public void testListViewTypeEquals() { try (final ListViewVector right = ListViewVector.empty("listview", allocator); final ListViewVector left1 = ListViewVector.empty("listview", allocator); final ListViewVector left2 = ListViewVector.empty("listview", allocator)) { right.addOrGetVector(FieldType.nullable(new ArrowType.Utf8())); left1.addOrGetVector(FieldType.nullable(new ArrowType.Utf8())); left2.addOrGetVector(FieldType.nullable(new ArrowType.FixedSizeBinary(2))); TypeEqualsVisitor visitor = new TypeEqualsVisitor(right); assertTrue(visitor.equals(left1)); assertFalse(visitor.equals(left2)); } } @Test public void testLargeListViewTypeEquals() { try (final LargeListViewVector right = LargeListViewVector.empty("largelistview", allocator); final LargeListViewVector left1 = LargeListViewVector.empty("largelistview", allocator); final LargeListViewVector left2 = LargeListViewVector.empty("largelistview", allocator)) { right.addOrGetVector(FieldType.nullable(new ArrowType.Utf8())); left1.addOrGetVector(FieldType.nullable(new ArrowType.Utf8())); left2.addOrGetVector(FieldType.nullable(new ArrowType.FixedSizeBinary(2))); TypeEqualsVisitor visitor = new TypeEqualsVisitor(right); assertTrue(visitor.equals(left1)); assertFalse(visitor.equals(left2)); } } @Test public void testStructTypeEquals() { try (final StructVector right = StructVector.empty("struct", allocator); final StructVector left1 = StructVector.empty("struct", allocator); final StructVector left2 = StructVector.empty("struct", allocator)) { right.addOrGet("child", FieldType.nullable(new ArrowType.Utf8()), VarCharVector.class); left1.addOrGet("child", FieldType.nullable(new ArrowType.Utf8()), VarCharVector.class); left2.addOrGet("child2", FieldType.nullable(new ArrowType.Utf8()), VarCharVector.class); TypeEqualsVisitor visitor = new TypeEqualsVisitor(right); assertTrue(visitor.equals(left1)); assertFalse(visitor.equals(left2)); } } @Test public void testUnionTypeEquals() { try (final UnionVector right = new UnionVector("union", allocator, /* field type */ null, /* call-back */ null); final UnionVector left1 = new UnionVector("union", allocator, /* field type */ null, /* call-back */ null); final UnionVector left2 = new UnionVector("union", allocator, /* field type */ null, /* call-back */ null)) { right.addVector(new IntVector("int", allocator)); left1.addVector(new IntVector("int", allocator)); left2.addVector(new BigIntVector("bigint", allocator)); TypeEqualsVisitor visitor = new TypeEqualsVisitor(right); assertTrue(visitor.equals(left1)); assertFalse(visitor.equals(left2)); } } @Test public void testDenseUnionTypeEquals() { try (DenseUnionVector vector1 = new DenseUnionVector("vector1", allocator, null, null); DenseUnionVector vector2 = new DenseUnionVector("vector2", allocator, null, null)) { vector1.allocateNew(); vector2.allocateNew(); // set children for vector1 byte intTypeId = vector1.registerNewTypeId(Field.nullable("int", Types.MinorType.INT.getType())); byte longTypeId = vector1.registerNewTypeId(Field.nullable("long", Types.MinorType.BIGINT.getType())); byte floatTypeId = vector1.registerNewTypeId(Field.nullable("float", Types.MinorType.FLOAT4.getType())); byte doubleTypeId = vector1.registerNewTypeId(Field.nullable("double", Types.MinorType.FLOAT8.getType())); vector1.addVector(floatTypeId, new Float4Vector("", allocator)); vector1.addVector(longTypeId, new BigIntVector("", allocator)); vector1.addVector(intTypeId, new IntVector("", allocator)); vector1.addVector(doubleTypeId, new Float8Vector("", allocator)); // set children for vector2 intTypeId = vector2.registerNewTypeId(Field.nullable("int", Types.MinorType.INT.getType())); longTypeId = vector2.registerNewTypeId(Field.nullable("long", Types.MinorType.BIGINT.getType())); floatTypeId = vector2.registerNewTypeId(Field.nullable("float", Types.MinorType.FLOAT4.getType())); doubleTypeId = vector2.registerNewTypeId(Field.nullable("double", Types.MinorType.FLOAT8.getType())); // add vectors in a different order vector2.addVector(intTypeId, new IntVector("", allocator)); vector2.addVector(floatTypeId, new Float4Vector("", allocator)); vector2.addVector(doubleTypeId, new Float8Vector("", allocator)); vector2.addVector(longTypeId, new BigIntVector("", allocator)); // compare ranges TypeEqualsVisitor typeVisitor = new TypeEqualsVisitor(vector2, /* check name */ false, /* check meta data */ true); assertTrue(typeVisitor.equals(vector1)); // if we check names, the types should be different typeVisitor = new TypeEqualsVisitor(vector2, /* check name */ true, /* check meta data */ true); assertFalse(typeVisitor.equals(vector1)); } } @Test public void testStringViewTypeEquals() { try (final ViewVarCharVector varchar1 = new ViewVarCharVector("varchar1", allocator); final ViewVarCharVector varchar2 = new ViewVarCharVector("varchar2", allocator); final ViewVarBinaryVector binary = new ViewVarBinaryVector("binary", allocator)) { final int valueCount = 2; final byte[] str0 = "apache".getBytes(StandardCharsets.UTF_8); final byte[] str1 = "arrow".getBytes(StandardCharsets.UTF_8); // add elements for varchar1 varchar1.allocateNew(48, valueCount); varchar1.set(0, str0); varchar1.set(1, str1); varchar1.setValueCount(valueCount); // add elements for varchar2 in a difference order varchar2.allocateNew(48, valueCount); varchar2.set(0, str1); varchar2.set(1, str0); varchar2.setValueCount(valueCount); // add elements for binary binary.allocateNew(48, valueCount); binary.set(0, str0); binary.set(1, str1); binary.setValueCount(valueCount); // compare ignore check name TypeEqualsVisitor visitor = new TypeEqualsVisitor(varchar1, /* check name */ false, /* check meta data */ true); assertTrue(visitor.equals(varchar2)); assertFalse(visitor.equals(binary)); // if we check names, the types should be different visitor = new TypeEqualsVisitor(varchar1, /* check name */ true, /* check meta data */ true); assertFalse(visitor.equals(varchar2)); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/complex/000077500000000000000000000000001515451615200272025ustar00rootroot00000000000000TestDenseUnionBufferSize.java000066400000000000000000000065671515451615200347000ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/complex/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BaseValueVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.holders.NullableIntHolder; import org.apache.arrow.vector.holders.NullableVarBinaryHolder; import org.apache.arrow.vector.types.UnionMode; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.FieldType; import org.junit.jupiter.api.Test; public class TestDenseUnionBufferSize { @Test public void testBufferSize() { try (BufferAllocator allocator = new RootAllocator(); DenseUnionVector duv = new DenseUnionVector( "duv", allocator, FieldType.nullable(new ArrowType.Union(UnionMode.Dense, null)), null)) { byte aTypeId = 42; byte bTypeId = 7; duv.addVector( aTypeId, new IntVector("a", FieldType.notNullable(new ArrowType.Int(32, true)), allocator)); duv.addVector( bTypeId, new VarBinaryVector("b", FieldType.notNullable(new ArrowType.Binary()), allocator)); NullableIntHolder intHolder = new NullableIntHolder(); NullableVarBinaryHolder varBinaryHolder = new NullableVarBinaryHolder(); int aCount = BaseValueVector.INITIAL_VALUE_ALLOCATION + 1; for (int i = 0; i < aCount; i++) { duv.setTypeId(i, aTypeId); duv.setSafe(i, intHolder); } int bCount = 2; for (int i = 0; i < bCount; i++) { duv.setTypeId(i + aCount, bTypeId); duv.setSafe(i + aCount, varBinaryHolder); } int count = aCount + bCount; duv.setValueCount(count); // will not necessarily see an error unless bounds checking is on. assertDoesNotThrow(duv::getBufferSize); IntVector intVector = duv.getIntVector(aTypeId); VarBinaryVector varBinaryVector = duv.getVarBinaryVector(bTypeId); long overhead = DenseUnionVector.TYPE_WIDTH + DenseUnionVector.OFFSET_WIDTH; assertEquals( overhead * count + intVector.getBufferSize() + varBinaryVector.getBufferSize(), duv.getBufferSize()); assertEquals( overhead * (aCount + 1) + intVector.getBufferSizeFor(aCount) + varBinaryVector.getBufferSizeFor(1), duv.getBufferSizeFor(aCount + 1)); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/complex/impl/000077500000000000000000000000001515451615200301435ustar00rootroot00000000000000TestComplexCopier.java000066400000000000000000000746201515451615200343510ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 static org.junit.jupiter.api.Assertions.assertThrows; 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; import org.apache.arrow.vector.compare.VectorEqualsVisitor; 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.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; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.DecimalUtility; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestComplexCopier { private BufferAllocator allocator; private static final int COUNT = 100; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testCopyFixedSizeListVector() { try (FixedSizeListVector from = FixedSizeListVector.empty("v", 3, allocator); FixedSizeListVector to = FixedSizeListVector.empty("v", 3, allocator)) { from.addOrGetVector(FieldType.nullable(Types.MinorType.INT.getType())); to.addOrGetVector(FieldType.nullable(Types.MinorType.INT.getType())); // populate from vector UnionFixedSizeListWriter writer = from.getWriter(); for (int i = 0; i < COUNT; i++) { writer.startList(); writer.integer().writeInt(i); writer.integer().writeInt(i * 2); writer.integer().writeInt(i * 3); writer.endList(); } from.setValueCount(COUNT); to.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); } // validate equals assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); } } @Test public void testInvalidCopyFixedSizeListVector() { try (FixedSizeListVector from = FixedSizeListVector.empty("v", 3, allocator); FixedSizeListVector to = FixedSizeListVector.empty("v", 2, allocator)) { from.addOrGetVector(FieldType.nullable(Types.MinorType.INT.getType())); to.addOrGetVector(FieldType.nullable(Types.MinorType.INT.getType())); // populate from vector UnionFixedSizeListWriter writer = from.getWriter(); for (int i = 0; i < COUNT; i++) { writer.startList(); writer.integer().writeInt(i); writer.integer().writeInt(i * 2); writer.integer().writeInt(i * 3); writer.endList(); } from.setValueCount(COUNT); to.setValueCount(COUNT); // copy values FieldReader in = from.getReader(); FieldWriter out = to.getWriter(); IllegalStateException e = assertThrows(IllegalStateException.class, () -> ComplexCopier.copy(in, out)); assertTrue(e.getMessage().contains("greater than listSize")); } } @Test public void testCopyMapVector() { 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(); mapWriter.key().integer().writeInt(i); mapWriter.value().integer().writeInt(i); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().decimal().writeDecimal(BigDecimal.valueOf(i * 2)); mapWriter.value().decimal().writeDecimal(BigDecimal.valueOf(i * 2)); 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 testCopyListVector() { 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(); listWriter.integer().writeInt(i); listWriter.integer().writeInt(i * 2); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(i); listWriter.list().bigInt().writeBigInt(i * 2); listWriter.list().bigInt().writeBigInt(i * 3); listWriter.list().endList(); listWriter.list().startList(); listWriter.list().decimal().writeDecimal(BigDecimal.valueOf(i * 4)); listWriter.list().decimal().writeDecimal(BigDecimal.valueOf(i * 5)); listWriter.list().endList(); 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 testCopyListVectorToANonEmptyList() { 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(); listWriter.integer().writeInt(i); listWriter.integer().writeInt(i * 2); 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)); // Copy again to the target vector which is non-empty 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)); // copy using copyFromSafe method for (int i = 0; i < COUNT; i++) { to.copyFromSafe(i, i, from); } to.setValueCount(COUNT); // validate equals assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); } } @Test public void testCopyListVectorWithNulls() { try (ListVector from = ListVector.empty("v", allocator); ListVector to = ListVector.empty("v", allocator)) { UnionListWriter listWriter = from.getWriter(); listWriter.allocate(); // writer null, [null,i,null,i*2,null] alternatively for (int i = 0; i < COUNT; i++) { listWriter.setPosition(i); if (i % 2 == 0) { listWriter.writeNull(); continue; } listWriter.startList(); listWriter.integer().writeNull(); listWriter.integer().writeInt(i); listWriter.integer().writeNull(); listWriter.integer().writeInt(i * 2); listWriter.integer().writeNull(); 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 testCopyListOfListVectorWithNulls() { try (ListVector from = ListVector.empty("v", allocator); ListVector to = ListVector.empty("v", allocator); ) { UnionListWriter listWriter = from.getWriter(); listWriter.allocate(); // write null, [null,[50,100,null,200],null, // [null,50,null,100,null,200,null],null] alternatively for (int i = 0; i < COUNT; i++) { listWriter.setPosition(i); if (i % 2 == 0) { listWriter.writeNull(); continue; } listWriter.startList(); listWriter.list().writeNull(); listWriter.list().startList(); listWriter.list().bigInt().writeBigInt(50); listWriter.list().bigInt().writeBigInt(100); listWriter.list().bigInt().writeNull(); listWriter.list().bigInt().writeBigInt(200); listWriter.list().endList(); listWriter.list().writeNull(); listWriter.list().startList(); listWriter.list().bigInt().writeNull(); listWriter.list().bigInt().writeBigInt(50); listWriter.list().bigInt().writeNull(); listWriter.list().bigInt().writeBigInt(100); listWriter.list().bigInt().writeNull(); listWriter.list().bigInt().writeBigInt(200); listWriter.list().bigInt().writeNull(); listWriter.list().endList(); listWriter.list().writeNull(); 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 testCopyListOStructVectorWithNulls() { try (ListVector from = ListVector.empty("v", allocator); ListVector to = ListVector.empty("v", allocator); ) { UnionListWriter listWriter = from.getWriter(); listWriter.allocate(); // write null, [null,{"f1":1,"f2":2},null, // {"f1":1,"f2":2},null] alternatively for (int i = 0; i < COUNT; i++) { listWriter.setPosition(i); if (i % 2 == 0) { listWriter.writeNull(); continue; } listWriter.startList(); listWriter.struct().writeNull(); listWriter.struct().start(); listWriter.struct().integer("f1").writeInt(1); listWriter.struct().integer("f2").writeInt(2); listWriter.struct().integer("f3").writeNull(); listWriter.struct().end(); listWriter.struct().writeNull(); listWriter.struct().start(); listWriter.struct().integer("f1").writeInt(1); listWriter.struct().integer("f2").writeInt(2); listWriter.struct().integer("f3").writeNull(); listWriter.struct().end(); listWriter.struct().writeNull(); 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 testCopyListOfListOfStructVectorWithNulls() { try (ListVector from = ListVector.empty("v", allocator); ListVector to = ListVector.empty("v", allocator); ) { UnionListWriter listWriter = from.getWriter(); listWriter.allocate(); // write null, // [null,[{"f1":50},null,{"f1":100},null,{"f1":200}],null, // [null,{"f1":50},null,{"f1":100},null,{"f1":200},null],null] // alternatively for (int i = 0; i < COUNT; i++) { listWriter.setPosition(i); if (i % 2 == 0) { listWriter.writeNull(); continue; } listWriter.startList(); listWriter.list().writeNull(); listWriter.list().startList(); listWriter.list().struct().start(); listWriter.list().struct().bigInt("f1").writeBigInt(50); listWriter.list().struct().end(); listWriter.list().struct().writeNull(); listWriter.list().struct().start(); listWriter.list().struct().bigInt("f1").writeBigInt(100); listWriter.list().struct().end(); listWriter.list().struct().writeNull(); listWriter.list().struct().start(); listWriter.list().struct().bigInt("f1").writeBigInt(200); listWriter.list().struct().end(); listWriter.list().endList(); listWriter.list().writeNull(); listWriter.list().startList(); listWriter.list().struct().writeNull(); listWriter.list().struct().start(); listWriter.list().struct().bigInt("f1").writeBigInt(50); listWriter.list().struct().end(); listWriter.list().struct().writeNull(); listWriter.list().struct().start(); listWriter.list().struct().bigInt("f1").writeBigInt(100); listWriter.list().struct().end(); listWriter.list().struct().writeNull(); listWriter.list().struct().start(); listWriter.list().struct().bigInt("f1").writeBigInt(200); listWriter.list().struct().end(); listWriter.list().struct().writeNull(); listWriter.list().endList(); listWriter.list().writeNull(); 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 testMapWithListValue() throws Exception { try (MapVector from = MapVector.empty("map", allocator, false); MapVector to = MapVector.empty("map", allocator, false)) { UnionMapWriter mapWriter = from.getWriter(); BaseWriter.ListWriter valueWriter; /* allocate memory */ mapWriter.allocate(); // write null, [{}, // {"value":[]},{"key":1,"value":[null,50,null,100,null,200,null]}, // null,{"key":2,"value":[null,75,null,125,null,150,null,175,null]}] // alternatively for (int i = 0; i < COUNT; i++) { mapWriter.setPosition(i); if (i % 2 == 0) { mapWriter.writeNull(); continue; } mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().bigInt().writeNull(); mapWriter.value().list().writeNull(); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeNull(); valueWriter = mapWriter.value().list(); valueWriter.startList(); valueWriter.endList(); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(1); valueWriter = mapWriter.value().list(); valueWriter.startList(); valueWriter.bigInt().writeNull(); valueWriter.bigInt().writeBigInt(50); valueWriter.bigInt().writeNull(); valueWriter.bigInt().writeBigInt(100); valueWriter.bigInt().writeNull(); valueWriter.bigInt().writeBigInt(200); valueWriter.bigInt().writeNull(); valueWriter.endList(); mapWriter.endEntry(); mapWriter.writeNull(); mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(2); valueWriter = mapWriter.value().list(); valueWriter.startList(); valueWriter.bigInt().writeNull(); valueWriter.bigInt().writeBigInt(75); valueWriter.bigInt().writeNull(); valueWriter.bigInt().writeBigInt(125); valueWriter.bigInt().writeNull(); valueWriter.bigInt().writeBigInt(150); valueWriter.bigInt().writeNull(); valueWriter.bigInt().writeBigInt(175); valueWriter.bigInt().writeNull(); valueWriter.endList(); mapWriter.endEntry(); mapWriter.endMap(); } mapWriter.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 testCopyFixedSizedListOfDecimalsVector() { try (FixedSizeListVector from = FixedSizeListVector.empty("v", 4, allocator); FixedSizeListVector to = FixedSizeListVector.empty("v", 4, allocator)) { from.addOrGetVector(FieldType.nullable(new ArrowType.Decimal(3, 0, 128))); to.addOrGetVector(FieldType.nullable(new ArrowType.Decimal(3, 0, 128))); DecimalHolder holder = new DecimalHolder(); holder.buffer = allocator.buffer(DecimalVector.TYPE_WIDTH); ArrowType arrowType = new ArrowType.Decimal(3, 0, 128); // populate from vector UnionFixedSizeListWriter writer = from.getWriter(); for (int i = 0; i < COUNT; i++) { writer.startList(); writer.decimal().writeDecimal(BigDecimal.valueOf(i)); DecimalUtility.writeBigDecimalToArrowBuf( new BigDecimal(i * 2), holder.buffer, 0, DecimalVector.TYPE_WIDTH); holder.start = 0; holder.scale = 0; holder.precision = 3; writer.decimal().write(holder); DecimalUtility.writeBigDecimalToArrowBuf( new BigDecimal(i * 3), holder.buffer, 0, DecimalVector.TYPE_WIDTH); writer.decimal().writeDecimal(0, holder.buffer, arrowType); writer .decimal() .writeBigEndianBytesToDecimal( BigDecimal.valueOf(i * 4).unscaledValue().toByteArray(), arrowType); writer.endList(); } from.setValueCount(COUNT); to.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); } // validate equals assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); holder.buffer.close(); } } @Test public void testCopyUnionListWithDecimal() { 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(); listWriter.decimal().writeDecimal(BigDecimal.valueOf(i * 2)); listWriter.integer().writeInt(i); listWriter .decimal() .writeBigEndianBytesToDecimal( BigDecimal.valueOf(i * 3).unscaledValue().toByteArray(), new ArrowType.Decimal(3, 0, 128)); 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 testCopyStructVector() { 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(); structWriter.integer("int").writeInt(i); structWriter.decimal("dec", 0, 38).writeDecimal(BigDecimal.valueOf(i * 2)); StructWriter innerStructWriter = structWriter.struct("struc"); innerStructWriter.start(); innerStructWriter.integer("innerint").writeInt(i * 3); innerStructWriter.decimal("innerdec", 0, 38).writeDecimal(BigDecimal.valueOf(i * 4)); innerStructWriter .decimal("innerdec", 0, 38) .writeBigEndianBytesToDecimal( BigDecimal.valueOf(i * 4).unscaledValue().toByteArray(), new ArrowType.Decimal(3, 0, 128)); innerStructWriter.end(); 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)); } } @Test public void testCopyDecimalVectorWrongScale() { try (FixedSizeListVector from = FixedSizeListVector.empty("v", 3, allocator); FixedSizeListVector to = FixedSizeListVector.empty("v", 3, allocator)) { from.addOrGetVector(FieldType.nullable(new ArrowType.Decimal(3, 2, 128))); to.addOrGetVector(FieldType.nullable(new ArrowType.Decimal(3, 1, 128))); // populate from vector UnionFixedSizeListWriter writer = from.getWriter(); for (int i = 0; i < COUNT; i++) { writer.startList(); writer.decimal().writeDecimal(BigDecimal.valueOf(1.23)); writer.decimal().writeDecimal(BigDecimal.valueOf(2.45)); writer.endList(); } from.setValueCount(COUNT); to.setValueCount(COUNT); // copy values FieldReader in = from.getReader(); FieldWriter out = to.getWriter(); UnsupportedOperationException e = assertThrows(UnsupportedOperationException.class, () -> ComplexCopier.copy(in, out)); assertTrue( e.getMessage().contains("BigDecimal scale must equal that in the Arrow vector: 2 != 1")); } } @Test public void testCopyStructVectorWithNulls() { try (StructVector from = StructVector.empty("v", allocator); StructVector to = StructVector.empty("v", allocator)) { NullableStructWriter writer = from.getWriter(); for (int i = 0; i < COUNT; ++i) { writer.setPosition(i); writer.start(); writer.integer("int").writeInt(i); if (i % 3 == 0) { writer.float4("child").writeFloat4(12.3f); } else if (i % 3 == 1) { writer.integer("child").writeInt(123); } else { writer.integer("child").writeNull(); } writer.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)); } } @Test public void testCopyStructOfMap() { try (final StructVector from = StructVector.empty("v", allocator); final StructVector to = StructVector.empty("v", allocator); ) { from.allocateNew(); NullableStructWriter structWriter = from.getWriter(); for (int i = 0; i < COUNT; i++) { structWriter.setPosition(i); structWriter.start(); BaseWriter.MapWriter innerMapWriter = structWriter.map("f1"); innerMapWriter.startMap(); innerMapWriter.startEntry(); innerMapWriter.key().integer().writeInt(i); innerMapWriter.value().integer().writeInt(i); innerMapWriter.endEntry(); innerMapWriter.endMap(); 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)); } } @Test public void testCopyMapVectorWithMapValue() { 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(); mapWriter.key().integer().writeInt(i); BaseWriter.MapWriter innerMapWriter = mapWriter.value().map(false); innerMapWriter.startMap(); innerMapWriter.startEntry(); innerMapWriter.key().integer().writeInt(i); innerMapWriter.value().integer().writeInt(i); innerMapWriter.endEntry(); innerMapWriter.endMap(); 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 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)); } } } TestPromotableWriter.java000066400000000000000000000765321515451615200351050ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/complex/impl/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 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 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; import org.apache.arrow.vector.complex.NonNullableStructVector; 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; public class TestPromotableWriter { private static final String EMPTY_SCHEMA_PATH = ""; private BufferAllocator allocator; @BeforeEach public void init() { allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testPromoteToUnion() throws Exception { try (final NonNullableStructVector container = NonNullableStructVector.empty(EMPTY_SCHEMA_PATH, allocator); final StructVector v = container.addOrGetStruct("test"); final PromotableWriter writer = new PromotableWriter(v, container)) { container.allocateNew(); writer.start(); writer.setPosition(0); writer.bit("A").writeBit(0); writer.setPosition(1); writer.bit("A").writeBit(1); writer.decimal("dec", 10, 10); writer.setPosition(2); writer.integer("A").writeInt(10); // we don't write anything in 3 writer.setPosition(4); writer.integer("A").writeInt(100); writer.setPosition(5); writer.timeStampMilliTZ("A").writeTimeStampMilliTZ(123123); // Also try the holder version for timeStampMilliTZ writer.setPosition(6); TimeStampMilliTZHolder tsmtzHolder = new TimeStampMilliTZHolder(); // This has to be UTC since the vector above was initialized using the non holder // version that defaults to UTC. tsmtzHolder.timezone = "UTC"; tsmtzHolder.value = 12345L; writer.timeStampMilliTZ("A").write(tsmtzHolder); writer.setPosition(7); DurationHolder durationHolder = new DurationHolder(); durationHolder.unit = TimeUnit.SECOND; durationHolder.value = 444413; writer.duration("A").write(durationHolder); writer.setPosition(8); ArrowBuf buf = allocator.buffer(4); buf.setInt(0, 18978); FixedSizeBinaryHolder binHolder = new FixedSizeBinaryHolder(); binHolder.byteWidth = 4; 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(11); final UnionVector uv = v.getChild("A", UnionVector.class); assertFalse(uv.isNull(0), "0 shouldn't be null"); assertEquals(false, uv.getObject(0)); assertFalse(uv.isNull(1), "1 shouldn't be null"); assertEquals(true, uv.getObject(1)); assertFalse(uv.isNull(2), "2 shouldn't be null"); assertEquals(10, uv.getObject(2)); assertNull(uv.getObject(3), "3 should be null"); assertFalse(uv.isNull(4), "4 shouldn't be null"); assertEquals(100, uv.getObject(4)); assertFalse(uv.isNull(5), "5 shouldn't be null"); assertEquals(123123L, uv.getObject(5)); assertFalse(uv.isNull(6), "6 shouldn't be null"); NullableTimeStampMilliTZHolder readBackHolder = new NullableTimeStampMilliTZHolder(); uv.getTimeStampMilliTZVector().get(6, readBackHolder); assertEquals(12345L, readBackHolder.value); assertEquals("UTC", readBackHolder.timezone); assertFalse(uv.isNull(7), "7 shouldn't be null"); assertEquals(444413L, ((java.time.Duration) uv.getObject(7)).getSeconds()); assertFalse(uv.isNull(8), "8 shouldn't be null"); assertEquals( 18978, ByteBuffer.wrap(uv.getFixedSizeBinaryVector().get(8)) .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(); ComplexWriterImpl newWriter = new ComplexWriterImpl(EMPTY_SCHEMA_PATH, container); StructWriter newStructWriter = newWriter.rootAsStruct(); newStructWriter.start(); newStructWriter.setPosition(2); newStructWriter.integer("A").writeInt(10); Field childField1 = container.getField().getChildren().get(0).getChildren().get(0); Field childField2 = container.getField().getChildren().get(0).getChildren().get(1); assertEquals( ArrowTypeID.Union, childField1.getType().getTypeID(), "Child field should be union type: " + childField1.getName()); assertEquals( ArrowTypeID.Decimal, childField2.getType().getTypeID(), "Child field should be decimal type: " + childField2.getName()); buf.close(); } } @Test public void testNoPromoteFloat4ToUnionWithNull() throws Exception { try (final NonNullableStructVector container = NonNullableStructVector.empty(EMPTY_SCHEMA_PATH, allocator); final StructVector v = container.addOrGetStruct("test"); final PromotableWriter writer = new PromotableWriter(v, container)) { container.allocateNew(); writer.start(); writer.list("list").startList(); writer.list("list").endList(); writer.end(); FieldType childTypeOfListInContainer = container .getField() .getChildren() .get(0) .getChildren() .get(0) .getChildren() .get(0) .getFieldType(); // create a listvector with same type as list in container to, say, hold a copy // this will be a nullvector ListVector lv = ListVector.empty("name", allocator); lv.addOrGetVector(childTypeOfListInContainer); assertEquals(childTypeOfListInContainer.getType(), Types.MinorType.NULL.getType()); assertEquals( lv.getChildrenFromFields().get(0).getMinorType().getType(), Types.MinorType.NULL.getType()); writer.start(); writer.list("list").startList(); writer.list("list").float4().writeFloat4(1.36f); writer.list("list").endList(); writer.end(); container.setValueCount(2); childTypeOfListInContainer = container .getField() .getChildren() .get(0) .getChildren() .get(0) .getChildren() .get(0) .getFieldType(); // repeat but now the type in container has been changed from null to float // we expect same behaviour from listvector lv.addOrGetVector(childTypeOfListInContainer); assertEquals(childTypeOfListInContainer.getType(), Types.MinorType.FLOAT4.getType()); assertEquals( lv.getChildrenFromFields().get(0).getMinorType().getType(), Types.MinorType.FLOAT4.getType()); lv.close(); } } @Test public void testNoPromoteTimeStampMilliTZToUnionWithNull() throws Exception { try (final NonNullableStructVector container = NonNullableStructVector.empty(EMPTY_SCHEMA_PATH, allocator); final StructVector v = container.addOrGetStruct("test"); final PromotableWriter writer = new PromotableWriter(v, container)) { container.allocateNew(); writer.start(); writer.list("list").startList(); writer.list("list").endList(); writer.end(); FieldType childTypeOfListInContainer = container .getField() .getChildren() .get(0) .getChildren() .get(0) .getChildren() .get(0) .getFieldType(); // create a listvector with same type as list in container to, say, hold a copy // this will be a nullvector ListVector lv = ListVector.empty("name", allocator); lv.addOrGetVector(childTypeOfListInContainer); assertEquals(childTypeOfListInContainer.getType(), Types.MinorType.NULL.getType()); assertEquals( lv.getChildrenFromFields().get(0).getMinorType().getType(), Types.MinorType.NULL.getType()); writer.start(); writer.list("list").startList(); TimeStampMilliTZHolder holder = new TimeStampMilliTZHolder(); holder.value = 12341234L; holder.timezone = "FakeTimeZone"; writer.list("list").timeStampMilliTZ().write(holder); // Test that we get an exception when the timezone doesn't match holder.timezone = "SomeTimeZone"; IllegalArgumentException ex = assertThrows( IllegalArgumentException.class, () -> writer.list("list").timeStampMilliTZ().write(holder)); assertEquals( "holder.timezone: SomeTimeZone not equal to vector timezone: FakeTimeZone", ex.getMessage()); writer.list("list").endList(); writer.end(); container.setValueCount(2); childTypeOfListInContainer = container .getField() .getChildren() .get(0) .getChildren() .get(0) .getChildren() .get(0) .getFieldType(); // repeat but now the type in container has been changed from null to float // we expect same behaviour from listvector lv.addOrGetVector(childTypeOfListInContainer); assertEquals( childTypeOfListInContainer.getType(), new ArrowType.Timestamp(TimeUnit.MILLISECOND, "FakeTimeZone")); assertEquals( lv.getChildrenFromFields().get(0).getField().getType(), new ArrowType.Timestamp(TimeUnit.MILLISECOND, "FakeTimeZone")); lv.close(); } } @Test public void testNoPromoteDurationToUnionWithNull() throws Exception { try (final NonNullableStructVector container = NonNullableStructVector.empty(EMPTY_SCHEMA_PATH, allocator); final StructVector v = container.addOrGetStruct("test"); final PromotableWriter writer = new PromotableWriter(v, container)) { container.allocateNew(); writer.start(); writer.list("list").startList(); writer.list("list").endList(); writer.end(); FieldType childTypeOfListInContainer = container .getField() .getChildren() .get(0) .getChildren() .get(0) .getChildren() .get(0) .getFieldType(); // create a listvector with same type as list in container to, say, hold a copy // this will be a nullvector ListVector lv = ListVector.empty("name", allocator); lv.addOrGetVector(childTypeOfListInContainer); assertEquals(childTypeOfListInContainer.getType(), Types.MinorType.NULL.getType()); assertEquals( lv.getChildrenFromFields().get(0).getMinorType().getType(), Types.MinorType.NULL.getType()); writer.start(); writer.list("list").startList(); DurationHolder holder = new DurationHolder(); holder.unit = TimeUnit.NANOSECOND; holder.value = 567657L; writer.list("list").duration().write(holder); // Test that we get an exception when the unit doesn't match holder.unit = TimeUnit.MICROSECOND; IllegalArgumentException ex = assertThrows( IllegalArgumentException.class, () -> writer.list("list").duration().write(holder)); assertEquals( "holder.unit: MICROSECOND not equal to vector unit: NANOSECOND", ex.getMessage()); writer.list("list").endList(); writer.end(); container.setValueCount(2); childTypeOfListInContainer = container .getField() .getChildren() .get(0) .getChildren() .get(0) .getChildren() .get(0) .getFieldType(); // repeat but now the type in container has been changed from null to float // we expect same behaviour from listvector lv.addOrGetVector(childTypeOfListInContainer); assertEquals( childTypeOfListInContainer.getType(), new ArrowType.Duration(TimeUnit.NANOSECOND)); assertEquals( lv.getChildrenFromFields().get(0).getField().getType(), new ArrowType.Duration(TimeUnit.NANOSECOND)); lv.close(); } } @Test public void testNoPromoteFixedSizeBinaryToUnionWithNull() throws Exception { try (final NonNullableStructVector container = NonNullableStructVector.empty(EMPTY_SCHEMA_PATH, allocator); final StructVector v = container.addOrGetStruct("test"); final PromotableWriter writer = new PromotableWriter(v, container)) { container.allocateNew(); writer.start(); writer.list("list").startList(); writer.list("list").endList(); writer.end(); FieldType childTypeOfListInContainer = container .getField() .getChildren() .get(0) .getChildren() .get(0) .getChildren() .get(0) .getFieldType(); // create a listvector with same type as list in container to, say, hold a copy // this will be a nullvector ListVector lv = ListVector.empty("name", allocator); lv.addOrGetVector(childTypeOfListInContainer); assertEquals(childTypeOfListInContainer.getType(), Types.MinorType.NULL.getType()); assertEquals( lv.getChildrenFromFields().get(0).getMinorType().getType(), Types.MinorType.NULL.getType()); writer.start(); writer.list("list").startList(); ArrowBuf buf = allocator.buffer(4); buf.setInt(0, 22222); FixedSizeBinaryHolder holder = new FixedSizeBinaryHolder(); holder.byteWidth = 4; holder.buffer = buf; writer.list("list").fixedSizeBinary().write(holder); // Test that we get an exception when the unit doesn't match holder.byteWidth = 7; IllegalArgumentException ex = assertThrows( IllegalArgumentException.class, () -> writer.list("list").fixedSizeBinary().write(holder)); assertEquals("holder.byteWidth: 7 not equal to vector byteWidth: 4", ex.getMessage()); writer.list("list").endList(); writer.end(); container.setValueCount(2); childTypeOfListInContainer = container .getField() .getChildren() .get(0) .getChildren() .get(0) .getChildren() .get(0) .getFieldType(); // repeat but now the type in container has been changed from null to float // we expect same behaviour from listvector lv.addOrGetVector(childTypeOfListInContainer); assertEquals(childTypeOfListInContainer.getType(), new ArrowType.FixedSizeBinary(4)); assertEquals( lv.getChildrenFromFields().get(0).getField().getType(), new ArrowType.FixedSizeBinary(4)); lv.close(); buf.close(); } } @Test public void testPromoteLargeVarCharHelpersOnStruct() throws Exception { try (final NonNullableStructVector container = NonNullableStructVector.empty(EMPTY_SCHEMA_PATH, allocator); final StructVector v = container.addOrGetStruct("test"); final PromotableWriter writer = new PromotableWriter(v, container)) { container.allocateNew(); writer.start(); writer.setPosition(0); writer.largeVarChar("c").writeLargeVarChar(new Text("foo")); writer.setPosition(1); writer.largeVarChar("c").writeLargeVarChar("foo2"); writer.end(); final LargeVarCharVector uv = v.getChild("c", LargeVarCharVector.class); assertEquals("foo", Objects.requireNonNull(uv.getObject(0)).toString()); assertEquals("foo2", Objects.requireNonNull(uv.getObject(1)).toString()); } } @Test public void testPromoteVarCharHelpersOnStruct() throws Exception { try (final NonNullableStructVector container = NonNullableStructVector.empty(EMPTY_SCHEMA_PATH, allocator); final StructVector v = container.addOrGetStruct("test"); final PromotableWriter writer = new PromotableWriter(v, container)) { container.allocateNew(); writer.start(); writer.setPosition(0); writer.varChar("c").writeVarChar(new Text("foo")); writer.setPosition(1); writer.varChar("c").writeVarChar("foo2"); writer.end(); final VarCharVector uv = v.getChild("c", VarCharVector.class); assertEquals("foo", Objects.requireNonNull(uv.getObject(0)).toString()); assertEquals("foo2", Objects.requireNonNull(uv.getObject(1)).toString()); } } @Test public void testPromoteVarCharHelpersDirect() throws Exception { try (final NonNullableStructVector container = NonNullableStructVector.empty(EMPTY_SCHEMA_PATH, allocator); final StructVector v = container.addOrGetStruct("test"); final PromotableWriter writer = new PromotableWriter(v, container)) { container.allocateNew(); writer.start(); writer.setPosition(0); writer.writeVarChar(new Text("foo")); writer.setPosition(1); writer.writeVarChar("foo2"); writer.end(); // The "test" vector in the parent container should have been replaced with a UnionVector. UnionVector promotedVector = container.getChild("test", UnionVector.class); VarCharVector vector = promotedVector.getVarCharVector(); assertEquals("foo", Objects.requireNonNull(vector.getObject(0)).toString()); assertEquals("foo2", Objects.requireNonNull(vector.getObject(1)).toString()); } } @Test public void testPromoteLargeVarCharHelpersDirect() throws Exception { try (final NonNullableStructVector container = NonNullableStructVector.empty(EMPTY_SCHEMA_PATH, allocator); final StructVector v = container.addOrGetStruct("test"); final PromotableWriter writer = new PromotableWriter(v, container)) { container.allocateNew(); writer.start(); writer.setPosition(0); writer.writeLargeVarChar(new Text("foo")); writer.setPosition(1); writer.writeLargeVarChar("foo2"); writer.end(); // The "test" vector in the parent container should have been replaced with a UnionVector. UnionVector promotedVector = container.getChild("test", UnionVector.class); LargeVarCharVector vector = promotedVector.getLargeVarCharVector(); assertEquals("foo", Objects.requireNonNull(vector.getObject(0)).toString()); assertEquals("foo2", Objects.requireNonNull(vector.getObject(1)).toString()); } } @Test public void testPromoteVarBinaryHelpersOnStruct() throws Exception { try (final NonNullableStructVector container = NonNullableStructVector.empty(EMPTY_SCHEMA_PATH, allocator); final StructVector v = container.addOrGetStruct("test"); final PromotableWriter writer = new PromotableWriter(v, container)) { container.allocateNew(); writer.start(); writer.setPosition(0); writer.varBinary("c").writeVarBinary("row1".getBytes(StandardCharsets.UTF_8)); writer.setPosition(1); writer .varBinary("c") .writeVarBinary( "row2".getBytes(StandardCharsets.UTF_8), 0, "row2".getBytes(StandardCharsets.UTF_8).length); writer.setPosition(2); writer .varBinary("c") .writeVarBinary(ByteBuffer.wrap("row3".getBytes(StandardCharsets.UTF_8))); writer.setPosition(3); writer .varBinary("c") .writeVarBinary( ByteBuffer.wrap("row4".getBytes(StandardCharsets.UTF_8)), 0, "row4".getBytes(StandardCharsets.UTF_8).length); writer.end(); final VarBinaryVector uv = v.getChild("c", VarBinaryVector.class); assertEquals("row1", new String(Objects.requireNonNull(uv.get(0)), StandardCharsets.UTF_8)); assertEquals("row2", new String(Objects.requireNonNull(uv.get(1)), StandardCharsets.UTF_8)); assertEquals("row3", new String(Objects.requireNonNull(uv.get(2)), StandardCharsets.UTF_8)); assertEquals("row4", new String(Objects.requireNonNull(uv.get(3)), StandardCharsets.UTF_8)); } } @Test public void testPromoteVarBinaryHelpersDirect() throws Exception { try (final NonNullableStructVector container = NonNullableStructVector.empty(EMPTY_SCHEMA_PATH, allocator); final StructVector v = container.addOrGetStruct("test"); final PromotableWriter writer = new PromotableWriter(v, container)) { container.allocateNew(); writer.start(); writer.setPosition(0); writer.writeVarBinary("row1".getBytes(StandardCharsets.UTF_8)); writer.setPosition(1); writer.writeVarBinary( "row2".getBytes(StandardCharsets.UTF_8), 0, "row2".getBytes(StandardCharsets.UTF_8).length); writer.setPosition(2); writer.writeVarBinary(ByteBuffer.wrap("row3".getBytes(StandardCharsets.UTF_8))); writer.setPosition(3); writer.writeVarBinary( ByteBuffer.wrap("row4".getBytes(StandardCharsets.UTF_8)), 0, "row4".getBytes(StandardCharsets.UTF_8).length); writer.end(); // The "test" vector in the parent container should have been replaced with a UnionVector. UnionVector promotedVector = container.getChild("test", UnionVector.class); VarBinaryVector uv = promotedVector.getVarBinaryVector(); assertEquals("row1", new String(Objects.requireNonNull(uv.get(0)), StandardCharsets.UTF_8)); assertEquals("row2", new String(Objects.requireNonNull(uv.get(1)), StandardCharsets.UTF_8)); assertEquals("row3", new String(Objects.requireNonNull(uv.get(2)), StandardCharsets.UTF_8)); assertEquals("row4", new String(Objects.requireNonNull(uv.get(3)), StandardCharsets.UTF_8)); } } @Test public void testPromoteLargeVarBinaryHelpersOnStruct() throws Exception { try (final NonNullableStructVector container = NonNullableStructVector.empty(EMPTY_SCHEMA_PATH, allocator); final StructVector v = container.addOrGetStruct("test"); final PromotableWriter writer = new PromotableWriter(v, container)) { container.allocateNew(); writer.start(); writer.setPosition(0); writer.largeVarBinary("c").writeLargeVarBinary("row1".getBytes(StandardCharsets.UTF_8)); writer.setPosition(1); writer .largeVarBinary("c") .writeLargeVarBinary( "row2".getBytes(StandardCharsets.UTF_8), 0, "row2".getBytes(StandardCharsets.UTF_8).length); writer.setPosition(2); writer .largeVarBinary("c") .writeLargeVarBinary(ByteBuffer.wrap("row3".getBytes(StandardCharsets.UTF_8))); writer.setPosition(3); writer .largeVarBinary("c") .writeLargeVarBinary( ByteBuffer.wrap("row4".getBytes(StandardCharsets.UTF_8)), 0, "row4".getBytes(StandardCharsets.UTF_8).length); writer.end(); final LargeVarBinaryVector uv = v.getChild("c", LargeVarBinaryVector.class); assertEquals("row1", new String(Objects.requireNonNull(uv.get(0)), StandardCharsets.UTF_8)); assertEquals("row2", new String(Objects.requireNonNull(uv.get(1)), StandardCharsets.UTF_8)); assertEquals("row3", new String(Objects.requireNonNull(uv.get(2)), StandardCharsets.UTF_8)); assertEquals("row4", new String(Objects.requireNonNull(uv.get(3)), StandardCharsets.UTF_8)); } } @Test public void testPromoteLargeVarBinaryHelpersDirect() throws Exception { try (final NonNullableStructVector container = NonNullableStructVector.empty(EMPTY_SCHEMA_PATH, allocator); final StructVector v = container.addOrGetStruct("test"); final PromotableWriter writer = new PromotableWriter(v, container)) { container.allocateNew(); writer.start(); writer.setPosition(0); writer.writeLargeVarBinary("row1".getBytes(StandardCharsets.UTF_8)); writer.setPosition(1); writer.writeLargeVarBinary( "row2".getBytes(StandardCharsets.UTF_8), 0, "row2".getBytes(StandardCharsets.UTF_8).length); writer.setPosition(2); writer.writeLargeVarBinary(ByteBuffer.wrap("row3".getBytes(StandardCharsets.UTF_8))); writer.setPosition(3); writer.writeLargeVarBinary( ByteBuffer.wrap("row4".getBytes(StandardCharsets.UTF_8)), 0, "row4".getBytes(StandardCharsets.UTF_8).length); writer.end(); // The "test" vector in the parent container should have been replaced with a UnionVector. UnionVector promotedVector = container.getChild("test", UnionVector.class); LargeVarBinaryVector uv = promotedVector.getLargeVarBinaryVector(); assertEquals("row1", new String(Objects.requireNonNull(uv.get(0)), StandardCharsets.UTF_8)); assertEquals("row2", new String(Objects.requireNonNull(uv.get(1)), StandardCharsets.UTF_8)); assertEquals("row3", new String(Objects.requireNonNull(uv.get(2)), StandardCharsets.UTF_8)); 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)); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/complex/writer/000077500000000000000000000000001515451615200305165ustar00rootroot00000000000000TestComplexWriter.java000066400000000000000000003057261515451615200347630ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/complex/writer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.writer; 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; import java.math.BigDecimal; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.time.LocalDateTime; import java.util.ArrayList; 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; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.LargeVarCharVector; import org.apache.arrow.vector.SchemaChangeCallBack; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.ListViewVector; import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.NonNullableStructVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.complex.impl.ComplexWriterImpl; import org.apache.arrow.vector.complex.impl.NullableStructReaderImpl; import org.apache.arrow.vector.complex.impl.NullableStructWriter; import org.apache.arrow.vector.complex.impl.SingleStructReaderImpl; import org.apache.arrow.vector.complex.impl.SingleStructWriter; import org.apache.arrow.vector.complex.impl.UnionListReader; import org.apache.arrow.vector.complex.impl.UnionListViewReader; import org.apache.arrow.vector.complex.impl.UnionListViewWriter; import org.apache.arrow.vector.complex.impl.UnionListWriter; import org.apache.arrow.vector.complex.impl.UnionMapReader; import org.apache.arrow.vector.complex.impl.UnionReader; import org.apache.arrow.vector.complex.impl.UnionWriter; import org.apache.arrow.vector.complex.reader.BaseReader.StructReader; import org.apache.arrow.vector.complex.reader.BigIntReader; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.complex.reader.Float4Reader; 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; import org.apache.arrow.vector.holders.IntHolder; import org.apache.arrow.vector.holders.NullableDurationHolder; 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; import org.apache.arrow.vector.types.pojo.ArrowType.Struct; import org.apache.arrow.vector.types.pojo.ArrowType.Timestamp; import org.apache.arrow.vector.types.pojo.ArrowType.Union; import org.apache.arrow.vector.types.pojo.ArrowType.Utf8; 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.DecimalUtility; import org.apache.arrow.vector.util.JsonStringArrayList; 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; public class TestComplexWriter { private BufferAllocator allocator; private static final int COUNT = 100; @BeforeEach public void init() { allocator = new RootAllocator(Integer.MAX_VALUE); } @AfterEach public void terminate() throws Exception { allocator.close(); } /* Test Utils */ private void checkNullableStruct(NonNullableStructVector structVector) { StructReader rootReader = new SingleStructReaderImpl(structVector).reader("root"); for (int i = 0; i < COUNT; i++) { rootReader.setPosition(i); assertTrue(rootReader.isSet(), "index is set: " + i); FieldReader struct = rootReader.reader("struct"); if (i % 2 == 0) { assertTrue(struct.isSet(), "index is set: " + i); assertNotNull(struct.readObject(), "index is set: " + i); assertEquals(i, struct.reader("nested").readLong().longValue()); } else { assertFalse(struct.isSet(), "index is not set: " + i); assertNull(struct.readObject(), "index is not set: " + i); } } } private void createListTypeVectorWithScalarType(FieldWriter writer) { for (int i = 0; i < COUNT; i++) { writer.startList(); for (int j = 0; j < i % 7; j++) { if (j % 2 == 0) { writer.writeInt(j); } else { IntHolder holder = new IntHolder(); holder.value = j; writer.write(holder); } } writer.endList(); } } private void checkListTypeVectorWithScalarType(FieldReader reader) { for (int i = 0; i < COUNT; i++) { reader.setPosition(i); for (int j = 0; j < i % 7; j++) { reader.next(); assertEquals(j, reader.reader().readInteger().intValue()); } } } private void createListTypeVectorWithScalarNull(FieldWriter writer) { for (int i = 0; i < COUNT; i++) { writer.startList(); for (int j = 0; j < i % 7; j++) { if (j % 2 == 0) { writer.writeNull(); } else { IntHolder holder = new IntHolder(); holder.value = j; writer.write(holder); } } writer.endList(); } } private void checkListTypeVectorWithScalarNull(FieldReader reader) { for (int i = 0; i < COUNT; i++) { reader.setPosition(i); for (int j = 0; j < i % 7; j++) { reader.next(); if (j % 2 == 0) { assertFalse(reader.reader().isSet(), "index is set: " + j); } else { assertTrue(reader.reader().isSet(), "index is not set: " + j); assertEquals(j, reader.reader().readInteger().intValue()); } } } } private void createListTypeVectorWithDecimalType(FieldWriter writer, DecimalHolder holder) { holder.buffer = allocator.buffer(DecimalVector.TYPE_WIDTH); ArrowType arrowType = new ArrowType.Decimal(10, 0, 128); for (int i = 0; i < COUNT; i++) { writer.startList(); for (int j = 0; j < i % 7; j++) { if (j % 4 == 0) { writer.writeDecimal(new BigDecimal(j)); } else if (j % 4 == 1) { DecimalUtility.writeBigDecimalToArrowBuf( new BigDecimal(j), holder.buffer, 0, DecimalVector.TYPE_WIDTH); holder.start = 0; holder.scale = 0; holder.precision = 10; writer.write(holder); } else if (j % 4 == 2) { DecimalUtility.writeBigDecimalToArrowBuf( new BigDecimal(j), holder.buffer, 0, DecimalVector.TYPE_WIDTH); writer.writeDecimal(0, holder.buffer, arrowType); } else { byte[] value = BigDecimal.valueOf(j).unscaledValue().toByteArray(); writer.writeBigEndianBytesToDecimal(value, arrowType); } } writer.endList(); } } private void checkListTypeVectorWithDecimalType(FieldReader reader) { for (int i = 0; i < COUNT; i++) { reader.setPosition(i); for (int j = 0; j < i % 7; j++) { reader.next(); Object expected = new BigDecimal(j); Object actual = reader.reader().readBigDecimal(); assertEquals(expected, actual); } } } private void createListTypeVectorWithTimeStampMilliTZType(FieldWriter writer) { for (int i = 0; i < COUNT; i++) { writer.startList(); for (int j = 0; j < i % 7; j++) { if (j % 2 == 0) { writer.writeNull(); } else { TimeStampMilliTZHolder holder = new TimeStampMilliTZHolder(); holder.timezone = "FakeTimeZone"; holder.value = j; writer.timeStampMilliTZ().write(holder); } } writer.endList(); } } private void checkListTypeVectorWithTimeStampMilliTZType(FieldReader reader) { for (int i = 0; i < COUNT; i++) { reader.setPosition(i); for (int j = 0; j < i % 7; j++) { reader.next(); if (j % 2 == 0) { assertFalse(reader.reader().isSet(), "index is set: " + j); } else { NullableTimeStampMilliTZHolder actual = new NullableTimeStampMilliTZHolder(); reader.reader().read(actual); assertEquals(j, actual.value); assertEquals("FakeTimeZone", actual.timezone); } } } } private void createNullsWithListWriters(FieldWriter writer) { for (int i = 0; i < COUNT; i++) { writer.setPosition(i); if (i % 2 == 0) { writer.startList(); if (i % 4 == 0) { writer.integer().writeNull(); } else { writer.integer().writeInt(i); writer.integer().writeInt(i * 2); } writer.endList(); } else { writer.writeNull(); } } } private void checkNullsWithListWriters(FieldReader reader) { for (int i = 0; i < COUNT; i++) { reader.setPosition(i); if (i % 2 == 0) { assertTrue(reader.isSet()); reader.next(); if (i % 4 == 0) { assertNull(reader.reader().readInteger()); } else { assertEquals(i, reader.reader().readInteger().intValue()); reader.next(); assertEquals(i * 2, reader.reader().readInteger().intValue()); } } else { assertFalse(reader.isSet()); } } } /* Test Cases */ @Test public void simpleNestedTypes() { NonNullableStructVector parent = populateStructVector(null); StructReader rootReader = new SingleStructReaderImpl(parent).reader("root"); for (int i = 0; i < COUNT; i++) { rootReader.setPosition(i); assertEquals(i, rootReader.reader("int").readInteger().intValue()); assertEquals(i, rootReader.reader("bigInt").readLong().longValue()); } parent.close(); } @Test public void transferPairSchemaChange() { SchemaChangeCallBack callBack1 = new SchemaChangeCallBack(); try (NonNullableStructVector parent = populateStructVector(callBack1)) { ComplexWriter writer = new ComplexWriterImpl("newWriter", parent); StructWriter rootWriter = writer.rootAsStruct(); IntWriter intWriter = rootWriter.integer("newInt"); intWriter.writeInt(1); writer.setValueCount(1); assertTrue(callBack1.getSchemaChangedAndReset()); // The second vector should not have registered a schema change assertFalse(callBack1.getSchemaChangedAndReset()); } } private NonNullableStructVector populateStructVector(CallBack callBack) { NonNullableStructVector parent = new NonNullableStructVector( "parent", allocator, new FieldType(false, Struct.INSTANCE, null, null), callBack); ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); IntWriter intWriter = rootWriter.integer("int"); BigIntWriter bigIntWriter = rootWriter.bigInt("bigInt"); for (int i = 0; i < COUNT; i++) { rootWriter.start(); intWriter.writeInt(i); bigIntWriter.writeBigInt(i); rootWriter.end(); } writer.setValueCount(COUNT); return parent; } @Test public void nullableStruct() { try (NonNullableStructVector structVector = NonNullableStructVector.empty("parent", allocator)) { ComplexWriter writer = new ComplexWriterImpl("root", structVector); StructWriter rootWriter = writer.rootAsStruct(); for (int i = 0; i < COUNT; i++) { rootWriter.start(); if (i % 2 == 0) { StructWriter structWriter = rootWriter.struct("struct"); structWriter.setPosition(i); structWriter.start(); structWriter.bigInt("nested").writeBigInt(i); structWriter.end(); } rootWriter.end(); } writer.setValueCount(COUNT); checkNullableStruct(structVector); } } /** * This test is similar to {@link #nullableStruct()} ()} but we get the inner struct writer once * at the beginning. */ @Test public void nullableStruct2() { try (NonNullableStructVector structVector = NonNullableStructVector.empty("parent", allocator)) { ComplexWriter writer = new ComplexWriterImpl("root", structVector); StructWriter rootWriter = writer.rootAsStruct(); StructWriter structWriter = rootWriter.struct("struct"); for (int i = 0; i < COUNT; i++) { rootWriter.start(); if (i % 2 == 0) { structWriter.setPosition(i); structWriter.start(); structWriter.bigInt("nested").writeBigInt(i); structWriter.end(); } rootWriter.end(); } writer.setValueCount(COUNT); checkNullableStruct(structVector); } } @Test public void testList() { try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); rootWriter.start(); rootWriter.bigInt("int").writeBigInt(0); rootWriter.list("list").startList(); rootWriter.list("list").bigInt().writeBigInt(0); rootWriter.list("list").endList(); rootWriter.end(); rootWriter.start(); rootWriter.bigInt("int").writeBigInt(1); rootWriter.end(); writer.setValueCount(2); StructReader rootReader = new SingleStructReaderImpl(parent).reader("root"); rootReader.setPosition(0); assertTrue(rootReader.reader("list").isSet(), "row 0 list is not set"); assertEquals(Long.valueOf(0), rootReader.reader("list").reader().readLong()); rootReader.setPosition(1); assertFalse(rootReader.reader("list").isSet(), "row 1 list is set"); } } private void createListTypeVectorWithDurationType(FieldWriter writer) { for (int i = 0; i < COUNT; i++) { writer.startList(); for (int j = 0; j < i % 7; j++) { if (j % 2 == 0) { writer.writeNull(); } else { DurationHolder holder = new DurationHolder(); holder.unit = TimeUnit.MICROSECOND; holder.value = j; writer.duration().write(holder); } } writer.endList(); } } private void checkListTypeVectorWithDurationType(FieldReader reader) { for (int i = 0; i < COUNT; i++) { reader.setPosition(i); for (int j = 0; j < i % 7; j++) { reader.next(); if (j % 2 == 0) { assertFalse(reader.reader().isSet(), "index is set: " + j); } else { NullableDurationHolder actual = new NullableDurationHolder(); reader.reader().read(actual); assertEquals(TimeUnit.MICROSECOND, actual.unit); assertEquals(j, actual.value); } } } } private void createScalarTypeVectorWithNullableType(FieldWriter writer) { for (int i = 0; i < COUNT; i++) { if (i % 2 == 0) { writer.setPosition(i); writer.startList(); for (int j = 0; j < i % 7; j++) { writer.writeInt(j); } writer.endList(); } } } private void checkScalarTypeVectorWithNullableType(FieldReader reader) { for (int i = 0; i < COUNT; i++) { reader.setPosition(i); if (i % 2 == 0) { assertTrue(reader.isSet(), "index is set: " + i); assertEquals(i % 7, ((List) reader.readObject()).size(), "correct length at: " + i); } else { assertFalse(reader.isSet(), "index is not set: " + i); assertNull(reader.readObject(), "index is not set: " + i); } } } private void createListTypeVectorWithStructType( FieldWriter fieldWriter, StructWriter structWriter) { for (int i = 0; i < COUNT; i++) { fieldWriter.startList(); for (int j = 0; j < i % 7; j++) { structWriter.start(); structWriter.integer("int").writeInt(j); structWriter.bigInt("bigInt").writeBigInt(j); structWriter.end(); } fieldWriter.endList(); } } private void checkListTypeVectorWithStructType(FieldReader reader) { for (int i = 0; i < COUNT; i++) { reader.setPosition(i); for (int j = 0; j < i % 7; j++) { reader.next(); assertEquals(j, reader.reader().reader("int").readInteger().intValue(), "record: " + i); assertEquals(j, reader.reader().reader("bigInt").readLong().longValue()); } } } private void checkListOfListTypes(final FieldReader reader) { for (int i = 0; i < COUNT; i++) { reader.setPosition(i); for (int j = 0; j < i % 7; j++) { reader.next(); FieldReader innerListReader = reader.reader(); for (int k = 0; k < i % 13; k++) { innerListReader.next(); assertEquals(k, innerListReader.reader().readInteger().intValue(), "record: " + i); } } } } private void checkUnionListType(FieldReader reader) { for (int i = 0; i < COUNT; i++) { reader.setPosition(i); for (int j = 0; j < i % 7; j++) { reader.next(); FieldReader innerListReader = reader.reader(); for (int k = 0; k < i % 13; k++) { innerListReader.next(); if (k % 2 == 0) { assertEquals(k, innerListReader.reader().readInteger().intValue(), "record: " + i); } else { assertEquals(k, innerListReader.reader().readLong().longValue(), "record: " + i); } } } } } private static void createListTypeVectorWithMapType(FieldWriter writer) { MapWriter innerMapWriter = writer.map(true); for (int i = 0; i < COUNT; i++) { writer.startList(); for (int j = 0; j < i % 7; j++) { innerMapWriter.startMap(); for (int k = 0; k < i % 13; k++) { innerMapWriter.startEntry(); innerMapWriter.key().integer().writeInt(k); if (k % 2 == 0) { innerMapWriter.value().bigInt().writeBigInt(k); } innerMapWriter.endEntry(); } innerMapWriter.endMap(); } writer.endList(); } } private void checkListTypeMap(FieldReader reader) { for (int i = 0; i < COUNT; i++) { reader.setPosition(i); for (int j = 0; j < i % 7; j++) { reader.next(); UnionMapReader mapReader = (UnionMapReader) reader.reader(); for (int k = 0; k < i % 13; k++) { mapReader.next(); assertEquals(k, mapReader.key().readInteger().intValue(), "record key: " + i); if (k % 2 == 0) { assertEquals(k, mapReader.value().readLong().longValue(), "record value: " + i); } else { assertNull(mapReader.value().readLong(), "record value: " + i); } } } } } /* Test Cases */ private void createListTypeVectorWithFixedSizeBinaryType( FieldWriter writer, List buffers) { for (int i = 0; i < COUNT; i++) { writer.startList(); for (int j = 0; j < i % 7; j++) { if (j % 2 == 0) { writer.writeNull(); } else { ArrowBuf buf = allocator.buffer(4); buf.setInt(0, j); FixedSizeBinaryHolder holder = new FixedSizeBinaryHolder(); holder.byteWidth = 4; holder.buffer = buf; writer.fixedSizeBinary().write(holder); buffers.add(buf); } } writer.endList(); } } private void checkListTypeVectorWithFixedSizeBinaryType(FieldReader reader) { for (int i = 0; i < COUNT; i++) { reader.setPosition(i); for (int j = 0; j < i % 7; j++) { reader.next(); if (j % 2 == 0) { assertFalse(reader.reader().isSet(), "index is set: " + j); } else { NullableFixedSizeBinaryHolder actual = new NullableFixedSizeBinaryHolder(); reader.reader().read(actual); assertEquals(j, actual.buffer.getInt(0)); assertEquals(4, actual.byteWidth); } } } } @Test public void listScalarType() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); createListTypeVectorWithScalarType(listWriter); listWriter.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); // validate checkListTypeVectorWithScalarType(listReader); } } @Test public void listViewScalarType() { try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { listViewVector.allocateNew(); UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); createListTypeVectorWithScalarType(listViewWriter); listViewWriter.setValueCount(COUNT); UnionListViewReader listViewReader = new UnionListViewReader(listViewVector); // validate checkListTypeVectorWithScalarType(listViewReader); } } @Test public void testListScalarNull() { /* Write to an integer list vector * each list of size 8 * and having its data values alternating between null and a non-null. * Read and verify */ try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); createListTypeVectorWithScalarNull(listWriter); listWriter.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); checkListTypeVectorWithScalarNull(listReader); } } @Test public void testListViewScalarNull() { /* Write to an integer list vector * each list of size 8 * and having its data values alternating between null and a non-null. * Read and verify */ try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { listViewVector.allocateNew(); UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); createListTypeVectorWithScalarNull(listViewWriter); listViewWriter.setValueCount(COUNT); UnionListViewReader listViewReader = new UnionListViewReader(listViewVector); checkListTypeVectorWithScalarNull(listViewReader); } } @Test public void listDecimalType() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); DecimalHolder holder = new DecimalHolder(); createListTypeVectorWithDecimalType(listWriter, holder); listWriter.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); checkListTypeVectorWithDecimalType(listReader); holder.buffer.close(); } } @Test public void listViewDecimalType() { try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { listViewVector.allocateNew(); UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); DecimalHolder holder = new DecimalHolder(); createListTypeVectorWithDecimalType(listViewWriter, holder); listViewWriter.setValueCount(COUNT); UnionListViewReader listViewReader = new UnionListViewReader(listViewVector); checkListTypeVectorWithDecimalType(listViewReader); holder.buffer.close(); } } @Test public void listTimeStampMilliTZType() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); createListTypeVectorWithTimeStampMilliTZType(listWriter); listWriter.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); checkListTypeVectorWithTimeStampMilliTZType(listReader); } } @Test public void listViewTimeStampMilliTZType() { try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { listViewVector.allocateNew(); UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); createListTypeVectorWithTimeStampMilliTZType(listViewWriter); listViewWriter.setValueCount(COUNT); UnionListViewReader listViewReader = new UnionListViewReader(listViewVector); checkListTypeVectorWithTimeStampMilliTZType(listViewReader); } } @Test public void listDurationType() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); createListTypeVectorWithDurationType(listWriter); listWriter.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); checkListTypeVectorWithDurationType(listReader); } } @Test public void listViewDurationType() { try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { listViewVector.allocateNew(); UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); createListTypeVectorWithDurationType(listViewWriter); listViewWriter.setValueCount(COUNT); UnionListViewReader listReader = new UnionListViewReader(listViewVector); checkListTypeVectorWithDurationType(listReader); } } @Test public void listFixedSizeBinaryType() throws Exception { List buffers = new ArrayList<>(); try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); createListTypeVectorWithFixedSizeBinaryType(listWriter, buffers); listWriter.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); checkListTypeVectorWithFixedSizeBinaryType(listReader); } AutoCloseables.close(buffers); } @Test public void listViewFixedSizeBinaryType() throws Exception { List buffers = new ArrayList<>(); try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { listViewVector.allocateNew(); UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); createListTypeVectorWithFixedSizeBinaryType(listViewWriter, buffers); listViewWriter.setValueCount(COUNT); UnionListViewReader listReader = new UnionListViewReader(listViewVector); checkListTypeVectorWithFixedSizeBinaryType(listReader); } AutoCloseables.close(buffers); } @Test public void listScalarTypeNullable() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); createScalarTypeVectorWithNullableType(listWriter); listWriter.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); checkScalarTypeVectorWithNullableType(listReader); } } @Test public void listViewScalarTypeNullable() { try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { listViewVector.allocateNew(); UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); createScalarTypeVectorWithNullableType(listViewWriter); listViewWriter.setValueCount(COUNT); UnionListViewReader listReader = new UnionListViewReader(listViewVector); checkScalarTypeVectorWithNullableType(listReader); } } @Test public void listStructType() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listViewWriter = new UnionListWriter(listVector); StructWriter structWriter = listViewWriter.struct(); createListTypeVectorWithStructType(listViewWriter, structWriter); listViewWriter.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); checkListTypeVectorWithStructType(listReader); } } @Test public void listViewStructType() { try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { listViewVector.allocateNew(); UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); StructWriter structWriter = listViewWriter.struct(); createListTypeVectorWithStructType(listViewWriter, structWriter); listViewWriter.setValueCount(COUNT); UnionListViewReader listReader = new UnionListViewReader(listViewVector); checkListTypeVectorWithStructType(listReader); } } @Test public void listListType() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); for (int i = 0; i < COUNT; i++) { listWriter.startList(); for (int j = 0; j < i % 7; j++) { ListWriter innerListWriter = listWriter.list(); innerListWriter.startList(); for (int k = 0; k < i % 13; k++) { innerListWriter.integer().writeInt(k); } innerListWriter.endList(); } listWriter.endList(); } listWriter.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); checkListOfListTypes(listReader); } } @Test public void listViewListType() { try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { listViewVector.allocateNew(); UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); for (int i = 0; i < COUNT; i++) { listViewWriter.startListView(); for (int j = 0; j < i % 7; j++) { ListWriter innerListWriter = listViewWriter.listView(); innerListWriter.startListView(); for (int k = 0; k < i % 13; k++) { innerListWriter.integer().writeInt(k); } innerListWriter.endListView(); } listViewWriter.endListView(); } listViewWriter.setValueCount(COUNT); UnionListViewReader listReader = new UnionListViewReader(listViewVector); checkListOfListTypes(listReader); } } /** * This test is similar to {@link #listListType()} but we get the inner list writer once at the * beginning. */ @Test public void listListType2() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); ListWriter innerListWriter = listWriter.list(); for (int i = 0; i < COUNT; i++) { listWriter.startList(); for (int j = 0; j < i % 7; j++) { innerListWriter.startList(); for (int k = 0; k < i % 13; k++) { innerListWriter.integer().writeInt(k); } innerListWriter.endList(); } listWriter.endList(); } listWriter.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); checkListOfListTypes(listReader); } } @Test public void listViewListType2() { try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { listViewVector.allocateNew(); UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); ListWriter innerListWriter = listViewWriter.list(); for (int i = 0; i < COUNT; i++) { listViewWriter.startListView(); for (int j = 0; j < i % 7; j++) { innerListWriter.startListView(); for (int k = 0; k < i % 13; k++) { innerListWriter.integer().writeInt(k); } innerListWriter.endListView(); } listViewWriter.endListView(); } listViewWriter.setValueCount(COUNT); UnionListViewReader listReader = new UnionListViewReader(listViewVector); checkListOfListTypes(listReader); } } @Test public void unionListListType() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); for (int i = 0; i < COUNT; i++) { listWriter.startList(); for (int j = 0; j < i % 7; j++) { ListWriter innerListWriter = listWriter.list(); innerListWriter.startList(); for (int k = 0; k < i % 13; k++) { if (k % 2 == 0) { innerListWriter.integer().writeInt(k); } else { innerListWriter.bigInt().writeBigInt(k); } } innerListWriter.endList(); } listWriter.endList(); } listWriter.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); checkUnionListType(listReader); } } @Test public void unionListViewListType() { try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { listViewVector.allocateNew(); UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); for (int i = 0; i < COUNT; i++) { listViewWriter.startList(); for (int j = 0; j < i % 7; j++) { ListWriter innerListWriter = listViewWriter.listView(); innerListWriter.startListView(); for (int k = 0; k < i % 13; k++) { if (k % 2 == 0) { innerListWriter.integer().writeInt(k); } else { innerListWriter.bigInt().writeBigInt(k); } } innerListWriter.endListView(); } listViewWriter.endListView(); } listViewWriter.setValueCount(COUNT); UnionListViewReader listViewReader = new UnionListViewReader(listViewVector); checkUnionListType(listViewReader); } } /** * This test is similar to {@link #unionListListType()} but we get the inner list writer once at * the beginning. */ @Test public void unionListListType2() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); ListWriter innerListWriter = listWriter.listView(); for (int i = 0; i < COUNT; i++) { listWriter.startList(); for (int j = 0; j < i % 7; j++) { innerListWriter.startList(); for (int k = 0; k < i % 13; k++) { if (k % 2 == 0) { innerListWriter.integer().writeInt(k); } else { innerListWriter.bigInt().writeBigInt(k); } } innerListWriter.endList(); } listWriter.endList(); } listWriter.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); checkUnionListType(listReader); } } /** * This test is similar to {@link #unionListViewListType()} but we get the inner list writer once * at the beginning. */ @Test public void unionListViewListType2() { try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { listViewVector.allocateNew(); UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); ListWriter innerListWriter = listViewWriter.listView(); for (int i = 0; i < COUNT; i++) { listViewWriter.startListView(); for (int j = 0; j < i % 7; j++) { innerListWriter.startListView(); for (int k = 0; k < i % 13; k++) { if (k % 2 == 0) { innerListWriter.integer().writeInt(k); } else { innerListWriter.bigInt().writeBigInt(k); } } innerListWriter.endListView(); } listViewWriter.endListView(); } listViewWriter.setValueCount(COUNT); UnionListViewReader listViewReader = new UnionListViewReader(listViewVector); checkUnionListType(listViewReader); } } @Test public void testListMapType() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); createListTypeVectorWithMapType(listWriter); listWriter.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); checkListTypeMap(listReader); // Verify that the map vector has keysSorted = true MapVector mapVector = (MapVector) listVector.getDataVector(); ArrowType arrowType = mapVector.getField().getFieldType().getType(); assertTrue(((ArrowType.Map) arrowType).getKeysSorted()); } } @Test public void testListViewMapType() { try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { listViewVector.allocateNew(); UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); createListTypeVectorWithMapType(listViewWriter); listViewWriter.setValueCount(COUNT); UnionListViewReader listViewReader = new UnionListViewReader(listViewVector); checkListTypeMap(listViewReader); // Verify that the map vector has keysSorted = true MapVector mapVector = (MapVector) listViewVector.getDataVector(); ArrowType arrowType = mapVector.getField().getFieldType().getType(); assertTrue(((ArrowType.Map) arrowType).getKeysSorted()); } } @Test public void simpleUnion() throws Exception { List bufs = new ArrayList(); UnionVector vector = 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) { unionWriter.writeInt(i); } else if (i % 5 == 1) { TimeStampMilliTZHolder holder = new TimeStampMilliTZHolder(); holder.value = (long) i; holder.timezone = "AsdfTimeZone"; unionWriter.write(holder); } else if (i % 5 == 2) { DurationHolder holder = new DurationHolder(); holder.value = (long) i; holder.unit = TimeUnit.NANOSECOND; unionWriter.write(holder); } else if (i % 5 == 3) { FixedSizeBinaryHolder holder = new FixedSizeBinaryHolder(); ArrowBuf buf = allocator.buffer(4); buf.setInt(0, i); holder.byteWidth = 4; 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); } } vector.setValueCount(COUNT); UnionReader unionReader = new UnionReader(vector); for (int i = 0; i < COUNT; i++) { unionReader.setPosition(i); if (i % 5 == 0) { assertEquals(i, unionReader.readInteger().intValue()); } else if (i % 5 == 1) { NullableTimeStampMilliTZHolder holder = new NullableTimeStampMilliTZHolder(); unionReader.read(holder); assertEquals(i, holder.value); assertEquals("AsdfTimeZone", holder.timezone); } else if (i % 5 == 2) { NullableDurationHolder holder = new NullableDurationHolder(); unionReader.read(holder); assertEquals(i, holder.value); assertEquals(TimeUnit.NANOSECOND, holder.unit); } else if (i % 5 == 3) { NullableFixedSizeBinaryHolder holder = new NullableFixedSizeBinaryHolder(); 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); } } vector.close(); AutoCloseables.close(bufs); } @Test public void promotableWriter() { try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); for (int i = 0; i < 100; i++) { BigIntWriter bigIntWriter = rootWriter.bigInt("a"); bigIntWriter.setPosition(i); bigIntWriter.writeBigInt(i); } Field field = parent.getField().getChildren().get(0).getChildren().get(0); assertEquals("a", field.getName()); assertEquals(Int.TYPE_TYPE, field.getType().getTypeID()); Int intType = (Int) field.getType(); assertEquals(64, intType.getBitWidth()); assertTrue(intType.getIsSigned()); for (int i = 100; i < 200; i++) { VarCharWriter varCharWriter = rootWriter.varChar("a"); varCharWriter.setPosition(i); byte[] bytes = Integer.toString(i).getBytes(StandardCharsets.UTF_8); ArrowBuf tempBuf = allocator.buffer(bytes.length); tempBuf.setBytes(0, bytes); varCharWriter.writeVarChar(0, bytes.length, tempBuf); tempBuf.close(); } field = parent.getField().getChildren().get(0).getChildren().get(0); assertEquals("a", field.getName()); assertEquals(Union.TYPE_TYPE, field.getType().getTypeID()); assertEquals(Int.TYPE_TYPE, field.getChildren().get(0).getType().getTypeID()); assertEquals(Utf8.TYPE_TYPE, field.getChildren().get(1).getType().getTypeID()); StructReader rootReader = new SingleStructReaderImpl(parent).reader("root"); for (int i = 0; i < 100; i++) { rootReader.setPosition(i); FieldReader reader = rootReader.reader("a"); Long value = reader.readLong(); assertNotNull(value, "index: " + i); assertEquals(i, value.intValue()); } for (int i = 100; i < 200; i++) { rootReader.setPosition(i); FieldReader reader = rootReader.reader("a"); Text value = reader.readText(); assertEquals(Integer.toString(i), value.toString()); } } } /** Even without writing to the writer, the union schema is created correctly. */ @Test public void promotableWriterSchema() { try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); rootWriter.bigInt("a"); rootWriter.varChar("a"); Field field = parent.getField().getChildren().get(0).getChildren().get(0); assertEquals("a", field.getName()); assertEquals(ArrowTypeID.Union, field.getType().getTypeID()); assertEquals(ArrowTypeID.Int, field.getChildren().get(0).getType().getTypeID()); Int intType = (Int) field.getChildren().get(0).getType(); assertEquals(64, intType.getBitWidth()); assertTrue(intType.getIsSigned()); assertEquals(ArrowTypeID.Utf8, field.getChildren().get(1).getType().getTypeID()); } } private Set getFieldNames(List fields) { Set fieldNames = new HashSet<>(); for (Field field : fields) { fieldNames.add(field.getName()); if (!field.getChildren().isEmpty()) { for (String name : getFieldNames(field.getChildren())) { fieldNames.add(field.getName() + "::" + name); } } } return fieldNames; } @Test public void structWriterMixedCaseFieldNames() { try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { // test case-sensitive StructWriter ComplexWriter writer = new ComplexWriterImpl("rootCaseSensitive", parent, false, true); StructWriter rootWriterCaseSensitive = writer.rootAsStruct(); rootWriterCaseSensitive.bigInt("int_field"); rootWriterCaseSensitive.bigInt("Int_Field"); rootWriterCaseSensitive.float4("float_field"); rootWriterCaseSensitive.float4("Float_Field"); StructWriter structFieldWriterCaseSensitive = rootWriterCaseSensitive.struct("struct_field"); structFieldWriterCaseSensitive.varChar("char_field"); structFieldWriterCaseSensitive.varChar("Char_Field"); ListWriter listFieldWriterCaseSensitive = rootWriterCaseSensitive.list("list_field"); StructWriter listStructFieldWriterCaseSensitive = listFieldWriterCaseSensitive.struct(); listStructFieldWriterCaseSensitive.bit("bit_field"); listStructFieldWriterCaseSensitive.bit("Bit_Field"); List fieldsCaseSensitive = parent.getField().getChildren().get(0).getChildren(); Set fieldNamesCaseSensitive = getFieldNames(fieldsCaseSensitive); assertEquals(11, fieldNamesCaseSensitive.size()); assertTrue(fieldNamesCaseSensitive.contains("int_field")); assertTrue(fieldNamesCaseSensitive.contains("Int_Field")); assertTrue(fieldNamesCaseSensitive.contains("float_field")); assertTrue(fieldNamesCaseSensitive.contains("Float_Field")); assertTrue(fieldNamesCaseSensitive.contains("struct_field")); assertTrue(fieldNamesCaseSensitive.contains("struct_field::char_field")); assertTrue(fieldNamesCaseSensitive.contains("struct_field::Char_Field")); assertTrue(fieldNamesCaseSensitive.contains("list_field")); assertTrue(fieldNamesCaseSensitive.contains("list_field::$data$")); assertTrue(fieldNamesCaseSensitive.contains("list_field::$data$::bit_field")); assertTrue(fieldNamesCaseSensitive.contains("list_field::$data$::Bit_Field")); // test case-insensitive StructWriter ComplexWriter writerCaseInsensitive = new ComplexWriterImpl("rootCaseInsensitive", parent, false, false); StructWriter rootWriterCaseInsensitive = writerCaseInsensitive.rootAsStruct(); rootWriterCaseInsensitive.bigInt("int_field"); rootWriterCaseInsensitive.bigInt("Int_Field"); rootWriterCaseInsensitive.float4("float_field"); rootWriterCaseInsensitive.float4("Float_Field"); StructWriter structFieldWriterCaseInsensitive = rootWriterCaseInsensitive.struct("struct_field"); structFieldWriterCaseInsensitive.varChar("char_field"); structFieldWriterCaseInsensitive.varChar("Char_Field"); ListWriter listFieldWriterCaseInsensitive = rootWriterCaseInsensitive.list("list_field"); StructWriter listStructFieldWriterCaseInsensitive = listFieldWriterCaseInsensitive.struct(); listStructFieldWriterCaseInsensitive.bit("bit_field"); listStructFieldWriterCaseInsensitive.bit("Bit_Field"); List fieldsCaseInsensitive = parent.getField().getChildren().get(1).getChildren(); Set fieldNamesCaseInsensitive = getFieldNames(fieldsCaseInsensitive); assertEquals(7, fieldNamesCaseInsensitive.size()); assertTrue(fieldNamesCaseInsensitive.contains("int_field")); assertTrue(fieldNamesCaseInsensitive.contains("float_field")); assertTrue(fieldNamesCaseInsensitive.contains("struct_field")); assertTrue(fieldNamesCaseInsensitive.contains("struct_field::char_field")); assertTrue(fieldNamesCaseSensitive.contains("list_field")); assertTrue(fieldNamesCaseSensitive.contains("list_field::$data$")); assertTrue(fieldNamesCaseSensitive.contains("list_field::$data$::bit_field")); } } @Test public void timeStampSecWriter() throws Exception { // test values final long expectedSecs = 981173106L; final LocalDateTime expectedSecDateTime = LocalDateTime.of(2001, 2, 3, 4, 5, 6, 0); try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { // write ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); { TimeStampSecWriter timeStampSecWriter = rootWriter.timeStampSec("sec"); timeStampSecWriter.setPosition(0); timeStampSecWriter.writeTimeStampSec(expectedSecs); } { TimeStampSecTZWriter timeStampSecTZWriter = rootWriter.timeStampSecTZ("secTZ", "UTC"); timeStampSecTZWriter.setPosition(1); timeStampSecTZWriter.writeTimeStampSecTZ(expectedSecs); } // schema List children = parent.getField().getChildren().get(0).getChildren(); checkTimestampField(children.get(0), "sec"); checkTimestampTZField(children.get(1), "secTZ", "UTC"); // read StructReader rootReader = new SingleStructReaderImpl(parent).reader("root"); { FieldReader secReader = rootReader.reader("sec"); secReader.setPosition(0); LocalDateTime secDateTime = secReader.readLocalDateTime(); assertEquals(expectedSecDateTime, secDateTime); long secLong = secReader.readLong(); assertEquals(expectedSecs, secLong); } { FieldReader secTZReader = rootReader.reader("secTZ"); secTZReader.setPosition(1); long secTZLong = secTZReader.readLong(); assertEquals(expectedSecs, secTZLong); } } } @Test public void timeStampMilliWriters() throws Exception { // test values final long expectedMillis = 981173106123L; final LocalDateTime expectedMilliDateTime = LocalDateTime.of(2001, 2, 3, 4, 5, 6, 123 * 1_000_000); try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator); ) { // write ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); { TimeStampMilliWriter timeStampWriter = rootWriter.timeStampMilli("milli"); timeStampWriter.setPosition(0); timeStampWriter.writeTimeStampMilli(expectedMillis); } String tz = "UTC"; { TimeStampMilliTZWriter timeStampTZWriter = rootWriter.timeStampMilliTZ("milliTZ", tz); timeStampTZWriter.setPosition(0); timeStampTZWriter.writeTimeStampMilliTZ(expectedMillis); } // schema List children = parent.getField().getChildren().get(0).getChildren(); checkTimestampField(children.get(0), "milli"); checkTimestampTZField(children.get(1), "milliTZ", tz); // read StructReader rootReader = new SingleStructReaderImpl(parent).reader("root"); { FieldReader milliReader = rootReader.reader("milli"); milliReader.setPosition(0); LocalDateTime milliDateTime = milliReader.readLocalDateTime(); assertEquals(expectedMilliDateTime, milliDateTime); long milliLong = milliReader.readLong(); assertEquals(expectedMillis, milliLong); } { FieldReader milliTZReader = rootReader.reader("milliTZ"); milliTZReader.setPosition(0); long milliTZLong = milliTZReader.readLong(); assertEquals(expectedMillis, milliTZLong); } } } private void checkTimestampField(Field field, String name) { assertEquals(name, field.getName()); assertEquals(ArrowType.Timestamp.TYPE_TYPE, field.getType().getTypeID()); } private void checkTimestampTZField(Field field, String name, String tz) { checkTimestampField(field, name); assertEquals(tz, ((Timestamp) field.getType()).getTimezone()); } @Test public void timeStampMicroWriters() throws Exception { // test values final long expectedMicros = 981173106123456L; final LocalDateTime expectedMicroDateTime = LocalDateTime.of(2001, 2, 3, 4, 5, 6, 123456 * 1000); try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { // write ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); { TimeStampMicroWriter timeStampMicroWriter = rootWriter.timeStampMicro("micro"); timeStampMicroWriter.setPosition(0); timeStampMicroWriter.writeTimeStampMicro(expectedMicros); } String tz = "UTC"; { TimeStampMicroTZWriter timeStampMicroWriter = rootWriter.timeStampMicroTZ("microTZ", tz); timeStampMicroWriter.setPosition(1); timeStampMicroWriter.writeTimeStampMicroTZ(expectedMicros); } // schema List children = parent.getField().getChildren().get(0).getChildren(); checkTimestampField(children.get(0), "micro"); checkTimestampTZField(children.get(1), "microTZ", tz); // read StructReader rootReader = new SingleStructReaderImpl(parent).reader("root"); { FieldReader microReader = rootReader.reader("micro"); microReader.setPosition(0); LocalDateTime microDateTime = microReader.readLocalDateTime(); assertEquals(expectedMicroDateTime, microDateTime); long microLong = microReader.readLong(); assertEquals(expectedMicros, microLong); } { FieldReader microReader = rootReader.reader("microTZ"); microReader.setPosition(1); long microLong = microReader.readLong(); assertEquals(expectedMicros, microLong); } } } @Test public void timeStampNanoWriters() throws Exception { // test values final long expectedNanos = 981173106123456789L; final LocalDateTime expectedNanoDateTime = LocalDateTime.of(2001, 2, 3, 4, 5, 6, 123456789); try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { // write ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); { TimeStampNanoWriter timeStampNanoWriter = rootWriter.timeStampNano("nano"); timeStampNanoWriter.setPosition(0); timeStampNanoWriter.writeTimeStampNano(expectedNanos); } String tz = "UTC"; { TimeStampNanoTZWriter timeStampNanoWriter = rootWriter.timeStampNanoTZ("nanoTZ", tz); timeStampNanoWriter.setPosition(0); timeStampNanoWriter.writeTimeStampNanoTZ(expectedNanos); } // schema List children = parent.getField().getChildren().get(0).getChildren(); checkTimestampField(children.get(0), "nano"); checkTimestampTZField(children.get(1), "nanoTZ", tz); // read StructReader rootReader = new SingleStructReaderImpl(parent).reader("root"); { FieldReader nanoReader = rootReader.reader("nano"); nanoReader.setPosition(0); LocalDateTime nanoDateTime = nanoReader.readLocalDateTime(); assertEquals(expectedNanoDateTime, nanoDateTime); long nanoLong = nanoReader.readLong(); assertEquals(expectedNanos, nanoLong); } { FieldReader nanoReader = rootReader.reader("nanoTZ"); nanoReader.setPosition(0); long nanoLong = nanoReader.readLong(); assertEquals(expectedNanos, nanoLong); NullableTimeStampNanoTZHolder h = new NullableTimeStampNanoTZHolder(); nanoReader.read(h); assertEquals(expectedNanos, h.value); } } } @Test public void fixedSizeBinaryWriters() throws Exception { // test values int numValues = 10; int byteWidth = 9; byte[][] values = new byte[numValues][byteWidth]; for (int i = 0; i < numValues; i++) { for (int j = 0; j < byteWidth; j++) { values[i][j] = ((byte) i); } } ArrowBuf[] bufs = new ArrowBuf[numValues]; for (int i = 0; i < numValues; i++) { bufs[i] = allocator.buffer(byteWidth); bufs[i].setBytes(0, values[i]); } try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { // write ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); String fieldName = "fixedSizeBinary"; FixedSizeBinaryWriter fixedSizeBinaryWriter = rootWriter.fixedSizeBinary(fieldName, byteWidth); for (int i = 0; i < numValues; i++) { fixedSizeBinaryWriter.setPosition(i); fixedSizeBinaryWriter.writeFixedSizeBinary(bufs[i]); } // schema List children = parent.getField().getChildren().get(0).getChildren(); assertEquals(fieldName, children.get(0).getName()); assertEquals(ArrowType.FixedSizeBinary.TYPE_TYPE, children.get(0).getType().getTypeID()); // read StructReader rootReader = new SingleStructReaderImpl(parent).reader("root"); FieldReader fixedSizeBinaryReader = rootReader.reader(fieldName); for (int i = 0; i < numValues; i++) { fixedSizeBinaryReader.setPosition(i); byte[] readValues = fixedSizeBinaryReader.readByteArray(); assertArrayEquals(values[i], readValues); } } AutoCloseables.close(bufs); } @Test public void complexCopierWithList() { try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); ListWriter listWriter = rootWriter.list("list"); StructWriter innerStructWriter = listWriter.struct(); IntWriter outerIntWriter = listWriter.integer(); rootWriter.start(); listWriter.startList(); outerIntWriter.writeInt(1); outerIntWriter.writeInt(2); innerStructWriter.start(); IntWriter intWriter = innerStructWriter.integer("a"); intWriter.writeInt(1); innerStructWriter.end(); innerStructWriter.start(); intWriter = innerStructWriter.integer("a"); intWriter.writeInt(2); innerStructWriter.end(); listWriter.endList(); rootWriter.end(); writer.setValueCount(1); StructVector structVector = (StructVector) parent.getChild("root"); TransferPair tp = structVector.getTransferPair(allocator); tp.splitAndTransfer(0, 1); NonNullableStructVector toStructVector = (NonNullableStructVector) tp.getTo(); JsonStringHashMap toMapValue = (JsonStringHashMap) toStructVector.getObject(0); JsonStringArrayList object = (JsonStringArrayList) toMapValue.get("list"); assertEquals(1, object.get(0)); assertEquals(2, object.get(1)); JsonStringHashMap innerStruct = (JsonStringHashMap) object.get(2); assertEquals(1, innerStruct.get("a")); innerStruct = (JsonStringHashMap) object.get(3); assertEquals(2, innerStruct.get("a")); toStructVector.close(); } } @Test public void complexCopierWithListView() { try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); ListWriter listViewWriter = rootWriter.listView("listView"); StructWriter innerStructWriter = listViewWriter.struct(); IntWriter outerIntWriter = listViewWriter.integer(); rootWriter.start(); listViewWriter.startListView(); outerIntWriter.writeInt(1); outerIntWriter.writeInt(2); innerStructWriter.start(); IntWriter intWriter = innerStructWriter.integer("a"); intWriter.writeInt(1); innerStructWriter.end(); innerStructWriter.start(); intWriter = innerStructWriter.integer("a"); intWriter.writeInt(2); innerStructWriter.end(); listViewWriter.endListView(); rootWriter.end(); writer.setValueCount(1); StructVector structVector = (StructVector) parent.getChild("root"); TransferPair tp = structVector.getTransferPair(allocator); tp.splitAndTransfer(0, 1); NonNullableStructVector toStructVector = (NonNullableStructVector) tp.getTo(); JsonStringHashMap toMapValue = (JsonStringHashMap) toStructVector.getObject(0); JsonStringArrayList object = (JsonStringArrayList) toMapValue.get("listView"); assertEquals(1, object.get(0)); assertEquals(2, object.get(1)); JsonStringHashMap innerStruct = (JsonStringHashMap) object.get(2); assertEquals(1, innerStruct.get("a")); innerStruct = (JsonStringHashMap) object.get(3); assertEquals(2, innerStruct.get("a")); toStructVector.close(); } } @Test public void testSingleStructWriter1() { /* initialize a SingleStructWriter with empty StructVector and then lazily * create all vectors with expected initialCapacity. */ try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { SingleStructWriter singleStructWriter = new SingleStructWriter(parent); int initialCapacity = 1024; singleStructWriter.setInitialCapacity(initialCapacity); IntWriter intWriter = singleStructWriter.integer("intField"); BigIntWriter bigIntWriter = singleStructWriter.bigInt("bigIntField"); Float4Writer float4Writer = singleStructWriter.float4("float4Field"); Float8Writer float8Writer = singleStructWriter.float8("float8Field"); ListWriter listWriter = singleStructWriter.list("listField"); ListWriter listViewWriter = singleStructWriter.listView("listViewField"); MapWriter mapWriter = singleStructWriter.map("mapField", false); int intValue = 100; long bigIntValue = 10000; float float4Value = 100.5f; double float8Value = 100.375; for (int i = 0; i < initialCapacity; i++) { singleStructWriter.start(); intWriter.writeInt(intValue + i); bigIntWriter.writeBigInt(bigIntValue + (long) i); float4Writer.writeFloat4(float4Value + (float) i); float8Writer.writeFloat8(float8Value + (double) i); listWriter.setPosition(i); listWriter.startList(); listWriter.integer().writeInt(intValue + i); listWriter.integer().writeInt(intValue + i + 1); listWriter.integer().writeInt(intValue + i + 2); listWriter.integer().writeInt(intValue + i + 3); listWriter.endList(); listViewWriter.setPosition(i); listViewWriter.startListView(); listViewWriter.integer().writeInt(intValue + i); listViewWriter.integer().writeInt(intValue + i + 1); listViewWriter.integer().writeInt(intValue + i + 2); listViewWriter.integer().writeInt(intValue + i + 3); listViewWriter.endListView(); mapWriter.setPosition(i); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().integer().writeInt(intValue + i); mapWriter.value().integer().writeInt(intValue + i + 1); mapWriter.endEntry(); mapWriter.startEntry(); mapWriter.key().integer().writeInt(intValue + i + 2); mapWriter.value().integer().writeInt(intValue + i + 3); mapWriter.endEntry(); mapWriter.endMap(); singleStructWriter.end(); } IntVector intVector = (IntVector) parent.getChild("intField"); BigIntVector bigIntVector = (BigIntVector) parent.getChild("bigIntField"); Float4Vector float4Vector = (Float4Vector) parent.getChild("float4Field"); Float8Vector float8Vector = (Float8Vector) parent.getChild("float8Field"); int capacity = singleStructWriter.getValueCapacity(); assertTrue(capacity >= initialCapacity && capacity < initialCapacity * 2); capacity = intVector.getValueCapacity(); assertTrue(capacity >= initialCapacity && capacity < initialCapacity * 2); capacity = bigIntVector.getValueCapacity(); assertTrue(capacity >= initialCapacity && capacity < initialCapacity * 2); capacity = float4Vector.getValueCapacity(); assertTrue(capacity >= initialCapacity && capacity < initialCapacity * 2); capacity = float8Vector.getValueCapacity(); assertTrue(capacity >= initialCapacity && capacity < initialCapacity * 2); StructReader singleStructReader = new SingleStructReaderImpl(parent); IntReader intReader = singleStructReader.reader("intField"); BigIntReader bigIntReader = singleStructReader.reader("bigIntField"); Float4Reader float4Reader = singleStructReader.reader("float4Field"); Float8Reader float8Reader = singleStructReader.reader("float8Field"); UnionListReader listReader = (UnionListReader) singleStructReader.reader("listField"); UnionListViewReader listViewReader = (UnionListViewReader) singleStructReader.reader("listViewField"); UnionMapReader mapReader = (UnionMapReader) singleStructReader.reader("mapField"); for (int i = 0; i < initialCapacity; i++) { intReader.setPosition(i); bigIntReader.setPosition(i); float4Reader.setPosition(i); float8Reader.setPosition(i); listReader.setPosition(i); listViewReader.setPosition(i); mapReader.setPosition(i); assertEquals(intValue + i, intReader.readInteger().intValue()); assertEquals(bigIntValue + (long) i, bigIntReader.readLong().longValue()); assertEquals(float4Value + (float) i, float4Reader.readFloat().floatValue(), 0); assertEquals(float8Value + (double) i, float8Reader.readDouble().doubleValue(), 0); for (int j = 0; j < 4; j++) { listReader.next(); assertEquals(intValue + i + j, listReader.reader().readInteger().intValue()); } for (int j = 0; j < 4; j++) { listViewReader.next(); assertEquals(intValue + i + j, listViewReader.reader().readInteger().intValue()); } for (int k = 0; k < 4; k += 2) { mapReader.next(); assertEquals(intValue + k + i, mapReader.key().readInteger().intValue()); assertEquals(intValue + k + i + 1, mapReader.value().readInteger().intValue()); } } } } @Test public void testListWriterWithNulls() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.setInitialCapacity(COUNT); listVector.allocateNew(); listVector.getValidityBuffer().setOne(0, (int) listVector.getValidityBuffer().capacity()); UnionListWriter listWriter = listVector.getWriter(); // expected listVector : [[null], null, [2, 4], null, [null], null, [6, 12], ...] createNullsWithListWriters(listWriter); listVector.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); checkNullsWithListWriters(listReader); } } @Test public void testListViewWriterWithNulls() { try (ListViewVector listViewVector = ListViewVector.empty("listView", allocator)) { listViewVector.setInitialCapacity(COUNT); listViewVector.allocateNew(); listViewVector .getValidityBuffer() .setOne(0, (int) listViewVector.getValidityBuffer().capacity()); UnionListViewWriter listWriter = listViewVector.getWriter(); // expected listVector : [[null], null, [2, 4], null, [null], null, [6, 12], ...] createNullsWithListWriters(listWriter); listViewVector.setValueCount(COUNT); UnionListViewReader listReader = new UnionListViewReader(listViewVector); checkNullsWithListWriters(listReader); } } @Test public void testListOfListWriterWithNulls() { try (ListVector listVector = ListVector.empty("listoflist", allocator)) { listVector.setInitialCapacity(COUNT); listVector.allocateNew(); listVector.getValidityBuffer().setOne(0, (int) listVector.getValidityBuffer().capacity()); UnionListWriter listWriter = listVector.getWriter(); // create list : [ [null], null, [[null, 2, 4]], null, [null], null, [[null, 6, 12]], ... ] for (int i = 0; i < COUNT; i++) { listWriter.setPosition(i); if (i % 2 == 0) { listWriter.startList(); if (i % 4 == 0) { listWriter.list().writeNull(); } else { listWriter.list().startList(); listWriter.list().integer().writeNull(); listWriter.list().integer().writeInt(i); listWriter.list().integer().writeInt(i * 2); listWriter.list().endList(); } listWriter.endList(); } else { listWriter.writeNull(); } } listVector.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); for (int i = 0; i < COUNT; i++) { listReader.setPosition(i); if (i % 2 == 0) { assertTrue(listReader.isSet()); listReader.next(); if (i % 4 == 0) { assertFalse(listReader.reader().isSet()); } else { listReader.reader().next(); assertFalse(listReader.reader().reader().isSet()); listReader.reader().next(); assertEquals(i, listReader.reader().reader().readInteger().intValue()); listReader.reader().next(); assertEquals(i * 2, listReader.reader().reader().readInteger().intValue()); } } else { assertFalse(listReader.isSet()); } } } } @Test public void testListViewOfListViewWriterWithNulls() { try (ListViewVector listViewVector = ListViewVector.empty("listViewoflistView", allocator)) { listViewVector.setInitialCapacity(COUNT); listViewVector.allocateNew(); listViewVector .getValidityBuffer() .setOne(0, (int) listViewVector.getValidityBuffer().capacity()); UnionListViewWriter listViewWriter = listViewVector.getWriter(); // create list : [ [null], null, [[null, 2, 4]], null, [null], null, [[null, 6, 12]], ... ] for (int i = 0; i < COUNT; i++) { listViewWriter.setPosition(i); if (i % 2 == 0) { listViewWriter.startListView(); if (i % 4 == 0) { listViewWriter.listView().writeNull(); } else { listViewWriter.listView().startListView(); listViewWriter.listView().integer().writeNull(); listViewWriter.listView().integer().writeInt(i); listViewWriter.listView().integer().writeInt(i * 2); listViewWriter.listView().endListView(); } listViewWriter.endListView(); } else { listViewWriter.writeNull(); } } listViewVector.setValueCount(COUNT); UnionListViewReader listViewReader = new UnionListViewReader(listViewVector); for (int i = 0; i < COUNT; i++) { listViewReader.setPosition(i); if (i % 2 == 0) { assertTrue(listViewReader.isSet()); listViewReader.next(); if (i % 4 == 0) { assertFalse(listViewReader.reader().isSet()); } else { listViewReader.reader().next(); assertFalse(listViewReader.reader().reader().isSet()); listViewReader.reader().next(); assertEquals(i, listViewReader.reader().reader().readInteger().intValue()); listViewReader.reader().next(); assertEquals(i * 2, listViewReader.reader().reader().readInteger().intValue()); } } else { assertFalse(listViewReader.isSet()); } } } } @Test public void testListOfListOfListWriterWithNulls() { try (ListVector listVector = ListVector.empty("listoflistoflist", allocator)) { listVector.setInitialCapacity(COUNT); listVector.allocateNew(); listVector.getValidityBuffer().setOne(0, (int) listVector.getValidityBuffer().capacity()); UnionListWriter listWriter = listVector.getWriter(); // create list : [ null, [null], [[null]], [[[null, 1, 2]]], null, [null], ... for (int i = 0; i < COUNT; i++) { listWriter.setPosition(i); if (i % 4 == 0) { listWriter.writeNull(); } else { listWriter.startList(); if (i % 4 == 1) { listWriter.list().writeNull(); } else if (i % 4 == 2) { listWriter.list().startList(); listWriter.list().list().writeNull(); listWriter.list().endList(); } else { listWriter.list().startList(); listWriter.list().list().startList(); listWriter.list().list().integer().writeNull(); listWriter.list().list().integer().writeInt(i); listWriter.list().list().integer().writeInt(i * 2); listWriter.list().list().endList(); listWriter.list().endList(); } listWriter.endList(); } } listVector.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); for (int i = 0; i < COUNT; i++) { listReader.setPosition(i); if (i % 4 == 0) { assertFalse(listReader.isSet()); } else { assertTrue(listReader.isSet()); listReader.next(); if (i % 4 == 1) { assertFalse(listReader.reader().isSet()); } else if (i % 4 == 2) { listReader.reader().next(); assertFalse(listReader.reader().reader().isSet()); } else { listReader.reader().next(); listReader.reader().reader().next(); assertFalse(listReader.reader().reader().reader().isSet()); listReader.reader().reader().next(); assertEquals(i, listReader.reader().reader().reader().readInteger().intValue()); listReader.reader().reader().next(); assertEquals(i * 2, listReader.reader().reader().reader().readInteger().intValue()); } } } } } @Test public void testListViewOfListViewOfListViewWriterWithNulls() { try (ListViewVector listViewVector = ListViewVector.empty("listViewoflistViewoflistView", allocator)) { listViewVector.setInitialCapacity(COUNT); listViewVector.allocateNew(); listViewVector .getValidityBuffer() .setOne(0, (int) listViewVector.getValidityBuffer().capacity()); UnionListViewWriter listViewWriter = listViewVector.getWriter(); // create list : [ null, [null], [[null]], [[[null, 1, 2]]], null, [null], ... for (int i = 0; i < COUNT; i++) { listViewWriter.setPosition(i); if (i % 4 == 0) { listViewWriter.writeNull(); } else { listViewWriter.startListView(); if (i % 4 == 1) { listViewWriter.listView().writeNull(); } else if (i % 4 == 2) { listViewWriter.listView().startListView(); listViewWriter.listView().listView().writeNull(); listViewWriter.listView().endListView(); } else { listViewWriter.listView().startListView(); listViewWriter.listView().listView().startListView(); listViewWriter.listView().listView().integer().writeNull(); listViewWriter.listView().listView().integer().writeInt(i); listViewWriter.listView().listView().integer().writeInt(i * 2); listViewWriter.listView().listView().endListView(); listViewWriter.listView().endListView(); } listViewWriter.endListView(); } } listViewVector.setValueCount(COUNT); UnionListViewReader listViewReader = new UnionListViewReader(listViewVector); for (int i = 0; i < COUNT; i++) { listViewReader.setPosition(i); if (i % 4 == 0) { assertFalse(listViewReader.isSet()); } else { assertTrue(listViewReader.isSet()); listViewReader.next(); if (i % 4 == 1) { assertFalse(listViewReader.reader().isSet()); } else if (i % 4 == 2) { listViewReader.reader().next(); assertFalse(listViewReader.reader().reader().isSet()); } else { listViewReader.reader().next(); listViewReader.reader().reader().next(); assertFalse(listViewReader.reader().reader().reader().isSet()); listViewReader.reader().reader().next(); assertEquals(i, listViewReader.reader().reader().reader().readInteger().intValue()); listViewReader.reader().reader().next(); assertEquals(i * 2, listViewReader.reader().reader().reader().readInteger().intValue()); } } } } } @Test public void testStructOfList() { try (StructVector structVector = StructVector.empty("struct1", allocator)) { structVector.addOrGetList("childList1"); NullableStructReaderImpl structReader = structVector.getReader(); FieldReader childListReader = structReader.reader("childList1"); assertNotNull(childListReader); } try (StructVector structVector = StructVector.empty("struct2", allocator)) { structVector.addOrGetList("childList2"); NullableStructWriter structWriter = structVector.getWriter(); structWriter.start(); ListWriter listWriter = structWriter.list("childList2"); listWriter.startList(); listWriter.integer().writeInt(10); listWriter.endList(); structWriter.end(); NullableStructReaderImpl structReader = structVector.getReader(); FieldReader childListReader = structReader.reader("childList2"); int size = childListReader.size(); assertEquals(1, size); int data = childListReader.reader().readInteger(); assertEquals(10, data); } try (StructVector structVector = StructVector.empty("struct3", allocator)) { structVector.addOrGetList("childList3"); NullableStructWriter structWriter = structVector.getWriter(); for (int i = 0; i < 5; ++i) { structWriter.setPosition(i); structWriter.start(); ListWriter listWriter = structWriter.list("childList3"); listWriter.startList(); listWriter.integer().writeInt(i); listWriter.endList(); structWriter.end(); } NullableStructReaderImpl structReader = structVector.getReader(); structReader.setPosition(3); FieldReader childListReader = structReader.reader("childList3"); int size = childListReader.size(); assertEquals(1, size); int data = ((List) childListReader.readObject()).get(0); assertEquals(3, data); } try (StructVector structVector = StructVector.empty("struct4", allocator)) { structVector.addOrGetList("childList4"); NullableStructWriter structWriter = structVector.getWriter(); for (int i = 0; i < 5; ++i) { structWriter.setPosition(i); structWriter.start(); structWriter.writeNull(); structWriter.end(); } NullableStructReaderImpl structReader = structVector.getReader(); structReader.setPosition(3); FieldReader childListReader = structReader.reader("childList4"); int size = childListReader.size(); assertEquals(0, size); } } @Test public void testMap() { try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { ComplexWriter writer = new ComplexWriterImpl("root", parent); MapWriter mapWriter = writer.rootAsMap(false); for (int i = 0; i < COUNT; i++) { mapWriter.startMap(); for (int j = 0; j < i % 7; j++) { mapWriter.startEntry(); if (j % 2 == 0) { mapWriter.key().integer().writeInt(j); mapWriter.value().integer().writeInt(j + 1); } else { IntHolder keyHolder = new IntHolder(); keyHolder.value = j; IntHolder valueHolder = new IntHolder(); valueHolder.value = j + 1; mapWriter.key().integer().write(keyHolder); mapWriter.value().integer().write(valueHolder); } mapWriter.endEntry(); } mapWriter.endMap(); } writer.setValueCount(COUNT); UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root"); for (int i = 0; i < COUNT; i++) { mapReader.setPosition(i); for (int j = 0; j < i % 7; j++) { mapReader.next(); assertEquals(j, mapReader.key().readInteger().intValue()); assertEquals(j + 1, mapReader.value().readInteger().intValue()); } } } } @Test public void testMapWithNulls() { try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { ComplexWriter writer = new ComplexWriterImpl("root", parent); MapWriter mapWriter = writer.rootAsMap(false); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().integer().writeNull(); mapWriter.value().integer().writeInt(1); mapWriter.endEntry(); mapWriter.endMap(); writer.setValueCount(1); UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root"); assertNull(mapReader.key().readInteger()); assertEquals(1, mapReader.value().readInteger().intValue()); } } @Test public void testMapWithListKey() { try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { ComplexWriter writer = new ComplexWriterImpl("root", parent); MapWriter mapWriter = writer.rootAsMap(false); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().list().startList(); for (int i = 0; i < 3; i++) { mapWriter.key().list().integer().writeInt(i); } mapWriter.key().list().endList(); mapWriter.value().integer().writeInt(1); mapWriter.endEntry(); mapWriter.endMap(); writer.setValueCount(1); UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root"); mapReader.key().next(); assertEquals(0, mapReader.key().reader().readInteger().intValue()); mapReader.key().next(); assertEquals(1, mapReader.key().reader().readInteger().intValue()); mapReader.key().next(); assertEquals(2, mapReader.key().reader().readInteger().intValue()); assertEquals(1, mapReader.value().readInteger().intValue()); } } @Test public void testMapWithStructKey() { try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { ComplexWriter writer = new ComplexWriterImpl("root", parent); MapWriter mapWriter = writer.rootAsMap(false); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().struct().start(); mapWriter.key().struct().integer("value1").writeInt(1); mapWriter.key().struct().integer("value2").writeInt(2); mapWriter.key().struct().end(); mapWriter.value().integer().writeInt(1); mapWriter.endEntry(); mapWriter.endMap(); writer.setValueCount(1); UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root"); assertEquals(1, mapReader.key().reader("value1").readInteger().intValue()); assertEquals(2, mapReader.key().reader("value2").readInteger().intValue()); assertEquals(1, mapReader.value().readInteger().intValue()); } } @Test public void structWriterVarCharHelpers() { try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { ComplexWriter writer = new ComplexWriterImpl("root", parent, false, true); StructWriter rootWriter = writer.rootAsStruct(); rootWriter.start(); rootWriter.setPosition(0); rootWriter.varChar("c").writeVarChar(new Text("row1")); rootWriter.setPosition(1); rootWriter.varChar("c").writeVarChar("row2"); rootWriter.end(); VarCharVector vector = parent.getChild("root", StructVector.class).getChild("c", VarCharVector.class); assertEquals("row1", vector.getObject(0).toString()); assertEquals("row2", vector.getObject(1).toString()); } } @Test public void structWriterVarCharViewHelpers() { try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { ComplexWriter writer = new ComplexWriterImpl("root", parent, false, true); StructWriter rootWriter = writer.rootAsStruct(); rootWriter.start(); rootWriter.setPosition(0); rootWriter.viewVarChar("c").writeViewVarChar(new Text("row1")); rootWriter.setPosition(1); rootWriter.viewVarChar("c").writeViewVarChar("row2"); rootWriter.end(); ViewVarCharVector vector = parent.getChild("root", StructVector.class).getChild("c", ViewVarCharVector.class); assertEquals("row1", vector.getObject(0).toString()); assertEquals("row2", vector.getObject(1).toString()); } } @Test public void structWriterLargeVarCharHelpers() { try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { ComplexWriter writer = new ComplexWriterImpl("root", parent, false, true); StructWriter rootWriter = writer.rootAsStruct(); rootWriter.start(); rootWriter.setPosition(0); rootWriter.largeVarChar("c").writeLargeVarChar(new Text("row1")); rootWriter.setPosition(1); rootWriter.largeVarChar("c").writeLargeVarChar("row2"); rootWriter.end(); LargeVarCharVector vector = parent.getChild("root", StructVector.class).getChild("c", LargeVarCharVector.class); assertEquals("row1", vector.getObject(0).toString()); assertEquals("row2", vector.getObject(1).toString()); } } @Test public void structWriterVarBinaryHelpers() { try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { ComplexWriter writer = new ComplexWriterImpl("root", parent, false, true); StructWriter rootWriter = writer.rootAsStruct(); rootWriter.start(); rootWriter.setPosition(0); rootWriter.varBinary("c").writeVarBinary("row1".getBytes(StandardCharsets.UTF_8)); rootWriter.setPosition(1); rootWriter .varBinary("c") .writeVarBinary( "row2".getBytes(StandardCharsets.UTF_8), 0, "row2".getBytes(StandardCharsets.UTF_8).length); rootWriter.setPosition(2); rootWriter .varBinary("c") .writeVarBinary(ByteBuffer.wrap("row3".getBytes(StandardCharsets.UTF_8))); rootWriter.setPosition(3); rootWriter .varBinary("c") .writeVarBinary( ByteBuffer.wrap("row4".getBytes(StandardCharsets.UTF_8)), 0, "row4".getBytes(StandardCharsets.UTF_8).length); rootWriter.end(); VarBinaryVector uv = parent.getChild("root", StructVector.class).getChild("c", VarBinaryVector.class); assertEquals("row1", new String(uv.get(0), StandardCharsets.UTF_8)); assertEquals("row2", new String(uv.get(1), StandardCharsets.UTF_8)); assertEquals("row3", new String(uv.get(2), StandardCharsets.UTF_8)); assertEquals("row4", new String(uv.get(3), StandardCharsets.UTF_8)); } } @Test public void structWriterLargeVarBinaryHelpers() { try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { ComplexWriter writer = new ComplexWriterImpl("root", parent, false, true); StructWriter rootWriter = writer.rootAsStruct(); rootWriter.start(); rootWriter.setPosition(0); rootWriter.largeVarBinary("c").writeLargeVarBinary("row1".getBytes(StandardCharsets.UTF_8)); rootWriter.setPosition(1); rootWriter .largeVarBinary("c") .writeLargeVarBinary( "row2".getBytes(StandardCharsets.UTF_8), 0, "row2".getBytes(StandardCharsets.UTF_8).length); rootWriter.setPosition(2); rootWriter .largeVarBinary("c") .writeLargeVarBinary(ByteBuffer.wrap("row3".getBytes(StandardCharsets.UTF_8))); rootWriter.setPosition(3); rootWriter .largeVarBinary("c") .writeLargeVarBinary( ByteBuffer.wrap("row4".getBytes(StandardCharsets.UTF_8)), 0, "row4".getBytes(StandardCharsets.UTF_8).length); rootWriter.end(); LargeVarBinaryVector uv = parent.getChild("root", StructVector.class).getChild("c", LargeVarBinaryVector.class); assertEquals("row1", new String(uv.get(0), StandardCharsets.UTF_8)); assertEquals("row2", new String(uv.get(1), StandardCharsets.UTF_8)); assertEquals("row3", new String(uv.get(2), StandardCharsets.UTF_8)); assertEquals("row4", new String(uv.get(3), StandardCharsets.UTF_8)); } } @Test public void listVarCharHelpers() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); listWriter.startList(); listWriter.writeVarChar("row1"); listWriter.writeVarChar(new Text("row2")); listWriter.endList(); listWriter.setValueCount(1); assertEquals("row1", listVector.getObject(0).get(0).toString()); assertEquals("row2", listVector.getObject(0).get(1).toString()); } } @Test public void listLargeVarCharHelpers() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); listWriter.startList(); listWriter.writeLargeVarChar("row1"); listWriter.writeLargeVarChar(new Text("row2")); listWriter.endList(); listWriter.setValueCount(1); assertEquals("row1", listVector.getObject(0).get(0).toString()); assertEquals("row2", listVector.getObject(0).get(1).toString()); } } @Test public void listVarBinaryHelpers() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); listWriter.startList(); listWriter.writeVarBinary("row1".getBytes(StandardCharsets.UTF_8)); listWriter.writeVarBinary( "row2".getBytes(StandardCharsets.UTF_8), 0, "row2".getBytes(StandardCharsets.UTF_8).length); listWriter.writeVarBinary(ByteBuffer.wrap("row3".getBytes(StandardCharsets.UTF_8))); listWriter.writeVarBinary( ByteBuffer.wrap("row4".getBytes(StandardCharsets.UTF_8)), 0, "row4".getBytes(StandardCharsets.UTF_8).length); listWriter.endList(); listWriter.setValueCount(1); assertEquals( "row1", new String((byte[]) listVector.getObject(0).get(0), StandardCharsets.UTF_8)); assertEquals( "row2", new String((byte[]) listVector.getObject(0).get(1), StandardCharsets.UTF_8)); assertEquals( "row3", new String((byte[]) listVector.getObject(0).get(2), StandardCharsets.UTF_8)); assertEquals( "row4", new String((byte[]) listVector.getObject(0).get(3), StandardCharsets.UTF_8)); } } @Test public void listLargeVarBinaryHelpers() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); listWriter.startList(); listWriter.writeLargeVarBinary("row1".getBytes(StandardCharsets.UTF_8)); listWriter.writeLargeVarBinary( "row2".getBytes(StandardCharsets.UTF_8), 0, "row2".getBytes(StandardCharsets.UTF_8).length); listWriter.writeLargeVarBinary(ByteBuffer.wrap("row3".getBytes(StandardCharsets.UTF_8))); listWriter.writeLargeVarBinary( ByteBuffer.wrap("row4".getBytes(StandardCharsets.UTF_8)), 0, "row4".getBytes(StandardCharsets.UTF_8).length); listWriter.endList(); listWriter.setValueCount(1); assertEquals( "row1", new String((byte[]) listVector.getObject(0).get(0), StandardCharsets.UTF_8)); assertEquals( "row2", new String((byte[]) listVector.getObject(0).get(1), StandardCharsets.UTF_8)); assertEquals( "row3", new String((byte[]) listVector.getObject(0).get(2), StandardCharsets.UTF_8)); assertEquals( "row4", new String((byte[]) listVector.getObject(0).get(3), StandardCharsets.UTF_8)); } } @Test public void unionWithVarCharAndBinaryHelpers() throws Exception { try (UnionVector vector = new UnionVector("union", allocator, /* field type */ null, /* call-back */ null)) { UnionWriter unionWriter = new UnionWriter(vector); unionWriter.allocate(); unionWriter.start(); unionWriter.setPosition(0); unionWriter.writeVarChar("row1"); unionWriter.setPosition(1); unionWriter.writeVarChar(new Text("row2")); unionWriter.setPosition(2); unionWriter.writeLargeVarChar("row3"); unionWriter.setPosition(3); unionWriter.writeLargeVarChar(new Text("row4")); unionWriter.setPosition(4); unionWriter.writeVarBinary("row5".getBytes(StandardCharsets.UTF_8)); unionWriter.setPosition(5); unionWriter.writeVarBinary( "row6".getBytes(StandardCharsets.UTF_8), 0, "row6".getBytes(StandardCharsets.UTF_8).length); unionWriter.setPosition(6); unionWriter.writeVarBinary(ByteBuffer.wrap("row7".getBytes(StandardCharsets.UTF_8))); unionWriter.setPosition(7); unionWriter.writeVarBinary( ByteBuffer.wrap("row8".getBytes(StandardCharsets.UTF_8)), 0, "row8".getBytes(StandardCharsets.UTF_8).length); unionWriter.setPosition(8); unionWriter.writeLargeVarBinary("row9".getBytes(StandardCharsets.UTF_8)); unionWriter.setPosition(9); unionWriter.writeLargeVarBinary( "row10".getBytes(StandardCharsets.UTF_8), 0, "row10".getBytes(StandardCharsets.UTF_8).length); unionWriter.setPosition(10); unionWriter.writeLargeVarBinary(ByteBuffer.wrap("row11".getBytes(StandardCharsets.UTF_8))); unionWriter.setPosition(11); unionWriter.writeLargeVarBinary( ByteBuffer.wrap("row12".getBytes(StandardCharsets.UTF_8)), 0, "row12".getBytes(StandardCharsets.UTF_8).length); unionWriter.end(); assertEquals("row1", new String(vector.getVarCharVector().get(0), StandardCharsets.UTF_8)); assertEquals("row2", new String(vector.getVarCharVector().get(1), StandardCharsets.UTF_8)); assertEquals( "row3", new String(vector.getLargeVarCharVector().get(2), StandardCharsets.UTF_8)); assertEquals( "row4", new String(vector.getLargeVarCharVector().get(3), StandardCharsets.UTF_8)); assertEquals("row5", new String(vector.getVarBinaryVector().get(4), StandardCharsets.UTF_8)); assertEquals("row6", new String(vector.getVarBinaryVector().get(5), StandardCharsets.UTF_8)); assertEquals("row7", new String(vector.getVarBinaryVector().get(6), StandardCharsets.UTF_8)); assertEquals("row8", new String(vector.getVarBinaryVector().get(7), StandardCharsets.UTF_8)); assertEquals( "row9", new String(vector.getLargeVarBinaryVector().get(8), StandardCharsets.UTF_8)); assertEquals( "row10", new String(vector.getLargeVarBinaryVector().get(9), StandardCharsets.UTF_8)); assertEquals( "row11", new String(vector.getLargeVarBinaryVector().get(10), StandardCharsets.UTF_8)); assertEquals( "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)); } } } TestSimpleWriter.java000066400000000000000000000156151515451615200346000ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/complex/writer/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.writer; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import java.nio.ByteBuffer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.LargeVarCharVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.complex.impl.LargeVarBinaryWriterImpl; import org.apache.arrow.vector.complex.impl.LargeVarCharWriterImpl; import org.apache.arrow.vector.complex.impl.VarBinaryWriterImpl; import org.apache.arrow.vector.complex.impl.VarCharWriterImpl; 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; public class TestSimpleWriter { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Integer.MAX_VALUE); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testWriteByteArrayToVarBinary() throws Exception { try (VarBinaryVector vector = new VarBinaryVector("test", allocator); VarBinaryWriter writer = new VarBinaryWriterImpl(vector)) { byte[] input = new byte[] {0x01, 0x02}; writer.writeVarBinary(input); byte[] result = vector.get(0); assertArrayEquals(input, result); } } @Test public void testWriteByteArrayWithOffsetToVarBinary() throws Exception { try (VarBinaryVector vector = new VarBinaryVector("test", allocator); VarBinaryWriter writer = new VarBinaryWriterImpl(vector)) { byte[] input = new byte[] {0x01, 0x02}; writer.writeVarBinary(input, 1, 1); byte[] result = vector.get(0); assertArrayEquals(new byte[] {0x02}, result); } } @Test public void testWriteByteBufferToVarBinary() throws Exception { try (VarBinaryVector vector = new VarBinaryVector("test", allocator); VarBinaryWriter writer = new VarBinaryWriterImpl(vector)) { byte[] input = new byte[] {0x01, 0x02}; ByteBuffer buffer = ByteBuffer.wrap(input); writer.writeVarBinary(buffer); byte[] result = vector.get(0); assertArrayEquals(input, result); } } @Test public void testWriteByteBufferWithOffsetToVarBinary() throws Exception { try (VarBinaryVector vector = new VarBinaryVector("test", allocator); VarBinaryWriter writer = new VarBinaryWriterImpl(vector)) { byte[] input = new byte[] {0x01, 0x02}; ByteBuffer buffer = ByteBuffer.wrap(input); writer.writeVarBinary(buffer, 1, 1); byte[] result = vector.get(0); assertArrayEquals(new byte[] {0x02}, result); } } @Test public void testWriteByteArrayToLargeVarBinary() throws Exception { try (LargeVarBinaryVector vector = new LargeVarBinaryVector("test", allocator); LargeVarBinaryWriter writer = new LargeVarBinaryWriterImpl(vector)) { byte[] input = new byte[] {0x01, 0x02}; writer.writeLargeVarBinary(input); byte[] result = vector.get(0); assertArrayEquals(input, result); } } @Test public void testWriteByteArrayWithOffsetToLargeVarBinary() throws Exception { try (LargeVarBinaryVector vector = new LargeVarBinaryVector("test", allocator); LargeVarBinaryWriter writer = new LargeVarBinaryWriterImpl(vector)) { byte[] input = new byte[] {0x01, 0x02}; writer.writeLargeVarBinary(input, 1, 1); byte[] result = vector.get(0); assertArrayEquals(new byte[] {0x02}, result); } } @Test public void testWriteByteBufferToLargeVarBinary() throws Exception { try (LargeVarBinaryVector vector = new LargeVarBinaryVector("test", allocator); LargeVarBinaryWriter writer = new LargeVarBinaryWriterImpl(vector)) { byte[] input = new byte[] {0x01, 0x02}; ByteBuffer buffer = ByteBuffer.wrap(input); writer.writeLargeVarBinary(buffer); byte[] result = vector.get(0); assertArrayEquals(input, result); } } @Test public void testWriteByteBufferWithOffsetToLargeVarBinary() throws Exception { try (LargeVarBinaryVector vector = new LargeVarBinaryVector("test", allocator); LargeVarBinaryWriter writer = new LargeVarBinaryWriterImpl(vector)) { byte[] input = new byte[] {0x01, 0x02}; ByteBuffer buffer = ByteBuffer.wrap(input); writer.writeLargeVarBinary(buffer, 1, 1); byte[] result = vector.get(0); assertArrayEquals(new byte[] {0x02}, result); } } @Test public void testWriteStringToVarChar() throws Exception { try (VarCharVector vector = new VarCharVector("test", allocator); VarCharWriter writer = new VarCharWriterImpl(vector)) { String input = "testInput"; writer.writeVarChar(input); String result = vector.getObject(0).toString(); assertEquals(input, result); } } @Test public void testWriteTextToVarChar() throws Exception { try (VarCharVector vector = new VarCharVector("test", allocator); VarCharWriter writer = new VarCharWriterImpl(vector)) { String input = "testInput"; writer.writeVarChar(new Text(input)); String result = vector.getObject(0).toString(); assertEquals(input, result); } } @Test public void testWriteStringToLargeVarChar() throws Exception { try (LargeVarCharVector vector = new LargeVarCharVector("test", allocator); LargeVarCharWriter writer = new LargeVarCharWriterImpl(vector)) { String input = "testInput"; writer.writeLargeVarChar(input); String result = vector.getObject(0).toString(); assertEquals(input, result); } } @Test public void testWriteTextToLargeVarChar() throws Exception { try (LargeVarCharVector vector = new LargeVarCharVector("test", allocator); LargeVarCharWriter writer = new LargeVarCharWriterImpl(vector)) { String input = "testInput"; writer.writeLargeVarChar(new Text(input)); String result = vector.getObject(0).toString(); assertEquals(input, result); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/ipc/000077500000000000000000000000001515451615200263065ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/ipc/BaseFileTest.java000066400000000000000000001070441515451615200314710ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import static org.apache.arrow.vector.TestUtils.newVarCharVector; 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.assertTrue; import java.io.IOException; import java.math.BigDecimal; import java.math.BigInteger; import java.nio.charset.StandardCharsets; import java.time.LocalDateTime; import java.time.LocalTime; import java.time.ZoneId; import java.time.ZoneOffset; import java.util.Arrays; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.Collections2; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.DateMilliVector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.NullVector; import org.apache.arrow.vector.TimeMilliVector; 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.ListVector; import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.impl.ComplexWriterImpl; import org.apache.arrow.vector.complex.impl.UnionListWriter; 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.ComplexWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.ListWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter; import org.apache.arrow.vector.complex.writer.BigIntWriter; import org.apache.arrow.vector.complex.writer.DateMilliWriter; import org.apache.arrow.vector.complex.writer.Float4Writer; import org.apache.arrow.vector.complex.writer.IntWriter; import org.apache.arrow.vector.complex.writer.TimeMilliWriter; import org.apache.arrow.vector.complex.writer.TimeStampMilliTZWriter; import org.apache.arrow.vector.complex.writer.TimeStampMilliWriter; import org.apache.arrow.vector.complex.writer.TimeStampNanoWriter; import org.apache.arrow.vector.complex.writer.UInt1Writer; import org.apache.arrow.vector.complex.writer.UInt2Writer; import org.apache.arrow.vector.complex.writer.UInt4Writer; import org.apache.arrow.vector.complex.writer.UInt8Writer; 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.holders.NullableTimeStampMilliHolder; 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.Text; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** Helps testing the file formats. */ public class BaseFileTest { private static final Logger LOGGER = LoggerFactory.getLogger(BaseFileTest.class); protected static final int COUNT = 10; protected BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Integer.MAX_VALUE); } @AfterEach public void tearDown() { allocator.close(); } private static short[] uint1Values = new short[] {0, 255, 1, 128, 2}; private static char[] uint2Values = new char[] {0, Character.MAX_VALUE, 1, Short.MAX_VALUE * 2, 2}; private static long[] uint4Values = new long[] {0, Integer.MAX_VALUE + 1L, 1, Integer.MAX_VALUE * 2L, 2}; private static BigInteger[] uint8Values = new BigInteger[] { BigInteger.valueOf(0), BigInteger.valueOf(Long.MAX_VALUE).multiply(BigInteger.valueOf(2)), BigInteger.valueOf(2), BigInteger.valueOf(Long.MAX_VALUE).add(BigInteger.valueOf(1)), BigInteger.valueOf(2) }; protected void writeData(int count, StructVector parent) { ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); IntWriter intWriter = rootWriter.integer("int"); UInt1Writer uint1Writer = rootWriter.uInt1("uint1"); UInt2Writer uint2Writer = rootWriter.uInt2("uint2"); UInt4Writer uint4Writer = rootWriter.uInt4("uint4"); UInt8Writer uint8Writer = rootWriter.uInt8("uint8"); BigIntWriter bigIntWriter = rootWriter.bigInt("bigInt"); Float4Writer float4Writer = rootWriter.float4("float"); for (int i = 0; i < count; i++) { intWriter.setPosition(i); intWriter.writeInt(i); uint1Writer.setPosition(i); // TODO: Fix add safe write methods on uint methods. uint1Writer.setPosition(i); uint1Writer.writeUInt1((byte) uint1Values[i % uint1Values.length]); uint2Writer.setPosition(i); uint2Writer.writeUInt2((char) uint2Values[i % uint2Values.length]); uint4Writer.setPosition(i); uint4Writer.writeUInt4((int) uint4Values[i % uint4Values.length]); uint8Writer.setPosition(i); uint8Writer.writeUInt8(uint8Values[i % uint8Values.length].longValue()); bigIntWriter.setPosition(i); bigIntWriter.writeBigInt(i); float4Writer.setPosition(i); float4Writer.writeFloat4(i == 0 ? Float.NaN : i); } writer.setValueCount(count); } protected void validateContent(int count, VectorSchemaRoot root) { for (int i = 0; i < count; i++) { assertEquals(i, root.getVector("int").getObject(i)); assertEquals( (Short) uint1Values[i % uint1Values.length], ((UInt1Vector) root.getVector("uint1")).getObjectNoOverflow(i)); assertEquals( (Character) uint2Values[i % uint2Values.length], (Character) ((UInt2Vector) root.getVector("uint2")).get(i), "Failed for index: " + i); assertEquals( (Long) uint4Values[i % uint4Values.length], ((UInt4Vector) root.getVector("uint4")).getObjectNoOverflow(i), "Failed for index: " + i); assertEquals( uint8Values[i % uint8Values.length], ((UInt8Vector) root.getVector("uint8")).getObjectNoOverflow(i), "Failed for index: " + i); assertEquals(Long.valueOf(i), root.getVector("bigInt").getObject(i)); assertEquals(i == 0 ? Float.NaN : i, root.getVector("float").getObject(i)); } } protected void writeComplexData(int count, StructVector parent) { ArrowBuf varchar = allocator.buffer(3); varchar.readerIndex(0); varchar.setByte(0, 'a'); varchar.setByte(1, 'b'); varchar.setByte(2, 'c'); varchar.writerIndex(3); ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); IntWriter intWriter = rootWriter.integer("int"); BigIntWriter bigIntWriter = rootWriter.bigInt("bigInt"); ListWriter listWriter = rootWriter.list("list"); StructWriter structWriter = rootWriter.struct("struct"); for (int i = 0; i < count; i++) { if (i % 5 != 3) { intWriter.setPosition(i); intWriter.writeInt(i); } bigIntWriter.setPosition(i); bigIntWriter.writeBigInt(i); listWriter.setPosition(i); listWriter.startList(); for (int j = 0; j < i % 3; j++) { listWriter.varChar().writeVarChar(0, 3, varchar); } listWriter.endList(); structWriter.setPosition(i); structWriter.start(); structWriter.timeStampMilli("timestamp").writeTimeStampMilli(i); structWriter.end(); } writer.setValueCount(count); varchar.getReferenceManager().release(); } public void printVectors(List vectors) { for (FieldVector vector : vectors) { LOGGER.debug(vector.getField().getName()); int valueCount = vector.getValueCount(); for (int i = 0; i < valueCount; i++) { LOGGER.debug(String.valueOf(vector.getObject(i))); } } } protected void validateComplexContent(int count, VectorSchemaRoot root) { assertEquals(count, root.getRowCount()); printVectors(root.getFieldVectors()); for (int i = 0; i < count; i++) { Object intVal = root.getVector("int").getObject(i); if (i % 5 != 3) { assertEquals(i, intVal); } else { assertNull(intVal); } assertEquals(Long.valueOf(i), root.getVector("bigInt").getObject(i)); assertEquals(i % 3, ((List) root.getVector("list").getObject(i)).size()); NullableTimeStampMilliHolder h = new NullableTimeStampMilliHolder(); FieldReader structReader = root.getVector("struct").getReader(); structReader.setPosition(i); structReader.reader("timestamp").read(h); assertEquals(i, h.value); } } private LocalDateTime makeDateTimeFromCount(int i) { return LocalDateTime.of(2000 + i, 1 + i, 1 + i, i, i, i, i * 100_000_000 + i); } protected void writeDateTimeData(int count, StructVector parent) { assertTrue(count < 100); ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); DateMilliWriter dateWriter = rootWriter.dateMilli("date"); TimeMilliWriter timeWriter = rootWriter.timeMilli("time"); TimeStampMilliWriter timeStampMilliWriter = rootWriter.timeStampMilli("timestamp-milli"); TimeStampMilliTZWriter timeStampMilliTZWriter = rootWriter.timeStampMilliTZ("timestamp-milliTZ", "Europe/Paris"); TimeStampNanoWriter timeStampNanoWriter = rootWriter.timeStampNano("timestamp-nano"); for (int i = 0; i < count; i++) { LocalDateTime dt = makeDateTimeFromCount(i); // Number of days in milliseconds since epoch, stored as 64-bit integer, only date part is // used dateWriter.setPosition(i); long dateLong = dt.toLocalDate().atStartOfDay().toInstant(ZoneOffset.UTC).toEpochMilli(); dateWriter.writeDateMilli(dateLong); // Time is a value in milliseconds since midnight, stored as 32-bit integer timeWriter.setPosition(i); int milliOfDay = (int) java.util.concurrent.TimeUnit.NANOSECONDS.toMillis(dt.toLocalTime().toNanoOfDay()); timeWriter.writeTimeMilli(milliOfDay); // Timestamp as milliseconds since the epoch, stored as 64-bit integer timeStampMilliWriter.setPosition(i); timeStampMilliWriter.writeTimeStampMilli(dt.toInstant(ZoneOffset.UTC).toEpochMilli()); // Timestamp as milliseconds since epoch with timezone timeStampMilliTZWriter.setPosition(i); timeStampMilliTZWriter.writeTimeStampMilliTZ( dt.atZone(ZoneId.of("Europe/Paris")).toInstant().toEpochMilli()); // Timestamp as nanoseconds since epoch timeStampNanoWriter.setPosition(i); long tsNanos = dt.toInstant(ZoneOffset.UTC).toEpochMilli() * 1_000_000 + i; // need to add back in nano val timeStampNanoWriter.writeTimeStampNano(tsNanos); } writer.setValueCount(count); } protected void validateDateTimeContent(int count, VectorSchemaRoot root) { assertEquals(count, root.getRowCount()); printVectors(root.getFieldVectors()); for (int i = 0; i < count; i++) { LocalDateTime dt = makeDateTimeFromCount(i); LocalDateTime dtMilli = dt.minusNanos(i); LocalDateTime dateVal = ((DateMilliVector) root.getVector("date")).getObject(i); LocalDateTime dateExpected = dt.toLocalDate().atStartOfDay(); assertEquals(dateExpected, dateVal); LocalTime timeVal = ((TimeMilliVector) root.getVector("time")).getObject(i).toLocalTime(); assertEquals(dtMilli.toLocalTime(), timeVal); Object timestampMilliVal = root.getVector("timestamp-milli").getObject(i); assertEquals(dtMilli, timestampMilliVal); Object timestampMilliTZVal = root.getVector("timestamp-milliTZ").getObject(i); assertEquals( dt.atZone(ZoneId.of("Europe/Paris")).toInstant().toEpochMilli(), timestampMilliTZVal); Object timestampNanoVal = root.getVector("timestamp-nano").getObject(i); assertEquals(dt, timestampNanoVal); } } protected VectorSchemaRoot writeFlatDictionaryData( BufferAllocator bufferAllocator, DictionaryProvider.MapDictionaryProvider provider) { // Define dictionaries and add to provider VarCharVector dictionary1Vector = newVarCharVector("D1", bufferAllocator); dictionary1Vector.allocateNewSafe(); dictionary1Vector.set(0, "foo".getBytes(StandardCharsets.UTF_8)); dictionary1Vector.set(1, "bar".getBytes(StandardCharsets.UTF_8)); dictionary1Vector.set(2, "baz".getBytes(StandardCharsets.UTF_8)); dictionary1Vector.setValueCount(3); Dictionary dictionary1 = new Dictionary(dictionary1Vector, new DictionaryEncoding(1L, false, null)); provider.put(dictionary1); VarCharVector dictionary2Vector = newVarCharVector("D2", bufferAllocator); dictionary2Vector.allocateNewSafe(); dictionary2Vector.set(0, "micro".getBytes(StandardCharsets.UTF_8)); dictionary2Vector.set(1, "small".getBytes(StandardCharsets.UTF_8)); dictionary2Vector.set(2, "large".getBytes(StandardCharsets.UTF_8)); dictionary2Vector.setValueCount(3); Dictionary dictionary2 = new Dictionary(dictionary2Vector, new DictionaryEncoding(2L, false, null)); provider.put(dictionary2); // Populate the vectors VarCharVector vector1A = newVarCharVector("varcharA", bufferAllocator); vector1A.allocateNewSafe(); vector1A.set(0, "foo".getBytes(StandardCharsets.UTF_8)); vector1A.set(1, "bar".getBytes(StandardCharsets.UTF_8)); vector1A.set(3, "baz".getBytes(StandardCharsets.UTF_8)); vector1A.set(4, "bar".getBytes(StandardCharsets.UTF_8)); vector1A.set(5, "baz".getBytes(StandardCharsets.UTF_8)); vector1A.setValueCount(6); FieldVector encodedVector1A = (FieldVector) DictionaryEncoder.encode(vector1A, dictionary1); vector1A.close(); // Done with this vector after encoding // Write this vector using indices instead of encoding IntVector encodedVector1B = new IntVector("varcharB", bufferAllocator); encodedVector1B.allocateNewSafe(); encodedVector1B.set(0, 2); // "baz" encodedVector1B.set(1, 1); // "bar" encodedVector1B.set(2, 2); // "baz" encodedVector1B.set(4, 1); // "bar" encodedVector1B.set(5, 0); // "foo" encodedVector1B.setValueCount(6); VarCharVector vector2 = newVarCharVector("sizes", bufferAllocator); vector2.allocateNewSafe(); vector2.set(1, "large".getBytes(StandardCharsets.UTF_8)); vector2.set(2, "small".getBytes(StandardCharsets.UTF_8)); vector2.set(3, "small".getBytes(StandardCharsets.UTF_8)); vector2.set(4, "large".getBytes(StandardCharsets.UTF_8)); vector2.setValueCount(6); FieldVector encodedVector2 = (FieldVector) DictionaryEncoder.encode(vector2, dictionary2); vector2.close(); // Done with this vector after encoding List fields = Arrays.asList( encodedVector1A.getField(), encodedVector1B.getField(), encodedVector2.getField()); List vectors = Collections2.asImmutableList(encodedVector1A, encodedVector1B, encodedVector2); return new VectorSchemaRoot(fields, vectors, encodedVector1A.getValueCount()); } protected void validateFlatDictionary(VectorSchemaRoot root, DictionaryProvider provider) { FieldVector vector1A = root.getVector("varcharA"); assertNotNull(vector1A); DictionaryEncoding encoding1A = vector1A.getField().getDictionary(); assertNotNull(encoding1A); assertEquals(1L, encoding1A.getId()); assertEquals(6, vector1A.getValueCount()); assertEquals(0, vector1A.getObject(0)); assertEquals(1, vector1A.getObject(1)); assertEquals(null, vector1A.getObject(2)); assertEquals(2, vector1A.getObject(3)); assertEquals(1, vector1A.getObject(4)); assertEquals(2, vector1A.getObject(5)); FieldVector vector1B = root.getVector("varcharB"); assertNotNull(vector1B); DictionaryEncoding encoding1B = vector1A.getField().getDictionary(); assertNotNull(encoding1B); assertTrue(encoding1A.equals(encoding1B)); assertEquals(1L, encoding1B.getId()); assertEquals(6, vector1B.getValueCount()); assertEquals(2, vector1B.getObject(0)); assertEquals(1, vector1B.getObject(1)); assertEquals(2, vector1B.getObject(2)); assertEquals(null, vector1B.getObject(3)); assertEquals(1, vector1B.getObject(4)); assertEquals(0, vector1B.getObject(5)); FieldVector vector2 = root.getVector("sizes"); assertNotNull(vector2); DictionaryEncoding encoding2 = vector2.getField().getDictionary(); assertNotNull(encoding2); assertEquals(2L, encoding2.getId()); assertEquals(6, vector2.getValueCount()); assertEquals(null, vector2.getObject(0)); assertEquals(2, vector2.getObject(1)); assertEquals(1, vector2.getObject(2)); assertEquals(1, vector2.getObject(3)); assertEquals(2, vector2.getObject(4)); assertEquals(null, vector2.getObject(5)); Dictionary dictionary1 = provider.lookup(1L); assertNotNull(dictionary1); VarCharVector dictionaryVector = ((VarCharVector) dictionary1.getVector()); assertEquals(3, dictionaryVector.getValueCount()); assertEquals(new Text("foo"), dictionaryVector.getObject(0)); assertEquals(new Text("bar"), dictionaryVector.getObject(1)); assertEquals(new Text("baz"), dictionaryVector.getObject(2)); Dictionary dictionary2 = provider.lookup(2L); assertNotNull(dictionary2); dictionaryVector = ((VarCharVector) dictionary2.getVector()); assertEquals(3, dictionaryVector.getValueCount()); assertEquals(new Text("micro"), dictionaryVector.getObject(0)); assertEquals(new Text("small"), dictionaryVector.getObject(1)); assertEquals(new Text("large"), dictionaryVector.getObject(2)); } protected VectorSchemaRoot writeNestedDictionaryData( BufferAllocator bufferAllocator, DictionaryProvider.MapDictionaryProvider provider) { // Define the dictionary and add to the provider VarCharVector dictionaryVector = newVarCharVector("D2", bufferAllocator); dictionaryVector.allocateNewSafe(); dictionaryVector.set(0, "foo".getBytes(StandardCharsets.UTF_8)); dictionaryVector.set(1, "bar".getBytes(StandardCharsets.UTF_8)); dictionaryVector.setValueCount(2); Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(2L, false, null)); provider.put(dictionary); // Write the vector data using dictionary indices ListVector listVector = ListVector.empty("list", bufferAllocator); DictionaryEncoding encoding = dictionary.getEncoding(); listVector.addOrGetVector(new FieldType(true, encoding.getIndexType(), encoding)); listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); listWriter.startList(); listWriter.writeInt(0); listWriter.writeInt(1); listWriter.endList(); listWriter.startList(); listWriter.writeInt(0); listWriter.endList(); listWriter.startList(); listWriter.writeInt(1); listWriter.endList(); listWriter.setValueCount(3); List fields = Collections2.asImmutableList(listVector.getField()); List vectors = Collections2.asImmutableList(listVector); return new VectorSchemaRoot(fields, vectors, 3); } protected void validateNestedDictionary(VectorSchemaRoot root, DictionaryProvider provider) { FieldVector vector = root.getFieldVectors().get(0); assertNotNull(vector); assertNull(vector.getField().getDictionary()); Field nestedField = vector.getField().getChildren().get(0); DictionaryEncoding encoding = nestedField.getDictionary(); assertNotNull(encoding); assertEquals(2L, encoding.getId()); assertEquals(new ArrowType.Int(32, true), encoding.getIndexType()); assertEquals(3, vector.getValueCount()); assertEquals(Arrays.asList(0, 1), vector.getObject(0)); assertEquals(Arrays.asList(0), vector.getObject(1)); assertEquals(Arrays.asList(1), vector.getObject(2)); Dictionary dictionary = provider.lookup(2L); assertNotNull(dictionary); VarCharVector dictionaryVector = ((VarCharVector) dictionary.getVector()); assertEquals(2, dictionaryVector.getValueCount()); assertEquals(new Text("foo"), dictionaryVector.getObject(0)); assertEquals(new Text("bar"), dictionaryVector.getObject(1)); } protected VectorSchemaRoot writeDecimalData(BufferAllocator bufferAllocator) { DecimalVector decimalVector1 = new DecimalVector("decimal1", bufferAllocator, 10, 3); DecimalVector decimalVector2 = new DecimalVector("decimal2", bufferAllocator, 4, 2); DecimalVector decimalVector3 = new DecimalVector("decimal3", bufferAllocator, 16, 8); int count = 10; decimalVector1.allocateNew(count); decimalVector2.allocateNew(count); decimalVector3.allocateNew(count); for (int i = 0; i < count; i++) { decimalVector1.setSafe(i, new BigDecimal(BigInteger.valueOf(i), 3)); decimalVector2.setSafe(i, new BigDecimal(BigInteger.valueOf(i * (1 << 10)), 2)); decimalVector3.setSafe(i, new BigDecimal(BigInteger.valueOf(i * 1111111111111111L), 8)); } decimalVector1.setValueCount(count); decimalVector2.setValueCount(count); decimalVector3.setValueCount(count); List fields = Collections2.asImmutableList( decimalVector1.getField(), decimalVector2.getField(), decimalVector3.getField()); List vectors = Collections2.asImmutableList(decimalVector1, decimalVector2, decimalVector3); return new VectorSchemaRoot(fields, vectors, count); } protected void validateDecimalData(VectorSchemaRoot root) { DecimalVector decimalVector1 = (DecimalVector) root.getVector("decimal1"); DecimalVector decimalVector2 = (DecimalVector) root.getVector("decimal2"); DecimalVector decimalVector3 = (DecimalVector) root.getVector("decimal3"); int count = 10; assertEquals(count, root.getRowCount()); for (int i = 0; i < count; i++) { // Verify decimal 1 vector BigDecimal readValue = decimalVector1.getObject(i); ArrowType.Decimal type = (ArrowType.Decimal) decimalVector1.getField().getType(); BigDecimal genValue = new BigDecimal(BigInteger.valueOf(i), type.getScale()); assertEquals(genValue, readValue); // Verify decimal 2 vector readValue = decimalVector2.getObject(i); type = (ArrowType.Decimal) decimalVector2.getField().getType(); genValue = new BigDecimal(BigInteger.valueOf(i * (1 << 10)), type.getScale()); assertEquals(genValue, readValue); // Verify decimal 3 vector readValue = decimalVector3.getObject(i); type = (ArrowType.Decimal) decimalVector3.getField().getType(); genValue = new BigDecimal(BigInteger.valueOf(i * 1111111111111111L), type.getScale()); assertEquals(genValue, readValue); } } protected VectorSchemaRoot writeNullData(int valueCount) { NullVector nullVector1 = new NullVector("vector1"); NullVector nullVector2 = new NullVector("vector2"); nullVector1.setValueCount(valueCount); nullVector2.setValueCount(valueCount); List fields = Collections2.asImmutableList(nullVector1.getField(), nullVector2.getField()); List vectors = Collections2.asImmutableList(nullVector1, nullVector2); return new VectorSchemaRoot(fields, vectors, valueCount); } protected void validateNullData(VectorSchemaRoot root, int valueCount) { NullVector vector1 = (NullVector) root.getFieldVectors().get(0); NullVector vector2 = (NullVector) root.getFieldVectors().get(1); assertEquals(valueCount, vector1.getValueCount()); assertEquals(valueCount, vector2.getValueCount()); } public void validateUnionData(int count, VectorSchemaRoot root) { FieldReader unionReader = root.getVector("union").getReader(); for (int i = 0; i < count; i++) { unionReader.setPosition(i); switch (i % 5) { case 0: assertEquals(i, unionReader.readInteger().intValue()); break; case 1: assertEquals(i, unionReader.readLong().longValue()); break; case 2: assertEquals(i % 3, unionReader.size()); break; case 3: assertEquals(3, unionReader.size()); break; case 4: NullableTimeStampMilliHolder h = new NullableTimeStampMilliHolder(); unionReader.reader("timestamp").read(h); assertEquals(i, h.value); break; default: assert false : "Unexpected value in switch statement: " + i; } } } public void writeUnionData(int count, StructVector parent) { ArrowBuf varchar = allocator.buffer(3); varchar.readerIndex(0); varchar.setByte(0, 'a'); varchar.setByte(1, 'b'); varchar.setByte(2, 'c'); varchar.writerIndex(3); ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); IntWriter intWriter = rootWriter.integer("union"); BigIntWriter bigIntWriter = rootWriter.bigInt("union"); ListWriter listWriter = rootWriter.list("union"); ListWriter listViewWriter = rootWriter.listView("union"); StructWriter structWriter = rootWriter.struct("union"); for (int i = 0; i < count; i++) { switch (i % 5) { case 0: intWriter.setPosition(i); intWriter.writeInt(i); break; case 1: bigIntWriter.setPosition(i); bigIntWriter.writeBigInt(i); break; case 2: listWriter.setPosition(i); listWriter.startList(); for (int j = 0; j < i % 3; j++) { listWriter.varChar().writeVarChar(0, 3, varchar); } listWriter.endList(); break; case 3: listViewWriter.setPosition(i); listViewWriter.startListView(); for (int j = 0; j < i % 5; j++) { listViewWriter.varChar().writeVarChar(0, 3, varchar); } listViewWriter.endListView(); break; case 4: structWriter.setPosition(i); structWriter.start(); structWriter.timeStampMilli("timestamp").writeTimeStampMilli(i); structWriter.end(); break; default: assert false : "Unexpected value in switch statement: " + i; } } writer.setValueCount(count); varchar.getReferenceManager().release(); } protected void writeVarBinaryData(int count, StructVector parent) { assertTrue(count < 100); ComplexWriter writer = new ComplexWriterImpl("root", parent); StructWriter rootWriter = writer.rootAsStruct(); ListWriter listWriter = rootWriter.list("list"); ArrowBuf varbin = allocator.buffer(count); for (int i = 0; i < count; i++) { varbin.setByte(i, i); listWriter.setPosition(i); listWriter.startList(); for (int j = 0; j < i % 3; j++) { listWriter.varBinary().writeVarBinary(0, i + 1, varbin); } listWriter.endList(); } writer.setValueCount(count); varbin.getReferenceManager().release(); } protected void validateVarBinary(int count, VectorSchemaRoot root) { assertEquals(count, root.getRowCount()); ListVector listVector = (ListVector) root.getVector("list"); byte[] expectedArray = new byte[count]; int numVarBinaryValues = 0; for (int i = 0; i < count; i++) { expectedArray[i] = (byte) i; List objList = listVector.getObject(i); if (i % 3 == 0) { assertTrue(objList.isEmpty()); } else { byte[] expected = Arrays.copyOfRange(expectedArray, 0, i + 1); for (int j = 0; j < i % 3; j++) { byte[] result = (byte[]) objList.get(j); assertArrayEquals(result, expected); numVarBinaryValues++; } } } // ListVector lastSet should be the index of last value + 1 assertEquals(listVector.getLastSet(), count - 1); // VarBinaryVector lastSet should be the index of last value VarBinaryVector binaryVector = (VarBinaryVector) listVector.getChildrenFromFields().get(0); assertEquals(binaryVector.getLastSet(), numVarBinaryValues - 1); } protected void writeBatchData(ArrowWriter writer, IntVector vector, VectorSchemaRoot root) throws IOException { writer.start(); vector.setNull(0); vector.setSafe(1, 1); vector.setSafe(2, 2); vector.setNull(3); vector.setSafe(4, 1); vector.setValueCount(5); root.setRowCount(5); writer.writeBatch(); vector.setNull(0); vector.setSafe(1, 1); vector.setSafe(2, 2); vector.setValueCount(3); root.setRowCount(3); writer.writeBatch(); writer.end(); } protected void validateBatchData(ArrowReader reader, IntVector vector) throws IOException { reader.loadNextBatch(); assertEquals(5, vector.getValueCount()); assertTrue(vector.isNull(0)); assertEquals(1, vector.get(1)); assertEquals(2, vector.get(2)); assertTrue(vector.isNull(3)); assertEquals(1, vector.get(4)); reader.loadNextBatch(); assertEquals(3, vector.getValueCount()); assertTrue(vector.isNull(0)); assertEquals(1, vector.get(1)); assertEquals(2, vector.get(2)); } protected VectorSchemaRoot writeMapData(BufferAllocator bufferAllocator) { MapVector mapVector = MapVector.empty("map", bufferAllocator, false); MapVector sortedMapVector = MapVector.empty("mapSorted", bufferAllocator, true); mapVector.allocateNew(); sortedMapVector.allocateNew(); UnionMapWriter mapWriter = mapVector.getWriter(); UnionMapWriter sortedMapWriter = sortedMapVector.getWriter(); final int count = 10; for (int i = 0; i < count; i++) { // Write mapVector with NULL values // i == 1 is a NULL if (i != 1) { mapWriter.setPosition(i); mapWriter.startMap(); // i == 3 is an empty map if (i != 3) { for (int j = 0; j < i + 1; j++) { mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(j); // i == 5 maps to a NULL value if (i != 5) { mapWriter.value().integer().writeInt(j); } mapWriter.endEntry(); } } mapWriter.endMap(); } // Write sortedMapVector sortedMapWriter.setPosition(i); sortedMapWriter.startMap(); for (int j = 0; j < i + 1; j++) { sortedMapWriter.startEntry(); sortedMapWriter.key().bigInt().writeBigInt(j); sortedMapWriter.value().integer().writeInt(j); sortedMapWriter.endEntry(); } sortedMapWriter.endMap(); } mapWriter.setValueCount(COUNT); sortedMapWriter.setValueCount(COUNT); List fields = Collections2.asImmutableList(mapVector.getField(), sortedMapVector.getField()); List vectors = Collections2.asImmutableList(mapVector, sortedMapVector); return new VectorSchemaRoot(fields, vectors, count); } protected void validateMapData(VectorSchemaRoot root) { MapVector mapVector = (MapVector) root.getVector("map"); MapVector sortedMapVector = (MapVector) root.getVector("mapSorted"); final int count = 10; assertEquals(count, root.getRowCount()); UnionMapReader mapReader = new UnionMapReader(mapVector); UnionMapReader sortedMapReader = new UnionMapReader(sortedMapVector); for (int i = 0; i < count; i++) { // Read mapVector with NULL values mapReader.setPosition(i); if (i == 1) { assertFalse(mapReader.isSet()); } else { if (i == 3) { JsonStringArrayList result = (JsonStringArrayList) mapReader.readObject(); assertTrue(result.isEmpty()); } else { for (int j = 0; j < i + 1; j++) { mapReader.next(); assertEquals(j, mapReader.key().readLong().longValue()); if (i == 5) { assertFalse(mapReader.value().isSet()); } else { assertEquals(j, mapReader.value().readInteger().intValue()); } } } } // Read sortedMapVector sortedMapReader.setPosition(i); for (int j = 0; j < i + 1; j++) { sortedMapReader.next(); assertEquals(j, sortedMapReader.key().readLong().longValue()); assertEquals(j, sortedMapReader.value().readInteger().intValue()); } } } protected VectorSchemaRoot writeListAsMapData(BufferAllocator bufferAllocator) { ListVector mapEntryList = ListVector.empty("entryList", bufferAllocator); FieldType mapEntryType = new FieldType(false, ArrowType.Struct.INSTANCE, null, null); StructVector mapEntryData = new StructVector("entryData", bufferAllocator, mapEntryType, null); mapEntryData.addOrGet( "myKey", new FieldType(false, new ArrowType.Int(64, true), null), BigIntVector.class); mapEntryData.addOrGet( "myValue", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); mapEntryList.initializeChildrenFromFields( Collections2.asImmutableList(mapEntryData.getField())); UnionListWriter entryWriter = mapEntryList.getWriter(); entryWriter.allocate(); final int count = 10; for (int i = 0; i < count; i++) { entryWriter.setPosition(i); entryWriter.startList(); for (int j = 0; j < i + 1; j++) { entryWriter.struct().start(); entryWriter.struct().bigInt("myKey").writeBigInt(j); entryWriter.struct().integer("myValue").writeInt(j); entryWriter.struct().end(); } entryWriter.endList(); } entryWriter.setValueCount(COUNT); MapVector mapVector = MapVector.empty("map", bufferAllocator, false); mapEntryList.makeTransferPair(mapVector).transfer(); List fields = Collections2.asImmutableList(mapVector.getField()); List vectors = Collections2.asImmutableList(mapVector); return new VectorSchemaRoot(fields, vectors, count); } protected void validateListAsMapData(VectorSchemaRoot root) { MapVector sortedMapVector = (MapVector) root.getVector("map"); final int count = 10; assertEquals(count, root.getRowCount()); UnionMapReader sortedMapReader = new UnionMapReader(sortedMapVector); sortedMapReader.setKeyValueNames("myKey", "myValue"); for (int i = 0; i < count; i++) { sortedMapReader.setPosition(i); for (int j = 0; j < i + 1; j++) { sortedMapReader.next(); assertEquals(j, sortedMapReader.key().readLong().longValue()); assertEquals(j, sortedMapReader.value().readInteger().intValue()); } } } } ITTestIPCWithLargeArrowBuffers.java000066400000000000000000000172411515451615200347460ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/ipc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; 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 java.io.File; import java.io.FileInputStream; import java.io.FileOutputStream; import java.io.IOException; import java.util.Arrays; import java.util.Map; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.IntVector; 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.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.types.pojo.Schema; import org.junit.jupiter.api.Test; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * Integration test for reading/writing {@link org.apache.arrow.vector.VectorSchemaRoot} with large * (more than 2GB) buffers by {@link ArrowReader} and {@link ArrowWriter}.. To run this test, please * make sure there is at least 8GB free memory, and 8GB free.disk space in the system. */ public class ITTestIPCWithLargeArrowBuffers { private static final Logger logger = LoggerFactory.getLogger(ITTestIPCWithLargeArrowBuffers.class); // 4GB buffer size static final long BUFFER_SIZE = 4 * 1024 * 1024 * 1024L; static final int DICTIONARY_VECTOR_SIZE = (int) (BUFFER_SIZE / BigIntVector.TYPE_WIDTH); static final int ENCODED_VECTOR_SIZE = (int) (BUFFER_SIZE / IntVector.TYPE_WIDTH); static final String FILE_NAME = "largeArrowData.data"; static final long DICTIONARY_ID = 123L; static final ArrowType.Int ENCODED_VECTOR_TYPE = new ArrowType.Int(32, true); static final DictionaryEncoding DICTIONARY_ENCODING = new DictionaryEncoding(DICTIONARY_ID, false, ENCODED_VECTOR_TYPE); static final FieldType ENCODED_FIELD_TYPE = new FieldType(true, ENCODED_VECTOR_TYPE, DICTIONARY_ENCODING, null); static final Field ENCODED_VECTOR_FIELD = new Field("encoded vector", ENCODED_FIELD_TYPE, null); private void testWriteLargeArrowData(boolean streamMode) throws IOException { // simulate encoding big int as int try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); BigIntVector dictVector = new BigIntVector("dic vector", allocator); FileOutputStream out = new FileOutputStream(FILE_NAME); IntVector encodedVector = (IntVector) ENCODED_VECTOR_FIELD.createVector(allocator)) { // prepare dictionary provider. DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); Dictionary dictionary = new Dictionary(dictVector, DICTIONARY_ENCODING); provider.put(dictionary); // populate the dictionary vector dictVector.allocateNew(DICTIONARY_VECTOR_SIZE); for (int i = 0; i < DICTIONARY_VECTOR_SIZE; i++) { dictVector.set(i, i); } dictVector.setValueCount(DICTIONARY_VECTOR_SIZE); assertTrue(dictVector.getDataBuffer().capacity() > Integer.MAX_VALUE); logger.trace("Populating dictionary vector finished"); // populate the encoded vector encodedVector.allocateNew(ENCODED_VECTOR_SIZE); for (int i = 0; i < ENCODED_VECTOR_SIZE; i++) { encodedVector.set(i, i % DICTIONARY_VECTOR_SIZE); } encodedVector.setValueCount(ENCODED_VECTOR_SIZE); assertTrue(encodedVector.getDataBuffer().capacity() > Integer.MAX_VALUE); logger.trace("Populating encoded vector finished"); // build vector schema root and write data. try (VectorSchemaRoot root = new VectorSchemaRoot( Arrays.asList(ENCODED_VECTOR_FIELD), Arrays.asList(encodedVector), ENCODED_VECTOR_SIZE); ArrowWriter writer = streamMode ? new ArrowStreamWriter(root, provider, out) : new ArrowFileWriter(root, provider, out.getChannel())) { writer.start(); writer.writeBatch(); writer.end(); logger.trace("Writing data finished"); } } assertTrue(new File(FILE_NAME).exists()); } private void testReadLargeArrowData(boolean streamMode) throws IOException { try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); FileInputStream in = new FileInputStream(FILE_NAME); ArrowReader reader = streamMode ? new ArrowStreamReader(in, allocator) : new ArrowFileReader(in.getChannel(), allocator)) { // verify schema Schema readSchema = reader.getVectorSchemaRoot().getSchema(); assertEquals(1, readSchema.getFields().size()); assertEquals(ENCODED_VECTOR_FIELD, readSchema.getFields().get(0)); logger.trace("Verifying schema finished"); // verify vector schema root assertTrue(reader.loadNextBatch()); VectorSchemaRoot root = reader.getVectorSchemaRoot(); assertEquals(ENCODED_VECTOR_SIZE, root.getRowCount()); assertEquals(1, root.getFieldVectors().size()); assertTrue(root.getFieldVectors().get(0) instanceof IntVector); IntVector encodedVector = (IntVector) root.getVector(0); for (int i = 0; i < ENCODED_VECTOR_SIZE; i++) { assertEquals(i % DICTIONARY_VECTOR_SIZE, encodedVector.get(i)); } logger.trace("Verifying encoded vector finished"); // verify dictionary Map dictVectors = reader.getDictionaryVectors(); assertEquals(1, dictVectors.size()); Dictionary dictionary = dictVectors.get(DICTIONARY_ID); assertNotNull(dictionary); assertTrue(dictionary.getVector() instanceof BigIntVector); BigIntVector dictVector = (BigIntVector) dictionary.getVector(); assertEquals(DICTIONARY_VECTOR_SIZE, dictVector.getValueCount()); for (int i = 0; i < DICTIONARY_VECTOR_SIZE; i++) { assertEquals(i, dictVector.get(i)); } logger.trace("Verifying dictionary vector finished"); // ensure no more data available assertFalse(reader.loadNextBatch()); } finally { File dataFile = new File(FILE_NAME); dataFile.delete(); assertFalse(dataFile.exists()); } } @Test public void testIPC() throws IOException { logger.trace("Start testing reading/writing large arrow stream data"); testWriteLargeArrowData(true); testReadLargeArrowData(true); logger.trace("Finish testing reading/writing large arrow stream data"); logger.trace("Start testing reading/writing large arrow file data"); testWriteLargeArrowData(false); testReadLargeArrowData(false); logger.trace("Finish testing reading/writing large arrow file data"); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/ipc/MessageSerializerTest.java000066400000000000000000000234411515451615200334330ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import static java.util.Arrays.asList; import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import static org.junit.jupiter.api.Assertions.assertArrayEquals; 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.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.nio.channels.Channels; import java.util.Arrays; import java.util.Collections; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.ipc.message.ArrowMessage; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.ipc.message.IpcOption; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.MetadataVersion; 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.types.pojo.Schema; import org.junit.jupiter.api.Test; public class MessageSerializerTest { public static ArrowBuf buf(BufferAllocator alloc, byte[] bytes) { ArrowBuf buffer = alloc.buffer(bytes.length); buffer.writeBytes(bytes); return buffer; } public static byte[] array(ArrowBuf buf) { byte[] bytes = new byte[checkedCastToInt(buf.readableBytes())]; buf.readBytes(bytes); return bytes; } private int intToByteRoundtrip(int v, byte[] bytes) { MessageSerializer.intToBytes(v, bytes); return MessageSerializer.bytesToInt(bytes); } @Test public void testIntToBytes() { byte[] bytes = new byte[4]; int[] values = new int[] {1, 15, 1 << 8, 1 << 16, Integer.MAX_VALUE}; for (int v : values) { assertEquals(intToByteRoundtrip(v, bytes), v); } } @Test public void testWriteMessageBufferAligned() throws IOException { ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); WriteChannel out = new WriteChannel(Channels.newChannel(outputStream)); // This is not a valid Arrow Message, only to test writing and alignment ByteBuffer buffer = ByteBuffer.allocate(8).order(ByteOrder.nativeOrder()); buffer.putInt(1); buffer.putInt(2); buffer.flip(); int bytesWritten = MessageSerializer.writeMessageBuffer(out, 8, buffer); assertEquals(16, bytesWritten); buffer.rewind(); buffer.putInt(3); buffer.flip(); bytesWritten = MessageSerializer.writeMessageBuffer(out, 4, buffer); assertEquals(16, bytesWritten); ByteArrayInputStream inputStream = new ByteArrayInputStream(outputStream.toByteArray()); ReadChannel in = new ReadChannel(Channels.newChannel(inputStream)); ByteBuffer result = ByteBuffer.allocate(32).order(ByteOrder.nativeOrder()); in.readFully(result); result.rewind(); // First message continuation, size, and 2 int values assertEquals(MessageSerializer.IPC_CONTINUATION_TOKEN, result.getInt()); // message length is represented in little endian result.order(ByteOrder.LITTLE_ENDIAN); assertEquals(8, result.getInt()); result.order(ByteOrder.nativeOrder()); assertEquals(1, result.getInt()); assertEquals(2, result.getInt()); // Second message continuation, size, 1 int value and 4 bytes padding assertEquals(MessageSerializer.IPC_CONTINUATION_TOKEN, result.getInt()); // message length is represented in little endian result.order(ByteOrder.LITTLE_ENDIAN); assertEquals(8, result.getInt()); result.order(ByteOrder.nativeOrder()); assertEquals(3, result.getInt()); assertEquals(0, result.getInt()); } @Test public void testSchemaMessageSerialization() throws IOException { Schema schema = testSchema(); ByteArrayOutputStream out = new ByteArrayOutputStream(); long size = MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), schema); assertEquals(size, out.toByteArray().length); ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray()); Schema deserialized = MessageSerializer.deserializeSchema(new ReadChannel(Channels.newChannel(in))); assertEquals(schema, deserialized); assertEquals(1, deserialized.getFields().size()); } @Test public void testSchemaDictionaryMessageSerialization() throws IOException { DictionaryEncoding dictionary = new DictionaryEncoding(9L, false, new ArrowType.Int(8, true)); Field field = new Field("test", new FieldType(true, ArrowType.Utf8.INSTANCE, dictionary, null), null); Schema schema = new Schema(Collections.singletonList(field)); ByteArrayOutputStream out = new ByteArrayOutputStream(); long size = MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), schema); assertEquals(size, out.toByteArray().length); ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray()); Schema deserialized = MessageSerializer.deserializeSchema(new ReadChannel(Channels.newChannel(in))); assertEquals(schema, deserialized); } @Test public void testSerializeRecordBatchV4() throws IOException { byte[] validity = new byte[] {(byte) 255, 0}; // second half is "undefined" byte[] values = new byte[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}; BufferAllocator alloc = new RootAllocator(Long.MAX_VALUE); ArrowBuf validityb = buf(alloc, validity); ArrowBuf valuesb = buf(alloc, values); ArrowRecordBatch batch = new ArrowRecordBatch(16, asList(new ArrowFieldNode(16, 8)), asList(validityb, valuesb)); // avoid writing legacy ipc format by default IpcOption option = new IpcOption(false, MetadataVersion.V4); ByteArrayOutputStream out = new ByteArrayOutputStream(); MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), batch, option); ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray()); ReadChannel channel = new ReadChannel(Channels.newChannel(in)); ArrowMessage deserialized = MessageSerializer.deserializeMessageBatch(channel, alloc); assertEquals(ArrowRecordBatch.class, deserialized.getClass()); verifyBatch((ArrowRecordBatch) deserialized, validity, values); } @Test public void testSerializeRecordBatchV5() throws Exception { byte[] validity = new byte[] {(byte) 255, 0}; // second half is "undefined" byte[] values = new byte[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}; BufferAllocator alloc = new RootAllocator(Long.MAX_VALUE); ArrowBuf validityb = buf(alloc, validity); ArrowBuf valuesb = buf(alloc, values); ArrowRecordBatch batch = new ArrowRecordBatch(16, asList(new ArrowFieldNode(16, 8)), asList(validityb, valuesb)); // avoid writing legacy ipc format by default IpcOption option = new IpcOption(false, MetadataVersion.V5); ByteArrayOutputStream out = new ByteArrayOutputStream(); MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), batch, option); validityb.close(); valuesb.close(); batch.close(); { ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray()); ReadChannel channel = new ReadChannel(Channels.newChannel(in)); ArrowMessage deserialized = MessageSerializer.deserializeMessageBatch(channel, alloc); assertEquals(ArrowRecordBatch.class, deserialized.getClass()); verifyBatch((ArrowRecordBatch) deserialized, validity, values); deserialized.close(); } { byte[] validBytes = out.toByteArray(); byte[] missingBytes = Arrays.copyOfRange(validBytes, /* from= */ 0, validBytes.length - 1); ByteArrayInputStream in = new ByteArrayInputStream(missingBytes); ReadChannel channel = new ReadChannel(Channels.newChannel(in)); assertThrows( IOException.class, () -> MessageSerializer.deserializeMessageBatch(channel, alloc)); } alloc.close(); } public static Schema testSchema() { return new Schema( asList( new Field( "testField", FieldType.nullable(new ArrowType.Int(8, true)), Collections.emptyList()))); } // Verifies batch contents matching test schema. public static void verifyBatch(ArrowRecordBatch batch, byte[] validity, byte[] values) { assertTrue(batch != null); List nodes = batch.getNodes(); assertEquals(1, nodes.size()); ArrowFieldNode node = nodes.get(0); assertEquals(16, node.getLength()); assertEquals(8, node.getNullCount()); List buffers = batch.getBuffers(); assertEquals(2, buffers.size()); assertArrayEquals(validity, MessageSerializerTest.array(buffers.get(0))); assertArrayEquals(values, MessageSerializerTest.array(buffers.get(1))); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/ipc/TestArrowFile.java000066400000000000000000000123361515451615200317100ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import static java.nio.channels.Channels.newChannel; import static org.apache.arrow.vector.TestUtils.newVarCharVector; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.File; import java.io.FileOutputStream; import java.io.IOException; import java.io.OutputStream; import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.List; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Collections2; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.types.pojo.Field; import org.junit.jupiter.api.Test; import org.slf4j.Logger; import org.slf4j.LoggerFactory; public class TestArrowFile extends BaseFileTest { private static final Logger LOGGER = LoggerFactory.getLogger(TestArrowFile.class); @Test public void testWrite() throws IOException { File file = new File("target/mytest_write.arrow"); int count = COUNT; try (BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); StructVector parent = StructVector.empty("parent", vectorAllocator)) { writeData(count, parent); write(parent.getChild("root"), file, new ByteArrayOutputStream()); } } @Test public void testWriteComplex() throws IOException { File file = new File("target/mytest_write_complex.arrow"); int count = COUNT; try (BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); StructVector parent = StructVector.empty("parent", vectorAllocator)) { writeComplexData(count, parent); FieldVector root = parent.getChild("root"); validateComplexContent(count, new VectorSchemaRoot(root)); write(root, file, new ByteArrayOutputStream()); } } /** * Writes the contents of parents to file. If outStream is non-null, also writes it to outStream * in the streaming serialized format. */ private void write(FieldVector parent, File file, OutputStream outStream) throws IOException { VectorSchemaRoot root = new VectorSchemaRoot(parent); try (FileOutputStream fileOutputStream = new FileOutputStream(file); ArrowFileWriter arrowWriter = new ArrowFileWriter(root, null, fileOutputStream.getChannel()); ) { LOGGER.debug("writing schema: " + root.getSchema()); arrowWriter.start(); arrowWriter.writeBatch(); arrowWriter.end(); } // Also try serializing to the stream writer. if (outStream != null) { try (ArrowStreamWriter arrowWriter = new ArrowStreamWriter(root, null, outStream)) { arrowWriter.start(); arrowWriter.writeBatch(); arrowWriter.end(); } } } @Test public void testFileStreamHasEos() throws IOException { try (VarCharVector vector1 = newVarCharVector("varchar1", allocator)) { vector1.allocateNewSafe(); vector1.set(0, "foo".getBytes(StandardCharsets.UTF_8)); vector1.set(1, "bar".getBytes(StandardCharsets.UTF_8)); vector1.set(3, "baz".getBytes(StandardCharsets.UTF_8)); vector1.set(4, "bar".getBytes(StandardCharsets.UTF_8)); vector1.set(5, "baz".getBytes(StandardCharsets.UTF_8)); vector1.setValueCount(6); List fields = Arrays.asList(vector1.getField()); List vectors = Collections2.asImmutableList(vector1); VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, vector1.getValueCount()); // write data ByteArrayOutputStream out = new ByteArrayOutputStream(); ArrowFileWriter writer = new ArrowFileWriter(root, null, newChannel(out)); writer.start(); writer.writeBatch(); writer.end(); byte[] bytes = out.toByteArray(); byte[] bytesWithoutMagic = new byte[bytes.length - 8]; System.arraycopy(bytes, 8, bytesWithoutMagic, 0, bytesWithoutMagic.length); try (ArrowStreamReader reader = new ArrowStreamReader(new ByteArrayInputStream(bytesWithoutMagic), allocator)) { assertTrue(reader.loadNextBatch()); // here will throw exception if read footer instead of eos. assertFalse(reader.loadNextBatch()); } } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/ipc/TestArrowFooter.java000066400000000000000000000050101515451615200322560ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import static java.util.Arrays.asList; import static org.junit.jupiter.api.Assertions.assertEquals; import com.google.flatbuffers.FlatBufferBuilder; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Collections; import java.util.List; import org.apache.arrow.flatbuf.Footer; import org.apache.arrow.vector.ipc.message.ArrowBlock; import org.apache.arrow.vector.ipc.message.ArrowFooter; 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.junit.jupiter.api.Test; public class TestArrowFooter { @Test public void test() { Schema schema = new Schema( asList( new Field( "a", FieldType.nullable(new ArrowType.Int(8, true)), Collections.emptyList()))); ArrowFooter footer = new ArrowFooter( schema, Collections.emptyList(), Collections.emptyList()); ArrowFooter newFooter = roundTrip(footer); assertEquals(footer, newFooter); List ids = new ArrayList<>(); ids.add(new ArrowBlock(0, 1, 2)); ids.add(new ArrowBlock(4, 5, 6)); footer = new ArrowFooter(schema, ids, ids); assertEquals(footer, roundTrip(footer)); } private ArrowFooter roundTrip(ArrowFooter footer) { FlatBufferBuilder builder = new FlatBufferBuilder(); int i = footer.writeTo(builder); builder.finish(i); ByteBuffer dataBuffer = builder.dataBuffer(); ArrowFooter newFooter = new ArrowFooter(Footer.getRootAsFooter(dataBuffer)); return newFooter; } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/ipc/TestArrowReaderWriter.java000066400000000000000000001243611515451615200334320ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import static java.nio.channels.Channels.newChannel; import static java.util.Arrays.asList; import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import static org.apache.arrow.vector.TestUtils.newVarCharVector; import static org.apache.arrow.vector.TestUtils.newVector; import static org.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector; 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.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.nio.ByteOrder; import java.nio.channels.Channels; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.function.BiFunction; import java.util.stream.Collectors; import org.apache.arrow.flatbuf.FieldNode; import org.apache.arrow.flatbuf.Message; import org.apache.arrow.flatbuf.RecordBatch; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Collections2; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.NullVector; import org.apache.arrow.vector.TestUtils; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; import org.apache.arrow.vector.compare.Range; import org.apache.arrow.vector.compare.RangeEqualsVisitor; import org.apache.arrow.vector.compare.TypeEqualsVisitor; import org.apache.arrow.vector.compare.VectorEqualsVisitor; import org.apache.arrow.vector.complex.StructVector; 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.ipc.message.ArrowBlock; import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.ipc.message.IpcOption; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.MetadataVersion; import org.apache.arrow.vector.types.Types.MinorType; 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.types.pojo.Schema; import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel; import org.apache.arrow.vector.util.DictionaryUtility; 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 TestArrowReaderWriter { private BufferAllocator allocator; private VarCharVector dictionaryVector1; private VarCharVector dictionaryVector2; private VarCharVector dictionaryVector3; private StructVector dictionaryVector4; private Dictionary dictionary1; private Dictionary dictionary2; private Dictionary dictionary3; private Dictionary dictionary4; private Schema schema; private Schema encodedSchema; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); dictionaryVector1 = newVarCharVector("D1", allocator); setVector( dictionaryVector1, "foo".getBytes(StandardCharsets.UTF_8), "bar".getBytes(StandardCharsets.UTF_8), "baz".getBytes(StandardCharsets.UTF_8)); dictionaryVector2 = newVarCharVector("D2", allocator); setVector( dictionaryVector2, "aa".getBytes(StandardCharsets.UTF_8), "bb".getBytes(StandardCharsets.UTF_8), "cc".getBytes(StandardCharsets.UTF_8)); dictionaryVector3 = newVarCharVector("D3", allocator); setVector( dictionaryVector3, "foo".getBytes(StandardCharsets.UTF_8), "bar".getBytes(StandardCharsets.UTF_8), "baz".getBytes(StandardCharsets.UTF_8), "aa".getBytes(StandardCharsets.UTF_8), "bb".getBytes(StandardCharsets.UTF_8), "cc".getBytes(StandardCharsets.UTF_8)); dictionaryVector4 = newVector(StructVector.class, "D4", MinorType.STRUCT, allocator); final Map> dictionaryValues4 = new HashMap<>(); dictionaryValues4.put("a", Arrays.asList(1, 2, 3)); dictionaryValues4.put("b", Arrays.asList(4, 5, 6)); setVector(dictionaryVector4, dictionaryValues4); dictionary1 = new Dictionary( dictionaryVector1, new DictionaryEncoding(/* id= */ 1L, /* ordered= */ false, /* indexType= */ null)); dictionary2 = new Dictionary( dictionaryVector2, new DictionaryEncoding(/* id= */ 2L, /* ordered= */ false, /* indexType= */ null)); dictionary3 = new Dictionary( dictionaryVector3, new DictionaryEncoding(/* id= */ 1L, /* ordered= */ false, /* indexType= */ null)); dictionary4 = new Dictionary( dictionaryVector4, new DictionaryEncoding(/* id= */ 3L, /* ordered= */ false, /* indexType= */ null)); } @AfterEach public void terminate() throws Exception { dictionaryVector1.close(); dictionaryVector2.close(); dictionaryVector3.close(); dictionaryVector4.close(); allocator.close(); } ArrowBuf buf(byte[] bytes) { ArrowBuf buffer = allocator.buffer(bytes.length); buffer.writeBytes(bytes); return buffer; } byte[] array(ArrowBuf buf) { byte[] bytes = new byte[checkedCastToInt(buf.readableBytes())]; buf.readBytes(bytes); return bytes; } @Test public void test() throws IOException { Schema schema = new Schema( asList( new Field( "testField", FieldType.nullable(new ArrowType.Int(8, true)), Collections.emptyList()))); ArrowType type = schema.getFields().get(0).getType(); FieldVector vector = TestUtils.newVector(FieldVector.class, "testField", type, allocator); vector.initializeChildrenFromFields(schema.getFields().get(0).getChildren()); byte[] validity = new byte[] {(byte) 255, 0}; // second half is "undefined" byte[] values = new byte[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}; ByteArrayOutputStream out = new ByteArrayOutputStream(); try (VectorSchemaRoot root = new VectorSchemaRoot(schema.getFields(), asList(vector), 16); ArrowFileWriter writer = new ArrowFileWriter(root, null, newChannel(out))) { ArrowBuf validityb = buf(validity); ArrowBuf valuesb = buf(values); ArrowRecordBatch batch = new ArrowRecordBatch(16, asList(new ArrowFieldNode(16, 8)), asList(validityb, valuesb)); VectorLoader loader = new VectorLoader(root); loader.load(batch); writer.writeBatch(); validityb.close(); valuesb.close(); batch.close(); } byte[] byteArray = out.toByteArray(); try (SeekableReadChannel channel = new SeekableReadChannel(new ByteArrayReadableSeekableByteChannel(byteArray)); ArrowFileReader reader = new ArrowFileReader(channel, allocator)) { Schema readSchema = reader.getVectorSchemaRoot().getSchema(); assertEquals(schema, readSchema); // TODO: dictionaries List recordBatches = reader.getRecordBlocks(); assertEquals(1, recordBatches.size()); reader.loadNextBatch(); VectorUnloader unloader = new VectorUnloader(reader.getVectorSchemaRoot()); ArrowRecordBatch recordBatch = unloader.getRecordBatch(); List nodes = recordBatch.getNodes(); assertEquals(1, nodes.size()); ArrowFieldNode node = nodes.get(0); assertEquals(16, node.getLength()); assertEquals(8, node.getNullCount()); List buffers = recordBatch.getBuffers(); assertEquals(2, buffers.size()); assertArrayEquals(validity, array(buffers.get(0))); assertArrayEquals(values, array(buffers.get(1))); // Read just the header. This demonstrates being able to read without need to // deserialize the buffer. ByteBuffer headerBuffer = ByteBuffer.allocate(recordBatches.get(0).getMetadataLength()); headerBuffer.put(byteArray, (int) recordBatches.get(0).getOffset(), headerBuffer.capacity()); // new format prefix_size ==8 headerBuffer.position(8); Message messageFB = Message.getRootAsMessage(headerBuffer); RecordBatch recordBatchFB = (RecordBatch) messageFB.header(new RecordBatch()); assertEquals(2, recordBatchFB.buffersLength()); assertEquals(1, recordBatchFB.nodesLength()); FieldNode nodeFB = recordBatchFB.nodes(0); assertEquals(16, nodeFB.length()); assertEquals(8, nodeFB.nullCount()); recordBatch.close(); } } @Test public void testWriteReadNullVector() throws IOException { int valueCount = 3; NullVector nullVector = new NullVector("vector"); nullVector.setValueCount(valueCount); Schema schema = new Schema(asList(nullVector.getField())); ByteArrayOutputStream out = new ByteArrayOutputStream(); try (VectorSchemaRoot root = new VectorSchemaRoot(schema.getFields(), asList(nullVector), valueCount); ArrowFileWriter writer = new ArrowFileWriter(root, null, newChannel(out))) { ArrowRecordBatch batch = new ArrowRecordBatch( valueCount, asList(new ArrowFieldNode(valueCount, 0)), Collections.emptyList()); VectorLoader loader = new VectorLoader(root); loader.load(batch); writer.writeBatch(); } byte[] byteArray = out.toByteArray(); try (SeekableReadChannel channel = new SeekableReadChannel(new ByteArrayReadableSeekableByteChannel(byteArray)); ArrowFileReader reader = new ArrowFileReader(channel, allocator)) { Schema readSchema = reader.getVectorSchemaRoot().getSchema(); assertEquals(schema, readSchema); List recordBatches = reader.getRecordBlocks(); assertEquals(1, recordBatches.size()); assertTrue(reader.loadNextBatch()); assertEquals(1, reader.getVectorSchemaRoot().getFieldVectors().size()); NullVector readNullVector = (NullVector) reader.getVectorSchemaRoot().getFieldVectors().get(0); assertEquals(valueCount, readNullVector.getValueCount()); } } @Test public void testWriteReadWithDictionaries() throws IOException { DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); provider.put(dictionary1); VarCharVector vector1 = newVarCharVector("varchar1", allocator); vector1.allocateNewSafe(); vector1.set(0, "foo".getBytes(StandardCharsets.UTF_8)); vector1.set(1, "bar".getBytes(StandardCharsets.UTF_8)); vector1.set(3, "baz".getBytes(StandardCharsets.UTF_8)); vector1.set(4, "bar".getBytes(StandardCharsets.UTF_8)); vector1.set(5, "baz".getBytes(StandardCharsets.UTF_8)); vector1.setValueCount(6); FieldVector encodedVector1 = (FieldVector) DictionaryEncoder.encode(vector1, dictionary1); vector1.close(); VarCharVector vector2 = newVarCharVector("varchar2", allocator); vector2.allocateNewSafe(); vector2.set(0, "bar".getBytes(StandardCharsets.UTF_8)); vector2.set(1, "baz".getBytes(StandardCharsets.UTF_8)); vector2.set(2, "foo".getBytes(StandardCharsets.UTF_8)); vector2.set(3, "foo".getBytes(StandardCharsets.UTF_8)); vector2.set(4, "foo".getBytes(StandardCharsets.UTF_8)); vector2.set(5, "bar".getBytes(StandardCharsets.UTF_8)); vector2.setValueCount(6); FieldVector encodedVector2 = (FieldVector) DictionaryEncoder.encode(vector2, dictionary1); vector2.close(); List fields = Arrays.asList(encodedVector1.getField(), encodedVector2.getField()); List vectors = Collections2.asImmutableList(encodedVector1, encodedVector2); try (VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, encodedVector1.getValueCount()); ByteArrayOutputStream out = new ByteArrayOutputStream(); ArrowFileWriter writer = new ArrowFileWriter(root, provider, newChannel(out)); ) { writer.start(); writer.writeBatch(); writer.end(); try (SeekableReadChannel channel = new SeekableReadChannel(new ByteArrayReadableSeekableByteChannel(out.toByteArray())); ArrowFileReader reader = new ArrowFileReader(channel, allocator)) { Schema readSchema = reader.getVectorSchemaRoot().getSchema(); assertEquals(root.getSchema(), readSchema); assertEquals(1, reader.getDictionaryBlocks().size()); assertEquals(1, reader.getRecordBlocks().size()); reader.loadNextBatch(); assertEquals(2, reader.getVectorSchemaRoot().getFieldVectors().size()); } } } @Test public void testWriteReadWithStructDictionaries() throws IOException { DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); provider.put(dictionary4); try (final StructVector vector = newVector(StructVector.class, "D4", MinorType.STRUCT, allocator)) { final Map> values = new HashMap<>(); // Index: 0, 2, 1, 2, 1, 0, 0 values.put("a", Arrays.asList(1, 3, 2, 3, 2, 1, 1)); values.put("b", Arrays.asList(4, 6, 5, 6, 5, 4, 4)); setVector(vector, values); FieldVector encodedVector = (FieldVector) DictionaryEncoder.encode(vector, dictionary4); List fields = Arrays.asList(encodedVector.getField()); List vectors = Collections2.asImmutableList(encodedVector); try (VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, encodedVector.getValueCount()); ByteArrayOutputStream out = new ByteArrayOutputStream(); ArrowFileWriter writer = new ArrowFileWriter(root, provider, newChannel(out)); ) { writer.start(); writer.writeBatch(); writer.end(); try (SeekableReadChannel channel = new SeekableReadChannel( new ByteArrayReadableSeekableByteChannel(out.toByteArray())); ArrowFileReader reader = new ArrowFileReader(channel, allocator)) { final VectorSchemaRoot readRoot = reader.getVectorSchemaRoot(); final Schema readSchema = readRoot.getSchema(); assertEquals(root.getSchema(), readSchema); assertEquals(1, reader.getDictionaryBlocks().size()); assertEquals(1, reader.getRecordBlocks().size()); reader.loadNextBatch(); assertEquals(1, readRoot.getFieldVectors().size()); assertEquals(1, reader.getDictionaryVectors().size()); // Read the encoded vector and check it final FieldVector readEncoded = readRoot.getVector(0); assertEquals(encodedVector.getValueCount(), readEncoded.getValueCount()); assertTrue( new RangeEqualsVisitor(encodedVector, readEncoded) .rangeEquals(new Range(0, 0, encodedVector.getValueCount()))); // Read the dictionary final Map readDictionaryMap = reader.getDictionaryVectors(); final Dictionary readDictionary = readDictionaryMap.get(readEncoded.getField().getDictionary().getId()); assertNotNull(readDictionary); // Assert the dictionary vector is correct final FieldVector readDictionaryVector = readDictionary.getVector(); assertEquals(dictionaryVector4.getValueCount(), readDictionaryVector.getValueCount()); final BiFunction typeComparatorIgnoreName = (v1, v2) -> new TypeEqualsVisitor(v1, false, true).equals(v2); assertTrue( new RangeEqualsVisitor( dictionaryVector4, readDictionaryVector, typeComparatorIgnoreName) .rangeEquals(new Range(0, 0, dictionaryVector4.getValueCount())), "Dictionary vectors are not equal"); // Assert the decoded vector is correct try (final ValueVector readVector = DictionaryEncoder.decode(readEncoded, readDictionary)) { assertEquals(vector.getValueCount(), readVector.getValueCount()); assertTrue( new RangeEqualsVisitor(vector, readVector, typeComparatorIgnoreName) .rangeEquals(new Range(0, 0, vector.getValueCount())), "Decoded vectors are not equal"); } } } } } @Test public void testEmptyStreamInFileIPC() throws IOException { DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); provider.put(dictionary1); VarCharVector vector = newVarCharVector("varchar", allocator); vector.allocateNewSafe(); vector.set(0, "foo".getBytes(StandardCharsets.UTF_8)); vector.set(1, "bar".getBytes(StandardCharsets.UTF_8)); vector.set(3, "baz".getBytes(StandardCharsets.UTF_8)); vector.set(4, "bar".getBytes(StandardCharsets.UTF_8)); vector.set(5, "baz".getBytes(StandardCharsets.UTF_8)); vector.setValueCount(6); FieldVector encodedVector1A = (FieldVector) DictionaryEncoder.encode(vector, dictionary1); vector.close(); List fields = Arrays.asList(encodedVector1A.getField()); List vectors = Collections2.asImmutableList(encodedVector1A); try (VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, encodedVector1A.getValueCount()); ByteArrayOutputStream out = new ByteArrayOutputStream(); ArrowFileWriter writer = new ArrowFileWriter(root, provider, newChannel(out))) { writer.start(); writer.end(); try (SeekableReadChannel channel = new SeekableReadChannel(new ByteArrayReadableSeekableByteChannel(out.toByteArray())); ArrowFileReader reader = new ArrowFileReader(channel, allocator)) { Schema readSchema = reader.getVectorSchemaRoot().getSchema(); assertEquals(root.getSchema(), readSchema); assertEquals(1, reader.getDictionaryVectors().size()); assertEquals(0, reader.getDictionaryBlocks().size()); assertEquals(0, reader.getRecordBlocks().size()); } } } @Test public void testEmptyStreamInStreamingIPC() throws IOException { DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); provider.put(dictionary1); VarCharVector vector = newVarCharVector("varchar", allocator); vector.allocateNewSafe(); vector.set(0, "foo".getBytes(StandardCharsets.UTF_8)); vector.set(1, "bar".getBytes(StandardCharsets.UTF_8)); vector.set(3, "baz".getBytes(StandardCharsets.UTF_8)); vector.set(4, "bar".getBytes(StandardCharsets.UTF_8)); vector.set(5, "baz".getBytes(StandardCharsets.UTF_8)); vector.setValueCount(6); FieldVector encodedVector = (FieldVector) DictionaryEncoder.encode(vector, dictionary1); vector.close(); List fields = Arrays.asList(encodedVector.getField()); try (VectorSchemaRoot root = new VectorSchemaRoot( fields, Arrays.asList(encodedVector), encodedVector.getValueCount()); ByteArrayOutputStream out = new ByteArrayOutputStream(); ArrowStreamWriter writer = new ArrowStreamWriter(root, provider, newChannel(out))) { writer.start(); writer.end(); try (ArrowStreamReader reader = new ArrowStreamReader( new ByteArrayReadableSeekableByteChannel(out.toByteArray()), allocator)) { Schema readSchema = reader.getVectorSchemaRoot().getSchema(); assertEquals(root.getSchema(), readSchema); assertEquals(1, reader.getDictionaryVectors().size()); assertFalse(reader.loadNextBatch()); } } } @Test public void testDictionaryReplacement() throws Exception { VarCharVector vector1 = newVarCharVector("varchar1", allocator); setVector( vector1, "foo".getBytes(StandardCharsets.UTF_8), "bar".getBytes(StandardCharsets.UTF_8), "baz".getBytes(StandardCharsets.UTF_8), "bar".getBytes(StandardCharsets.UTF_8)); FieldVector encodedVector1 = (FieldVector) DictionaryEncoder.encode(vector1, dictionary1); VarCharVector vector2 = newVarCharVector("varchar2", allocator); setVector( vector2, "foo".getBytes(StandardCharsets.UTF_8), "foo".getBytes(StandardCharsets.UTF_8), "foo".getBytes(StandardCharsets.UTF_8), "foo".getBytes(StandardCharsets.UTF_8)); FieldVector encodedVector2 = (FieldVector) DictionaryEncoder.encode(vector2, dictionary1); DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); provider.put(dictionary1); List schemaFields = new ArrayList<>(); schemaFields.add( DictionaryUtility.toMessageFormat(encodedVector1.getField(), provider, new HashSet<>())); schemaFields.add( DictionaryUtility.toMessageFormat(encodedVector2.getField(), provider, new HashSet<>())); Schema schema = new Schema(schemaFields); ByteArrayOutputStream outStream = new ByteArrayOutputStream(); WriteChannel out = new WriteChannel(newChannel(outStream)); // write schema MessageSerializer.serialize(out, schema); List closeableList = new ArrayList<>(); // write non-delta dictionary with id=1 serializeDictionaryBatch(out, dictionary3, false, closeableList); // write non-delta dictionary with id=1 serializeDictionaryBatch(out, dictionary1, false, closeableList); // write recordBatch2 serializeRecordBatch(out, Arrays.asList(encodedVector1, encodedVector2), closeableList); // write eos out.writeIntLittleEndian(0); try (ArrowStreamReader reader = new ArrowStreamReader( new ByteArrayReadableSeekableByteChannel(outStream.toByteArray()), allocator)) { assertEquals(1, reader.getDictionaryVectors().size()); assertTrue(reader.loadNextBatch()); FieldVector dictionaryVector = reader.getDictionaryVectors().get(1L).getVector(); // make sure the delta dictionary is concatenated. assertTrue(VectorEqualsVisitor.vectorEquals(dictionaryVector, dictionaryVector1, null)); assertFalse(reader.loadNextBatch()); } vector1.close(); vector2.close(); AutoCloseables.close(closeableList); } @Test public void testDeltaDictionary() throws Exception { VarCharVector vector1 = newVarCharVector("varchar1", allocator); setVector( vector1, "foo".getBytes(StandardCharsets.UTF_8), "bar".getBytes(StandardCharsets.UTF_8), "baz".getBytes(StandardCharsets.UTF_8), "bar".getBytes(StandardCharsets.UTF_8)); FieldVector encodedVector1 = (FieldVector) DictionaryEncoder.encode(vector1, dictionary1); VarCharVector vector2 = newVarCharVector("varchar2", allocator); setVector( vector2, "foo".getBytes(StandardCharsets.UTF_8), "aa".getBytes(StandardCharsets.UTF_8), "bb".getBytes(StandardCharsets.UTF_8), "cc".getBytes(StandardCharsets.UTF_8)); FieldVector encodedVector2 = (FieldVector) DictionaryEncoder.encode(vector2, dictionary3); DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); provider.put(dictionary1); provider.put(dictionary3); List schemaFields = new ArrayList<>(); schemaFields.add( DictionaryUtility.toMessageFormat(encodedVector1.getField(), provider, new HashSet<>())); schemaFields.add( DictionaryUtility.toMessageFormat(encodedVector2.getField(), provider, new HashSet<>())); Schema schema = new Schema(schemaFields); ByteArrayOutputStream outStream = new ByteArrayOutputStream(); WriteChannel out = new WriteChannel(newChannel(outStream)); // write schema MessageSerializer.serialize(out, schema); List closeableList = new ArrayList<>(); // write non-delta dictionary with id=1 serializeDictionaryBatch(out, dictionary1, false, closeableList); // write delta dictionary with id=1 Dictionary deltaDictionary = new Dictionary(dictionaryVector2, new DictionaryEncoding(1L, false, null)); serializeDictionaryBatch(out, deltaDictionary, true, closeableList); deltaDictionary.getVector().close(); // write recordBatch2 serializeRecordBatch(out, Arrays.asList(encodedVector1, encodedVector2), closeableList); // write eos out.writeIntLittleEndian(0); try (ArrowStreamReader reader = new ArrowStreamReader( new ByteArrayReadableSeekableByteChannel(outStream.toByteArray()), allocator)) { assertEquals(1, reader.getDictionaryVectors().size()); assertTrue(reader.loadNextBatch()); FieldVector dictionaryVector = reader.getDictionaryVectors().get(1L).getVector(); // make sure the delta dictionary is concatenated. assertTrue(VectorEqualsVisitor.vectorEquals(dictionaryVector, dictionaryVector3, null)); assertFalse(reader.loadNextBatch()); } vector1.close(); vector2.close(); AutoCloseables.close(closeableList); } // Tests that the ArrowStreamWriter re-emits dictionaries when they change @Test public void testWriteReadStreamWithDictionaryReplacement() throws Exception { DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); provider.put(dictionary1); String[] batch0 = {"foo", "bar", "baz", "bar", "baz"}; String[] batch1 = {"foo", "aa", "bar", "bb", "baz", "cc"}; VarCharVector vector = newVarCharVector("varchar", allocator); vector.allocateNewSafe(); for (int i = 0; i < batch0.length; ++i) { vector.set(i, batch0[i].getBytes(StandardCharsets.UTF_8)); } vector.setValueCount(batch0.length); FieldVector encodedVector1 = (FieldVector) DictionaryEncoder.encode(vector, dictionary1); List fields = Arrays.asList(encodedVector1.getField()); try (ByteArrayOutputStream out = new ByteArrayOutputStream()) { try (VectorSchemaRoot root = new VectorSchemaRoot( fields, Arrays.asList(encodedVector1), encodedVector1.getValueCount()); ArrowStreamWriter writer = new ArrowStreamWriter(root, provider, newChannel(out))) { writer.start(); // Write batch with initial data and dictionary writer.writeBatch(); // Create data for the next batch, using an extended dictionary with the same id vector.reset(); for (int i = 0; i < batch1.length; ++i) { vector.set(i, batch1[i].getBytes(StandardCharsets.UTF_8)); } vector.setValueCount(batch1.length); // Re-encode and move encoded data into the vector schema root provider.put(dictionary3); FieldVector encodedVector2 = (FieldVector) DictionaryEncoder.encode(vector, dictionary3); TransferPair transferPair = encodedVector2.makeTransferPair(root.getVector(0)); transferPair.transfer(); // Write second batch root.setRowCount(batch1.length); writer.writeBatch(); writer.end(); } try (ArrowStreamReader reader = new ArrowStreamReader( new ByteArrayReadableSeekableByteChannel(out.toByteArray()), allocator)) { VectorSchemaRoot root = reader.getVectorSchemaRoot(); // Read and verify first batch assertTrue(reader.loadNextBatch()); assertEquals(batch0.length, root.getRowCount()); FieldVector readEncoded1 = root.getVector(0); long dictionaryId = readEncoded1.getField().getDictionary().getId(); try (VarCharVector decodedValues = (VarCharVector) DictionaryEncoder.decode(readEncoded1, reader.lookup(dictionaryId))) { for (int i = 0; i < batch0.length; ++i) { assertEquals(batch0[i], new String(decodedValues.get(i), StandardCharsets.UTF_8)); } } // Read and verify second batch assertTrue(reader.loadNextBatch()); assertEquals(batch1.length, root.getRowCount()); FieldVector readEncoded2 = root.getVector(0); dictionaryId = readEncoded2.getField().getDictionary().getId(); try (VarCharVector decodedValues = (VarCharVector) DictionaryEncoder.decode(readEncoded2, reader.lookup(dictionaryId))) { for (int i = 0; i < batch1.length; ++i) { assertEquals(batch1[i], new String(decodedValues.get(i), StandardCharsets.UTF_8)); } } assertFalse(reader.loadNextBatch()); } } vector.close(); } private void serializeDictionaryBatch( WriteChannel out, Dictionary dictionary, boolean isDelta, List closeables) throws IOException { FieldVector dictVector = dictionary.getVector(); VectorSchemaRoot root = new VectorSchemaRoot( Collections.singletonList(dictVector.getField()), Collections.singletonList(dictVector), dictVector.getValueCount()); ArrowDictionaryBatch batch = new ArrowDictionaryBatch( dictionary.getEncoding().getId(), new VectorUnloader(root).getRecordBatch(), isDelta); MessageSerializer.serialize(out, batch); closeables.add(batch); closeables.add(root); } private void serializeRecordBatch( WriteChannel out, List vectors, List closeables) throws IOException { List fields = vectors.stream().map(v -> v.getField()).collect(Collectors.toList()); VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, vectors.get(0).getValueCount()); VectorUnloader unloader = new VectorUnloader(root); ArrowRecordBatch batch = unloader.getRecordBatch(); MessageSerializer.serialize(out, batch); closeables.add(batch); closeables.add(root); } @Test public void testReadInterleavedData() throws IOException { List batches = createRecordBatches(); ByteArrayOutputStream outStream = new ByteArrayOutputStream(); WriteChannel out = new WriteChannel(newChannel(outStream)); // write schema MessageSerializer.serialize(out, schema); // write dictionary1 FieldVector dictVector1 = dictionary1.getVector(); VectorSchemaRoot dictRoot1 = new VectorSchemaRoot( Collections.singletonList(dictVector1.getField()), Collections.singletonList(dictVector1), dictVector1.getValueCount()); ArrowDictionaryBatch dictionaryBatch1 = new ArrowDictionaryBatch(1, new VectorUnloader(dictRoot1).getRecordBatch()); MessageSerializer.serialize(out, dictionaryBatch1); dictionaryBatch1.close(); dictRoot1.close(); // write recordBatch1 MessageSerializer.serialize(out, batches.get(0)); // write dictionary2 FieldVector dictVector2 = dictionary2.getVector(); VectorSchemaRoot dictRoot2 = new VectorSchemaRoot( Collections.singletonList(dictVector2.getField()), Collections.singletonList(dictVector2), dictVector2.getValueCount()); ArrowDictionaryBatch dictionaryBatch2 = new ArrowDictionaryBatch(2, new VectorUnloader(dictRoot2).getRecordBatch()); MessageSerializer.serialize(out, dictionaryBatch2); dictionaryBatch2.close(); dictRoot2.close(); // write recordBatch1 MessageSerializer.serialize(out, batches.get(1)); // write eos out.writeIntLittleEndian(0); try (ArrowStreamReader reader = new ArrowStreamReader( new ByteArrayReadableSeekableByteChannel(outStream.toByteArray()), allocator)) { Schema readSchema = reader.getVectorSchemaRoot().getSchema(); assertEquals(encodedSchema, readSchema); assertEquals(2, reader.getDictionaryVectors().size()); assertTrue(reader.loadNextBatch()); assertTrue(reader.loadNextBatch()); assertFalse(reader.loadNextBatch()); } batches.forEach(batch -> batch.close()); } private List createRecordBatches() { List batches = new ArrayList<>(); DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); provider.put(dictionary1); provider.put(dictionary2); VarCharVector vectorA1 = newVarCharVector("varcharA1", allocator); vectorA1.allocateNewSafe(); vectorA1.set(0, "foo".getBytes(StandardCharsets.UTF_8)); vectorA1.set(1, "bar".getBytes(StandardCharsets.UTF_8)); vectorA1.set(3, "baz".getBytes(StandardCharsets.UTF_8)); vectorA1.set(4, "bar".getBytes(StandardCharsets.UTF_8)); vectorA1.set(5, "baz".getBytes(StandardCharsets.UTF_8)); vectorA1.setValueCount(6); VarCharVector vectorA2 = newVarCharVector("varcharA2", allocator); vectorA2.setValueCount(6); FieldVector encodedVectorA1 = (FieldVector) DictionaryEncoder.encode(vectorA1, dictionary1); vectorA1.close(); FieldVector encodedVectorA2 = (FieldVector) DictionaryEncoder.encode(vectorA1, dictionary2); vectorA2.close(); List fields = Arrays.asList(encodedVectorA1.getField(), encodedVectorA2.getField()); List vectors = Collections2.asImmutableList(encodedVectorA1, encodedVectorA2); VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, encodedVectorA1.getValueCount()); VectorUnloader unloader = new VectorUnloader(root); batches.add(unloader.getRecordBatch()); root.close(); VarCharVector vectorB1 = newVarCharVector("varcharB1", allocator); vectorB1.setValueCount(6); VarCharVector vectorB2 = newVarCharVector("varcharB2", allocator); vectorB2.allocateNew(); vectorB2.setValueCount(6); vectorB2.set(0, "aa".getBytes(StandardCharsets.UTF_8)); vectorB2.set(1, "aa".getBytes(StandardCharsets.UTF_8)); vectorB2.set(3, "bb".getBytes(StandardCharsets.UTF_8)); vectorB2.set(4, "bb".getBytes(StandardCharsets.UTF_8)); vectorB2.set(5, "cc".getBytes(StandardCharsets.UTF_8)); vectorB2.setValueCount(6); FieldVector encodedVectorB1 = (FieldVector) DictionaryEncoder.encode(vectorB1, dictionary1); vectorB1.close(); FieldVector encodedVectorB2 = (FieldVector) DictionaryEncoder.encode(vectorB2, dictionary2); vectorB2.close(); List fieldsB = Arrays.asList(encodedVectorB1.getField(), encodedVectorB2.getField()); List vectorsB = Collections2.asImmutableList(encodedVectorB1, encodedVectorB2); VectorSchemaRoot rootB = new VectorSchemaRoot(fieldsB, vectorsB, 6); VectorUnloader unloaderB = new VectorUnloader(rootB); batches.add(unloaderB.getRecordBatch()); rootB.close(); List schemaFields = new ArrayList<>(); schemaFields.add( DictionaryUtility.toMessageFormat(encodedVectorA1.getField(), provider, new HashSet<>())); schemaFields.add( DictionaryUtility.toMessageFormat(encodedVectorA2.getField(), provider, new HashSet<>())); schema = new Schema(schemaFields); encodedSchema = new Schema(Arrays.asList(encodedVectorA1.getField(), encodedVectorA2.getField())); return batches; } @Test public void testLegacyIpcBackwardsCompatibility() throws Exception { Schema schema = new Schema(asList(Field.nullable("field", new ArrowType.Int(32, true)))); IntVector vector = new IntVector("vector", allocator); final int valueCount = 2; vector.setValueCount(valueCount); vector.setSafe(0, 1); vector.setSafe(1, 2); ArrowRecordBatch batch = new ArrowRecordBatch( valueCount, asList(new ArrowFieldNode(valueCount, 0)), asList(vector.getValidityBuffer(), vector.getDataBuffer())); ByteArrayOutputStream outStream = new ByteArrayOutputStream(); WriteChannel out = new WriteChannel(newChannel(outStream)); // write legacy ipc format IpcOption option = new IpcOption(true, MetadataVersion.DEFAULT); MessageSerializer.serialize(out, schema, option); MessageSerializer.serialize(out, batch); ReadChannel in = new ReadChannel(newChannel(new ByteArrayInputStream(outStream.toByteArray()))); Schema readSchema = MessageSerializer.deserializeSchema(in); assertEquals(schema, readSchema); ArrowRecordBatch readBatch = MessageSerializer.deserializeRecordBatch(in, allocator); assertEquals(batch.getLength(), readBatch.getLength()); assertEquals(batch.computeBodyLength(), readBatch.computeBodyLength()); readBatch.close(); // write ipc format with continuation option = IpcOption.DEFAULT; MessageSerializer.serialize(out, schema, option); MessageSerializer.serialize(out, batch); ReadChannel in2 = new ReadChannel(newChannel(new ByteArrayInputStream(outStream.toByteArray()))); Schema readSchema2 = MessageSerializer.deserializeSchema(in2); assertEquals(schema, readSchema2); ArrowRecordBatch readBatch2 = MessageSerializer.deserializeRecordBatch(in2, allocator); assertEquals(batch.getLength(), readBatch2.getLength()); assertEquals(batch.computeBodyLength(), readBatch2.computeBodyLength()); readBatch2.close(); batch.close(); vector.close(); } @Test public void testChannelReadFully() throws IOException { final ByteBuffer buf = ByteBuffer.allocate(4).order(ByteOrder.nativeOrder()); buf.putInt(200); buf.rewind(); try (ReadChannel channel = new ReadChannel(Channels.newChannel(new ByteArrayInputStream(buf.array()))); ArrowBuf arrBuf = allocator.buffer(8)) { arrBuf.setInt(0, 100); arrBuf.writerIndex(4); assertEquals(4, arrBuf.writerIndex()); long n = channel.readFully(arrBuf, 4); assertEquals(4, n); assertEquals(8, arrBuf.writerIndex()); assertEquals(100, arrBuf.getInt(0)); assertEquals(200, arrBuf.getInt(4)); } } @Test public void testChannelReadFullyEos() throws IOException { final ByteBuffer buf = ByteBuffer.allocate(4).order(ByteOrder.nativeOrder()); buf.putInt(10); buf.rewind(); try (ReadChannel channel = new ReadChannel(Channels.newChannel(new ByteArrayInputStream(buf.array()))); ArrowBuf arrBuf = allocator.buffer(8)) { int n = channel.readFully(arrBuf.nioBuffer(0, 8)); assertEquals(4, n); // the input has only 4 bytes, so the number of bytes read should be 4 assertEquals(4, channel.bytesRead()); // the first 4 bytes have been read successfully. assertEquals(10, arrBuf.getInt(0)); } } @Test public void testCustomMetaData() throws IOException { VarCharVector vector = newVarCharVector("varchar1", allocator); List fields = Arrays.asList(vector.getField()); List vectors = Collections2.asImmutableList(vector); Map metadata = new HashMap<>(); metadata.put("key1", "value1"); metadata.put("key2", "value2"); try (VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, vector.getValueCount()); ByteArrayOutputStream out = new ByteArrayOutputStream(); ArrowFileWriter writer = new ArrowFileWriter(root, null, newChannel(out), metadata); ) { writer.start(); writer.end(); try (SeekableReadChannel channel = new SeekableReadChannel(new ByteArrayReadableSeekableByteChannel(out.toByteArray())); ArrowFileReader reader = new ArrowFileReader(channel, allocator)) { reader.getVectorSchemaRoot(); Map readMeta = reader.getMetaData(); assertEquals(2, readMeta.size()); assertEquals("value1", readMeta.get("key1")); assertEquals("value2", readMeta.get("key2")); } } } /** * This test case covers the case for which the footer size is extremely large (much larger than * the file size). Due to integer overflow, our implementation fails detect the problem, which * leads to extremely large memory allocation and eventually causing an OutOfMemoryError. */ @Test public void testFileFooterSizeOverflow() { // copy of org.apache.arrow.vector.ipc.ArrowMagic#MAGIC final byte[] magicBytes = "ARROW1".getBytes(StandardCharsets.UTF_8); // prepare input data byte[] data = new byte[30]; System.arraycopy(magicBytes, 0, data, 0, ArrowMagic.MAGIC_LENGTH); int footerLength = Integer.MAX_VALUE; byte[] footerLengthBytes = ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN).putInt(footerLength).array(); int footerOffset = data.length - ArrowMagic.MAGIC_LENGTH - 4; System.arraycopy(footerLengthBytes, 0, data, footerOffset, 4); System.arraycopy(magicBytes, 0, data, footerOffset + 4, ArrowMagic.MAGIC_LENGTH); // test file reader InvalidArrowFileException e = assertThrows( InvalidArrowFileException.class, () -> { try (SeekableReadChannel channel = new SeekableReadChannel(new ByteArrayReadableSeekableByteChannel(data)); ArrowFileReader reader = new ArrowFileReader(channel, allocator)) { reader.getVectorSchemaRoot().getSchema(); } }); assertEquals("invalid footer length: " + footerLength, e.getMessage()); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/ipc/TestArrowStream.java000066400000000000000000000132741515451615200322660ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; 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.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.channels.Channels; import java.util.Collections; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.TinyIntVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.api.Test; public class TestArrowStream extends BaseFileTest { @Test public void testEmptyStream() throws IOException { Schema schema = MessageSerializerTest.testSchema(); VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); // Write the stream. ByteArrayOutputStream out = new ByteArrayOutputStream(); ArrowStreamWriter writer = new ArrowStreamWriter(root, null, out); writer.close(); assertTrue(out.size() > 0); ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray()); try (ArrowStreamReader reader = new ArrowStreamReader(in, allocator)) { assertEquals(schema, reader.getVectorSchemaRoot().getSchema()); // Empty should return false assertFalse(reader.loadNextBatch()); assertEquals(0, reader.getVectorSchemaRoot().getRowCount()); assertFalse(reader.loadNextBatch()); assertEquals(0, reader.getVectorSchemaRoot().getRowCount()); } } @Test public void testStreamZeroLengthBatch() throws IOException { ByteArrayOutputStream os = new ByteArrayOutputStream(); try (IntVector vector = new IntVector("foo", allocator); ) { Schema schema = new Schema(Collections.singletonList(vector.getField())); try (VectorSchemaRoot root = new VectorSchemaRoot( schema, Collections.singletonList(vector), vector.getValueCount()); ArrowStreamWriter writer = new ArrowStreamWriter(root, null, Channels.newChannel(os)); ) { vector.setValueCount(0); root.setRowCount(0); writer.writeBatch(); writer.end(); } } ByteArrayInputStream in = new ByteArrayInputStream(os.toByteArray()); try (ArrowStreamReader reader = new ArrowStreamReader(in, allocator); ) { VectorSchemaRoot root = reader.getVectorSchemaRoot(); IntVector vector = (IntVector) root.getFieldVectors().get(0); reader.loadNextBatch(); assertEquals(0, vector.getValueCount()); assertEquals(0, root.getRowCount()); } } @Test public void testReadWrite() throws IOException { Schema schema = MessageSerializerTest.testSchema(); try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { int numBatches = 1; root.getFieldVectors().get(0).allocateNew(); TinyIntVector vector = (TinyIntVector) root.getFieldVectors().get(0); for (int i = 0; i < 16; i++) { vector.set(i, i < 8 ? 1 : 0, (byte) (i + 1)); } vector.setValueCount(16); root.setRowCount(16); ByteArrayOutputStream out = new ByteArrayOutputStream(); long bytesWritten = 0; try (ArrowStreamWriter writer = new ArrowStreamWriter(root, null, out)) { writer.start(); for (int i = 0; i < numBatches; i++) { writer.writeBatch(); } writer.end(); bytesWritten = writer.bytesWritten(); } ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray()); try (ArrowStreamReader reader = new ArrowStreamReader(in, allocator)) { Schema readSchema = reader.getVectorSchemaRoot().getSchema(); assertEquals(schema, readSchema); for (int i = 0; i < numBatches; i++) { assertTrue(reader.loadNextBatch()); } // TODO figure out why reader isn't getting padding bytes assertEquals(bytesWritten, reader.bytesRead() + 8); assertFalse(reader.loadNextBatch()); assertEquals(0, reader.getVectorSchemaRoot().getRowCount()); } } } @Test public void testReadWriteMultipleBatches() throws IOException { ByteArrayOutputStream os = new ByteArrayOutputStream(); try (IntVector vector = new IntVector("foo", allocator); ) { Schema schema = new Schema(Collections.singletonList(vector.getField())); try (VectorSchemaRoot root = new VectorSchemaRoot( schema, Collections.singletonList(vector), vector.getValueCount()); ArrowStreamWriter writer = new ArrowStreamWriter(root, null, Channels.newChannel(os)); ) { writeBatchData(writer, vector, root); } } ByteArrayInputStream in = new ByteArrayInputStream(os.toByteArray()); try (ArrowStreamReader reader = new ArrowStreamReader(in, allocator); ) { IntVector vector = (IntVector) reader.getVectorSchemaRoot().getFieldVectors().get(0); validateBatchData(reader, vector); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/ipc/TestArrowStreamPipe.java000066400000000000000000000122111515451615200330720ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; import java.io.IOException; import java.nio.channels.Pipe; import java.nio.channels.ReadableByteChannel; import java.nio.channels.WritableByteChannel; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.TinyIntVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.api.Test; public class TestArrowStreamPipe { Schema schema = MessageSerializerTest.testSchema(); BufferAllocator alloc = new RootAllocator(Long.MAX_VALUE); private final class WriterThread extends Thread { private final int numBatches; private final ArrowStreamWriter writer; private final VectorSchemaRoot root; public WriterThread(int numBatches, WritableByteChannel sinkChannel) throws IOException { this.numBatches = numBatches; BufferAllocator allocator = alloc.newChildAllocator("writer thread", 0, Integer.MAX_VALUE); root = VectorSchemaRoot.create(schema, allocator); writer = new ArrowStreamWriter(root, null, sinkChannel); } @Override public void run() { try { writer.start(); for (int j = 0; j < numBatches; j++) { root.getFieldVectors().get(0).allocateNew(); TinyIntVector vector = (TinyIntVector) root.getFieldVectors().get(0); // Send a changing batch id first vector.set(0, j); for (int i = 1; i < 16; i++) { vector.set(i, i < 8 ? 1 : 0, (byte) (i + 1)); } vector.setValueCount(16); root.setRowCount(16); writer.writeBatch(); } writer.close(); root.close(); } catch (IOException e) { e.printStackTrace(); fail(e.toString()); // have to explicitly fail since we're in a separate thread } } public long bytesWritten() { return writer.bytesWritten(); } } private final class ReaderThread extends Thread { private int batchesRead = 0; private final ArrowStreamReader reader; private final BufferAllocator alloc = new RootAllocator(Long.MAX_VALUE); private boolean done = false; public ReaderThread(ReadableByteChannel sourceChannel) throws IOException { reader = new ArrowStreamReader(sourceChannel, alloc) { @Override public boolean loadNextBatch() throws IOException { if (super.loadNextBatch()) { batchesRead++; } else { done = true; return false; } VectorSchemaRoot root = getVectorSchemaRoot(); assertEquals(16, root.getRowCount()); TinyIntVector vector = (TinyIntVector) root.getFieldVectors().get(0); assertEquals((byte) (batchesRead - 1), vector.get(0)); for (int i = 1; i < 16; i++) { if (i < 8) { assertEquals((byte) (i + 1), vector.get(i)); } else { assertTrue(vector.isNull(i)); } } return true; } }; } @Override public void run() { try { assertEquals(schema, reader.getVectorSchemaRoot().getSchema()); while (!done) { assertTrue(reader.loadNextBatch() != done); } reader.close(); } catch (IOException e) { e.printStackTrace(); fail(e.toString()); // have to explicitly fail since we're in a separate thread } } public int getBatchesRead() { return batchesRead; } public long bytesRead() { return reader.bytesRead(); } } // Starts up a producer and consumer thread to read/write batches. @Test public void pipeTest() throws IOException, InterruptedException { final int NUM_BATCHES = 10; Pipe pipe = Pipe.open(); WriterThread writer = new WriterThread(NUM_BATCHES, pipe.sink()); ReaderThread reader = new ReaderThread(pipe.source()); writer.start(); reader.start(); reader.join(); writer.join(); assertEquals(NUM_BATCHES, reader.getBatchesRead()); assertEquals(writer.bytesWritten(), reader.bytesRead()); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java000066400000000000000000000466611515451615200313770ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import java.io.File; import java.io.IOException; import java.net.URL; import java.nio.file.Files; import java.nio.file.Path; import java.util.Arrays; import java.util.Collections; import java.util.List; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.impl.ComplexWriterImpl; import org.apache.arrow.vector.complex.writer.BaseWriter; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider; import org.apache.arrow.vector.types.UnionMode; 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.Validator; import org.junit.jupiter.api.Test; import org.slf4j.Logger; import org.slf4j.LoggerFactory; public class TestJSONFile extends BaseFileTest { private static final Logger LOGGER = LoggerFactory.getLogger(TestJSONFile.class); @Test public void testNoBatches() throws IOException { File file = new File("target/no_batches.json"); try (BufferAllocator originalVectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); StructVector parent = StructVector.empty("parent", originalVectorAllocator)) { BaseWriter.ComplexWriter writer = new ComplexWriterImpl("root", parent); BaseWriter.StructWriter rootWriter = writer.rootAsStruct(); rootWriter.integer("int"); rootWriter.uInt1("uint1"); rootWriter.bigInt("bigInt"); rootWriter.float4("float"); JsonFileWriter jsonWriter = new JsonFileWriter(file, JsonFileWriter.config().pretty(true)); jsonWriter.start(new VectorSchemaRoot(parent.getChild("root")).getSchema(), null); jsonWriter.close(); } // read try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE); JsonFileReader reader = new JsonFileReader(file, readerAllocator)) { Schema schema = reader.start(); LOGGER.debug("reading schema: " + schema); } } @Test public void testWriteRead() throws IOException { File file = new File("target/mytest.json"); int count = COUNT; // write try (BufferAllocator originalVectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); StructVector parent = StructVector.empty("parent", originalVectorAllocator)) { writeData(count, parent); writeJSON(file, new VectorSchemaRoot(parent.getChild("root")), null); } // read try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE); JsonFileReader reader = new JsonFileReader(file, readerAllocator)) { Schema schema = reader.start(); LOGGER.debug("reading schema: " + schema); // initialize vectors try (VectorSchemaRoot root = reader.read(); ) { validateContent(count, root); } } } @Test public void testWriteReadComplexJSON() throws IOException { File file = new File("target/mytest_complex.json"); int count = COUNT; // write try (BufferAllocator originalVectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); StructVector parent = StructVector.empty("parent", originalVectorAllocator)) { writeComplexData(count, parent); writeJSON(file, new VectorSchemaRoot(parent.getChild("root")), null); } // read try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE); JsonFileReader reader = new JsonFileReader(file, readerAllocator); ) { Schema schema = reader.start(); LOGGER.debug("reading schema: " + schema); // initialize vectors try (VectorSchemaRoot root = reader.read(); ) { validateComplexContent(count, root); } } } @Test public void testWriteComplexJSON() throws IOException { File file = new File("target/mytest_write_complex.json"); int count = COUNT; try (BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); StructVector parent = StructVector.empty("parent", vectorAllocator)) { writeComplexData(count, parent); VectorSchemaRoot root = new VectorSchemaRoot(parent.getChild("root")); validateComplexContent(root.getRowCount(), root); writeJSON(file, root, null); } } public void writeJSON(File file, VectorSchemaRoot root, DictionaryProvider provider) throws IOException { JsonFileWriter writer = new JsonFileWriter(file, JsonFileWriter.config().pretty(true)); writer.start(root.getSchema(), provider); writer.write(root); writer.close(); } @Test public void testWriteReadUnionJSON() throws IOException { File file = new File("target/mytest_write_union.json"); int count = COUNT; try (BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); StructVector parent = StructVector.empty("parent", vectorAllocator)) { writeUnionData(count, parent); printVectors(parent.getChildrenFromFields()); try (VectorSchemaRoot root = new VectorSchemaRoot(parent.getChild("root"))) { validateUnionData(count, root); writeJSON(file, root, null); // read try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE)) { JsonFileReader reader = new JsonFileReader(file, readerAllocator); Schema schema = reader.start(); LOGGER.debug("reading schema: " + schema); try (VectorSchemaRoot rootFromJson = reader.read(); ) { validateUnionData(count, rootFromJson); Validator.compareVectorSchemaRoot(root, rootFromJson); } } } } } @Test public void testWriteReadDateTimeJSON() throws IOException { File file = new File("target/mytest_datetime.json"); int count = COUNT; // write try (BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); StructVector parent = StructVector.empty("parent", vectorAllocator)) { writeDateTimeData(count, parent); printVectors(parent.getChildrenFromFields()); VectorSchemaRoot root = new VectorSchemaRoot(parent.getChild("root")); validateDateTimeContent(count, root); writeJSON(file, new VectorSchemaRoot(parent.getChild("root")), null); } // read try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE); JsonFileReader reader = new JsonFileReader(file, readerAllocator)) { Schema schema = reader.start(); LOGGER.debug("reading schema: " + schema); // initialize vectors try (VectorSchemaRoot root = reader.read(); ) { validateDateTimeContent(count, root); } } } @Test public void testWriteReadDictionaryJSON() throws IOException { File file = new File("target/mytest_dictionary.json"); // write try (BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE)) { MapDictionaryProvider provider = new MapDictionaryProvider(); try (VectorSchemaRoot root = writeFlatDictionaryData(vectorAllocator, provider)) { printVectors(root.getFieldVectors()); validateFlatDictionary(root, provider); writeJSON(file, root, provider); } // Need to close dictionary vectors for (long id : provider.getDictionaryIds()) { provider.lookup(id).getVector().close(); } } // read try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE); JsonFileReader reader = new JsonFileReader(file, readerAllocator)) { Schema schema = reader.start(); LOGGER.debug("reading schema: " + schema); // initialize vectors try (VectorSchemaRoot root = reader.read(); ) { validateFlatDictionary(root, reader); } } } @Test public void testWriteReadNestedDictionaryJSON() throws IOException { File file = new File("target/mytest_dict_nested.json"); // data being written: // [['foo', 'bar'], ['foo'], ['bar']] -> [[0, 1], [0], [1]] // write try (BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE)) { MapDictionaryProvider provider = new MapDictionaryProvider(); try (VectorSchemaRoot root = writeNestedDictionaryData(vectorAllocator, provider)) { printVectors(root.getFieldVectors()); validateNestedDictionary(root, provider); writeJSON(file, root, provider); } // Need to close dictionary vectors for (long id : provider.getDictionaryIds()) { provider.lookup(id).getVector().close(); } } // read try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE); JsonFileReader reader = new JsonFileReader(file, readerAllocator)) { Schema schema = reader.start(); LOGGER.debug("reading schema: " + schema); // initialize vectors try (VectorSchemaRoot root = reader.read(); ) { validateNestedDictionary(root, reader); } } } @Test public void testWriteReadDecimalJSON() throws IOException { File file = new File("target/mytest_decimal.json"); // write try (BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); VectorSchemaRoot root = writeDecimalData(vectorAllocator)) { printVectors(root.getFieldVectors()); validateDecimalData(root); writeJSON(file, root, null); } // read try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE); JsonFileReader reader = new JsonFileReader(file, readerAllocator)) { Schema schema = reader.start(); LOGGER.debug("reading schema: " + schema); // initialize vectors try (VectorSchemaRoot root = reader.read(); ) { validateDecimalData(root); } } } @Test public void testSetStructLength() throws IOException { URL resource = getClass().getResource("/integration_json_struct.json"); assertNotNull(resource); File file = new File(resource.getFile()); try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE); JsonFileReader reader = new JsonFileReader(file, readerAllocator)) { Schema schema = reader.start(); LOGGER.debug("reading schema: " + schema); // initialize vectors try (VectorSchemaRoot root = reader.read(); ) { FieldVector vector = root.getVector("struct_nullable"); assertEquals(7, vector.getValueCount()); } } } @Test public void testWriteReadVarBinJSON() throws IOException { File file = new File("target/mytest_varbin.json"); int count = COUNT; // write try (BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); StructVector parent = StructVector.empty("parent", vectorAllocator)) { writeVarBinaryData(count, parent); VectorSchemaRoot root = new VectorSchemaRoot(parent.getChild("root")); validateVarBinary(count, root); writeJSON(file, new VectorSchemaRoot(parent.getChild("root")), null); } // read try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE); JsonFileReader reader = new JsonFileReader(file, readerAllocator)) { Schema schema = reader.start(); LOGGER.debug("reading schema: " + schema); // initialize vectors try (VectorSchemaRoot root = reader.read(); ) { validateVarBinary(count, root); } } } @Test public void testWriteReadMapJSON() throws IOException { File file = new File("target/mytest_map.json"); // write try (BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); VectorSchemaRoot root = writeMapData(vectorAllocator)) { printVectors(root.getFieldVectors()); validateMapData(root); writeJSON(file, root, null); } // read try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE); JsonFileReader reader = new JsonFileReader(file, readerAllocator)) { Schema schema = reader.start(); LOGGER.debug("reading schema: " + schema); // initialize vectors try (VectorSchemaRoot root = reader.read(); ) { validateMapData(root); } } } @Test public void testWriteReadNullJSON() throws IOException { File file = new File("target/mytest_null.json"); int valueCount = 10; // write try (BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE); VectorSchemaRoot root = writeNullData(valueCount)) { printVectors(root.getFieldVectors()); validateNullData(root, valueCount); writeJSON(file, root, null); } // read try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE); JsonFileReader reader = new JsonFileReader(file, readerAllocator)) { Schema schema = reader.start(); LOGGER.debug("reading schema: " + schema); // initialize vectors try (VectorSchemaRoot root = reader.read(); ) { validateNullData(root, valueCount); } } } /** Regression test for ARROW-17107. */ @Test public void testRoundtripEmptyVector() throws Exception { final List fields = Arrays.asList( Field.nullable("utf8", ArrowType.Utf8.INSTANCE), Field.nullable("largeutf8", ArrowType.LargeUtf8.INSTANCE), Field.nullable("binary", ArrowType.Binary.INSTANCE), Field.nullable("largebinary", ArrowType.LargeBinary.INSTANCE), Field.nullable("fixedsizebinary", new ArrowType.FixedSizeBinary(2)), Field.nullable("decimal128", new ArrowType.Decimal(3, 2, 128)), Field.nullable("decimal128", new ArrowType.Decimal(3, 2, 256)), new Field( "list", FieldType.nullable(ArrowType.List.INSTANCE), Collections.singletonList(Field.nullable("items", new ArrowType.Int(32, true)))), new Field( "listview", FieldType.nullable(ArrowType.ListView.INSTANCE), Collections.singletonList(Field.nullable("items", new ArrowType.Int(32, true)))), new Field( "largelist", FieldType.nullable(ArrowType.LargeList.INSTANCE), Collections.singletonList(Field.nullable("items", new ArrowType.Int(32, true)))), new Field( "largelistview", FieldType.nullable(ArrowType.LargeListView.INSTANCE), Collections.singletonList(Field.nullable("items", new ArrowType.Int(32, true)))), new Field( "map", FieldType.nullable(new ArrowType.Map(/*keyssorted*/ false)), Collections.singletonList( new Field( "items", FieldType.notNullable(ArrowType.Struct.INSTANCE), Arrays.asList( Field.notNullable("keys", new ArrowType.Int(32, true)), Field.nullable("values", new ArrowType.Int(32, true)))))), new Field( "fixedsizelist", FieldType.nullable(new ArrowType.FixedSizeList(2)), Collections.singletonList(Field.nullable("items", new ArrowType.Int(32, true)))), new Field( "denseunion", FieldType.nullable(new ArrowType.Union(UnionMode.Dense, new int[] {0})), Collections.singletonList(Field.nullable("items", new ArrowType.Int(32, true)))), new Field( "sparseunion", FieldType.nullable(new ArrowType.Union(UnionMode.Sparse, new int[] {0})), Collections.singletonList(Field.nullable("items", new ArrowType.Int(32, true))))); for (final Field field : fields) { final Schema schema = new Schema(Collections.singletonList(field)); try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { Path outputPath = Files.createTempFile("arrow-", ".json"); File outputFile = outputPath.toFile(); outputFile.deleteOnExit(); // Try with no allocation try (final JsonFileWriter jsonWriter = new JsonFileWriter(outputFile, JsonFileWriter.config().pretty(true))) { jsonWriter.start(schema, null); jsonWriter.write(root); } catch (Exception e) { throw new RuntimeException("Test failed for empty vector of type " + field, e); } try (JsonFileReader reader = new JsonFileReader(outputFile, allocator)) { final Schema readSchema = reader.start(); assertEquals(schema, readSchema); try (final VectorSchemaRoot data = reader.read()) { assertNotNull(data); assertEquals(0, data.getRowCount()); } assertNull(reader.read()); } // Try with an explicit allocation root.allocateNew(); root.setRowCount(0); try (final JsonFileWriter jsonWriter = new JsonFileWriter(outputFile, JsonFileWriter.config().pretty(true))) { jsonWriter.start(schema, null); jsonWriter.write(root); } catch (Exception e) { throw new RuntimeException("Test failed for empty vector of type " + field, e); } try (JsonFileReader reader = new JsonFileReader(outputFile, allocator)) { final Schema readSchema = reader.start(); assertEquals(schema, readSchema); try (final VectorSchemaRoot data = reader.read()) { assertNotNull(data); assertEquals(0, data.getRowCount()); } assertNull(reader.read()); } } } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/ipc/TestRoundTrip.java000066400000000000000000000715221515451615200317460ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import static org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider; 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.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assumptions.assumeTrue; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.File; import java.io.FileInputStream; import java.io.FileOutputStream; import java.io.IOException; import java.nio.channels.Channels; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiConsumer; import java.util.stream.Stream; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.Collections2; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.FixedSizeBinaryVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.TinyIntVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.ipc.message.ArrowBlock; import org.apache.arrow.vector.ipc.message.ArrowBuffer; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.ipc.message.IpcOption; import org.apache.arrow.vector.ipc.message.MessageMetadataResult; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.MetadataVersion; 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.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.MethodSource; import org.slf4j.Logger; import org.slf4j.LoggerFactory; public class TestRoundTrip extends BaseFileTest { private static final Logger LOGGER = LoggerFactory.getLogger(TestRoundTrip.class); private static BufferAllocator allocator; static Stream getWriteOption() { final IpcOption legacy = new IpcOption(true, MetadataVersion.V4); final IpcOption version4 = new IpcOption(false, MetadataVersion.V4); return Stream.of( new Object[] {"V4Legacy", legacy}, new Object[] {"V4", version4}, new Object[] {"V5", IpcOption.DEFAULT}); } @BeforeAll public static void setUpClass() { allocator = new RootAllocator(Integer.MAX_VALUE); } @AfterAll public static void tearDownClass() { allocator.close(); } @ParameterizedTest(name = "options = {0}") @MethodSource("getWriteOption") public void testStruct(String name, IpcOption writeOption) throws Exception { try (final BufferAllocator originalVectorAllocator = allocator.newChildAllocator("original vectors", 0, allocator.getLimit()); final StructVector parent = StructVector.empty("parent", originalVectorAllocator)) { writeData(COUNT, parent); roundTrip( name, writeOption, new VectorSchemaRoot(parent.getChild("root")), /* dictionaryProvider */ null, TestRoundTrip::writeSingleBatch, validateFileBatches(new int[] {COUNT}, this::validateContent), validateStreamBatches(new int[] {COUNT}, this::validateContent)); } } @ParameterizedTest(name = "options = {0}") @MethodSource("getWriteOption") public void testComplex(String name, IpcOption writeOption) throws Exception { try (final BufferAllocator originalVectorAllocator = allocator.newChildAllocator("original vectors", 0, allocator.getLimit()); final StructVector parent = StructVector.empty("parent", originalVectorAllocator)) { writeComplexData(COUNT, parent); roundTrip( name, writeOption, new VectorSchemaRoot(parent.getChild("root")), /* dictionaryProvider */ null, TestRoundTrip::writeSingleBatch, validateFileBatches(new int[] {COUNT}, this::validateComplexContent), validateStreamBatches(new int[] {COUNT}, this::validateComplexContent)); } } @ParameterizedTest(name = "options = {0}") @MethodSource("getWriteOption") public void testMultipleRecordBatches(String name, IpcOption writeOption) throws Exception { int[] counts = {10, 5}; try (final BufferAllocator originalVectorAllocator = allocator.newChildAllocator("original vectors", 0, allocator.getLimit()); final StructVector parent = StructVector.empty("parent", originalVectorAllocator)) { writeData(counts[0], parent); roundTrip( name, writeOption, new VectorSchemaRoot(parent.getChild("root")), /* dictionaryProvider */ null, (root, writer) -> { writer.start(); parent.allocateNew(); writeData(counts[0], parent); root.setRowCount(counts[0]); writer.writeBatch(); parent.allocateNew(); // if we write the same data we don't catch that the metadata is stored in the wrong // order. writeData(counts[1], parent); root.setRowCount(counts[1]); writer.writeBatch(); writer.end(); }, validateFileBatches(counts, this::validateContent), validateStreamBatches(counts, this::validateContent)); } } @ParameterizedTest(name = "options = {0}") @MethodSource("getWriteOption") public void testUnionV4(String name, IpcOption writeOption) throws Exception { assumeTrue(writeOption.metadataVersion == MetadataVersion.V4); final File temp = File.createTempFile("arrow-test-" + name + "-", ".arrow"); temp.deleteOnExit(); final ByteArrayOutputStream memoryStream = new ByteArrayOutputStream(); try (final BufferAllocator originalVectorAllocator = allocator.newChildAllocator("original vectors", 0, allocator.getLimit()); final StructVector parent = StructVector.empty("parent", originalVectorAllocator)) { writeUnionData(COUNT, parent); final VectorSchemaRoot root = new VectorSchemaRoot(parent.getChild("root")); IllegalArgumentException e = assertThrows( IllegalArgumentException.class, () -> { try (final FileOutputStream fileStream = new FileOutputStream(temp)) { new ArrowFileWriter(root, null, fileStream.getChannel(), writeOption); new ArrowStreamWriter(root, null, Channels.newChannel(memoryStream), writeOption); } }); assertTrue(e.getMessage().contains("Cannot write union with V4 metadata"), e.getMessage()); e = assertThrows( IllegalArgumentException.class, () -> { new ArrowStreamWriter(root, null, Channels.newChannel(memoryStream), writeOption); }); assertTrue(e.getMessage().contains("Cannot write union with V4 metadata"), e.getMessage()); } } @ParameterizedTest(name = "options = {0}") @MethodSource("getWriteOption") public void testUnionV5(String name, IpcOption writeOption) throws Exception { assumeTrue(writeOption.metadataVersion == MetadataVersion.V5); try (final BufferAllocator originalVectorAllocator = allocator.newChildAllocator("original vectors", 0, allocator.getLimit()); final StructVector parent = StructVector.empty("parent", originalVectorAllocator)) { writeUnionData(COUNT, parent); VectorSchemaRoot root = new VectorSchemaRoot(parent.getChild("root")); validateUnionData(COUNT, root); roundTrip( name, writeOption, root, /* dictionaryProvider */ null, TestRoundTrip::writeSingleBatch, validateFileBatches(new int[] {COUNT}, this::validateUnionData), validateStreamBatches(new int[] {COUNT}, this::validateUnionData)); } } @ParameterizedTest(name = "options = {0}") @MethodSource("getWriteOption") public void testTiny(String name, IpcOption writeOption) throws Exception { try (final VectorSchemaRoot root = VectorSchemaRoot.create(MessageSerializerTest.testSchema(), allocator)) { root.getFieldVectors().get(0).allocateNew(); int count = 16; TinyIntVector vector = (TinyIntVector) root.getFieldVectors().get(0); for (int i = 0; i < count; i++) { vector.set(i, i < 8 ? 1 : 0, (byte) (i + 1)); } vector.setValueCount(count); root.setRowCount(count); roundTrip( name, writeOption, root, /* dictionaryProvider */ null, TestRoundTrip::writeSingleBatch, validateFileBatches(new int[] {count}, this::validateTinyData), validateStreamBatches(new int[] {count}, this::validateTinyData)); } } private void validateTinyData(int count, VectorSchemaRoot root) { assertEquals(count, root.getRowCount()); TinyIntVector vector = (TinyIntVector) root.getFieldVectors().get(0); for (int i = 0; i < count; i++) { if (i < 8) { assertEquals((byte) (i + 1), vector.get(i)); } else { assertTrue(vector.isNull(i)); } } } @ParameterizedTest(name = "options = {0}") @MethodSource("getWriteOption") public void testMetadata(String name, IpcOption writeOption) throws Exception { List childFields = new ArrayList<>(); childFields.add( new Field( "varchar-child", new FieldType(true, ArrowType.Utf8.INSTANCE, null, metadata(1)), null)); childFields.add( new Field( "float-child", new FieldType( true, new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE), null, metadata(2)), null)); childFields.add( new Field( "int-child", new FieldType(false, new ArrowType.Int(32, true), null, metadata(3)), null)); childFields.add( new Field( "list-child", new FieldType(true, ArrowType.List.INSTANCE, null, metadata(4)), Collections2.asImmutableList( new Field("l1", FieldType.nullable(new ArrowType.Int(16, true)), null)))); Field field = new Field( "meta", new FieldType(true, ArrowType.Struct.INSTANCE, null, metadata(0)), childFields); Map metadata = new HashMap<>(); metadata.put("s1", "v1"); metadata.put("s2", "v2"); Schema originalSchema = new Schema(Collections2.asImmutableList(field), metadata); assertEquals(metadata, originalSchema.getCustomMetadata()); try (final BufferAllocator originalVectorAllocator = allocator.newChildAllocator("original vectors", 0, allocator.getLimit()); final StructVector vector = (StructVector) field.createVector(originalVectorAllocator)) { vector.allocateNewSafe(); vector.setValueCount(0); List vectors = Collections2.asImmutableList(vector); VectorSchemaRoot root = new VectorSchemaRoot(originalSchema, vectors, 0); BiConsumer validate = (count, readRoot) -> { Schema schema = readRoot.getSchema(); assertEquals(originalSchema, schema); assertEquals(originalSchema.getCustomMetadata(), schema.getCustomMetadata()); Field top = schema.getFields().get(0); assertEquals(metadata(0), top.getMetadata()); for (int i = 0; i < 4; i++) { assertEquals(metadata(i + 1), top.getChildren().get(i).getMetadata()); } }; roundTrip( name, writeOption, root, /* dictionaryProvider */ null, TestRoundTrip::writeSingleBatch, validateFileBatches(new int[] {0}, validate), validateStreamBatches(new int[] {0}, validate)); } } private Map metadata(int i) { Map map = new HashMap<>(); map.put("k_" + i, "v_" + i); map.put("k2_" + i, "v2_" + i); return Collections.unmodifiableMap(map); } @ParameterizedTest(name = "options = {0}") @MethodSource("getWriteOption") public void testFlatDictionary(String name, IpcOption writeOption) throws Exception { AtomicInteger numDictionaryBlocksWritten = new AtomicInteger(); MapDictionaryProvider provider = new MapDictionaryProvider(); try (final BufferAllocator originalVectorAllocator = allocator.newChildAllocator("original vectors", 0, allocator.getLimit()); final VectorSchemaRoot root = writeFlatDictionaryData(originalVectorAllocator, provider)) { roundTrip( name, writeOption, root, provider, (ignored, writer) -> { writer.start(); writer.writeBatch(); writer.end(); if (writer instanceof ArrowFileWriter) { numDictionaryBlocksWritten.set( ((ArrowFileWriter) writer).getDictionaryBlocks().size()); } }, (fileReader) -> { VectorSchemaRoot readRoot = fileReader.getVectorSchemaRoot(); Schema schema = readRoot.getSchema(); LOGGER.debug("reading schema: " + schema); assertTrue(fileReader.loadNextBatch()); validateFlatDictionary(readRoot, fileReader); assertEquals(numDictionaryBlocksWritten.get(), fileReader.getDictionaryBlocks().size()); }, (streamReader) -> { VectorSchemaRoot readRoot = streamReader.getVectorSchemaRoot(); Schema schema = readRoot.getSchema(); LOGGER.debug("reading schema: " + schema); assertTrue(streamReader.loadNextBatch()); validateFlatDictionary(readRoot, streamReader); }); // Need to close dictionary vectors for (long id : provider.getDictionaryIds()) { provider.lookup(id).getVector().close(); } } } @ParameterizedTest(name = "options = {0}") @MethodSource("getWriteOption") public void testNestedDictionary(String name, IpcOption writeOption) throws Exception { AtomicInteger numDictionaryBlocksWritten = new AtomicInteger(); MapDictionaryProvider provider = new MapDictionaryProvider(); // data being written: // [['foo', 'bar'], ['foo'], ['bar']] -> [[0, 1], [0], [1]] try (final BufferAllocator originalVectorAllocator = allocator.newChildAllocator("original vectors", 0, allocator.getLimit()); final VectorSchemaRoot root = writeNestedDictionaryData(originalVectorAllocator, provider)) { CheckedConsumer validateDictionary = (streamReader) -> { VectorSchemaRoot readRoot = streamReader.getVectorSchemaRoot(); Schema schema = readRoot.getSchema(); LOGGER.debug("reading schema: " + schema); assertTrue(streamReader.loadNextBatch()); validateNestedDictionary(readRoot, streamReader); }; roundTrip( name, writeOption, root, provider, (ignored, writer) -> { writer.start(); writer.writeBatch(); writer.end(); if (writer instanceof ArrowFileWriter) { numDictionaryBlocksWritten.set( ((ArrowFileWriter) writer).getDictionaryBlocks().size()); } }, validateDictionary, validateDictionary); // Need to close dictionary vectors for (long id : provider.getDictionaryIds()) { provider.lookup(id).getVector().close(); } } } @ParameterizedTest(name = "options = {0}") @MethodSource("getWriteOption") public void testFixedSizeBinary(String name, IpcOption writeOption) throws Exception { final int count = 10; final int typeWidth = 11; byte[][] byteValues = new byte[count][typeWidth]; for (int i = 0; i < count; i++) { for (int j = 0; j < typeWidth; j++) { byteValues[i][j] = ((byte) i); } } BiConsumer validator = (expectedCount, root) -> { for (int i = 0; i < expectedCount; i++) { assertArrayEquals( byteValues[i], ((byte[]) root.getVector("fixed-binary").getObject(i))); } }; try (final BufferAllocator originalVectorAllocator = allocator.newChildAllocator("original vectors", 0, allocator.getLimit()); final StructVector parent = StructVector.empty("parent", originalVectorAllocator)) { FixedSizeBinaryVector fixedSizeBinaryVector = parent.addOrGet( "fixed-binary", FieldType.nullable(new ArrowType.FixedSizeBinary(typeWidth)), FixedSizeBinaryVector.class); parent.allocateNew(); for (int i = 0; i < count; i++) { fixedSizeBinaryVector.set(i, byteValues[i]); } parent.setValueCount(count); roundTrip( name, writeOption, new VectorSchemaRoot(parent), /* dictionaryProvider */ null, TestRoundTrip::writeSingleBatch, validateFileBatches(new int[] {count}, validator), validateStreamBatches(new int[] {count}, validator)); } } @ParameterizedTest(name = "options = {0}") @MethodSource("getWriteOption") public void testFixedSizeList(String name, IpcOption writeOption) throws Exception { BiConsumer validator = (expectedCount, root) -> { for (int i = 0; i < expectedCount; i++) { assertEquals( Collections2.asImmutableList(i + 0.1f, i + 10.1f), root.getVector("float-pairs").getObject(i)); assertEquals(i, root.getVector("ints").getObject(i)); } }; try (final BufferAllocator originalVectorAllocator = allocator.newChildAllocator("original vectors", 0, allocator.getLimit()); final StructVector parent = StructVector.empty("parent", originalVectorAllocator)) { FixedSizeListVector tuples = parent.addOrGet( "float-pairs", FieldType.nullable(new ArrowType.FixedSizeList(2)), FixedSizeListVector.class); Float4Vector floats = (Float4Vector) tuples .addOrGetVector(FieldType.nullable(Types.MinorType.FLOAT4.getType())) .getVector(); IntVector ints = parent.addOrGet("ints", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); parent.allocateNew(); for (int i = 0; i < COUNT; i++) { tuples.setNotNull(i); floats.set(i * 2, i + 0.1f); floats.set(i * 2 + 1, i + 10.1f); ints.set(i, i); } parent.setValueCount(COUNT); roundTrip( name, writeOption, new VectorSchemaRoot(parent), /* dictionaryProvider */ null, TestRoundTrip::writeSingleBatch, validateFileBatches(new int[] {COUNT}, validator), validateStreamBatches(new int[] {COUNT}, validator)); } } @ParameterizedTest(name = "options = {0}") @MethodSource("getWriteOption") public void testVarBinary(String name, IpcOption writeOption) throws Exception { try (final BufferAllocator originalVectorAllocator = allocator.newChildAllocator("original vectors", 0, allocator.getLimit()); final StructVector parent = StructVector.empty("parent", originalVectorAllocator)) { writeVarBinaryData(COUNT, parent); VectorSchemaRoot root = new VectorSchemaRoot(parent.getChild("root")); validateVarBinary(COUNT, root); roundTrip( name, writeOption, root, /* dictionaryProvider */ null, TestRoundTrip::writeSingleBatch, validateFileBatches(new int[] {COUNT}, this::validateVarBinary), validateStreamBatches(new int[] {COUNT}, this::validateVarBinary)); } } @ParameterizedTest(name = "options = {0}") @MethodSource("getWriteOption") public void testReadWriteMultipleBatches(String name, IpcOption writeOption) throws IOException { File file = new File("target/mytest_nulls_multibatch.arrow"); int numBlocksWritten = 0; try (IntVector vector = new IntVector("foo", allocator); ) { Schema schema = new Schema(Collections.singletonList(vector.getField())); try (FileOutputStream fileOutputStream = new FileOutputStream(file); VectorSchemaRoot root = new VectorSchemaRoot( schema, Collections.singletonList((FieldVector) vector), vector.getValueCount()); ArrowFileWriter writer = new ArrowFileWriter(root, null, fileOutputStream.getChannel(), writeOption)) { writeBatchData(writer, vector, root); numBlocksWritten = writer.getRecordBlocks().size(); } } try (FileInputStream fileInputStream = new FileInputStream(file); ArrowFileReader reader = new ArrowFileReader(fileInputStream.getChannel(), allocator); ) { IntVector vector = (IntVector) reader.getVectorSchemaRoot().getFieldVectors().get(0); validateBatchData(reader, vector); assertEquals(numBlocksWritten, reader.getRecordBlocks().size()); } } @ParameterizedTest(name = "options = {0}") @MethodSource("getWriteOption") public void testMap(String name, IpcOption writeOption) throws Exception { try (final BufferAllocator originalVectorAllocator = allocator.newChildAllocator("original vectors", 0, allocator.getLimit()); final VectorSchemaRoot root = writeMapData(originalVectorAllocator)) { roundTrip( name, writeOption, root, /* dictionaryProvider */ null, TestRoundTrip::writeSingleBatch, validateFileBatches( new int[] {root.getRowCount()}, (count, readRoot) -> validateMapData(readRoot)), validateStreamBatches( new int[] {root.getRowCount()}, (count, readRoot) -> validateMapData(readRoot))); } } @ParameterizedTest(name = "options = {0}") @MethodSource("getWriteOption") public void testListAsMap(String name, IpcOption writeOption) throws Exception { try (final BufferAllocator originalVectorAllocator = allocator.newChildAllocator("original vectors", 0, allocator.getLimit()); final VectorSchemaRoot root = writeListAsMapData(originalVectorAllocator)) { roundTrip( name, writeOption, root, /* dictionaryProvider */ null, TestRoundTrip::writeSingleBatch, validateFileBatches( new int[] {root.getRowCount()}, (count, readRoot) -> validateListAsMapData(readRoot)), validateStreamBatches( new int[] {root.getRowCount()}, (count, readRoot) -> validateListAsMapData(readRoot))); } } // Generic test helpers private static void writeSingleBatch(VectorSchemaRoot root, ArrowWriter writer) throws IOException { writer.start(); writer.writeBatch(); writer.end(); } private CheckedConsumer validateFileBatches( int[] counts, BiConsumer validator) { return (arrowReader) -> { VectorSchemaRoot root = arrowReader.getVectorSchemaRoot(); VectorUnloader unloader = new VectorUnloader(root); Schema schema = root.getSchema(); LOGGER.debug("reading schema: " + schema); int i = 0; List recordBatches = arrowReader.getRecordBlocks(); assertEquals(counts.length, recordBatches.size()); long previousOffset = 0; for (ArrowBlock rbBlock : recordBatches) { assertTrue( rbBlock.getOffset() > previousOffset, rbBlock.getOffset() + " > " + previousOffset); previousOffset = rbBlock.getOffset(); arrowReader.loadRecordBatch(rbBlock); assertEquals(counts[i], root.getRowCount(), "RB #" + i); validator.accept(counts[i], root); try (final ArrowRecordBatch batch = unloader.getRecordBatch()) { List buffersLayout = batch.getBuffersLayout(); for (ArrowBuffer arrowBuffer : buffersLayout) { assertEquals(0, arrowBuffer.getOffset() % 8); } } ++i; } }; } private CheckedConsumer validateStreamBatches( int[] counts, BiConsumer validator) { return (arrowReader) -> { VectorSchemaRoot root = arrowReader.getVectorSchemaRoot(); VectorUnloader unloader = new VectorUnloader(root); Schema schema = root.getSchema(); LOGGER.debug("reading schema: " + schema); int i = 0; for (int n = 0; n < counts.length; n++) { assertTrue(arrowReader.loadNextBatch()); assertEquals(counts[i], root.getRowCount(), "RB #" + i); validator.accept(counts[i], root); try (final ArrowRecordBatch batch = unloader.getRecordBatch()) { final List buffersLayout = batch.getBuffersLayout(); for (ArrowBuffer arrowBuffer : buffersLayout) { assertEquals(0, arrowBuffer.getOffset() % 8); } } ++i; } assertFalse(arrowReader.loadNextBatch()); }; } @FunctionalInterface interface CheckedConsumer { void accept(T t) throws Exception; } @FunctionalInterface interface CheckedBiConsumer { void accept(T t, U u) throws Exception; } private void roundTrip( String name, IpcOption writeOption, VectorSchemaRoot root, DictionaryProvider provider, CheckedBiConsumer writer, CheckedConsumer fileValidator, CheckedConsumer streamValidator) throws Exception { final File temp = File.createTempFile("arrow-test-" + name + "-", ".arrow"); temp.deleteOnExit(); final ByteArrayOutputStream memoryStream = new ByteArrayOutputStream(); final Map metadata = new HashMap<>(); metadata.put("foo", "bar"); try (final FileOutputStream fileStream = new FileOutputStream(temp); final ArrowFileWriter fileWriter = new ArrowFileWriter(root, provider, fileStream.getChannel(), metadata, writeOption); final ArrowStreamWriter streamWriter = new ArrowStreamWriter(root, provider, Channels.newChannel(memoryStream), writeOption)) { writer.accept(root, fileWriter); writer.accept(root, streamWriter); } MessageMetadataResult metadataResult = MessageSerializer.readMessage( new ReadChannel( Channels.newChannel(new ByteArrayInputStream(memoryStream.toByteArray())))); assertNotNull(metadataResult); assertEquals(writeOption.metadataVersion.toFlatbufID(), metadataResult.getMessage().version()); try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, allocator.getLimit()); FileInputStream fileInputStream = new FileInputStream(temp); ByteArrayInputStream inputStream = new ByteArrayInputStream(memoryStream.toByteArray()); ArrowFileReader fileReader = new ArrowFileReader(fileInputStream.getChannel(), readerAllocator); ArrowStreamReader streamReader = new ArrowStreamReader(inputStream, readerAllocator)) { fileValidator.accept(fileReader); streamValidator.accept(streamReader); assertEquals(writeOption.metadataVersion, fileReader.getFooter().getMetadataVersion()); assertEquals(metadata, fileReader.getMetaData()); } } } TestUIntDictionaryRoundTrip.java000066400000000000000000000242651515451615200345170ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/ipc/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc; import static org.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.channels.Channels; import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.Map; import java.util.function.ToIntBiFunction; import java.util.stream.Stream; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.FieldVector; 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.ValueVector; import org.apache.arrow.vector.VarCharVector; 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.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.types.pojo.Schema; import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; /** Test the round-trip of dictionary encoding, with unsigned integer as indices. */ public class TestUIntDictionaryRoundTrip { private BufferAllocator allocator; private DictionaryProvider.MapDictionaryProvider dictionaryProvider; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); dictionaryProvider = new DictionaryProvider.MapDictionaryProvider(); } @AfterEach public void terminate() throws Exception { allocator.close(); } private byte[] writeData(boolean streamMode, FieldVector encodedVector) throws IOException { ByteArrayOutputStream out = new ByteArrayOutputStream(); VectorSchemaRoot root = new VectorSchemaRoot( Arrays.asList(encodedVector.getField()), Arrays.asList(encodedVector), encodedVector.getValueCount()); try (ArrowWriter writer = streamMode ? new ArrowStreamWriter(root, dictionaryProvider, out) : new ArrowFileWriter(root, dictionaryProvider, Channels.newChannel(out))) { writer.start(); writer.writeBatch(); writer.end(); return out.toByteArray(); } } private void readData( boolean streamMode, byte[] data, Field expectedField, ToIntBiFunction valGetter, long dictionaryID, int[] expectedIndices, String[] expectedDictItems) throws IOException { try (ArrowReader reader = streamMode ? new ArrowStreamReader(new ByteArrayInputStream(data), allocator) : new ArrowFileReader( new SeekableReadChannel(new ByteArrayReadableSeekableByteChannel(data)), allocator)) { // verify schema Schema readSchema = reader.getVectorSchemaRoot().getSchema(); assertEquals(1, readSchema.getFields().size()); assertEquals(expectedField, readSchema.getFields().get(0)); // verify vector schema root assertTrue(reader.loadNextBatch()); VectorSchemaRoot root = reader.getVectorSchemaRoot(); assertEquals(1, root.getFieldVectors().size()); ValueVector encodedVector = root.getVector(0); assertEquals(expectedIndices.length, encodedVector.getValueCount()); for (int i = 0; i < expectedIndices.length; i++) { assertEquals(expectedIndices[i], valGetter.applyAsInt(encodedVector, i)); } // verify dictionary Map dictVectors = reader.getDictionaryVectors(); assertEquals(1, dictVectors.size()); Dictionary dictionary = dictVectors.get(dictionaryID); assertNotNull(dictionary); assertTrue(dictionary.getVector() instanceof VarCharVector); VarCharVector dictVector = (VarCharVector) dictionary.getVector(); assertEquals(expectedDictItems.length, dictVector.getValueCount()); for (int i = 0; i < dictVector.getValueCount(); i++) { assertArrayEquals(expectedDictItems[i].getBytes(StandardCharsets.UTF_8), dictVector.get(i)); } } } private ValueVector createEncodedVector(int bitWidth, VarCharVector dictionaryVector) { final DictionaryEncoding dictionaryEncoding = new DictionaryEncoding(bitWidth, false, new ArrowType.Int(bitWidth, false)); Dictionary dictionary = new Dictionary(dictionaryVector, dictionaryEncoding); dictionaryProvider.put(dictionary); final FieldType type = new FieldType(true, dictionaryEncoding.getIndexType(), dictionaryEncoding, null); final Field field = new Field("encoded", type, null); return field.createVector(allocator); } @ParameterizedTest(name = "stream mode = {0}") @MethodSource("getRepeat") public void testUInt1RoundTrip(boolean streamMode) throws IOException { final int vectorLength = UInt1Vector.MAX_UINT1 & UInt1Vector.PROMOTION_MASK; try (VarCharVector dictionaryVector = new VarCharVector("dictionary", allocator); UInt1Vector encodedVector1 = (UInt1Vector) createEncodedVector(8, dictionaryVector)) { int[] indices = new int[vectorLength]; String[] dictionaryItems = new String[vectorLength]; for (int i = 0; i < vectorLength; i++) { encodedVector1.setSafe(i, (byte) i); indices[i] = i; dictionaryItems[i] = String.valueOf(i); } encodedVector1.setValueCount(vectorLength); setVector(dictionaryVector, dictionaryItems); byte[] data = writeData(streamMode, encodedVector1); readData( streamMode, data, encodedVector1.getField(), (vector, index) -> (int) ((UInt1Vector) vector).getValueAsLong(index), 8L, indices, dictionaryItems); } } @ParameterizedTest(name = "stream mode = {0}") @MethodSource("getRepeat") public void testUInt2RoundTrip(boolean streamMode) throws IOException { try (VarCharVector dictionaryVector = new VarCharVector("dictionary", allocator); UInt2Vector encodedVector2 = (UInt2Vector) createEncodedVector(16, dictionaryVector)) { int[] indices = new int[] {1, 3, 5, 7, 9, UInt2Vector.MAX_UINT2}; String[] dictItems = new String[UInt2Vector.MAX_UINT2]; for (int i = 0; i < UInt2Vector.MAX_UINT2; i++) { dictItems[i] = String.valueOf(i); } setVector( encodedVector2, (char) 1, (char) 3, (char) 5, (char) 7, (char) 9, UInt2Vector.MAX_UINT2); setVector(dictionaryVector, dictItems); byte[] data = writeData(streamMode, encodedVector2); readData( streamMode, data, encodedVector2.getField(), (vector, index) -> (int) ((UInt2Vector) vector).getValueAsLong(index), 16L, indices, dictItems); } } @ParameterizedTest(name = "stream mode = {0}") @MethodSource("getRepeat") public void testUInt4RoundTrip(boolean streamMode) throws IOException { final int dictLength = 10; try (VarCharVector dictionaryVector = new VarCharVector("dictionary", allocator); UInt4Vector encodedVector4 = (UInt4Vector) createEncodedVector(32, dictionaryVector)) { int[] indices = new int[] {1, 3, 5, 7, 9}; String[] dictItems = new String[dictLength]; for (int i = 0; i < dictLength; i++) { dictItems[i] = String.valueOf(i); } setVector(encodedVector4, 1, 3, 5, 7, 9); setVector(dictionaryVector, dictItems); setVector(encodedVector4, 1, 3, 5, 7, 9); byte[] data = writeData(streamMode, encodedVector4); readData( streamMode, data, encodedVector4.getField(), (vector, index) -> (int) ((UInt4Vector) vector).getValueAsLong(index), 32L, indices, dictItems); } } @ParameterizedTest(name = "stream mode = {0}") @MethodSource("getRepeat") public void testUInt8RoundTrip(boolean streamMode) throws IOException { final int dictLength = 10; try (VarCharVector dictionaryVector = new VarCharVector("dictionary", allocator); UInt8Vector encodedVector8 = (UInt8Vector) createEncodedVector(64, dictionaryVector)) { int[] indices = new int[] {1, 3, 5, 7, 9}; String[] dictItems = new String[dictLength]; for (int i = 0; i < dictLength; i++) { dictItems[i] = String.valueOf(i); } setVector(encodedVector8, 1L, 3L, 5L, 7L, 9L); setVector(dictionaryVector, dictItems); byte[] data = writeData(streamMode, encodedVector8); readData( streamMode, data, encodedVector8.getField(), (vector, index) -> (int) ((UInt8Vector) vector).getValueAsLong(index), 64L, indices, dictItems); } } static Stream getRepeat() { return Stream.of(Arguments.of(true), Arguments.of(false)); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/ipc/message/000077500000000000000000000000001515451615200277325ustar00rootroot00000000000000TestMessageMetadataResult.java000066400000000000000000000024461515451615200356100ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/ipc/message/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.ipc.message; import static org.junit.jupiter.api.Assertions.assertEquals; import java.nio.ByteBuffer; import org.junit.jupiter.api.Test; public class TestMessageMetadataResult { @Test public void getMessageLength_returnsConstructValue() { // This API is used by spark. MessageMetadataResult result = new MessageMetadataResult( 1, ByteBuffer.allocate(0), new org.apache.arrow.flatbuf.Message()); assertEquals(1, result.getMessageLength()); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/pojo/000077500000000000000000000000001515451615200265025ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/pojo/TestConvert.java000066400000000000000000000201011515451615200316170ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.pojo; import static org.apache.arrow.vector.types.FloatingPointPrecision.DOUBLE; import static org.apache.arrow.vector.types.FloatingPointPrecision.SINGLE; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import com.google.flatbuffers.FlatBufferBuilder; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; import java.util.Map; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.Collections2; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.UnionMode; import org.apache.arrow.vector.types.pojo.ArrowType.FloatingPoint; import org.apache.arrow.vector.types.pojo.ArrowType.Int; import org.apache.arrow.vector.types.pojo.ArrowType.List; import org.apache.arrow.vector.types.pojo.ArrowType.Struct; import org.apache.arrow.vector.types.pojo.ArrowType.Timestamp; import org.apache.arrow.vector.types.pojo.ArrowType.Union; import org.apache.arrow.vector.types.pojo.ArrowType.Utf8; 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.Test; /** Test conversion between Flatbuf and Pojo field representations. */ public class TestConvert { @Test public void simple() { Field initialField = new Field("a", FieldType.nullable(new Int(32, true)), null); run(initialField); } @Test public void complex() { java.util.List children = new ArrayList<>(); children.add(new Field("child1", FieldType.nullable(Utf8.INSTANCE), null)); children.add( new Field( "child2", FieldType.nullable(new FloatingPoint(SINGLE)), Collections.emptyList())); Field initialField = new Field("a", FieldType.nullable(Struct.INSTANCE), children); run(initialField); } @Test public void list() throws Exception { java.util.List children = new ArrayList<>(); try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); ListVector writeVector = ListVector.empty("list", allocator); FixedSizeListVector writeFixedVector = FixedSizeListVector.empty("fixedlist", 5, allocator)) { Field listVectorField = writeVector.getField(); children.add(listVectorField); Field listFixedVectorField = writeFixedVector.getField(); children.add(listFixedVectorField); } Field initialField = new Field("a", FieldType.nullable(Struct.INSTANCE), children); java.util.List parent = new ArrayList<>(); parent.add(initialField); FlatBufferBuilder builder = new FlatBufferBuilder(); builder.finish(initialField.getField(builder)); org.apache.arrow.flatbuf.Field flatBufField = org.apache.arrow.flatbuf.Field.getRootAsField(builder.dataBuffer()); Field finalField = Field.convertField(flatBufField); assertEquals(initialField, finalField); assertFalse(finalField.toString().contains("[DEFAULT]")); Schema initialSchema = new Schema(parent); String jsonSchema = initialSchema.toJson(); String modifiedSchema = jsonSchema.replace("$data$", "[DEFAULT]"); Schema tempSchema = Schema.fromJSON(modifiedSchema); FlatBufferBuilder schemaBuilder = new FlatBufferBuilder(); org.apache.arrow.vector.types.pojo.Schema schema = new org.apache.arrow.vector.types.pojo.Schema(tempSchema.getFields()); schemaBuilder.finish(schema.getSchema(schemaBuilder)); Schema finalSchema = Schema.deserialize(ByteBuffer.wrap(schemaBuilder.sizedByteArray())); assertFalse(finalSchema.toString().contains("[DEFAULT]")); } @Test public void schema() { java.util.List children = new ArrayList<>(); children.add(new Field("child1", FieldType.nullable(Utf8.INSTANCE), null)); children.add( new Field( "child2", FieldType.nullable(new FloatingPoint(SINGLE)), Collections.emptyList())); Schema initialSchema = new Schema(children); run(initialSchema); } @Test public void schemaMetadata() { java.util.List children = new ArrayList<>(); children.add(new Field("child1", FieldType.nullable(Utf8.INSTANCE), null)); children.add( new Field( "child2", FieldType.nullable(new FloatingPoint(SINGLE)), Collections.emptyList())); Map metadata = new HashMap<>(); metadata.put("key1", "value1"); metadata.put("key2", "value2"); Schema initialSchema = new Schema(children, metadata); run(initialSchema); } @Test public void nestedSchema() { java.util.List children = new ArrayList<>(); children.add(new Field("child1", FieldType.nullable(Utf8.INSTANCE), null)); children.add( new Field( "child2", FieldType.nullable(new FloatingPoint(SINGLE)), Collections.emptyList())); children.add( new Field( "child3", FieldType.nullable(new Struct()), Collections2.asImmutableList( new Field("child3.1", FieldType.nullable(Utf8.INSTANCE), null), new Field( "child3.2", FieldType.nullable(new FloatingPoint(DOUBLE)), Collections.emptyList())))); children.add( new Field( "child4", FieldType.nullable(new List()), Collections2.asImmutableList( new Field("child4.1", FieldType.nullable(Utf8.INSTANCE), null)))); children.add( new Field( "child5", FieldType.nullable( new Union( UnionMode.Sparse, new int[] {MinorType.TIMESTAMPMILLI.ordinal(), MinorType.FLOAT8.ordinal()})), Collections2.asImmutableList( new Field( "child5.1", FieldType.nullable(new Timestamp(TimeUnit.MILLISECOND, null)), null), new Field( "child5.2", FieldType.nullable(new FloatingPoint(DOUBLE)), Collections.emptyList()), new Field( "child5.3", FieldType.nullable(new Timestamp(TimeUnit.MILLISECOND, "UTC")), null)))); Schema initialSchema = new Schema(children); run(initialSchema); } private void run(Field initialField) { FlatBufferBuilder builder = new FlatBufferBuilder(); builder.finish(initialField.getField(builder)); org.apache.arrow.flatbuf.Field flatBufField = org.apache.arrow.flatbuf.Field.getRootAsField(builder.dataBuffer()); Field finalField = Field.convertField(flatBufField); assertEquals(initialField, finalField); } private void run(Schema initialSchema) { FlatBufferBuilder builder = new FlatBufferBuilder(); builder.finish(initialSchema.getSchema(builder)); org.apache.arrow.flatbuf.Schema flatBufSchema = org.apache.arrow.flatbuf.Schema.getRootAsSchema(builder.dataBuffer()); Schema finalSchema = Schema.convertSchema(flatBufSchema); assertEquals(initialSchema, finalSchema); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/table/000077500000000000000000000000001515451615200266225ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/table/BaseTableTest.java000066400000000000000000000265331515451615200321600ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.table; import static org.apache.arrow.vector.table.TestUtils.INT_VECTOR_NAME; import static org.apache.arrow.vector.table.TestUtils.INT_VECTOR_NAME_1; import static org.apache.arrow.vector.table.TestUtils.INT_VECTOR_NAME_2; import static org.apache.arrow.vector.table.TestUtils.intPlusVarcharColumns; import static org.apache.arrow.vector.table.TestUtils.twoIntColumns; 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.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; import java.util.Objects; 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.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryEncoder; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.DictionaryEncoding; import org.apache.arrow.vector.types.pojo.FieldType; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; class BaseTableTest { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @Test void getReaderByName() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { assertNotNull(t.getReader(INT_VECTOR_NAME_1)); } } @Test void getReaderByIndex() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { assertNotNull(t.getReader(0)); } } @Test void getReaderByField() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { assertNotNull(t.getReader(t.getField(INT_VECTOR_NAME_1))); } } @Test void getSchema() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { assertNotNull(t.getSchema()); assertEquals(2, t.getSchema().getFields().size()); } } @Test void insertVector() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { ArrowType intArrowType = new ArrowType.Int(32, true); FieldType intFieldType = new FieldType(true, intArrowType, null); IntVector v3 = new IntVector("3", intFieldType, allocator); List revisedVectors = t.insertVector(2, v3); assertEquals(3, revisedVectors.size()); assertEquals(v3, revisedVectors.get(2)); } } @Test void insertVectorFirstPosition() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { ArrowType intArrowType = new ArrowType.Int(32, true); FieldType intFieldType = new FieldType(true, intArrowType, null); IntVector v3 = new IntVector("3", intFieldType, allocator); List revisedVectors = t.insertVector(0, v3); assertEquals(3, revisedVectors.size()); assertEquals(v3, revisedVectors.get(0)); } } @Test void extractVector() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { List revisedVectors = t.extractVector(0); assertEquals(2, t.getVectorCount()); // vector not removed from table yet assertEquals(1, revisedVectors.size()); } } @Test void close() { IntVector v = new IntVector(INT_VECTOR_NAME, allocator); v.setSafe(0, 132); List vectors = new ArrayList<>(); vectors.add(v); v.setValueCount(1); try (Table t = new Table(vectors)) { t.close(); for (FieldVector fieldVector : t.fieldVectors) { assertEquals(0, fieldVector.getValueCount()); } } } @Test void getRowCount() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { assertEquals(2, t.getRowCount()); } } @Test void toVectorSchemaRoot() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { assertNotNull(t.getVector(INT_VECTOR_NAME_1)); assertNotNull(t.getVector(INT_VECTOR_NAME_2)); VectorSchemaRoot vsr = t.toVectorSchemaRoot(); assertNotNull(vsr.getVector(INT_VECTOR_NAME_1)); assertNotNull(vsr.getVector(INT_VECTOR_NAME_2)); assertEquals( t.getSchema().findField(INT_VECTOR_NAME_1), vsr.getSchema().findField(INT_VECTOR_NAME_1)); } } @Test void getVector() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { assertNotNull(t.getVector(0)); } } @Test void testGetVector() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { assertNotNull(t.getVector(INT_VECTOR_NAME_1)); assertThrows(IllegalArgumentException.class, () -> t.getVector("wrong name")); } } @Test void getVectorCopyByIndex() { List vectorList = twoIntColumns(allocator); List vectorList2 = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { // compare value by value for (int vIdx = 0; vIdx < vectorList.size(); vIdx++) { IntVector original = (IntVector) vectorList2.get(vIdx); IntVector copy = (IntVector) t.getVectorCopy(vIdx); assertNotNull(copy); assertEquals(2, copy.getValueCount()); assertEquals(0, copy.getNullCount()); for (int i = 0; i < t.getRowCount(); i++) { assertEquals(original.getObject(i), copy.getObject(i)); } } assertThrows(IllegalArgumentException.class, () -> t.getVector("wrong name")); } } @Test void getVectorCopyByName() { List vectorList = twoIntColumns(allocator); List vectorList2 = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { assertNotNull(t.getVectorCopy(INT_VECTOR_NAME_1)); for (int vIdx = 0; vIdx < vectorList.size(); vIdx++) { IntVector original = (IntVector) vectorList2.get(vIdx); IntVector copy = (IntVector) t.getVectorCopy(original.getName()); assertEquals(2, copy.getValueCount()); assertEquals(0, copy.getNullCount()); for (int i = 0; i < t.getRowCount(); i++) { assertEquals(original.getObject(i), copy.getObject(i)); } } assertThrows(IllegalArgumentException.class, () -> t.getVector("wrong name")); } } @Test void immutableCursor() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { assertNotNull(t.immutableRow()); } } @Test void contentToTsvString() { IntVector v = new IntVector(INT_VECTOR_NAME, allocator); v.setSafe(0, 1); v.setSafe(1, 2); v.setSafe(2, 3); v.setValueCount(3); try (Table t = Table.of(v)) { assertEquals(3, t.rowCount); List values = new ArrayList<>(); for (Row r : t) { values.add(r.getInt(INT_VECTOR_NAME)); } assertEquals(3, values.size()); List intList = new ArrayList<>(); intList.add(1); intList.add(2); intList.add(3); assertTrue(values.containsAll(intList)); String printed = "intCol\n" + "1\n" + "2\n" + "3\n"; assertEquals(printed, t.contentToTSVString()); } } @Test void isDeletedRow() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { assertFalse(t.isRowDeleted(0)); assertFalse(t.isRowDeleted(1)); } } @Test void testEncode() { List vectorList = intPlusVarcharColumns(allocator); VarCharVector original = (VarCharVector) vectorList.get(1); DictionaryProvider provider = getDictionary(); try (Table t = new Table(vectorList, vectorList.get(0).getValueCount(), provider)) { IntVector v = (IntVector) t.encode(original.getName(), 1L); assertNotNull(v); assertEquals(0, v.get(0)); assertEquals(1, v.get(1)); } } @Test void testDecode() { List vectorList = intPlusVarcharColumns(allocator); VarCharVector original = (VarCharVector) vectorList.get(1); VarCharVector dictionaryVector = new VarCharVector("dictionary", allocator); dictionaryVector.allocateNew(2); dictionaryVector.set(0, "one".getBytes(StandardCharsets.UTF_8)); dictionaryVector.set(1, "two".getBytes(StandardCharsets.UTF_8)); dictionaryVector.setValueCount(2); Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); DictionaryEncoder encoder = new DictionaryEncoder(dictionary, allocator); IntVector encoded = (IntVector) encoder.encode(original); vectorList.remove(original); vectorList.add(encoded); DictionaryProvider provider = getDictionary(); try (Table t = new Table(vectorList, vectorList.get(0).getValueCount(), provider)) { VarCharVector v = (VarCharVector) t.decode(encoded.getName(), 1L); assertNotNull(v); assertEquals("one", new String(Objects.requireNonNull(v.get(0)), StandardCharsets.UTF_8)); assertEquals("two", new String(Objects.requireNonNull(v.get(1)), StandardCharsets.UTF_8)); } } @Test void getProvider() { List vectorList = intPlusVarcharColumns(allocator); DictionaryProvider provider = getDictionary(); try (Table t = new Table(vectorList, vectorList.get(0).getValueCount(), provider)) { assertEquals(provider, t.getDictionaryProvider()); } } private DictionaryProvider getDictionary() { DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); DictionaryEncoding encoding = new DictionaryEncoding(1L, false, null); VarCharVector dictionaryVector = new VarCharVector("dictionary", allocator); dictionaryVector.allocateNew(2); dictionaryVector.set(0, "one".getBytes(StandardCharsets.UTF_8)); dictionaryVector.set(1, "two".getBytes(StandardCharsets.UTF_8)); dictionaryVector.setValueCount(2); Dictionary dictionary = new Dictionary(dictionaryVector, encoding); provider.put(dictionary); return provider; } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/table/RowTest.java000066400000000000000000001053231515451615200311000ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.table; import static org.apache.arrow.vector.table.TestUtils.BIGINT_INT_MAP_VECTOR_NAME; import static org.apache.arrow.vector.table.TestUtils.FIXEDBINARY_VECTOR_NAME_1; import static org.apache.arrow.vector.table.TestUtils.INT_LIST_VECTOR_NAME; import static org.apache.arrow.vector.table.TestUtils.INT_VECTOR_NAME_1; import static org.apache.arrow.vector.table.TestUtils.STRUCT_VECTOR_NAME; import static org.apache.arrow.vector.table.TestUtils.UNION_VECTOR_NAME; import static org.apache.arrow.vector.table.TestUtils.VARBINARY_VECTOR_NAME_1; import static org.apache.arrow.vector.table.TestUtils.VARCHAR_VECTOR_NAME_1; import static org.apache.arrow.vector.table.TestUtils.fixedWidthVectors; import static org.apache.arrow.vector.table.TestUtils.intPlusFixedBinaryColumns; import static org.apache.arrow.vector.table.TestUtils.intPlusLargeVarBinaryColumns; import static org.apache.arrow.vector.table.TestUtils.intPlusLargeVarcharColumns; import static org.apache.arrow.vector.table.TestUtils.intPlusVarBinaryColumns; import static org.apache.arrow.vector.table.TestUtils.intPlusVarcharColumns; import static org.apache.arrow.vector.table.TestUtils.simpleDenseUnionVector; import static org.apache.arrow.vector.table.TestUtils.simpleListVector; import static org.apache.arrow.vector.table.TestUtils.simpleMapVector; import static org.apache.arrow.vector.table.TestUtils.simpleStructVector; import static org.apache.arrow.vector.table.TestUtils.simpleUnionVector; import static org.apache.arrow.vector.table.TestUtils.timezoneTemporalVectors; import static org.apache.arrow.vector.table.TestUtils.twoIntColumns; 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 java.math.BigDecimal; import java.nio.charset.StandardCharsets; import java.time.Duration; import java.time.LocalDateTime; import java.time.Period; import java.util.ArrayList; import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntervalDayVector; import org.apache.arrow.vector.IntervalMonthDayNanoVector; import org.apache.arrow.vector.IntervalYearVector; import org.apache.arrow.vector.PeriodDuration; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.complex.DenseUnionVector; 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.UnionVector; import org.apache.arrow.vector.holders.NullableBigIntHolder; import org.apache.arrow.vector.holders.NullableBitHolder; import org.apache.arrow.vector.holders.NullableDecimalHolder; import org.apache.arrow.vector.holders.NullableDurationHolder; import org.apache.arrow.vector.holders.NullableFloat4Holder; import org.apache.arrow.vector.holders.NullableFloat8Holder; import org.apache.arrow.vector.holders.NullableIntHolder; import org.apache.arrow.vector.holders.NullableIntervalDayHolder; import org.apache.arrow.vector.holders.NullableIntervalMonthDayNanoHolder; import org.apache.arrow.vector.holders.NullableIntervalYearHolder; import org.apache.arrow.vector.holders.NullableSmallIntHolder; import org.apache.arrow.vector.holders.NullableTimeMicroHolder; import org.apache.arrow.vector.holders.NullableTimeMilliHolder; import org.apache.arrow.vector.holders.NullableTimeNanoHolder; import org.apache.arrow.vector.holders.NullableTimeSecHolder; import org.apache.arrow.vector.holders.NullableTimeStampMicroHolder; import org.apache.arrow.vector.holders.NullableTimeStampMicroTZHolder; import org.apache.arrow.vector.holders.NullableTimeStampMilliHolder; import org.apache.arrow.vector.holders.NullableTimeStampMilliTZHolder; import org.apache.arrow.vector.holders.NullableTimeStampNanoHolder; import org.apache.arrow.vector.holders.NullableTimeStampNanoTZHolder; import org.apache.arrow.vector.holders.NullableTimeStampSecHolder; import org.apache.arrow.vector.holders.NullableTimeStampSecTZHolder; import org.apache.arrow.vector.holders.NullableTinyIntHolder; import org.apache.arrow.vector.holders.NullableUInt1Holder; import org.apache.arrow.vector.holders.NullableUInt2Holder; import org.apache.arrow.vector.holders.NullableUInt4Holder; import org.apache.arrow.vector.holders.NullableUInt8Holder; import org.apache.arrow.vector.types.IntervalUnit; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.TestExtensionType; import org.apache.arrow.vector.util.JsonStringHashMap; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; class RowTest { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void terminate() { allocator.close(); } @Test void constructor() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { Row c = t.immutableRow(); assertEquals(StandardCharsets.UTF_8, c.getDefaultCharacterSet()); } } @Test void at() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { Row c = t.immutableRow(); assertEquals(c.getRowNumber(), -1); c.setPosition(1); assertEquals(c.getRowNumber(), 1); } } @Test void getIntByVectorIndex() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { Row c = t.immutableRow(); c.setPosition(1); assertEquals(2, c.getInt(0)); } } @Test void getIntByVectorName() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { Row c = t.immutableRow(); c.setPosition(1); assertEquals(2, c.getInt(INT_VECTOR_NAME_1)); } } @Test void testNameNotFound() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { Row c = t.immutableRow(); c.setPosition(1); assertThrows(IllegalArgumentException.class, () -> c.getVarCharObj("wrong name")); } } @Test void testWrongType() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { Row c = t.immutableRow(); c.setPosition(1); assertThrows(ClassCastException.class, () -> c.getVarCharObj(INT_VECTOR_NAME_1)); } } @Test void getDecimal() { List vectors = new ArrayList<>(); DecimalVector decimalVector = new DecimalVector("decimal_vector", allocator, 55, 10); vectors.add(decimalVector); decimalVector.setSafe(0, new BigDecimal("0.0543278923")); decimalVector.setSafe(1, new BigDecimal("2.0543278923")); decimalVector.setValueCount(2); BigDecimal one = decimalVector.getObject(1); NullableDecimalHolder holder1 = new NullableDecimalHolder(); NullableDecimalHolder holder2 = new NullableDecimalHolder(); try (Table t = new Table(vectors)) { Row c = t.immutableRow(); c.setPosition(1); assertEquals(one, c.getDecimalObj("decimal_vector")); assertEquals(one, c.getDecimalObj(0)); c.getDecimal(0, holder1); c.getDecimal("decimal_vector", holder2); assertEquals(holder1.buffer, holder2.buffer); assertEquals(c.getDecimal(0).memoryAddress(), c.getDecimal("decimal_vector").memoryAddress()); } } @Test void getDuration() { List vectors = new ArrayList<>(); TimeUnit unit = TimeUnit.SECOND; final FieldType fieldType = FieldType.nullable(new ArrowType.Duration(unit)); DurationVector durationVector = new DurationVector("duration_vector", fieldType, allocator); NullableDurationHolder holder1 = new NullableDurationHolder(); NullableDurationHolder holder2 = new NullableDurationHolder(); holder1.value = 100; holder1.unit = TimeUnit.SECOND; holder1.isSet = 1; holder2.value = 200; holder2.unit = TimeUnit.SECOND; holder2.isSet = 1; vectors.add(durationVector); durationVector.setSafe(0, holder1); durationVector.setSafe(1, holder2); durationVector.setValueCount(2); Duration one = durationVector.getObject(1); try (Table t = new Table(vectors)) { Row c = t.immutableRow(); c.setPosition(1); assertEquals(one, c.getDurationObj("duration_vector")); assertEquals(one, c.getDurationObj(0)); c.getDuration(0, holder1); c.getDuration("duration_vector", holder2); assertEquals(holder1.value, holder2.value); ArrowBuf durationBuf1 = c.getDuration(0); ArrowBuf durationBuf2 = c.getDuration("duration_vector"); assertEquals(durationBuf1.memoryAddress(), durationBuf2.memoryAddress()); } } @Test void getIntervalDay() { List vectors = new ArrayList<>(); IntervalUnit unit = IntervalUnit.DAY_TIME; final FieldType fieldType = FieldType.nullable(new ArrowType.Interval(unit)); IntervalDayVector intervalDayVector = new IntervalDayVector("intervalDay_vector", fieldType, allocator); NullableIntervalDayHolder holder1 = new NullableIntervalDayHolder(); NullableIntervalDayHolder holder2 = new NullableIntervalDayHolder(); holder1.days = 100; holder1.milliseconds = 1000; holder1.isSet = 1; holder2.days = 200; holder2.milliseconds = 2000; holder2.isSet = 1; vectors.add(intervalDayVector); intervalDayVector.setSafe(0, holder1); intervalDayVector.setSafe(1, holder2); intervalDayVector.setValueCount(2); Duration one = intervalDayVector.getObject(1); try (Table t = new Table(vectors)) { Row c = t.immutableRow(); c.setPosition(1); assertEquals(one, c.getIntervalDayObj("intervalDay_vector")); assertEquals(one, c.getIntervalDayObj(0)); c.getIntervalDay(0, holder1); c.getIntervalDay("intervalDay_vector", holder2); assertEquals(holder1.days, holder2.days); assertEquals(holder1.milliseconds, holder2.milliseconds); ArrowBuf intDayBuf1 = c.getIntervalDay(0); ArrowBuf intDayBuf2 = c.getIntervalDay("intervalDay_vector"); assertEquals(intDayBuf1.memoryAddress(), intDayBuf2.memoryAddress()); } } @Test void getIntervalMonth() { List vectors = new ArrayList<>(); IntervalUnit unit = IntervalUnit.MONTH_DAY_NANO; final FieldType fieldType = FieldType.nullable(new ArrowType.Interval(unit)); IntervalMonthDayNanoVector intervalMonthVector = new IntervalMonthDayNanoVector("intervalMonth_vector", fieldType, allocator); NullableIntervalMonthDayNanoHolder holder1 = new NullableIntervalMonthDayNanoHolder(); NullableIntervalMonthDayNanoHolder holder2 = new NullableIntervalMonthDayNanoHolder(); holder1.days = 1; holder1.months = 10; holder1.isSet = 1; holder2.days = 2; holder2.months = 20; holder2.isSet = 1; vectors.add(intervalMonthVector); intervalMonthVector.setSafe(0, holder1); intervalMonthVector.setSafe(1, holder2); intervalMonthVector.setValueCount(2); PeriodDuration one = intervalMonthVector.getObject(1); try (Table t = new Table(vectors)) { Row c = t.immutableRow(); c.setPosition(1); assertEquals(one, c.getIntervalMonthDayNanoObj("intervalMonth_vector")); assertEquals(one, c.getIntervalMonthDayNanoObj(0)); c.getIntervalMonthDayNano(0, holder1); c.getIntervalMonthDayNano("intervalMonth_vector", holder2); assertEquals(holder1.days, holder2.days); assertEquals(holder1.months, holder2.months); ArrowBuf intMonthBuf1 = c.getIntervalMonthDayNano(0); ArrowBuf intMonthBuf2 = c.getIntervalMonthDayNano("intervalMonth_vector"); assertEquals(intMonthBuf1.memoryAddress(), intMonthBuf2.memoryAddress()); } } @Test void getIntervalYear() { List vectors = new ArrayList<>(); IntervalUnit unit = IntervalUnit.YEAR_MONTH; final FieldType fieldType = FieldType.nullable(new ArrowType.Interval(unit)); IntervalYearVector intervalYearVector = new IntervalYearVector("intervalYear_vector", fieldType, allocator); NullableIntervalYearHolder holder1 = new NullableIntervalYearHolder(); NullableIntervalYearHolder holder2 = new NullableIntervalYearHolder(); holder1.value = 1; holder1.isSet = 1; holder2.value = 2; holder2.isSet = 1; vectors.add(intervalYearVector); intervalYearVector.setSafe(0, holder1); intervalYearVector.setSafe(1, holder2); intervalYearVector.setValueCount(2); Period one = intervalYearVector.getObject(1); try (Table t = new Table(vectors)) { Row c = t.immutableRow(); c.setPosition(1); assertEquals(one, c.getIntervalYearObj("intervalYear_vector")); assertEquals(one, c.getIntervalYearObj(0)); c.getIntervalYear(0, holder1); c.getIntervalYear("intervalYear_vector", holder2); assertEquals(holder1.value, holder2.value); int intYear1 = c.getIntervalYear(0); int intYear2 = c.getIntervalYear("intervalYear_vector"); assertEquals(2, intYear1); assertEquals(intYear1, intYear2); } } @Test void getBit() { List vectors = new ArrayList<>(); BitVector bitVector = new BitVector("bit_vector", allocator); NullableBitHolder holder1 = new NullableBitHolder(); NullableBitHolder holder2 = new NullableBitHolder(); vectors.add(bitVector); bitVector.setSafe(0, 0); bitVector.setSafe(1, 1); bitVector.setValueCount(2); int one = bitVector.get(1); try (Table t = new Table(vectors)) { Row c = t.immutableRow(); c.setPosition(1); assertEquals(one, c.getBit("bit_vector")); assertEquals(one, c.getBit(0)); c.getBit(0, holder1); c.getBit("bit_vector", holder2); assertEquals(holder1.value, holder2.value); } } @Test void hasNext() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { Row c = t.immutableRow(); assertTrue(c.hasNext()); c.setPosition(1); assertFalse(c.hasNext()); } } @Test void next() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { Row c = t.immutableRow(); c.setPosition(0); c.next(); assertEquals(1, c.getRowNumber()); } } @Test void isNull() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { Row c = t.immutableRow(); c.setPosition(1); assertFalse(c.isNull(0)); } } @Test void isNullByFieldName() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { Row c = t.immutableRow(); c.setPosition(1); assertFalse(c.isNull(INT_VECTOR_NAME_1)); } } @Test void fixedWidthVectorTest() { List vectorList = fixedWidthVectors(allocator, 2); try (Table t = new Table(vectorList)) { Row c = t.immutableRow(); c.setPosition(1); // integer tests using vector name and index assertFalse(c.isNull("bigInt_vector")); assertEquals(c.getInt("int_vector"), c.getInt(0)); assertEquals(c.getBigInt("bigInt_vector"), c.getBigInt(1)); assertEquals(c.getSmallInt("smallInt_vector"), c.getSmallInt(2)); assertEquals(c.getTinyInt("tinyInt_vector"), c.getTinyInt(3)); // integer tests using Nullable Holders NullableIntHolder int4Holder = new NullableIntHolder(); NullableTinyIntHolder int1Holder = new NullableTinyIntHolder(); NullableSmallIntHolder int2Holder = new NullableSmallIntHolder(); NullableBigIntHolder int8Holder = new NullableBigIntHolder(); c.getInt(0, int4Holder); c.getBigInt(1, int8Holder); c.getSmallInt(2, int2Holder); c.getTinyInt(3, int1Holder); assertEquals(c.getInt("int_vector"), int4Holder.value); assertEquals(c.getBigInt("bigInt_vector"), int8Holder.value); assertEquals(c.getSmallInt("smallInt_vector"), int2Holder.value); assertEquals(c.getTinyInt("tinyInt_vector"), int1Holder.value); c.getInt("int_vector", int4Holder); c.getBigInt("bigInt_vector", int8Holder); c.getSmallInt("smallInt_vector", int2Holder); c.getTinyInt("tinyInt_vector", int1Holder); assertEquals(c.getInt("int_vector"), int4Holder.value); assertEquals(c.getBigInt("bigInt_vector"), int8Holder.value); assertEquals(c.getSmallInt("smallInt_vector"), int2Holder.value); assertEquals(c.getTinyInt("tinyInt_vector"), int1Holder.value); // uint tests using vector name and index assertEquals(c.getUInt1("uInt1_vector"), c.getUInt1(4)); assertEquals(c.getUInt2("uInt2_vector"), c.getUInt2(5)); assertEquals(c.getUInt4("uInt4_vector"), c.getUInt4(6)); assertEquals(c.getUInt8("uInt8_vector"), c.getUInt8(7)); // UInt tests using Nullable Holders NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); NullableUInt1Holder uInt1Holder = new NullableUInt1Holder(); NullableUInt2Holder uInt2Holder = new NullableUInt2Holder(); NullableUInt8Holder uInt8Holder = new NullableUInt8Holder(); // fill the holders using vector index and test c.getUInt1(4, uInt1Holder); c.getUInt2(5, uInt2Holder); c.getUInt4(6, uInt4Holder); c.getUInt8(7, uInt8Holder); assertEquals(c.getUInt1("uInt1_vector"), uInt1Holder.value); assertEquals(c.getUInt2("uInt2_vector"), uInt2Holder.value); assertEquals(c.getUInt4("uInt4_vector"), uInt4Holder.value); assertEquals(c.getUInt8("uInt8_vector"), uInt8Holder.value); // refill the holders using vector name and retest c.getUInt1("uInt1_vector", uInt1Holder); c.getUInt2("uInt2_vector", uInt2Holder); c.getUInt4("uInt4_vector", uInt4Holder); c.getUInt8("uInt8_vector", uInt8Holder); assertEquals(c.getUInt1("uInt1_vector"), uInt1Holder.value); assertEquals(c.getUInt2("uInt2_vector"), uInt2Holder.value); assertEquals(c.getUInt4("uInt4_vector"), uInt4Holder.value); assertEquals(c.getUInt8("uInt8_vector"), uInt8Holder.value); // tests floating point assertEquals(c.getFloat4("float4_vector"), c.getFloat4(8)); assertEquals(c.getFloat8("float8_vector"), c.getFloat8(9)); // floating point tests using Nullable Holders NullableFloat4Holder float4Holder = new NullableFloat4Holder(); NullableFloat8Holder float8Holder = new NullableFloat8Holder(); // fill the holders using vector index and test c.getFloat4(8, float4Holder); c.getFloat8(9, float8Holder); assertEquals(c.getFloat4("float4_vector"), float4Holder.value); assertEquals(c.getFloat8("float8_vector"), float8Holder.value); // refill the holders using vector name and retest c.getFloat4("float4_vector", float4Holder); c.getFloat8("float8_vector", float8Holder); assertEquals(c.getFloat4("float4_vector"), float4Holder.value); assertEquals(c.getFloat8("float8_vector"), float8Holder.value); // test time values using vector name versus vector index assertEquals(c.getTimeSec("timeSec_vector"), c.getTimeSec(10)); assertEquals(c.getTimeMilli("timeMilli_vector"), c.getTimeMilli(11)); assertEquals(c.getTimeMicro("timeMicro_vector"), c.getTimeMicro(12)); assertEquals(c.getTimeNano("timeNano_vector"), c.getTimeNano(13)); // time tests using Nullable Holders NullableTimeSecHolder timeSecHolder = new NullableTimeSecHolder(); NullableTimeMilliHolder timeMilliHolder = new NullableTimeMilliHolder(); NullableTimeMicroHolder timeMicroHolder = new NullableTimeMicroHolder(); NullableTimeNanoHolder timeNanoHolder = new NullableTimeNanoHolder(); // fill the holders using vector index and test c.getTimeSec(10, timeSecHolder); c.getTimeMilli(11, timeMilliHolder); c.getTimeMicro(12, timeMicroHolder); c.getTimeNano(13, timeNanoHolder); assertEquals(c.getTimeSec("timeSec_vector"), timeSecHolder.value); assertEquals(c.getTimeMilli("timeMilli_vector"), timeMilliHolder.value); assertEquals(c.getTimeMicro("timeMicro_vector"), timeMicroHolder.value); assertEquals(c.getTimeNano("timeNano_vector"), timeNanoHolder.value); LocalDateTime milliDT = c.getTimeMilliObj(11); assertNotNull(milliDT); assertEquals(milliDT, c.getTimeMilliObj("timeMilli_vector")); // refill the holders using vector name and retest c.getTimeSec("timeSec_vector", timeSecHolder); c.getTimeMilli("timeMilli_vector", timeMilliHolder); c.getTimeMicro("timeMicro_vector", timeMicroHolder); c.getTimeNano("timeNano_vector", timeNanoHolder); assertEquals(c.getTimeSec("timeSec_vector"), timeSecHolder.value); assertEquals(c.getTimeMilli("timeMilli_vector"), timeMilliHolder.value); assertEquals(c.getTimeMicro("timeMicro_vector"), timeMicroHolder.value); assertEquals(c.getTimeNano("timeNano_vector"), timeNanoHolder.value); assertEquals(c.getTimeStampSec("timeStampSec_vector"), c.getTimeStampSec(14)); assertEquals(c.getTimeStampMilli("timeStampMilli_vector"), c.getTimeStampMilli(15)); assertEquals(c.getTimeStampMicro("timeStampMicro_vector"), c.getTimeStampMicro(16)); assertEquals(c.getTimeStampNano("timeStampNano_vector"), c.getTimeStampNano(17)); // time stamp tests using Nullable Holders NullableTimeStampSecHolder timeStampSecHolder = new NullableTimeStampSecHolder(); NullableTimeStampMilliHolder timeStampMilliHolder = new NullableTimeStampMilliHolder(); NullableTimeStampMicroHolder timeStampMicroHolder = new NullableTimeStampMicroHolder(); NullableTimeStampNanoHolder timeStampNanoHolder = new NullableTimeStampNanoHolder(); // fill the holders using vector index and test c.getTimeStampSec(14, timeStampSecHolder); c.getTimeStampMilli(15, timeStampMilliHolder); c.getTimeStampMicro(16, timeStampMicroHolder); c.getTimeStampNano(17, timeStampNanoHolder); assertEquals(c.getTimeStampSec("timeStampSec_vector"), timeStampSecHolder.value); assertEquals(c.getTimeStampMilli("timeStampMilli_vector"), timeStampMilliHolder.value); assertEquals(c.getTimeStampMicro("timeStampMicro_vector"), timeStampMicroHolder.value); assertEquals(c.getTimeStampNano("timeStampNano_vector"), timeStampNanoHolder.value); LocalDateTime secDT = c.getTimeStampSecObj(14); assertNotNull(secDT); assertEquals(secDT, c.getTimeStampSecObj("timeStampSec_vector")); LocalDateTime milliDT1 = c.getTimeStampMilliObj(15); assertNotNull(milliDT1); assertEquals(milliDT1, c.getTimeStampMilliObj("timeStampMilli_vector")); LocalDateTime microDT = c.getTimeStampMicroObj(16); assertNotNull(microDT); assertEquals(microDT, c.getTimeStampMicroObj("timeStampMicro_vector")); LocalDateTime nanoDT = c.getTimeStampNanoObj(17); assertNotNull(nanoDT); assertEquals(nanoDT, c.getTimeStampNanoObj("timeStampNano_vector")); // refill the holders using vector name and retest c.getTimeStampSec("timeStampSec_vector", timeStampSecHolder); c.getTimeStampMilli("timeStampMilli_vector", timeStampMilliHolder); c.getTimeStampMicro("timeStampMicro_vector", timeStampMicroHolder); c.getTimeStampNano("timeStampNano_vector", timeStampNanoHolder); assertEquals(c.getTimeStampSec("timeStampSec_vector"), timeStampSecHolder.value); assertEquals(c.getTimeStampMilli("timeStampMilli_vector"), timeStampMilliHolder.value); assertEquals(c.getTimeStampMicro("timeStampMicro_vector"), timeStampMicroHolder.value); assertEquals(c.getTimeStampNano("timeStampNano_vector"), timeStampNanoHolder.value); } } @Test void timestampsWithTimezones() { List vectorList = timezoneTemporalVectors(allocator, 2); try (Table t = new Table(vectorList)) { Row c = t.immutableRow(); c.setPosition(1); assertEquals(c.getTimeStampSecTZ("timeStampSecTZ_vector"), c.getTimeStampSecTZ(0)); assertEquals(c.getTimeStampMilliTZ("timeStampMilliTZ_vector"), c.getTimeStampMilliTZ(1)); assertEquals(c.getTimeStampMicroTZ("timeStampMicroTZ_vector"), c.getTimeStampMicroTZ(2)); assertEquals(c.getTimeStampNanoTZ("timeStampNanoTZ_vector"), c.getTimeStampNanoTZ(3)); // time stamp tests using Nullable Holders NullableTimeStampSecTZHolder timeStampSecHolder = new NullableTimeStampSecTZHolder(); NullableTimeStampMilliTZHolder timeStampMilliHolder = new NullableTimeStampMilliTZHolder(); NullableTimeStampMicroTZHolder timeStampMicroHolder = new NullableTimeStampMicroTZHolder(); NullableTimeStampNanoTZHolder timeStampNanoHolder = new NullableTimeStampNanoTZHolder(); // fill the holders using vector index and test c.getTimeStampSecTZ(0, timeStampSecHolder); c.getTimeStampMilliTZ(1, timeStampMilliHolder); c.getTimeStampMicroTZ(2, timeStampMicroHolder); c.getTimeStampNanoTZ(3, timeStampNanoHolder); long tsSec = timeStampSecHolder.value; long tsMil = timeStampMilliHolder.value; long tsMic = timeStampMicroHolder.value; long tsNan = timeStampNanoHolder.value; assertEquals(c.getTimeStampSecTZ("timeStampSecTZ_vector"), timeStampSecHolder.value); assertEquals(c.getTimeStampMilliTZ("timeStampMilliTZ_vector"), timeStampMilliHolder.value); assertEquals(c.getTimeStampMicroTZ("timeStampMicroTZ_vector"), timeStampMicroHolder.value); assertEquals(c.getTimeStampNanoTZ("timeStampNanoTZ_vector"), timeStampNanoHolder.value); // fill the holders using vector index and test c.getTimeStampSecTZ("timeStampSecTZ_vector", timeStampSecHolder); c.getTimeStampMilliTZ("timeStampMilliTZ_vector", timeStampMilliHolder); c.getTimeStampMicroTZ("timeStampMicroTZ_vector", timeStampMicroHolder); c.getTimeStampNanoTZ("timeStampNanoTZ_vector", timeStampNanoHolder); assertEquals(tsSec, timeStampSecHolder.value); assertEquals(tsMil, timeStampMilliHolder.value); assertEquals(tsMic, timeStampMicroHolder.value); assertEquals(tsNan, timeStampNanoHolder.value); } } @Test void getVarChar() { List vectorList = intPlusVarcharColumns(allocator); try (Table t = new Table(vectorList)) { Row c = t.immutableRow(); c.setPosition(1); assertEquals(c.getVarCharObj(1), "two"); assertEquals(c.getVarCharObj(1), c.getVarCharObj(VARCHAR_VECTOR_NAME_1)); assertArrayEquals( "two".getBytes(StandardCharsets.UTF_8), c.getVarChar(VARCHAR_VECTOR_NAME_1)); assertArrayEquals("two".getBytes(StandardCharsets.UTF_8), c.getVarChar(1)); } } @Test void getVarBinary() { List vectorList = intPlusVarBinaryColumns(allocator); try (Table t = new Table(vectorList)) { Row c = t.immutableRow(); c.setPosition(1); assertArrayEquals(c.getVarBinary(1), "two".getBytes(StandardCharsets.UTF_8)); assertArrayEquals(c.getVarBinary(1), c.getVarBinary(VARBINARY_VECTOR_NAME_1)); } } @Test void getLargeVarBinary() { List vectorList = intPlusLargeVarBinaryColumns(allocator); try (Table t = new Table(vectorList)) { Row c = t.immutableRow(); c.setPosition(1); assertArrayEquals(c.getLargeVarBinary(1), "two".getBytes(StandardCharsets.UTF_8)); assertArrayEquals(c.getLargeVarBinary(1), c.getLargeVarBinary(VARBINARY_VECTOR_NAME_1)); } } @Test void getLargeVarChar() { List vectorList = intPlusLargeVarcharColumns(allocator); try (Table t = new Table(vectorList)) { Row c = t.immutableRow(); c.setPosition(1); assertEquals(c.getLargeVarCharObj(1), "two"); assertEquals(c.getLargeVarCharObj(1), c.getLargeVarCharObj(VARCHAR_VECTOR_NAME_1)); assertArrayEquals( "two".getBytes(StandardCharsets.UTF_8), c.getLargeVarChar(VARCHAR_VECTOR_NAME_1)); assertArrayEquals("two".getBytes(StandardCharsets.UTF_8), c.getLargeVarChar(1)); } } @Test void getFixedBinary() { List vectorList = intPlusFixedBinaryColumns(allocator); try (Table t = new Table(vectorList)) { Row c = t.immutableRow(); c.setPosition(1); assertArrayEquals(c.getFixedSizeBinary(1), "two".getBytes(StandardCharsets.UTF_8)); assertArrayEquals(c.getFixedSizeBinary(1), c.getFixedSizeBinary(FIXEDBINARY_VECTOR_NAME_1)); } } @Test void testSimpleListVector1() { try (ListVector listVector = simpleListVector(allocator); VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.of(listVector); Table table = new Table(vectorSchemaRoot)) { for (Row c : table) { @SuppressWarnings("unchecked") List list = (List) c.getList(INT_LIST_VECTOR_NAME); assertEquals(10, list.size()); } } } @Test void testSimpleListVector2() { try (ListVector listVector = simpleListVector(allocator); VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.of(listVector); Table table = new Table(vectorSchemaRoot)) { for (Row c : table) { @SuppressWarnings("unchecked") List list = (List) c.getList(0); assertEquals(10, list.size()); } } } @Test void testSimpleStructVector1() { try (StructVector structVector = simpleStructVector(allocator); VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.of(structVector); Table table = new Table(vectorSchemaRoot)) { for (Row c : table) { @SuppressWarnings("unchecked") JsonStringHashMap struct = (JsonStringHashMap) c.getStruct(STRUCT_VECTOR_NAME); @SuppressWarnings("unchecked") JsonStringHashMap struct1 = (JsonStringHashMap) c.getStruct(0); int a = (int) struct.get("struct_int_child"); double b = (double) struct.get("struct_flt_child"); int a1 = (int) struct1.get("struct_int_child"); double b1 = (double) struct1.get("struct_flt_child"); assertNotNull(struct); assertEquals(a, a1); assertEquals(b, b1); assertTrue(a >= 0); assertTrue(b <= a, String.format("a = %s and b = %s", a, b)); } } } @Test void testSimpleUnionVector() { try (UnionVector unionVector = simpleUnionVector(allocator); VectorSchemaRoot vsr = VectorSchemaRoot.of(unionVector); Table table = new Table(vsr)) { Row c = table.immutableRow(); c.setPosition(0); Object object0 = c.getUnion(UNION_VECTOR_NAME); Object object1 = c.getUnion(0); assertEquals(object0, object1); c.setPosition(1); assertNull(c.getUnion(UNION_VECTOR_NAME)); c.setPosition(2); Object object2 = c.getUnion(UNION_VECTOR_NAME); assertEquals(100, object0); assertEquals(100, object2); } } @Test void testSimpleDenseUnionVector() { try (DenseUnionVector unionVector = simpleDenseUnionVector(allocator); VectorSchemaRoot vsr = VectorSchemaRoot.of(unionVector); Table table = new Table(vsr)) { Row c = table.immutableRow(); c.setPosition(0); Object object0 = c.getDenseUnion(UNION_VECTOR_NAME); Object object1 = c.getDenseUnion(0); assertEquals(object0, object1); c.setPosition(1); assertNull(c.getDenseUnion(UNION_VECTOR_NAME)); c.setPosition(2); Object object2 = c.getDenseUnion(UNION_VECTOR_NAME); assertEquals(100, object0); assertEquals(100, object2); } } @Test void testExtensionTypeVector() { TestExtensionType.LocationVector vector = new TestExtensionType.LocationVector("location", allocator); vector.allocateNew(); vector.set(0, 34.073814f, -118.240784f); vector.setValueCount(1); try (VectorSchemaRoot vsr = VectorSchemaRoot.of(vector); Table table = new Table(vsr)) { Row c = table.immutableRow(); c.setPosition(0); Object object0 = c.getExtensionType("location"); Object object1 = c.getExtensionType(0); assertEquals(object0, object1); @SuppressWarnings("unchecked") JsonStringHashMap struct0 = (JsonStringHashMap) object0; assertEquals(34.073814f, struct0.get("Latitude")); } } @Test void testSimpleMapVector1() { try (MapVector mapVector = simpleMapVector(allocator); Table table = Table.of(mapVector)) { int i = 1; for (Row c : table) { @SuppressWarnings("unchecked") List> list = (List>) c.getMap(BIGINT_INT_MAP_VECTOR_NAME); @SuppressWarnings("unchecked") List> list1 = (List>) c.getMap(0); for (int j = 0; j < list1.size(); j++) { assertEquals(list.get(j), list1.get(j)); } if (list != null && !list.isEmpty()) { assertEquals(i, list.size()); for (JsonStringHashMap sv : list) { assertEquals(2, sv.size()); Long o1 = (Long) sv.get("key"); Integer o2 = (Integer) sv.get("value"); assertEquals(o1, o2.longValue()); } } i++; } } } @Test void resetPosition() { try (ListVector listVector = simpleListVector(allocator); VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.of(listVector); Table table = new Table(vectorSchemaRoot)) { Row row = table.immutableRow(); row.next(); assertEquals(0, row.rowNumber); row.resetPosition(); assertEquals(-1, row.rowNumber); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/table/TableTest.java000066400000000000000000000220701515451615200313550ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.table; import static org.apache.arrow.vector.table.TestUtils.INT_VECTOR_NAME_1; import static org.apache.arrow.vector.table.TestUtils.twoIntColumns; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.HashSet; import java.util.Iterator; import java.util.List; 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 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.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; class TableTest { private final ArrowType intArrowType = new ArrowType.Int(32, true); private final FieldType intFieldType = new FieldType(true, intArrowType, null); private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @Test void of() { List vectorList = twoIntColumns(allocator); try (Table t = Table.of(vectorList.toArray(new FieldVector[2]))) { Row c = t.immutableRow(); assertEquals(2, t.getRowCount()); assertEquals(2, t.getVectorCount()); IntVector intVector1 = (IntVector) vectorList.get(0); assertEquals(INT_VECTOR_NAME_1, intVector1.getName()); c.setPosition(0); // Now test changes to the first vector // first Table value is 1 assertEquals(1, c.getInt(INT_VECTOR_NAME_1)); // original vector is updated to set first value to 44 intVector1.setSafe(0, 44); assertEquals(44, intVector1.get(0)); // first Table value is still 1 for the zeroth vector assertEquals(1, c.getInt(0)); } } @Test void constructor() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList, 2)) { assertEquals(2, t.getRowCount()); assertEquals(2, t.getVectorCount()); Row c = t.immutableRow(); IntVector intVector1 = (IntVector) vectorList.get(0); c.setPosition(0); // Now test changes to the first vector // first Table value is 1 assertEquals(1, c.getInt(INT_VECTOR_NAME_1)); // original vector is updated to set first value to 44 intVector1.setSafe(0, 44); assertEquals(44, intVector1.get(0)); assertEquals(44, ((IntVector) vectorList.get(0)).get(0)); // first Table value is still 1 for the zeroth vector assertEquals(1, c.getInt(INT_VECTOR_NAME_1)); } } /** * Tests construction with an iterable that's not a list (there is a specialty constructor for * Lists). */ @Test void constructor2() { List vectorList = twoIntColumns(allocator); Iterable iterable = new HashSet<>(vectorList); try (Table t = new Table(iterable)) { assertEquals(2, t.getRowCount()); assertEquals(2, t.getVectorCount()); Row c = t.immutableRow(); IntVector intVector1 = (IntVector) vectorList.get(0); c.setPosition(0); // Now test changes to the first vector // first Table value is 1 assertEquals(1, c.getInt(INT_VECTOR_NAME_1)); // original vector is updated to set first value to 44 intVector1.setSafe(0, 44); assertEquals(44, intVector1.get(0)); assertEquals(44, ((IntVector) vectorList.get(0)).get(0)); // first Table value is still 1 for the zeroth vector assertEquals(1, c.getInt(INT_VECTOR_NAME_1)); } } @Test void copy() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { assertEquals(2, t.getVectorCount()); try (Table copy = t.copy()) { for (FieldVector v : t.fieldVectors) { FieldVector vCopy = copy.getVector(v.getName()); assertNotNull(vCopy); assertEquals(v.getValueCount(), vCopy.getValueCount()); for (int i = 0; i < v.getValueCount(); i++) { Integer vValue = ((IntVector) v).getObject(i); Integer vCopyValue = ((IntVector) vCopy).getObject(i); assertEquals(vValue, vCopyValue); } } } } } @Test void addVector() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { IntVector v3 = new IntVector("3", intFieldType, allocator); Table t2 = t.addVector(2, v3); assertEquals(3, t2.fieldVectors.size()); assertTrue(t2.getVector("3").isNull(0)); assertTrue(t2.getVector("3").isNull(1)); t2.close(); } } @Test void removeVector() { List vectorList = twoIntColumns(allocator); IntVector v2 = (IntVector) vectorList.get(1); int val1 = v2.get(0); int val2 = v2.get(1); try (Table t = new Table(vectorList)) { Table t2 = t.removeVector(0); assertEquals(1, t2.fieldVectors.size()); assertEquals(val1, ((IntVector) t2.getVector(0)).get(0)); assertEquals(val2, ((IntVector) t2.getVector(0)).get(1)); } } /** Tests table iterator in enhanced for loop. */ @Test void iterator1() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { Iterator iterator = t.iterator(); assertNotNull(iterator); assertTrue(iterator.hasNext()); int sum = 0; for (Row row : t) { sum += row.getInt(0); } assertEquals(3, sum); } } /** Tests explicit iterator. */ @SuppressWarnings("WhileLoopReplaceableByForEach") @Test void iterator2() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { Iterator iterator = t.iterator(); assertNotNull(iterator); assertTrue(iterator.hasNext()); int sum = 0; Iterator it = t.iterator(); while (it.hasNext()) { Row row = it.next(); sum += row.getInt(0); } assertEquals(3, sum); } } /** * Tests a slice operation where no length is provided, so the range extends to the end of the * table. */ @Test void sliceToEnd() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { Table slice = t.slice(1); assertEquals(1, slice.rowCount); assertEquals(2, t.rowCount); // memory is copied for slice, not transferred slice.close(); } } /** Tests a slice operation with a given length parameter. */ @Test void sliceRange() { List vectorList = twoIntColumns(allocator); try (Table t = new Table(vectorList)) { Table slice = t.slice(1, 1); assertEquals(1, slice.rowCount); assertEquals(2, t.rowCount); // memory is copied for slice, not transferred slice.close(); } } /** * Tests creation of a table from a VectorSchemaRoot. * *

    Also tests that updates to the source Vectors do not impact the values in the Table */ @Test void constructFromVsr() { List vectorList = twoIntColumns(allocator); try (VectorSchemaRoot vsr = new VectorSchemaRoot(vectorList)) { Table t = new Table(vsr); Row c = t.immutableRow(); assertEquals(2, t.rowCount); assertEquals(0, vsr.getRowCount()); // memory is copied for slice, not transferred IntVector intVector1 = (IntVector) vectorList.get(0); c.setPosition(0); // Now test changes to the first vector // first Table value is 1 assertEquals(1, c.getInt(INT_VECTOR_NAME_1)); // original vector is updated to set first value to 44 intVector1.setSafe(0, 44); assertEquals(44, intVector1.get(0)); assertEquals(44, ((IntVector) vsr.getVector(0)).get(0)); // first Table value is still 1 for the zeroth vector assertEquals(1, c.getInt(INT_VECTOR_NAME_1)); // TEST FIELDS // Schema schema = t.schema; Field f1 = t.getField(INT_VECTOR_NAME_1); FieldVector fv1 = vectorList.get(0); assertEquals(f1, fv1.getField()); assertEquals(f1, schema.findField(INT_VECTOR_NAME_1)); t.close(); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/table/TestUtils.java000066400000000000000000000442011515451615200314260ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.table; import static org.apache.arrow.vector.complex.BaseRepeatedValueVector.OFFSET_WIDTH; import java.math.BigDecimal; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.BitVectorHelper; 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.FieldVector; import org.apache.arrow.vector.FixedSizeBinaryVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.GenerateSampleData; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.IntervalDayVector; import org.apache.arrow.vector.IntervalMonthDayNanoVector; import org.apache.arrow.vector.IntervalYearVector; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.LargeVarCharVector; 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.DenseUnionVector; 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.UnionVector; import org.apache.arrow.vector.complex.impl.NullableStructWriter; import org.apache.arrow.vector.complex.impl.UnionMapWriter; import org.apache.arrow.vector.complex.writer.Float8Writer; import org.apache.arrow.vector.complex.writer.IntWriter; import org.apache.arrow.vector.holders.NullableUInt4Holder; 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.types.pojo.FieldType; public class TestUtils { public static final String INT_VECTOR_NAME = "intCol"; public static final String INT_VECTOR_NAME_1 = "intCol1"; public static final String VARCHAR_VECTOR_NAME_1 = "varcharCol1"; public static final String VARBINARY_VECTOR_NAME_1 = "varbinaryCol1"; public static final String FIXEDBINARY_VECTOR_NAME_1 = "varbinaryCol1"; public static final String INT_VECTOR_NAME_2 = "intCol2"; public static final String INT_LIST_VECTOR_NAME = "int list vector"; public static final String BIGINT_INT_MAP_VECTOR_NAME = "bigint-int map vector"; public static final String STRUCT_VECTOR_NAME = "struct_vector"; public static final String UNION_VECTOR_NAME = "union_vector"; /** * Returns a list of two IntVectors to be used to instantiate Tables for testing. Each IntVector * has two values set. */ static List twoIntColumns(BufferAllocator allocator) { List vectorList = new ArrayList<>(); IntVector v1 = getSimpleIntVector(allocator); IntVector v2 = new IntVector(INT_VECTOR_NAME_2, allocator); v2.allocateNew(2); v2.set(0, 3); v2.set(1, 4); v2.setValueCount(2); vectorList.add(v1); vectorList.add(v2); return vectorList; } /** * Returns a list of two FieldVectors to be used to instantiate Tables for testing. The first * vector is an IntVector and the second is a VarCharVector. Each vector has two values set. */ static List intPlusVarcharColumns(BufferAllocator allocator) { List vectorList = new ArrayList<>(); IntVector v1 = getSimpleIntVector(allocator); VarCharVector v2 = new VarCharVector(VARCHAR_VECTOR_NAME_1, allocator); v2.allocateNew(2); v2.set(0, "one".getBytes(StandardCharsets.UTF_8)); v2.set(1, "two".getBytes(StandardCharsets.UTF_8)); v2.setValueCount(2); vectorList.add(v1); vectorList.add(v2); return vectorList; } /** * Returns a list of two FieldVectors to be used to instantiate Tables for testing. The first * vector is an IntVector and the second is a LargeVarCharVector. Each vector has two values set. */ static List intPlusLargeVarcharColumns(BufferAllocator allocator) { List vectorList = new ArrayList<>(); IntVector v1 = getSimpleIntVector(allocator); LargeVarCharVector v2 = new LargeVarCharVector(VARCHAR_VECTOR_NAME_1, allocator); v2.allocateNew(2); v2.set(0, "one".getBytes(StandardCharsets.UTF_8)); v2.set(1, "two".getBytes(StandardCharsets.UTF_8)); v2.setValueCount(2); vectorList.add(v1); vectorList.add(v2); return vectorList; } /** * Returns a list of two FieldVectors to be used to instantiate Tables for testing. The first * vector is an IntVector and the second is a VarBinaryVector. Each vector has two values set. The * large binary vectors values are "one" and "two" encoded with UTF-8 */ static List intPlusVarBinaryColumns(BufferAllocator allocator) { List vectorList = new ArrayList<>(); IntVector v1 = getSimpleIntVector(allocator); VarBinaryVector v2 = new VarBinaryVector(VARBINARY_VECTOR_NAME_1, allocator); v2.allocateNew(2); v2.set(0, "one".getBytes(StandardCharsets.UTF_8)); v2.set(1, "two".getBytes(StandardCharsets.UTF_8)); v2.setValueCount(2); vectorList.add(v1); vectorList.add(v2); return vectorList; } /** * Returns a list of two FieldVectors to be used to instantiate Tables for testing. The first * vector is an IntVector and the second is a VarBinaryVector. Each vector has two values set. The * large binary vectors values are "one" and "two" encoded with UTF-8 */ static List intPlusLargeVarBinaryColumns(BufferAllocator allocator) { List vectorList = new ArrayList<>(); IntVector v1 = getSimpleIntVector(allocator); LargeVarBinaryVector v2 = new LargeVarBinaryVector(VARBINARY_VECTOR_NAME_1, allocator); v2.allocateNew(2); v2.set(0, "one".getBytes(StandardCharsets.UTF_8)); v2.set(1, "two".getBytes(StandardCharsets.UTF_8)); v2.setValueCount(2); vectorList.add(v1); vectorList.add(v2); return vectorList; } /** * Returns a list of two FieldVectors to be used to instantiate Tables for testing. The first * vector is an IntVector and the second is a FixedSizeBinary vector. Each vector has two values * set. The large binary vectors values are "one" and "two" encoded with UTF-8 */ static List intPlusFixedBinaryColumns(BufferAllocator allocator) { List vectorList = new ArrayList<>(); IntVector v1 = getSimpleIntVector(allocator); FixedSizeBinaryVector v2 = new FixedSizeBinaryVector(FIXEDBINARY_VECTOR_NAME_1, allocator, 3); v2.allocateNew(2); v2.set(0, "one".getBytes(StandardCharsets.UTF_8)); v2.set(1, "two".getBytes(StandardCharsets.UTF_8)); v2.setValueCount(2); vectorList.add(v1); vectorList.add(v2); return vectorList; } private static IntVector getSimpleIntVector(BufferAllocator allocator) { IntVector v1 = new IntVector(INT_VECTOR_NAME_1, allocator); v1.allocateNew(2); v1.set(0, 1); v1.set(1, 2); v1.setValueCount(2); return v1; } /** * Returns a list of fixed-width vectors for testing. It includes * *

      *
    1. all integral and floating point types *
    2. all basic times and timestamps (second, milli, micro, nano *
    * * The vector names are based on their type name (e.g. BigIntVector is called "bigInt_vector" */ static List fixedWidthVectors(BufferAllocator allocator, int rowCount) { List vectors = new ArrayList<>(); numericVectors(vectors, allocator, rowCount); simpleTemporalVectors(vectors, allocator, rowCount); return vectors; } /** * Returns a list of all integral and floating point vectors. The vector names are based on their * type name (e.g. BigIntVector is called "bigInt_vector" */ static List numericVectors( List vectors, BufferAllocator allocator, int rowCount) { vectors.add(new IntVector("int_vector", allocator)); vectors.add(new BigIntVector("bigInt_vector", allocator)); vectors.add(new SmallIntVector("smallInt_vector", allocator)); vectors.add(new TinyIntVector("tinyInt_vector", allocator)); vectors.add(new UInt1Vector("uInt1_vector", allocator)); vectors.add(new UInt2Vector("uInt2_vector", allocator)); vectors.add(new UInt4Vector("uInt4_vector", allocator)); vectors.add(new UInt8Vector("uInt8_vector", allocator)); vectors.add(new Float4Vector("float4_vector", allocator)); vectors.add(new Float8Vector("float8_vector", allocator)); vectors.forEach(vec -> GenerateSampleData.generateTestData(vec, rowCount)); return vectors; } static List numericVectors(BufferAllocator allocator, int rowCount) { List vectors = new ArrayList<>(); return numericVectors(vectors, allocator, rowCount); } static List simpleTemporalVectors( List vectors, BufferAllocator allocator, int rowCount) { vectors.add(new TimeSecVector("timeSec_vector", allocator)); vectors.add(new TimeMilliVector("timeMilli_vector", allocator)); vectors.add(new TimeMicroVector("timeMicro_vector", allocator)); vectors.add(new TimeNanoVector("timeNano_vector", allocator)); vectors.add(new TimeStampSecVector("timeStampSec_vector", allocator)); vectors.add(new TimeStampMilliVector("timeStampMilli_vector", allocator)); vectors.add(new TimeStampMicroVector("timeStampMicro_vector", allocator)); vectors.add(new TimeStampNanoVector("timeStampNano_vector", allocator)); vectors.add(new DateMilliVector("dateMilli_vector", allocator)); vectors.add(new DateDayVector("dateDay_vector", allocator)); vectors.forEach(vec -> GenerateSampleData.generateTestData(vec, rowCount)); return vectors; } static List simpleTemporalVectors(BufferAllocator allocator, int rowCount) { List vectors = new ArrayList<>(); return simpleTemporalVectors(vectors, allocator, rowCount); } static List timezoneTemporalVectors(BufferAllocator allocator, int rowCount) { List vectors = new ArrayList<>(); vectors.add(new TimeStampSecTZVector("timeStampSecTZ_vector", allocator, "UTC")); vectors.add(new TimeStampMilliTZVector("timeStampMilliTZ_vector", allocator, "UTC")); vectors.add(new TimeStampMicroTZVector("timeStampMicroTZ_vector", allocator, "UTC")); vectors.add(new TimeStampNanoTZVector("timeStampNanoTZ_vector", allocator, "UTC")); vectors.forEach(vec -> GenerateSampleData.generateTestData(vec, rowCount)); return vectors; } static List intervalVectors(BufferAllocator allocator, int rowCount) { List vectors = new ArrayList<>(); vectors.add(new IntervalDayVector("intervalDay_vector", allocator)); vectors.add(new IntervalYearVector("intervalYear_vector", allocator)); vectors.add(new IntervalMonthDayNanoVector("intervalMonthDayNano_vector", allocator)); vectors.add( new DurationVector( "duration_vector", new FieldType(true, new ArrowType.Duration(TimeUnit.SECOND), null), allocator)); vectors.forEach(vec -> GenerateSampleData.generateTestData(vec, rowCount)); return vectors; } /** Returns a list vector of ints. */ static ListVector simpleListVector(BufferAllocator allocator) { ListVector listVector = ListVector.empty(INT_LIST_VECTOR_NAME, allocator); final int innerCount = 80; // total number of values final int outerCount = 8; // total number of values in the list vector itself final int listLength = innerCount / outerCount; // length of an individual list Types.MinorType type = Types.MinorType.INT; listVector.addOrGetVector(FieldType.nullable(type.getType())); listVector.allocateNew(); IntVector dataVector = (IntVector) listVector.getDataVector(); for (int i = 0; i < innerCount; i++) { dataVector.set(i, i); } dataVector.setValueCount(innerCount); for (int i = 0; i < outerCount; i++) { BitVectorHelper.setBit(listVector.getValidityBuffer(), i); listVector.getOffsetBuffer().setInt(i * OFFSET_WIDTH, i * listLength); listVector.getOffsetBuffer().setInt((i + 1) * OFFSET_WIDTH, (i + 1) * listLength); } listVector.setLastSet(outerCount - 1); listVector.setValueCount(outerCount); return listVector; } static StructVector simpleStructVector(BufferAllocator allocator) { final String INT_COL = "struct_int_child"; final String FLT_COL = "struct_flt_child"; StructVector structVector = StructVector.empty(STRUCT_VECTOR_NAME, allocator); final int size = 6; // number of structs NullableStructWriter structWriter = structVector.getWriter(); structVector.addOrGet( INT_COL, FieldType.nullable(Types.MinorType.INT.getType()), IntVector.class); structVector.addOrGet( FLT_COL, FieldType.nullable(Types.MinorType.INT.getType()), IntVector.class); structVector.allocateNew(); IntWriter intWriter = structWriter.integer(INT_COL); Float8Writer float8Writer = structWriter.float8(FLT_COL); for (int i = 0; i < size; i++) { structWriter.setPosition(i); structWriter.start(); intWriter.writeInt(i); float8Writer.writeFloat8(i * .1); structWriter.end(); } structWriter.setValueCount(size); return structVector; } /** Returns a MapVector of longs to doubles. */ static MapVector simpleMapVector(BufferAllocator allocator) { MapVector mapVector = MapVector.empty(BIGINT_INT_MAP_VECTOR_NAME, allocator, false); mapVector.allocateNew(); int count = 5; UnionMapWriter mapWriter = mapVector.getWriter(); for (int i = 0; i < count; i++) { mapWriter.startMap(); for (int j = 0; j < i + 1; j++) { mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(j); mapWriter.value().integer().writeInt(j); mapWriter.endEntry(); } mapWriter.endMap(); } mapWriter.setValueCount(count); return mapVector; } static List decimalVector(BufferAllocator allocator, int rowCount) { List vectors = new ArrayList<>(); vectors.add( new DecimalVector( "decimal_vector", new FieldType(true, new ArrowType.Decimal(38, 10, 128), null), allocator)); vectors.forEach(vec -> generateDecimalData((DecimalVector) vec, rowCount)); return vectors; } static List bitVector(BufferAllocator allocator, int rowCount) { List vectors = new ArrayList<>(); vectors.add(new BitVector("bit_vector", allocator)); vectors.forEach(vec -> GenerateSampleData.generateTestData(vec, rowCount)); return vectors; } /** Returns a UnionVector. */ static UnionVector simpleUnionVector(BufferAllocator allocator) { final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); uInt4Holder.value = 100; uInt4Holder.isSet = 1; UnionVector unionVector = new UnionVector(UNION_VECTOR_NAME, allocator, null, null); unionVector.allocateNew(); // write some data unionVector.setType(0, Types.MinorType.UINT4); unionVector.setSafe(0, uInt4Holder); unionVector.setType(2, Types.MinorType.UINT4); unionVector.setSafe(2, uInt4Holder); unionVector.setValueCount(4); return unionVector; } /** Returns a DenseUnionVector. */ static DenseUnionVector simpleDenseUnionVector(BufferAllocator allocator) { final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); uInt4Holder.value = 100; uInt4Holder.isSet = 1; DenseUnionVector unionVector = new DenseUnionVector(UNION_VECTOR_NAME, allocator, null, null); unionVector.allocateNew(); // write some data byte uint4TypeId = unionVector.registerNewTypeId(Field.nullable("", Types.MinorType.UINT4.getType())); unionVector.setTypeId(0, uint4TypeId); unionVector.setSafe(0, uInt4Holder); unionVector.setTypeId(2, uint4TypeId); unionVector.setSafe(2, uInt4Holder); unionVector.setValueCount(4); return unionVector; } private static void generateDecimalData(DecimalVector vector, int valueCount) { final BigDecimal even = new BigDecimal("0.0543278923"); final BigDecimal odd = new BigDecimal("2.0543278923"); for (int i = 0; i < valueCount; i++) { if (i % 2 == 0) { vector.setSafe(i, even); } else { vector.setSafe(i, odd); } } vector.setValueCount(valueCount); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/test/000077500000000000000000000000001515451615200265125ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/test/util/000077500000000000000000000000001515451615200274675ustar00rootroot00000000000000ArrowTestDataUtil.java000066400000000000000000000032011515451615200336310ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/test/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.test.util; import java.nio.file.Path; import java.nio.file.Paths; import java.util.Objects; /** Utility methods and constants for working with the arrow-testing repo. */ public final class ArrowTestDataUtil { public static final String TEST_DATA_ENV_VAR = "ARROW_TEST_DATA"; public static final String TEST_DATA_PROPERTY = "arrow.test.dataRoot"; public static Path getTestDataRoot() { String path = System.getenv(TEST_DATA_ENV_VAR); if (path == null) { path = System.getProperty(TEST_DATA_PROPERTY); } return Paths.get( Objects.requireNonNull( path, String.format( "Could not find test data path. Set the environment variable %s or the JVM property %s.", TEST_DATA_ENV_VAR, TEST_DATA_PROPERTY))); } private ArrowTestDataUtil() {} } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/testing/000077500000000000000000000000001515451615200272105ustar00rootroot00000000000000RandomDataGenerator.java000066400000000000000000000030551515451615200336600ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/testing/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.testing; import java.util.Random; import java.util.function.Supplier; /** Utility for generating random data. */ public class RandomDataGenerator { static final Random random = new Random(0); public static final Supplier TINY_INT_GENERATOR = () -> (byte) random.nextInt(); public static final Supplier SMALL_INT_GENERATOR = () -> (short) random.nextInt(); public static final Supplier INT_GENERATOR = () -> random.nextInt(); public static final Supplier LONG_GENERATOR = () -> random.nextLong(); public static final Supplier FLOAT_GENERATOR = () -> random.nextFloat(); public static final Supplier DOUBLE_GENERATOR = () -> random.nextDouble(); private RandomDataGenerator() {} } TestValueVectorPopulator.java000066400000000000000000000477031515451615200350140ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/testing/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.testing; import static org.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector; import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.charset.StandardCharsets; 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.DateMilliVector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; 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.IntervalYearVector; 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.TimeStampMicroVector; import org.apache.arrow.vector.TimeStampMilliVector; import org.apache.arrow.vector.TimeStampNanoVector; 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.compare.VectorEqualsVisitor; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.FieldType; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestValueVectorPopulator { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testPopulateBigIntVector() { try (final BigIntVector vector1 = new BigIntVector("vector", allocator); final BigIntVector vector2 = new BigIntVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i); } } vector1.setValueCount(10); setVector(vector2, null, 1L, null, 3L, null, 5L, null, 7L, null, 9L); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateBitVector() { try (final BitVector vector1 = new BitVector("vector", allocator); final BitVector vector2 = new BitVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i > 5 ? 0 : 1); } } vector1.setValueCount(10); setVector(vector2, null, 1, null, 1, null, 1, null, 0, null, 0); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateDateDayVector() { try (final DateDayVector vector1 = new DateDayVector("vector", allocator); final DateDayVector vector2 = new DateDayVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i * 10); } } vector1.setValueCount(10); setVector(vector2, null, 10, null, 30, null, 50, null, 70, null, 90); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateDateMilliVector() { try (final DateMilliVector vector1 = new DateMilliVector("vector", allocator); final DateMilliVector vector2 = new DateMilliVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i * 1000); } } vector1.setValueCount(10); setVector(vector2, null, 1000L, null, 3000L, null, 5000L, null, 7000L, null, 9000L); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateDecimalVector() { try (final DecimalVector vector1 = new DecimalVector("vector", allocator, 10, 3); final DecimalVector vector2 = new DecimalVector("vector", allocator, 10, 3)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i); } } vector1.setValueCount(10); setVector(vector2, null, 1L, null, 3L, null, 5L, null, 7L, null, 9L); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateDurationVector() { final FieldType fieldType = FieldType.nullable(new ArrowType.Duration(TimeUnit.SECOND)); try (final DurationVector vector1 = new DurationVector("vector", fieldType, allocator); final DurationVector vector2 = new DurationVector("vector", fieldType, allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i); } } vector1.setValueCount(10); setVector(vector2, null, 1L, null, 3L, null, 5L, null, 7L, null, 9L); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateFixedSizeBinaryVector() { try (final FixedSizeBinaryVector vector1 = new FixedSizeBinaryVector("vector", allocator, 5); final FixedSizeBinaryVector vector2 = new FixedSizeBinaryVector("vector", allocator, 5)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, ("test" + i).getBytes(StandardCharsets.UTF_8)); } } vector1.setValueCount(10); setVector( vector2, null, "test1".getBytes(StandardCharsets.UTF_8), null, "test3".getBytes(StandardCharsets.UTF_8), null, "test5".getBytes(StandardCharsets.UTF_8), null, "test7".getBytes(StandardCharsets.UTF_8), null, "test9".getBytes(StandardCharsets.UTF_8)); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateFloat4Vector() { try (final Float4Vector vector1 = new Float4Vector("vector", allocator); final Float4Vector vector2 = new Float4Vector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i); } } vector1.setValueCount(10); setVector(vector2, null, 1f, null, 3f, null, 5f, null, 7f, null, 9f); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateFloat8Vector() { try (final Float8Vector vector1 = new Float8Vector("vector", allocator); final Float8Vector vector2 = new Float8Vector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i); } } vector1.setValueCount(10); setVector(vector2, null, 1d, null, 3d, null, 5d, null, 7d, null, 9d); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateIntVector() { try (final IntVector vector1 = new IntVector("vector", allocator); final IntVector vector2 = new IntVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i); } } vector1.setValueCount(10); ValueVectorDataPopulator.setVector(vector2, null, 1, null, 3, null, 5, null, 7, null, 9); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateSmallIntVector() { try (final SmallIntVector vector1 = new SmallIntVector("vector", allocator); final SmallIntVector vector2 = new SmallIntVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i); } } vector1.setValueCount(10); ValueVectorDataPopulator.setVector( vector2, null, (short) 1, null, (short) 3, null, (short) 5, null, (short) 7, null, (short) 9); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateIntervalDayVector() { try (final IntervalYearVector vector1 = new IntervalYearVector("vector", allocator); final IntervalYearVector vector2 = new IntervalYearVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i); } } vector1.setValueCount(10); ValueVectorDataPopulator.setVector(vector2, null, 1, null, 3, null, 5, null, 7, null, 9); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateTimeMicroVector() { try (final TimeMicroVector vector1 = new TimeMicroVector("vector", allocator); final TimeMicroVector vector2 = new TimeMicroVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i * 10000); } } vector1.setValueCount(10); setVector(vector2, null, 10000L, null, 30000L, null, 50000L, null, 70000L, null, 90000L); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateTimeMilliVector() { try (final TimeMilliVector vector1 = new TimeMilliVector("vector", allocator); final TimeMilliVector vector2 = new TimeMilliVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i * 100); } } vector1.setValueCount(10); setVector(vector2, null, 100, null, 300, null, 500, null, 700, null, 900); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateTimeNanoVector() { try (final TimeNanoVector vector1 = new TimeNanoVector("vector", allocator); final TimeNanoVector vector2 = new TimeNanoVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i * 10000); } } vector1.setValueCount(10); setVector(vector2, null, 10000L, null, 30000L, null, 50000L, null, 70000L, null, 90000L); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateTimeSecVector() { try (final TimeSecVector vector1 = new TimeSecVector("vector", allocator); final TimeSecVector vector2 = new TimeSecVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i * 100); } } vector1.setValueCount(10); setVector(vector2, null, 100, null, 300, null, 500, null, 700, null, 900); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateTimeStampMicroVector() { try (final TimeStampMicroVector vector1 = new TimeStampMicroVector("vector", allocator); final TimeStampMicroVector vector2 = new TimeStampMicroVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i * 10000); } } vector1.setValueCount(10); setVector(vector2, null, 10000L, null, 30000L, null, 50000L, null, 70000L, null, 90000L); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateTimeStampMilliVector() { try (final TimeStampMilliVector vector1 = new TimeStampMilliVector("vector", allocator); final TimeStampMilliVector vector2 = new TimeStampMilliVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i * 10000); } } vector1.setValueCount(10); setVector(vector2, null, 10000L, null, 30000L, null, 50000L, null, 70000L, null, 90000L); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateTimeStampNanoVector() { try (final TimeStampNanoVector vector1 = new TimeStampNanoVector("vector", allocator); final TimeStampNanoVector vector2 = new TimeStampNanoVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i * 10000); } } vector1.setValueCount(10); setVector(vector2, null, 10000L, null, 30000L, null, 50000L, null, 70000L, null, 90000L); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateTimeStampSecVector() { try (final TimeStampSecVector vector1 = new TimeStampSecVector("vector", allocator); final TimeStampSecVector vector2 = new TimeStampSecVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i * 100); } } vector1.setValueCount(10); setVector(vector2, null, 100L, null, 300L, null, 500L, null, 700L, null, 900L); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateTinyIntVector() { try (final TinyIntVector vector1 = new TinyIntVector("vector", allocator); final TinyIntVector vector2 = new TinyIntVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i); } } vector1.setValueCount(10); setVector( vector2, null, (byte) 1, null, (byte) 3, null, (byte) 5, null, (byte) 7, null, (byte) 9); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateUInt1Vector() { try (final UInt1Vector vector1 = new UInt1Vector("vector", allocator); final UInt1Vector vector2 = new UInt1Vector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i); } } vector1.setValueCount(10); setVector( vector2, null, (byte) 1, null, (byte) 3, null, (byte) 5, null, (byte) 7, null, (byte) 9); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateUInt2Vector() { try (final UInt2Vector vector1 = new UInt2Vector("vector", allocator); final UInt2Vector vector2 = new UInt2Vector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i); } } vector1.setValueCount(10); setVector( vector2, null, (char) 1, null, (char) 3, null, (char) 5, null, (char) 7, null, (char) 9); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateUInt4Vector() { try (final UInt4Vector vector1 = new UInt4Vector("vector", allocator); final UInt4Vector vector2 = new UInt4Vector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i); } } vector1.setValueCount(10); setVector(vector2, null, 1, null, 3, null, 5, null, 7, null, 9); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateUInt8Vector() { try (final UInt8Vector vector1 = new UInt8Vector("vector", allocator); final UInt8Vector vector2 = new UInt8Vector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, i); } } vector1.setValueCount(10); setVector(vector2, null, 1L, null, 3L, null, 5L, null, 7L, null, 9L); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateVarBinaryVector() { try (final VarBinaryVector vector1 = new VarBinaryVector("vector", allocator); final VarBinaryVector vector2 = new VarBinaryVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, ("test" + i).getBytes(StandardCharsets.UTF_8)); } } vector1.setValueCount(10); setVector( vector2, null, "test1".getBytes(StandardCharsets.UTF_8), null, "test3".getBytes(StandardCharsets.UTF_8), null, "test5".getBytes(StandardCharsets.UTF_8), null, "test7".getBytes(StandardCharsets.UTF_8), null, "test9".getBytes(StandardCharsets.UTF_8)); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } @Test public void testPopulateVarCharVector() { try (final VarCharVector vector1 = new VarCharVector("vector", allocator); final VarCharVector vector2 = new VarCharVector("vector", allocator)) { vector1.allocateNew(10); for (int i = 0; i < 10; i++) { if (i % 2 == 0) { vector1.setNull(i); } else { vector1.set(i, ("test" + i).getBytes(StandardCharsets.UTF_8)); } } vector1.setValueCount(10); setVector(vector2, null, "test1", null, "test3", null, "test5", null, "test7", null, "test9"); assertTrue(VectorEqualsVisitor.vectorEquals(vector1, vector2)); } } } ValueVectorDataPopulator.java000066400000000000000000000641051515451615200347410ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/testing/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.testing; import static org.junit.jupiter.api.Assertions.assertEquals; import java.math.BigDecimal; import java.nio.charset.StandardCharsets; import java.util.List; import java.util.Map; import java.util.Map.Entry; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.BitVectorHelper; 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.DurationVector; 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.IntervalDayVector; import org.apache.arrow.vector.IntervalYearVector; import org.apache.arrow.vector.LargeVarCharVector; 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.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; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.LargeListVector; 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.holders.IntervalDayHolder; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; /** Utility for populating {@link org.apache.arrow.vector.ValueVector}. */ public class ValueVectorDataPopulator { private ValueVectorDataPopulator() {} /** Populate values for BigIntVector. */ public static void setVector(BigIntVector vector, Long... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for BitVector. */ public static void setVector(BitVector vector, Integer... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** * Populate values for DateDayVector. * * @param values numbers of days since UNIX epoch */ public static void setVector(DateDayVector vector, Integer... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** * Populate values for DateMilliVector. * * @param values numbers of milliseconds since UNIX epoch */ public static void setVector(DateMilliVector vector, Long... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for DecimalVector. */ public static void setVector(DecimalVector vector, Long... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for Decimal256Vector. */ public static void setVector(Decimal256Vector vector, Long... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for Decimal256Vector. */ public static void setVector(Decimal256Vector vector, BigDecimal... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** * Populate values for DurationVector. * * @param values values of elapsed time in either seconds, milliseconds, microseconds or * nanoseconds. */ public static void setVector(DurationVector vector, Long... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for FixedSizeBinaryVector. */ public static void setVector(FixedSizeBinaryVector vector, byte[]... values) { final int length = values.length; vector.allocateNewSafe(); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for Float2Vector. */ public static void setVector(Float2Vector vector, Float... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.setWithPossibleTruncate(i, values[i]); } } vector.setValueCount(length); } /** Populate values for Float4Vector. */ public static void setVector(Float4Vector vector, Float... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for Float8Vector. */ public static void setVector(Float8Vector vector, Double... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for IntVector. */ public static void setVector(IntVector vector, Integer... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** * Populate values for IntervalDayVector. * * @param values holders witch holds days and milliseconds values which represents interval in SQL * style. */ public static void setVector(IntervalDayVector vector, IntervalDayHolder... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i].days, values[i].milliseconds); } } vector.setValueCount(length); } /** * Populate values for IntervalYearVector. * * @param values total month intervals in SQL style. */ public static void setVector(IntervalYearVector vector, Integer... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for SmallIntVector. */ public static void setVector(SmallIntVector vector, Short... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** * Populate values for TimeMicroVector. * * @param values numbers of microseconds since UNIX epoch */ public static void setVector(TimeMicroVector vector, Long... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** * Populate values for TimeMicroVector. * * @param values numbers of milliseconds since UNIX epoch */ public static void setVector(TimeMilliVector vector, Integer... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** * Populate values for TimeNanoVector. * * @param values numbers of nanoseconds since UNIX epoch */ public static void setVector(TimeNanoVector vector, Long... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** * Populate values for TimeSecVector. * * @param values numbers of seconds since UNIX epoch */ public static void setVector(TimeSecVector vector, Integer... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** * Populate values for TimeStampMicroTZVector. * * @param values numbers of microseconds since UNIX epoch */ public static void setVector(TimeStampMicroTZVector vector, Long... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** * Populate values for TimeStampMicroVector. * * @param values numbers of microseconds since UNIX epoch */ public static void setVector(TimeStampMicroVector vector, Long... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** * Populate values for TimeStampMilliTZVector. * * @param values numbers of milliseconds since UNIX epoch */ public static void setVector(TimeStampMilliTZVector vector, Long... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** * Populate values for TimeStampMilliVector. * * @param values numbers of milliseconds since UNIX epoch */ public static void setVector(TimeStampMilliVector vector, Long... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** * Populate values for TimeStampNanoTZVector. * * @param values numbers of nanoseconds since UNIX epoch */ public static void setVector(TimeStampNanoTZVector vector, Long... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** * Populate values for TimeStampNanoVector. * * @param values numbers of nanoseconds since UNIX epoch */ public static void setVector(TimeStampNanoVector vector, Long... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** * Populate values for TimeStampSecTZVector. * * @param values numbers of seconds since UNIX epoch */ public static void setVector(TimeStampSecTZVector vector, Long... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** * Populate values for TimeStampSecVector. * * @param values numbers of seconds since UNIX epoch */ public static void setVector(TimeStampSecVector vector, Long... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for TinyIntVector. */ public static void setVector(TinyIntVector vector, Byte... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for UInt1Vector. */ public static void setVector(UInt1Vector vector, Byte... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for UInt2Vector. */ public static void setVector(UInt2Vector vector, Character... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for UInt4Vector. */ public static void setVector(UInt4Vector vector, Integer... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for UInt8Vector. */ public static void setVector(UInt8Vector vector, Long... values) { final int length = values.length; vector.allocateNew(length); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for VarBinaryVector. */ public static void setVector(VarBinaryVector vector, byte[]... values) { final int length = values.length; vector.allocateNewSafe(); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for VarCharVector. */ public static void setVector(VarCharVector vector, byte[]... values) { final int length = values.length; vector.allocateNewSafe(); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } public static void setVector(VariableWidthFieldVector vector, byte[]... values) { final int length = values.length; vector.allocateNewSafe(); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for LargeVarCharVector. */ public static void setVector(LargeVarCharVector vector, byte[]... values) { final int length = values.length; vector.allocateNewSafe(); for (int i = 0; i < length; i++) { if (values[i] != null) { vector.set(i, values[i]); } } vector.setValueCount(length); } /** Populate values for VarCharVector. */ public static void setVector(VarCharVector 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 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; 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 {@link ListVector}. */ public static void setVector(ListVector vector, List... values) { vector.allocateNewSafe(); Types.MinorType type = Types.MinorType.INT; vector.addOrGetVector(FieldType.nullable(type.getType())); IntVector dataVector = (IntVector) vector.getDataVector(); dataVector.allocateNew(); // set underlying vectors int curPos = 0; vector.getOffsetBuffer().setInt(0, curPos); for (int i = 0; i < values.length; i++) { if (values[i] == null) { BitVectorHelper.unsetBit(vector.getValidityBuffer(), i); } else { BitVectorHelper.setBit(vector.getValidityBuffer(), i); for (int value : values[i]) { dataVector.setSafe(curPos, value); curPos += 1; } } vector.getOffsetBuffer().setInt((i + 1) * BaseRepeatedValueVector.OFFSET_WIDTH, curPos); } dataVector.setValueCount(curPos); vector.setLastSet(values.length - 1); vector.setValueCount(values.length); } /** Populate values for {@link LargeListVector}. */ public static void setVector(LargeListVector vector, List... values) { vector.allocateNewSafe(); Types.MinorType type = Types.MinorType.INT; vector.addOrGetVector(FieldType.nullable(type.getType())); IntVector dataVector = (IntVector) vector.getDataVector(); dataVector.allocateNew(); // set underlying vectors int curPos = 0; vector.getOffsetBuffer().setLong(0, curPos); for (int i = 0; i < values.length; i++) { if (values[i] == null) { BitVectorHelper.unsetBit(vector.getValidityBuffer(), i); } else { BitVectorHelper.setBit(vector.getValidityBuffer(), i); for (int value : values[i]) { dataVector.setSafe(curPos, value); curPos += 1; } } vector.getOffsetBuffer().setLong((long) (i + 1) * LargeListVector.OFFSET_WIDTH, curPos); } dataVector.setValueCount(curPos); vector.setLastSet(values.length - 1); vector.setValueCount(values.length); } /** Populate values for {@link FixedSizeListVector}. */ public static void setVector(FixedSizeListVector vector, List... values) { vector.allocateNewSafe(); for (int i = 0; i < values.length; i++) { if (values[i] != null) { assertEquals(vector.getListSize(), values[i].size()); } } Types.MinorType type = Types.MinorType.INT; vector.addOrGetVector(FieldType.nullable(type.getType())); IntVector dataVector = (IntVector) vector.getDataVector(); dataVector.allocateNew(); // set underlying vectors int curPos = 0; for (int i = 0; i < values.length; i++) { if (values[i] == null) { BitVectorHelper.unsetBit(vector.getValidityBuffer(), i); } else { BitVectorHelper.setBit(vector.getValidityBuffer(), i); for (int value : values[i]) { dataVector.setSafe(curPos, value); curPos += 1; } } } dataVector.setValueCount(curPos); vector.setValueCount(values.length); } /** Populate values for {@link StructVector}. */ public static void setVector(StructVector vector, Map> values) { vector.allocateNewSafe(); int valueCount = 0; for (final Entry> entry : values.entrySet()) { // Add the child final IntVector child = vector.addOrGet( entry.getKey(), FieldType.nullable(MinorType.INT.getType()), IntVector.class); // Write the values to the child child.allocateNew(); final List v = entry.getValue(); for (int i = 0; i < v.size(); i++) { if (v.get(i) != null) { child.set(i, v.get(i)); vector.setIndexDefined(i); } else { child.setNull(i); } } valueCount = Math.max(valueCount, v.size()); } vector.setValueCount(valueCount); } /** Populate values for {@link ListViewVector}. */ public static void setVector(ListViewVector vector, List... values) { vector.allocateNewSafe(); Types.MinorType type = Types.MinorType.INT; vector.addOrGetVector(FieldType.nullable(type.getType())); IntVector dataVector = (IntVector) vector.getDataVector(); dataVector.allocateNew(); // set underlying vectors int curPos = 0; for (int i = 0; i < values.length; i++) { vector.getOffsetBuffer().setInt((long) i * BaseRepeatedValueViewVector.OFFSET_WIDTH, curPos); if (values[i] == null) { BitVectorHelper.unsetBit(vector.getValidityBuffer(), i); } else { BitVectorHelper.setBit(vector.getValidityBuffer(), i); for (int value : values[i]) { dataVector.setSafe(curPos, value); curPos += 1; } } vector .getSizeBuffer() .setInt((long) i * BaseRepeatedValueViewVector.SIZE_WIDTH, values[i].size()); } dataVector.setValueCount(curPos); vector.setValueCount(values.length); } /** Populate values for {@link ListViewVector}. */ public static void setVector(LargeListViewVector vector, List... values) { vector.allocateNewSafe(); Types.MinorType type = Types.MinorType.INT; vector.addOrGetVector(FieldType.nullable(type.getType())); IntVector dataVector = (IntVector) vector.getDataVector(); dataVector.allocateNew(); // set underlying vectors int curPos = 0; for (int i = 0; i < values.length; i++) { vector .getOffsetBuffer() .setInt((long) i * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH, curPos); if (values[i] == null) { BitVectorHelper.unsetBit(vector.getValidityBuffer(), i); } else { BitVectorHelper.setBit(vector.getValidityBuffer(), i); for (int value : values[i]) { dataVector.setSafe(curPos, value); curPos += 1; } } vector .getSizeBuffer() .setInt((long) i * BaseRepeatedValueViewVector.SIZE_WIDTH, values[i].size()); } dataVector.setValueCount(curPos); vector.setValueCount(values.length); } public static void setVector( RunEndEncodedVector vector, List runEnds, List values) { int runCount = runEnds.size(); assert runCount == values.size(); final FieldType valueType = FieldType.notNullable(MinorType.INT.getType()); final FieldType runEndType = FieldType.notNullable(Types.MinorType.INT.getType()); final Field valueField = new Field("value", valueType, null); final Field runEndField = new Field("ree", runEndType, null); vector.initializeChildrenFromFields(List.of(runEndField, valueField)); IntVector runEndsVector = (IntVector) vector.getRunEndsVector(); runEndsVector.setValueCount(runCount); for (int i = 0; i < runCount; i++) { if (runEnds.get(i) == null) { runEndsVector.setNull(i); } else { runEndsVector.set(i, runEnds.get(i)); } } IntVector valuesVector = (IntVector) vector.getValuesVector(); valuesVector.setValueCount(runCount); for (int i = 0; i < runCount; i++) { if (runEnds.get(i) == null) { valuesVector.setNull(i); } else { valuesVector.set(i, values.get(i)); } } if (runCount > 0) { vector.setValueCount(runEnds.get(runCount - 1)); } else { vector.setValueCount(0); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/types/000077500000000000000000000000001515451615200266775ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/types/pojo/000077500000000000000000000000001515451615200276465ustar00rootroot00000000000000TestExtensionType.java000066400000000000000000000364441515451615200341230ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/types/pojo/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.arrow.vector.types.pojo; import 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; 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.ByteBuffer; 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 java.util.UUID; 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.ExtensionTypeVector; 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; import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; import org.apache.arrow.vector.util.VectorBatchAppender; import org.apache.arrow.vector.validate.ValidateVectorVisitor; import org.junit.jupiter.api.Test; public class TestExtensionType { /** Test that a custom UUID type can be round-tripped through a temporary file. */ @Test public void roundtripUuid() throws IOException { ensureRegistered(UuidType.INSTANCE); final Schema schema = 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(); UUID u2 = UUID.randomUUID(); UuidVector vector = (UuidVector) root.getVector("a"); vector.setValueCount(2); vector.set(0, u1); vector.set(1, u2); root.setRowCount(2); final File file = File.createTempFile("uuidtest", ".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 UuidType expectedType = UuidType.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 that a custom UUID type can be read as its underlying type. */ @Test public void readUnderlyingType() throws IOException { ensureRegistered(UuidType.INSTANCE); final Schema schema = 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(); UUID u2 = UUID.randomUUID(); UuidVector vector = (UuidVector) root.getVector("a"); vector.setValueCount(2); vector.set(0, u1); vector.set(1, u2); root.setRowCount(2); final File file = File.createTempFile("uuidtest", ".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(UuidType.INSTANCE); 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(1, readerRoot.getSchema().getFields().size()); assertEquals("a", readerRoot.getSchema().getFields().get(0).getName()); assertTrue( readerRoot.getSchema().getFields().get(0).getType() instanceof ArrowType.FixedSizeBinary); assertEquals( 16, ((ArrowType.FixedSizeBinary) readerRoot.getSchema().getFields().get(0).getType()) .getByteWidth()); final Field field = readerRoot.getSchema().getFields().get(0); final UuidType expectedType = UuidType.INSTANCE; assertEquals( field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME), expectedType.extensionName()); assertEquals( field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_METADATA), expectedType.serialize()); final FixedSizeBinaryVector deserialized = (FixedSizeBinaryVector) 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)) { final UUID uuid = vector.getObject(i); final ByteBuffer bb = ByteBuffer.allocate(16); bb.putLong(uuid.getMostSignificantBits()); bb.putLong(uuid.getLeastSignificantBits()); assertArrayEquals(bb.array(), deserialized.get(i)); } } } } } @Test public void testNullCheck() { NullPointerException e = assertThrows( NullPointerException.class, () -> { try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); final ExtensionTypeVector vector = new UuidVector("uuid", allocator, null)) { vector.getField(); vector.allocateNewSafe(); } }); assertTrue(e.getMessage().contains("underlyingVector cannot be null.")); } /** Test that a custom Location type can be round-tripped through a temporary file. */ @Test public void roundtripLocation() throws IOException { ExtensionTypeRegistry.register(new LocationType()); final Schema schema = new Schema(Collections.singletonList(Field.nullable("location", new LocationType()))); try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { LocationVector vector = (LocationVector) root.getVector("location"); vector.allocateNew(); vector.set(0, 34.073814f, -118.240784f); vector.set(2, 37.768056f, -122.3875f); vector.set(3, 40.739716f, -73.840782f); vector.setValueCount(4); root.setRowCount(4); final File file = File.createTempFile("locationtest", ".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 LocationType expectedType = new LocationType(); 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); assertTrue(deserialized instanceof LocationVector); assertEquals("location", deserialized.getName()); StructVector deserStruct = (StructVector) deserialized.getUnderlyingVector(); assertNotNull(deserStruct.getChild("Latitude")); assertNotNull(deserStruct.getChild("Longitude")); 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 testVectorCompare() { UuidType uuidType = UuidType.INSTANCE; ExtensionTypeRegistry.register(uuidType); try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); UuidVector a1 = (UuidVector) uuidType.getNewVector("a", FieldType.nullable(uuidType), allocator); UuidVector a2 = (UuidVector) uuidType.getNewVector("a", FieldType.nullable(uuidType), allocator); UuidVector bb = (UuidVector) uuidType.getNewVector("a", FieldType.nullable(uuidType), allocator)) { UUID u1 = UUID.randomUUID(); UUID u2 = UUID.randomUUID(); // Test out type and vector validation visitors for an ExtensionTypeVector ValidateVectorVisitor validateVisitor = new ValidateVectorVisitor(); validateVisitor.visit(a1, null); a1.setValueCount(2); a1.set(0, u1); a1.set(1, u2); a2.setValueCount(2); a2.set(0, u1); a2.set(1, u2); bb.setValueCount(2); bb.set(0, u2); bb.set(1, u1); 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)); // Test out vector appender VectorBatchAppender.batchAppend(a1, a2, bb); assertEquals(6, a1.getValueCount()); validateVisitor.visit(a1, null); } } static class LocationType extends ExtensionType { @Override public ArrowType storageType() { return Struct.INSTANCE; } @Override public String extensionName() { return "location"; } @Override public boolean extensionEquals(ExtensionType other) { return other instanceof LocationType; } @Override public ArrowType deserialize(ArrowType storageType, String serializedData) { if (!storageType.equals(storageType())) { throw new UnsupportedOperationException( "Cannot construct LocationType from underlying type " + storageType); } return new LocationType(); } @Override public String serialize() { return ""; } @Override 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 implements ValueIterableVector> { private static StructVector buildUnderlyingVector(String name, BufferAllocator allocator) { final StructVector underlyingVector = new StructVector(name, allocator, FieldType.nullable(ArrowType.Struct.INSTANCE), null); underlyingVector.addOrGet( "Latitude", FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), Float4Vector.class); underlyingVector.addOrGet( "Longitude", FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), Float4Vector.class); return underlyingVector; } public LocationVector(String name, BufferAllocator allocator) { super(name, allocator, buildUnderlyingVector(name, allocator)); } @Override public int hashCode(int index) { return hashCode(index, null); } @Override public int hashCode(int index, ArrowBufHasher hasher) { return getUnderlyingVector().hashCode(index, hasher); } @Override public java.util.Map getObject(int index) { return getUnderlyingVector().getObject(index); } public void set(int index, float latitude, float longitude) { getUnderlyingVector().getChild("Latitude", Float4Vector.class).set(index, latitude); getUnderlyingVector().getChild("Longitude", Float4Vector.class).set(index, longitude); getUnderlyingVector().setIndexDefined(index); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/types/pojo/TestField.java000066400000000000000000000045561515451615200324060ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.arrow.vector.types.pojo; import static org.apache.arrow.vector.types.pojo.Schema.METADATA_KEY; import static org.apache.arrow.vector.types.pojo.Schema.METADATA_VALUE; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.IOException; import java.util.Collections; import java.util.HashMap; import java.util.Map; import org.apache.arrow.vector.types.pojo.ArrowType.Int; import org.junit.jupiter.api.Test; public class TestField { private static Field field( String name, boolean nullable, ArrowType type, Map metadata) { return new Field(name, new FieldType(nullable, type, null, metadata), Collections.emptyList()); } @Test public void testMetadata() throws IOException { Map metadata = new HashMap<>(1); metadata.put("testKey", "testValue"); Schema schema = new Schema(Collections.singletonList(field("a", false, new Int(8, true), metadata))); String json = schema.toJson(); Schema actual = Schema.fromJSON(json); jsonContains( json, "\"" + METADATA_KEY + "\" : \"testKey\"", "\"" + METADATA_VALUE + "\" : \"testValue\""); Map actualMetadata = actual.getFields().get(0).getMetadata(); assertEquals(1, actualMetadata.size()); assertEquals("testValue", actualMetadata.get("testKey")); } private void jsonContains(String json, String... strings) { for (String string : strings) { assertTrue(json.contains(string), json + " contains " + string); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/types/pojo/TestSchema.java000066400000000000000000000306511515451615200325560ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.arrow.vector.types.pojo; import static java.util.Arrays.asList; import static org.apache.arrow.vector.types.pojo.Schema.METADATA_KEY; import static org.apache.arrow.vector.types.pojo.Schema.METADATA_VALUE; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.IOException; import java.nio.ByteBuffer; import java.util.HashMap; import java.util.Map; import org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.IntervalUnit; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.UnionMode; import org.apache.arrow.vector.types.pojo.ArrowType.Binary; import org.apache.arrow.vector.types.pojo.ArrowType.Bool; import org.apache.arrow.vector.types.pojo.ArrowType.Date; import org.apache.arrow.vector.types.pojo.ArrowType.Decimal; import org.apache.arrow.vector.types.pojo.ArrowType.Duration; import org.apache.arrow.vector.types.pojo.ArrowType.FixedSizeBinary; import org.apache.arrow.vector.types.pojo.ArrowType.FloatingPoint; import org.apache.arrow.vector.types.pojo.ArrowType.Int; import org.apache.arrow.vector.types.pojo.ArrowType.Interval; import org.apache.arrow.vector.types.pojo.ArrowType.List; import org.apache.arrow.vector.types.pojo.ArrowType.Null; import org.apache.arrow.vector.types.pojo.ArrowType.Struct; import org.apache.arrow.vector.types.pojo.ArrowType.Time; import org.apache.arrow.vector.types.pojo.ArrowType.Timestamp; import org.apache.arrow.vector.types.pojo.ArrowType.Union; import org.apache.arrow.vector.types.pojo.ArrowType.Utf8; import org.junit.jupiter.api.Test; public class TestSchema { private static Field field(String name, boolean nullable, ArrowType type, Field... children) { return new Field(name, new FieldType(nullable, type, null, null), asList(children)); } private static Field field(String name, ArrowType type, Field... children) { return field(name, true, type, children); } @Test public void testComplex() throws IOException { Schema schema = new Schema( asList( field("a", false, new Int(8, true)), field("b", new Struct(), field("c", new Int(16, true)), field("d", new Utf8())), field("e", new List(), field(null, new Date(DateUnit.MILLISECOND))), field("f", new FloatingPoint(FloatingPointPrecision.SINGLE)), field("g", new Timestamp(TimeUnit.MILLISECOND, "UTC")), field("h", new Timestamp(TimeUnit.MICROSECOND, null)), field("i", new Interval(IntervalUnit.DAY_TIME)), field("j", new ArrowType.Duration(TimeUnit.SECOND)))); roundTrip(schema); assertEquals( "Schema, e: List, " + "f: FloatingPoint(SINGLE), g: Timestamp(MILLISECOND, UTC), h: Timestamp(MICROSECOND, null), " + "i: Interval(DAY_TIME), j: Duration(SECOND)>", schema.toString()); } @Test public void testAll() throws IOException { Schema schema = new Schema( asList( field("a", false, new Null()), field("b", new Struct(), field("ba", new Null())), field("c", new List(), field("ca", new Null())), field( "d", new Union(UnionMode.Sparse, new int[] {1, 2, 3}), field("da", new Null())), field("e", new Int(8, true)), field("f", new FloatingPoint(FloatingPointPrecision.SINGLE)), field("g", new Utf8()), field("h", new Binary()), field("i", new Bool()), field("j", new Decimal(5, 5, 128)), field("k", new Date(DateUnit.DAY)), field("l", new Date(DateUnit.MILLISECOND)), field("m", new Time(TimeUnit.SECOND, 32)), field("n", new Time(TimeUnit.MILLISECOND, 32)), field("o", new Time(TimeUnit.MICROSECOND, 64)), field("p", new Time(TimeUnit.NANOSECOND, 64)), field("q", new Timestamp(TimeUnit.MILLISECOND, "UTC")), field("r", new Timestamp(TimeUnit.MICROSECOND, null)), field("s", new Interval(IntervalUnit.DAY_TIME)), field("t", new FixedSizeBinary(100)), field("u", new Duration(TimeUnit.SECOND)), field("v", new Duration(TimeUnit.MICROSECOND)))); roundTrip(schema); } @Test public void testUnion() throws IOException { Schema schema = new Schema( asList( field( "d", new Union(UnionMode.Sparse, new int[] {1, 2, 3}), field("da", new Null())))); roundTrip(schema); contains(schema, "Sparse"); } @Test public void testDate() throws IOException { Schema schema = new Schema( asList(field("a", new Date(DateUnit.DAY)), field("b", new Date(DateUnit.MILLISECOND)))); roundTrip(schema); assertEquals("Schema", schema.toString()); } @Test public void testTime() throws IOException { Schema schema = new Schema( asList( field("a", new Time(TimeUnit.SECOND, 32)), field("b", new Time(TimeUnit.MILLISECOND, 32)), field("c", new Time(TimeUnit.MICROSECOND, 64)), field("d", new Time(TimeUnit.NANOSECOND, 64)))); roundTrip(schema); assertEquals( "Schema", schema.toString()); } @Test public void testTS() throws IOException { Schema schema = new Schema( asList( field("a", new Timestamp(TimeUnit.SECOND, "UTC")), field("b", new Timestamp(TimeUnit.MILLISECOND, "UTC")), field("c", new Timestamp(TimeUnit.MICROSECOND, "UTC")), field("d", new Timestamp(TimeUnit.NANOSECOND, "UTC")), field("e", new Timestamp(TimeUnit.SECOND, null)), field("f", new Timestamp(TimeUnit.MILLISECOND, null)), field("g", new Timestamp(TimeUnit.MICROSECOND, null)), field("h", new Timestamp(TimeUnit.NANOSECOND, null)))); roundTrip(schema); assertEquals( "Schema", schema.toString()); } @Test public void testInterval() throws IOException { Schema schema = new Schema( asList( field("a", new Interval(IntervalUnit.YEAR_MONTH)), field("b", new Interval(IntervalUnit.DAY_TIME)))); roundTrip(schema); contains(schema, "YEAR_MONTH", "DAY_TIME"); } @Test public void testRoundTripDurationInterval() throws IOException { Schema schema = new Schema( asList( field("a", new Duration(TimeUnit.SECOND)), field("b", new Duration(TimeUnit.MILLISECOND)), field("c", new Duration(TimeUnit.MICROSECOND)), field("d", new Duration(TimeUnit.NANOSECOND)))); roundTrip(schema); contains(schema, "SECOND", "MILLI", "MICRO", "NANO"); } @Test public void testFP() throws IOException { Schema schema = new Schema( asList( field("a", new FloatingPoint(FloatingPointPrecision.HALF)), field("b", new FloatingPoint(FloatingPointPrecision.SINGLE)), field("c", new FloatingPoint(FloatingPointPrecision.DOUBLE)))); roundTrip(schema); contains(schema, "HALF", "SINGLE", "DOUBLE"); } @Test public void testMetadata() throws IOException { Map metadata = new HashMap<>(1); metadata.put("testKey", "testValue"); java.util.List fields = asList( field("a", false, new Int(8, true)), field("b", new Struct(), field("c", new Int(16, true)), field("d", new Utf8())), field("e", new List(), field(null, new Date(DateUnit.MILLISECOND)))); Schema schema = new Schema(fields, metadata); roundTrip(schema); contains( schema, "\"" + METADATA_KEY + "\" : \"testKey\"", "\"" + METADATA_VALUE + "\" : \"testValue\""); } @Test public void testMessageSerialization() { Schema schema = new Schema( asList( field("a", false, new Null()), field("b", new Struct(), field("ba", new Null())), field("c", new List(), field("ca", new Null())), field( "d", new Union(UnionMode.Sparse, new int[] {1, 2, 3}), field("da", new Null())), field("e", new Int(8, true)), field("f", new FloatingPoint(FloatingPointPrecision.SINGLE)), field("g", new Utf8()), field("h", new Binary()), field("i", new Bool()), field("j", new Decimal(5, 5, 128)), field("k", new Date(DateUnit.DAY)), field("l", new Date(DateUnit.MILLISECOND)), field("m", new Time(TimeUnit.SECOND, 32)), field("n", new Time(TimeUnit.MILLISECOND, 32)), field("o", new Time(TimeUnit.MICROSECOND, 64)), field("p", new Time(TimeUnit.NANOSECOND, 64)), field("q", new Timestamp(TimeUnit.MILLISECOND, "UTC")), field("r", new Timestamp(TimeUnit.MICROSECOND, null)), field("s", new Interval(IntervalUnit.DAY_TIME)), field("t", new FixedSizeBinary(100)), field("u", new Duration(TimeUnit.SECOND)), field("v", new Duration(TimeUnit.MICROSECOND)))); roundTripMessage(schema); } private void roundTrip(Schema schema) throws IOException { String json = schema.toJson(); Schema actual = Schema.fromJSON(json); assertEquals(schema.toJson(), actual.toJson()); assertEquals(schema, actual); validateFieldsHashcode(schema.getFields(), actual.getFields()); assertEquals(schema.hashCode(), actual.hashCode()); } private void roundTripMessage(Schema schema) { byte[] bytes = schema.serializeAsMessage(); Schema actual = Schema.deserializeMessage(ByteBuffer.wrap(bytes)); assertEquals(schema.toJson(), actual.toJson()); assertEquals(schema, actual); validateFieldsHashcode(schema.getFields(), actual.getFields()); assertEquals(schema.hashCode(), actual.hashCode()); } private void validateFieldsHashcode( java.util.List schemaFields, java.util.List actualFields) { assertEquals(schemaFields.size(), actualFields.size()); if (schemaFields.size() == 0) { return; } for (int i = 0; i < schemaFields.size(); i++) { Field schemaField = schemaFields.get(i); Field actualField = actualFields.get(i); validateFieldsHashcode(schemaField.getChildren(), actualField.getChildren()); validateHashCode(schemaField.getType(), actualField.getType()); validateHashCode(schemaField, actualField); } } private void validateHashCode(Object o1, Object o2) { assertEquals(o1, o2); assertEquals(o1.hashCode(), o2.hashCode(), o1 + " == " + o2); } private void contains(Schema schema, String... s) { String json = schema.toJson(); for (String string : s) { assertTrue(json.contains(string), json + " contains " + string); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/util/000077500000000000000000000000001515451615200265105ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/util/DecimalUtilityTest.java000066400000000000000000000130241515451615200331350ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.junit.jupiter.api.Assertions.assertEquals; import java.math.BigDecimal; import java.math.BigInteger; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.Test; public class DecimalUtilityTest { private static final BigInteger[] MAX_BIG_INT = new BigInteger[] { BigInteger.valueOf(10).pow(38).subtract(java.math.BigInteger.ONE), java.math.BigInteger.valueOf(10).pow(76) }; private static final BigInteger[] MIN_BIG_INT = new BigInteger[] { MAX_BIG_INT[0].multiply(BigInteger.valueOf(-1)), MAX_BIG_INT[1].multiply(BigInteger.valueOf(-1)) }; @Test public void testSetLongInDecimalArrowBuf() { int[] byteLengths = new int[] {16, 32}; for (int x = 0; x < 2; x++) { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(byteLengths[x]); ) { int[] intValues = new int[] {Integer.MAX_VALUE, Integer.MIN_VALUE, 0}; for (int val : intValues) { buf.clear(); DecimalUtility.writeLongToArrowBuf((long) val, buf, 0, byteLengths[x]); BigDecimal actual = DecimalUtility.getBigDecimalFromArrowBuf(buf, 0, 0, byteLengths[x]); BigDecimal expected = BigDecimal.valueOf(val); assertEquals(expected, actual); } } } } @Test public void testSetByteArrayInDecimalArrowBuf() { int[] byteLengths = new int[] {16, 32}; for (int x = 0; x < 2; x++) { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(byteLengths[x]); ) { int[] intValues = new int[] {Integer.MAX_VALUE, Integer.MIN_VALUE, 0}; for (int val : intValues) { buf.clear(); DecimalUtility.writeByteArrayToArrowBuf( BigInteger.valueOf(val).toByteArray(), buf, 0, byteLengths[x]); BigDecimal actual = DecimalUtility.getBigDecimalFromArrowBuf(buf, 0, 0, byteLengths[x]); BigDecimal expected = BigDecimal.valueOf(val); assertEquals(expected, actual); } long[] longValues = new long[] {Long.MIN_VALUE, 0, Long.MAX_VALUE}; for (long val : longValues) { buf.clear(); DecimalUtility.writeByteArrayToArrowBuf( BigInteger.valueOf(val).toByteArray(), buf, 0, byteLengths[x]); BigDecimal actual = DecimalUtility.getBigDecimalFromArrowBuf(buf, 0, 0, byteLengths[x]); BigDecimal expected = BigDecimal.valueOf(val); assertEquals(expected, actual); } BigInteger[] decimals = new BigInteger[] {MAX_BIG_INT[x], new BigInteger("0"), MIN_BIG_INT[x]}; for (BigInteger val : decimals) { buf.clear(); DecimalUtility.writeByteArrayToArrowBuf(val.toByteArray(), buf, 0, byteLengths[x]); BigDecimal actual = DecimalUtility.getBigDecimalFromArrowBuf(buf, 0, 0, byteLengths[x]); BigDecimal expected = new BigDecimal(val); assertEquals(expected, actual); } } } } @Test public void testSetBigDecimalInDecimalArrowBuf() { int[] byteLengths = new int[] {16, 32}; for (int x = 0; x < 2; x++) { try (BufferAllocator allocator = new RootAllocator(128); ArrowBuf buf = allocator.buffer(byteLengths[x]); ) { int[] intValues = new int[] {Integer.MAX_VALUE, Integer.MIN_VALUE, 0}; for (int val : intValues) { buf.clear(); DecimalUtility.writeBigDecimalToArrowBuf(BigDecimal.valueOf(val), buf, 0, byteLengths[x]); BigDecimal actual = DecimalUtility.getBigDecimalFromArrowBuf(buf, 0, 0, byteLengths[x]); BigDecimal expected = BigDecimal.valueOf(val); assertEquals(expected, actual); } long[] longValues = new long[] {Long.MIN_VALUE, 0, Long.MAX_VALUE}; for (long val : longValues) { buf.clear(); DecimalUtility.writeBigDecimalToArrowBuf(BigDecimal.valueOf(val), buf, 0, byteLengths[x]); BigDecimal actual = DecimalUtility.getBigDecimalFromArrowBuf(buf, 0, 0, byteLengths[x]); BigDecimal expected = BigDecimal.valueOf(val); assertEquals(expected, actual); } BigInteger[] decimals = new BigInteger[] {MAX_BIG_INT[x], new BigInteger("0"), MIN_BIG_INT[x]}; for (BigInteger val : decimals) { buf.clear(); DecimalUtility.writeBigDecimalToArrowBuf(new BigDecimal(val), buf, 0, byteLengths[x]); BigDecimal actual = DecimalUtility.getBigDecimalFromArrowBuf(buf, 0, 0, byteLengths[x]); BigDecimal expected = new BigDecimal(val); assertEquals(expected, actual); } } } } } TestDataSizeRoundingUtil.java000066400000000000000000000056351515451615200342150ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.junit.jupiter.api.Assertions.assertEquals; import org.junit.jupiter.api.Test; /** Test cases for {@link DataSizeRoundingUtil}. */ public class TestDataSizeRoundingUtil { @Test public void testRoundUpTo8MultipleInt() { assertEquals(0, DataSizeRoundingUtil.roundUpTo8Multiple(0)); assertEquals(16, DataSizeRoundingUtil.roundUpTo8Multiple(9)); assertEquals(24, DataSizeRoundingUtil.roundUpTo8Multiple(20)); assertEquals(128, DataSizeRoundingUtil.roundUpTo8Multiple(128)); } @Test public void testRoundUpTo8MultipleLong() { assertEquals(0L, DataSizeRoundingUtil.roundUpTo8Multiple(0L)); assertEquals(40L, DataSizeRoundingUtil.roundUpTo8Multiple(37L)); assertEquals(32L, DataSizeRoundingUtil.roundUpTo8Multiple(29L)); assertEquals(512L, DataSizeRoundingUtil.roundUpTo8Multiple(512L)); } @Test public void testRoundDownTo8MultipleInt() { assertEquals(0, DataSizeRoundingUtil.roundDownTo8Multiple(0)); assertEquals(16, DataSizeRoundingUtil.roundDownTo8Multiple(23)); assertEquals(24, DataSizeRoundingUtil.roundDownTo8Multiple(27)); assertEquals(128, DataSizeRoundingUtil.roundDownTo8Multiple(128)); } @Test public void testRoundDownTo8MultipleLong() { assertEquals(0L, DataSizeRoundingUtil.roundDownTo8Multiple(0L)); assertEquals(40L, DataSizeRoundingUtil.roundDownTo8Multiple(45L)); assertEquals(32L, DataSizeRoundingUtil.roundDownTo8Multiple(39L)); assertEquals(512L, DataSizeRoundingUtil.roundDownTo8Multiple(512L)); } @Test public void testDivideBy8CeilInt() { assertEquals(0, DataSizeRoundingUtil.divideBy8Ceil(0)); assertEquals(3, DataSizeRoundingUtil.divideBy8Ceil(23)); assertEquals(5, DataSizeRoundingUtil.divideBy8Ceil(35)); assertEquals(24, DataSizeRoundingUtil.divideBy8Ceil(192)); } @Test public void testDivideBy8CeilLong() { assertEquals(0L, DataSizeRoundingUtil.divideBy8Ceil(0L)); assertEquals(5L, DataSizeRoundingUtil.divideBy8Ceil(37L)); assertEquals(10L, DataSizeRoundingUtil.divideBy8Ceil(73L)); assertEquals(25L, DataSizeRoundingUtil.divideBy8Ceil(200L)); } } TestElementAddressableVectorIterator.java000066400000000000000000000100051515451615200365500ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNull; import java.nio.charset.StandardCharsets; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.memory.util.ArrowBufPointer; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VarCharVector; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link ElementAddressableVectorIterator}. */ public class TestElementAddressableVectorIterator { private final int VECTOR_LENGTH = 100; private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testIterateIntVector() { try (IntVector intVector = new IntVector("", allocator)) { intVector.allocateNew(VECTOR_LENGTH); intVector.setValueCount(VECTOR_LENGTH); // prepare data in sorted order for (int i = 0; i < VECTOR_LENGTH; i++) { if (i == 0) { intVector.setNull(i); } else { intVector.set(i, i); } } // iterate ElementAddressableVectorIterator it = new ElementAddressableVectorIterator<>(intVector); int index = 0; while (it.hasNext()) { ArrowBufPointer pt; if (index % 2 == 0) { // use populated pointer. pt = new ArrowBufPointer(); it.next(pt); } else { // use iterator inner pointer pt = it.next(); } if (index == 0) { assertNull(pt.getBuf()); } else { assertEquals(index, pt.getBuf().getInt(pt.getOffset())); } index += 1; } } } @Test public void testIterateVarCharVector() { try (VarCharVector strVector = new VarCharVector("", allocator)) { strVector.allocateNew(VECTOR_LENGTH * 10, VECTOR_LENGTH); strVector.setValueCount(VECTOR_LENGTH); // prepare data in sorted order for (int i = 0; i < VECTOR_LENGTH; i++) { if (i == 0) { strVector.setNull(i); } else { strVector.set(i, String.valueOf(i).getBytes(StandardCharsets.UTF_8)); } } // iterate ElementAddressableVectorIterator it = new ElementAddressableVectorIterator<>(strVector); int index = 0; while (it.hasNext()) { ArrowBufPointer pt; if (index % 2 == 0) { // use populated pointer. pt = new ArrowBufPointer(); it.next(pt); } else { // use iterator inner pointer pt = it.next(); } if (index == 0) { assertNull(pt.getBuf()); } else { String expected = String.valueOf(index); byte[] actual = new byte[expected.length()]; assertEquals(expected.length(), pt.getLength()); pt.getBuf().getBytes(pt.getOffset(), actual); assertEquals(expected, new String(actual, StandardCharsets.UTF_8)); } index += 1; } } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/util/TestMapWithOrdinal.java000066400000000000000000000054341515451615200331030ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.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.assertTrue; import java.util.Collection; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestMapWithOrdinal { private MapWithOrdinal map; @BeforeEach public void setUp() { map = new MapWithOrdinalImpl<>(); } @Test public void testGetByOrdinal() { map.put("key0", "val0", true); assertEquals("val0", map.getByOrdinal(0)); map.put("key1", "val1", true); assertEquals("val1", map.getByOrdinal(1)); assertEquals("val0", map.getByOrdinal(0)); } @Test public void testGetByKey() { map.put("key0", "val0", true); assertEquals("val0", map.get("key0")); map.put("key1", "val1", true); assertEquals("val1", map.get("key1")); assertEquals("val0", map.get("key0")); } @Test public void testInvalidOrdinal() { map.put("key0", "val0", true); assertNull(map.getByOrdinal(1)); map.removeAll("key0"); assertNull(map.getByOrdinal(0)); } @Test public void testInvalidKey() { MapWithOrdinalImpl map = new MapWithOrdinalImpl<>(); map.put("key0", "val0", true); assertNull(map.get("fake_key")); map.removeAll("key0"); assertNull(map.get("key0")); } @Test public void testValues() { map.put("key0", "val0", true); map.put("key1", "val1", true); Collection values = map.values(); assertTrue(values.contains("val0")); assertTrue(values.contains("val1")); map.put("key1", "new_val1", true); values = map.values(); assertTrue(values.contains("val0")); assertTrue(values.contains("new_val1")); assertFalse(values.contains("val1")); map.removeAll("key0"); assertTrue(values.contains("new_val1")); assertFalse(values.contains("val0")); } } TestMultiMapWithOrdinal.java000066400000000000000000000043041515451615200340320ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import org.junit.jupiter.api.Test; public class TestMultiMapWithOrdinal { @Test public void test() { MultiMapWithOrdinal map = new MultiMapWithOrdinal<>(); map.put("x", "1", false); assertEquals(1, map.size()); map.remove("x", "1"); assertTrue(map.isEmpty()); map.put("x", "1", false); map.put("x", "2", false); map.put("y", "0", false); assertEquals(3, map.size()); assertEquals(2, map.getAll("x").size()); assertEquals("1", map.getAll("x").stream().findFirst().get()); assertEquals("1", map.getByOrdinal(0)); assertEquals("2", map.getByOrdinal(1)); assertEquals("0", map.getByOrdinal(2)); assertTrue(map.remove("x", "1")); assertFalse(map.remove("x", "1")); assertEquals("0", map.getByOrdinal(0)); assertEquals(2, map.size()); map.put("x", "3", true); assertEquals(1, map.getAll("x").size()); assertEquals("3", map.getAll("x").stream().findFirst().get()); map.put("z", "4", false); assertEquals(3, map.size()); map.put("z", "5", false); map.put("z", "6", false); assertEquals(5, map.size()); map.removeAll("z"); assertEquals(2, map.size()); assertFalse(map.containsKey("z")); } } TestReusableByteArray.java000066400000000000000000000115521515451615200335250ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.Base64; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BaseValueVector; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestReusableByteArray { private BufferAllocator allocator; @BeforeEach public void prepare() { // Permit allocating 4 vectors of max size. allocator = new RootAllocator(4 * BaseValueVector.MAX_ALLOCATION_SIZE); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testSetByteArrayRepeatedly() { ReusableByteArray byteArray = new ReusableByteArray(); try (ArrowBuf workingBuf = allocator.buffer(100)) { final String str = "test"; workingBuf.setBytes(0, str.getBytes(StandardCharsets.UTF_8)); byteArray.set(workingBuf, 0, str.getBytes(StandardCharsets.UTF_8).length); assertEquals(str.getBytes(StandardCharsets.UTF_8).length, byteArray.getLength()); assertArrayEquals( str.getBytes(StandardCharsets.UTF_8), Arrays.copyOfRange(byteArray.getBuffer(), 0, (int) byteArray.getLength())); assertEquals( Base64.getEncoder().encodeToString(str.getBytes(StandardCharsets.UTF_8)), byteArray.toString()); assertEquals(new ReusableByteArray(str.getBytes(StandardCharsets.UTF_8)), byteArray); assertEquals( new ReusableByteArray(str.getBytes(StandardCharsets.UTF_8)).hashCode(), byteArray.hashCode()); // Test a longer string. Should require reallocation. final String str2 = "test_longer"; byte[] oldBuffer = byteArray.getBuffer(); workingBuf.clear(); workingBuf.setBytes(0, str2.getBytes(StandardCharsets.UTF_8)); byteArray.set(workingBuf, 0, str2.getBytes(StandardCharsets.UTF_8).length); assertEquals(str2.getBytes(StandardCharsets.UTF_8).length, byteArray.getLength()); assertArrayEquals( str2.getBytes(StandardCharsets.UTF_8), Arrays.copyOfRange(byteArray.getBuffer(), 0, (int) byteArray.getLength())); assertEquals( Base64.getEncoder().encodeToString(str2.getBytes(StandardCharsets.UTF_8)), byteArray.toString()); assertEquals(new ReusableByteArray(str2.getBytes(StandardCharsets.UTF_8)), byteArray); assertEquals( new ReusableByteArray(str2.getBytes(StandardCharsets.UTF_8)).hashCode(), byteArray.hashCode()); // Verify reallocation needed. assertNotSame(oldBuffer, byteArray.getBuffer()); assertTrue(byteArray.getBuffer().length > oldBuffer.length); // Test writing a shorter string. Should not require reallocation. final String str3 = "short"; oldBuffer = byteArray.getBuffer(); workingBuf.clear(); workingBuf.setBytes(0, str3.getBytes(StandardCharsets.UTF_8)); byteArray.set(workingBuf, 0, str3.getBytes(StandardCharsets.UTF_8).length); assertEquals(str3.getBytes(StandardCharsets.UTF_8).length, byteArray.getLength()); assertArrayEquals( str3.getBytes(StandardCharsets.UTF_8), Arrays.copyOfRange(byteArray.getBuffer(), 0, (int) byteArray.getLength())); assertEquals( Base64.getEncoder().encodeToString(str3.getBytes(StandardCharsets.UTF_8)), byteArray.toString()); assertEquals(new ReusableByteArray(str3.getBytes(StandardCharsets.UTF_8)), byteArray); assertEquals( new ReusableByteArray(str3.getBytes(StandardCharsets.UTF_8)).hashCode(), byteArray.hashCode()); // Verify reallocation was not needed. assertSame(oldBuffer, byteArray.getBuffer()); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/util/TestSchemaUtil.java000066400000000000000000000037141515451615200322560ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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 java.util.Arrays.asList; import static org.junit.jupiter.api.Assertions.assertEquals; import java.io.IOException; import org.apache.arrow.memory.RootAllocator; 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.junit.jupiter.api.Test; public class TestSchemaUtil { private static Field field(String name, boolean nullable, ArrowType type, Field... children) { return new Field(name, new FieldType(nullable, type, null, null), asList(children)); } @Test public void testSerializationAndDeserialization() throws IOException { Schema schema = new Schema( asList( field("a", false, new ArrowType.Null()), field("b", true, new ArrowType.Utf8()), field("c", true, new ArrowType.Binary()))); byte[] serialized = SchemaUtility.serialize(schema); Schema deserialized = SchemaUtility.deserialize(serialized, new RootAllocator(Long.MAX_VALUE)); assertEquals(schema, deserialized); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/util/TestValidator.java000066400000000000000000000051371515451615200321460ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.Validator.equalEnough; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import org.junit.jupiter.api.Test; public class TestValidator { @Test public void testFloatComp() { assertTrue(equalEnough(912.4140000000002F, 912.414F)); assertTrue(equalEnough(912.4140000000002D, 912.414D)); assertTrue(equalEnough(912.414F, 912.4140000000002F)); assertTrue(equalEnough(912.414D, 912.4140000000002D)); assertFalse(equalEnough(912.414D, 912.4140001D)); assertFalse(equalEnough(null, 912.414D)); assertTrue(equalEnough((Float) null, null)); assertTrue(equalEnough((Double) null, null)); assertFalse(equalEnough(912.414D, null)); assertFalse(equalEnough(Double.MAX_VALUE, Double.MIN_VALUE)); assertFalse(equalEnough(Double.MIN_VALUE, Double.MAX_VALUE)); assertTrue(equalEnough(Double.MAX_VALUE, Double.MAX_VALUE)); assertTrue(equalEnough(Double.MIN_VALUE, Double.MIN_VALUE)); assertTrue(equalEnough(Double.NEGATIVE_INFINITY, Double.NEGATIVE_INFINITY)); assertFalse(equalEnough(Double.NEGATIVE_INFINITY, Double.POSITIVE_INFINITY)); assertTrue(equalEnough(Double.NaN, Double.NaN)); assertFalse(equalEnough(1.0, Double.NaN)); assertFalse(equalEnough(Float.MAX_VALUE, Float.MIN_VALUE)); assertFalse(equalEnough(Float.MIN_VALUE, Float.MAX_VALUE)); assertTrue(equalEnough(Float.MAX_VALUE, Float.MAX_VALUE)); assertTrue(equalEnough(Float.MIN_VALUE, Float.MIN_VALUE)); assertTrue(equalEnough(Float.NEGATIVE_INFINITY, Float.NEGATIVE_INFINITY)); assertFalse(equalEnough(Float.NEGATIVE_INFINITY, Float.POSITIVE_INFINITY)); assertTrue(equalEnough(Float.NaN, Float.NaN)); assertFalse(equalEnough(1.0F, Float.NaN)); } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java000066400000000000000000001211561515451615200331420ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.charset.StandardCharsets; 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.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.memory.util.CommonUtil; import org.apache.arrow.vector.BaseValueVector; 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; import org.apache.arrow.vector.complex.DenseUnionVector; 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.testing.ValueVectorDataPopulator; 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.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link VectorAppender}. */ public class TestVectorAppender { private BufferAllocator allocator; @BeforeEach public void prepare() { // Permit allocating 4 vectors of max size. allocator = new RootAllocator(4 * BaseValueVector.MAX_ALLOCATION_SIZE); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testAppendFixedWidthVector() { final int length1 = 10; final int length2 = 5; try (IntVector target = new IntVector("", allocator); IntVector delta = new IntVector("", allocator)) { target.allocateNew(length1); delta.allocateNew(length2); ValueVectorDataPopulator.setVector(target, 0, 1, 2, 3, 4, 5, 6, null, 8, 9); ValueVectorDataPopulator.setVector(delta, null, 11, 12, 13, 14); VectorAppender appender = new VectorAppender(target); delta.accept(appender, null); assertEquals(length1 + length2, target.getValueCount()); try (IntVector expected = new IntVector("expected", allocator)) { expected.allocateNew(); ValueVectorDataPopulator.setVector( expected, 0, 1, 2, 3, 4, 5, 6, null, 8, 9, null, 11, 12, 13, 14); assertVectorsEqual(expected, target); } } } @Test public void testAppendBitVector() { final int length1 = 10; final int length2 = 5; try (BitVector target = new BitVector("", allocator); BitVector delta = new BitVector("", allocator)) { target.allocateNew(length1); delta.allocateNew(length2); ValueVectorDataPopulator.setVector(target, 0, 1, 0, 1, 0, 1, 0, null, 0, 1); ValueVectorDataPopulator.setVector(delta, null, 1, 1, 0, 0); VectorAppender appender = new VectorAppender(target); delta.accept(appender, null); assertEquals(length1 + length2, target.getValueCount()); try (BitVector expected = new BitVector("expected", allocator)) { expected.allocateNew(); ValueVectorDataPopulator.setVector( expected, 0, 1, 0, 1, 0, 1, 0, null, 0, 1, null, 1, 1, 0, 0); assertVectorsEqual(expected, target); } } } @Test public void testAppendEmptyFixedWidthVector() { try (IntVector target = new IntVector("", allocator); IntVector delta = new IntVector("", allocator)) { ValueVectorDataPopulator.setVector(target, 0, 1, 2, 3, 4, 5, 6, null, 8, 9); VectorAppender appender = new VectorAppender(target); delta.accept(appender, null); assertEquals(10, target.getValueCount()); try (IntVector expected = new IntVector("expected", allocator)) { ValueVectorDataPopulator.setVector(expected, 0, 1, 2, 3, 4, 5, 6, null, 8, 9); assertVectorsEqual(expected, target); } } } @Test public void testAppendVariableWidthVector() { final int length1 = 10; final int length2 = 5; try (VarCharVector target = new VarCharVector("", allocator); VarCharVector delta = new VarCharVector("", 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 (VarCharVector expected = new VarCharVector("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 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); VarCharVector delta = new VarCharVector("", allocator)) { ValueVectorDataPopulator.setVector( target, "a0", "a1", "a2", "a3", null, "a5", "a6", "a7", "a8", "a9"); VectorAppender appender = new VectorAppender(target); delta.accept(appender, null); try (VarCharVector expected = new VarCharVector("expected", allocator)) { ValueVectorDataPopulator.setVector( expected, "a0", "a1", "a2", "a3", null, "a5", "a6", "a7", "a8", "a9"); assertVectorsEqual(expected, target); } } } @Test public void testAppendLargeAndSmallVariableVectorsWithinLimit() { // 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); new VectorAppender(target).visit(delta, null); } } private VarCharVector makeVarCharVec(int numElements, int bytesPerElement) { VarCharVector v = new VarCharVector("text", allocator); v.allocateNew((long) numElements * bytesPerElement, numElements); for (int i = 0; i < numElements; i++) { String s = String.join("", Collections.nCopies(bytesPerElement, "a")); v.setSafe(i, s.getBytes(StandardCharsets.US_ASCII)); } v.setValueCount(numElements); return v; } @Test public void testAppendLargeVariableWidthVector() { final int length1 = 5; final int length2 = 10; try (LargeVarCharVector target = new LargeVarCharVector("", allocator); LargeVarCharVector delta = new LargeVarCharVector("", allocator)) { target.allocateNew(5, length1); delta.allocateNew(5, length2); ValueVectorDataPopulator.setVector(target, "a0", null, "a2", "a3", null); ValueVectorDataPopulator.setVector( delta, "a5", "a6", "a7", null, null, "a10", "a11", "a12", "a13", null); VectorAppender appender = new VectorAppender(target); delta.accept(appender, null); try (LargeVarCharVector expected = new LargeVarCharVector("expected", allocator)) { expected.allocateNew(); ValueVectorDataPopulator.setVector( expected, "a0", null, "a2", "a3", null, "a5", "a6", "a7", null, null, "a10", "a11", "a12", "a13", null); assertVectorsEqual(expected, target); } } } @Test public void testAppendEmptyLargeVariableWidthVector() { try (LargeVarCharVector target = new LargeVarCharVector("", allocator); LargeVarCharVector delta = new LargeVarCharVector("", allocator)) { ValueVectorDataPopulator.setVector(target, "a0", null, "a2", "a3", null); VectorAppender appender = new VectorAppender(target); delta.accept(appender, null); try (LargeVarCharVector expected = new LargeVarCharVector("expected", allocator)) { ValueVectorDataPopulator.setVector(expected, "a0", null, "a2", "a3", null); assertVectorsEqual(expected, target); } } } @Test public void testAppendListVector() { final int length1 = 5; final int length2 = 2; 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), null, Arrays.asList(6, 7), Arrays.asList(8, 9)); assertEquals(length1, target.getValueCount()); delta.allocateNew(); ValueVectorDataPopulator.setVector( delta, Arrays.asList(10, 11, 12, 13, 14), Arrays.asList(15, 16, 17, 18, 19)); assertEquals(length2, delta.getValueCount()); VectorAppender appender = new VectorAppender(target); delta.accept(appender, null); assertEquals(7, target.getValueCount()); List expected = Arrays.asList(0, 1); assertEquals(expected, target.getObject(0)); expected = Arrays.asList(2, 3); assertEquals(expected, target.getObject(1)); assertTrue(target.isNull(2)); expected = Arrays.asList(6, 7); assertEquals(expected, target.getObject(3)); expected = Arrays.asList(8, 9); assertEquals(expected, target.getObject(4)); expected = Arrays.asList(10, 11, 12, 13, 14); assertEquals(expected, target.getObject(5)); expected = Arrays.asList(15, 16, 17, 18, 19); assertEquals(expected, target.getObject(6)); } } @Test public void testAppendEmptyListVector() { try (ListVector target = ListVector.empty("target", allocator); ListVector delta = ListVector.empty("delta", allocator)) { // populate target with data ValueVectorDataPopulator.setVector( target, Arrays.asList(0, 1), Arrays.asList(2, 3), null, Arrays.asList(6, 7)); assertEquals(4, target.getValueCount()); // leave delta vector empty and unallocated delta.addOrGetVector(FieldType.nullable(Types.MinorType.INT.getType())); VectorAppender appender = new VectorAppender(target); delta.accept(appender, null); // verify delta vector has original data assertEquals(4, target.getValueCount()); List expected = Arrays.asList(0, 1); assertEquals(expected, target.getObject(0)); expected = Arrays.asList(2, 3); assertEquals(expected, target.getObject(1)); assertTrue(target.isNull(2)); expected = Arrays.asList(6, 7); assertEquals(expected, target.getObject(3)); } } @Test public void testAppendFixedSizeListVector() { try (FixedSizeListVector target = FixedSizeListVector.empty("target", 5, allocator); FixedSizeListVector delta = FixedSizeListVector.empty("delta", 5, allocator)) { target.allocateNew(); ValueVectorDataPopulator.setVector(target, Arrays.asList(0, 1, 2, 3, 4), null); assertEquals(2, target.getValueCount()); delta.allocateNew(); ValueVectorDataPopulator.setVector( delta, Arrays.asList(10, 11, 12, 13, 14), Arrays.asList(15, 16, 17, 18, 19)); assertEquals(2, delta.getValueCount()); VectorAppender appender = new VectorAppender(target); delta.accept(appender, null); assertEquals(4, target.getValueCount()); assertEquals(Arrays.asList(0, 1, 2, 3, 4), target.getObject(0)); assertTrue(target.isNull(1)); assertEquals(Arrays.asList(10, 11, 12, 13, 14), target.getObject(2)); assertEquals(Arrays.asList(15, 16, 17, 18, 19), target.getObject(3)); } } @Test public void testAppendEmptyFixedSizeListVector() { try (FixedSizeListVector target = FixedSizeListVector.empty("target", 5, allocator); FixedSizeListVector delta = FixedSizeListVector.empty("delta", 5, allocator)) { ValueVectorDataPopulator.setVector(target, Arrays.asList(0, 1, 2, 3, 4), null); assertEquals(2, target.getValueCount()); // leave delta vector empty and unallocated delta.addOrGetVector(FieldType.nullable(Types.MinorType.INT.getType())); VectorAppender appender = new VectorAppender(target); delta.accept(appender, null); assertEquals(2, target.getValueCount()); assertEquals(Arrays.asList(0, 1, 2, 3, 4), target.getObject(0)); assertTrue(target.isNull(1)); } } @Test public void testAppendEmptyLargeListVector() { try (LargeListVector target = LargeListVector.empty("target", allocator); LargeListVector delta = LargeListVector.empty("delta", allocator)) { ValueVectorDataPopulator.setVector(target, Arrays.asList(0, 1, 2, 3, 4), null); assertEquals(2, target.getValueCount()); // leave delta vector empty and unallocated delta.addOrGetVector(FieldType.nullable(Types.MinorType.INT.getType())); VectorAppender appender = new VectorAppender(target); delta.accept(appender, null); assertEquals(2, target.getValueCount()); assertEquals(Arrays.asList(0, 1, 2, 3, 4), target.getObject(0)); assertTrue(target.isNull(1)); } } @Test public void testAppendStructVector() { final int length1 = 10; final int length2 = 5; try (final StructVector target = StructVector.empty("target", allocator); final StructVector delta = StructVector.empty("delta", allocator)) { IntVector targetChild1 = target.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); VarCharVector targetChild2 = target.addOrGet("f1", FieldType.nullable(new ArrowType.Utf8()), VarCharVector.class); targetChild1.allocateNew(); targetChild2.allocateNew(); ValueVectorDataPopulator.setVector(targetChild1, 0, 1, 2, 3, 4, null, 6, 7, 8, 9); ValueVectorDataPopulator.setVector( targetChild2, "a0", "a1", "a2", "a3", "a4", "a5", "a6", null, "a8", "a9"); target.setValueCount(length1); IntVector deltaChild1 = delta.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); VarCharVector deltaChild2 = delta.addOrGet("f1", FieldType.nullable(new ArrowType.Utf8()), VarCharVector.class); deltaChild1.allocateNew(); deltaChild2.allocateNew(); ValueVectorDataPopulator.setVector(deltaChild1, 10, 11, 12, null, 14); ValueVectorDataPopulator.setVector(deltaChild2, "a10", "a11", "a12", "a13", "a14"); delta.setValueCount(length2); VectorAppender appender = new VectorAppender(target); delta.accept(appender, null); assertEquals(length1 + length2, target.getValueCount()); try (IntVector expected1 = new IntVector("expected1", allocator); VarCharVector expected2 = new VarCharVector("expected2", allocator)) { expected1.allocateNew(); expected2.allocateNew(); ValueVectorDataPopulator.setVector( expected1, 0, 1, 2, 3, 4, null, 6, 7, 8, 9, 10, 11, 12, null, 14); ValueVectorDataPopulator.setVector( expected2, "a0", "a1", "a2", "a3", "a4", "a5", "a6", null, "a8", "a9", "a10", "a11", "a12", "a13", "a14"); assertVectorsEqual(expected1, target.getChild("f0")); assertVectorsEqual(expected2, target.getChild("f1")); } } } @Test public void testAppendEmptyStructVector() { try (final StructVector target = StructVector.empty("target", allocator); final StructVector delta = StructVector.empty("delta", allocator)) { IntVector targetChild1 = target.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); VarCharVector targetChild2 = target.addOrGet("f1", FieldType.nullable(new ArrowType.Utf8()), VarCharVector.class); ValueVectorDataPopulator.setVector(targetChild1, 0, 1, 2, 3, 4, null, 6, 7, 8, 9); ValueVectorDataPopulator.setVector( targetChild2, "a0", "a1", "a2", "a3", "a4", "a5", "a6", null, "a8", "a9"); target.setValueCount(10); // leave delta vector fields empty and unallocated delta.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); delta.addOrGet("f1", FieldType.nullable(new ArrowType.Utf8()), VarCharVector.class); VectorAppender appender = new VectorAppender(target); delta.accept(appender, null); assertEquals(10, target.getValueCount()); try (IntVector expected1 = new IntVector("expected1", allocator); VarCharVector expected2 = new VarCharVector("expected2", allocator)) { ValueVectorDataPopulator.setVector(expected1, 0, 1, 2, 3, 4, null, 6, 7, 8, 9); ValueVectorDataPopulator.setVector( expected2, "a0", "a1", "a2", "a3", "a4", "a5", "a6", null, "a8", "a9"); assertVectorsEqual(expected1, target.getChild("f0")); assertVectorsEqual(expected2, target.getChild("f1")); } } } @Test public void testAppendUnionVector() { final int length1 = 10; final int length2 = 5; try (final UnionVector target = UnionVector.empty("target", allocator); final UnionVector delta = UnionVector.empty("delta", allocator)) { // alternating ints and big ints target.setType(0, Types.MinorType.INT); target.setType(1, Types.MinorType.BIGINT); target.setType(2, Types.MinorType.INT); target.setType(3, Types.MinorType.BIGINT); target.setType(4, Types.MinorType.INT); target.setType(5, Types.MinorType.BIGINT); target.setType(6, Types.MinorType.INT); target.setType(7, Types.MinorType.BIGINT); target.setType(8, Types.MinorType.INT); target.setType(9, Types.MinorType.BIGINT); target.setType(10, Types.MinorType.INT); target.setType(11, Types.MinorType.BIGINT); target.setType(12, Types.MinorType.INT); target.setType(13, Types.MinorType.BIGINT); target.setType(14, Types.MinorType.INT); target.setType(15, Types.MinorType.BIGINT); target.setType(16, Types.MinorType.INT); target.setType(17, Types.MinorType.BIGINT); target.setType(18, Types.MinorType.INT); target.setType(19, Types.MinorType.BIGINT); IntVector targetIntVec = target.getIntVector(); targetIntVec.allocateNew(); ValueVectorDataPopulator.setVector( targetIntVec, 0, null, 1, null, 2, null, 3, null, 4, null, 5, null, 6, null, 7, null, 8, null, 9, null); assertEquals(length1 * 2, targetIntVec.getValueCount()); BigIntVector targetBigIntVec = target.getBigIntVector(); targetBigIntVec.allocateNew(); ValueVectorDataPopulator.setVector( targetBigIntVec, null, 0L, null, 1L, null, 2L, null, 3L, null, 4L, null, 5L, null, 6L, null, 7L, null, 8L, null, 9L); assertEquals(length1 * 2, targetBigIntVec.getValueCount()); target.setValueCount(length1 * 2); // populate the delta vector delta.setType(0, Types.MinorType.FLOAT4); delta.setType(1, Types.MinorType.FLOAT4); delta.setType(2, Types.MinorType.FLOAT4); delta.setType(3, Types.MinorType.FLOAT4); delta.setType(4, Types.MinorType.FLOAT4); Float4Vector deltaFloatVector = delta.getFloat4Vector(); deltaFloatVector.allocateNew(); ValueVectorDataPopulator.setVector(deltaFloatVector, 10f, 11f, 12f, 13f, 14f); assertEquals(length2, deltaFloatVector.getValueCount()); delta.setValueCount(length2); VectorAppender appender = new VectorAppender(target); delta.accept(appender, null); assertEquals(length1 * 2 + length2, target.getValueCount()); for (int i = 0; i < length1; i++) { Object intObj = target.getObject(i * 2); assertTrue(intObj instanceof Integer); assertEquals(i, ((Integer) intObj).intValue()); Object longObj = target.getObject(i * 2 + 1); assertTrue(longObj instanceof Long); assertEquals(i, ((Long) longObj).longValue()); } for (int i = 0; i < length2; i++) { Object floatObj = target.getObject(length1 * 2 + i); assertTrue(floatObj instanceof Float); assertEquals(i + length1, ((Float) floatObj).intValue()); } } } @Test public void testAppendEmptyUnionVector() { final int length1 = 10; try (final UnionVector target = UnionVector.empty("target", allocator); final UnionVector delta = UnionVector.empty("delta", allocator)) { // alternating ints and big ints target.setType(0, Types.MinorType.INT); target.setType(1, Types.MinorType.BIGINT); target.setType(2, Types.MinorType.INT); target.setType(3, Types.MinorType.BIGINT); target.setType(4, Types.MinorType.INT); target.setType(5, Types.MinorType.BIGINT); target.setType(6, Types.MinorType.INT); target.setType(7, Types.MinorType.BIGINT); target.setType(8, Types.MinorType.INT); target.setType(9, Types.MinorType.BIGINT); target.setType(10, Types.MinorType.INT); target.setType(11, Types.MinorType.BIGINT); target.setType(12, Types.MinorType.INT); target.setType(13, Types.MinorType.BIGINT); target.setType(14, Types.MinorType.INT); target.setType(15, Types.MinorType.BIGINT); target.setType(16, Types.MinorType.INT); target.setType(17, Types.MinorType.BIGINT); target.setType(18, Types.MinorType.INT); target.setType(19, Types.MinorType.BIGINT); IntVector targetIntVec = target.getIntVector(); ValueVectorDataPopulator.setVector( targetIntVec, 0, null, 1, null, 2, null, 3, null, 4, null, 5, null, 6, null, 7, null, 8, null, 9, null); assertEquals(length1 * 2, targetIntVec.getValueCount()); BigIntVector targetBigIntVec = target.getBigIntVector(); ValueVectorDataPopulator.setVector( targetBigIntVec, null, 0L, null, 1L, null, 2L, null, 3L, null, 4L, null, 5L, null, 6L, null, 7L, null, 8L, null, 9L); assertEquals(length1 * 2, targetBigIntVec.getValueCount()); target.setValueCount(length1 * 2); // initialize the delta vector but leave it empty and unallocated delta.setType(0, Types.MinorType.FLOAT4); delta.setType(1, Types.MinorType.FLOAT4); delta.setType(2, Types.MinorType.FLOAT4); delta.setType(3, Types.MinorType.FLOAT4); delta.setType(4, Types.MinorType.FLOAT4); VectorAppender appender = new VectorAppender(target); delta.accept(appender, null); assertEquals(length1 * 2, target.getValueCount()); for (int i = 0; i < length1; i++) { Object intObj = target.getObject(i * 2); assertTrue(intObj instanceof Integer); assertEquals(i, ((Integer) intObj).intValue()); Object longObj = target.getObject(i * 2 + 1); assertTrue(longObj instanceof Long); assertEquals(i, ((Long) longObj).longValue()); } } } private DenseUnionVector getTargetVector() { // create a vector, and populate it with values {1, 2, null, 10L} final NullableIntHolder intHolder = new NullableIntHolder(); intHolder.isSet = 1; final NullableBigIntHolder longHolder = new NullableBigIntHolder(); longHolder.isSet = 1; final NullableFloat4Holder floatHolder = new NullableFloat4Holder(); floatHolder.isSet = 1; DenseUnionVector targetVector = new DenseUnionVector("target vector", allocator, null, null); targetVector.allocateNew(); while (targetVector.getValueCapacity() < 4) { targetVector.reAlloc(); } byte intTypeId = targetVector.registerNewTypeId(Field.nullable("", Types.MinorType.INT.getType())); targetVector.setTypeId(0, intTypeId); intHolder.value = 1; targetVector.setSafe(0, intHolder); targetVector.setTypeId(1, intTypeId); intHolder.value = 2; targetVector.setSafe(1, intHolder); byte longTypeId = targetVector.registerNewTypeId(Field.nullable("", Types.MinorType.BIGINT.getType())); targetVector.setTypeId(3, longTypeId); longHolder.value = 10L; targetVector.setSafe(3, longHolder); targetVector.setValueCount(4); assertVectorValuesEqual(targetVector, new Object[] {1, 2, null, 10L}); return targetVector; } private DenseUnionVector getDeltaVector() { // create a vector, and populate it with values {7, null, 8L, 9.0f} final NullableIntHolder intHolder = new NullableIntHolder(); intHolder.isSet = 1; final NullableBigIntHolder longHolder = new NullableBigIntHolder(); longHolder.isSet = 1; final NullableFloat4Holder floatHolder = new NullableFloat4Holder(); floatHolder.isSet = 1; DenseUnionVector deltaVector = new DenseUnionVector("target vector", allocator, null, null); while (deltaVector.getValueCapacity() < 4) { deltaVector.reAlloc(); } byte intTypeId = deltaVector.registerNewTypeId(Field.nullable("", Types.MinorType.INT.getType())); deltaVector.setTypeId(0, intTypeId); intHolder.value = 7; deltaVector.setSafe(0, intHolder); byte longTypeId = deltaVector.registerNewTypeId(Field.nullable("", Types.MinorType.BIGINT.getType())); deltaVector.setTypeId(2, longTypeId); longHolder.value = 8L; deltaVector.setSafe(2, longHolder); byte floatTypeId = deltaVector.registerNewTypeId(Field.nullable("", Types.MinorType.FLOAT4.getType())); deltaVector.setTypeId(3, floatTypeId); floatHolder.value = 9.0f; deltaVector.setSafe(3, floatHolder); deltaVector.setValueCount(4); assertVectorValuesEqual(deltaVector, new Object[] {7, null, 8L, 9.0f}); return deltaVector; } @Test public void testAppendDenseUnionVector() { try (DenseUnionVector targetVector = getTargetVector(); DenseUnionVector deltaVector = getDeltaVector()) { // append VectorAppender appender = new VectorAppender(targetVector); deltaVector.accept(appender, null); assertVectorValuesEqual(targetVector, new Object[] {1, 2, null, 10L, 7, null, 8L, 9.0f}); } // test reverse append try (DenseUnionVector targetVector = getTargetVector(); DenseUnionVector deltaVector = getDeltaVector()) { // append VectorAppender appender = new VectorAppender(deltaVector); targetVector.accept(appender, null); assertVectorValuesEqual(deltaVector, new Object[] {7, null, 8L, 9.0f, 1, 2, null, 10L}); } } private DenseUnionVector getEmptyDeltaVector() { // create a vector, but leave it empty and uninitialized DenseUnionVector deltaVector = new DenseUnionVector("target vector", allocator, null, null); byte intTypeId = deltaVector.registerNewTypeId(Field.nullable("", Types.MinorType.INT.getType())); deltaVector.setTypeId(0, intTypeId); byte longTypeId = deltaVector.registerNewTypeId(Field.nullable("", Types.MinorType.BIGINT.getType())); deltaVector.setTypeId(2, longTypeId); byte floatTypeId = deltaVector.registerNewTypeId(Field.nullable("", Types.MinorType.FLOAT4.getType())); deltaVector.setTypeId(3, floatTypeId); return deltaVector; } @Test public void testAppendEmptyDenseUnionVector() { try (DenseUnionVector targetVector = getTargetVector(); DenseUnionVector deltaVector = getEmptyDeltaVector()) { // append VectorAppender appender = new VectorAppender(targetVector); deltaVector.accept(appender, null); assertVectorValuesEqual(targetVector, new Object[] {1, 2, null, 10L}); } } /** Test appending dense union vectors where the child vectors do not match. */ @Test public void testAppendDenseUnionVectorMismatch() { final NullableIntHolder intHolder = new NullableIntHolder(); intHolder.isSet = 1; final NullableBigIntHolder longHolder = new NullableBigIntHolder(); longHolder.isSet = 1; final NullableFloat4Holder floatHolder = new NullableFloat4Holder(); floatHolder.isSet = 1; try (DenseUnionVector targetVector = new DenseUnionVector("target vector", allocator, null, null); DenseUnionVector deltaVector = new DenseUnionVector("target vector", allocator, null, null)) { targetVector.allocateNew(); deltaVector.allocateNew(); // populate the target vector with values {1, 2L} while (targetVector.getValueCapacity() < 2) { targetVector.reAlloc(); } byte intTypeId = targetVector.registerNewTypeId(Field.nullable("", Types.MinorType.INT.getType())); targetVector.setTypeId(0, intTypeId); intHolder.value = 1; targetVector.setSafe(0, intHolder); byte longTypeId = targetVector.registerNewTypeId(Field.nullable("", Types.MinorType.BIGINT.getType())); targetVector.setTypeId(1, longTypeId); longHolder.value = 2L; targetVector.setSafe(1, longHolder); targetVector.setValueCount(2); assertVectorValuesEqual(targetVector, new Object[] {1, 2L}); // populate the delta vector with values {3, 5.0f} while (deltaVector.getValueCapacity() < 2) { deltaVector.reAlloc(); } intTypeId = deltaVector.registerNewTypeId(Field.nullable("", Types.MinorType.INT.getType())); deltaVector.setTypeId(0, intTypeId); intHolder.value = 3; deltaVector.setSafe(0, intHolder); byte floatTypeId = deltaVector.registerNewTypeId(Field.nullable("", Types.MinorType.FLOAT4.getType())); deltaVector.setTypeId(1, floatTypeId); floatHolder.value = 5.0f; deltaVector.setSafe(1, floatHolder); deltaVector.setValueCount(2); assertVectorValuesEqual(deltaVector, new Object[] {3, 5.0f}); // append VectorAppender appender = new VectorAppender(targetVector); assertThrows(IllegalArgumentException.class, () -> deltaVector.accept(appender, null)); } } @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; try (IntVector target = new IntVector("", allocator); VarCharVector delta = new VarCharVector("", allocator)) { target.allocateNew(vectorLength); delta.allocateNew(vectorLength); VectorAppender appender = new VectorAppender(target); assertThrows(IllegalArgumentException.class, () -> delta.accept(appender, null)); } } private void assertVectorValuesEqual(ValueVector vector, Object[] values) { assertEquals(vector.getValueCount(), values.length); for (int i = 0; i < values.length; i++) { assertEquals(vector.getObject(i), values[i]); } } public static void assertVectorsEqual(ValueVector vector1, ValueVector vector2) { assertEquals(vector1.getValueCount(), vector2.getValueCount()); TypeEqualsVisitor typeEqualsVisitor = new TypeEqualsVisitor(vector1, false, false); RangeEqualsVisitor equalsVisitor = new RangeEqualsVisitor(vector1, vector2, (v1, v2) -> typeEqualsVisitor.equals(vector2)); assertTrue(equalsVisitor.rangeEquals(new Range(0, 0, vector1.getValueCount()))); } } TestVectorBatchAppender.java000066400000000000000000000045471515451615200340310ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.junit.jupiter.api.Assertions.assertEquals; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link VectorBatchAppender}. */ public class TestVectorBatchAppender { private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testBatchAppendIntVector() { final int length1 = 10; final int length2 = 5; final int length3 = 7; try (IntVector target = new IntVector("", allocator); IntVector delta1 = new IntVector("", allocator); IntVector delta2 = new IntVector("", allocator)) { target.allocateNew(length1); delta1.allocateNew(length2); delta2.allocateNew(length3); ValueVectorDataPopulator.setVector(target, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9); ValueVectorDataPopulator.setVector(delta1, 10, 11, 12, 13, 14); ValueVectorDataPopulator.setVector(delta2, 15, 16, 17, 18, 19, 20, 21); VectorBatchAppender.batchAppend(target, delta1, delta2); assertEquals(length1 + length2 + length3, target.getValueCount()); for (int i = 0; i < target.getValueCount(); i++) { assertEquals(i, target.get(i)); } } } } TestVectorSchemaRootAppender.java000066400000000000000000000156771515451615200350620ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/util/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.util.TestVectorAppender.assertVectorsEqual; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link VectorSchemaRootAppender}. */ public class TestVectorSchemaRootAppender { private BufferAllocator allocator; @BeforeEach public void prepare() { allocator = new RootAllocator(1024 * 1024); } @AfterEach public void shutdown() { allocator.close(); } @Test public void testVectorSchemaRootAppend() { final int length1 = 5; final int length2 = 3; final int length3 = 2; try (IntVector targetChild1 = new IntVector("t1", allocator); VarCharVector targetChild2 = new VarCharVector("t2", allocator); BigIntVector targetChild3 = new BigIntVector("t3", allocator); IntVector deltaChildOne1 = new IntVector("do1", allocator); VarCharVector deltaChildOne2 = new VarCharVector("do2", allocator); BigIntVector deltaChildOne3 = new BigIntVector("do3", allocator); IntVector deltaChildTwo1 = new IntVector("dt1", allocator); VarCharVector deltaChildTwo2 = new VarCharVector("dt2", allocator); BigIntVector deltaChildTwo3 = new BigIntVector("dt3", allocator)) { ValueVectorDataPopulator.setVector(targetChild1, 0, 1, null, 3, 4); ValueVectorDataPopulator.setVector(targetChild2, "zero", "one", null, "three", "four"); ValueVectorDataPopulator.setVector(targetChild3, 0L, 10L, null, 30L, 40L); VectorSchemaRoot root1 = VectorSchemaRoot.of(targetChild1, targetChild2, targetChild3); root1.setRowCount(length1); ValueVectorDataPopulator.setVector(deltaChildOne1, 5, 6, 7); ValueVectorDataPopulator.setVector(deltaChildOne2, "five", "six", "seven"); ValueVectorDataPopulator.setVector(deltaChildOne3, 50L, 60L, 70L); VectorSchemaRoot root2 = VectorSchemaRoot.of(deltaChildOne1, deltaChildOne2, deltaChildOne3); root2.setRowCount(length2); ValueVectorDataPopulator.setVector(deltaChildTwo1, null, 9); ValueVectorDataPopulator.setVector(deltaChildTwo2, null, "nine"); ValueVectorDataPopulator.setVector(deltaChildTwo3, null, 90L); VectorSchemaRoot root3 = VectorSchemaRoot.of(deltaChildTwo1, deltaChildTwo2, deltaChildTwo3); root3.setRowCount(length3); VectorSchemaRootAppender.append(root1, root2, root3); assertEquals(length1 + length2 + length3, root1.getRowCount()); assertEquals(3, root1.getFieldVectors().size()); try (IntVector expected1 = new IntVector("", allocator); VarCharVector expected2 = new VarCharVector("", allocator); BigIntVector expected3 = new BigIntVector("", allocator)) { ValueVectorDataPopulator.setVector(expected1, 0, 1, null, 3, 4, 5, 6, 7, null, 9); ValueVectorDataPopulator.setVector( expected2, "zero", "one", null, "three", "four", "five", "six", "seven", null, "nine"); ValueVectorDataPopulator.setVector( expected3, 0L, 10L, null, 30L, 40L, 50L, 60L, 70L, null, 90L); assertVectorsEqual(expected1, root1.getVector(0)); assertVectorsEqual(expected2, root1.getVector(1)); assertVectorsEqual(expected3, root1.getVector(2)); } } } @Test public void testRootWithDifferentChildCounts() { try (IntVector targetChild1 = new IntVector("t1", allocator); VarCharVector targetChild2 = new VarCharVector("t2", allocator); BigIntVector targetChild3 = new BigIntVector("t3", allocator); IntVector deltaChild1 = new IntVector("d1", allocator); VarCharVector deltaChild2 = new VarCharVector("d2", allocator)) { ValueVectorDataPopulator.setVector(targetChild1, 0, 1, null, 3, 4); ValueVectorDataPopulator.setVector(targetChild2, "zero", "one", null, "three", "four"); ValueVectorDataPopulator.setVector(targetChild3, 0L, 10L, null, 30L, 40L); VectorSchemaRoot root1 = VectorSchemaRoot.of(targetChild1, targetChild2, targetChild3); root1.setRowCount(5); ValueVectorDataPopulator.setVector(deltaChild1, 5, 6, 7); ValueVectorDataPopulator.setVector(deltaChild2, "five", "six", "seven"); VectorSchemaRoot root2 = VectorSchemaRoot.of(deltaChild1, deltaChild2); root2.setRowCount(3); IllegalArgumentException exp = assertThrows( IllegalArgumentException.class, () -> VectorSchemaRootAppender.append(root1, root2)); assertEquals( "Vector schema roots have different numbers of child vectors.", exp.getMessage()); } } @Test public void testRootWithDifferentChildTypes() { try (IntVector targetChild1 = new IntVector("t1", allocator); VarCharVector targetChild2 = new VarCharVector("t2", allocator); IntVector deltaChild1 = new IntVector("d1", allocator); VarCharVector deltaChild2 = new VarCharVector("d2", allocator)) { ValueVectorDataPopulator.setVector(targetChild1, 0, 1, null, 3, 4); ValueVectorDataPopulator.setVector(targetChild2, "zero", "one", null, "three", "four"); VectorSchemaRoot root1 = VectorSchemaRoot.of(targetChild1, targetChild2); root1.setRowCount(5); ValueVectorDataPopulator.setVector(deltaChild1, 5, 6, 7); ValueVectorDataPopulator.setVector(deltaChild2, "five", "six", "seven"); // note that the child vectors are in reverse order VectorSchemaRoot root2 = VectorSchemaRoot.of(deltaChild2, deltaChild1); root2.setRowCount(3); IllegalArgumentException exp = assertThrows( IllegalArgumentException.class, () -> VectorSchemaRootAppender.append(root1, root2)); assertEquals("Vector schema roots have different schemas.", exp.getMessage()); } } } apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/validate/000077500000000000000000000000001515451615200273245ustar00rootroot00000000000000TestValidateVector.java000066400000000000000000000271221515451615200336700ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/validate/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.validate; import static org.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector; import static org.apache.arrow.vector.util.ValueVectorUtility.validate; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; 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; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.LargeVarCharVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.complex.DenseUnionVector; 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.complex.impl.NullableStructWriter; import org.apache.arrow.vector.holders.NullableFloat4Holder; import org.apache.arrow.vector.holders.NullableFloat8Holder; import org.apache.arrow.vector.types.Types; 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; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestValidateVector { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } private static final Charset utf8Charset = Charset.forName("UTF-8"); private static final byte[] STR1 = "AAAAA1".getBytes(utf8Charset); private static final byte[] STR2 = "BBBBBBBBB2".getBytes(utf8Charset); private static final byte[] STR3 = "CCCC3".getBytes(utf8Charset); @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testBaseFixedWidthVector() { try (final IntVector vector = new IntVector("v", allocator)) { validate(vector); setVector(vector, 1, 2, 3); validate(vector); vector.getDataBuffer().capacity(0); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> validate(vector)); assertTrue(e.getMessage().contains("Not enough capacity for fixed width data buffer")); } } @Test public void testBaseVariableWidthVector() { try (final VarCharVector vector = new VarCharVector("v", allocator)) { validate(vector); setVector(vector, STR1, STR2, STR3); validate(vector); vector.getDataBuffer().capacity(0); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> validate(vector)); assertTrue(e.getMessage().contains("Not enough capacity for data buffer")); } } @Test public void testBaseLargeVariableWidthVector() { try (final LargeVarCharVector vector = new LargeVarCharVector("v", allocator)) { validate(vector); setVector(vector, STR1, STR2, null, STR3); validate(vector); vector.getDataBuffer().capacity(0); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> validate(vector)); assertTrue(e.getMessage().contains("Not enough capacity for data buffer")); } } @Test public void testListVector() { try (final ListVector vector = ListVector.empty("v", allocator)) { validate(vector); setVector(vector, Arrays.asList(1, 2, 3), Arrays.asList(4, 5)); validate(vector); vector.getDataVector().setValueCount(3); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> validate(vector)); assertTrue(e.getMessage().contains("Inner vector does not contain enough elements.")); } } @Test public void testLargeListVector() { try (final LargeListVector vector = LargeListVector.empty("v", allocator)) { validate(vector); setVector(vector, Arrays.asList(1, 2, 3, 4), Arrays.asList(5, 6)); validate(vector); vector.getDataVector().setValueCount(4); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> validate(vector)); assertTrue(e.getMessage().contains("Inner vector does not contain enough elements.")); } } @Test public void testFixedSizeListVector() { try (final FixedSizeListVector vector = FixedSizeListVector.empty("v", 3, allocator)) { validate(vector); setVector(vector, Arrays.asList(1, 2, 3), Arrays.asList(4, 5, 6)); validate(vector); vector.getDataVector().setValueCount(3); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> validate(vector)); assertTrue(e.getMessage().contains("Inner vector does not contain enough elements.")); } } @Test public void testStructVectorRangeEquals() { try (final StructVector vector = StructVector.empty("struct", allocator)) { vector.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); vector.addOrGet("f1", FieldType.nullable(new ArrowType.Int(64, true)), BigIntVector.class); validate(vector); NullableStructWriter writer = vector.getWriter(); writer.allocate(); writeStructVector(writer, 1, 10L); writeStructVector(writer, 2, 20L); writeStructVector(writer, 3, 30L); writeStructVector(writer, 4, 40L); writeStructVector(writer, 5, 50L); writer.setValueCount(5); vector.getChild("f0").setValueCount(2); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> validate(vector)); assertTrue(e.getMessage().contains("Struct vector length not equal to child vector length")); vector.getChild("f0").setValueCount(5); validate(vector); vector.getChild("f0").getDataBuffer().capacity(0); ValidateUtil.ValidateException e2 = assertThrows(ValidateUtil.ValidateException.class, () -> validate(vector)); assertTrue(e2.getMessage().contains("Not enough capacity for fixed width data buffer")); } } @Test public void testUnionVector() { try (final UnionVector vector = UnionVector.empty("union", allocator)) { validate(vector); final NullableFloat4Holder float4Holder = new NullableFloat4Holder(); float4Holder.value = 1.01f; float4Holder.isSet = 1; final NullableFloat8Holder float8Holder = new NullableFloat8Holder(); float8Holder.value = 2.02f; float8Holder.isSet = 1; vector.setType(0, Types.MinorType.FLOAT4); vector.setSafe(0, float4Holder); vector.setType(1, Types.MinorType.FLOAT8); vector.setSafe(1, float8Holder); vector.setValueCount(2); validate(vector); vector.getChildrenFromFields().get(0).setValueCount(1); ValidateUtil.ValidateException e1 = assertThrows(ValidateUtil.ValidateException.class, () -> validate(vector)); assertTrue(e1.getMessage().contains("Union vector length not equal to child vector length")); vector.getChildrenFromFields().get(0).setValueCount(2); validate(vector); vector.getChildrenFromFields().get(0).getDataBuffer().capacity(0); ValidateUtil.ValidateException e2 = assertThrows(ValidateUtil.ValidateException.class, () -> validate(vector)); assertTrue(e2.getMessage().contains("Not enough capacity for fixed width data buffer")); } } @Test public void testDenseUnionVector() { try (final DenseUnionVector vector = DenseUnionVector.empty("union", allocator)) { validate(vector); final NullableFloat4Holder float4Holder = new NullableFloat4Holder(); float4Holder.value = 1.01f; float4Holder.isSet = 1; final NullableFloat8Holder float8Holder = new NullableFloat8Holder(); float8Holder.value = 2.02f; float8Holder.isSet = 1; byte float4TypeId = vector.registerNewTypeId(Field.nullable("", Types.MinorType.FLOAT4.getType())); byte float8TypeId = vector.registerNewTypeId(Field.nullable("", Types.MinorType.FLOAT8.getType())); vector.setTypeId(0, float4TypeId); vector.setSafe(0, float4Holder); vector.setTypeId(1, float8TypeId); vector.setSafe(1, float8Holder); vector.setValueCount(2); validate(vector); vector.getChildrenFromFields().get(0).getDataBuffer().capacity(0); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> validate(vector)); assertTrue(e.getMessage().contains("Not enough capacity for fixed width data buffer")); } } @Test public void testBaseFixedWidthVectorInstanceMethod() { try (final IntVector vector = new IntVector("v", allocator)) { vector.validate(); setVector(vector, 1, 2, 3); vector.validate(); vector.getDataBuffer().capacity(0); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> vector.validate()); assertTrue(e.getMessage().contains("Not enough capacity for fixed width data buffer")); } } @Test public void testRunEndEncodedVector() { final FieldType valueType = FieldType.notNullable(Types.MinorType.BIGINT.getType()); final FieldType runEndType = FieldType.notNullable(Types.MinorType.INT.getType()); final Field valueField = new Field("value", valueType, null); final Field runEndField = new Field("ree", runEndType, null); try (RunEndEncodedVector vector = new RunEndEncodedVector( new Field( "ree", FieldType.notNullable(RunEndEncoded.INSTANCE), List.of(runEndField, valueField)), allocator, null)) { vector.validate(); int runCount = 1; vector.allocateNew(); ((BigIntVector) vector.getValuesVector()).set(0, 1); ((IntVector) vector.getRunEndsVector()).set(0, 10); vector.getValuesVector().setValueCount(runCount); vector.getRunEndsVector().setValueCount(runCount); vector.setValueCount(10); vector.validate(); vector.getRunEndsVector().setValueCount(0); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> vector.validate()); assertTrue(e.getMessage().contains("Run end vector does not contain enough elements")); } } private void writeStructVector(NullableStructWriter writer, int value1, long value2) { writer.start(); writer.integer("f0").writeInt(value1); writer.bigInt("f1").writeBigInt(value2); writer.end(); } } TestValidateVectorFull.java000066400000000000000000000273401515451615200345150ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/validate/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.validate; import static org.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector; import static org.apache.arrow.vector.util.ValueVectorUtility.validateFull; 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.nio.charset.StandardCharsets; import java.util.Arrays; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.Decimal256Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.LargeVarCharVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.LargeListVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.holders.NullableFloat4Holder; import org.apache.arrow.vector.holders.NullableFloat8Holder; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; 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.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestValidateVectorFull { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testBaseVariableWidthVector() { try (final VarCharVector vector = new VarCharVector("v", allocator)) { validateFull(vector); setVector(vector, "aaa", "bbb", "ccc"); validateFull(vector); ArrowBuf offsetBuf = vector.getOffsetBuffer(); offsetBuf.setInt(0, 100); offsetBuf.setInt(4, 50); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> validateFull(vector)); assertTrue( e.getMessage() .contains("The values in positions 0 and 1 of the offset buffer are decreasing")); } } @Test public void testBaseLargeVariableWidthVector() { try (final LargeVarCharVector vector = new LargeVarCharVector("v", allocator)) { validateFull(vector); setVector(vector, "aaa", "bbb", null, "ccc"); validateFull(vector); ArrowBuf offsetBuf = vector.getOffsetBuffer(); offsetBuf.setLong(0, 100); offsetBuf.setLong(8, 50); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> validateFull(vector)); assertTrue( e.getMessage() .contains( "The values in positions 0 and 1 of the large offset buffer are decreasing")); } } @Test public void testListVector() { try (final ListVector vector = ListVector.empty("v", allocator)) { validateFull(vector); setVector(vector, Arrays.asList(1, 2, 3), Arrays.asList(4, 5), Arrays.asList(6, 7, 8, 9)); validateFull(vector); ArrowBuf offsetBuf = vector.getOffsetBuffer(); offsetBuf.setInt(0, 100); offsetBuf.setInt(8, 50); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> validateFull(vector)); assertTrue( e.getMessage() .contains("The values in positions 0 and 1 of the offset buffer are decreasing")); } } @Test public void testLargeListVector() { try (final LargeListVector vector = LargeListVector.empty("v", allocator)) { validateFull(vector); setVector(vector, Arrays.asList(1, 2, 3), Arrays.asList(4, 5), Arrays.asList(6, 7, 8, 9)); validateFull(vector); ArrowBuf offsetBuf = vector.getOffsetBuffer(); offsetBuf.setLong(0, 100); offsetBuf.setLong(16, 50); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> validateFull(vector)); assertTrue( e.getMessage() .contains( "The values in positions 0 and 1 of the large offset buffer are decreasing")); } } @Test public void testStructVectorRangeEquals() { try (final StructVector vector = StructVector.empty("struct", allocator)) { IntVector intVector = vector.addOrGet("f0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class); VarCharVector strVector = vector.addOrGet("f1", FieldType.nullable(new ArrowType.Utf8()), VarCharVector.class); validateFull(vector); validateFull(intVector); validateFull(strVector); ValueVectorDataPopulator.setVector(intVector, 1, 2, 3, 4, 5); ValueVectorDataPopulator.setVector(strVector, "a", "b", "c", "d", "e"); vector.setValueCount(5); validateFull(vector); validateFull(intVector); validateFull(strVector); ArrowBuf offsetBuf = strVector.getOffsetBuffer(); offsetBuf.setInt(0, 100); offsetBuf.setInt(8, 50); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> validateFull(strVector)); assertTrue( e.getMessage() .contains("The values in positions 0 and 1 of the offset buffer are decreasing")); e = assertThrows(ValidateUtil.ValidateException.class, () -> validateFull(vector)); assertTrue( e.getMessage() .contains("The values in positions 0 and 1 of the offset buffer are decreasing")); } } @Test public void testUnionVector() { try (final UnionVector vector = UnionVector.empty("union", allocator)) { validateFull(vector); final NullableFloat4Holder float4Holder = new NullableFloat4Holder(); float4Holder.value = 1.01f; float4Holder.isSet = 1; final NullableFloat8Holder float8Holder = new NullableFloat8Holder(); float8Holder.value = 2.02f; float8Holder.isSet = 1; vector.setType(0, Types.MinorType.FLOAT4); vector.setSafe(0, float4Holder); vector.setType(1, Types.MinorType.FLOAT8); vector.setSafe(1, float8Holder); vector.setValueCount(2); validateFull(vector); // negative type id vector.getTypeBuffer().setByte(0, -1); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> validateFull(vector)); assertTrue(e.getMessage().contains("The type id at position 0 is negative")); } } @Test public void testDenseUnionVector() { try (final DenseUnionVector vector = DenseUnionVector.empty("union", allocator)) { validateFull(vector); final NullableFloat4Holder float4Holder = new NullableFloat4Holder(); float4Holder.value = 1.01f; float4Holder.isSet = 1; final NullableFloat8Holder float8Holder = new NullableFloat8Holder(); float8Holder.value = 2.02f; float8Holder.isSet = 1; byte float4TypeId = vector.registerNewTypeId(Field.nullable("", Types.MinorType.FLOAT4.getType())); byte float8TypeId = vector.registerNewTypeId(Field.nullable("", Types.MinorType.FLOAT8.getType())); vector.setTypeId(0, float4TypeId); vector.setSafe(0, float4Holder); vector.setTypeId(1, float8TypeId); vector.setSafe(1, float8Holder); vector.setValueCount(2); validateFull(vector); ValueVector subVector = vector.getVectorByType(float4TypeId); assertTrue(subVector instanceof Float4Vector); assertEquals(1, subVector.getValueCount()); // shrink sub-vector subVector.setValueCount(0); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> validateFull(vector)); assertTrue(e.getMessage().contains("Dense union vector offset exceeds sub-vector boundary")); } } @Test public void testBaseVariableWidthVectorInstanceMethod() { try (final VarCharVector vector = new VarCharVector("v", allocator)) { vector.validateFull(); setVector(vector, "aaa", "bbb", "ccc"); vector.validateFull(); ArrowBuf offsetBuf = vector.getOffsetBuffer(); offsetBuf.setInt(0, 100); offsetBuf.setInt(4, 50); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, vector::validateFull); assertTrue( e.getMessage() .contains("The values in positions 0 and 1 of the offset buffer are decreasing")); } } @Test public void testValidateVarCharUTF8() { try (final VarCharVector vector = new VarCharVector("v", allocator)) { vector.validateFull(); setVector( vector, "aaa".getBytes(StandardCharsets.UTF_8), "bbb".getBytes(StandardCharsets.UTF_8), new byte[] {(byte) 0xFF, (byte) 0xFE}); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, vector::validateFull); assertTrue(e.getMessage().contains("UTF")); } } @Test public void testValidateLargeVarCharUTF8() { try (final LargeVarCharVector vector = new LargeVarCharVector("v", allocator)) { vector.validateFull(); setVector( vector, "aaa".getBytes(StandardCharsets.UTF_8), "bbb".getBytes(StandardCharsets.UTF_8), new byte[] {(byte) 0xFF, (byte) 0xFE}); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, vector::validateFull); assertTrue(e.getMessage().contains("UTF")); } } @Test public void testValidateDecimal() { try (final DecimalVector vector = new DecimalVector( Field.nullable("v", new ArrowType.Decimal(2, 0, DecimalVector.TYPE_WIDTH * 8)), allocator)) { vector.validateFull(); setVector(vector, 1L); vector.validateFull(); vector.clear(); setVector(vector, Long.MAX_VALUE); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, vector::validateFull); assertTrue(e.getMessage().contains("Decimal")); } } @Test public void testValidateDecimal256() { try (final Decimal256Vector vector = new Decimal256Vector( Field.nullable("v", new ArrowType.Decimal(2, 0, DecimalVector.TYPE_WIDTH * 8)), allocator)) { vector.validateFull(); setVector(vector, 1L); vector.validateFull(); vector.clear(); setVector(vector, Long.MAX_VALUE); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, vector::validateFull); assertTrue(e.getMessage().contains("Decimal")); } } } TestValidateVectorSchemaRoot.java000066400000000000000000000074641515451615200356640ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/validate/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.validate; import static org.apache.arrow.vector.util.ValueVectorUtility.validate; import static org.apache.arrow.vector.util.ValueVectorUtility.validateFull; 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.vector.IntVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class TestValidateVectorSchemaRoot { private BufferAllocator allocator; @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void terminate() throws Exception { allocator.close(); } @Test public void testValidatePositive() { try (IntVector intVector = new IntVector("int vector", allocator); VarCharVector strVector = new VarCharVector("var char vector", allocator)) { VectorSchemaRoot root = VectorSchemaRoot.of(intVector, strVector); validate(root); validateFull(root); ValueVectorDataPopulator.setVector(intVector, 1, 2, 3, 4, 5); ValueVectorDataPopulator.setVector(strVector, "a", "b", "c", "d", "e"); root.setRowCount(5); validate(root); validateFull(root); } } @Test public void testValidateNegative() { try (IntVector intVector = new IntVector("int vector", allocator); VarCharVector strVector = new VarCharVector("var char vector", allocator)) { VectorSchemaRoot root = VectorSchemaRoot.of(intVector, strVector); ValueVectorDataPopulator.setVector(intVector, 1, 2, 3, 4, 5); ValueVectorDataPopulator.setVector(strVector, "a", "b", "c", "d", "e"); // validate mismatching value counts root.setRowCount(4); intVector.setValueCount(5); strVector.setValueCount(5); ValidateUtil.ValidateException e = assertThrows(ValidateUtil.ValidateException.class, () -> validate(root)); assertTrue( e.getMessage() .contains("Child vector and vector schema root have different value counts")); e = assertThrows(ValidateUtil.ValidateException.class, () -> validateFull(root)); assertTrue( e.getMessage() .contains("Child vector and vector schema root have different value counts")); // valid problems with the child vector root.setRowCount(5); ArrowBuf offsetBuf = strVector.getOffsetBuffer(); offsetBuf.setInt(0, 100); offsetBuf.setInt(8, 50); validate(root); e = assertThrows(ValidateUtil.ValidateException.class, () -> validateFull(root)); assertTrue( e.getMessage() .contains("The values in positions 0 and 1 of the offset buffer are decreasing")); } } } TestValidateVectorTypeVisitor.java000066400000000000000000000363671515451615200361250ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/java/org/apache/arrow/vector/validate/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS 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.validate; import static org.junit.jupiter.api.Assertions.assertThrows; import java.util.function.Supplier; 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.DateMilliVector; import org.apache.arrow.vector.Decimal256Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; 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.IntervalDayVector; import org.apache.arrow.vector.IntervalYearVector; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.LargeVarCharVector; 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.ValueVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.complex.DenseUnionVector; 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.StructVector; import org.apache.arrow.vector.complex.UnionVector; 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.FieldType; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; /** Test cases for {@link ValidateVectorTypeVisitor}. */ public class TestValidateVectorTypeVisitor { private BufferAllocator allocator; private ValidateVectorTypeVisitor visitor = new ValidateVectorTypeVisitor(); @BeforeEach public void init() { allocator = new RootAllocator(Long.MAX_VALUE); } @AfterEach public void terminate() throws Exception { allocator.close(); } private void testPositiveCase(Supplier vectorGenerator) { try (ValueVector vector = vectorGenerator.get(); ) { vector.accept(visitor, null); } } private void testNegativeCase(Supplier vectorGenerator) { try (ValueVector vector = vectorGenerator.get()) { assertThrows( ValidateUtil.ValidateException.class, () -> { vector.accept(visitor, null); }); } } @Test public void testFixedWidthVectorsPositive() { // integer vectors testPositiveCase(() -> new TinyIntVector("vector", allocator)); testPositiveCase(() -> new SmallIntVector("vector", allocator)); testPositiveCase(() -> new IntVector("vector", allocator)); testPositiveCase(() -> new BigIntVector("vector", allocator)); testPositiveCase(() -> new UInt1Vector("vector", allocator)); testPositiveCase(() -> new UInt2Vector("vector", allocator)); testPositiveCase(() -> new UInt4Vector("vector", allocator)); testPositiveCase(() -> new UInt8Vector("vector", allocator)); testPositiveCase(() -> new BitVector("vector", allocator)); testPositiveCase(() -> new DecimalVector("vector", allocator, 30, 16)); // date vectors testPositiveCase(() -> new DateDayVector("vector", allocator)); testPositiveCase(() -> new DateMilliVector("vector", allocator)); testPositiveCase( () -> new DurationVector( "vector", FieldType.nullable(new ArrowType.Duration(TimeUnit.SECOND)), allocator)); // float vectors testPositiveCase(() -> new Float4Vector("vector", allocator)); testPositiveCase(() -> new Float8Vector("vector", allocator)); // interval vectors testPositiveCase(() -> new IntervalDayVector("vector", allocator)); testPositiveCase(() -> new IntervalYearVector("vector", allocator)); // time vectors testPositiveCase(() -> new TimeMicroVector("vector", allocator)); testPositiveCase(() -> new TimeMilliVector("vector", allocator)); testPositiveCase(() -> new TimeMicroVector("vector", allocator)); testPositiveCase(() -> new TimeSecVector("vector", allocator)); // time stamp vectors testPositiveCase(() -> new TimeStampMicroTZVector("vector", allocator, "cn")); testPositiveCase(() -> new TimeStampMicroVector("vector", allocator)); testPositiveCase(() -> new TimeStampMilliTZVector("vector", allocator, "cn")); testPositiveCase(() -> new TimeStampMilliVector("vector", allocator)); testPositiveCase(() -> new TimeStampNanoTZVector("vector", allocator, "cn")); testPositiveCase(() -> new TimeStampNanoVector("vector", allocator)); testPositiveCase(() -> new TimeStampSecTZVector("vector", allocator, "cn")); testPositiveCase(() -> new TimeStampSecVector("vector", allocator)); testPositiveCase(() -> new FixedSizeBinaryVector("vector", allocator, 5)); } @Test public void testFixedWidthVectorsNegative() { // integer vectors testNegativeCase( () -> new TinyIntVector( "vector", FieldType.nullable(Types.MinorType.INT.getType()), allocator)); testNegativeCase( () -> new SmallIntVector( "vector", FieldType.nullable(Types.MinorType.INT.getType()), allocator)); testNegativeCase( () -> new BigIntVector( "vector", FieldType.nullable(Types.MinorType.SMALLINT.getType()), allocator)); testNegativeCase( () -> new BigIntVector( "vector", FieldType.nullable(Types.MinorType.SMALLINT.getType()), allocator)); testNegativeCase( () -> new UInt1Vector( "vector", FieldType.nullable(Types.MinorType.SMALLINT.getType()), allocator)); testNegativeCase( () -> new UInt2Vector( "vector", FieldType.nullable(Types.MinorType.SMALLINT.getType()), allocator)); testNegativeCase( () -> new UInt4Vector( "vector", FieldType.nullable(Types.MinorType.SMALLINT.getType()), allocator)); testNegativeCase( () -> new UInt8Vector( "vector", FieldType.nullable(Types.MinorType.SMALLINT.getType()), allocator)); testNegativeCase( () -> new BitVector( "vector", FieldType.nullable(Types.MinorType.BIGINT.getType()), allocator)); testNegativeCase(() -> new DecimalVector("vector", allocator, 30, -16)); // date vectors testNegativeCase( () -> new DateDayVector( "vector", FieldType.nullable(Types.MinorType.FLOAT4.getType()), allocator)); testNegativeCase( () -> new DateMilliVector( "vector", FieldType.nullable(Types.MinorType.BIGINT.getType()), allocator)); // float pont vectors testNegativeCase( () -> new Float4Vector( "vector", FieldType.nullable(Types.MinorType.BIGINT.getType()), allocator)); testNegativeCase( () -> new Float8Vector( "vector", FieldType.nullable(Types.MinorType.BIGINT.getType()), allocator)); // interval vectors testNegativeCase( () -> new IntervalDayVector( "vector", FieldType.nullable(Types.MinorType.INT.getType()), allocator)); testNegativeCase( () -> new IntervalYearVector( "vector", FieldType.nullable(Types.MinorType.BIGINT.getType()), allocator)); // time vectors testNegativeCase( () -> new TimeMilliVector( "vector", FieldType.nullable(Types.MinorType.BIGINT.getType()), allocator)); testNegativeCase( () -> new TimeMicroVector( "vector", FieldType.nullable(Types.MinorType.BIGINT.getType()), allocator)); testNegativeCase( () -> new TimeNanoVector( "vector", FieldType.nullable(Types.MinorType.BIGINT.getType()), allocator)); testNegativeCase( () -> new TimeSecVector( "vector", FieldType.nullable(Types.MinorType.BIGINT.getType()), allocator)); // time stamp vectors testNegativeCase(() -> new TimeStampMicroTZVector("vector", allocator, null)); testNegativeCase( () -> new TimeStampMicroVector( "vector", FieldType.nullable(Types.MinorType.BIGINT.getType()), allocator)); testNegativeCase(() -> new TimeStampMilliTZVector("vector", allocator, null)); testNegativeCase( () -> new TimeStampMilliVector( "vector", FieldType.nullable(Types.MinorType.BIGINT.getType()), allocator)); testNegativeCase(() -> new TimeStampNanoTZVector("vector", allocator, null)); testNegativeCase( () -> new TimeStampNanoVector( "vector", FieldType.nullable(Types.MinorType.BIGINT.getType()), allocator)); testNegativeCase(() -> new TimeStampSecTZVector("vector", allocator, null)); testNegativeCase( () -> new TimeStampSecVector( "vector", FieldType.nullable(Types.MinorType.BIGINT.getType()), allocator)); } @Test public void testDecimalVector() { testPositiveCase( () -> new DecimalVector( "dec", FieldType.nullable(ArrowType.Decimal.createDecimal(10, 10, 128)), allocator)); testPositiveCase( () -> new DecimalVector( "dec", FieldType.nullable(ArrowType.Decimal.createDecimal(38, 10, 128)), allocator)); testPositiveCase( () -> new Decimal256Vector( "dec", FieldType.nullable(ArrowType.Decimal.createDecimal(50, 10, 256)), allocator)); testPositiveCase( () -> new Decimal256Vector( "dec", FieldType.nullable(ArrowType.Decimal.createDecimal(76, 10, 256)), allocator)); testNegativeCase( () -> new DecimalVector( "dec", FieldType.nullable(ArrowType.Decimal.createDecimal(50, 10, 128)), allocator)); testNegativeCase( () -> new Decimal256Vector( "dec", FieldType.nullable(ArrowType.Decimal.createDecimal(100, 10, 256)), allocator)); testNegativeCase( () -> new DecimalVector( "dec", FieldType.nullable(ArrowType.Decimal.createDecimal(0, 10, 128)), allocator)); testNegativeCase( () -> new Decimal256Vector( "dec", FieldType.nullable(ArrowType.Decimal.createDecimal(-1, 10, 256)), allocator)); testNegativeCase( () -> new Decimal256Vector( "dec", FieldType.nullable(ArrowType.Decimal.createDecimal(30, 10, 64)), allocator)); testNegativeCase( () -> new Decimal256Vector( "dec", FieldType.nullable(ArrowType.Decimal.createDecimal(10, 20, 256)), allocator)); } @Test public void testVariableWidthVectorsPositive() { testPositiveCase(() -> new VarCharVector("vector", allocator)); testPositiveCase(() -> new VarBinaryVector("vector", allocator)); } @Test public void testVariableWidthVectorsNegative() { testNegativeCase( () -> new VarCharVector( "vector", FieldType.nullable(Types.MinorType.INT.getType()), allocator)); testNegativeCase( () -> new VarBinaryVector( "vector", FieldType.nullable(Types.MinorType.INT.getType()), allocator)); } @Test public void testLargeVariableWidthVectorsPositive() { testPositiveCase(() -> new LargeVarCharVector("vector", allocator)); testPositiveCase(() -> new LargeVarBinaryVector("vector", allocator)); } @Test public void testLargeVariableWidthVectorsNegative() { testNegativeCase( () -> new LargeVarCharVector( "vector", FieldType.nullable(Types.MinorType.INT.getType()), allocator)); testNegativeCase( () -> new LargeVarBinaryVector( "vector", FieldType.nullable(Types.MinorType.INT.getType()), allocator)); } @Test public void testListVector() { testPositiveCase(() -> ListVector.empty("vector", allocator)); testNegativeCase( () -> new ListVector( "vector", allocator, FieldType.nullable(Types.MinorType.INT.getType()), null)); } @Test public void testLargeListVector() { testPositiveCase(() -> LargeListVector.empty("vector", allocator)); testNegativeCase( () -> new LargeListVector( "vector", allocator, FieldType.nullable(Types.MinorType.INT.getType()), null)); } @Test public void testFixedSizeListVector() { testPositiveCase(() -> FixedSizeListVector.empty("vector", 10, allocator)); } @Test public void testStructVector() { testPositiveCase(() -> StructVector.empty("vector", allocator)); testNegativeCase( () -> new StructVector( "vector", allocator, FieldType.nullable(Types.MinorType.INT.getType()), null)); } @Test public void testUnionVector() { testPositiveCase(() -> UnionVector.empty("vector", allocator)); } @Test public void testDenseUnionVector() { testPositiveCase(() -> DenseUnionVector.empty("vector", allocator)); } @Test public void testNullVector() { testPositiveCase(() -> new NullVector("null vec")); } } apache-arrow-java-19.0.0/vector/src/test/resources/000077500000000000000000000000001515451615200221605ustar00rootroot00000000000000apache-arrow-java-19.0.0/vector/src/test/resources/integration_json_struct.json000066400000000000000000000075761515451615200300520ustar00rootroot00000000000000{ "schema": { "fields": [ { "name": "struct_nullable", "type": { "name": "struct" }, "nullable": true, "children": [ { "name": "f1", "type": { "name": "int", "isSigned": true, "bitWidth": 32 }, "nullable": true, "children": [] }, { "name": "f2", "type": { "name": "utf8" }, "nullable": true, "children": [] } ] } ] }, "batches": [ { "count": 7, "columns": [ { "name": "struct_nullable", "count": 7, "VALIDITY": [ 0, 1, 1, 1, 0, 1, 0 ], "children": [ { "name": "f1", "count": 7, "VALIDITY": [ 1, 0, 1, 1, 1, 0, 0 ], "DATA": [ 1402032511, 290876774, 137773603, 410361374, 1959836418, 1995074679, -163525262 ] }, { "name": "f2", "count": 7, "VALIDITY": [ 0, 1, 1, 1, 0, 1, 0 ], "OFFSET": [ 0, 0, 7, 14, 21, 21, 28, 28 ], "DATA": [ "", "MhRNxD4", "3F9HBxK", "aVd88fp", "", "3loZrRf", "" ] } ] } ] }, { "count": 10, "columns": [ { "name": "struct_nullable", "count": 10, "VALIDITY": [ 0, 1, 1, 0, 1, 0, 0, 1, 1, 1 ], "children": [ { "name": "f1", "count": 10, "VALIDITY": [ 0, 0, 0, 0, 0, 0, 1, 0, 0, 0 ], "DATA": [ -2041500147, 1715692943, -35444996, 1425496657, 112765084, 1760754983, 413888857, 2039738337, -1924327700, 670528518 ] }, { "name": "f2", "count": 10, "VALIDITY": [ 1, 0, 0, 1, 1, 1, 1, 1, 1, 0 ], "OFFSET": [ 0, 7, 7, 7, 14, 21, 28, 35, 42, 49, 49 ], "DATA": [ "AS5oARE", "", "", "JGdagcX", "78SLiRw", "vbGf7OY", "5uh5fTs", "0ilsf82", "LjS9MbU", "" ] } ] } ] } ] } apache-arrow-java-19.0.0/vector/src/test/resources/logback.xml000066400000000000000000000025221515451615200243050ustar00rootroot00000000000000 %d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n